From 7b2546883e990e4c08b782fddac58245bdaec19d Mon Sep 17 00:00:00 2001
From: davidmezzetti <561939+davidmezzetti@users.noreply.github.com>
Date: Tue, 7 Nov 2023 11:00:11 -0500
Subject: [PATCH] Add support for SQL bind parameters, closes #564

---
 src/python/txtai/database/base.py           | 10 +++--
 src/python/txtai/database/client.py         |  9 ++--
 src/python/txtai/database/duckdb.py         | 49 ++++++++++++++++++++-
 src/python/txtai/database/rdbms.py          |  5 ++-
 src/python/txtai/database/sql/expression.py |  5 ++-
 src/python/txtai/embeddings/base.py         | 10 +++--
 src/python/txtai/embeddings/search/base.py  | 12 ++---
 src/python/txtai/embeddings/search/scan.py  | 40 ++++++++++++++---
 test/python/testdatabase/testdatabase.py    |  2 +-
 test/python/testdatabase/testrdbms.py       | 26 ++++++++++-
 10 files changed, 137 insertions(+), 31 deletions(-)

diff --git a/src/python/txtai/database/base.py b/src/python/txtai/database/base.py
index 58afff04..9e10c3bd 100644
--- a/src/python/txtai/database/base.py
+++ b/src/python/txtai/database/base.py
@@ -113,7 +113,7 @@ class Database:
 
         raise NotImplementedError
 
-    def search(self, query, similarity=None, limit=None):
+    def search(self, query, similarity=None, limit=None, parameters=None):
         """
         Runs a search against the database. Supports the following methods:
 
@@ -133,6 +133,7 @@ class Database:
             query: input query
             similarity: similarity results as [(indexid, score)]
             limit: maximum number of results to return
+            parameters: dict of named parameters to bind to placeholders
 
         Returns:
             query results as a list of dicts
@@ -159,7 +160,7 @@ class Database:
         query["where"] = where
 
         # Run query
-        return self.query(query, limit)
+        return self.query(query, limit, parameters)
 
     def parse(self, query):
         """
@@ -200,13 +201,14 @@ class Database:
 
         raise NotImplementedError
 
-    def query(self, query, limit):
+    def query(self, query, limit, parameters):
         """
         Executes query against database.
 
         Args:
             query: input query
             limit: maximum number of results to return
+            parameters: dict of named parameters to bind to placeholders
 
         Returns:
             query results
@@ -312,7 +314,7 @@ class Database:
 
         try:
             # Debug log SQL
-            logger.debug(*args)
+            logger.debug(" ".join(["%s"] * len(args)), *args)
 
             return function(*args)
         except Exception as ex:
diff --git a/src/python/txtai/database/client.py b/src/python/txtai/database/client.py
index e4ed8b46..638f3720 100644
--- a/src/python/txtai/database/client.py
+++ b/src/python/txtai/database/client.py
@@ -61,7 +61,7 @@ class Client(RDBMS):
         d = aliased(Document, name="d")
 
         # Build JSON column expression for column
-        return str(cast(d.data[name], Text).compile(dialect=self.connection.bind.dialect, compile_kwargs={"literal_binds": True}))
+        return str(cast(d.data[name].as_string(), Text).compile(dialect=self.connection.bind.dialect, compile_kwargs={"literal_binds": True}))
 
     def createtables(self):
         # Create tables
@@ -111,7 +111,7 @@ class Client(RDBMS):
         content = os.environ.get("CLIENT_URL") if content == "client" else content
 
         # Create engine using database URL
-        engine = create_engine(content, poolclass=StaticPool, echo=False)
+        engine = create_engine(content, poolclass=StaticPool, echo=False, json_serializer=lambda x: x)
 
         # Create database session
         return Session(engine)
@@ -138,18 +138,19 @@ class Cursor:
     def __iter__(self):
         return self.result
 
-    def execute(self, statement):
+    def execute(self, statement, parameters=None):
         """
         Executes statement.
 
         Args:
             statement: statement to execute
+            parameters: optional dictionary with bind parameters
         """
 
         if isinstance(statement, str):
             statement = textsql(statement)
 
-        self.result = self.connection.execute(statement)
+        self.result = self.connection.execute(statement, parameters)
 
     def fetchall(self):
         """
