refactor(sdk): make queryweaver_sdk pip-installable with shared core#544
Open
galshubeli wants to merge 18 commits intostagingfrom
Open
refactor(sdk): make queryweaver_sdk pip-installable with shared core#544galshubeli wants to merge 18 commits intostagingfrom
galshubeli wants to merge 18 commits intostagingfrom
Conversation
- Add detailed assertions for query results (customer names, counts, etc.) - Add tests for filter queries, count aggregation, and joins - Validate SQL query structure and result data - Add session-scoped event loop to fix pytest-asyncio issues - Handle async event loop cleanup errors gracefully with skip - Expand model serialization tests
Disable warnings that are intentional architectural choices: - C0415: import-outside-toplevel (lazy imports for SDK) - W0718: broad-exception-caught (error handling) - R0902: too-many-instance-attributes (dataclasses) - R0903: too-few-public-methods - R0911: too-many-return-statements - R0913/R0917: too-many-arguments (SDK API design) - C0302: too-many-lines
- Extract SDK sync functions to new api/core/text2sql_sync.py module - Split QueryResult into composition: QueryResult + QueryMetadata + QueryAnalysis - Reduce local variables in query_database_sync with helper functions - Fix broad exception handling - use specific Redis/Connection/OS errors - Refactor query method to accept Union[str, QueryRequest] - Add compatibility properties to QueryResult for backwards compatibility - Document lazy imports in client.py module docstring Pylint score improved from 9.81/10 to 9.91/10 Remaining E0401 errors are missing dependencies in venv, not code issues
Resolve merge conflicts in .github/workflows/tests.yml, Makefile, and pyproject.toml. Standardize on uv (from staging migration) while preserving SDK additions from the PR branch: - tests.yml: Use uv instead of pipenv in sdk-tests job, keep SDK test exclusion in unit-tests, update action versions to v6 - Makefile: Use 'uv run' directly instead of pipenv/uv fallback, keep SDK targets (test-sdk, build-package, docker-test-services) - pyproject.toml: Keep SDK optional-dependencies structure (server/dev/all), use staging's dependency versions (~= notation), merge pytest/pylint config from both sides, keep [tool.uv] and [dependency-groups] Co-authored-by: Copilot <[email protected]>
- Use 'uv sync --all-extras' in tests, playwright, and pylint workflows to install server/dev optional dependencies (fastapi, uvicorn, etc.) - Fix imports in api/routes/graphs.py: import GENERAL_PREFIX and graph_name from text2sql_common, errors from api.core.errors - Remove unused variable 'last_message' in api/core/schema_loader.py - Add pylint disable comments for too-many-arguments/locals in api/core/text2sql_sync.py Co-authored-by: Copilot <[email protected]>
- Use uv sync --all-extras in CI (needed for SDK optional deps) - Update dependency versions to staging's latest (litellm 1.82, falkordb 1.6, etc.) - Keep SDK optional dependency structure (core vs server vs dev) - Align playwright version across dependency groups (~=1.58.0) - Merge pylint config (max-line-length, ignore-patterns, messages_control) - Regenerate uv.lock Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>
Bumps SDK to 0.2.0. Removes the api.extensions.db global mutation that
made two QueryWeaver instances collide, fixes the un-prefixed
database_id contract (breaking), and deduplicates orchestration
between the streaming and sync text2sql paths via shared helpers in
api.core.text2sql_common.
Key changes:
- api/core/db_resolver.py (new): lazy server-default fallback for db
- api/core/text2sql_common.py: 6 new shared helpers (execute_with_healing,
refresh_schema_if_modified, quote_identifiers_from_graph,
save_memory_background, format_ai_response,
build_destructive_confirmation_message); MESSAGE_DELIMITER moved here
- api/core/{result,request}_models.py (new): dataclasses extracted so
api.core no longer depends on queryweaver_sdk
- api/graph.py, loaders/*, memory/graphiti_tool.py: accept optional
db param; routes unchanged (use resolver fallback)
- queryweaver_sdk/client.py: no more api.extensions.db mutation;
close() awaits in-flight memory writes via contextvar-scoped sink
- queryweaver_sdk/models.py: re-export shim over api.core
- tests/test_sdk/: async-with fixtures, test_two_instances_isolated,
no more skip-on-RuntimeError wrappers
- CI: restored --locked --all-extras in all three workflows
Breaking (SDK 0.1.0 -> 0.2.0):
- DatabaseConnection.database_id is now un-prefixed (fixes double-
namespacing when passed back into query/delete/get_schema)
Fixes verified end-to-end: NL -> SQL -> results against live
PostgreSQL via SDK; FastAPI server boots cleanly with 21 routes and
HSTS middleware; two QueryWeaver instances coexist with independent
FalkorDB handles; uv build produces a working wheel; clean-venv
install succeeds.
Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
Completed Working on "Code Review"✅ Review publishing completed successfully. Posted comments from all chunks and submitted final review: COMMENT. Total comments: 8 across 8 files. ✅ Workflow completed successfully. |
|
🚅 Deployed to the QueryWeaver-pr-544 environment in queryweaver
|
Contributor
📝 WalkthroughWalkthroughA Python SDK for Text2SQL is introduced alongside core refactoring that centralizes shared utilities, implements dependency injection for database connections, adds sync/async API functions for SDK usage, and includes integration test infrastructure for multi-service validation. Changes
Sequence DiagramsequenceDiagram
participant App as QueryWeaver SDK
participant Core as text2sql_sync
participant Analysis as Analysis Agents
participant Exec as execute_with_healing
participant Loader as BaseLoader
participant DB as FalkorDB/SQL
participant Memory as MemoryTool
App->>Core: query_database_sync(user_id, graph_id, chat_data)
Core->>Analysis: validate_relevancy + get_table_discovery
Analysis-->>Core: relevance/tables result
alt Off-topic
Core-->>App: QueryResult(missing_information="...")
else Needs confirmation
Core-->>App: QueryResult(requires_confirmation=true)
else Proceed
Core->>Analysis: analyze_and_generate_sql
Analysis-->>Core: sql_query + analysis_output
Core->>Exec: execute_with_healing(sql_query, ...)
Exec->>Loader: execute SQL
Loader->>DB: run query
DB-->>Loader: results
Exec->>DB: check health/healing
Exec-->>Core: (result, analysis)
Core->>Memory: save_memory_background(question, sql, result, ...)
Core-->>App: QueryResult(sql_query, results, analysis, metadata)
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~75 minutes
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![overcut-ai[bot]](https://avatars.githubusercontent.com/in/1367736?s=60&v=4){kind=small}
There was a problem hiding this comment.
Thanks for the substantial refactor and test expansion. Consolidating the posted findings:
- 8 MAJOR comments (no BLOCKER/CRITICAL/MINOR/SUGGESTION/PRAISE in this review set)
- 8 files affected
Key themes
- SDK packaging/runtime boundary regressions: several findings indicate minimal
pip install queryweavermay still fail due to transitive server-only imports and dependency coupling. - Behavioral parity/correctness gaps: differences between SDK and streaming paths (memory persistence, error wrapping) can produce inconsistent outcomes and contracts.
- Test reliability/coverage risks: current fixture scoping and API-key gating reduce deterministic coverage for non-LLM SDK behaviors.
Recommended next steps
- Harden package boundaries so SDK imports and core SDK operations run on minimal install; add CI coverage for minimal-install smoke tests.
- Restore contract consistency by aligning error handling and memory-write behavior between SDK and streaming paths.
- Improve test determinism by removing unnecessary LLM-key gating for non-LLM tests and isolating mutable DB state per test.
Addressing these should materially reduce runtime surprises and improve confidence in the 0.2.0 SDK surface.
| analysis = _parse_analysis_result(answer_an) | ||
|
|
||
| if not analysis.is_valid: | ||
| return _invalid_sql_result(ctx, analysis, answer_an) |
There was a problem hiding this comment.
[major]: When AnalysisAgent returns is_sql_translatable=false, the SDK path returns early via _invalid_sql_result without persisting query/memory, while the streaming path still saves memory for the turn. This breaks parity and loses conversation traces for failed/ambiguous questions when use_memory is enabled.
Before returning from the invalid-SQL branch, call save_memory_background(...) similarly to other exit paths (include success=false and the follow-up response payload).
| if success: | ||
| # SDK callers pass the un-prefixed database_id back into query/delete/etc., | ||
| # where graph_name(user_id, db_name) re-applies the user_id prefix. | ||
| db_name = urlparse(url).path.lstrip("/") |
There was a problem hiding this comment.
[major]: load_database_sync now derives database_id using urlparse(url).path.lstrip('/'), but for URLs like .../mydb?sslmode=require this can propagate malformed IDs back to SDK callers. Round-tripping this value into query/delete/get_schema/refresh_schema can target a non-existent graph and break connect→query flows.
Derive the DB name from parsed URL components (path only) and never include query/fragment; reject empty/invalid names explicitly instead of returning potentially malformed IDs.
| except ResponseError as re: | ||
| raise GraphNotFoundError("Failed to delete graph, Graph not found") from re | ||
| except Exception as e: # pylint: disable=broad-exception-caught | ||
| except (RedisError, ConnectionError) as e: |
There was a problem hiding this comment.
[major]: Replacing the previous broad fallback with only (RedisError, ConnectionError) can let other backend/driver failures escape unwrapped, causing inconsistent API/SDK error contracts on delete failures.
Restore a final broad catch that wraps unexpected failures into InternalError after logging, while preserving specific handling for known Redis exceptions.
| """ | ||
|
|
||
| from api.core.request_models import QueryRequest | ||
| from api.core.result_models import ( |
There was a problem hiding this comment.
[major]: queryweaver_sdk.models re-exports from api.core.*, and the api.core import chain can transitively require server extras (e.g., graph/fastapi-adjacent modules). This can break the documented minimal SDK install by failing at import time.
Decouple shared dataclasses from api package side effects (or make api.core.__init__ minimal/lazy) so importing SDK models does not require server extras.
| Raises: | ||
| ValueError: If the database URL format is invalid. | ||
| """ | ||
| from api.core.schema_loader import load_database_sync |
There was a problem hiding this comment.
[major]: SDK methods route through api.core.schema_loader / api.core.text2sql, which can transitively import loader code requiring graphiti_core. With graphiti-core moved to optional extras, normal SDK operations can raise ModuleNotFoundError under pip install queryweaver.
Either keep required runtime deps in base install or isolate SDK execution paths from server-only imports; add packaging CI that validates minimal install can run core SDK operations.
|
|
||
| ```bash | ||
| # SDK only (minimal dependencies) | ||
| pip install queryweaver |
There was a problem hiding this comment.
[major]: README claims pip install queryweaver is sufficient, but current import/runtime paths indicate optional extras may still be required for documented SDK flows. This can produce runtime/import errors for users following quick start instructions.
Align docs with actual requirements, or preferably fix package boundaries so minimal install works as documented and verify this in CI.
|
|
||
| @pytest.mark.asyncio | ||
| @pytest.mark.requires_postgres | ||
| async def test_connect_postgres(self, falkordb_url, postgres_url, has_llm_key): |
There was a problem hiding this comment.
[major]: Deterministic SDK lifecycle tests (e.g., connect/get_schema/delete) are skipped when no LLM key is set, reducing regression coverage for non-LLM behavior and teardown reliability.
Remove LLM-key gating from tests that do not call query() and keep it only for truly LLM-dependent tests.
| ) | ||
|
|
||
|
|
||
| @pytest.fixture(scope="session") |
There was a problem hiding this comment.
[major]: Session-scoped mutable DB fixtures can make test outcomes order-dependent when tests write data, risking flaky counts/results across reruns or parallelization.
Use function-scoped isolation (or per-test schema/db names with teardown) for deterministic SDK integration assertions.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 20
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (5)
api/loaders/graph_loader.py (1)
19-32:⚠️ Potential issue | 🟡 MinorAnnotate the new optional graph DB handle.
load_to_graphnow exposes a dependency-injection parameter; keep the loader API typed. As per coding guidelines,api/**/*.py: Include type hints throughout Python code.Proposed fix
import json +from typing import Any import tqdm @@ - db=None, + db: Any | None = None, ) -> None:🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@api/loaders/graph_loader.py` around lines 19 - 32, The new optional DB handle parameter on load_to_graph needs a type annotation; update the function signature for load_to_graph to type the db parameter as Optional[...] (e.g., Optional[FalkorDB] or the concrete DB interface used by resolve_db) and import Optional from typing (or use typing.Optional), keeping the existing -> None return type; ensure any references to resolve_db(db) still type-check with the chosen DB type and update any stub/type imports (FalkorDB or GraphDB interface) so the annotation is valid.api/graph.py (1)
337-350:⚠️ Potential issue | 🟠 MajorPreserve column-only search results.
When the LLM returns only column descriptions,
main_taskscontains one task, buttables_by_columns_desbecomes[]; those matches are silently discarded.Proposed fix
# Execute the main embedding-based searches in parallel results = await asyncio.gather(*main_tasks) # Unpack results based on what tasks we ran - tables_des = results[0] if table_embeddings else [] - tables_by_columns_des = results[1] if (table_embeddings and column_embeddings) else [] + result_index = 0 + tables_des = [] + tables_by_columns_des = [] + if table_embeddings: + tables_des = results[result_index] + result_index += 1 + if column_embeddings: + tables_by_columns_des = results[result_index]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@api/graph.py` around lines 337 - 350, The current unpacking logic assumes fixed positions in results and drops column-only searches; instead iterate results in insertion order and assign them based on which tasks were appended to main_tasks: after building main_tasks with _find_tables and/or _find_tables_by_columns, consume results sequentially (use an index or pop from a list) and set tables_des when table_embeddings was requested and tables_by_columns_des when column_embeddings was requested so that a single-task run (column-only or table-only) maps correctly to the right variable (refer to main_tasks, _find_tables, _find_tables_by_columns, results, tables_des, tables_by_columns_des).api/loaders/mysql_loader.py (1)
469-482:⚠️ Potential issue | 🔴 CriticalFix the async generator consumption in
refresh_graph_schema.
MySQLLoader.load()is an async generator (containsyieldstatements), soawait MySQLLoader.load(...)at line 482 raisesTypeError. Additionally, deleting the graph at line 470 before reloading creates a risky window where the schema could be lost if reload fails. Useasync forto properly consume the async generator:Proposed fix
- success, message = await MySQLLoader.load(prefix, db_url, db=db) + success = False + message = "Failed to reload schema" + async for success, message in MySQLLoader.load(prefix, db_url, db=db): + pass🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@api/loaders/mysql_loader.py` around lines 469 - 482, The code calls await MySQLLoader.load(...) but MySQLLoader.load is an async generator, so replace the await with an async for loop to consume yielded items (e.g., "async for result in MySQLLoader.load(prefix, db_url, db=db):" and capture success/message from the yielded values); also avoid deleting the existing graph before a successful reload—first run the reload into a temporary/staging graph or buffer and only call graph.delete() and replace the old graph after the async for confirms success (adjust the logic in refresh_graph_schema to use the async-for consumption of MySQLLoader.load and perform the delete/replace only after a successful load).api/loaders/postgres_loader.py (2)
510-523:⚠️ Potential issue | 🔴 CriticalConsume the async generator instead of awaiting it.
PostgresLoader.load()yields progress tuples, soawait PostgresLoader.load(...)raisesTypeError. Since line 511 deletes the graph first, schema refresh can drop the graph and then fail to reload it.Proposed fix
- success, message = await PostgresLoader.load(prefix, db_url, db=db) + success = False + message = "Failed to reload schema" + async for success, message in PostgresLoader.load(prefix, db_url, db=db): + pass🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@api/loaders/postgres_loader.py` around lines 510 - 523, The code awaits PostgresLoader.load(...) even though load is an async generator and also deletes the graph before reloading which can cause reload failures; fix by first running the reload by iterating the async generator (use "async for progress in PostgresLoader.load(prefix, db_url, db=db)" to consume progress tuples and capture the final (success, message)), then only call await graph.delete() or delete the old graph after confirming success; reference resolve_db.select_graph(graph_id), graph.delete(), graph_id, prefix, and PostgresLoader.load to locate the code to change.
177-180:⚠️ Potential issue | 🟠 MajorUse
urlparsefor PostgreSQL database-name extraction.The current split-based extraction fails for valid PostgreSQL URLs with Unix-socket parameters, e.g.,
postgresql:///mydb?host=/var/run/postgresql, producing the scheme name instead of the database name. This causes incorrect graph name generation.Proposed fix
- # Extract database name from connection URL - db_name = connection_url.split('/')[-1] - if '?' in db_name: - db_name = db_name.split('?')[0] + # Extract database name from connection URL + parsed_url = urlparse(connection_url) + db_name = unquote(parsed_url.path.lstrip("/")) + if not db_name: + raise ValueError("PostgreSQL connection URL must include a database name")Note:
urlparseandunquoteare already imported.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@api/loaders/postgres_loader.py` around lines 177 - 180, Replace the brittle split-based extraction of db_name with urlparse/unquote: parse connection_url via urlparse(connection_url), take parsed.path, strip any leading '/' (parsed.path.lstrip('/')), unquote the result to handle percent-encoding, and assign that to db_name; keep a fallback (e.g., original split logic) only if the parsed path is empty. Update the code that sets db_name (the variable assigned from connection_url) to use this parsing logic so Unix-socket and query-parameter URLs yield the correct database name.
🧹 Nitpick comments (7)
.github/workflows/tests.yml (1)
121-124: Pinuvto a consistent version across jobs.Other steps in this workflow (and the wider repo) pin
version: "0.7.12"(see line 42 in this file, and.github/workflows/playwright.yml). This newsdk-testsjob uses"latest", which can drift and make CI non-reproducible.🔧 Suggested fix
- name: Install uv uses: astral-sh/setup-uv@d4b2f3b6ecc6e67c4457f6d3e41ec42d3d0fcb86 # v5.4.2 with: - version: "latest" + version: "0.7.12"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In @.github/workflows/tests.yml around lines 121 - 124, The "Install uv" step uses astral-sh/setup-uv with version: "latest" which can cause CI drift; change the with.version value to the pinned release used elsewhere (e.g., "0.7.12") so the sdk-tests job matches other workflows; locate the step that references astral-sh/setup-uv (the "Install uv" step) and replace version: "latest" with version: "0.7.12".Makefile (2)
83-86: Prefer--waitover a fixed sleep.
docker-compose.test.ymlalready declares healthchecks for falkordb/postgres/mysql, sodocker compose up -d --waitwill block until each service is healthy. The currentsleep 10is flaky on slow runners and unnecessarily slow on fast ones.🔧 Suggested fix
docker-test-services: ## Start all test services (FalkorDB + PostgreSQL + MySQL) - docker compose -f docker-compose.test.yml up -d - `@echo` "Waiting for services to be ready..." - `@sleep` 10 + docker compose -f docker-compose.test.yml up -d --wait --wait-timeout 120🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Makefile` around lines 83 - 86, The docker-test-services Makefile target uses a fixed sleep which is flaky; update the docker compose invocation in the docker-test-services recipe (the `docker compose -f docker-compose.test.yml up -d` line) to use `--wait` (i.e., `docker compose -f docker-compose.test.yml up -d --wait`) so it blocks until the healthchecks in docker-compose.test.yml pass, and remove the subsequent `sleep 10` (and optional "Waiting for services to be ready..." echo) to rely on the built-in wait behavior.
1-1: Add new targets to.PHONY.
test-sdk-quick(line 94) andtest-all(line 97) are new phony targets not listed in.PHONY, somakewill treat them as file targets if a file of the same name appears.🔧 Suggested fix
-.PHONY: help install test test-unit test-e2e test-e2e-headed lint format clean setup-dev build lint-frontend test-sdk docker-test-services docker-test-stop build-package +.PHONY: help install test test-unit test-e2e test-e2e-headed lint format clean setup-dev build lint-frontend test-sdk test-sdk-quick test-all docker-test-services docker-test-stop build-package🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Makefile` at line 1, The Makefile's .PHONY declaration is missing the newly added phony targets test-sdk-quick and test-all; update the .PHONY line (the .PHONY entry at the top of the Makefile) to include test-sdk-quick and test-all so make treats them as phony targets rather than files, ensuring the targets defined later in the file are always executed.queryweaver_sdk/client.py (1)
207-215: Redundant per-lineimport-outside-topleveldisables.The module-level
# pylint: disable=import-outside-toplevelat line 24 already covers the whole file, so the inline disables on lines 213 and 214 are noise. Drop them for consistency with the other lazy imports in this file (e.g. lines 91, 120, 156‑157, 200, 232, 251, 273‑274 which rely on the module-level directive).🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@queryweaver_sdk/client.py` around lines 207 - 215, Remove the redundant inline pylint disables on the lazy imports inside the async method list_databases: delete the trailing "# pylint: disable=import-outside-toplevel" comments from the two import statements that import list_databases from api.core.schema_loader and GENERAL_PREFIX from api.core.text2sql_common so the method uses the module-level disable already present; ensure the imports remain inside list_databases (preserving lazy import behavior) and run linters to confirm no new warnings.tests/test_sdk/conftest.py (1)
147-154:queryweaverfixture swallows close errors silently.If
await qw.close()raises (e.g. Redis pool already disconnected by a test), the exception becomes a teardown error that masks the actual test failure. Consider wrapping teardown in try/except-and-log, or useasync with QueryWeaver(...)inside the fixture for symmetric lifecycle handling:-@pytest.fixture -async def queryweaver(falkordb_url): - """Provide initialized QueryWeaver instance with proper teardown.""" - from queryweaver_sdk import QueryWeaver - - qw = QueryWeaver(falkordb_url=falkordb_url, user_id="test_user") - yield qw - await qw.close() +@pytest.fixture +async def queryweaver(falkordb_url): + """Provide initialized QueryWeaver instance with proper teardown.""" + from queryweaver_sdk import QueryWeaver + + async with QueryWeaver(falkordb_url=falkordb_url, user_id="test_user") as qw: + yield qw🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_sdk/conftest.py` around lines 147 - 154, The queryweaver fixture currently yields a QueryWeaver instance and calls await qw.close() in teardown which can raise and mask test failures; update the fixture to either use symmetric lifecycle by creating the instance inside an async context (async with QueryWeaver(falkordb_url=..., user_id="test_user") as qw: yield qw) or wrap the teardown call in a try/except that catches exceptions from qw.close() and logs them (referencing the fixture name queryweaver, the QueryWeaver constructor, and the qw.close() method) so close errors don't obscure test failures.tests/test_sdk/test_queryweaver.py (1)
24-30: Sync init tests leak aFalkorDBConnectionobject.
test_init_defaultsandtest_init_with_custom_user_idconstructQueryWeaverbut neverclose()it. Because_dbis lazy this won't open a real Redis pool today, but the pattern is fragile — if__init__ever starts eagerly probing the connection (e.g. for a version check), these tests will start leaking live connections. Useasync with(with@pytest.mark.asyncio) or explicitlyawait qw.close()to be future‑proof.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tests/test_sdk/test_queryweaver.py` around lines 24 - 30, The tests test_init_defaults and test_init_with_custom_user_id create QueryWeaver instances that may allocate a FalkorDBConnection; update each test to ensure the connection is closed by either turning the test into an async test (add `@pytest.mark.asyncio`) and using "async with QueryWeaver(...)" or by explicitly awaiting qw.close() at the end of the test; reference QueryWeaver and its close() method (or the context-manager protocol for QueryWeaver) so the FalkorDBConnection held in _db is properly cleaned up to prevent leaking live connections.queryweaver_sdk/connection.py (1)
83-86: Avoid direct access to FalkorDB client's internal connection attribute — use the publicaclose()API instead.When constructed via
FalkorDB(host=host, port=port)(lines 84),self._poolremainsNoneandclose()falls into theelifbranch that callsself._db.connection.aclose()(line 118). This directly accesses an undocumented internal attribute of the FalkorDB client. The public API providesaclose()method on the FalkorDB client itself, which properly handles the underlying connection cleanup.Consider either:
- Simpler fix: Replace
await self._db.connection.aclose()withawait self._db.aclose()(the public API), or- Unified approach: Build a connection URL from host/port and use
BlockingConnectionPool.from_url()for both code paths, ensuring consistent teardown viapool.disconnect().🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@queryweaver_sdk/connection.py` around lines 83 - 86, The close path currently calls the FalkorDB client's undocumented internal attribute (self._db.connection.aclose()); update the teardown to use the public API by replacing that call with await self._db.aclose() in the close() method of the connection class (or alternatively refactor construction to use BlockingConnectionPool.from_url() for both creation branches and call pool.disconnect() for a unified cleanup); locate references to FalkorDB, close(), self._db.connection.aclose(), and adjust them to use the public aclose() or the BlockingConnectionPool.from_url()/pool.disconnect() approach so the client is torn down via its supported API.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In @.github/workflows/tests.yml:
- Around line 136-143: The "Run SDK tests" job will run without OPENAI_API_KEY
for forked PRs causing LLM tests to fail; update the workflow so LLM-dependent
tests only run when the secret is present by making one of these changes: add a
job-level condition like if: ${{ secrets.OPENAI_API_KEY != '' }} to the "Run SDK
tests" job (or split the job into a non-LLM "test-sdk-quick" job and a separate
"Run SDK LLM tests" job gated by the secret or by a label/pull_request_target
with safeguards), or add an early pre-flight step in "Run SDK tests" that checks
OPENAI_API_KEY and fails with a clear message if empty; ensure the change
references the OPENAI_API_KEY env var used in the job so LLM tests are skipped
when the secret is unavailable.
In `@api/core/db_resolver.py`:
- Around line 12-18: Add explicit type hints to the resolve_db boundary: import
typing (e.g., "from typing import Optional, Any") and change the signature of
resolve_db to accept db: Optional[Any] = None and return Any, and annotate the
lazily imported api.extensions.db (imported as _default_db) as the same type so
the function returns a typed handle; update any local comments or # type: ignore
as needed when importing _default_db to satisfy the type checker.
In `@api/core/text2sql_common.py`:
- Around line 140-143: The function detect_destructive_operation currently uses
sql_query.strip().split()[0] which misclassifies queries with leading comments
(e.g., "-- comment\nDROP TABLE ..."); update detect_destructive_operation to
first remove leading SQL comments and whitespace before extracting the first
token: skip or strip initial single-line comments beginning with "--" and
initial block comments "/* ... */" (repeated until the first non-comment token),
then extract the verb, uppercase it, and check against DESTRUCTIVE_OPS while
safely handling empty/whitespace-only input.
- Around line 67-71: The code currently raises GraphNotFoundError for an
invalid/blank graph_id; change this to raise InvalidArgumentError instead so
callers treat it as a 400 Bad Request. Locate the block where graph_id is
normalized (graph_id = graph_id.strip()[:200]) and replace the
GraphNotFoundError instantiation with an InvalidArgumentError containing the
same or similar message ("Invalid graph_id, must be less than 200 characters.").
Ensure the exception import or reference for InvalidArgumentError is
present/used in text2sql_common.py so the new exception type is available.
In `@api/core/text2sql_sync.py`:
- Around line 412-416: The SDK sync path currently allows destructive queries
against general/demo graphs; update the destructive-op guard in the sync
handlers so destructive operations on general graphs are intercepted: in
query_database_sync() (around the block at 412-416) and in
execute_destructive_operation_sync() (the block around 452-480) change/add the
check to return _confirmation_required_result(ctx, analysis) when
analysis.is_destructive and is_general_graph(ctx.db.graph_id) so that
destructive operations against demo/general graphs require confirmation instead
of being executed.
- Around line 152-158: Add the same per-request override validation used in the
streaming path before constructing _ChatContext: validate getattr(chat_data,
'custom_model', None) follows the "vendor/model" pattern with vendor one of
(openai, anthropic, gemini, azure, ollama, cohere) and validate
getattr(chat_data, 'custom_api_key', None) is either None or has length >= 10;
if either check fails, raise a clear ValueError (or similar) so invalid
overrides are rejected before creating chat_ctx in text2sql_sync.py where
_ChatContext is instantiated.
In `@api/graph.py`:
- Around line 39-41: Annotate all parameters that default to None with explicit
Optional types and add the required typing imports; e.g., update
get_db_description(graph_id: str, db=None) -> tuple[str, str] to something like
get_db_description(graph_id: str, db: Optional[Any] = None) -> tuple[str, str]
(import Optional, Any from typing) and similarly change any parameter like
db_description: str = None to db_description: Optional[str] = None; apply the
same pattern to the other affected functions (the ones around lines 57-59,
73-76, 87, and 279-283) and ensure typing imports are added at the top of the
module.
In `@api/memory/graphiti_tool.py`:
- Around line 60-64: The constructor currently accepts an untyped db injection;
add explicit type hints to the public API by typing the db parameter and the
instance attribute: annotate GraphitiTool.__init__(self, user_id: str, graph_id:
str, db: Optional[DatabaseType] = None) (choose the concrete alias used in the
project, e.g., Optional[Mapping[str, Any]] or the project's DB protocol), and
annotate self._db and self.memory_db_name where appropriate; update resolve_db
return type or cast to the chosen DatabaseType and also apply the same typing
fix to the other constructor/initializer block referenced around lines 82-90 so
all injected db parameters and stored attributes are consistently typed
(reference symbols: __init__, self._db, self.memory_db_name, resolve_db,
FalkorDriver).
In `@docker-compose.test.yml`:
- Around line 5-13: The docker-compose service "falkordb" currently uses an
unpinned image ("falkordb/falkordb:latest"), which breaks test reproducibility;
change the "image" field in docker-compose.test.yml for the falkordb service to
a specific, tested tag (e.g., "falkordb/falkordb:<version>") instead of :latest,
and ensure that this pinned version matches the one used in CI or test docs
(update CI/test config if necessary).
In `@pyproject.toml`:
- Around line 120-123: The pyproject.toml pylint section ([tool.pylint.main])
currently sets ignore-patterns to ["test_.*\\.py", "conftest\\.py"], which
excludes test files from linting; remove or clear the ignore-patterns setting so
pylint runs over all Python files (update or delete the ignore-patterns key
under [tool.pylint.main], leaving max-line-length = 120 intact) to ensure tests
are included in the project-wide lint job.
- Around line 67-68: The console-script entry "queryweaver = \"api.index:main\""
points to a non-existent function; either add a main() function in api/index.py
that initializes and runs the CLI entrypoint, or change the entry to
"api.index:app" if the intended symbol is the existing ASGI/FastAPI app; update
the pyproject.toml entry accordingly and ensure the referenced symbol (main or
app) is exported at module level in api.index.
In `@queryweaver_sdk/client.py`:
- Around line 287-295: The current close() snapshots _pending_tasks once and
awaits them, missing any tasks scheduled during that await; change close() to
repeatedly drain until no tasks remain by looping: while self._pending_tasks:
take a snapshot (e.g., tasks = list(self._pending_tasks)), await
asyncio.gather(*tasks, return_exceptions=True) (optionally await
asyncio.sleep(0) to allow scheduling), and repeat until the set is empty, then
call await self._connection.close(); ensure you reference the close method, the
_pending_tasks collection, and _connection when making this change.
- Around line 254-285: The execute_confirmed method currently ignores
per-request LLM overrides and must accept and forward them: update
execute_confirmed's signature to accept custom_api_key and custom_model (or a
QueryRequest-style wrapper) and populate the ConfirmRequest with those fields
when constructing confirm_data; then pass the updated confirm_data through to
execute_destructive_operation_sync (and/or pass through any wrapper) so the same
per-request LLM credentials used by query()/ChatRequest are honored during
confirmation. Ensure to reference execute_confirmed, ConfirmRequest, and
execute_destructive_operation_sync when making the changes.
- Around line 103-121: connect_database currently calls load_database_sync
directly and therefore does not register any background work with the client's
task sink; to make behavior consistent with query() and execute_confirmed(),
wrap the load_database_sync call with the client's _bind_task_sink so any
background tasks get added to _pending_tasks and will be awaited in close();
locate the connect_database method and replace the direct await
load_database_sync(...) call with an awaited call into _bind_task_sink(...)
passing the same coroutine (referencing connect_database, _bind_task_sink,
_pending_tasks, close, and load_database_sync), or if the omission is
intentional add a clear inline comment in connect_database explaining why schema
loading should not be tracked.
In `@queryweaver_sdk/connection.py`:
- Around line 111-119: The close() method currently clears _db/_pool but allows
the db property to lazily recreate connections; add a boolean attribute _closed
on the Connection class, set self._closed = True at the end of close() (and
ensure close() is idempotent), and update the db property getter to check
self._closed and raise RuntimeError("Connection is closed") if true so
post-__aexit__ calls cannot recreate resources; reference the existing close()
method and the db property in your changes.
In `@README.md`:
- Around line 299-304: The README example shows using QueryWeaver instances
(QueryWeaver(...)) and then calling a.query(...) / b.query(...) without
connecting databases first; update the snippet to call the instance method
connect_database (e.g., await a.connect_database("sales") and await
b.connect_database("ops")) before invoking query, or alternatively add a clear
"illustrative only" note above the block; reference QueryWeaver,
connect_database, and query so the change is applied to the correct example.
In `@tests/test_sdk/conftest.py`:
- Around line 8-18: The pytest test markers registered in pytest_configure
currently add only dependency gates (requires_llm, requires_postgres,
requires_mysql); add guideline-required markers by registering "integration" and
"slow" (and optionally "e2e", "auth", "unit") via additional
config.addinivalue_line calls inside pytest_configure, then annotate the real
integration tests (e.g., module tests/test_sdk/test_queryweaver.py) with
module-level pytest marks such as pytestmark = [pytest.mark.integration,
pytest.mark.slow] so the suite is correctly categorized; update the
pytest_configure function and add the pytestmark declaration in
test_queryweaver.py.
- Around line 49-91: The try block creating and using psycopg2's conn and cursor
can leak resources if an exception is raised; update the block that opens conn
and cursor to use context managers (with conn, with conn.cursor()) or ensure
conn.close() and cursor.close() are called in a finally block so the DB
connection/session is always closed before calling pytest.skip; apply the same
fix pattern to the MySQL fixture that uses conn and cursor to prevent leaking
connections on exceptions.
In `@tests/test_sdk/test_queryweaver.py`:
- Around line 269-357: The TestModels class in
tests/test_sdk/test_queryweaver.py lacks a pytest marker; add a unit marker by
importing pytest (if not present) and either decorate the TestModels class with
`@pytest.mark.unit` or add pytestmark = [pytest.mark.unit] at the top of that
class to mark these pure dataclass tests as unit tests (refer to the TestModels
class name to locate the change).
- Around line 87-90: The test test_connect_invalid_url currently uses a broad
pytest.raises(Exception); update it to assert the documented exception by using
pytest.raises(ValueError) when calling
QueryWeaver.connect_database("invalid://url") so the test specifically expects
ValueError (the exception noted in QueryWeaver.connect_database docstring) and
will fail if a different exception type is raised.
---
Outside diff comments:
In `@api/graph.py`:
- Around line 337-350: The current unpacking logic assumes fixed positions in
results and drops column-only searches; instead iterate results in insertion
order and assign them based on which tasks were appended to main_tasks: after
building main_tasks with _find_tables and/or _find_tables_by_columns, consume
results sequentially (use an index or pop from a list) and set tables_des when
table_embeddings was requested and tables_by_columns_des when column_embeddings
was requested so that a single-task run (column-only or table-only) maps
correctly to the right variable (refer to main_tasks, _find_tables,
_find_tables_by_columns, results, tables_des, tables_by_columns_des).
In `@api/loaders/graph_loader.py`:
- Around line 19-32: The new optional DB handle parameter on load_to_graph needs
a type annotation; update the function signature for load_to_graph to type the
db parameter as Optional[...] (e.g., Optional[FalkorDB] or the concrete DB
interface used by resolve_db) and import Optional from typing (or use
typing.Optional), keeping the existing -> None return type; ensure any
references to resolve_db(db) still type-check with the chosen DB type and update
any stub/type imports (FalkorDB or GraphDB interface) so the annotation is
valid.
In `@api/loaders/mysql_loader.py`:
- Around line 469-482: The code calls await MySQLLoader.load(...) but
MySQLLoader.load is an async generator, so replace the await with an async for
loop to consume yielded items (e.g., "async for result in
MySQLLoader.load(prefix, db_url, db=db):" and capture success/message from the
yielded values); also avoid deleting the existing graph before a successful
reload—first run the reload into a temporary/staging graph or buffer and only
call graph.delete() and replace the old graph after the async for confirms
success (adjust the logic in refresh_graph_schema to use the async-for
consumption of MySQLLoader.load and perform the delete/replace only after a
successful load).
In `@api/loaders/postgres_loader.py`:
- Around line 510-523: The code awaits PostgresLoader.load(...) even though load
is an async generator and also deletes the graph before reloading which can
cause reload failures; fix by first running the reload by iterating the async
generator (use "async for progress in PostgresLoader.load(prefix, db_url,
db=db)" to consume progress tuples and capture the final (success, message)),
then only call await graph.delete() or delete the old graph after confirming
success; reference resolve_db.select_graph(graph_id), graph.delete(), graph_id,
prefix, and PostgresLoader.load to locate the code to change.
- Around line 177-180: Replace the brittle split-based extraction of db_name
with urlparse/unquote: parse connection_url via urlparse(connection_url), take
parsed.path, strip any leading '/' (parsed.path.lstrip('/')), unquote the result
to handle percent-encoding, and assign that to db_name; keep a fallback (e.g.,
original split logic) only if the parsed path is empty. Update the code that
sets db_name (the variable assigned from connection_url) to use this parsing
logic so Unix-socket and query-parameter URLs yield the correct database name.
---
Nitpick comments:
In @.github/workflows/tests.yml:
- Around line 121-124: The "Install uv" step uses astral-sh/setup-uv with
version: "latest" which can cause CI drift; change the with.version value to the
pinned release used elsewhere (e.g., "0.7.12") so the sdk-tests job matches
other workflows; locate the step that references astral-sh/setup-uv (the
"Install uv" step) and replace version: "latest" with version: "0.7.12".
In `@Makefile`:
- Around line 83-86: The docker-test-services Makefile target uses a fixed sleep
which is flaky; update the docker compose invocation in the docker-test-services
recipe (the `docker compose -f docker-compose.test.yml up -d` line) to use
`--wait` (i.e., `docker compose -f docker-compose.test.yml up -d --wait`) so it
blocks until the healthchecks in docker-compose.test.yml pass, and remove the
subsequent `sleep 10` (and optional "Waiting for services to be ready..." echo)
to rely on the built-in wait behavior.
- Line 1: The Makefile's .PHONY declaration is missing the newly added phony
targets test-sdk-quick and test-all; update the .PHONY line (the .PHONY entry at
the top of the Makefile) to include test-sdk-quick and test-all so make treats
them as phony targets rather than files, ensuring the targets defined later in
the file are always executed.
In `@queryweaver_sdk/client.py`:
- Around line 207-215: Remove the redundant inline pylint disables on the lazy
imports inside the async method list_databases: delete the trailing "# pylint:
disable=import-outside-toplevel" comments from the two import statements that
import list_databases from api.core.schema_loader and GENERAL_PREFIX from
api.core.text2sql_common so the method uses the module-level disable already
present; ensure the imports remain inside list_databases (preserving lazy import
behavior) and run linters to confirm no new warnings.
In `@queryweaver_sdk/connection.py`:
- Around line 83-86: The close path currently calls the FalkorDB client's
undocumented internal attribute (self._db.connection.aclose()); update the
teardown to use the public API by replacing that call with await
self._db.aclose() in the close() method of the connection class (or
alternatively refactor construction to use BlockingConnectionPool.from_url() for
both creation branches and call pool.disconnect() for a unified cleanup); locate
references to FalkorDB, close(), self._db.connection.aclose(), and adjust them
to use the public aclose() or the
BlockingConnectionPool.from_url()/pool.disconnect() approach so the client is
torn down via its supported API.
In `@tests/test_sdk/conftest.py`:
- Around line 147-154: The queryweaver fixture currently yields a QueryWeaver
instance and calls await qw.close() in teardown which can raise and mask test
failures; update the fixture to either use symmetric lifecycle by creating the
instance inside an async context (async with QueryWeaver(falkordb_url=...,
user_id="test_user") as qw: yield qw) or wrap the teardown call in a try/except
that catches exceptions from qw.close() and logs them (referencing the fixture
name queryweaver, the QueryWeaver constructor, and the qw.close() method) so
close errors don't obscure test failures.
In `@tests/test_sdk/test_queryweaver.py`:
- Around line 24-30: The tests test_init_defaults and
test_init_with_custom_user_id create QueryWeaver instances that may allocate a
FalkorDBConnection; update each test to ensure the connection is closed by
either turning the test into an async test (add `@pytest.mark.asyncio`) and using
"async with QueryWeaver(...)" or by explicitly awaiting qw.close() at the end of
the test; reference QueryWeaver and its close() method (or the context-manager
protocol for QueryWeaver) so the FalkorDBConnection held in _db is properly
cleaned up to prevent leaking live connections.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e4ea8b59-e4f1-420a-b1e0-1cab480fd0e5
⛔ Files ignored due to path filters (2)
app/package-lock.jsonis excluded by!**/package-lock.jsonuv.lockis excluded by!**/*.lock
📒 Files selected for processing (30)
.github/wordlist.txt.github/workflows/playwright.yml.github/workflows/pylint.yml.github/workflows/tests.ymlMakefileREADME.mdapi/core/__init__.pyapi/core/db_resolver.pyapi/core/request_models.pyapi/core/result_models.pyapi/core/schema_loader.pyapi/core/text2sql.pyapi/core/text2sql_common.pyapi/core/text2sql_sync.pyapi/graph.pyapi/loaders/graph_loader.pyapi/loaders/mysql_loader.pyapi/loaders/postgres_loader.pyapi/memory/graphiti_tool.pyapi/routes/database.pyapi/routes/graphs.pydocker-compose.test.ymlpyproject.tomlqueryweaver_sdk/__init__.pyqueryweaver_sdk/client.pyqueryweaver_sdk/connection.pyqueryweaver_sdk/models.pytests/test_sdk/__init__.pytests/test_sdk/conftest.pytests/test_sdk/test_queryweaver.py
💤 Files with no reviewable changes (1)
- api/routes/database.py
Comment on lines
+136
to
+143
| - name: Run SDK tests | ||
| env: | ||
| FALKORDB_URL: redis://localhost:6379 | ||
| TEST_POSTGRES_URL: postgresql://postgres:postgres@localhost:5432/testdb | ||
| TEST_MYSQL_URL: mysql://root:root@localhost:3306/testdb | ||
| OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }} | ||
| run: | | ||
| uv run python -m pytest tests/test_sdk/ -v |
Contributor
There was a problem hiding this comment.
SDK tests will fail on PRs from forks due to missing OPENAI_API_KEY.
secrets.OPENAI_API_KEY is not available to workflows triggered by pull_request from forked repositories, so this job will run with an empty key and the tests that hit the LLM will fail. Also, if an internal developer forgets to set the secret, the only signal is a late test failure.
Consider either:
- Skipping this job when the secret is empty (e.g.,
if: ${{ secrets.OPENAI_API_KEY != '' }}at job level, or a pre-flight step that fails with a clearer message). - Splitting into a non-LLM subset that can run on all PRs (similar to
test-sdk-quickin the Makefile) and gating the LLM subset behind a label orpull_request_targetwith appropriate safeguards.
🧰 Tools
🪛 Checkov (3.2.519)
[medium] 139-140: Basic Auth Credentials
(CKV_SECRET_4)
[medium] 140-141: Basic Auth Credentials
(CKV_SECRET_4)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In @.github/workflows/tests.yml around lines 136 - 143, The "Run SDK tests" job
will run without OPENAI_API_KEY for forked PRs causing LLM tests to fail; update
the workflow so LLM-dependent tests only run when the secret is present by
making one of these changes: add a job-level condition like if: ${{
secrets.OPENAI_API_KEY != '' }} to the "Run SDK tests" job (or split the job
into a non-LLM "test-sdk-quick" job and a separate "Run SDK LLM tests" job gated
by the secret or by a label/pull_request_target with safeguards), or add an
early pre-flight step in "Run SDK tests" that checks OPENAI_API_KEY and fails
with a clear message if empty; ensure the change references the OPENAI_API_KEY
env var used in the job so LLM tests are skipped when the secret is unavailable.
Comment on lines
+12
to
+18
| def resolve_db(db=None): | ||
| """Return the given ``db`` handle, or lazily import the server default.""" | ||
| if db is not None: | ||
| return db | ||
| # pylint: disable=import-outside-toplevel | ||
| from api.extensions import db as _default_db | ||
| return _default_db |
Contributor
There was a problem hiding this comment.
Add type hints to the resolver boundary.
resolve_db is now shared core API; typing the injected handle keeps the new dependency-injection boundary explicit. As per coding guidelines, api/**/*.py: Include type hints throughout Python code.
Proposed fix
+from typing import Any
+
+
-def resolve_db(db=None):
+def resolve_db(db: Any | None = None) -> Any:
"""Return the given ``db`` handle, or lazily import the server default."""
if db is not None:
return db🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@api/core/db_resolver.py` around lines 12 - 18, Add explicit type hints to the
resolve_db boundary: import typing (e.g., "from typing import Optional, Any")
and change the signature of resolve_db to accept db: Optional[Any] = None and
return Any, and annotate the lazily imported api.extensions.db (imported as
_default_db) as the same type so the function returns a typed handle; update any
local comments or # type: ignore as needed when importing _default_db to satisfy
the type checker.
Comment on lines
+67
to
+71
| graph_id = graph_id.strip()[:200] | ||
| if not graph_id: | ||
| raise GraphNotFoundError( | ||
| "Invalid graph_id, must be less than 200 characters." | ||
| ) |
Contributor
There was a problem hiding this comment.
Raise InvalidArgumentError for invalid graph IDs.
An empty/blank graph_id is a bad request, not a missing graph. Several callers catch InvalidArgumentError for 400 handling, so raising GraphNotFoundError can produce wrong SDK/HTTP behavior.
Proposed fix
graph_id = graph_id.strip()[:200]
if not graph_id:
- raise GraphNotFoundError(
- "Invalid graph_id, must be less than 200 characters."
- )
+ raise InvalidArgumentError("graph_id cannot be empty")🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@api/core/text2sql_common.py` around lines 67 - 71, The code currently raises
GraphNotFoundError for an invalid/blank graph_id; change this to raise
InvalidArgumentError instead so callers treat it as a 400 Bad Request. Locate
the block where graph_id is normalized (graph_id = graph_id.strip()[:200]) and
replace the GraphNotFoundError instantiation with an InvalidArgumentError
containing the same or similar message ("Invalid graph_id, must be less than 200
characters."). Ensure the exception import or reference for InvalidArgumentError
is present/used in text2sql_common.py so the new exception type is available.
Comment on lines
+140
to
+143
| def detect_destructive_operation(sql_query: str) -> tuple[str, bool]: | ||
| """Return ``(sql_type, is_destructive)`` for a SQL statement.""" | ||
| sql_type = sql_query.strip().split()[0].upper() if sql_query else "" | ||
| return sql_type, sql_type in DESTRUCTIVE_OPS |
Contributor
There was a problem hiding this comment.
Harden destructive-SQL detection against leading comments.
strip().split()[0] classifies -- comment\nDROP TABLE ... as non-destructive, bypassing confirmation/demo-graph protections before execution. Strip leading SQL comments before extracting the statement verb.
Proposed direction
+import re
+
+_LEADING_SQL_COMMENTS = re.compile(
+ r"^\s*(?:(?:--[^\n]*(?:\n|$))|(?:/\*.*?\*/))*",
+ re.DOTALL,
+)
+
def detect_destructive_operation(sql_query: str) -> tuple[str, bool]:
"""Return ``(sql_type, is_destructive)`` for a SQL statement."""
- sql_type = sql_query.strip().split()[0].upper() if sql_query else ""
+ normalized = _LEADING_SQL_COMMENTS.sub("", sql_query or "").strip()
+ sql_type = normalized.split()[0].upper() if normalized else ""
return sql_type, sql_type in DESTRUCTIVE_OPS🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@api/core/text2sql_common.py` around lines 140 - 143, The function
detect_destructive_operation currently uses sql_query.strip().split()[0] which
misclassifies queries with leading comments (e.g., "-- comment\nDROP TABLE
..."); update detect_destructive_operation to first remove leading SQL comments
and whitespace before extracting the first token: skip or strip initial
single-line comments beginning with "--" and initial block comments "/* ... */"
(repeated until the first non-comment token), then extract the verb, uppercase
it, and check against DESTRUCTIVE_OPS while safely handling
empty/whitespace-only input.
Comment on lines
+152
to
+158
| chat_ctx = _ChatContext( | ||
| queries_history=queries_history, | ||
| result_history=result_history, | ||
| instructions=instructions, | ||
| custom_api_key=getattr(chat_data, 'custom_api_key', None), | ||
| custom_model=getattr(chat_data, 'custom_model', None), | ||
| ) |
Contributor
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Verify whether SDK/request models already validate custom_model and custom_api_key.
rg -n -C3 "custom_model|custom_api_key|vendor/model|API key is too short|Unsupported vendor"Repository: FalkorDB/QueryWeaver
Length of output: 20746
🏁 Script executed:
cat -n api/core/request_models.py | head -100Repository: FalkorDB/QueryWeaver
Length of output: 1505
🏁 Script executed:
sed -n '140,170p' api/core/text2sql_sync.pyRepository: FalkorDB/QueryWeaver
Length of output: 1104
🏁 Script executed:
sed -n '110,165p' api/core/text2sql_sync.pyRepository: FalkorDB/QueryWeaver
Length of output: 2030
🏁 Script executed:
rg -n "_initialize_query_context|validate_and_truncate_chat" api/core/text2sql_sync.pyRepository: FalkorDB/QueryWeaver
Length of output: 279
🏁 Script executed:
sed -n '335,365p' api/core/text2sql_sync.pyRepository: FalkorDB/QueryWeaver
Length of output: 1202
🏁 Script executed:
rg -n "query_database_sync" api/Repository: FalkorDB/QueryWeaver
Length of output: 170
🏁 Script executed:
sed -n '170,190p' queryweaver_sdk/client.pyRepository: FalkorDB/QueryWeaver
Length of output: 787
🏁 Script executed:
rg -n "ChatRequest|query_database_sync" api/routes/ api/api.py 2>/dev/null | head -30Repository: FalkorDB/QueryWeaver
Length of output: 302
🏁 Script executed:
sed -n '130,200p' api/routes/graphs.pyRepository: FalkorDB/QueryWeaver
Length of output: 2943
🏁 Script executed:
rg -n "async def query_database" api/core/Repository: FalkorDB/QueryWeaver
Length of output: 322
🏁 Script executed:
sed -n '169,220p' api/core/text2sql.pyRepository: FalkorDB/QueryWeaver
Length of output: 2518
Mirror per-request override validation in the SDK path.
The streaming path (api/core/text2sql.py, lines 202–216) validates custom_model format and custom_api_key length, but the sync path (api/core/text2sql_sync.py, lines 152–158) passes them directly to agents without validation. Add the same validation in the sync path initializer before constructing _ChatContext.
custom_model: must be "vendor/model" with vendor in (openai, anthropic, gemini, azure, ollama, cohere)
custom_api_key: must be ≥10 characters if provided
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@api/core/text2sql_sync.py` around lines 152 - 158, Add the same per-request
override validation used in the streaming path before constructing _ChatContext:
validate getattr(chat_data, 'custom_model', None) follows the "vendor/model"
pattern with vendor one of (openai, anthropic, gemini, azure, ollama, cohere)
and validate getattr(chat_data, 'custom_api_key', None) is either None or has
length >= 10; if either check fails, raise a clear ValueError (or similar) so
invalid overrides are rejected before creating chat_ctx in text2sql_sync.py
where _ChatContext is instantiated.
Comment on lines
+299
to
+304
| ```python | ||
| async with QueryWeaver(falkordb_url="redis://host-a:6379", user_id="tenant_a") as a, \ | ||
| QueryWeaver(falkordb_url="redis://host-b:6379", user_id="tenant_b") as b: | ||
| await a.query("sales", "Show top customers") | ||
| await b.query("ops", "Count open tickets") | ||
| ``` |
Contributor
There was a problem hiding this comment.
Illustrative snippet passes un-connected database IDs.
The example calls a.query("sales", ...) / b.query("ops", ...) without a prior connect_database, which will fail against the real SDK. Consider either adding a connect step or explicitly labeling the block as illustrative so users don't copy it as-is.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@README.md` around lines 299 - 304, The README example shows using QueryWeaver
instances (QueryWeaver(...)) and then calling a.query(...) / b.query(...)
without connecting databases first; update the snippet to call the instance
method connect_database (e.g., await a.connect_database("sales") and await
b.connect_database("ops")) before invoking query, or alternatively add a clear
"illustrative only" note above the block; reference QueryWeaver,
connect_database, and query so the change is applied to the correct example.
Comment on lines
+8
to
+18
| def pytest_configure(config): | ||
| """Configure pytest with custom markers.""" | ||
| config.addinivalue_line( | ||
| "markers", "requires_llm: mark test as requiring LLM API key" | ||
| ) | ||
| config.addinivalue_line( | ||
| "markers", "requires_postgres: mark test as requiring PostgreSQL" | ||
| ) | ||
| config.addinivalue_line( | ||
| "markers", "requires_mysql: mark test as requiring MySQL" | ||
| ) |
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Add guideline-required test markers.
The registered markers (requires_llm, requires_postgres, requires_mysql) are dependency gates and are fine, but per coding guidelines tests should additionally be marked with one of e2e, slow, auth, integration, unit. Since this whole suite talks to real FalkorDB/Postgres/MySQL + an LLM, these tests are integration (arguably e2e + slow). Consider registering those markers here (if not already in root conftest.py/pytest.ini) and applying them in tests/test_sdk/test_queryweaver.py — e.g. via pytestmark = [pytest.mark.integration, pytest.mark.slow] at the module level.
As per coding guidelines: "Mark test functions with pytest markers: e2e, slow, auth, integration, unit".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_sdk/conftest.py` around lines 8 - 18, The pytest test markers
registered in pytest_configure currently add only dependency gates
(requires_llm, requires_postgres, requires_mysql); add guideline-required
markers by registering "integration" and "slow" (and optionally "e2e", "auth",
"unit") via additional config.addinivalue_line calls inside pytest_configure,
then annotate the real integration tests (e.g., module
tests/test_sdk/test_queryweaver.py) with module-level pytest marks such as
pytestmark = [pytest.mark.integration, pytest.mark.slow] so the suite is
correctly categorized; update the pytest_configure function and add the
pytestmark declaration in test_queryweaver.py.
Comment on lines
+49
to
+91
| try: | ||
| import psycopg2 | ||
| conn = psycopg2.connect(url) | ||
| cursor = conn.cursor() | ||
|
|
||
| # Create test tables (DROP + CREATE ensures a clean slate) | ||
| cursor.execute(""" | ||
| DROP TABLE IF EXISTS orders CASCADE; | ||
| DROP TABLE IF EXISTS customers CASCADE; | ||
|
|
||
| CREATE TABLE customers ( | ||
| id SERIAL PRIMARY KEY, | ||
| name VARCHAR(100) NOT NULL, | ||
| email VARCHAR(100) UNIQUE, | ||
| city VARCHAR(100) | ||
| ); | ||
|
|
||
| CREATE TABLE orders ( | ||
| id SERIAL PRIMARY KEY, | ||
| customer_id INTEGER REFERENCES customers(id), | ||
| product VARCHAR(100), | ||
| amount DECIMAL(10,2), | ||
| order_date DATE | ||
| ); | ||
|
|
||
| -- Insert test data (UNIQUE on email allows ON CONFLICT) | ||
| INSERT INTO customers (name, email, city) VALUES | ||
| ('Alice Smith', '[email protected]', 'New York'), | ||
| ('Bob Jones', '[email protected]', 'Los Angeles'), | ||
| ('Carol White', '[email protected]', 'New York') | ||
| ON CONFLICT (email) DO NOTHING; | ||
|
|
||
| INSERT INTO orders (customer_id, product, amount, order_date) VALUES | ||
| (1, 'Widget', 29.99, '2024-01-15'), | ||
| (1, 'Gadget', 49.99, '2024-01-20'), | ||
| (2, 'Widget', 29.99, '2024-02-01'); | ||
| """) | ||
| conn.commit() | ||
| conn.close() | ||
| except Exception as e: | ||
| pytest.skip(f"PostgreSQL not available: {e}") | ||
|
|
||
| return url |
Contributor
There was a problem hiding this comment.
Leaked psycopg2 connection/cursor on exception.
If any of the DDL/DML statements raise, control jumps to the except block and conn.close() is never called — the connection (and its server-side session) leaks until garbage-collected. Same pattern in the MySQL fixture (lines 111‑144).
🛠 Proposed fix using a `with` block
- try:
- import psycopg2
- conn = psycopg2.connect(url)
- cursor = conn.cursor()
-
- # Create test tables (DROP + CREATE ensures a clean slate)
- cursor.execute("""
+ try:
+ import psycopg2
+ with psycopg2.connect(url) as conn:
+ with conn.cursor() as cursor:
+ cursor.execute("""
DROP TABLE IF EXISTS orders CASCADE;
...
- """)
- conn.commit()
- conn.close()
+ """)
+ conn.commit()
except Exception as e:
pytest.skip(f"PostgreSQL not available: {e}")🧰 Tools
🪛 Ruff (0.15.10)
[warning] 88-88: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_sdk/conftest.py` around lines 49 - 91, The try block creating and
using psycopg2's conn and cursor can leak resources if an exception is raised;
update the block that opens conn and cursor to use context managers (with conn,
with conn.cursor()) or ensure conn.close() and cursor.close() are called in a
finally block so the DB connection/session is always closed before calling
pytest.skip; apply the same fix pattern to the MySQL fixture that uses conn and
cursor to prevent leaking connections on exceptions.
Comment on lines
+87
to
+90
| @pytest.mark.asyncio | ||
| async def test_connect_invalid_url(self, queryweaver): | ||
| with pytest.raises(Exception): # noqa: B017 - expect InvalidArgumentError or subclass | ||
| await queryweaver.connect_database("invalid://url") |
Contributor
There was a problem hiding this comment.
Narrow pytest.raises(Exception) to the expected error type.
The # noqa: B017 suppression and the comment both acknowledge this is suboptimal. If the SDK is documented to raise ValueError (see QueryWeaver.connect_database docstring: "Raises: ValueError"), tighten this to pytest.raises(ValueError) so the test fails if the SDK starts raising e.g. a bare ConnectionError for an unparsable URL — which would otherwise slip past this test.
- async def test_connect_invalid_url(self, queryweaver):
- with pytest.raises(Exception): # noqa: B017 - expect InvalidArgumentError or subclass
- await queryweaver.connect_database("invalid://url")
+ async def test_connect_invalid_url(self, queryweaver):
+ with pytest.raises(ValueError):
+ await queryweaver.connect_database("invalid://url")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| @pytest.mark.asyncio | |
| async def test_connect_invalid_url(self, queryweaver): | |
| with pytest.raises(Exception): # noqa: B017 - expect InvalidArgumentError or subclass | |
| await queryweaver.connect_database("invalid://url") | |
| `@pytest.mark.asyncio` | |
| async def test_connect_invalid_url(self, queryweaver): | |
| with pytest.raises(ValueError): | |
| await queryweaver.connect_database("invalid://url") |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_sdk/test_queryweaver.py` around lines 87 - 90, The test
test_connect_invalid_url currently uses a broad pytest.raises(Exception); update
it to assert the documented exception by using pytest.raises(ValueError) when
calling QueryWeaver.connect_database("invalid://url") so the test specifically
expects ValueError (the exception noted in QueryWeaver.connect_database
docstring) and will fail if a different exception type is raised.
Comment on lines
+269
to
+357
| class TestModels: | ||
| """Dataclass serialization / defaults.""" | ||
|
|
||
| def test_query_result_to_dict(self): | ||
| result = QueryResult( | ||
| sql_query="SELECT * FROM customers", | ||
| results=[{"id": 1, "name": "Alice"}], | ||
| ai_response="Found 1 customer", | ||
| metadata=QueryMetadata( | ||
| confidence=0.95, | ||
| is_destructive=False, | ||
| requires_confirmation=False, | ||
| execution_time=0.5, | ||
| ), | ||
| ) | ||
|
|
||
| d = result.to_dict() | ||
| assert d["sql_query"] == "SELECT * FROM customers" | ||
| assert d["confidence"] == 0.95 | ||
| assert d["results"] == [{"id": 1, "name": "Alice"}] | ||
| assert d["ai_response"] == "Found 1 customer" | ||
| assert d["is_destructive"] is False | ||
| assert d["requires_confirmation"] is False | ||
| assert d["execution_time"] == 0.5 | ||
|
|
||
| def test_schema_result_to_dict(self): | ||
| result = SchemaResult( | ||
| nodes=[{"id": "customers", "name": "customers"}], | ||
| links=[{"source": "orders", "target": "customers"}], | ||
| ) | ||
|
|
||
| d = result.to_dict() | ||
| assert d["nodes"][0]["name"] == "customers" | ||
| assert d["links"][0]["source"] == "orders" | ||
| assert d["links"][0]["target"] == "customers" | ||
|
|
||
| def test_database_connection_to_dict(self): | ||
| result = DatabaseConnection( | ||
| database_id="testdb", | ||
| success=True, | ||
| tables_loaded=5, | ||
| message="Connected successfully", | ||
| ) | ||
|
|
||
| d = result.to_dict() | ||
| assert d["database_id"] == "testdb" | ||
| assert d["success"] is True | ||
| assert d["tables_loaded"] == 5 | ||
| assert d["message"] == "Connected successfully" | ||
|
|
||
| def test_query_result_default_values(self): | ||
| result = QueryResult( | ||
| sql_query="SELECT 1", | ||
| results=[], | ||
| ai_response="Test", | ||
| metadata=QueryMetadata(confidence=0.8), | ||
| ) | ||
|
|
||
| assert result.is_destructive is False | ||
| assert result.requires_confirmation is False | ||
| assert result.execution_time == 0.0 | ||
| assert result.is_valid is True | ||
| assert result.missing_information == "" | ||
| assert result.ambiguities == "" | ||
| assert result.explanation == "" | ||
|
|
||
| def test_database_connection_failure(self): | ||
| result = DatabaseConnection( | ||
| database_id="", | ||
| success=False, | ||
| tables_loaded=0, | ||
| message="Connection refused", | ||
| ) | ||
|
|
||
| d = result.to_dict() | ||
| assert d["database_id"] == "" | ||
| assert d["success"] is False | ||
| assert d["tables_loaded"] == 0 | ||
| assert "refused" in d["message"].lower() | ||
|
|
||
| def test_query_request_custom_model_fields(self): | ||
| """custom_api_key/custom_model are threaded through to agents.""" | ||
| req = QueryRequest( | ||
| question="test", | ||
| custom_api_key="sk-test-123", | ||
| custom_model="openai/gpt-4.1", | ||
| ) | ||
| assert req.custom_api_key == "sk-test-123" | ||
| assert req.custom_model == "openai/gpt-4.1" |
Contributor
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
Pure‑unit model tests should be explicitly marked unit.
TestModels (and only TestModels) doesn't require FalkorDB/Postgres/MySQL/LLM — it's a pure dataclass test. Per coding guidelines every test function should carry one of the standard markers (e2e, slow, auth, integration, unit). Mark this class with @pytest.mark.unit and the rest of the classes (which talk to real services) with @pytest.mark.integration (+ @pytest.mark.slow where LLM is involved) — ideally via a module/class pytestmark = [...] to avoid per-test noise.
As per coding guidelines: "Mark test functions with pytest markers: e2e, slow, auth, integration, unit".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/test_sdk/test_queryweaver.py` around lines 269 - 357, The TestModels
class in tests/test_sdk/test_queryweaver.py lacks a pytest marker; add a unit
marker by importing pytest (if not present) and either decorate the TestModels
class with `@pytest.mark.unit` or add pytestmark = [pytest.mark.unit] at the top
of that class to mark these pure dataclass tests as unit tests (refer to the
TestModels class name to locate the change).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Supersedes #384. Fixes the correctness bugs surfaced in review, eliminates the
api.extensions.dbglobal mutation, deduplicates orchestration between the streaming and SDK text2sql paths, and makesqueryweaver_sdkstructurally ready forpip install queryweaver. Bumps SDK to 0.2.0 (one breaking change:database_idis now un-prefixed).What changed
Correctness
schema_loader.load_database_syncreturns un-prefixeddatabase_id— fixes double-namespacing when the value is passed back intoquery/delete/get_schema.exceptclause broadened to match streaming path, sopsycopg2.Error/pymysql.Errorfrom bad SQL actually trigger healing.tables_loadedsubstring-counting heuristic (it was double-counting).urllib.parse.urlparsefor DB-name extraction instead of brittlesplit.uv sync --locked --all-extras.Architecture — no more hidden globals
api/core/db_resolver.py::resolve_db()— lazy server-default fallback.api/graph.py(find,get_db_description,get_user_rules,set_user_rules),api/memory/graphiti_tool.py::MemoryTool,api/loaders/{graph,postgres,mysql}_loader.py— all accept optionaldb=Noneparam. Routes unchanged (use resolver fallback, keepfrom api.extensions import db).QueryWeaver.__init__no longer mutatesapi.extensions.db. Multiple SDK instances in one process are now independent.api.extensionsload (no side-effect FalkorDB connect at import time).Deduplication — one text2sql implementation
api/core/text2sql_common.pygained 6 shared helpers:execute_with_healing,refresh_schema_if_modified,quote_identifiers_from_graph,save_memory_background,format_ai_response,build_destructive_confirmation_message.MESSAGE_DELIMITERmoved to the common module (was duplicated in 3 files).text2sql_sync.pyrewritten to call the helpers; streamingtext2sql.pyuses the leaf helpers (inline healing + schema-refresh remain for now — full unification needs queue-based emit, deferred).Layering
api/core/result_models.py+api/core/request_models.py(new) now own the dataclasses.queryweaver_sdk.modelsis a thin re-export shim. Fixes the invertedapi → queryweaver_sdkimport.SDK additions
QueryRequest.custom_api_key/custom_model— per-request LLM overrides (parity with streaming path).QueryWeaver.close()awaits in-flight memory writes via a contextvar-scoped task sink._graph_name, unused_ChatContext.use_user_rules).Tests
tests/test_sdk/test_queryweaver.pyto useasync with QueryWeaver(...)for proper teardown — no more skip-on-RuntimeError wrappers.test_two_instances_isolated(regression guard for the global-mutation fix).test_query_with_history.Breaking changes (SDK 0.1.0 → 0.2.0)
DatabaseConnection.database_idis now the un-prefixed database name (e.g."mydb", not"user_mydb"). Pass this value directly back intoquery,delete_database,get_schema,refresh_schema— the SDK applies namespacing internally.Verification
test_two_instances_isolated). Pylint 10.00/10.uv buildproducesqueryweaver-0.2.0-py3-none-any.whl; installed cleanly in a Python 3.13 fresh venv; import does not loadapi.extensions; two instances have independent_dbhandles.connect_database→get_schema→ natural-language query → SQL (SELECT id, name FROM customers WHERE city = 'New York') → correct filtered results (2 NY customers) →delete_database. All 13 checks passed.uv run uvicorn api.index:app) boots cleanly; 21 routes registered; HSTS middleware active; unchanged behavior from the streaming path's perspective.Test plan
sdk-package-cleanupOPENAI_API_KEY+ PG/MySQL services)🤖 Generated with Claude Code
Summary by CodeRabbit
Release Notes
New Features
Documentation
Chores