finish implementing the query builder

docker-compose.yml

@@ -2,7 +2,7 @@ version: '3.3'
 
 services:
   db:
-    image: juxt/xtdb-in-memory:1.20.0
+    image: juxt/xtdb-in-memory:1.22.1
   py:
     build: .
     volumes:

pyxtdb.py

@@ -26,47 +26,153 @@ def __init__(self, desc, code):
         self.description = desc
         self.status_code = code
 
+
+# Query Builder
+
+class Symbol(edn_format.Symbol):
+    pass
+
+class Keyword(edn_format.Keyword):
+    pass
+
+class Char(edn_format.Char):
+    pass
+
 class Query:
     def __init__(self, uri="http://localhost:3000", node=None):
         if node:
             self.node = node
         else:
             self.node = Node(uri)
         self._find_clause = None
+        self._in_clause = None
         self._where_clauses = []
+        self._rules_clauses = []
+        self._order_by_clauses = []
+        self._limit = None
+        self._offset = None
         self._values = None
+        self._timeout = None
+        self._in_args = None
         self.error = None
 
-    def find(self, clause):
+    def find(self, *args):
         if self._values:
             raise AlreadySent()
-        self._find_clause = clause
+        if len(args) == 0:
+            raise ValueError('missing find arguments')
+        if len(args) == 1 and type(args[0]) == str:
+            self._find_clause = edn_format.loads('['+args[0]+']')
+        else:
+            self._find_clause = list(args)
         return self
 
-    def where(self, clause):
+    def in_(self, *args):
         if self._values:
             raise AlreadySent()
-        self._where_clauses.append(clause)
+        if len(args) == 0:
+            raise ValueError('missing in arguments')
+        if len(args) == 1 and type(args[0]) == str:
+            self._in_clause = edn_format.loads('['+args[0]+']')
+        else:
+            self._in_clause = list(args)
         return self
 
+    def where(self, *args):
+        if self._values:
+            raise AlreadySent()
+        if len(args) == 0:
+            raise ValueError('missing where arguments')
+        if len(args) == 1 and type(args[0]) == str:
+            self._where_clauses.append(edn_format.loads('['+args[0]+']'))
+        else:
+            self._where_clauses.append(list(args))
+        return self
+
+    def rules(self, *args):
+        if self._values:
+            raise AlreadySent()
+        if len(args) == 0:
+            raise ValueError('missing rules arguments')
+        if len(args) == 1 and type(args[0]) == str:
+            self._rules_clauses.append(edn_format.loads('['+args[0]+']'))
+        else:
+            self._rules_clauses.append(list(args))
+        return self
+
+    def order_by(self, *args):
+        if self._values:
+            raise AlreadySent()
+        if len(args) == 0:
+            raise ValueError('missing order_by arguments')
+        if len(args) == 1 and type(args[0]) == str:
+            self._order_by_clauses.append(edn_format.loads('['+args[0]+']'))
+        else:
+            self._order_by_clauses.append(list(args))
+        return self
+
+    def limit(self, limit):
+        if self._values:
+            raise AlreadySent()
+        self._limit = limit
+        return self
+
+    def offset(self, offset):
+        if self._values:
+            raise AlreadySent()
+        self._offset = offset
+        return self
+
+    def timeout(self, timeout):
+        if self._values:
+            raise AlreadySent()
+        self._timeout = timeout
+        return self
+
+    def in_args(self, *args):
+        if self._values:
+            raise AlreadySent()
+        if len(args) == 0:
+            raise ValueError('missing in-args arguments')
+        self._in_args = list(args)
+        return self
+
+    def query(self):
+        q = {
+            Keyword('find'):  self._find_clause,
+            Keyword('where'): self._where_clauses
+        }
+        if self._in_clause:
+            q[Keyword('in')] = self._in_clause
+        if self._rules_clauses:
+            q[Keyword('rules')] = self._rules_clauses
+        if self._order_by_clauses:
+            q[Keyword('order-by')] = self._order_by_clauses
+        if self._limit:
+            q[Keyword('limit')] = self._limit
+        if self._offset:
+            q[Keyword('offset')] = self._offset
+        if self._timeout:
+            q[Keyword('timeout')] = self._timeout
+        return q
+
     def values(self):
+        if not self._find_clause:
+            raise BadQuery("query has no find clause")
         if not self._where_clauses:
-            raise BadQuery("No Where Clause")
-        _query = """
-        {
-            :find [%s]
-            :where [
-                [%s]
-            ]
-        }
-        """ % (self._find_clause, ']\n['.join(self._where_clauses))
-
-        result = self.node.query(_query)
+            raise BadQuery("query has no where clause")
+        query = self.query()
+        result = self.node.query(query, in_args=self._in_args)
         if type(result) is dict:
             self.error = json.dumps(result, indent=4, sort_keys=True)
             return []
         return result
 
+    def __str__(self):
+        v = {Keyword('query'): self.query(),
+             Keyword('in-args'): self._in_args}
+        return edn_format.dumps(v)
+
     def __iter__(self):
         self._values = self.values()
         return self

tests.py

@@ -2,6 +2,7 @@
 import pytest
 import os
 import pyxtdb
+from pyxtdb import Symbol, Keyword
 import edn_format
 
 XTDB_URL = os.environ.get('XTDB_URL', 'http://localhost:3000')
@@ -105,21 +106,7 @@ def test_query_with_args(billies_db):
                         in_args='[[["Billy" "Joel"] ["Billy" "Idol"]]]')
     assert len(result) == 2
 
