Fix PR #3 integration regressions in query helpers

Author: ankitml

paradedb/sqlalchemy/pdb.py

@@ -25,6 +25,12 @@ def snippet(
     if (start_tag is None) != (end_tag is None):
         raise InvalidArgumentError("start_tag and end_tag must be provided together")
 
+    if max_num_chars is not None and start_tag is None and end_tag is None:
+        # ParadeDB versions in CI don't support pdb.snippet(field, max_num_chars)
+        # directly. Supplying default tags targets the supported 4-arg form.
+        start_tag = "<b>"
+        end_tag = "</b>"
+
     args: list[Any] = [field]
     if start_tag is not None and end_tag is not None:
         args.extend([start_tag, end_tag])
@@ -36,12 +42,21 @@ def snippet(
 def snippets(
     field: ColumnElement,
     *,
+    start_tag: str | None = None,
+    end_tag: str | None = None,
     max_num_chars: int | None = None,
     limit: int | None = None,
     offset: int | None = None,
     sort_by: str | None = None,
 ) -> ClauseElement:
+    if (start_tag is None) != (end_tag is None):
+        raise InvalidArgumentError("start_tag and end_tag must be provided together")
+
     named_args: list[tuple[str, Any]] = []
+    if start_tag is not None:
+        named_args.append(("start_tag", start_tag))
+    if end_tag is not None:
+        named_args.append(("end_tag", end_tag))
     if max_num_chars is not None:
         named_args.append(("max_num_chars", max_num_chars))
     if limit is not None:

paradedb/sqlalchemy/search.py

@@ -8,6 +8,7 @@
 from sqlalchemy.sql import operators
 from sqlalchemy.sql.elements import ClauseElement, ColumnElement
 
+from ._functions import PDBFunctionWithNamedArgs
 from ._pdb_cast import PDBCast
 from .errors import InvalidArgumentError, InvalidMoreLikeThisOptionsError
 from .validation import (
@@ -238,6 +239,8 @@ def range_term(
         raise InvalidArgumentError(
             f"relation must be one of: {', '.join(sorted(_VALID_RANGE_RELATIONS))}"
         )
+    escaped_relation = relation.replace("'", "''")
+    relation_arg: ClauseElement = literal_column(f"'{escaped_relation}'")
     if range_type is not None:
         if range_type not in _VALID_RANGE_TYPES:
             raise InvalidArgumentError(
@@ -247,7 +250,7 @@ def range_term(
         bounds_arg: ClauseElement = literal_column(f"'{escaped}'::{range_type}")
     else:
         bounds_arg = literal(bounds)
-    return field.operate(_QUERY, func.pdb.range_term(bounds_arg, relation))
+    return field.operate(_QUERY, func.pdb.range_term(bounds_arg, relation_arg))
 
 
 def more_like_this(
@@ -309,58 +312,36 @@ def more_like_this(
     if stopwords is not None:
         require_non_empty_strings(stopwords, field_name="stopwords", error_cls=error_cls)
 
-    options_provided = any(
-        option is not None
-        for option in (
-            min_term_frequency,
-            max_query_terms,
-            min_doc_frequency,
-            max_doc_frequency,
-            min_word_length,
-            max_word_length,
-            boost_factor,
-            stopwords,
-        )
-    )
-
-    def _build_id_args(doc_id: Any) -> list[Any]:
-        id_args: list[Any] = [doc_id]
-        if fields is not None or options_provided:
-            id_args.append(array(fields or [], type_=Text()))
-        if options_provided:
-            id_args.extend(
-                [
-                    1 if min_term_frequency is None else min_term_frequency,
-                    25 if max_query_terms is None else max_query_terms,
-                    0 if min_doc_frequency is None else min_doc_frequency,
-                    1_000_000 if max_doc_frequency is None else max_doc_frequency,
-                    0 if min_word_length is None else min_word_length,
-                    1000 if max_word_length is None else max_word_length,
-                    0.0 if boost_factor is None else boost_factor,
-                    array(stopwords or [], type_=Text()),
-                ]
-            )
-        return id_args
+    named_options: list[tuple[str, Any]] = []
+    if min_term_frequency is not None:
+        named_options.append(("min_term_frequency", min_term_frequency))
+    if max_query_terms is not None:
+        named_options.append(("max_query_terms", max_query_terms))
+    if min_doc_frequency is not None:
+        named_options.append(("min_doc_frequency", min_doc_frequency))
+    if max_doc_frequency is not None:
+        named_options.append(("max_doc_frequency", max_doc_frequency))
+    if min_word_length is not None:
+        named_options.append(("min_word_length", min_word_length))
+    if max_word_length is not None:
+        named_options.append(("max_word_length", max_word_length))
+    if boost_factor is not None:
+        named_options.append(("boost_factor", boost_factor))
+    if stopwords is not None:
+        named_options.append(("stopwords", array(stopwords, type_=Text())))
+
+    def _build_mlt_call(source_arg: ClauseElement, *, include_fields: bool) -> ClauseElement:
+        positional_args: list[ClauseElement] = [source_arg]
+        if include_fields and fields is not None:
+            positional_args.append(array(fields, type_=Text()))
+        return PDBFunctionWithNamedArgs("more_like_this", positional_args, named_options)
 
     if document_ids is not None:
-        return or_(*[field.operate(_QUERY, func.pdb.more_like_this(*_build_id_args(doc_id))) for doc_id in document_ids])
+        clauses = [field.operate(_QUERY, _build_mlt_call(literal(doc_id), include_fields=True)) for doc_id in document_ids]
+        return or_(*clauses)
 
     if document_id is not None:
-        return field.operate(_QUERY, func.pdb.more_like_this(*_build_id_args(document_id)))
+        return field.operate(_QUERY, _build_mlt_call(literal(document_id), include_fields=True))
 
     payload = document if isinstance(document, str) else json.dumps(document, separators=(",", ":"), sort_keys=True)
-    doc_args: list[Any] = [payload]
-    if options_provided:
-        doc_args.extend(
-            [
-                1 if min_term_frequency is None else min_term_frequency,
-                25 if max_query_terms is None else max_query_terms,
-                0 if min_doc_frequency is None else min_doc_frequency,
-                1_000_000 if max_doc_frequency is None else max_doc_frequency,
-                0 if min_word_length is None else min_word_length,
-                1000 if max_word_length is None else max_word_length,
-                0.0 if boost_factor is None else boost_factor,
-                array(stopwords or [], type_=Text()),
-            ]
-        )
-    return field.operate(_QUERY, func.pdb.more_like_this(*doc_args))
+    return field.operate(_QUERY, _build_mlt_call(literal(payload), include_fields=False))

paradedb/sqlalchemy/select_with.py

@@ -42,6 +42,8 @@ def snippets(
     field: ColumnElement,
     *,
     label: str = "snippets",
+    start_tag: str | None = None,
+    end_tag: str | None = None,
     max_num_chars: int | None = None,
     limit: int | None = None,
     offset: int | None = None,
@@ -51,6 +53,8 @@ def snippets(
     return stmt.add_columns(
         pdb.snippets(
             field,
+            start_tag=start_tag,
+            end_tag=end_tag,
             max_num_chars=max_num_chars,
             limit=limit,
             offset=offset,

tests/integration/conftest.py

@@ -103,7 +103,9 @@ def explain_plan_json(session: Session, stmt) -> dict[str, Any]:
             compile_kwargs={"literal_binds": True},
         )
     )
-    explain_result = session.execute(text(f"EXPLAIN (FORMAT JSON) {sql}")).scalar_one()
+    # Use driver-level execution so SQLAlchemy doesn't treat JSON fragments like
+    # `...:2...` inside string literals as bind placeholders.
+    explain_result = session.connection().exec_driver_sql(f"EXPLAIN (FORMAT JSON) {sql}").scalar_one()
     return _normalize_explain_plan(explain_result)
 
 

tests/integration/test_indexing_integration.py

@@ -182,8 +182,9 @@ def test_bm25_index_json_keys_when_supported(engine):
         ).one()
 
     assert "->>" in row.indexdef
-    assert "alias=metadata_color" in row.indexdef
-    assert "alias=metadata_location" in row.indexdef
+    assert "'color'" in row.indexdef
+    assert "'location'" in row.indexdef
+    assert row.indexdef.count("pdb.literal(") >= 2
 
     _drop_table_and_index(engine, table_name, index_name)
 

tests/unit/test_sql_compilation_unit.py

@@ -61,14 +61,24 @@ def test_pdb_helpers_compile():
     stmt = select(
         pdb.score(products.c.id).label("score"),
         pdb.snippet(products.c.description, start_tag="<mark>", end_tag="</mark>", max_num_chars=100).label("snippet"),
-        pdb.snippets(products.c.description, max_num_chars=15, limit=1, offset=0, sort_by="position").label("snippets"),
+        pdb.snippets(
+            products.c.description,
+            start_tag="[",
+            end_tag="]",
+            max_num_chars=15,
+            limit=1,
+            offset=0,
+            sort_by="position",
+        ).label("snippets"),
         pdb.snippet_positions(products.c.description).label("positions"),
     )
     sql = _sql(stmt)
 
     assert "pdb.score(products.id) AS score" in sql
     assert "pdb.snippet(products.description, '<mark>', '</mark>', 100) AS snippet" in sql
     assert "pdb.snippets(products.description" in sql
+    assert "start_tag => '['" in sql
+    assert "end_tag => ']'" in sql
     assert "max_num_chars => 15" in sql
     assert '"limit" => 1' in sql
     assert '"offset" => 0' in sql
@@ -167,7 +177,10 @@ def test_more_like_this_compile():
 
     assert "id @@@ pdb.more_like_this(3, ARRAY['description'])" in by_id_sql
     assert "id @@@ pdb.more_like_this('{\"description\":\"wireless earbuds\"}')" in by_doc_sql
-    assert "id @@@ pdb.more_like_this(3, ARRAY['description'], 2, 10" in with_opts_sql
+    assert "id @@@ pdb.more_like_this(3, ARRAY['description']" in with_opts_sql
+    assert "min_term_frequency => 2" in with_opts_sql
+    assert "max_query_terms => 10" in with_opts_sql
+    assert "stopwords => ARRAY['the', 'a']" in with_opts_sql
     assert "ARRAY['the', 'a']" in with_opts_sql