From 5a9dfb65d16c0f243ec1b8359e3a93d11891a759 Mon Sep 17 00:00:00 2001 From: Srimon Date: Thu, 11 Jun 2026 20:07:29 +0530 Subject: [PATCH 1/2] feat: add gRPC support with configurable options for Qdrant connection --- src/qql/__init__.py | 4 +++ src/qql/ast_nodes.py | 12 +++---- src/qql/cli.py | 52 ++++++++++++++++++++++++++--- src/qql/config.py | 4 +++ src/qql/connection.py | 18 ++++++++++- src/qql/executor.py | 57 +++++++++++++++++++++++++++++--- src/qql/parser.py | 13 +++++--- tests/test_connection.py | 18 ++++++++++- tests/test_executor.py | 70 ++++++++++++++++++++++++++++++++++++++++ 9 files changed, 225 insertions(+), 23 deletions(-) diff --git a/src/qql/__init__.py b/src/qql/__init__.py index da43725..cd8009f 100644 --- a/src/qql/__init__.py +++ b/src/qql/__init__.py @@ -43,6 +43,8 @@ def run_query( secret: str | None = None, default_model: str | None = None, verify: bool | str = True, + prefer_grpc: bool = False, + grpc_port: int = 6334, ) -> ExecutionResult: """One-shot convenience function kept for backward compatibility. @@ -61,5 +63,7 @@ def run_query( secret=secret, default_model=default_model, verify=verify, + prefer_grpc=prefer_grpc, + grpc_port=grpc_port, ) as conn: return conn.run_query(query) diff --git a/src/qql/ast_nodes.py b/src/qql/ast_nodes.py index 5f0b2f2..aa6f9ca 100644 --- a/src/qql/ast_nodes.py +++ b/src/qql/ast_nodes.py @@ -25,9 +25,9 @@ class QuantizationConfig: class SearchWith: """Query-time search params supported by Qdrant SearchParams.""" hnsw_ef: int | None = None - exact: bool = False - acorn: bool = False - indexed_only: bool = False + exact: bool | None = None + acorn: bool | None = None + indexed_only: bool | None = None quantization: "QuantizationSearchWith | None" = None mmr_diversity: float | None = None mmr_candidates: int | None = None @@ -99,7 +99,7 @@ class CompareExpr: """field op literal — covers =, !=, >, >=, <, <=""" field: str op: str # one of: "=", "!=", ">", ">=", "<", "<=" - value: str | int | float | bool + value: str | int | float | bool | None @dataclass(frozen=True) @@ -114,14 +114,14 @@ class BetweenExpr: class InExpr: """field IN (v1, v2, ...)""" field: str - values: tuple[str | int | float | bool, ...] + values: tuple[str | int | float | bool | None, ...] @dataclass(frozen=True) class NotInExpr: """field NOT IN (v1, v2, ...)""" field: str - values: tuple[str | int | float | bool, ...] + values: tuple[str | int | float | bool | None, ...] @dataclass(frozen=True) diff --git a/src/qql/cli.py b/src/qql/cli.py index 36f266a..62922cc 100644 --- a/src/qql/cli.py +++ b/src/qql/cli.py @@ -1,6 +1,7 @@ from __future__ import annotations import sys +from typing import Any import click from prompt_toolkit import PromptSession @@ -18,6 +19,20 @@ console = Console() err_console = Console(stderr=True) + +def _client_kwargs_from_cfg(cfg: QQLConfig) -> dict[str, Any]: + """Build QdrantClient keyword arguments from a QQLConfig.""" + kwargs: dict[str, Any] = { + "url": cfg.url, + "api_key": cfg.secret, + "verify": cfg.verify, + } + if cfg.prefer_grpc: + kwargs["prefer_grpc"] = True + kwargs["grpc_port"] = cfg.grpc_port + return kwargs + + HELP_TEXT = """ [bold cyan]QQL — Qdrant Query Language[/bold cyan] @@ -185,11 +200,26 @@ def main(ctx: click.Context) -> None: type=click.Path(exists=True, readable=True, dir_okay=False, resolve_path=True), help="Path to a custom CA certificate bundle (PEM).", ) +@click.option( + "--prefer-grpc", + is_flag=True, + default=False, + help="Connect via gRPC transport instead of HTTP.", +) +@click.option( + "--grpc-port", + type=int, + default=6334, + show_default=True, + help="gRPC port of the Qdrant instance.", +) def connect( url: str, secret: str | None, verify: bool, ca_cert: str | None, + prefer_grpc: bool, + grpc_port: int, ) -> None: """Connect to a Qdrant instance and launch the QQL shell.""" from qdrant_client import QdrantClient @@ -201,8 +231,17 @@ def connect( console.print(f"Connecting to [bold]{url}[/bold]...") + client_kwargs: dict[str, Any] = { + "url": url, + "api_key": secret, + "verify": verify_val, + } + if prefer_grpc: + client_kwargs["prefer_grpc"] = True + client_kwargs["grpc_port"] = grpc_port + try: - client = QdrantClient(url=url, api_key=secret, verify=verify_val) + client = QdrantClient(**client_kwargs) client.get_collections() except Exception as e: err_console.print(f"[bold red]Connection failed:[/bold red] {e}") @@ -210,7 +249,10 @@ def connect( else: client.close() - cfg = QQLConfig(url=url, secret=secret, verify=verify_val) + cfg = QQLConfig( + url=url, secret=secret, verify=verify_val, + prefer_grpc=prefer_grpc, grpc_port=grpc_port, + ) save_config(cfg) console.print("[bold green]Connected.[/bold green] Config saved to ~/.qql/config.json\n") _launch_repl(cfg) @@ -252,7 +294,7 @@ def execute(file: str, stop_on_error: bool) -> None: sys.exit(1) try: - client = QdrantClient(url=cfg.url, api_key=cfg.secret, verify=cfg.verify) + client = QdrantClient(**_client_kwargs_from_cfg(cfg)) client.get_collections() except Exception as e: err_console.print(f"[bold red]Connection failed:[/bold red] {e}") @@ -310,7 +352,7 @@ def dump(collection: str, output: str, batch_size: int) -> None: sys.exit(1) try: - client = QdrantClient(url=cfg.url, api_key=cfg.secret, verify=cfg.verify) + client = QdrantClient(**_client_kwargs_from_cfg(cfg)) client.get_collections() except Exception as e: err_console.print(f"[bold red]Connection failed:[/bold red] {e}") @@ -343,7 +385,7 @@ def _launch_repl(cfg: QQLConfig) -> None: from qdrant_client import QdrantClient try: - client = QdrantClient(url=cfg.url, api_key=cfg.secret, verify=cfg.verify) + client = QdrantClient(**_client_kwargs_from_cfg(cfg)) client.get_collections() except Exception as e: err_console.print(f"[bold red]Could not connect to {cfg.url}:[/bold red] {e}") diff --git a/src/qql/config.py b/src/qql/config.py index 652e89a..b052ac9 100644 --- a/src/qql/config.py +++ b/src/qql/config.py @@ -20,6 +20,8 @@ class QQLConfig: default_dense_vector_name: str = DEFAULT_DENSE_VECTOR_NAME default_sparse_vector_name: str = DEFAULT_SPARSE_VECTOR_NAME verify: bool | str = True + prefer_grpc: bool = False + grpc_port: int = 6334 def save_config(cfg: QQLConfig) -> None: @@ -45,6 +47,8 @@ def load_config() -> QQLConfig | None: "default_sparse_vector_name", DEFAULT_SPARSE_VECTOR_NAME ), verify=data.get("verify", True), + prefer_grpc=data.get("prefer_grpc", False), + grpc_port=data.get("grpc_port", 6334), ) diff --git a/src/qql/connection.py b/src/qql/connection.py index fe72824..8ed3e54 100644 --- a/src/qql/connection.py +++ b/src/qql/connection.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Any + from .config import DEFAULT_MODEL, QQLConfig from .executor import Executor, ExecutionResult from .lexer import Lexer @@ -52,6 +54,8 @@ def __init__( secret: str | None = None, default_model: str | None = None, verify: bool | str = True, + prefer_grpc: bool = False, + grpc_port: int = 6334, ) -> None: """Create a connection to a Qdrant instance. @@ -64,6 +68,8 @@ def __init__( verify: SSL certificate verification. Set to ``False`` to skip verification for self-signed/internal certificates, or pass a path to a custom CA bundle (default: ``True``). + prefer_grpc: Whether to connect via fast gRPC transport. + grpc_port: The gRPC port of Qdrant instance (default: 6334). """ from qdrant_client import QdrantClient @@ -72,8 +78,18 @@ def __init__( secret=secret, default_model=default_model or DEFAULT_MODEL, verify=verify, + prefer_grpc=prefer_grpc, + grpc_port=grpc_port, ) - self._client = QdrantClient(url=url, api_key=secret, verify=verify) + client_kwargs: dict[str, Any] = { + "url": url, + "api_key": secret, + "verify": verify, + } + if prefer_grpc: + client_kwargs["prefer_grpc"] = True + client_kwargs["grpc_port"] = grpc_port + self._client = QdrantClient(**client_kwargs) self._executor = Executor(self._client, self._config) # ── Public API ──────────────────────────────────────────────────────── diff --git a/src/qql/executor.py b/src/qql/executor.py index 975947d..5e6735b 100644 --- a/src/qql/executor.py +++ b/src/qql/executor.py @@ -240,12 +240,18 @@ def execute(self, node: ASTNode) -> ExecutionResult: # ── Statement executors ─────────────────────────────────────────────── + @staticmethod + def _is_grpc_not_found_error(error: BaseException) -> bool: + """Return True if *error* is a gRPC NOT_FOUND status.""" + from grpc import RpcError, StatusCode + return isinstance(error, RpcError) and error.code() == StatusCode.NOT_FOUND + def _fetch_collection_info(self, name: str): """Fetch full CollectionInfo for *name* in a single API call. Returns the CollectionInfo object when the collection exists, or - ``None`` when the collection is not found (HTTP 404). Any other - Qdrant error is re-raised as :class:`QQLRuntimeError`. + ``None`` when the collection is not found (HTTP 404 or gRPC NOT_FOUND). + Any other Qdrant error is re-raised as :class:`QQLRuntimeError`. """ try: return self._client.get_collection(name) @@ -255,6 +261,18 @@ def _fetch_collection_info(self, name: str): raise QQLRuntimeError( f"Qdrant error fetching collection '{name}': {e}" ) from e + except ValueError as e: + if f"Collection {name} not found" in str(e): + return None + raise QQLRuntimeError( + f"Qdrant error fetching collection '{name}': {e}" + ) from e + except Exception as e: + if self._is_grpc_not_found_error(e): + return None + raise QQLRuntimeError( + f"Qdrant error fetching collection '{name}': {e}" + ) from e def _topology_from_collection_info(self, info: Any) -> CollectionTopology: """Parse a CollectionInfo object into a :class:`CollectionTopology`. @@ -1835,6 +1853,15 @@ def _build_qdrant_filter(self, expr: FilterExpr) -> Any: # ── Comparison ──────────────────────────────────────────────────── if isinstance(expr, CompareExpr): + if expr.value is None: + null_condition = IsNullCondition(is_null=PayloadField(key=expr.field)) + if expr.op == "=": + return null_condition + if expr.op == "!=": + return Filter(must_not=[null_condition]) + raise QQLRuntimeError( + f"Cannot use operator '{expr.op}' with null for field '{expr.field}'" + ) if expr.op == "=": return FieldCondition( key=expr.field, match=MatchValue(value=expr.value) @@ -1858,14 +1885,34 @@ def _build_qdrant_filter(self, expr: FilterExpr) -> Any: # ── IN / NOT IN ─────────────────────────────────────────────────── if isinstance(expr, InExpr): - return FieldCondition( - key=expr.field, match=MatchAny(any=list(expr.values)) + non_nulls = [v for v in expr.values if v is not None] + if len(non_nulls) == len(expr.values): + return FieldCondition( + key=expr.field, match=MatchAny(any=non_nulls) + ) + null_condition = IsNullCondition(is_null=PayloadField(key=expr.field)) + if not non_nulls: + return null_condition + return Filter( + should=[ + null_condition, + FieldCondition(key=expr.field, match=MatchAny(any=non_nulls)), + ] ) if isinstance(expr, NotInExpr): + non_nulls = [v for v in expr.values if v is not None] + null_condition = IsNullCondition(is_null=PayloadField(key=expr.field)) + if len(non_nulls) != len(expr.values): + must_not = [null_condition] + if non_nulls: + must_not.append( + FieldCondition(key=expr.field, match=MatchAny(any=non_nulls)) + ) + return Filter(must_not=must_not) return FieldCondition( key=expr.field, - match=MatchExcept(**{"except": list(expr.values)}), + match=MatchExcept(**{"except": non_nulls}), ) # ── IS NULL / IS NOT NULL ───────────────────────────────────────── diff --git a/src/qql/parser.py b/src/qql/parser.py index b62b9b0..acc88c0 100644 --- a/src/qql/parser.py +++ b/src/qql/parser.py @@ -1165,12 +1165,15 @@ def _parse_field_path(self) -> str: f"Expected a field name, got '{tok.value}'", tok.pos ) - def _parse_literal(self) -> str | int | float | bool: - """STRING | INTEGER | FLOAT | boolean""" + def _parse_literal(self) -> str | int | float | bool | None: + """STRING | INTEGER | FLOAT | boolean | NULL""" tok = self._peek() if tok.kind == TokenKind.STRING: self._advance() return tok.value + if tok.kind == TokenKind.NULL: + self._advance() + return None if tok.kind == TokenKind.INTEGER: self._advance() return int(tok.value) @@ -1186,7 +1189,7 @@ def _parse_literal(self) -> str | int | float | bool: self._advance() return False raise QQLSyntaxError( - f"Expected a literal value (string, integer, float, or boolean), got '{tok.value}'", + f"Expected a literal value (string, integer, float, boolean, or null), got '{tok.value}'", tok.pos, ) @@ -1203,10 +1206,10 @@ def _parse_number(self) -> int | float: f"Expected a number, got '{tok.value}'", tok.pos ) - def _parse_literal_list(self) -> list[str | int | float | bool]: + def _parse_literal_list(self) -> list[str | int | float | bool | None]: """'(' literal { ',' literal } [','] ')' — used by IN / NOT IN.""" self._expect(TokenKind.LPAREN) - items: list[str | int | float | bool] = [] + items: list[str | int | float | bool | None] = [] if self._peek().kind == TokenKind.RPAREN: self._advance() return items diff --git a/tests/test_connection.py b/tests/test_connection.py index 2278002..ebcce74 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -44,6 +44,17 @@ def test_custom_ca_bundle_passed_to_qdrant_client(self, mocker): ) assert conn.config.verify == "/etc/ssl/internal-ca.pem" + def test_grpc_options_passed_to_qdrant_client(self, mocker): + mock_client_cls = mocker.patch("qdrant_client.QdrantClient") + Connection("http://localhost:6333", prefer_grpc=True, grpc_port=9999) + mock_client_cls.assert_called_once_with( + url="http://localhost:6333", + api_key=None, + verify=True, + prefer_grpc=True, + grpc_port=9999, + ) + def test_custom_default_model_stored_in_config(self, mocker): mocker.patch("qdrant_client.QdrantClient") conn = Connection("http://localhost:6333", default_model="BAAI/bge-small-en-v1.5") @@ -174,7 +185,12 @@ def test_run_query_delegates_to_connection(self, mocker): conn_cls = mocker.patch("qql.Connection", return_value=conn_instance) run_query("SHOW COLLECTIONS", url="http://localhost:6333") conn_cls.assert_called_once_with( - url="http://localhost:6333", secret=None, default_model=None, verify=True + url="http://localhost:6333", + secret=None, + default_model=None, + verify=True, + prefer_grpc=False, + grpc_port=6334, ) conn_instance.run_query.assert_called_once_with("SHOW COLLECTIONS") diff --git a/tests/test_executor.py b/tests/test_executor.py index 0a2a74f..1c7c897 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -86,6 +86,76 @@ def mock_embedder(mocker): return mock_embed +class TestFetchCollectionInfo: + def test_value_error_collection_not_found_returns_none(self, executor, mock_client): + mock_client.get_collection.side_effect = ValueError("Collection docs not found") + result = executor._fetch_collection_info("docs") + assert result is None + + def test_value_error_other_message_raises(self, executor, mock_client): + mock_client.get_collection.side_effect = ValueError("transport failed") + with pytest.raises(QQLRuntimeError, match="Qdrant error fetching collection"): + executor._fetch_collection_info("docs") + + def test_unexpected_response_404_returns_none(self, executor, mock_client): + mock_client.get_collection.side_effect = UnexpectedResponse( + status_code=404, reason_phrase="Not Found", content=b"Not Found", headers={}, + ) + result = executor._fetch_collection_info("docs") + assert result is None + + def test_unexpected_response_other_status_raises(self, executor, mock_client): + mock_client.get_collection.side_effect = UnexpectedResponse( + status_code=500, reason_phrase="Internal Error", content=b"", headers={}, + ) + with pytest.raises(QQLRuntimeError, match="Qdrant error fetching collection"): + executor._fetch_collection_info("docs") + + def test_runtime_error_wrapped_as_qql_runtime_error(self, executor, mock_client): + mock_client.get_collection.side_effect = RuntimeError("transport failed") + with pytest.raises(QQLRuntimeError, match="Qdrant error fetching collection"): + executor._fetch_collection_info("docs") + + def test_grpc_not_found_error_returns_none(self, executor, mock_client, mocker): + class FakeRpcError(Exception): + pass + + NOT_FOUND_CODE = 5 + fakeStatusCode = mocker.MagicMock() + fakeStatusCode.NOT_FOUND = NOT_FOUND_CODE + + fake_error = FakeRpcError("NOT_FOUND") + fake_error.code = mocker.MagicMock(return_value=NOT_FOUND_CODE) + + grpc_mod = mocker.MagicMock() + grpc_mod.RpcError = FakeRpcError + grpc_mod.StatusCode = fakeStatusCode + mocker.patch.dict("sys.modules", {"grpc": grpc_mod}) + mock_client.get_collection.side_effect = fake_error + + result = executor._fetch_collection_info("docs") + assert result is None + + def test_grpc_other_error_raises(self, executor, mock_client, mocker): + class FakeRpcError(Exception): + pass + + fakeStatusCode = mocker.MagicMock() + fakeStatusCode.NOT_FOUND = 5 + + fake_error = FakeRpcError("INTERNAL") + fake_error.code = mocker.MagicMock(return_value=13) + + grpc_mod = mocker.MagicMock() + grpc_mod.RpcError = FakeRpcError + grpc_mod.StatusCode = fakeStatusCode + mocker.patch.dict("sys.modules", {"grpc": grpc_mod}) + mock_client.get_collection.side_effect = fake_error + + with pytest.raises(QQLRuntimeError, match="Qdrant error fetching collection"): + executor._fetch_collection_info("docs") + + class TestInsert: def test_insert_creates_collection_when_missing(self, executor, mock_client): node = InsertStmt(collection="notes", values={"text": "hello"}, model=None) From caca74cd2c9e5007c90b773fefff9ac465b719a3 Mon Sep 17 00:00:00 2001 From: Srimon Date: Thu, 11 Jun 2026 20:35:46 +0530 Subject: [PATCH 2/2] feat: improve error handling and parameter validation in Qdrant connection and executor --- src/qql/cli.py | 3 +++ src/qql/executor.py | 4 ++-- src/qql/parser.py | 12 ++++++------ tests/test_executor.py | 4 +++- 4 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/qql/cli.py b/src/qql/cli.py index 62922cc..fd8d47a 100644 --- a/src/qql/cli.py +++ b/src/qql/cli.py @@ -240,11 +240,14 @@ def connect( client_kwargs["prefer_grpc"] = True client_kwargs["grpc_port"] = grpc_port + client = None try: client = QdrantClient(**client_kwargs) client.get_collections() except Exception as e: err_console.print(f"[bold red]Connection failed:[/bold red] {e}") + if client is not None: + client.close() sys.exit(1) else: client.close() diff --git a/src/qql/executor.py b/src/qql/executor.py index 5e6735b..af12465 100644 --- a/src/qql/executor.py +++ b/src/qql/executor.py @@ -1351,8 +1351,8 @@ def _build_search_params(self, with_clause: SearchWith | None) -> SearchParams | hnsw_ef=with_clause.hnsw_ef, exact=with_clause.exact, quantization=quantization, - indexed_only=True if with_clause.indexed_only else None, - acorn=AcornSearchParams(enable=True) if with_clause.acorn else None, + indexed_only=with_clause.indexed_only if with_clause.indexed_only is not None else None, + acorn=AcornSearchParams(enable=with_clause.acorn) if with_clause.acorn is not None else None, ) def _build_hnsw_config(self, config: CollectionConfig | None) -> HnswConfigDiff | None: diff --git a/src/qql/parser.py b/src/qql/parser.py index acc88c0..a0b4d25 100644 --- a/src/qql/parser.py +++ b/src/qql/parser.py @@ -806,9 +806,9 @@ def _parse_search(self) -> SearchStmt: else: with_clause = SearchWith( hnsw_ef=parsed_with.hnsw_ef or with_clause.hnsw_ef, - exact=parsed_with.exact or with_clause.exact, - acorn=parsed_with.acorn or with_clause.acorn, - indexed_only=parsed_with.indexed_only or with_clause.indexed_only, + exact=parsed_with.exact if parsed_with.exact is not None else with_clause.exact, + acorn=parsed_with.acorn if parsed_with.acorn is not None else with_clause.acorn, + indexed_only=parsed_with.indexed_only if parsed_with.indexed_only is not None else with_clause.indexed_only, quantization=parsed_with.quantization or with_clause.quantization, mmr_diversity=( parsed_with.mmr_diversity @@ -1363,9 +1363,9 @@ def _parse_value(self) -> Any: def _parse_with_clause(self) -> SearchWith: self._expect(TokenKind.LBRACE) hnsw_ef: int | None = None - exact: bool = False - acorn: bool = False - indexed_only: bool = False + exact: bool | None = None + acorn: bool | None = None + indexed_only: bool | None = None quantization: QuantizationSearchWith | None = None mmr_diversity: float | None = None mmr_candidates: int | None = None diff --git a/tests/test_executor.py b/tests/test_executor.py index 1c7c897..cf23f3e 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -1272,6 +1272,7 @@ def test_sparse_search_forwards_search_params(self, executor, mock_client, mocke mock_response = mocker.MagicMock() mock_response.points = [] mock_client.query_points.return_value = mock_response + mocker.patch("qql.executor.SparseEmbedder", return_value=mocker.MagicMock()) node = SearchStmt( collection="notes", @@ -3606,7 +3607,7 @@ def test_insert_existing_dense_collection_fetches_metadata_once( mock_client.get_collection.assert_called_once() def test_insert_existing_hybrid_collection_fetches_metadata_once( - self, executor, mock_client + self, executor, mock_client, mocker ): """Hybrid auto-detect INSERT must call get_collection() exactly once.""" from qdrant_client.models import Distance, SparseVectorParams, VectorParams @@ -3618,6 +3619,7 @@ def test_insert_existing_hybrid_collection_fetches_metadata_once( mock_client.get_collection.return_value.config.params.sparse_vectors = { "sparse": SparseVectorParams() } + mocker.patch("qql.executor.SparseEmbedder", return_value=mocker.MagicMock()) node = InsertStmt(collection="docs", values={"text": "hello"}, model=None) executor.execute(node)