-    # query can be edn parsed from a string
-    q = edn_format.loads('{:find [?e] :where [[?e :name first] [?e :last-name last]] :in [first last]}')
-    result = node.query(query=q, in_args=["Billy", "Joel"])
-    assert len(result) == 1
-
-    # or edn constructed by hand
-    result = node.query(
-        query= \
-        { edn_format.Keyword('find')  : [edn_format.Symbol('e')],
-          edn_format.Keyword('where') : [[edn_format.Symbol('e'), edn_format.Keyword('last-name'), edn_format.Symbol('name')]],
-          edn_format.Keyword('in')    : [edn_format.Symbol('name')] },
-        in_args=['Joel'])
-    assert len(result) == 1
-
-def test_query_model(billies_db):
+def test_query_model_strings(billies_db):
     node = billies_db
 
     # Fetch records with name "Billy"
@@ -132,6 +119,96 @@ def test_query_model(billies_db):
     assert len(list(result)) == 1
     assert result.error == None
 
+    # scalar argument
+    result = node.find('?e').where('?e :last-name last').in_('last').in_args('Joel')
+    assert str(result) == '{:query {:find [?e] :where [[?e :last-name last]] :in [last]} :in-args ["Joel"]}'
+    assert len(list(result)) == 1
+
+    # multiple scalars
+    result = node.find('?e') \
+                 .where('?e :name first') \
+                 .where('?e :last-name last') \
+                 .in_('first last') \
+                 .in_args('Billy', 'Joel')
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last]] :in [first last]} :in-args ["Billy" "Joel"]}'
+    assert len(list(result)) == 1
+
+    # tuple
+    result = node.find('?e') \
+                 .where('?e :name first') \
+                 .where('?e :last-name last') \
+                 .in_('[first last]') \
+                 .in_args(['Billy', 'Joel'])
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last]] :in [[first last]]} :in-args [["Billy" "Joel"]]}'
+    assert len(list(result)) == 1
+
+    # tuple and scalar
+    result = node.find('?e') \
+                 .where('?e :name first') \
+                 .where('?e :last-name last') \
+                 .where('?e :profession job') \
+                 .in_('[first last] job') \
+                 .in_args(['Billy', 'Joel'], 'singer')
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last] [?e :profession job]] :in [[first last] job]} :in-args [["Billy" "Joel"] "singer"]}'
+    assert len(list(result)) == 1
+
+    # collection
+    result = node.find('?e') \
+                 .where('?e :name first') \
+                 .where('?e :last-name last') \
+                 .in_('[[first last]]') \
+                 .in_args([['Billy', 'Joel'], ['Billy', 'Idol']])
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last]] :in [[[first last]]]} :in-args [[["Billy" "Joel"] ["Billy" "Idol"]]]}'
+    assert len(list(result)) == 2
+
+def test_query_model_edn(billies_db):
+    node = billies_db
+
+    # scalar argument
+    result = node.find(Symbol('?e')) \
+                 .where(Symbol('?e'), Keyword('last-name'), Symbol('last')) \
+                 .in_(Symbol('last')) \
+                 .in_args('Joel')
+    assert str(result) == '{:query {:find [?e] :where [[?e :last-name last]] :in [last]} :in-args ["Joel"]}'
+    assert len(list(result)) == 1
+
+    # multiple scalars
+    result = node.find(Symbol('?e')) \
+                 .where(Symbol('?e'), Keyword('name'), Symbol('first')) \
+                 .where(Symbol('?e'), Keyword('last-name'), Symbol('last')) \
+                 .in_(Symbol('first'), Symbol('last')) \
+                 .in_args('Billy', 'Joel')
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last]] :in [first last]} :in-args ["Billy" "Joel"]}'
+    assert len(list(result)) == 1
+
+    # tuple
+    result = node.find(Symbol('?e')) \
+                 .where(Symbol('?e'), Keyword('name'), Symbol('first')) \
+                 .where(Symbol('?e'), Keyword('last-name'), Symbol('last')) \
+                 .in_([Symbol('first'), Symbol('last')]) \
+                 .in_args(['Billy', 'Joel'])
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last]] :in [[first last]]} :in-args [["Billy" "Joel"]]}'
+    assert len(list(result)) == 1
+
+    # tuple and scalar
+    result = node.find(Symbol('?e')) \
+                 .where(Symbol('?e'), Keyword('name'), Symbol('first')) \
+                 .where(Symbol('?e'), Keyword('last-name'), Symbol('last')) \
+                 .where(Symbol('?e'), Keyword('profession'), Symbol('job')) \
+                 .in_([Symbol('first'), Symbol('last')], Symbol('job')) \
+                 .in_args(['Billy', 'Joel'], 'singer')
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last] [?e :profession job]] :in [[first last] job]} :in-args [["Billy" "Joel"] "singer"]}'
+    assert len(list(result)) == 1
+
+    # collection
+    result = node.find(Symbol('?e')) \
+                 .where(Symbol('?e'), Keyword('name'), Symbol('first')) \
+                 .where(Symbol('?e'), Keyword('last-name'), Symbol('last')) \
+                 .in_([[Symbol('first'), Symbol('last')]]) \
+                 .in_args([['Billy', 'Joel'], ['Billy', 'Idol']])
+    assert str(result) == '{:query {:find [?e] :where [[?e :name first] [?e :last-name last]] :in [[[first last]]]} :in-args [[["Billy" "Joel"] ["Billy" "Idol"]]]}'
+    assert len(list(result)) == 2
+
 def test_kwargs():
 
     known_args = ['my-foo', 'my-bar?', 'my-json', 'my-edn']