diff --git a/src/python/txtai/database/duckdb.py b/src/python/txtai/database/duckdb.py
index 22f00c91..2406591b 100644
--- a/src/python/txtai/database/duckdb.py
+++ b/src/python/txtai/database/duckdb.py
@@ -3,6 +3,7 @@ DuckDB module
 """
 
 import os
+import re
 
 from tempfile import TemporaryDirectory
 
@@ -33,18 +34,22 @@ class DuckDB(Embedded):
         if not DUCKDB:
             raise ImportError('DuckDB is not available - install "database" extra to enable')
 
+    def execute(self, function, *args):
+        # Call parent method with DuckDB compatible arguments
+        return super().execute(function, *self.formatargs(args))
+
     def insertdocument(self, uid, data, tags, entry):
         # Delete existing document
         self.cursor.execute(DuckDB.DELETE_DOCUMENT, [uid])
 
-        # Call parent logic
+        # Call parent method
         super().insertdocument(uid, data, tags, entry)
 
     def insertobject(self, uid, data, tags, entry):
         # Delete existing object
         self.cursor.execute(DuckDB.DELETE_OBJECT, [uid])
 
-        # Call parent logic
+        # Call parent method
         super().insertobject(uid, data, tags, entry)
 
     def connect(self, path=":memory:"):
@@ -58,6 +63,14 @@ class DuckDB(Embedded):
     def getcursor(self):
         return self.connection
 
+    def jsonprefix(self):
+        # Return json column prefix
+        return "json_extract_string(data"
+
+    def jsoncolumn(self, name):
+        # Generate json column using json_extract function
+        return f"json_extract_string(data, '$.{name}')"
+
     def rows(self):
         # Iteratively retrieve and yield rows
         batch = 256
@@ -103,3 +116,35 @@ class DuckDB(Embedded):
         connection.begin()
 
         return connection
+
+    def formatargs(self, args):
+        """
+        DuckDB doesn't support named parameters. This method replaces named parameters with question marks
+        and makes parameters a list.
+
+        Args:
+            args: input arguments
+
+        Returns:
+            DuckDB compatible args
+        """
+
+        if args and len(args) > 1:
+            # Unpack query args
+            query, parameters = args
+
+            # Iterate over parameters
+            #   - Replace named parameters with ?'s
+            #   - Build list of value with position indexes
+            params = []
+            for key, value in parameters.items():
+                pattern = rf"\:{key}(?=\s|$)"
+                match = re.search(pattern, query)
+                if match:
+                    query = re.sub(pattern, "?", query, count=1)
+                    params.append((match.start(), value))
+
+            # Repack query and parameter list
+            args = (query, [value for _, value in sorted(params, key=lambda x: x[0])])
+
+        return args
diff --git a/src/python/txtai/database/rdbms.py b/src/python/txtai/database/rdbms.py
index 30d6be3e..44a55e60 100644
--- a/src/python/txtai/database/rdbms.py
+++ b/src/python/txtai/database/rdbms.py
@@ -178,7 +178,7 @@ class RDBMS(Database):
         return Statement.IDS_CLAUSE % batch
 
     # pylint: disable=R0912
-    def query(self, query, limit):
+    def query(self, query, limit, parameters):
         # Extract query components
         select = query.get("select", self.defaults())
         where = query.get("where")
@@ -214,7 +214,8 @@ class RDBMS(Database):
             self.scores(None)
 
         # Runs a user query through execute method, which has common user query handling logic
-        self.execute(self.cursor.execute, query)
+        args = (query, parameters) if parameters else (query,)
+        self.execute(self.cursor.execute, *args)
 
         # Retrieve column list from query
         columns = [c[0] for c in self.cursor.description]
diff --git a/src/python/txtai/database/sql/expression.py b/src/python/txtai/database/sql/expression.py
index 647050d0..9bd92df6 100644
--- a/src/python/txtai/database/sql/expression.py
+++ b/src/python/txtai/database/sql/expression.py
@@ -387,7 +387,7 @@ class Expression:
 
     def resolve(self, token, aliases):
         """
-        Resolves this token's value if it is not an alias.
+        Resolves this token's value if it is not an alias or a bind parameter.
 
         Args:
             token: token to resolve
@@ -397,7 +397,8 @@ class Expression:
             resolved token value
         """
 
-        if aliases and Token.normalize(token) in aliases:
+        # Check for alias or bind parameter
+        if (aliases and Token.normalize(token) in aliases) or (token.startswith(":")):
             return token
 
         return self.resolver(token)
diff --git a/src/python/txtai/embeddings/base.py b/src/python/txtai/embeddings/base.py
index c8f4d5cb..3bbf9fe2 100644
--- a/src/python/txtai/embeddings/base.py
+++ b/src/python/txtai/embeddings/base.py
@@ -344,7 +344,7 @@ class Embeddings:
         # Default to 0 when no suitable method found
         return 0
 
-    def search(self, query, limit=None, weights=None, index=None):
+    def search(self, query, limit=None, weights=None, index=None, parameters=None):
         """
         Finds documents most similar to the input query. This method will run either an index search
         or an index + database search depending on if a database is available.
