diff --git a/piccolo/querystring.py b/piccolo/querystring.py index c2fc0f80f..1f7282a7d 100644 --- a/piccolo/querystring.py +++ b/piccolo/querystring.py @@ -282,12 +282,30 @@ def __lt__(self, value) -> QueryString: def __le__(self, value) -> QueryString: return QueryString("{} <= {}", self, value) + def __truediv__(self, value) -> QueryString: + return QueryString("{} / {}", self, value) + + def __mul__(self, value) -> QueryString: + return QueryString("{} * {}", self, value) + + def __pow__(self, value) -> QueryString: + return QueryString("{} ^ {}", self, value) + + def __mod__(self, value) -> QueryString: + return QueryString("{} % {}", self, value) + def is_in(self, value) -> QueryString: return QueryString("{} IN {}", self, value) def not_in(self, value) -> QueryString: return QueryString("{} NOT IN {}", self, value) + def like(self, value: str) -> QueryString: + return QueryString("{} LIKE {}", self, value) + + def ilike(self, value: str) -> QueryString: + return QueryString("{} ILIKE {}", self, value) + class Unquoted(QueryString): """ diff --git a/tests/query/test_querystring.py b/tests/query/test_querystring.py index 59631ab7f..58ca29495 100644 --- a/tests/query/test_querystring.py +++ b/tests/query/test_querystring.py @@ -1,6 +1,7 @@ from unittest import TestCase from piccolo.querystring import QueryString +from tests.base import postgres_only # TODO - add more extensive tests (increased nesting and argument count). @@ -28,3 +29,138 @@ def test_string(self): def test_querystring_with_no_args(self): qs = QueryString("SELECT name FROM band") self.assertEqual(qs.compile_string(), ("SELECT name FROM band", [])) + + +@postgres_only +class TestQueryStringOperators(TestCase): + """ + Make sure basic operations can be used on ``QueryString``. + """ + + def test_add(self): + query = QueryString("SELECT price") + 1 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price + $1", [1]), + ) + + def test_multiply(self): + query = QueryString("SELECT price") * 2 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price * $1", [2]), + ) + + def test_divide(self): + query = QueryString("SELECT price") / 1 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price / $1", [1]), + ) + + def test_power(self): + query = QueryString("SELECT price") ** 2 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price ^ $1", [2]), + ) + + def test_subtract(self): + query = QueryString("SELECT price") - 1 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price - $1", [1]), + ) + + def test_modulus(self): + query = QueryString("SELECT price") % 1 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price % $1", [1]), + ) + + def test_like(self): + query = QueryString("strip(name)").like("Python%") + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("strip(name) LIKE $1", ["Python%"]), + ) + + def test_ilike(self): + query = QueryString("strip(name)").ilike("Python%") + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("strip(name) ILIKE $1", ["Python%"]), + ) + + def test_greater_than(self): + query = QueryString("SELECT price") > 10 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price > $1", [10]), + ) + + def test_greater_equal_than(self): + query = QueryString("SELECT price") >= 10 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price >= $1", [10]), + ) + + def test_less_than(self): + query = QueryString("SELECT price") < 10 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price < $1", [10]), + ) + + def test_less_equal_than(self): + query = QueryString("SELECT price") <= 10 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price <= $1", [10]), + ) + + def test_equals(self): + query = QueryString("SELECT price") == 10 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price = $1", [10]), + ) + + def test_not_equals(self): + query = QueryString("SELECT price") != 10 + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price != $1", [10]), + ) + + def test_is_in(self): + query = QueryString("SELECT price").is_in([10, 20, 30]) + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price IN $1", [[10, 20, 30]]), + ) + + def test_not_in(self): + query = QueryString("SELECT price").not_in([10, 20, 30]) + self.assertIsInstance(query, QueryString) + self.assertEqual( + query.compile_string(), + ("SELECT price NOT IN $1", [[10, 20, 30]]), + )