mirror of
https://github.com/neuml/txtai.git
synced 2026-10-11 21:39:55 +00:00
Add support for SQL bind parameters, closes #564
This commit is contained in:
davidmezzetti
@@ -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:
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user