@@ -354,15 +354,16 @@ class Embeddings:
             limit: maximum results
             weights: hybrid score weights, if applicable
             index: index name, if applicable
+            parameters: dict of named parameters to bind to placeholders
 
         Returns:
             list of (id, score) for index search, list of dict for an index + database search
         """
 
-        results = self.batchsearch([query], limit, weights, index)
+        results = self.batchsearch([query], limit, weights, index, [parameters])
         return results[0] if results else results
 
-    def batchsearch(self, queries, limit=None, weights=None, index=None):
+    def batchsearch(self, queries, limit=None, weights=None, index=None, parameters=None):
         """
         Finds documents most similar to the input queries. This method will run either an index search
         or an index + database search depending on if a database is available.
@@ -372,12 +373,13 @@ class Embeddings:
             limit: maximum results
             weights: hybrid score weights, if applicable
             index: index name, if applicable
+            parameters: list of dicts of named parameters to bind to placeholders
 
         Returns:
             list of (id, score) per query for index search, list of dict per query for an index + database search
         """
 
-        return Search(self)(queries, limit, weights, index)
+        return Search(self)(queries, limit, weights, index, parameters)
 
     def similarity(self, query, data):
         """
diff --git a/src/python/txtai/embeddings/search/base.py b/src/python/txtai/embeddings/search/base.py
index 6063ab4f..5a16f0fd 100644
--- a/src/python/txtai/embeddings/search/base.py
+++ b/src/python/txtai/embeddings/search/base.py
@@ -37,7 +37,7 @@ class Search:
         self.query = embeddings.query
         self.scoring = embeddings.scoring if embeddings.issparse() else None
 
-    def __call__(self, queries, limit=None, weights=None, index=None):
+    def __call__(self, queries, limit=None, weights=None, index=None, parameters=None):
         """
         Executes a batch search for queries. This method will run either an index search or an index + database search
         depending on if a database is available.
@@ -47,6 +47,7 @@ class Search:
             limit: maximum results
             weights: hybrid score weights
             index: index name
+            parameters: list of dicts of named parameters to bind to placeholders
 
         Returns:
             list of (id, score) per query for index search, list of dict per query for an index + database search
@@ -66,7 +67,7 @@ class Search:
 
         # Database search
         if not self.indexids and self.database:
-            return self.dbsearch(queries, limit, weights, index)
+            return self.dbsearch(queries, limit, weights, index, parameters)
 
         # Default vector index query (sparse, dense or hybrid)
         return self.search(queries, limit, weights, index)
@@ -209,7 +210,7 @@ class Search:
 
         return results
 
-    def dbsearch(self, queries, limit, weights, index):
+    def dbsearch(self, queries, limit, weights, index, parameters):
         """
         Executes an index + database search.
 
@@ -218,6 +219,7 @@ class Search:
             limit: maximum results
             weights: default hybrid score weights
             index: default index name
+            parameters: list of dicts of named parameters to bind to placeholders
 
         Returns:
             list of dict per query
@@ -230,13 +232,13 @@ class Search:
         limit = max(limit, self.limit(queries))
 
         # Bulk index scan
-        scan = Scan(self.search, limit, weights, index)(queries)
+        scan = Scan(self.search, limit, weights, index)(queries, parameters)
 
         # Combine index search results with database search results
         results = []
         for x, query in enumerate(queries):
             # Run the database query, get matching bulk searches for current query
-            result = self.database.search(query, [r for y, r in scan if x == y], limit)
+            result = self.database.search(query, [r for y, r in scan if x == y], limit, parameters[x] if parameters and parameters[x] else None)
             results.append(result)
 
         return results
diff --git a/src/python/txtai/embeddings/search/scan.py b/src/python/txtai/embeddings/search/scan.py
index 076e9fb5..a7803e7b 100644
--- a/src/python/txtai/embeddings/search/scan.py
+++ b/src/python/txtai/embeddings/search/scan.py
@@ -34,12 +34,13 @@ class Scan:
         # Default index
         self.index = index
 
-    def __call__(self, queries):
+    def __call__(self, queries, parameters):
         """
         Executes a scan for a list of queries.
 
         Args:
             queries: list of queries to run
+            parameters: list of dicts of named parameters to bind to placeholders
 
         Returns:
             list of (id, score) per query
@@ -52,7 +53,7 @@ class Scan:
         default = None
 
         # Group by index and run
-        for index, iqueries in self.parse(queries).items():
+        for index, iqueries in self.parse(queries, parameters).items():
             # Query limit to pass to batch search
             candidates = [query.candidates for query in iqueries if query.candidates]
             if not candidates and not default:
@@ -75,12 +76,13 @@ class Scan:
         # Sort by query uid and return results
         return [result for _, result in sorted(results.items())]
 
-    def parse(self, queries):
+    def parse(self, queries, parameters):
         """
         Parse index query clauses from a list of parsed queries.
 
         Args:
             queries: list of parsed queries
