1013 Add more operators to QueryString (#1014)

* more quertstring operators * add modulus method * add tests * add more tests
piccolo-orm · Jun 10, 2024 · cfb674a · cfb674a
1 parent 28f871d
commit cfb674a
Show file tree

Hide file tree

Showing 2 changed files with 154 additions and 0 deletions.
diff --git 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
@@ -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]]),
+        )