+            parameters: list of dicts of named parameters to bind to placeholders
 
         Returns:
             index query clauses grouped by index
@@ -91,7 +93,11 @@ class Scan:
             if "similar" in query:
                 # Extract similar query clauses
                 for params in query["similar"]:
-                    # Parse query clause parameters
+                    # Resolve bind parameters
+                    if parameters and parameters[x]:
+                        params = self.bind(params, parameters[x])
+
+                    # Parse query clause
                     clause = Clause(uid, x, params)
 
                     # Create clause list for index
@@ -104,6 +110,28 @@ class Scan:
 
         return results
 
+    def bind(self, similar, parameters):
+        """
+        Resolves bind parameters for a similar function call.
+
+        Args:
+            similar: similar function call arguments
+            parameters: bind parameters
+
+        Returns:
+            similar function call arguments with resolved bind parameters
+        """
+
+        resolved = []
+        for p in similar:
+            # Resolve bind parameters
+            if isinstance(p, str) and p.startswith(":") and p[1:] in parameters:
+                resolved.append(parameters[p[1:]])
+            else:
+                resolved.append(p)
+
+        return resolved
+
     def default(self, queries):
         """
         Derives the default number of candidates. The number of candidates are the number of results to bring back
@@ -155,11 +183,11 @@ class Clause:
         """
 
         for param in params:
-            if param.isdigit():
+            if (isinstance(param, str) and param.isdigit()) or isinstance(param, int):
                 # Number of query candidates
                 self.candidates = int(param)
 
-            elif param.replace(".", "").isdigit():
+            elif (isinstance(param, str) and param.replace(".", "").isdigit()) or isinstance(param, float):
                 # Hybrid score weights
                 self.weights = float(param)
 
diff --git a/test/python/testdatabase/testdatabase.py b/test/python/testdatabase/testdatabase.py
index 248c2223..44a2331c 100644
--- a/test/python/testdatabase/testdatabase.py
+++ b/test/python/testdatabase/testdatabase.py
@@ -29,4 +29,4 @@ class TestDatabase(unittest.TestCase):
         self.assertRaises(NotImplementedError, database.count)
         self.assertRaises(NotImplementedError, database.resolve, None, None)
         self.assertRaises(NotImplementedError, database.embed, None, None)
-        self.assertRaises(NotImplementedError, database.query, None, None)
+        self.assertRaises(NotImplementedError, database.query, None, None, None)
diff --git a/test/python/testdatabase/testrdbms.py b/test/python/testdatabase/testrdbms.py
index e73040b5..04ac080f 100644
--- a/test/python/testdatabase/testrdbms.py
+++ b/test/python/testdatabase/testrdbms.py
@@ -615,7 +615,7 @@ class Common:
             """
 
             # Create an index for the list of text
-            self.embeddings.index([(uid, {"text": text, "length": len(text)}, None) for uid, text in enumerate(self.data)])
+            self.embeddings.index([(uid, {"text": text, "length": len(text), "attribute": f"ID{uid}"}, None) for uid, text in enumerate(self.data)])
 
             # Test similar
             result = self.embeddings.search(
@@ -643,10 +643,34 @@ class Common:
             result = self.embeddings.search("select id, text, length, data, entry from txtai")[0]
             self.assertEqual(sorted(result.keys()), ["data", "entry", "id", "length", "text"])
 
+            # Test column filtering
+            result = self.embeddings.search("select text from txtai where attribute = 'ID4'", 1)[0]
+            self.assertEqual(result["text"], self.data[4])
+
             # Test SQL parse error
             with self.assertRaises(SQLError):
                 self.embeddings.search("select * from txtai where bad,query")
 
+        def testSQLBind(self):
+            """
+            Test SQL statements with bind parameters
+            """
+
+            # Create an index for the list of text
+            self.embeddings.index([(uid, text, None) for uid, text in enumerate(self.data)])
+
+            # Test similar clause bind parameters
+            result = self.embeddings.search("select id, text, score from txtai where similar(:x)", parameters={"x": "feel good story"})[0]
+            self.assertEqual(result["text"], self.data[4])
+
+            # Test similar clause bind and non-bind parameters
+            result = self.embeddings.search("select id, text, score from txtai where similar(:x, 0.5)", parameters={"x": "feel good story"})[0]
+            self.assertEqual(result["text"], self.data[4])
+
+            # Test where filtering with bind parameters
+            result = self.embeddings.search("select * from txtai where text like :x", parameters={"x": "%iceberg%"})[0]
+            self.assertEqual(result["text"], self.data[1])
+
         def testSubindex(self):
             """
             Test subindex