diff --git a/WRENAI_LOCAL_ASK_HANDOFF.md b/WRENAI_LOCAL_ASK_HANDOFF.md new file mode 100644 index 0000000000..80d3277e03 --- /dev/null +++ b/WRENAI_LOCAL_ASK_HANDOFF.md @@ -0,0 +1,256 @@ +# WrenAI Ask Schema Grounding Handoff + +Date: 2026-09-01 + +## Current Branch + +- Workspace: `D:\WrenAI` +- Branch: `organization/ask-schema-grounding-20260820` +- Local HEAD before this handoff commit: `0ff1e6e23 Improve Ask schema grounding` +- Direct push from `D:\WrenAI` was rejected because the local branch was behind the remote branch. +- The push was prepared from clean integration worktree `D:\WrenAI-push-ask-20260901`. +- Integration base: remote `organization/ask-schema-grounding-20260820` at `2f4d8f360 Tighten schema-driven Ask semantic validation`. +- Local commits replayed onto the remote tip: `0ff1e6e23` and `5179591df`. +- Merge-resolution fixes were applied in the integration worktree and validated before push. +- Do not use `.codex-tmp` as runtime source. The AI service was restarted from `D:\WrenAI\wren-ai-service`. + +## Goal Continued + +Continue the WrenAI Ask schema-grounding work for CWPay and CW_GL while preserving Orders and PCB_DB behavior. The focus of this continuation was Ask speed and correctness on large schemas: + +- add timing logs for Ask stages +- identify slow stages +- reduce unnecessary LLM calls +- cache schema-derived metadata and ranking inputs +- keep SQL generation schema-verified +- return clear unsupported results instead of hallucinated SQL +- validate CWPay, CW_GL, Orders, and PCB_DB + +## What Was Done + +### Timing and Observability + +Ask timing logs now cover the key path across UI and AI service: + +- frontend request +- task creation +- schema retrieval support context +- schema retrieval +- candidate ranking +- LLM intent generation +- SQL generation / deterministic fast path +- SQL validation +- SQL execution +- answer formatting request and polling +- cancel request and cancellation point + +Relevant logs reviewed: + +- `.codex-tmp\ai-dev-after13.err.log` +- `.codex-tmp\ai-dev-after14.err.log` +- `.codex-tmp\ai-dev-after15.err.log` +- `.codex-tmp\ai-dev-after17.err.log` +- `.codex-tmp\ai-dev-after19.err.log` +- `.codex-tmp\ai-dev-after20.err.log` + +Final stage summary from `.codex-tmp\ai-dev-after20.err.log`: + +- `sql_generation_fast_path`: avg about 1.4s, max about 2.8s +- `schema_retrieval`: avg about 0.5s, max about 1.5s +- `schema_retrieval_support_context`: avg about 19ms +- task creation and frontend markers were effectively negligible +- no LLM intent generation was used by the final Ask validation tasks except the explicit cancel test + +There is one non-Ask background outlier in the same AI log: `schema_retrieval_total` around 142s from a column-pruning/question-recommendation path. It was not part of the final Ask task timing set. + +### Performance Improvements + +- Added deterministic schema-driven fast paths for clear count, grouping, top-N, latest, date/month bucket, distribution, listing, and same-thread group-result follow-up shapes. +- Added pre-intent unsupported handling for simple analytics requests when active-project schema coverage is missing. +- Added schema metadata/token/index caching for table/column/description-derived matching. +- Changed large-schema retrieval to broad candidate gathering and small top-K reranking. +- Changed generation-context limiting to skip oversized candidates and continue looking for smaller valid candidates within the token budget. +- Tightened same-thread follow-up grounding to use compact latest verified SQL/table identifiers instead of accumulating stale or oversized history. +- Preserved schema validation and dry-run validation. Unsupported cases return `NO_RELEVANT_SQL` instead of invalid SQL. +- Fixed SQL literal offset handling so extracted filter values are validated against the right columns. +- Fixed top record/listing behavior so row-level date questions order by verified date columns instead of unrelated numeric fields. +- Tightened answer formatting prompts to use executed SQL result rows only and not invent analysis, values, code, or examples. +- Added transient MSSQL deadlock retry around Ibis query execution. + +### Count Shape Fix + +The CWPay question `How many invoice records are there?` now returns a scalar aggregate: + +```sql +SELECT + COUNT(*) AS "record_count" +FROM + "dbo_View_Open_Invoices" +``` + +The result shape is one column, `record_count`, and one row. It no longer returns invoice detail columns for that simple count shape. + +### Same-Thread Follow-Ups + +Follow-up questions now retrieve exact prior verified tables from the latest SQL and use a compact grounding query. This fixed the slow same-thread path that previously fell back to LLM calls on large schemas. + +Final targeted follow-up waits: + +- CWPay: about 3.0s +- CW_GL: about 3.0s +- Orders: about 3.1s +- PCB_DB: about 2.0s + +## Validation Results + +Final artifacts: + +- `.codex-tmp\ask_perf_benchmark_after_final.json` +- `.codex-tmp\resume_schema_grounding_validation_after_final.json` +- `.codex-tmp\cancel_check_after20.json` + +Before/after benchmark: + +- Before avg Ask wait: 61,818.6ms +- Before max Ask wait: 143,366ms +- After avg Ask wait: 3,388ms +- After max Ask wait: 5,056ms + +Observed CWPay examples: + +- `How many invoice records are there?` + - before: 84,653ms + - after: 3,031ms + - final SQL uses scalar `COUNT(*) AS "record_count"` +- `Show invoices by business unit` + - before: 112,881ms + - after: 3,036ms + - final SQL groups by verified `bunit` + +Final full validation: + +- Project checks: 4/4 +- Ask cases: 18/18 +- Same-thread follow-ups: 4/4 + +Project checks: + +- CWPay: 364 datasource tables, 350 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` +- CW_GL: 223 datasource tables, 223 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` +- Orders: 103 datasource tables, 101 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` +- PCB_DB: 76 datasource tables, 68 models, modeling page passed, preview 3/3, deploy `SUCCESS`, sync `SYNCRONIZED` + +Regression status: + +- Orders by customer: passed +- Orders revenue/date/month/latest families: passed +- PCB_DB repairs/status/priority/month/latest families: passed +- Unsupported cross-project questions: passed with `NO_RELEVANT_SQL` +- No observed schema leakage between projects +- No observed `Failed to create asking task` +- No observed hallucinated tables or columns in final validation + +Cancel validation: + +- CWPay cancel task: `7a2e1247-2a73-4a52-9fdd-d41004afc3c7` +- cancel mutation returned `true` +- final status: `STOPPED` +- elapsed to terminal status: 520ms + +## Tests Run + +From `D:\WrenAI\wren-ai-service`: + +```powershell +.\venv\Scripts\python.exe -m pytest tests/pytest/services/test_ask.py tests/pytest/pipelines/generation/test_sql_schema_grounding.py tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py -q +``` + +Result: 114 passed. + +Additional focused Ask service test: + +```powershell +.\venv\Scripts\python.exe -m pytest tests/pytest/services/test_ask.py -q +``` + +Result: 4 passed. + +Post-integration focused test from `D:\WrenAI-push-ask-20260901\wren-ai-service` using the existing `D:\WrenAI` virtualenv: + +```powershell +D:\WrenAI\wren-ai-service\venv\Scripts\python.exe -m pytest tests/pytest/services/test_ask.py tests/pytest/pipelines/generation/test_sql_schema_grounding.py tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py tests/pytest/pipelines/generation/test_sql_answer_prompt.py tests/pytest/pipelines/indexing/test_db_schema.py tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py -q +``` + +Result: 139 passed, 1 skipped. + +Targeted Jest command was attempted through the repo Yarn release: + +```powershell +node .yarn/releases/yarn-4.5.3.cjs test src/apollo/server/services/tests/queryService.test.ts --runInBand +``` + +It did not reach the changed `QueryService` tests because TypeScript compilation failed first in existing `src/apollo/server/repositories/baseRepository.ts` type errors. + +Warnings were pre-existing Pydantic deprecation warnings and existing coroutine cleanup warnings in semantics-preparation tests. + +## Files To Include In Handoff Commit + +Include the Ask/UI source and focused tests: + +- `wren-ai-service/src/pipelines/generation/data_assistance.py` +- `wren-ai-service/src/pipelines/generation/followup_sql_generation.py` +- `wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py` +- `wren-ai-service/src/pipelines/generation/intent_classification.py` +- `wren-ai-service/src/pipelines/generation/sql_answer.py` +- `wren-ai-service/src/pipelines/generation/sql_correction.py` +- `wren-ai-service/src/pipelines/generation/sql_generation.py` +- `wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py` +- `wren-ai-service/src/pipelines/generation/utils/sql.py` +- `wren-ai-service/src/pipelines/indexing/db_schema.py` +- `wren-ai-service/src/pipelines/indexing/utils/helper.py` +- `wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py` +- `wren-ai-service/src/web/v1/routers/ask.py` +- `wren-ai-service/src/web/v1/services/ask.py` +- `wren-ai-service/src/web/v1/services/ask_feedback.py` +- `wren-ai-service/src/web/v1/services/sql_answer.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py` +- `wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py` +- `wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py` +- `wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py` +- `wren-ai-service/tests/pytest/services/test_ask.py` +- `wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts` +- `wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts` +- `wren-ui/next.config.js` +- `wren-ui/src/apollo/server/resolvers/askingResolver.ts` +- `wren-ui/src/apollo/server/resolvers/modelResolver.ts` +- `wren-ui/src/apollo/server/services/askingService.ts` +- `wren-ui/src/apollo/server/services/askingTaskTracker.ts` +- `wren-ui/src/apollo/server/services/queryService.ts` +- `wren-ui/src/apollo/server/services/tests/queryService.test.ts` +- `wren-ui/src/apollo/server/utils/manifest.ts` +- `WRENAI_LOCAL_ASK_HANDOFF.md` + +Do not include: + +- `.codex-tmp` +- local configs +- logs +- venv folders +- extracted datasource dumps +- Qdrant/storage runtime data +- `wren-engine` submodule pointer +- `wren-ui/.yarn/releases/yarn-4.5.3.cjs` mode-only churn +- `wren-ui/package-lock.json` unless intentionally changing package management + +## Remaining Blockers + +- Modeling AI Assistant generate semantics/relationships for CW_GL remains unresolved. Earlier evidence showed semantics omitted the selected model and relationships timed out. Final Ask performance validation skipped assistant generation checks. +- Direct push from the original dirty workspace was rejected as non-fast-forward; the final push was prepared by replaying the work onto the remote tip in a clean integration worktree. + +## Guardrails Preserved + +- No app logic hardcodes datasource, project, organization, table, column, filter-value, or prompt-specific mappings. +- SQL is generated only from verified retrieved schema and then validated. +- Answer formatting is grounded in executed SQL results only. +- Unsupported or weakly covered questions fail quickly with a clear unsupported/clarification result. diff --git a/deployment/kustomizations/base/cm.yaml b/deployment/kustomizations/base/cm.yaml index 0bbe6d5e1b..1817a848e2 100644 --- a/deployment/kustomizations/base/cm.yaml +++ b/deployment/kustomizations/base/cm.yaml @@ -222,13 +222,17 @@ data: column_indexing_batch_size: 50 table_retrieval_size: 10 table_column_retrieval_size: 100 - allow_intent_classification: true + allow_intent_classification: false allow_sql_generation_reasoning: true allow_sql_functions_retrieval: true enable_column_pruning: false max_sql_correction_retries: 3 query_cache_maxsize: 1000 query_cache_ttl: 3600 + semantics_description_max_models_per_batch: 4 + semantics_description_max_columns_per_batch: 50 + semantics_description_max_concurrent_tasks: 4 + semantics_description_generation_timeout_seconds: 120 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: DEBUG diff --git a/docker/config.example.yaml b/docker/config.example.yaml index 2a22e788e9..f7617e6e76 100644 --- a/docker/config.example.yaml +++ b/docker/config.example.yaml @@ -6,42 +6,42 @@ models: model: gpt-4.1-nano-2025-04-14 context_window_size: 1000000 kwargs: - max_tokens: 4096 + max_tokens: 8192 n: 1 seed: 0 temperature: 0 - model: gpt-4.1-mini-2025-04-14 context_window_size: 1000000 kwargs: - max_tokens: 4096 + max_tokens: 8192 n: 1 seed: 0 temperature: 0 - model: gpt-4.1-2025-04-14 context_window_size: 1000000 kwargs: - max_tokens: 4096 + max_tokens: 8192 n: 1 seed: 0 temperature: 0 - model: gpt-5-nano-2025-08-07 context_window_size: 380000 kwargs: - max_completion_tokens: 4096 + max_completion_tokens: 8192 n: 1 seed: 0 reasoning_effort: minimal - model: gpt-5-mini-2025-08-07 context_window_size: 380000 kwargs: - max_completion_tokens: 4096 + max_completion_tokens: 8192 n: 1 seed: 0 reasoning_effort: minimal - model: gpt-5-2025-08-07 context_window_size: 380000 kwargs: - max_completion_tokens: 4096 + max_completion_tokens: 8192 n: 1 seed: 0 reasoning_effort: minimal @@ -176,12 +176,16 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false max_sql_correction_retries: 3 query_cache_maxsize: 1000 query_cache_ttl: 3600 + semantics_description_max_models_per_batch: 4 + semantics_description_max_columns_per_batch: 50 + semantics_description_max_concurrent_tasks: 4 + semantics_description_generation_timeout_seconds: 120 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: DEBUG @@ -190,4 +194,4 @@ settings: sql_pairs_similarity_threshold: 0.7 sql_pairs_retrieval_max_size: 10 instructions_similarity_threshold: 0.7 - instructions_top_k: 10 \ No newline at end of file + instructions_top_k: 10 diff --git a/package-lock.json b/package-lock.json new file mode 100644 index 0000000000..f3e1aa42f9 --- /dev/null +++ b/package-lock.json @@ -0,0 +1,256 @@ +{ + "name": "WrenAI1", + "lockfileVersion": 3, + "requires": true, + "packages": { + "": { + "devDependencies": { + "@apollo/client": "^3.6.9" + } + }, + "node_modules/@apollo/client": { + "version": "3.6.9", + "resolved": "https://registry.npmjs.org/@apollo/client/-/client-3.6.9.tgz", + "integrity": "sha512-Y1yu8qa2YeaCUBVuw08x8NHenFi0sw2I3KCu7Kw9mDSu86HmmtHJkCAifKVrN2iPgDTW/BbP3EpSV8/EQCcxZA==", + "dev": true, + "license": "MIT", + "dependencies": { + "@graphql-typed-document-node/core": "^3.1.1", + "@wry/context": "^0.6.0", + "@wry/equality": "^0.5.0", + "@wry/trie": "^0.3.0", + "graphql-tag": "^2.12.6", + "hoist-non-react-statics": "^3.3.2", + "optimism": "^0.16.1", + "prop-types": "^15.7.2", + "symbol-observable": "^4.0.0", + "ts-invariant": "^0.10.3", + "tslib": "^2.3.0", + "zen-observable-ts": "^1.2.5" + }, + "peerDependencies": { + "graphql": "^14.0.0 || ^15.0.0 || ^16.0.0", + "graphql-ws": "^5.5.5", + "react": "^16.8.0 || ^17.0.0 || ^18.0.0", + "subscriptions-transport-ws": "^0.9.0 || ^0.11.0" + }, + "peerDependenciesMeta": { + "graphql-ws": { + "optional": true + }, + "react": { + "optional": true + }, + "subscriptions-transport-ws": { + "optional": true + } + } + }, + "node_modules/@graphql-typed-document-node/core": { + "version": "3.2.0", + "resolved": "https://registry.npmjs.org/@graphql-typed-document-node/core/-/core-3.2.0.tgz", + "integrity": "sha512-mB9oAsNCm9aM3/SOv4YtBMqZbYj10R7dkq8byBqxGY/ncFwhf2oQzMV+LCRlWoDSEBJ3COiR1yeDvMtsoOsuFQ==", + "dev": true, + "license": "MIT", + "peerDependencies": { + "graphql": "^0.8.0 || ^0.9.0 || ^0.10.0 || ^0.11.0 || ^0.12.0 || ^0.13.0 || ^14.0.0 || ^15.0.0 || ^16.0.0 || ^17.0.0" + } + }, + "node_modules/@wry/context": { + "version": "0.6.1", + "resolved": "https://registry.npmjs.org/@wry/context/-/context-0.6.1.tgz", + "integrity": "sha512-LOmVnY1iTU2D8tv4Xf6MVMZZ+juIJ87Kt/plMijjN20NMAXGmH4u8bS1t0uT74cZ5gwpocYueV58YwyI8y+GKw==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/@wry/equality": { + "version": "0.5.7", + "resolved": "https://registry.npmjs.org/@wry/equality/-/equality-0.5.7.tgz", + "integrity": "sha512-BRFORjsTuQv5gxcXsuDXx6oGRhuVsEGwZy6LOzRRfgu+eSfxbhUQ9L9YtSEIuIjY/o7g3iWFjrc5eSY1GXP2Dw==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/@wry/trie": { + "version": "0.3.2", + "resolved": "https://registry.npmjs.org/@wry/trie/-/trie-0.3.2.tgz", + "integrity": "sha512-yRTyhWSls2OY/pYLfwff867r8ekooZ4UI+/gxot5Wj8EFwSf2rG+n+Mo/6LoLQm1TKA4GRj2+LCpbfS937dClQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/graphql": { + "version": "16.14.1", + "resolved": "https://registry.npmjs.org/graphql/-/graphql-16.14.1.tgz", + "integrity": "sha512-cQOsSMS/IrDz82PVyRDvf/Q1F/bRbBVjJlh+xYOkI1qw2bWRvWGiWc+m2O0d6l4Bt1fyY+8kzJ8JFWGJqNeDBg==", + "dev": true, + "license": "MIT", + "peer": true, + "engines": { + "node": "^12.22.0 || ^14.16.0 || ^16.0.0 || >=17.0.0" + } + }, + "node_modules/graphql-tag": { + "version": "2.12.6", + "resolved": "https://registry.npmjs.org/graphql-tag/-/graphql-tag-2.12.6.tgz", + "integrity": "sha512-FdSNcu2QQcWnM2VNvSCCDCVS5PpPqpzgFT8+GXzqJuoDd0CBncxCY278u4mhRO7tMgo2JjgJA5aZ+nWSQ/Z+xg==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.1.0" + }, + "engines": { + "node": ">=10" + }, + "peerDependencies": { + "graphql": "^0.9.0 || ^0.10.0 || ^0.11.0 || ^0.12.0 || ^0.13.0 || ^14.0.0 || ^15.0.0 || ^16.0.0" + } + }, + "node_modules/hoist-non-react-statics": { + "version": "3.3.2", + "resolved": "https://registry.npmjs.org/hoist-non-react-statics/-/hoist-non-react-statics-3.3.2.tgz", + "integrity": "sha512-/gGivxi8JPKWNm/W0jSmzcMPpfpPLc3dY/6GxhX2hQ9iGj3aDfklV4ET7NjKpSinLpJ5vafa9iiGIEZg10SfBw==", + "dev": true, + "license": "BSD-3-Clause", + "dependencies": { + "react-is": "^16.7.0" + } + }, + "node_modules/js-tokens": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/js-tokens/-/js-tokens-4.0.0.tgz", + "integrity": "sha512-RdJUflcE3cUzKiMqQgsCu06FPu9UdIJO0beYbPhHN4k6apgJtifcoCtT9bcxOpYBtpD2kCM6Sbzg4CausW/PKQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/loose-envify": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/loose-envify/-/loose-envify-1.4.0.tgz", + "integrity": "sha512-lyuxPGr/Wfhrlem2CL/UcnUc1zcqKAImBDzukY7Y5F/yQiNdko6+fRLevlw1HgMySw7f611UIY408EtxRSoK3Q==", + "dev": true, + "license": "MIT", + "dependencies": { + "js-tokens": "^3.0.0 || ^4.0.0" + }, + "bin": { + "loose-envify": "cli.js" + } + }, + "node_modules/object-assign": { + "version": "4.1.1", + "resolved": "https://registry.npmjs.org/object-assign/-/object-assign-4.1.1.tgz", + "integrity": "sha512-rJgTQnkUnH1sFw8yT6VSU3zD3sWmu6sZhIseY8VX+GRu3P6F7Fu+JNDoXfklElbLJSnc3FUQHVe4cU5hj+BcUg==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10.0" + } + }, + "node_modules/optimism": { + "version": "0.16.2", + "resolved": "https://registry.npmjs.org/optimism/-/optimism-0.16.2.tgz", + "integrity": "sha512-zWNbgWj+3vLEjZNIh/okkY2EUfX+vB9TJopzIZwT1xxaMqC5hRLLraePod4c5n4He08xuXNH+zhKFFCu390wiQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@wry/context": "^0.7.0", + "@wry/trie": "^0.3.0" + } + }, + "node_modules/optimism/node_modules/@wry/context": { + "version": "0.7.4", + "resolved": "https://registry.npmjs.org/@wry/context/-/context-0.7.4.tgz", + "integrity": "sha512-jmT7Sb4ZQWI5iyu3lobQxICu2nC/vbUhP0vIdd6tHC9PTfenmRmuIFqktc6GH9cgi+ZHnsLWPvfSvc4DrYmKiQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.3.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/prop-types": { + "version": "15.8.1", + "resolved": "https://registry.npmjs.org/prop-types/-/prop-types-15.8.1.tgz", + "integrity": "sha512-oj87CgZICdulUohogVAR7AjlC0327U4el4L6eAvOqCeudMDVU0NThNaV+b9Df4dXgSP1gXMTnPdhfe/2qDH5cg==", + "dev": true, + "license": "MIT", + "dependencies": { + "loose-envify": "^1.4.0", + "object-assign": "^4.1.1", + "react-is": "^16.13.1" + } + }, + "node_modules/react-is": { + "version": "16.13.1", + "resolved": "https://registry.npmjs.org/react-is/-/react-is-16.13.1.tgz", + "integrity": "sha512-24e6ynE2H+OKt4kqsOvNd8kBpV65zoxbA4BVsEOB3ARVWQki/DHzaUoC5KuON/BiccDaCCTZBuOcfZs70kR8bQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/symbol-observable": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/symbol-observable/-/symbol-observable-4.0.0.tgz", + "integrity": "sha512-b19dMThMV4HVFynSAM1++gBHAbk2Tc/osgLIBZMKsyqh34jb2e8Os7T6ZW/Bt3pJFdBTd2JwAnAAEQV7rSNvcQ==", + "dev": true, + "license": "MIT", + "engines": { + "node": ">=0.10" + } + }, + "node_modules/ts-invariant": { + "version": "0.10.3", + "resolved": "https://registry.npmjs.org/ts-invariant/-/ts-invariant-0.10.3.tgz", + "integrity": "sha512-uivwYcQaxAucv1CzRp2n/QdYPo4ILf9VXgH19zEIjFx2EJufV16P0JtJVpYHy89DItG6Kwj2oIUjrcK5au+4tQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "tslib": "^2.1.0" + }, + "engines": { + "node": ">=8" + } + }, + "node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", + "dev": true, + "license": "0BSD" + }, + "node_modules/zen-observable": { + "version": "0.8.15", + "resolved": "https://registry.npmjs.org/zen-observable/-/zen-observable-0.8.15.tgz", + "integrity": "sha512-PQ2PC7R9rslx84ndNBZB/Dkv8V8fZEpk83RLgXtYd0fwUgEjseMn1Dgajh2x6S8QbZAFa9p2qVCEuYZNgve0dQ==", + "dev": true, + "license": "MIT" + }, + "node_modules/zen-observable-ts": { + "version": "1.2.5", + "resolved": "https://registry.npmjs.org/zen-observable-ts/-/zen-observable-ts-1.2.5.tgz", + "integrity": "sha512-QZWQekv6iB72Naeake9hS1KxHlotfRpe+WGNbNx5/ta+R3DNjVO2bswf63gXlWDcs+EMd7XY8HfVQyP1X6T4Zg==", + "dev": true, + "license": "MIT", + "dependencies": { + "zen-observable": "0.8.15" + } + } + } +} diff --git a/package.json b/package.json new file mode 100644 index 0000000000..10f0ca32dc --- /dev/null +++ b/package.json @@ -0,0 +1,5 @@ +{ + "devDependencies": { + "@apollo/client": "^3.6.9" + } +} diff --git a/patches/0001-Fix-MSSQL-raw-query-execution.patch b/patches/0001-Fix-MSSQL-raw-query-execution.patch new file mode 100644 index 0000000000..50290f2433 --- /dev/null +++ b/patches/0001-Fix-MSSQL-raw-query-execution.patch @@ -0,0 +1,338 @@ +From 5d53b50cd4cf498192a695677d2054864e0391f7 Mon Sep 17 00:00:00 2001 +From: Harshitha +Date: Mon, 22 Jun 2026 13:35:44 +0530 +Subject: [PATCH] Fix MSSQL raw query execution + +--- + ibis-server/app/model/connector.py | 31 +++-- + .../tests/model/test_mssql_connector.py | 94 +++++++++++++++ + wren/src/wren/connector/mssql.py | 111 ++++++++++++++---- + 3 files changed, 204 insertions(+), 32 deletions(-) + create mode 100644 ibis-server/tests/model/test_mssql_connector.py + +diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py +index f9dc009e..21bbb5a3 100644 +--- a/ibis-server/app/model/connector.py ++++ b/ibis-server/app/model/connector.py +@@ -595,10 +595,20 @@ class MSSqlConnector(IbisConnector): + try: + with closing(self.connection.raw_sql(sql)) as cur: + rows = cur.fetchall() +- columns = [column[0] for column in (cur.description or [])] ++ columns = [ ++ self._cursor_column_name(column, index) ++ for index, column in enumerate(cur.description or []) ++ ] ++ ++ df = pd.DataFrame(rows, columns=columns) ++ if limit is not None: ++ df = df.head(limit) ++ return pa.Table.from_pandas(df, preserve_index=False) + except AttributeError as e: +- # Workaround for ibis issue #10331 in the execution path. +- if e.args and e.args[0] == "'NoneType' object has no attribute 'lower'": ++ # Ibis' MSSQL schema probe can mask SQL Server describe errors by ++ # calling .lower() on a NULL system_type_name before checking ++ # error_message. Return the real database message when this leaks. ++ if self._is_none_lower_attribute_error(e): + error_message = self._describe_sql_for_error_message(sql) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, +@@ -608,11 +618,6 @@ class MSSqlConnector(IbisConnector): + ) from e + raise + +- df = pd.DataFrame(rows, columns=columns) +- if limit is not None: +- df = df.head(limit) +- return pa.Table.from_pandas(df, preserve_index=False) +- + def _round_decimal_columns(self, ibis_table: Table, scale: int = 9) -> pa.Table: + def round_decimal(val): + if val is None: +@@ -688,8 +693,8 @@ class MSSqlConnector(IbisConnector): + + return inner.sql(dialect="tsql") + +- except Exception as e: +- return f"Error: {e!s}" ++ except Exception: ++ return sql_query + + def _normalize_tsql_for_execution(self, sql: str) -> str: + replacements = ( +@@ -729,7 +734,7 @@ class MSSqlConnector(IbisConnector): + raise + except AttributeError as e: + # Workaround for ibis issue #10331 +- if e.args[0] == "'NoneType' object has no attribute 'lower'": ++ if self._is_none_lower_attribute_error(e): + error_message = self._describe_sql_for_error_message(normalized_sql) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, +@@ -765,6 +770,10 @@ class MSSqlConnector(IbisConnector): + return "" + return rows[0][0] or "" + ++ @staticmethod ++ def _is_none_lower_attribute_error(error: AttributeError) -> bool: ++ return "NoneType" in str(error) and "lower" in str(error) ++ + + class CannerConnector(IbisConnector): + def __init__(self, connection_info: ConnectionInfo): +diff --git a/ibis-server/tests/model/test_mssql_connector.py b/ibis-server/tests/model/test_mssql_connector.py +new file mode 100644 +index 00000000..7743afe1 +--- /dev/null ++++ b/ibis-server/tests/model/test_mssql_connector.py +@@ -0,0 +1,94 @@ ++from app.model.connector import MSSqlConnector ++from app.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError ++ ++ ++class FakeCursor: ++ def __init__(self, rows, description): ++ self._rows = rows ++ self.description = description ++ self.closed = False ++ ++ def fetchall(self): ++ return self._rows ++ ++ def close(self): ++ self.closed = True ++ ++ ++class FakeConnection: ++ def __init__(self, cursors=None, error=None): ++ self.cursors = list(cursors or []) ++ self.error = error ++ self.queries = [] ++ ++ def raw_sql(self, sql): ++ self.queries.append(sql) ++ if self.error: ++ error = self.error ++ self.error = None ++ raise error ++ return self.cursors.pop(0) ++ ++ ++def _connector(connection): ++ connector = MSSqlConnector.__new__(MSSqlConnector) ++ connector.connection = connection ++ return connector ++ ++ ++def test_query_uses_raw_sql_for_grouped_aggregate_results(): ++ connection = FakeConnection( ++ [ ++ FakeCursor( ++ rows=[("Widgets", 12), ("Gadgets", 8)], ++ description=[("ProdType",), ("TotalQty",)], ++ ) ++ ] ++ ) ++ connector = _connector(connection) ++ ++ result = connector.query( ++ 'SELECT "ProdType", SUM("Qty") AS "TotalQty" ' ++ 'FROM "dbo_tblSalesHistory" ' ++ 'GROUP BY "ProdType" ' ++ 'ORDER BY SUM("Qty") DESC NULLS LAST' ++ ) ++ ++ assert connection.queries == [ ++ 'SELECT "ProdType", SUM("Qty") AS "TotalQty" ' ++ 'FROM "dbo_tblSalesHistory" ' ++ 'GROUP BY "ProdType" ' ++ 'ORDER BY SUM("Qty") DESC' ++ ] ++ assert result.column_names == ["ProdType", "TotalQty"] ++ assert result.to_pylist() == [ ++ {"ProdType": "Widgets", "TotalQty": 12}, ++ {"ProdType": "Gadgets", "TotalQty": 8}, ++ ] ++ ++ ++def test_query_translates_masked_mssql_describe_error(): ++ connection = FakeConnection( ++ cursors=[ ++ FakeCursor( ++ rows=[("Invalid column name 'ProdType'.",)], ++ description=[("error_message",)], ++ ) ++ ], ++ error=AttributeError("'NoneType' object has no attribute 'lower'"), ++ ) ++ connector = _connector(connection) ++ ++ try: ++ connector.query('SELECT "ProdType" FROM "dbo_tblSalesHistory"') ++ except WrenError as exc: ++ error = exc ++ else: ++ raise AssertionError("Expected WrenError") ++ ++ assert error.error_code == ErrorCode.INVALID_SQL ++ assert error.phase == ErrorPhase.SQL_EXECUTION ++ assert error.metadata == { ++ DIALECT_SQL: 'SELECT "ProdType" FROM "dbo_tblSalesHistory"' ++ } ++ assert "Invalid column name 'ProdType'." in error.message +diff --git a/wren/src/wren/connector/mssql.py b/wren/src/wren/connector/mssql.py +index f7454676..1fd765d7 100644 +--- a/wren/src/wren/connector/mssql.py ++++ b/wren/src/wren/connector/mssql.py +@@ -1,8 +1,9 @@ + from contextlib import closing + from decimal import Decimal as PyDecimal ++import re + ++import pandas as pd + import pyarrow as pa +-import sqlglot.expressions as sge + from ibis.expr.datatypes import Decimal + from ibis.expr.types import Table + from sqlglot import exp, parse_one +@@ -17,12 +18,29 @@ class MSSqlConnector(IbisConnector): + super().__init__(DataSource.mssql, connection_info) + + def query(self, sql: str, limit: int | None = None) -> pa.Table: +- sql = self._flatten_pagination_limit(sql) +- ibis_table = self.connection.sql(sql) +- if limit is not None: +- ibis_table = ibis_table.limit(limit) +- ibis_table = self._handle_pyarrow_unsupported_type(ibis_table) +- return self._round_decimal_columns(ibis_table) ++ sql = self._flatten_pagination_limit(self._normalize_tsql_for_execution(sql)) ++ try: ++ with closing(self.connection.raw_sql(sql)) as cur: ++ rows = cur.fetchall() ++ columns = [ ++ self._cursor_column_name(column, index) ++ for index, column in enumerate(cur.description or []) ++ ] ++ ++ df = pd.DataFrame(rows, columns=columns) ++ if limit is not None: ++ df = df.head(limit) ++ return pa.Table.from_pandas(df, preserve_index=False) ++ except AttributeError as e: ++ if self._is_none_lower_attribute_error(e): ++ error_message = self._describe_sql_for_error_message(sql) ++ raise WrenError( ++ error_code=ErrorCode.INVALID_SQL, ++ message=f"The sql query failed. {error_message or str(e)}.", ++ phase=ErrorPhase.SQL_EXECUTION, ++ metadata={DIALECT_SQL: sql}, ++ ) from e ++ raise + + def _round_decimal_columns(self, ibis_table: Table, scale: int = 9) -> pa.Table: + def round_decimal(val): +@@ -78,35 +96,86 @@ class MSSqlConnector(IbisConnector): + except Exception: + return sql_query + ++ def _normalize_tsql_for_execution(self, sql: str) -> str: ++ replacements = ( ++ (r"DATE_PART\s*\(", "DATEPART("), ++ (r"DATEPART\(\s*YEAR\s*,", "DATEPART('YEAR',"), ++ (r"DATEPART\(\s*MONTH\s*,", "DATEPART('MONTH',"), ++ (r"DATEPART\(\s*DAY\s*,", "DATEPART('DAY',"), ++ (r"DATEDIFF\(\s*'SECOND'\s*,", "DATEDIFF(SECOND,"), ++ (r"DATEDIFF\(\s*'MINUTE'\s*,", "DATEDIFF(MINUTE,"), ++ (r"DATEDIFF\(\s*'HOUR'\s*,", "DATEDIFF(HOUR,"), ++ (r"DATEDIFF\(\s*'DAY'\s*,", "DATEDIFF(DAY,"), ++ (r"\s+NULLS\s+LAST\b", ""), ++ (r"\s+NULLS\s+FIRST\b", ""), ++ ) ++ ++ normalized = sql ++ for pattern, replacement in replacements: ++ normalized = re.sub(pattern, replacement, normalized, flags=re.IGNORECASE) ++ ++ return normalized ++ ++ def _quote_sql_literal(self, sql: str) -> str: ++ return "N'" + sql.replace("'", "''") + "'" ++ + def dry_run(self, sql: str) -> None: ++ normalized_sql = self._normalize_tsql_for_execution(sql) + try: +- super().dry_run(sql) ++ error_message = self._describe_sql_for_error_message(normalized_sql) ++ if error_message: ++ raise WrenError( ++ error_code=ErrorCode.INVALID_SQL, ++ message=f"The sql dry run failed. {error_message}.", ++ phase=ErrorPhase.SQL_DRY_RUN, ++ metadata={DIALECT_SQL: normalized_sql}, ++ ) ++ except WrenError: ++ raise + except AttributeError as e: +- if "NoneType" in str(e) and "lower" in str(e): +- error_message = self._describe_sql_for_error_message(sql) ++ if self._is_none_lower_attribute_error(e): ++ error_message = self._describe_sql_for_error_message(normalized_sql) + raise WrenError( + error_code=ErrorCode.INVALID_SQL, + message=f"The sql dry run failed. {error_message}.", + phase=ErrorPhase.SQL_DRY_RUN, +- metadata={DIALECT_SQL: sql}, ++ metadata={DIALECT_SQL: normalized_sql}, + ) from e + raise WrenError( + error_code=ErrorCode.IBIS_PROJECT_ERROR, + message=str(e), + phase=ErrorPhase.SQL_DRY_RUN, + ) from e ++ except Exception as e: ++ raise WrenError( ++ error_code=ErrorCode.INVALID_SQL, ++ message=f"The sql dry run failed. {e!s}.", ++ phase=ErrorPhase.SQL_DRY_RUN, ++ metadata={DIALECT_SQL: normalized_sql}, ++ ) from e + + def _describe_sql_for_error_message(self, sql: str) -> str: +- try: +- tsql = sge.convert(sql).sql("mssql") +- describe_sql = f"SELECT error_message FROM sys.dm_exec_describe_first_result_set({tsql}, NULL, 0)" +- with closing(self.connection.raw_sql(describe_sql)) as cur: +- rows = cur.fetchall() +- if not rows: +- return "Unknown reason" +- return rows[0][0] +- except Exception: +- return "Unknown reason" ++ describe_sql = ( ++ "SELECT error_message " ++ "FROM sys.dm_exec_describe_first_result_set(" ++ f"{self._quote_sql_literal(sql)}, NULL, 0)" ++ ) ++ with closing(self.connection.raw_sql(describe_sql)) as cur: ++ rows = cur.fetchall() ++ if rows is None or len(rows) == 0: ++ return "" ++ return rows[0][0] or "" ++ ++ @staticmethod ++ def _cursor_column_name(column, index: int) -> str: ++ name = getattr(column, "name", None) ++ if name is None and isinstance(column, (tuple, list)) and column: ++ name = column[0] ++ return str(name or f"column_{index + 1}") ++ ++ @staticmethod ++ def _is_none_lower_attribute_error(error: AttributeError) -> bool: ++ return "NoneType" in str(error) and "lower" in str(error) + + + def create_connector(connection_info) -> MSSqlConnector: +-- +2.53.0.windows.2 + diff --git a/patches/0002-Fix-CTE-rewrite-for-set-operations.patch b/patches/0002-Fix-CTE-rewrite-for-set-operations.patch new file mode 100644 index 0000000000..80a84bd81f --- /dev/null +++ b/patches/0002-Fix-CTE-rewrite-for-set-operations.patch @@ -0,0 +1,231 @@ +From 667da5a177c624a515ef9fc6d1f8f3d5a68793a4 Mon Sep 17 00:00:00 2001 +From: Harshitha +Date: Wed, 29 Jul 2026 01:34:39 +0530 +Subject: [PATCH] Fix CTE rewrite for set operations + +--- + wren/src/wren/mdl/cte_rewriter.py | 97 ++++++++++++++++++++-------- + wren/tests/unit/test_cte_rewriter.py | 29 +++++++-- + 2 files changed, 95 insertions(+), 31 deletions(-) + +diff --git a/wren/src/wren/mdl/cte_rewriter.py b/wren/src/wren/mdl/cte_rewriter.py +index 62be74ce..8f9f93e9 100644 +--- a/wren/src/wren/mdl/cte_rewriter.py ++++ b/wren/src/wren/mdl/cte_rewriter.py +@@ -16,6 +16,7 @@ from sqlglot import exp, parse_one + from sqlglot.optimizer.normalize_identifiers import normalize_identifiers + from sqlglot.optimizer.qualify_columns import qualify_columns + from sqlglot.optimizer.qualify_tables import qualify_tables ++from sqlglot.optimizer.scope import traverse_scope + from sqlglot.schema import MappingSchema + + # Ensure the Wren dialect is registered with sqlglot on import. +@@ -172,6 +173,13 @@ class CTERewriter: + if model_name: + used[model_name][col.name] = None + ++ for model_name, columns in self._collect_scope_model_columns( ++ copy, user_cte_names ++ ).items(): ++ model_columns = used.setdefault(model_name, {}) ++ for column in columns: ++ model_columns[column] = None ++ + return {m: None if m in star_models else list(cols) for m, cols in used.items()} + + # ------------------------------------------------------------------ +@@ -214,7 +222,8 @@ class CTERewriter: + if not model_ctes: + return + +- existing_with = ast.args.get("with_") ++ with_key = "with_" if "with_" in ast.arg_types else "with" ++ existing_with = ast.args.get(with_key) + + if existing_with: + # Prepend model CTEs before user CTEs +@@ -223,10 +232,10 @@ class CTERewriter: + existing_with.set("expressions", all_ctes) + else: + with_clause = exp.With(expressions=model_ctes) +- ast.set("with_", with_clause) ++ ast.set(with_key, with_clause) + + # Preserve RECURSIVE if the original WITH had it +- final_with = ast.args.get("with_") ++ final_with = ast.args.get(with_key) + if existing_with and existing_with.args.get("recursive"): + final_with.set("recursive", True) + +@@ -243,33 +252,69 @@ class CTERewriter: + the referenced model. + """ + star_models: set[str] = set() +- select = ast.find(exp.Select) +- if not select: +- return star_models ++ for scope in traverse_scope(ast): ++ select = scope.expression ++ if not isinstance(select, exp.Select): ++ continue ++ ++ alias_to_model = self._scope_alias_to_model(scope, user_cte_names) ++ for sel_expr in select.expressions: ++ if isinstance(sel_expr, exp.Star): ++ star_models.update(alias_to_model.values()) ++ elif isinstance(sel_expr, exp.Column) and isinstance( ++ sel_expr.this, exp.Star ++ ): ++ table_ref = sel_expr.table ++ if table_ref and table_ref in alias_to_model: ++ star_models.add(alias_to_model[table_ref]) ++ ++ return star_models ++ ++ def _collect_scope_model_columns( ++ self, ast: exp.Expression, user_cte_names: set[str] ++ ) -> dict[str, dict[str, None]]: ++ used: dict[str, dict[str, None]] = {} ++ ++ for scope in traverse_scope(ast): ++ alias_to_model = self._scope_alias_to_model(scope, user_cte_names) ++ for model_name in alias_to_model.values(): ++ used.setdefault(model_name, {}) ++ ++ if not alias_to_model: ++ continue ++ ++ for col in scope.columns: ++ table_ref = col.table ++ model_name = alias_to_model.get(table_ref) if table_ref else None ++ unique_models = list(dict.fromkeys(alias_to_model.values())) ++ if model_name is None and len(unique_models) == 1: ++ candidate_model = unique_models[0] ++ model_columns = self._col_orig_name.get(candidate_model, {}) ++ if col.name.lower() in model_columns: ++ model_name = candidate_model ++ ++ if model_name: ++ used.setdefault(model_name, {})[col.name] = None + +- # Build alias → model mapping from tables in FROM/JOIN ++ return used ++ ++ def _scope_alias_to_model( ++ self, scope, user_cte_names: set[str] ++ ) -> dict[str, str]: + alias_to_model: dict[str, str] = {} +- for table in ast.find_all(exp.Table): +- name = table.name +- if name not in self.model_dict or name in user_cte_names: ++ ++ for alias, (_, source) in scope.selected_sources.items(): ++ if not isinstance(source, exp.Table): + continue +- alias = table.alias or name +- alias_to_model[alias] = name +- alias_to_model[name] = name +- +- for sel_expr in select.expressions: +- if isinstance(sel_expr, exp.Star): +- # Bare * → all models +- star_models.update(alias_to_model.values()) +- elif isinstance(sel_expr, exp.Column) and isinstance( +- sel_expr.this, exp.Star +- ): +- # table.* → specific model +- table_ref = sel_expr.table +- if table_ref and table_ref in alias_to_model: +- star_models.add(alias_to_model[table_ref]) + +- return star_models ++ model_name = source.name ++ if model_name not in self.model_dict or model_name in user_cte_names: ++ continue ++ ++ alias_to_model[alias] = model_name ++ alias_to_model[model_name] = model_name ++ ++ return alias_to_model + + @staticmethod + def _collect_user_cte_names(ast: exp.Expression) -> set[str]: +diff --git a/wren/tests/unit/test_cte_rewriter.py b/wren/tests/unit/test_cte_rewriter.py +index 289bcd06..bcc17c46 100644 +--- a/wren/tests/unit/test_cte_rewriter.py ++++ b/wren/tests/unit/test_cte_rewriter.py +@@ -103,7 +103,7 @@ def _make_rewriter( + def _has_cte(sql: str, cte_name: str, dialect: str = "duckdb") -> bool: + """Return True if *sql* contains a CTE named *cte_name*.""" + ast = sqlglot.parse_one(sql, dialect=dialect) +- with_clause = ast.args.get("with_") ++ with_clause = ast.args.get("with_") or ast.args.get("with") + if not with_clause: + return False + for cte in with_clause.expressions: +@@ -115,7 +115,7 @@ def _has_cte(sql: str, cte_name: str, dialect: str = "duckdb") -> bool: + + def _count_ctes(sql: str, dialect: str = "duckdb") -> int: + ast = sqlglot.parse_one(sql, dialect=dialect) +- with_clause = ast.args.get("with_") ++ with_clause = ast.args.get("with_") or ast.args.get("with") + if not with_clause: + return 0 + return len(with_clause.expressions) +@@ -124,7 +124,7 @@ def _count_ctes(sql: str, dialect: str = "duckdb") -> int: + def _cte_body_sql(sql: str, cte_name: str, dialect: str = "duckdb") -> str | None: + """Return the SQL body of a named CTE, or None if not found.""" + ast = sqlglot.parse_one(sql, dialect=dialect) +- with_clause = ast.args.get("with_") ++ with_clause = ast.args.get("with_") or ast.args.get("with") + if not with_clause: + return None + for cte in with_clause.expressions: +@@ -237,8 +237,9 @@ class TestCTEEdgeCases: + assert _has_cte(result, "summary") + # Model CTE should come before user CTE + ast = sqlglot.parse_one(result, dialect="duckdb") ++ with_clause = ast.args.get("with_") or ast.args.get("with") + cte_names = [ +- cte.args["alias"].this.name for cte in ast.args["with_"].expressions ++ cte.args["alias"].this.name for cte in with_clause.expressions + ] + assert cte_names.index("orders") < cte_names.index("summary") + +@@ -299,7 +300,8 @@ class TestCTEEdgeCases: + assert _has_cte(result, "hierarchy") + # RECURSIVE keyword must be preserved + ast = sqlglot.parse_one(result, dialect="duckdb") +- assert ast.args["with_"].args.get("recursive") ++ with_clause = ast.args.get("with_") or ast.args.get("with") ++ assert with_clause.args.get("recursive") + + def test_no_model_references_fallback(self): + """Query referencing no models falls back to direct transform_sql.""" +@@ -309,6 +311,23 @@ class TestCTEEdgeCases: + with pytest.raises(Exception): + rw.rewrite("SELECT * FROM unknown_table") + ++ def test_union_collects_columns_from_each_branch(self): ++ """Set-operation branches should each contribute model columns.""" ++ rw = _make_rewriter(_MULTI_MODEL_MANIFEST) ++ result = rw.rewrite( ++ 'SELECT o_orderkey FROM "orders" ' ++ 'UNION SELECT c_custkey FROM "customer"' ++ ) ++ ++ orders_body = _cte_body_sql(result, "orders") ++ customer_body = _cte_body_sql(result, "customer") ++ ++ assert orders_body is not None ++ assert customer_body is not None ++ assert "o_orderkey" in orders_body ++ assert "c_custkey" in customer_body ++ assert "SELECT 1" not in customer_body ++ + + # --------------------------------------------------------------------------- + # Tests: correlated subquery (the key fix) +-- +2.53.0.windows.2 + diff --git a/patches/0003-Validate-Wren-SQL-against-active-manifest.patch b/patches/0003-Validate-Wren-SQL-against-active-manifest.patch new file mode 100644 index 0000000000..192ca1ae5f --- /dev/null +++ b/patches/0003-Validate-Wren-SQL-against-active-manifest.patch @@ -0,0 +1,459 @@ +From 0000000000000000000000000000000000000000 Mon Sep 17 00:00:00 2001 +From: Codex +Date: Sat, 8 Aug 2026 11:00:00 +0530 +Subject: [PATCH] Validate Wren SQL against active manifest before execution + +--- + ibis-server/app/mdl/rewriter.py | 2 + + ibis-server/app/mdl/sql_validator.py | 210 ++++++++++++++++++++ + ibis-server/app/routers/v2/connector.py | 2 + + ibis-server/app/routers/v3/connector.py | 2 + + ibis-server/tests/mdl/test_sql_validator.py | 135 +++++++++++++ + 5 files changed, 351 insertions(+) + create mode 100644 ibis-server/app/mdl/sql_validator.py + create mode 100644 ibis-server/tests/mdl/test_sql_validator.py + +diff --git a/ibis-server/app/mdl/rewriter.py b/ibis-server/app/mdl/rewriter.py +index 7398a72e..bb402d11 100644 +--- a/ibis-server/app/mdl/rewriter.py ++++ b/ibis-server/app/mdl/rewriter.py +@@ -17,6 +17,7 @@ from app.mdl.core import ( + to_json_base64, + ) + from app.mdl.java_engine import JavaEngineConnector ++from app.mdl.sql_validator import validate_wren_sql_against_manifest + from app.model.data_source import DataSource + from app.model.error import PLANNED_SQL, ErrorCode, ErrorPhase, WrenError + from app.util import base64_to_dict +@@ -180,6 +181,7 @@ class Rewriter: + + @tracer.start_as_current_span("rewrite", kind=trace.SpanKind.INTERNAL) + async def rewrite(self, sql: str) -> str: ++ validate_wren_sql_against_manifest(self.manifest_str, sql) + manifest_str = ( + self._extract_manifest(self.manifest_str, sql) or self.manifest_str + ) +diff --git a/ibis-server/app/mdl/sql_validator.py b/ibis-server/app/mdl/sql_validator.py +new file mode 100644 +index 00000000..1e60e3c2 +--- /dev/null ++++ b/ibis-server/app/mdl/sql_validator.py +@@ -0,0 +1,210 @@ ++from __future__ import annotations ++ ++from dataclasses import dataclass ++ ++import sqlglot ++from sqlglot import exp ++from sqlglot.optimizer.scope import build_scope ++ ++from app.custom_sqlglot.dialects.wren import Wren ++from app.model.error import ErrorCode, ErrorPhase, WrenError ++from app.util import base64_to_dict ++ ++ ++_INVALID_IDENTIFIER_MESSAGE = ( ++ "SQL references identifiers that are not present in the active manifest. " ++ "Regenerate the SQL using only deployed model and column metadata." ++) ++ ++ ++@dataclass(frozen=True) ++class ManifestIdentifierIndex: ++ relation_keys: frozenset[str] ++ columns_by_relation: dict[str, frozenset[str]] ++ ++ @property ++ def all_columns(self) -> frozenset[str]: ++ return frozenset( ++ column ++ for relation_columns in self.columns_by_relation.values() ++ for column in relation_columns ++ ) ++ ++ ++def validate_wren_sql_against_manifest(manifest_str: str, sql: str) -> None: ++ manifest = base64_to_dict(manifest_str) ++ index = _build_manifest_identifier_index(manifest) ++ ++ try: ++ ast = sqlglot.parse_one(sql, dialect=Wren) ++ except sqlglot.errors.SqlglotError as e: ++ raise WrenError( ++ ErrorCode.INVALID_SQL, ++ "SQL could not be parsed for manifest validation.", ++ phase=ErrorPhase.SQL_PARSING, ++ ) from e ++ ++ _validate_relation_identifiers(ast, index) ++ _validate_column_identifiers(ast, index) ++ ++ ++def _build_manifest_identifier_index(manifest: dict) -> ManifestIdentifierIndex: ++ catalog = _normalize_identifier(manifest.get("catalog")) ++ schema = _normalize_identifier(manifest.get("schema")) ++ relation_keys: set[str] = set() ++ columns_by_relation: dict[str, frozenset[str]] = {} ++ ++ for item in _iter_queryable_manifest_items(manifest): ++ name = _normalize_identifier(item.get("name")) ++ if not name: ++ continue ++ ++ relation_keys.update(_relation_key_variants(name, catalog, schema)) ++ columns_by_relation[name] = frozenset(_iter_item_column_names(item)) ++ ++ return ManifestIdentifierIndex( ++ relation_keys=frozenset(relation_keys), ++ columns_by_relation=columns_by_relation, ++ ) ++ ++ ++def _iter_queryable_manifest_items(manifest: dict): ++ yield from manifest.get("models", []) ++ yield from manifest.get("views", []) ++ yield from manifest.get("metrics", []) ++ ++ ++def _iter_item_column_names(item: dict): ++ for column in item.get("columns", []): ++ name = _normalize_identifier(column.get("name")) ++ if name: ++ yield name ++ ++ for key in ("measure", "measures", "dimension", "dimensions"): ++ for column in item.get(key, []): ++ name = _normalize_identifier(column.get("name")) ++ if name: ++ yield name ++ ++ ++def _relation_key_variants(name: str, catalog: str | None, schema: str | None): ++ yield name ++ if schema: ++ yield f"{schema}.{name}" ++ if catalog and schema: ++ yield f"{catalog}.{schema}.{name}" ++ ++ ++def _validate_relation_identifiers( ++ ast: exp.Expression, index: ManifestIdentifierIndex ++) -> None: ++ cte_names = _cte_names(ast) ++ for table in ast.find_all(exp.Table): ++ name = _normalize_identifier(table.name) ++ if not name or name in cte_names: ++ continue ++ ++ if _table_key(table) not in index.relation_keys: ++ raise WrenError( ++ ErrorCode.INVALID_SQL, ++ _INVALID_IDENTIFIER_MESSAGE, ++ phase=ErrorPhase.VALIDATION, ++ ) ++ ++ ++def _validate_column_identifiers( ++ ast: exp.Expression, index: ManifestIdentifierIndex ++) -> None: ++ root = build_scope(ast) ++ if root is None: ++ return ++ ++ select_aliases = _select_aliases(ast) ++ cte_names = _cte_names(ast) ++ all_columns = index.all_columns ++ ++ for scope in root.traverse(): ++ source_relations, has_derived_source = _scope_source_relations(scope, cte_names) ++ for column in scope.columns: ++ column_name = _normalize_identifier(column.name) ++ if ( ++ not column_name ++ or column_name == "*" ++ or column_name in select_aliases ++ ): ++ continue ++ ++ qualifier = _normalize_identifier(column.table) ++ if qualifier: ++ relation_name = source_relations.get(qualifier) or qualifier ++ if relation_name in cte_names or qualifier in cte_names: ++ continue ++ ++ relation_columns = index.columns_by_relation.get(relation_name) ++ if relation_columns and column_name not in relation_columns: ++ _raise_invalid_identifier() ++ continue ++ ++ if not source_relations or has_derived_source: ++ continue ++ ++ if len(source_relations) == 1: ++ relation_columns = index.columns_by_relation.get( ++ next(iter(source_relations.values())) ++ ) ++ if relation_columns and column_name not in relation_columns: ++ _raise_invalid_identifier() ++ continue ++ ++ if all_columns and column_name not in all_columns: ++ _raise_invalid_identifier() ++ ++ ++def _scope_source_relations(scope, cte_names: set[str]) -> tuple[dict[str, str], bool]: ++ source_relations: dict[str, str] = {} ++ has_derived_source = False ++ ++ for alias, (_node, source) in scope.selected_sources.items(): ++ alias_name = _normalize_identifier(alias) ++ if not alias_name: ++ continue ++ ++ if isinstance(source, exp.Table): ++ source_name = _normalize_identifier(source.name) ++ if not source_name or source_name in cte_names: ++ has_derived_source = True ++ continue ++ ++ source_relations[alias_name] = source_name ++ source_relations[source_name] = source_name ++ else: ++ has_derived_source = True ++ ++ return source_relations, has_derived_source ++ ++ ++def _table_key(table: exp.Table) -> str: ++ parts = [ ++ _normalize_identifier(part) ++ for part in (table.catalog, table.db, table.name) ++ if _normalize_identifier(part) ++ ] ++ return ".".join(parts) ++ ++ ++def _cte_names(ast: exp.Expression) -> set[str]: ++ return { ++ name ++ for cte in ast.find_all(exp.CTE) ++ if (name := _normalize_identifier(cte.alias_or_name)) ++ } ++ ++ ++def _select_aliases(ast: exp.Expression) -> set[str]: ++ return { ++ alias ++ for expression in ast.find_all(exp.Alias) ++ if (alias := _normalize_identifier(expression.alias)) ++ } ++ ++ ++def _normalize_identifier(value) -> str | None: ++ if value is None: ++ return None ++ text = str(value).strip() ++ return text.lower() if text else None ++ ++ ++def _raise_invalid_identifier() -> None: ++ raise WrenError( ++ ErrorCode.INVALID_SQL, ++ _INVALID_IDENTIFIER_MESSAGE, ++ phase=ErrorPhase.VALIDATION, ++ ) +diff --git a/ibis-server/app/routers/v2/connector.py b/ibis-server/app/routers/v2/connector.py +index c68f08de..1b1fb95e 100644 +--- a/ibis-server/app/routers/v2/connector.py ++++ b/ibis-server/app/routers/v2/connector.py +@@ -16,6 +16,7 @@ from app.dependencies import ( + ) + from app.mdl.java_engine import JavaEngineConnector + from app.mdl.rewriter import Rewriter ++from app.mdl.sql_validator import validate_wren_sql_against_manifest + from app.mdl.substitute import ModelSubstitute + from app.model import ( + BigQueryProjectConnectionInfo, +@@ -103,6 +104,7 @@ async def query( + name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers) + ) as span: + set_attribute(headers, span) ++ validate_wren_sql_against_manifest(dto.manifest_str, dto.sql) + try: + sql = pushdown_limit(dto.sql, limit) + except Exception as e: +diff --git a/ibis-server/app/routers/v3/connector.py b/ibis-server/app/routers/v3/connector.py +index f8d857bf..e8686f92 100644 +--- a/ibis-server/app/routers/v3/connector.py ++++ b/ibis-server/app/routers/v3/connector.py +@@ -22,6 +22,7 @@ from app.mdl import knowledge + from app.mdl.core import get_session_context + from app.mdl.java_engine import JavaEngineConnector + from app.mdl.rewriter import Rewriter ++from app.mdl.sql_validator import validate_wren_sql_against_manifest + from app.mdl.substitute import ModelSubstitute + from app.model import ( + DryPlanDTO, +@@ -100,6 +101,7 @@ async def query( + name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers) + ) as span: + set_attribute(headers, span) ++ validate_wren_sql_against_manifest(dto.manifest_str, dto.sql) + connection_info = data_source.get_connection_info( + resolve_connection_info(dto), dict(headers) + ) +diff --git a/ibis-server/tests/mdl/test_sql_validator.py b/ibis-server/tests/mdl/test_sql_validator.py +new file mode 100644 +index 00000000..8229706a +--- /dev/null ++++ b/ibis-server/tests/mdl/test_sql_validator.py +@@ -0,0 +1,135 @@ ++import base64 ++ ++import orjson ++import pytest ++ ++from app.mdl.sql_validator import validate_wren_sql_against_manifest ++from app.model.error import ErrorCode, ErrorPhase, WrenError ++ ++ ++def _manifest_str(manifest: dict) -> str: ++ return base64.b64encode(orjson.dumps(manifest)).decode("utf-8") ++ ++ ++@pytest.fixture ++def manifest_str(): ++ return _manifest_str( ++ { ++ "catalog": "semantic_catalog", ++ "schema": "semantic_schema", ++ "models": [ ++ { ++ "name": "model_alpha", ++ "columns": [ ++ {"name": "entity_id", "type": "integer"}, ++ {"name": "metric_value", "type": "integer"}, ++ {"name": "event_date", "type": "date"}, ++ ], ++ }, ++ { ++ "name": "model_beta", ++ "columns": [ ++ {"name": "entity_id", "type": "integer"}, ++ {"name": "category_name", "type": "varchar"}, ++ ], ++ }, ++ ], ++ "relationships": [ ++ { ++ "name": "relationship_alpha_beta", ++ "models": ["model_alpha", "model_beta"], ++ "condition": "model_alpha.entity_id = model_beta.entity_id", ++ "joinType": "MANY_TO_ONE", ++ } ++ ], ++ "views": [ ++ { ++ "name": "view_alpha", ++ "columns": [ ++ {"name": "category_name", "type": "varchar"}, ++ {"name": "metric_value", "type": "integer"}, ++ ], ++ } ++ ], ++ "metrics": [ ++ { ++ "name": "metric_alpha", ++ "measure": [ ++ {"name": "total_value", "type": "integer"}, ++ ], ++ "dimension": [ ++ {"name": "category_name", "type": "varchar"}, ++ ], ++ } ++ ], ++ } ++ ) ++ ++ ++def test_validate_wren_sql_accepts_manifest_models_columns_and_aliases(manifest_str): ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ """ ++ SELECT ++ a.entity_id, ++ b.category_name, ++ SUM(a.metric_value) AS total_value ++ FROM semantic_catalog.semantic_schema.model_alpha AS a ++ JOIN model_beta AS b ON a.entity_id = b.entity_id ++ WHERE a.event_date >= DATE '2026-01-01' ++ GROUP BY a.entity_id, b.category_name ++ ORDER BY total_value DESC ++ """, ++ ) ++ ++ ++def test_validate_wren_sql_accepts_views_and_metrics(manifest_str): ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ "SELECT category_name, metric_value FROM view_alpha", ++ ) ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ "SELECT category_name, total_value FROM metric_alpha", ++ ) ++ ++ ++def test_validate_wren_sql_rejects_unknown_relation_before_execution(manifest_str): ++ with pytest.raises(WrenError) as error: ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ "SELECT entity_id FROM missing_model", ++ ) ++ ++ assert error.value.error_code == ErrorCode.INVALID_SQL ++ assert error.value.phase == ErrorPhase.VALIDATION ++ assert "missing_model" not in error.value.message ++ ++ ++def test_validate_wren_sql_rejects_unknown_column_before_execution(manifest_str): ++ with pytest.raises(WrenError) as error: ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ "SELECT missing_column FROM model_alpha", ++ ) ++ ++ assert error.value.error_code == ErrorCode.INVALID_SQL ++ assert error.value.phase == ErrorPhase.VALIDATION ++ assert "missing_column" not in error.value.message ++ ++ ++def test_validate_wren_sql_does_not_reject_derived_cte_columns(manifest_str): ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ """ ++ WITH scoped AS ( ++ SELECT category_name, SUM(metric_value) AS scoped_total ++ FROM model_alpha ++ JOIN model_beta ON model_alpha.entity_id = model_beta.entity_id ++ GROUP BY category_name ++ ) ++ SELECT category_name, scoped_total FROM scoped ++ """, ++ ) ++ ++ ++def test_validate_wren_sql_rejects_invalid_column_inside_cte(manifest_str): ++ with pytest.raises(WrenError): ++ validate_wren_sql_against_manifest( ++ manifest_str, ++ """ ++ WITH scoped AS ( ++ SELECT invalid_metric AS scoped_total ++ FROM model_alpha ++ ) ++ SELECT scoped_total FROM scoped ++ """, ++ ) +-- +2.53.0.windows.2 diff --git a/patches/wren-engine-ibis-null-lower-fix.patch b/patches/wren-engine-ibis-null-lower-fix.patch new file mode 100644 index 0000000000..1d0be6b71b --- /dev/null +++ b/patches/wren-engine-ibis-null-lower-fix.patch @@ -0,0 +1,57 @@ +From ec835849487e5300fe54b2a5cfd7db4f31b85e1d Mon Sep 17 00:00:00 2001 +From: Harshitha +Date: Sun, 21 Jun 2026 20:48:16 +0530 +Subject: [PATCH] Fallback raw SQL for Ibis null type query errors + +--- + ibis-server/app/model/connector.py | 33 ++++++++++++++++++++++++++---- + 1 file changed, 29 insertions(+), 4 deletions(-) + +diff --git a/ibis-server/app/model/connector.py b/ibis-server/app/model/connector.py +index 397bb8c2..b4437159 100644 +--- a/ibis-server/app/model/connector.py ++++ b/ibis-server/app/model/connector.py +@@ -344,11 +344,36 @@ class IbisConnector(ConnectorABC): + + @tracer.start_as_current_span("connector_query", kind=trace.SpanKind.CLIENT) + def query(self, sql: str, limit: int | None = None) -> pa.Table: +- ibis_table = self.connection.sql(sql) ++ try: ++ ibis_table = self.connection.sql(sql) ++ if limit is not None: ++ ibis_table = ibis_table.limit(limit) ++ ibis_table = self._handle_pyarrow_unsupported_type(ibis_table) ++ return ibis_table.to_pyarrow() ++ except AttributeError as e: ++ if e.args and e.args[0] == "'NoneType' object has no attribute 'lower'": ++ return self._query_raw_sql(sql, limit) ++ raise ++ ++ def _query_raw_sql(self, sql: str, limit: int | None = None) -> pa.Table: ++ with closing(self.connection.raw_sql(sql)) as cur: ++ rows = cur.fetchall() ++ columns = [ ++ self._cursor_column_name(column, index) ++ for index, column in enumerate(cur.description or []) ++ ] ++ ++ df = pd.DataFrame(rows, columns=columns) + if limit is not None: +- ibis_table = ibis_table.limit(limit) +- ibis_table = self._handle_pyarrow_unsupported_type(ibis_table) +- return ibis_table.to_pyarrow() ++ df = df.head(limit) ++ return pa.Table.from_pandas(df, preserve_index=False) ++ ++ @staticmethod ++ def _cursor_column_name(column: Any, index: int) -> str: ++ name = getattr(column, "name", None) ++ if name is None and isinstance(column, (tuple, list)) and column: ++ name = column[0] ++ return str(name or f"column_{index + 1}") + + def _handle_pyarrow_unsupported_type(self, ibis_table: Table, **kwargs) -> Table: + result_table = ibis_table +-- +2.53.0.windows.2 + diff --git a/questions_1000.md b/questions_1000.md new file mode 100644 index 0000000000..20a0d03423 --- /dev/null +++ b/questions_1000.md @@ -0,0 +1,1004 @@ +# 1000 Test Questions + +Generated from the invoice, purchase order, supplier, journal entry, document, task, workflow, prepayment, and metadata-grounding patterns in the provided examples. + +1. Show the total number of invoices. +2. Show the total number of purchase orders. +3. Show the total number of purchase order line items. +4. Show the total number of invoice line items. +5. Show the total number of journal entries. +6. Show the total number of documents. +7. Show the total number of tasks. +8. Show the total number of suppliers. +9. Show the total number of prepayments. +10. Show the total number of workflow records. +11. Show all invoices. +12. Show all purchase orders. +13. Show all purchase order line items. +14. Show all invoice line items. +15. Show all journal entries. +16. Show all documents. +17. Show all tasks. +18. Show all suppliers. +19. Show all prepayments. +20. Show all workflow records. +21. Show the latest invoices by invoice date. +22. Show the latest purchase orders by purchase order date. +23. Show the latest purchase order line items by line creation date. +24. Show the latest invoice line items by invoice line date. +25. Show the latest journal entries by journal date. +26. Show the latest documents by load date. +27. Show the latest tasks by start date. +28. Show the latest suppliers by supplier creation date. +29. Show the latest prepayments by prepayment date. +30. Show the latest workflow records by workflow date. +31. Show the oldest invoices by invoice date. +32. Show the oldest purchase orders by purchase order date. +33. Show the oldest purchase order line items by line creation date. +34. Show the oldest invoice line items by invoice line date. +35. Show the oldest journal entries by journal date. +36. Show the oldest documents by load date. +37. Show the oldest tasks by start date. +38. Show the oldest suppliers by supplier creation date. +39. Show the oldest prepayments by prepayment date. +40. Show the oldest workflow records by workflow date. +41. Show the top 10 invoices by invoice amount. +42. Show the top 10 purchase orders by purchase order amount. +43. Show the top 10 purchase order line items by line amount. +44. Show the top 10 invoice line items by line amount. +45. Show the top 10 journal entries by ending balance. +46. Show the top 10 documents by document amount. +47. Show the top 10 tasks by task duration. +48. Show the top 10 suppliers by supplier spend. +49. Show the top 10 prepayments by prepayment amount. +50. Show the top 10 workflow records by workflow duration. +51. Show the number of invoices by business unit. +52. Show the number of invoices by supplier. +53. Show the number of invoices by product line. +54. Show the number of invoices by product. +55. Show the number of invoices by currency. +56. Show the number of invoices by cost center. +57. Show the number of invoices by profit center. +58. Show the number of invoices by customer. +59. Show the number of invoices by invoice type. +60. Show the number of invoices by priority. +61. Show the number of invoices by entry type. +62. Show the number of invoices by task owner. +63. Show the number of invoices by task status. +64. Show the number of invoices by posting status. +65. Show the number of invoices by journal status. +66. Show the number of invoices by GL account. +67. Show the number of invoices by preparer. +68. Show the number of invoices by reviewer. +69. Show the number of invoices by reason code. +70. Show the number of invoices by workflow status. +71. Show the number of purchase orders by business unit. +72. Show the number of purchase orders by supplier. +73. Show the number of purchase orders by product line. +74. Show the number of purchase orders by product. +75. Show the number of purchase orders by currency. +76. Show the number of purchase orders by cost center. +77. Show the number of purchase orders by profit center. +78. Show the number of purchase orders by customer. +79. Show the number of purchase orders by invoice type. +80. Show the number of purchase orders by priority. +81. Show the number of purchase orders by entry type. +82. Show the number of purchase orders by task owner. +83. Show the number of purchase orders by task status. +84. Show the number of purchase orders by posting status. +85. Show the number of purchase orders by journal status. +86. Show the number of purchase orders by GL account. +87. Show the number of purchase orders by preparer. +88. Show the number of purchase orders by reviewer. +89. Show the number of purchase orders by reason code. +90. Show the number of purchase orders by workflow status. +91. Show the number of purchase order line items by business unit. +92. Show the number of purchase order line items by supplier. +93. Show the number of purchase order line items by product line. +94. Show the number of purchase order line items by product. +95. Show the number of purchase order line items by currency. +96. Show the number of purchase order line items by cost center. +97. Show the number of purchase order line items by profit center. +98. Show the number of purchase order line items by customer. +99. Show the number of purchase order line items by invoice type. +100. Show the number of purchase order line items by priority. +101. Show the number of purchase order line items by entry type. +102. Show the number of purchase order line items by task owner. +103. Show the number of purchase order line items by task status. +104. Show the number of purchase order line items by posting status. +105. Show the number of purchase order line items by journal status. +106. Show the number of purchase order line items by GL account. +107. Show the number of purchase order line items by preparer. +108. Show the number of purchase order line items by reviewer. +109. Show the number of purchase order line items by reason code. +110. Show the number of purchase order line items by workflow status. +111. Show the number of invoice line items by business unit. +112. Show the number of invoice line items by supplier. +113. Show the number of invoice line items by product line. +114. Show the number of invoice line items by product. +115. Show the number of invoice line items by currency. +116. Show the number of invoice line items by cost center. +117. Show the number of invoice line items by profit center. +118. Show the number of invoice line items by customer. +119. Show the number of invoice line items by invoice type. +120. Show the number of invoice line items by priority. +121. Show the number of invoice line items by entry type. +122. Show the number of invoice line items by task owner. +123. Show the number of invoice line items by task status. +124. Show the number of invoice line items by posting status. +125. Show the number of invoice line items by journal status. +126. Show the number of invoice line items by GL account. +127. Show the number of invoice line items by preparer. +128. Show the number of invoice line items by reviewer. +129. Show the number of invoice line items by reason code. +130. Show the number of invoice line items by workflow status. +131. Show the number of journal entries by business unit. +132. Show the number of journal entries by supplier. +133. Show the number of journal entries by product line. +134. Show the number of journal entries by product. +135. Show the number of journal entries by currency. +136. Show the number of journal entries by cost center. +137. Show the number of journal entries by profit center. +138. Show the number of journal entries by customer. +139. Show the number of journal entries by invoice type. +140. Show the number of journal entries by priority. +141. Show the number of journal entries by entry type. +142. Show the number of journal entries by task owner. +143. Show the number of journal entries by task status. +144. Show the number of journal entries by posting status. +145. Show the number of journal entries by journal status. +146. Show the number of journal entries by GL account. +147. Show the number of journal entries by preparer. +148. Show the number of journal entries by reviewer. +149. Show the number of journal entries by reason code. +150. Show the number of journal entries by workflow status. +151. Show the number of documents by business unit. +152. Show the number of documents by supplier. +153. Show the number of documents by product line. +154. Show the number of documents by product. +155. Show the number of documents by currency. +156. Show the number of documents by cost center. +157. Show the number of documents by profit center. +158. Show the number of documents by customer. +159. Show the number of documents by invoice type. +160. Show the number of documents by priority. +161. Show the number of documents by entry type. +162. Show the number of documents by task owner. +163. Show the number of documents by task status. +164. Show the number of documents by posting status. +165. Show the number of documents by journal status. +166. Show the number of documents by GL account. +167. Show the number of documents by preparer. +168. Show the number of documents by reviewer. +169. Show the number of documents by reason code. +170. Show the number of documents by workflow status. +171. Show the number of tasks by business unit. +172. Show the number of tasks by supplier. +173. Show the number of tasks by product line. +174. Show the number of tasks by product. +175. Show the number of tasks by currency. +176. Show the number of tasks by cost center. +177. Show the number of tasks by profit center. +178. Show the number of tasks by customer. +179. Show the number of tasks by invoice type. +180. Show the number of tasks by priority. +181. Show the number of tasks by entry type. +182. Show the number of tasks by task owner. +183. Show the number of tasks by task status. +184. Show the number of tasks by posting status. +185. Show the number of tasks by journal status. +186. Show the number of tasks by GL account. +187. Show the number of tasks by preparer. +188. Show the number of tasks by reviewer. +189. Show the number of tasks by reason code. +190. Show the number of tasks by workflow status. +191. Show the number of suppliers by business unit. +192. Show the number of suppliers by supplier. +193. Show the number of suppliers by product line. +194. Show the number of suppliers by product. +195. Show the number of suppliers by currency. +196. Show the number of suppliers by cost center. +197. Show the number of suppliers by profit center. +198. Show the number of suppliers by customer. +199. Show the number of suppliers by invoice type. +200. Show the number of suppliers by priority. +201. Show the number of suppliers by entry type. +202. Show the number of suppliers by task owner. +203. Show the number of suppliers by task status. +204. Show the number of suppliers by posting status. +205. Show the number of suppliers by journal status. +206. Show the number of suppliers by GL account. +207. Show the number of suppliers by preparer. +208. Show the number of suppliers by reviewer. +209. Show the number of suppliers by reason code. +210. Show the number of suppliers by workflow status. +211. Show the number of prepayments by business unit. +212. Show the number of prepayments by supplier. +213. Show the number of prepayments by product line. +214. Show the number of prepayments by product. +215. Show the number of prepayments by currency. +216. Show the number of prepayments by cost center. +217. Show the number of prepayments by profit center. +218. Show the number of prepayments by customer. +219. Show the number of prepayments by invoice type. +220. Show the number of prepayments by priority. +221. Show the number of prepayments by entry type. +222. Show the number of prepayments by task owner. +223. Show the number of prepayments by task status. +224. Show the number of prepayments by posting status. +225. Show the number of prepayments by journal status. +226. Show the number of prepayments by GL account. +227. Show the number of prepayments by preparer. +228. Show the number of prepayments by reviewer. +229. Show the number of prepayments by reason code. +230. Show the number of prepayments by workflow status. +231. Show the number of workflow records by business unit. +232. Show the number of workflow records by supplier. +233. Show the number of workflow records by product line. +234. Show the number of workflow records by product. +235. Show the number of workflow records by currency. +236. Show the number of workflow records by cost center. +237. Show the number of workflow records by profit center. +238. Show the number of workflow records by customer. +239. Show the number of workflow records by invoice type. +240. Show the number of workflow records by priority. +241. Show the number of workflow records by entry type. +242. Show the number of workflow records by task owner. +243. Show the number of workflow records by task status. +244. Show the number of workflow records by posting status. +245. Show the number of workflow records by journal status. +246. Show the number of workflow records by GL account. +247. Show the number of workflow records by preparer. +248. Show the number of workflow records by reviewer. +249. Show the number of workflow records by reason code. +250. Show the number of workflow records by workflow status. +251. Show the total invoice amount by business unit. +252. Show the total invoice amount by supplier. +253. Show the total invoice amount by product line. +254. Show the total invoice amount by product. +255. Show the total invoice amount by currency. +256. Show the total invoice amount by cost center. +257. Show the total invoice amount by profit center. +258. Show the total invoice amount by customer. +259. Show the total invoice amount by invoice type. +260. Show the total invoice amount by priority. +261. Show the total invoice amount by entry type. +262. Show the total invoice amount by task owner. +263. Show the total invoice amount by task status. +264. Show the total invoice amount by posting status. +265. Show the total invoice amount by journal status. +266. Show the total invoice amount by GL account. +267. Show the total invoice amount by preparer. +268. Show the total invoice amount by reviewer. +269. Show the total invoice amount by reason code. +270. Show the total invoice amount by workflow status. +271. Show the total purchase order amount by business unit. +272. Show the total purchase order amount by supplier. +273. Show the total purchase order amount by product line. +274. Show the total purchase order amount by product. +275. Show the total purchase order amount by currency. +276. Show the total purchase order amount by cost center. +277. Show the total purchase order amount by profit center. +278. Show the total purchase order amount by customer. +279. Show the total purchase order amount by invoice type. +280. Show the total purchase order amount by priority. +281. Show the total purchase order amount by entry type. +282. Show the total purchase order amount by task owner. +283. Show the total purchase order amount by task status. +284. Show the total purchase order amount by posting status. +285. Show the total purchase order amount by journal status. +286. Show the total purchase order amount by GL account. +287. Show the total purchase order amount by preparer. +288. Show the total purchase order amount by reviewer. +289. Show the total purchase order amount by reason code. +290. Show the total purchase order amount by workflow status. +291. Show the total line amount by business unit. +292. Show the total line amount by supplier. +293. Show the total line amount by product line. +294. Show the total line amount by product. +295. Show the total line amount by currency. +296. Show the total line amount by cost center. +297. Show the total line amount by profit center. +298. Show the total line amount by customer. +299. Show the total line amount by invoice type. +300. Show the total line amount by priority. +301. Show the total line amount by entry type. +302. Show the total line amount by task owner. +303. Show the total line amount by task status. +304. Show the total line amount by posting status. +305. Show the total line amount by journal status. +306. Show the total line amount by GL account. +307. Show the total line amount by preparer. +308. Show the total line amount by reviewer. +309. Show the total line amount by reason code. +310. Show the total line amount by workflow status. +311. Show the total ending balance by business unit. +312. Show the total ending balance by supplier. +313. Show the total ending balance by product line. +314. Show the total ending balance by product. +315. Show the total ending balance by currency. +316. Show the total ending balance by cost center. +317. Show the total ending balance by profit center. +318. Show the total ending balance by customer. +319. Show the total ending balance by invoice type. +320. Show the total ending balance by priority. +321. Show the total ending balance by entry type. +322. Show the total ending balance by task owner. +323. Show the total ending balance by task status. +324. Show the total ending balance by posting status. +325. Show the total ending balance by journal status. +326. Show the total ending balance by GL account. +327. Show the total ending balance by preparer. +328. Show the total ending balance by reviewer. +329. Show the total ending balance by reason code. +330. Show the total ending balance by workflow status. +331. Show the total document amount by business unit. +332. Show the total document amount by supplier. +333. Show the total document amount by product line. +334. Show the total document amount by product. +335. Show the total document amount by currency. +336. Show the total document amount by cost center. +337. Show the total document amount by profit center. +338. Show the total document amount by customer. +339. Show the total document amount by invoice type. +340. Show the total document amount by priority. +341. Show the total document amount by entry type. +342. Show the total document amount by task owner. +343. Show the total document amount by task status. +344. Show the total document amount by posting status. +345. Show the total document amount by journal status. +346. Show the total document amount by GL account. +347. Show the total document amount by preparer. +348. Show the total document amount by reviewer. +349. Show the total document amount by reason code. +350. Show the total document amount by workflow status. +351. Show the total task duration by business unit. +352. Show the total task duration by supplier. +353. Show the total task duration by product line. +354. Show the total task duration by product. +355. Show the total task duration by currency. +356. Show the total task duration by cost center. +357. Show the total task duration by profit center. +358. Show the total task duration by customer. +359. Show the total task duration by invoice type. +360. Show the total task duration by priority. +361. Show the total task duration by entry type. +362. Show the total task duration by task owner. +363. Show the total task duration by task status. +364. Show the total task duration by posting status. +365. Show the total task duration by journal status. +366. Show the total task duration by GL account. +367. Show the total task duration by preparer. +368. Show the total task duration by reviewer. +369. Show the total task duration by reason code. +370. Show the total task duration by workflow status. +371. Show the total supplier spend by business unit. +372. Show the total supplier spend by supplier. +373. Show the total supplier spend by product line. +374. Show the total supplier spend by product. +375. Show the total supplier spend by currency. +376. Show the total supplier spend by cost center. +377. Show the total supplier spend by profit center. +378. Show the total supplier spend by customer. +379. Show the total supplier spend by invoice type. +380. Show the total supplier spend by priority. +381. Show the total supplier spend by entry type. +382. Show the total supplier spend by task owner. +383. Show the total supplier spend by task status. +384. Show the total supplier spend by posting status. +385. Show the total supplier spend by journal status. +386. Show the total supplier spend by GL account. +387. Show the total supplier spend by preparer. +388. Show the total supplier spend by reviewer. +389. Show the total supplier spend by reason code. +390. Show the total supplier spend by workflow status. +391. Show the total prepayment amount by business unit. +392. Show the total prepayment amount by supplier. +393. Show the total prepayment amount by product line. +394. Show the total prepayment amount by product. +395. Show the total prepayment amount by currency. +396. Show the total prepayment amount by cost center. +397. Show the total prepayment amount by profit center. +398. Show the total prepayment amount by customer. +399. Show the total prepayment amount by invoice type. +400. Show the total prepayment amount by priority. +401. Show the total prepayment amount by entry type. +402. Show the total prepayment amount by task owner. +403. Show the total prepayment amount by task status. +404. Show the total prepayment amount by posting status. +405. Show the total prepayment amount by journal status. +406. Show the total prepayment amount by GL account. +407. Show the total prepayment amount by preparer. +408. Show the total prepayment amount by reviewer. +409. Show the total prepayment amount by reason code. +410. Show the total prepayment amount by workflow status. +411. Show the total workflow duration by business unit. +412. Show the total workflow duration by supplier. +413. Show the total workflow duration by product line. +414. Show the total workflow duration by product. +415. Show the total workflow duration by currency. +416. Show the total workflow duration by cost center. +417. Show the total workflow duration by profit center. +418. Show the total workflow duration by customer. +419. Show the total workflow duration by invoice type. +420. Show the total workflow duration by priority. +421. Show the total workflow duration by entry type. +422. Show the total workflow duration by task owner. +423. Show the total workflow duration by task status. +424. Show the total workflow duration by posting status. +425. Show the total workflow duration by journal status. +426. Show the total workflow duration by GL account. +427. Show the total workflow duration by preparer. +428. Show the total workflow duration by reviewer. +429. Show the total workflow duration by reason code. +430. Show the total workflow duration by workflow status. +431. Show invoices grouped by business unit. +432. Show invoices grouped by supplier. +433. Show invoices grouped by product line. +434. Show invoices grouped by product. +435. Show invoices grouped by currency. +436. Show invoices grouped by cost center. +437. Show invoices grouped by profit center. +438. Show invoices grouped by customer. +439. Show invoices grouped by invoice type. +440. Show invoices grouped by priority. +441. Show invoices grouped by entry type. +442. Show invoices grouped by task owner. +443. Show invoices grouped by task status. +444. Show invoices grouped by posting status. +445. Show invoices grouped by journal status. +446. Show invoices grouped by GL account. +447. Show invoices grouped by preparer. +448. Show invoices grouped by reviewer. +449. Show invoices grouped by reason code. +450. Show invoices grouped by workflow status. +451. Show purchase orders grouped by business unit. +452. Show purchase orders grouped by supplier. +453. Show purchase orders grouped by product line. +454. Show purchase orders grouped by product. +455. Show purchase orders grouped by currency. +456. Show purchase orders grouped by cost center. +457. Show purchase orders grouped by profit center. +458. Show purchase orders grouped by customer. +459. Show purchase orders grouped by invoice type. +460. Show purchase orders grouped by priority. +461. Show purchase orders grouped by entry type. +462. Show purchase orders grouped by task owner. +463. Show purchase orders grouped by task status. +464. Show purchase orders grouped by posting status. +465. Show purchase orders grouped by journal status. +466. Show purchase orders grouped by GL account. +467. Show purchase orders grouped by preparer. +468. Show purchase orders grouped by reviewer. +469. Show purchase orders grouped by reason code. +470. Show purchase orders grouped by workflow status. +471. Show purchase order line items grouped by business unit. +472. Show purchase order line items grouped by supplier. +473. Show purchase order line items grouped by product line. +474. Show purchase order line items grouped by product. +475. Show purchase order line items grouped by currency. +476. Show purchase order line items grouped by cost center. +477. Show purchase order line items grouped by profit center. +478. Show purchase order line items grouped by customer. +479. Show purchase order line items grouped by invoice type. +480. Show purchase order line items grouped by priority. +481. Show purchase order line items grouped by entry type. +482. Show purchase order line items grouped by task owner. +483. Show purchase order line items grouped by task status. +484. Show purchase order line items grouped by posting status. +485. Show purchase order line items grouped by journal status. +486. Show purchase order line items grouped by GL account. +487. Show purchase order line items grouped by preparer. +488. Show purchase order line items grouped by reviewer. +489. Show purchase order line items grouped by reason code. +490. Show purchase order line items grouped by workflow status. +491. Show invoice line items grouped by business unit. +492. Show invoice line items grouped by supplier. +493. Show invoice line items grouped by product line. +494. Show invoice line items grouped by product. +495. Show invoice line items grouped by currency. +496. Show invoice line items grouped by cost center. +497. Show invoice line items grouped by profit center. +498. Show invoice line items grouped by customer. +499. Show invoice line items grouped by invoice type. +500. Show invoice line items grouped by priority. +501. Show invoice line items grouped by entry type. +502. Show invoice line items grouped by task owner. +503. Show invoice line items grouped by task status. +504. Show invoice line items grouped by posting status. +505. Show invoice line items grouped by journal status. +506. Show invoice line items grouped by GL account. +507. Show invoice line items grouped by preparer. +508. Show invoice line items grouped by reviewer. +509. Show invoice line items grouped by reason code. +510. Show invoice line items grouped by workflow status. +511. Show journal entries grouped by business unit. +512. Show journal entries grouped by supplier. +513. Show journal entries grouped by product line. +514. Show journal entries grouped by product. +515. Show journal entries grouped by currency. +516. Show journal entries grouped by cost center. +517. Show journal entries grouped by profit center. +518. Show journal entries grouped by customer. +519. Show journal entries grouped by invoice type. +520. Show journal entries grouped by priority. +521. Show journal entries grouped by entry type. +522. Show journal entries grouped by task owner. +523. Show journal entries grouped by task status. +524. Show journal entries grouped by posting status. +525. Show journal entries grouped by journal status. +526. Show journal entries grouped by GL account. +527. Show journal entries grouped by preparer. +528. Show journal entries grouped by reviewer. +529. Show journal entries grouped by reason code. +530. Show journal entries grouped by workflow status. +531. Show documents grouped by business unit. +532. Show documents grouped by supplier. +533. Show documents grouped by product line. +534. Show documents grouped by product. +535. Show documents grouped by currency. +536. Show documents grouped by cost center. +537. Show documents grouped by profit center. +538. Show documents grouped by customer. +539. Show documents grouped by invoice type. +540. Show documents grouped by priority. +541. Show documents grouped by entry type. +542. Show documents grouped by task owner. +543. Show documents grouped by task status. +544. Show documents grouped by posting status. +545. Show documents grouped by journal status. +546. Show documents grouped by GL account. +547. Show documents grouped by preparer. +548. Show documents grouped by reviewer. +549. Show documents grouped by reason code. +550. Show documents grouped by workflow status. +551. Show tasks grouped by business unit. +552. Show tasks grouped by supplier. +553. Show tasks grouped by product line. +554. Show tasks grouped by product. +555. Show tasks grouped by currency. +556. Show tasks grouped by cost center. +557. Show tasks grouped by profit center. +558. Show tasks grouped by customer. +559. Show tasks grouped by invoice type. +560. Show tasks grouped by priority. +561. Show tasks grouped by entry type. +562. Show tasks grouped by task owner. +563. Show tasks grouped by task status. +564. Show tasks grouped by posting status. +565. Show tasks grouped by journal status. +566. Show tasks grouped by GL account. +567. Show tasks grouped by preparer. +568. Show tasks grouped by reviewer. +569. Show tasks grouped by reason code. +570. Show tasks grouped by workflow status. +571. Show suppliers grouped by business unit. +572. Show suppliers grouped by supplier. +573. Show suppliers grouped by product line. +574. Show suppliers grouped by product. +575. Show suppliers grouped by currency. +576. Show suppliers grouped by cost center. +577. Show suppliers grouped by profit center. +578. Show suppliers grouped by customer. +579. Show suppliers grouped by invoice type. +580. Show suppliers grouped by priority. +581. Show suppliers grouped by entry type. +582. Show suppliers grouped by task owner. +583. Show suppliers grouped by task status. +584. Show suppliers grouped by posting status. +585. Show suppliers grouped by journal status. +586. Show suppliers grouped by GL account. +587. Show suppliers grouped by preparer. +588. Show suppliers grouped by reviewer. +589. Show suppliers grouped by reason code. +590. Show suppliers grouped by workflow status. +591. Show prepayments grouped by business unit. +592. Show prepayments grouped by supplier. +593. Show prepayments grouped by product line. +594. Show prepayments grouped by product. +595. Show prepayments grouped by currency. +596. Show prepayments grouped by cost center. +597. Show prepayments grouped by profit center. +598. Show prepayments grouped by customer. +599. Show prepayments grouped by invoice type. +600. Show prepayments grouped by priority. +601. Show prepayments grouped by entry type. +602. Show prepayments grouped by task owner. +603. Show prepayments grouped by task status. +604. Show prepayments grouped by posting status. +605. Show prepayments grouped by journal status. +606. Show prepayments grouped by GL account. +607. Show prepayments grouped by preparer. +608. Show prepayments grouped by reviewer. +609. Show prepayments grouped by reason code. +610. Show prepayments grouped by workflow status. +611. Show workflow records grouped by business unit. +612. Show workflow records grouped by supplier. +613. Show workflow records grouped by product line. +614. Show workflow records grouped by product. +615. Show workflow records grouped by currency. +616. Show workflow records grouped by cost center. +617. Show workflow records grouped by profit center. +618. Show workflow records grouped by customer. +619. Show workflow records grouped by invoice type. +620. Show workflow records grouped by priority. +621. Show workflow records grouped by entry type. +622. Show workflow records grouped by task owner. +623. Show workflow records grouped by task status. +624. Show workflow records grouped by posting status. +625. Show workflow records grouped by journal status. +626. Show workflow records grouped by GL account. +627. Show workflow records grouped by preparer. +628. Show workflow records grouped by reviewer. +629. Show workflow records grouped by reason code. +630. Show workflow records grouped by workflow status. +631. Show the top 10 business unit values by number of invoices. +632. Show the top 10 supplier values by number of invoices. +633. Show the top 10 product line values by number of invoices. +634. Show the top 10 product values by number of invoices. +635. Show the top 10 currency values by number of invoices. +636. Show the top 10 cost center values by number of invoices. +637. Show the top 10 profit center values by number of invoices. +638. Show the top 10 customer values by number of invoices. +639. Show the top 10 invoice type values by number of invoices. +640. Show the top 10 priority values by number of invoices. +641. Show the top 10 entry type values by number of invoices. +642. Show the top 10 task owner values by number of invoices. +643. Show the top 10 task status values by number of invoices. +644. Show the top 10 posting status values by number of invoices. +645. Show the top 10 journal status values by number of invoices. +646. Show the top 10 GL account values by number of invoices. +647. Show the top 10 preparer values by number of invoices. +648. Show the top 10 reviewer values by number of invoices. +649. Show the top 10 reason code values by number of invoices. +650. Show the top 10 workflow status values by number of invoices. +651. Show the top 10 business unit values by number of purchase orders. +652. Show the top 10 supplier values by number of purchase orders. +653. Show the top 10 product line values by number of purchase orders. +654. Show the top 10 product values by number of purchase orders. +655. Show the top 10 currency values by number of purchase orders. +656. Show the top 10 cost center values by number of purchase orders. +657. Show the top 10 profit center values by number of purchase orders. +658. Show the top 10 customer values by number of purchase orders. +659. Show the top 10 invoice type values by number of purchase orders. +660. Show the top 10 priority values by number of purchase orders. +661. Show the top 10 entry type values by number of purchase orders. +662. Show the top 10 task owner values by number of purchase orders. +663. Show the top 10 task status values by number of purchase orders. +664. Show the top 10 posting status values by number of purchase orders. +665. Show the top 10 journal status values by number of purchase orders. +666. Show the top 10 GL account values by number of purchase orders. +667. Show the top 10 preparer values by number of purchase orders. +668. Show the top 10 reviewer values by number of purchase orders. +669. Show the top 10 reason code values by number of purchase orders. +670. Show the top 10 workflow status values by number of purchase orders. +671. Show the top 10 business unit values by number of purchase order line items. +672. Show the top 10 supplier values by number of purchase order line items. +673. Show the top 10 product line values by number of purchase order line items. +674. Show the top 10 product values by number of purchase order line items. +675. Show the top 10 currency values by number of purchase order line items. +676. Show the top 10 cost center values by number of purchase order line items. +677. Show the top 10 profit center values by number of purchase order line items. +678. Show the top 10 customer values by number of purchase order line items. +679. Show the top 10 invoice type values by number of purchase order line items. +680. Show the top 10 priority values by number of purchase order line items. +681. Show the top 10 entry type values by number of purchase order line items. +682. Show the top 10 task owner values by number of purchase order line items. +683. Show the top 10 task status values by number of purchase order line items. +684. Show the top 10 posting status values by number of purchase order line items. +685. Show the top 10 journal status values by number of purchase order line items. +686. Show the top 10 GL account values by number of purchase order line items. +687. Show the top 10 preparer values by number of purchase order line items. +688. Show the top 10 reviewer values by number of purchase order line items. +689. Show the top 10 reason code values by number of purchase order line items. +690. Show the top 10 workflow status values by number of purchase order line items. +691. Show the top 10 business unit values by number of invoice line items. +692. Show the top 10 supplier values by number of invoice line items. +693. Show the top 10 product line values by number of invoice line items. +694. Show the top 10 product values by number of invoice line items. +695. Show the top 10 currency values by number of invoice line items. +696. Show the top 10 cost center values by number of invoice line items. +697. Show the top 10 profit center values by number of invoice line items. +698. Show the top 10 customer values by number of invoice line items. +699. Show the top 10 invoice type values by number of invoice line items. +700. Show the top 10 priority values by number of invoice line items. +701. Show the top 10 entry type values by number of invoice line items. +702. Show the top 10 task owner values by number of invoice line items. +703. Show the top 10 task status values by number of invoice line items. +704. Show the top 10 posting status values by number of invoice line items. +705. Show the top 10 journal status values by number of invoice line items. +706. Show the top 10 GL account values by number of invoice line items. +707. Show the top 10 preparer values by number of invoice line items. +708. Show the top 10 reviewer values by number of invoice line items. +709. Show the top 10 reason code values by number of invoice line items. +710. Show the top 10 workflow status values by number of invoice line items. +711. Show the top 10 business unit values by number of journal entries. +712. Show the top 10 supplier values by number of journal entries. +713. Show the top 10 product line values by number of journal entries. +714. Show the top 10 product values by number of journal entries. +715. Show the top 10 currency values by number of journal entries. +716. Show the top 10 cost center values by number of journal entries. +717. Show the top 10 profit center values by number of journal entries. +718. Show the top 10 customer values by number of journal entries. +719. Show the top 10 invoice type values by number of journal entries. +720. Show the top 10 priority values by number of journal entries. +721. Show the top 10 entry type values by number of journal entries. +722. Show the top 10 task owner values by number of journal entries. +723. Show the top 10 task status values by number of journal entries. +724. Show the top 10 posting status values by number of journal entries. +725. Show the top 10 journal status values by number of journal entries. +726. Show the top 10 GL account values by number of journal entries. +727. Show the top 10 preparer values by number of journal entries. +728. Show the top 10 reviewer values by number of journal entries. +729. Show the top 10 reason code values by number of journal entries. +730. Show the top 10 workflow status values by number of journal entries. +731. Show the top 10 business unit values by number of documents. +732. Show the top 10 supplier values by number of documents. +733. Show the top 10 product line values by number of documents. +734. Show the top 10 product values by number of documents. +735. Show the top 10 currency values by number of documents. +736. Show the top 10 cost center values by number of documents. +737. Show the top 10 profit center values by number of documents. +738. Show the top 10 customer values by number of documents. +739. Show the top 10 invoice type values by number of documents. +740. Show the top 10 priority values by number of documents. +741. Show the top 10 entry type values by number of documents. +742. Show the top 10 task owner values by number of documents. +743. Show the top 10 task status values by number of documents. +744. Show the top 10 posting status values by number of documents. +745. Show the top 10 journal status values by number of documents. +746. Show the top 10 GL account values by number of documents. +747. Show the top 10 preparer values by number of documents. +748. Show the top 10 reviewer values by number of documents. +749. Show the top 10 reason code values by number of documents. +750. Show the top 10 workflow status values by number of documents. +751. Show the top 10 business unit values by number of tasks. +752. Show the top 10 supplier values by number of tasks. +753. Show the top 10 product line values by number of tasks. +754. Show the top 10 product values by number of tasks. +755. Show the top 10 currency values by number of tasks. +756. Show the top 10 cost center values by number of tasks. +757. Show the top 10 profit center values by number of tasks. +758. Show the top 10 customer values by number of tasks. +759. Show the top 10 invoice type values by number of tasks. +760. Show the top 10 priority values by number of tasks. +761. Show the top 10 entry type values by number of tasks. +762. Show the top 10 task owner values by number of tasks. +763. Show the top 10 task status values by number of tasks. +764. Show the top 10 posting status values by number of tasks. +765. Show the top 10 journal status values by number of tasks. +766. Show the top 10 GL account values by number of tasks. +767. Show the top 10 preparer values by number of tasks. +768. Show the top 10 reviewer values by number of tasks. +769. Show the top 10 reason code values by number of tasks. +770. Show the top 10 workflow status values by number of tasks. +771. Show the top 10 business unit values by number of suppliers. +772. Show the top 10 supplier values by number of suppliers. +773. Show the top 10 product line values by number of suppliers. +774. Show the top 10 product values by number of suppliers. +775. Show the top 10 currency values by number of suppliers. +776. Show the top 10 cost center values by number of suppliers. +777. Show the top 10 profit center values by number of suppliers. +778. Show the top 10 customer values by number of suppliers. +779. Show the top 10 invoice type values by number of suppliers. +780. Show the top 10 priority values by number of suppliers. +781. Show the top 10 entry type values by number of suppliers. +782. Show the top 10 task owner values by number of suppliers. +783. Show the top 10 task status values by number of suppliers. +784. Show the top 10 posting status values by number of suppliers. +785. Show the top 10 journal status values by number of suppliers. +786. Show the top 10 GL account values by number of suppliers. +787. Show the top 10 preparer values by number of suppliers. +788. Show the top 10 reviewer values by number of suppliers. +789. Show the top 10 reason code values by number of suppliers. +790. Show the top 10 workflow status values by number of suppliers. +791. Show the top 10 business unit values by number of prepayments. +792. Show the top 10 supplier values by number of prepayments. +793. Show the top 10 product line values by number of prepayments. +794. Show the top 10 product values by number of prepayments. +795. Show the top 10 currency values by number of prepayments. +796. Show the top 10 cost center values by number of prepayments. +797. Show the top 10 profit center values by number of prepayments. +798. Show the top 10 customer values by number of prepayments. +799. Show the top 10 invoice type values by number of prepayments. +800. Show the top 10 priority values by number of prepayments. +801. Show the top 10 entry type values by number of prepayments. +802. Show the top 10 task owner values by number of prepayments. +803. Show the top 10 task status values by number of prepayments. +804. Show the top 10 posting status values by number of prepayments. +805. Show the top 10 journal status values by number of prepayments. +806. Show the top 10 GL account values by number of prepayments. +807. Show the top 10 preparer values by number of prepayments. +808. Show the top 10 reviewer values by number of prepayments. +809. Show the top 10 reason code values by number of prepayments. +810. Show the top 10 workflow status values by number of prepayments. +811. Show the top 10 business unit values by number of workflow records. +812. Show the top 10 supplier values by number of workflow records. +813. Show the top 10 product line values by number of workflow records. +814. Show the top 10 product values by number of workflow records. +815. Show the top 10 currency values by number of workflow records. +816. Show the top 10 cost center values by number of workflow records. +817. Show the top 10 profit center values by number of workflow records. +818. Show the top 10 customer values by number of workflow records. +819. Show the top 10 invoice type values by number of workflow records. +820. Show the top 10 priority values by number of workflow records. +821. Show the top 10 entry type values by number of workflow records. +822. Show the top 10 task owner values by number of workflow records. +823. Show the top 10 task status values by number of workflow records. +824. Show the top 10 posting status values by number of workflow records. +825. Show the top 10 journal status values by number of workflow records. +826. Show the top 10 GL account values by number of workflow records. +827. Show the top 10 preparer values by number of workflow records. +828. Show the top 10 reviewer values by number of workflow records. +829. Show the top 10 reason code values by number of workflow records. +830. Show the top 10 workflow status values by number of workflow records. +831. Show invoices for each business unit ordered by count from highest to lowest. +832. Show invoices for each supplier ordered by count from highest to lowest. +833. Show invoices for each product line ordered by count from highest to lowest. +834. Show invoices for each product ordered by count from highest to lowest. +835. Show invoices for each currency ordered by count from highest to lowest. +836. Show invoices for each cost center ordered by count from highest to lowest. +837. Show invoices for each profit center ordered by count from highest to lowest. +838. Show invoices for each customer ordered by count from highest to lowest. +839. Show invoices for each invoice type ordered by count from highest to lowest. +840. Show invoices for each priority ordered by count from highest to lowest. +841. Show invoices for each entry type ordered by count from highest to lowest. +842. Show invoices for each task owner ordered by count from highest to lowest. +843. Show invoices for each task status ordered by count from highest to lowest. +844. Show invoices for each posting status ordered by count from highest to lowest. +845. Show invoices for each journal status ordered by count from highest to lowest. +846. Show invoices for each GL account ordered by count from highest to lowest. +847. Show invoices for each preparer ordered by count from highest to lowest. +848. Show invoices for each reviewer ordered by count from highest to lowest. +849. Show invoices for each reason code ordered by count from highest to lowest. +850. Show invoices for each workflow status ordered by count from highest to lowest. +851. Show purchase orders for each business unit ordered by count from highest to lowest. +852. Show purchase orders for each supplier ordered by count from highest to lowest. +853. Show purchase orders for each product line ordered by count from highest to lowest. +854. Show purchase orders for each product ordered by count from highest to lowest. +855. Show purchase orders for each currency ordered by count from highest to lowest. +856. Show purchase orders for each cost center ordered by count from highest to lowest. +857. Show purchase orders for each profit center ordered by count from highest to lowest. +858. Show purchase orders for each customer ordered by count from highest to lowest. +859. Show purchase orders for each invoice type ordered by count from highest to lowest. +860. Show purchase orders for each priority ordered by count from highest to lowest. +861. Show purchase orders for each entry type ordered by count from highest to lowest. +862. Show purchase orders for each task owner ordered by count from highest to lowest. +863. Show purchase orders for each task status ordered by count from highest to lowest. +864. Show purchase orders for each posting status ordered by count from highest to lowest. +865. Show purchase orders for each journal status ordered by count from highest to lowest. +866. Show purchase orders for each GL account ordered by count from highest to lowest. +867. Show purchase orders for each preparer ordered by count from highest to lowest. +868. Show purchase orders for each reviewer ordered by count from highest to lowest. +869. Show purchase orders for each reason code ordered by count from highest to lowest. +870. Show purchase orders for each workflow status ordered by count from highest to lowest. +871. Show purchase order line items for each business unit ordered by count from highest to lowest. +872. Show purchase order line items for each supplier ordered by count from highest to lowest. +873. Show purchase order line items for each product line ordered by count from highest to lowest. +874. Show purchase order line items for each product ordered by count from highest to lowest. +875. Show purchase order line items for each currency ordered by count from highest to lowest. +876. Show purchase order line items for each cost center ordered by count from highest to lowest. +877. Show purchase order line items for each profit center ordered by count from highest to lowest. +878. Show purchase order line items for each customer ordered by count from highest to lowest. +879. Show purchase order line items for each invoice type ordered by count from highest to lowest. +880. Show purchase order line items for each priority ordered by count from highest to lowest. +881. Show purchase order line items for each entry type ordered by count from highest to lowest. +882. Show purchase order line items for each task owner ordered by count from highest to lowest. +883. Show purchase order line items for each task status ordered by count from highest to lowest. +884. Show purchase order line items for each posting status ordered by count from highest to lowest. +885. Show purchase order line items for each journal status ordered by count from highest to lowest. +886. Show purchase order line items for each GL account ordered by count from highest to lowest. +887. Show purchase order line items for each preparer ordered by count from highest to lowest. +888. Show purchase order line items for each reviewer ordered by count from highest to lowest. +889. Show purchase order line items for each reason code ordered by count from highest to lowest. +890. Show purchase order line items for each workflow status ordered by count from highest to lowest. +891. Show invoice line items for each business unit ordered by count from highest to lowest. +892. Show invoice line items for each supplier ordered by count from highest to lowest. +893. Show invoice line items for each product line ordered by count from highest to lowest. +894. Show invoice line items for each product ordered by count from highest to lowest. +895. Show invoice line items for each currency ordered by count from highest to lowest. +896. Show invoice line items for each cost center ordered by count from highest to lowest. +897. Show invoice line items for each profit center ordered by count from highest to lowest. +898. Show invoice line items for each customer ordered by count from highest to lowest. +899. Show invoice line items for each invoice type ordered by count from highest to lowest. +900. Show invoice line items for each priority ordered by count from highest to lowest. +901. Show invoice line items for each entry type ordered by count from highest to lowest. +902. Show invoice line items for each task owner ordered by count from highest to lowest. +903. Show invoice line items for each task status ordered by count from highest to lowest. +904. Show invoice line items for each posting status ordered by count from highest to lowest. +905. Show invoice line items for each journal status ordered by count from highest to lowest. +906. Show invoice line items for each GL account ordered by count from highest to lowest. +907. Show invoice line items for each preparer ordered by count from highest to lowest. +908. Show invoice line items for each reviewer ordered by count from highest to lowest. +909. Show invoice line items for each reason code ordered by count from highest to lowest. +910. Show invoice line items for each workflow status ordered by count from highest to lowest. +911. Show journal entries for each business unit ordered by count from highest to lowest. +912. Show journal entries for each supplier ordered by count from highest to lowest. +913. Show journal entries for each product line ordered by count from highest to lowest. +914. Show journal entries for each product ordered by count from highest to lowest. +915. Show journal entries for each currency ordered by count from highest to lowest. +916. Show journal entries for each cost center ordered by count from highest to lowest. +917. Show journal entries for each profit center ordered by count from highest to lowest. +918. Show journal entries for each customer ordered by count from highest to lowest. +919. Show journal entries for each invoice type ordered by count from highest to lowest. +920. Show journal entries for each priority ordered by count from highest to lowest. +921. Show journal entries for each entry type ordered by count from highest to lowest. +922. Show journal entries for each task owner ordered by count from highest to lowest. +923. Show journal entries for each task status ordered by count from highest to lowest. +924. Show journal entries for each posting status ordered by count from highest to lowest. +925. Show journal entries for each journal status ordered by count from highest to lowest. +926. Show journal entries for each GL account ordered by count from highest to lowest. +927. Show journal entries for each preparer ordered by count from highest to lowest. +928. Show journal entries for each reviewer ordered by count from highest to lowest. +929. Show journal entries for each reason code ordered by count from highest to lowest. +930. Show journal entries for each workflow status ordered by count from highest to lowest. +931. Show documents for each business unit ordered by count from highest to lowest. +932. Show documents for each supplier ordered by count from highest to lowest. +933. Show documents for each product line ordered by count from highest to lowest. +934. Show documents for each product ordered by count from highest to lowest. +935. Show documents for each currency ordered by count from highest to lowest. +936. Show documents for each cost center ordered by count from highest to lowest. +937. Show documents for each profit center ordered by count from highest to lowest. +938. Show documents for each customer ordered by count from highest to lowest. +939. Show documents for each invoice type ordered by count from highest to lowest. +940. Show documents for each priority ordered by count from highest to lowest. +941. Show documents for each entry type ordered by count from highest to lowest. +942. Show documents for each task owner ordered by count from highest to lowest. +943. Show documents for each task status ordered by count from highest to lowest. +944. Show documents for each posting status ordered by count from highest to lowest. +945. Show documents for each journal status ordered by count from highest to lowest. +946. Show documents for each GL account ordered by count from highest to lowest. +947. Show documents for each preparer ordered by count from highest to lowest. +948. Show documents for each reviewer ordered by count from highest to lowest. +949. Show documents for each reason code ordered by count from highest to lowest. +950. Show documents for each workflow status ordered by count from highest to lowest. +951. Show tasks for each business unit ordered by count from highest to lowest. +952. Show tasks for each supplier ordered by count from highest to lowest. +953. Show tasks for each product line ordered by count from highest to lowest. +954. Show tasks for each product ordered by count from highest to lowest. +955. Show tasks for each currency ordered by count from highest to lowest. +956. Show tasks for each cost center ordered by count from highest to lowest. +957. Show tasks for each profit center ordered by count from highest to lowest. +958. Show tasks for each customer ordered by count from highest to lowest. +959. Show tasks for each invoice type ordered by count from highest to lowest. +960. Show tasks for each priority ordered by count from highest to lowest. +961. Show tasks for each entry type ordered by count from highest to lowest. +962. Show tasks for each task owner ordered by count from highest to lowest. +963. Show tasks for each task status ordered by count from highest to lowest. +964. Show tasks for each posting status ordered by count from highest to lowest. +965. Show tasks for each journal status ordered by count from highest to lowest. +966. Show tasks for each GL account ordered by count from highest to lowest. +967. Show tasks for each preparer ordered by count from highest to lowest. +968. Show tasks for each reviewer ordered by count from highest to lowest. +969. Show tasks for each reason code ordered by count from highest to lowest. +970. Show tasks for each workflow status ordered by count from highest to lowest. +971. Show suppliers for each business unit ordered by count from highest to lowest. +972. Show suppliers for each supplier ordered by count from highest to lowest. +973. Show suppliers for each product line ordered by count from highest to lowest. +974. Show suppliers for each product ordered by count from highest to lowest. +975. Show suppliers for each currency ordered by count from highest to lowest. +976. Show suppliers for each cost center ordered by count from highest to lowest. +977. Show suppliers for each profit center ordered by count from highest to lowest. +978. Show suppliers for each customer ordered by count from highest to lowest. +979. Show suppliers for each invoice type ordered by count from highest to lowest. +980. Show suppliers for each priority ordered by count from highest to lowest. +981. Show suppliers for each entry type ordered by count from highest to lowest. +982. Show suppliers for each task owner ordered by count from highest to lowest. +983. Show suppliers for each task status ordered by count from highest to lowest. +984. Show suppliers for each posting status ordered by count from highest to lowest. +985. Show suppliers for each journal status ordered by count from highest to lowest. +986. Show suppliers for each GL account ordered by count from highest to lowest. +987. Show suppliers for each preparer ordered by count from highest to lowest. +988. Show suppliers for each reviewer ordered by count from highest to lowest. +989. Show suppliers for each reason code ordered by count from highest to lowest. +990. Show suppliers for each workflow status ordered by count from highest to lowest. +991. Show prepayments for each business unit ordered by count from highest to lowest. +992. Show prepayments for each supplier ordered by count from highest to lowest. +993. Show prepayments for each product line ordered by count from highest to lowest. +994. Show prepayments for each product ordered by count from highest to lowest. +995. Show prepayments for each currency ordered by count from highest to lowest. +996. Show prepayments for each cost center ordered by count from highest to lowest. +997. Show prepayments for each profit center ordered by count from highest to lowest. +998. Show prepayments for each customer ordered by count from highest to lowest. +999. Show prepayments for each invoice type ordered by count from highest to lowest. +1000. Show prepayments for each priority ordered by count from highest to lowest. diff --git a/wren-ai-service/docs/config_examples/config.anthropic.yaml b/wren-ai-service/docs/config_examples/config.anthropic.yaml index 76e5d96526..a781436cdd 100644 --- a/wren-ai-service/docs/config_examples/config.anthropic.yaml +++ b/wren-ai-service/docs/config_examples/config.anthropic.yaml @@ -152,10 +152,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.azure.yaml b/wren-ai-service/docs/config_examples/config.azure.yaml index 9319394727..bc3869a809 100644 --- a/wren-ai-service/docs/config_examples/config.azure.yaml +++ b/wren-ai-service/docs/config_examples/config.azure.yaml @@ -19,7 +19,7 @@ models: temperature: 0 n: 1 seed: 0 - max_tokens: 4096 + max_tokens: 8192 --- type: embedder @@ -165,10 +165,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.bedrock.yaml b/wren-ai-service/docs/config_examples/config.bedrock.yaml index a94474de67..02f0633eb3 100644 --- a/wren-ai-service/docs/config_examples/config.bedrock.yaml +++ b/wren-ai-service/docs/config_examples/config.bedrock.yaml @@ -168,10 +168,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.deepseek.yaml b/wren-ai-service/docs/config_examples/config.deepseek.yaml index 3a00b30ad8..7749844074 100644 --- a/wren-ai-service/docs/config_examples/config.deepseek.yaml +++ b/wren-ai-service/docs/config_examples/config.deepseek.yaml @@ -175,10 +175,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml index 9d087accb3..f54d85a082 100644 --- a/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.google_ai_studio.yaml @@ -161,10 +161,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.google_vertexai.yaml b/wren-ai-service/docs/config_examples/config.google_vertexai.yaml index 0b29acb7e5..7bbb6206ab 100644 --- a/wren-ai-service/docs/config_examples/config.google_vertexai.yaml +++ b/wren-ai-service/docs/config_examples/config.google_vertexai.yaml @@ -18,7 +18,7 @@ models: n: 1 temperature: 0 seed: 0 - max_tokens: 4096 + max_tokens: 8192 - model: vertex_ai/gemini-2.5-flash # vertex_ai/ alias: gemini-llm-for-chart timeout: 120 @@ -26,7 +26,7 @@ models: n: 1 temperature: 0 seed: 0 - max_tokens: 4096 + max_tokens: 8192 response_format: type: json_object @@ -169,10 +169,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.grok.yaml b/wren-ai-service/docs/config_examples/config.grok.yaml index b01de55b7e..e5270512b8 100644 --- a/wren-ai-service/docs/config_examples/config.grok.yaml +++ b/wren-ai-service/docs/config_examples/config.grok.yaml @@ -15,7 +15,7 @@ models: timeout: 120 kwargs: temperature: 0 - max_tokens: 4096 + max_tokens: 8192 n: 1 --- @@ -157,10 +157,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.groq.yaml b/wren-ai-service/docs/config_examples/config.groq.yaml index a07a577e9d..0f7c1fc05f 100644 --- a/wren-ai-service/docs/config_examples/config.groq.yaml +++ b/wren-ai-service/docs/config_examples/config.groq.yaml @@ -156,10 +156,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.lm_studio.yaml b/wren-ai-service/docs/config_examples/config.lm_studio.yaml index 131ac04384..ef29353d93 100644 --- a/wren-ai-service/docs/config_examples/config.lm_studio.yaml +++ b/wren-ai-service/docs/config_examples/config.lm_studio.yaml @@ -155,10 +155,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.ollama.yaml b/wren-ai-service/docs/config_examples/config.ollama.yaml index 5f8e6c4ea2..06c41f1e4d 100644 --- a/wren-ai-service/docs/config_examples/config.ollama.yaml +++ b/wren-ai-service/docs/config_examples/config.ollama.yaml @@ -14,6 +14,7 @@ models: timeout: 600 kwargs: n: 1 + num_predict: 8192 temperature: 0 --- @@ -155,10 +156,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.open_router.yaml b/wren-ai-service/docs/config_examples/config.open_router.yaml index ecbcaa0731..0bebc90d25 100644 --- a/wren-ai-service/docs/config_examples/config.open_router.yaml +++ b/wren-ai-service/docs/config_examples/config.open_router.yaml @@ -153,10 +153,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.qwen3.yaml b/wren-ai-service/docs/config_examples/config.qwen3.yaml index 0ebaed162a..1fa694ed4d 100644 --- a/wren-ai-service/docs/config_examples/config.qwen3.yaml +++ b/wren-ai-service/docs/config_examples/config.qwen3.yaml @@ -195,10 +195,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/config_examples/config.zhipu.yaml b/wren-ai-service/docs/config_examples/config.zhipu.yaml index 8d0db87d6c..7d762b33ce 100644 --- a/wren-ai-service/docs/config_examples/config.zhipu.yaml +++ b/wren-ai-service/docs/config_examples/config.zhipu.yaml @@ -203,10 +203,10 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 query_cache_maxsize: 1000 query_cache_ttl: 3600 langfuse_host: https://cloud.langfuse.com diff --git a/wren-ai-service/docs/configuration.md b/wren-ai-service/docs/configuration.md index c88c39ce8d..fc0a442bba 100644 --- a/wren-ai-service/docs/configuration.md +++ b/wren-ai-service/docs/configuration.md @@ -54,7 +54,7 @@ The configuration file (`config.yaml`) is structured into several sections, each kwargs: temperature: 0 n: 1 - max_tokens: 4096 + max_tokens: 8192 response_format: type: "json_object" - model: gpt-4o-mini diff --git a/wren-ai-service/src/__main__.py b/wren-ai-service/src/__main__.py index de141c3fde..6ff2d11165 100644 --- a/wren-ai-service/src/__main__.py +++ b/wren-ai-service/src/__main__.py @@ -1,10 +1,14 @@ from contextlib import asynccontextmanager +from importlib.util import find_spec +from pathlib import Path import uvicorn from fastapi import FastAPI from fastapi.exceptions import RequestValidationError from fastapi.middleware.cors import CORSMiddleware +from fastapi.openapi.docs import get_swagger_ui_html from fastapi.responses import ORJSONResponse, RedirectResponse +from fastapi.staticfiles import StaticFiles from langfuse.decorators import langfuse_context from src.config import settings @@ -24,6 +28,22 @@ ) +def get_local_swagger_static_dir() -> Path | None: + litellm_spec = find_spec("litellm") + if not litellm_spec or not litellm_spec.submodule_search_locations: + return None + + swagger_dir = ( + Path(next(iter(litellm_spec.submodule_search_locations))) / "proxy" / "swagger" + ) + required_assets = ("swagger-ui-bundle.js", "swagger-ui.css", "favicon.ico") + + if all((swagger_dir / asset).is_file() for asset in required_assets): + return swagger_dir + + return None + + # https://fastapi.tiangolo.com/advanced/events/#lifespan @asynccontextmanager async def lifespan(app: FastAPI): @@ -42,10 +62,19 @@ async def lifespan(app: FastAPI): app = FastAPI( title="wren-ai-service API Docs", lifespan=lifespan, + docs_url=None, redoc_url=None, default_response_class=ORJSONResponse, ) +swagger_static_dir = get_local_swagger_static_dir() +if swagger_static_dir: + app.mount( + "/_docs/static", + StaticFiles(directory=swagger_static_dir), + name="swagger-static", + ) + app.add_middleware( CORSMiddleware, allow_origins=["*"], @@ -81,6 +110,25 @@ def root(): return RedirectResponse(url="/docs") +@app.get("/docs", include_in_schema=False) +def swagger_ui_html(): + kwargs = { + "openapi_url": app.openapi_url, + "title": f"{app.title} - Swagger UI", + } + + if swagger_static_dir: + kwargs.update( + { + "swagger_js_url": "/_docs/static/swagger-ui-bundle.js", + "swagger_css_url": "/_docs/static/swagger-ui.css", + "swagger_favicon_url": "/_docs/static/favicon.ico", + } + ) + + return get_swagger_ui_html(**kwargs) + + @app.get("/health") def health(): return {"status": "ok"} @@ -95,6 +143,6 @@ def health(): reload_includes=["src/**/*.py", ".env.dev", "config.yaml"], reload_excludes=["tests/**/*.py", "eval/**/*.py"], workers=1, - loop="uvloop", - http="httptools", + loop="auto", + http="auto", ) diff --git a/wren-ai-service/src/config.py b/wren-ai-service/src/config.py index c5acf4ae47..7d9a686dd2 100644 --- a/wren-ai-service/src/config.py +++ b/wren-ai-service/src/config.py @@ -29,7 +29,7 @@ class Settings(BaseSettings): column_indexing_batch_size: int = Field(default=50) table_retrieval_size: int = Field(default=10) table_column_retrieval_size: int = Field(default=100) - enable_column_pruning: bool = Field(default=False) + enable_column_pruning: bool = Field(default=True) historical_question_retrieval_similarity_threshold: float = Field(default=0.9) sql_pairs_similarity_threshold: float = Field(default=0.7) sql_pairs_retrieval_max_size: int = Field(default=10) @@ -38,12 +38,22 @@ class Settings(BaseSettings): # generation config allow_intent_classification: bool = Field(default=True) - allow_sql_generation_reasoning: bool = Field(default=True) + allow_sql_generation_reasoning: bool = Field(default=False) allow_sql_functions_retrieval: bool = Field(default=True) allow_sql_diagnosis: bool = Field(default=True) allow_sql_knowledge_retrieval: bool = Field(default=False) max_histories: int = Field(default=5) - max_sql_correction_retries: int = Field(default=3) + max_sql_correction_retries: int = Field(default=0) + sql_generation_timeout_seconds: float = Field(default=45.0) + # Kept for compatibility with deployed configs. This controls the bounded + # generation/correction calls when present. + pipeline_timeout_seconds: float | None = Field(default=None) + relationship_recommendation_timeout_seconds: float = Field(default=180.0) + semantics_description_timeout_seconds: float | None = Field(default=None) + semantics_description_generation_timeout_seconds: float = Field(default=120.0) + semantics_description_max_models_per_batch: int = Field(default=4) + semantics_description_max_columns_per_batch: int = Field(default=50) + semantics_description_max_concurrent_tasks: int = Field(default=4) # engine config engine_timeout: float = Field(default=30.0) diff --git a/wren-ai-service/src/globals.py b/wren-ai-service/src/globals.py index 9343344616..6ab877b352 100644 --- a/wren-ai-service/src/globals.py +++ b/wren-ai-service/src/globals.py @@ -145,6 +145,9 @@ def create_service_container( **pipe_components["followup_sql_generation_reasoning"], ), "sql_correction": _sql_correction_pipeline, + "sql_tables_extraction": generation.SQLTablesExtraction( + **pipe_components["sql_tables_extraction"], + ), "followup_sql_generation": generation.FollowUpSQLGeneration( **pipe_components["followup_sql_generation"], ), @@ -219,6 +222,9 @@ def create_service_container( **pipe_components["relationship_recommendation"], ) }, + generation_timeout_seconds=( + settings.relationship_recommendation_timeout_seconds + ), **query_cache, ), question_recommendation=services.QuestionRecommendation( diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py index f6114d63b1..e7fe653e8d 100644 --- a/wren-ai-service/src/pipelines/common.py +++ b/wren-ai-service/src/pipelines/common.py @@ -1,13 +1,34 @@ +import asyncio import re from typing import Any, List, Optional, Tuple from haystack import Document, component -def get_engine_supported_data_type(data_type: str) -> str: +def build_project_deploy_filter( + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, +) -> dict[str, Any] | None: + conditions = [] + + if project_id: + conditions.append( + {"field": "project_id", "operator": "==", "value": project_id} + ) + + if mdl_hash: + conditions.append({"field": "mdl_hash", "operator": "==", "value": mdl_hash}) + + return {"operator": "AND", "conditions": conditions} if conditions else None + + +def get_engine_supported_data_type(data_type: str | None) -> str: """ This function makes sure downstream ai pipeline get column data types in a format that is supported by the data engine. """ + if not data_type: + return "UNKNOWN" + match data_type.upper(): case "BPCHAR" | "NAME" | "UUID" | "INET": return "VARCHAR" @@ -28,34 +49,59 @@ def get_engine_supported_data_type(data_type: str) -> str: def build_table_ddl( - content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None + content: dict, + columns: Optional[set[str]] = None, + tables: Optional[set[str]] = None, + include_semantic_comments: bool = True, ) -> Tuple[str, bool, bool]: columns_ddl = [] has_calculated_field = False has_json_field = False + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) for column in content["columns"]: if column["type"] == "COLUMN": + raw_data_type = column["data_type"] + supported_data_type = get_engine_supported_data_type(raw_data_type) if ( - (not columns or (columns and column["name"] in columns)) - and column["data_type"].lower() - != "unknown" # quick fix: filtering out UNKNOWN column type + ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) + and ( + raw_data_type is None + or supported_data_type.lower() + != "unknown" # quick fix: filtering out UNKNOWN column type + ) ): if "This column is a Calculated Field" in column["comment"]: has_calculated_field = True - if column["data_type"].lower() == "json": + if supported_data_type.lower() == "json": has_json_field = True - column_ddl = f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + column_comment = column["comment"] if include_semantic_comments else "" + column_ddl = f"{column_comment}{column['name']} {supported_data_type}" if column["is_primary_key"]: column_ddl += " PRIMARY KEY" columns_ddl.append(column_ddl) elif column["type"] == "FOREIGN_KEY": - if not tables or (tables and set(column["tables"]).issubset(tables)): - columns_ddl.append(f"{column['comment']}{column['constraint']}") + if not tables or (tables and set(column.get("tables", [])).issubset(tables)): + relationship_comment = ( + column["comment"] if include_semantic_comments else "" + ) + columns_ddl.append(f"{relationship_comment}{column['constraint']}") + table_comment = content["comment"] if include_semantic_comments else "" return ( ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{table_comment}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ), @@ -64,15 +110,42 @@ def build_table_ddl( ) -async def retrieve_metadata(project_id: str, retriever) -> dict[str, Any]: - filters = None - if project_id: - filters = { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } +async def retrieve_metadata( + project_id: str, + retriever, + mdl_hash: Optional[str] = None, +) -> dict[str, Any]: + cache_key = ( + id(retriever), + str(project_id), + str(mdl_hash), + ) + if project_id and mdl_hash: + if cache_key in _METADATA_CACHE: + return _METADATA_CACHE[cache_key] + + lock = _METADATA_CACHE_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + if cache_key in _METADATA_CACHE: + return _METADATA_CACHE[cache_key] + + metadata = await _retrieve_metadata_uncached(project_id, retriever, mdl_hash) + _METADATA_CACHE[cache_key] = metadata + return metadata + + return await _retrieve_metadata_uncached(project_id, retriever, mdl_hash) + + +_METADATA_CACHE: dict[tuple[int, str, str], dict[str, Any]] = {} +_METADATA_CACHE_LOCKS: dict[tuple[int, str, str], asyncio.Lock] = {} + + +async def _retrieve_metadata_uncached( + project_id: str, + retriever, + mdl_hash: Optional[str] = None, +) -> dict[str, Any]: + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) result = await retriever.run(query_embedding=[], filters=filters) documents = result["documents"] diff --git a/wren-ai-service/src/pipelines/generation/data_assistance.py b/wren-ai-service/src/pipelines/generation/data_assistance.py index 51b91197f9..e10ed4a530 100644 --- a/wren-ai-service/src/pipelines/generation/data_assistance.py +++ b/wren-ai-service/src/pipelines/generation/data_assistance.py @@ -29,6 +29,7 @@ - There should be proper line breaks, whitespace, and Markdown formatting(headers, lists, tables, etc.) in your response. - If the language is Traditional/Simplified Chinese, Korean, or Japanese, the maximum response length is 150 words; otherwise, the maximum response length is 110 words. - MUST NOT add SQL code in your response. +- Use only the provided DATABASE SCHEMA as context. Do not invent, assume, or name tables or columns that are not present there; do not provide hypothetical schema. - If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. ### OUTPUT FORMAT ### diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py index 35cfb8fccf..2ebec98c54 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation.py @@ -5,8 +5,8 @@ from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider @@ -16,10 +16,13 @@ SQLGenPostProcessor, construct_ask_history_messages, construct_instructions, + construct_schema_identifier_catalog, + generate_simple_analytics_sql, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, get_sql_generation_system_prompt, + unsupported_schema_generation_result, ) from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge @@ -34,6 +37,11 @@ Given the following user's follow-up question and previous SQL query and summary, generate one SQL query to best answer user's question. +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -60,11 +68,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} -Summary: -{{sample.summary}} -SQL: -{{sample.sql}} +Question: +{{sample.question}} {% endfor %} {% endif %} @@ -77,11 +84,9 @@ ### QUESTION ### User's Follow-up Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, relationships, and history only to understand meaning; the SQL must use exact declared table and column names from the WREN SQL IDENTIFIER CONTRACT and DATABASE SCHEMA. Treat source metadata, physical names, lineage names, semantic labels, user question words, and prior SQL as non-executable background unless the exact same identifier is declared in the contract. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. -### REASONING PLAN ### -{{ sql_generation_reasoning }} - -Let's think step by step. +Return only the final JSON SQL response. """ @@ -103,7 +108,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -120,6 +124,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -147,13 +152,21 @@ async def post_process( generate_sql_in_followup: dict, post_processor: SQLGenPostProcessor, data_source: str, + query: str | None = None, + grounding_query: str | None = None, + documents: list[str] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_in_followup.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, + contexts=validation_contexts or documents, + fallback_query=grounding_query or query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -201,6 +214,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -208,22 +222,41 @@ async def run( use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + grounding_query: str | None = None, + validation_contexts: list[str] | None = None, ): logger.info("Follow-Up SQL Generation pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) + if project_id or use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) else: metadata = {} + data_source = metadata.get("data_source", "local_file") + + effective_grounding_query = grounding_query or query + unsupported_result = unsupported_schema_generation_result( + effective_grounding_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "Follow-up SQL generation skipped before LLM because selected schema does not cover requested concepts." + ) + return {"post_process": unsupported_result} return await self._pipe.execute( ["post_process"], inputs={ "query": query, + "grounding_query": effective_grounding_query, "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "histories": histories, "project_id": project_id, + "mdl_hash": mdl_hash, "sql_samples": sql_samples, "instructions": instructions, "has_calculated_field": has_calculated_field, @@ -232,8 +265,67 @@ async def run( "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) + + async def run_deterministic_fast_path( + self, + query: str, + contexts: list[str], + project_id: str | None = None, + mdl_hash: str | None = None, + use_dry_plan: bool = False, + allow_dry_plan_fallback: bool = True, + grounding_query: str | None = None, + ) -> dict | None: + if use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) + else: + metadata = {} + data_source = metadata.get("data_source", "local_file") + effective_grounding_query = grounding_query or query + + unsupported_result = unsupported_schema_generation_result( + effective_grounding_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "Follow-up SQL deterministic fast path returned unsupported schema before LLM." + ) + return {"post_process": unsupported_result, "fast_path": "unsupported"} + + deterministic_sql = generate_simple_analytics_sql( + effective_grounding_query, + contexts, + ) + if not deterministic_sql: + return None + + logger.info("Follow-up SQL deterministic fast path produced a candidate.") + post_process = await self._components["post_processor"].run( + [deterministic_sql], + project_id=project_id, + mdl_hash=mdl_hash, + contexts=contexts, + fallback_query=effective_grounding_query, + use_dry_plan=use_dry_plan, + data_source=data_source, + allow_dry_plan_fallback=allow_dry_plan_fallback, + ) + if post_process.get("valid_generation_result"): + logger.info("Follow-up SQL deterministic fast path accepted candidate.") + return {"post_process": post_process, "fast_path": "deterministic"} + + logger.info( + "Follow-up SQL deterministic fast path rejected candidate; continuing to LLM. reason=%s", + post_process.get("invalid_generation_result", {}).get("error"), + ) + return None diff --git a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py index 42b28c5b8f..3831186d9f 100644 --- a/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/followup_sql_generation_reasoning.py @@ -13,6 +13,8 @@ from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_schema_identifier_catalog, + sanitize_sql_generation_reasoning, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -23,6 +25,13 @@ sql_generation_reasoning_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +Use exact model, column, and relationship names from that contract whenever reasoning names schema objects. +Treat source metadata, physical names, lineage names, semantic labels, prior SQL, and user wording as non-executable background unless the exact same identifier is declared in the contract. +If the contract does not contain a table, column, or relationship required by the user's intent, say the current schema context is insufficient instead of naming an assumed object. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -30,11 +39,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} @@ -71,16 +79,19 @@ def prompt( sql_samples: list[dict], instructions: list[dict], prompt_builder: PromptBuilder, + validation_contexts: Optional[list[str]] = None, configuration: Configuration | None = Configuration(), ) -> dict: + schema_documents = validation_contexts or documents _prompt = prompt_builder.run( query=query, - documents=documents, + documents=schema_documents, histories=histories, sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, ), + schema_identifier_catalog=construct_schema_identifier_catalog(schema_documents), language=configuration.language, current_time=configuration.show_current_time(), ) @@ -105,7 +116,7 @@ async def generate_sql_reasoning( def post_process( generate_sql_reasoning: dict, ) -> dict: - return generate_sql_reasoning.get("replies")[0] + return sanitize_sql_generation_reasoning(generate_sql_reasoning.get("replies")[0]) ## End of Pipeline @@ -174,6 +185,9 @@ async def run( histories: list[AskHistory], sql_samples: Optional[list[dict]] = None, instructions: Optional[list[dict]] = None, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + validation_contexts: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, ): @@ -186,6 +200,9 @@ async def run( "histories": histories, "sql_samples": sql_samples or [], "instructions": instructions or [], + "project_id": project_id, + "mdl_hash": mdl_hash, + "validation_contexts": validation_contexts, "configuration": configuration, "query_id": query_id, **self._components, diff --git a/wren-ai-service/src/pipelines/generation/intent_classification.py b/wren-ai-service/src/pipelines/generation/intent_classification.py index 4d6cd313cd..1ae2274aac 100644 --- a/wren-ai-service/src/pipelines/generation/intent_classification.py +++ b/wren-ai-service/src/pipelines/generation/intent_classification.py @@ -13,7 +13,11 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider -from src.pipelines.common import build_table_ddl, clean_up_new_lines +from src.pipelines.common import ( + build_project_deploy_filter, + build_table_ddl, + clean_up_new_lines, +) from src.pipelines.generation.utils.sql import construct_instructions from src.utils import trace_cost from src.web.v1.services import Configuration @@ -22,6 +26,17 @@ logger = logging.getLogger("wren-ai-service") +def _project_filter_conditions( + project_id: str | None, + mdl_hash: str | None = None, +) -> list[dict[str, Any]]: + project_deploy_filter = build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ) + return project_deploy_filter["conditions"] if project_deploy_filter else [] + + intent_classification_system_prompt = """ ### Task ### You are an expert detective specializing in intent classification. Combine the user's current question and previous questions to determine their true intent based on the provided database schema. Classify the intent into one of these categories: `MISLEADING_QUERY`, `TEXT_TO_SQL`, `GENERAL`, or `USER_GUIDE`. Additionally, provide a concise reasoning (maximum 20 words) for your classification. @@ -36,6 +51,7 @@ - **Vague Queries:** If the question is vague or does not related to a table or property from the schema, classify it as `MISLEADING_QUERY`. - **Incomplete Queries:** If the question is related to the database schema but references unspecified values (e.g., "the following", "these", "those") without providing them, classify as `GENERAL`. - **Time-related Queries:** Don't rephrase time-related information in the user's question. +- **Schema-resolvable references:** Treat user language as schema-resolvable references when it maps to provided table or column names, aliases, descriptions, or data details, even if the user did not type exact table or column names; do not require the user to write exact schema identifiers. Do not classify a data retrieval or analytics question as MISLEADING only because the user did not write exact table or column names. ### Intent Definitions ### @@ -43,13 +59,13 @@ **When to Use:** - The user's inputs are about modifying SQL from previous questions. - The user's inputs are related to the database schema and requires an SQL query. -- The question (or related previous query) includes references to specific tables, columns, or data details. -- The question includes **complete information** with specific tables, columns, or data values needed for execution. +- The question (or related previous query) includes references to specific tables, columns, data details, or schema-resolvable database concepts. +- The question includes **complete information** with specific or schema-resolvable tables, columns, concepts, or data values needed for execution. - The question provides **all necessary parameters** to generate executable SQL. **Requirements:** - Must have complete filter criteria, specific values, or clear references to previous context. -- Include specific table and column names from the schema in your reasoning or modifying SQL from previous questions. +- Include exact schema identifiers in your reasoning only when they are available; otherwise cite the user's schema-resolvable references. - Reference phrases from the user's inputs that clearly relate to the schema. **Examples:** @@ -175,7 +191,7 @@ async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> d @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, table_retriever: Any + embedding: dict, project_id: str, table_retriever: Any, mdl_hash: str | None = None ) -> dict: filters = { "operator": "AND", @@ -184,10 +200,7 @@ async def table_retrieval( ], } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) return await table_retriever.run( query_embedding=embedding.get("embedding"), @@ -197,7 +210,11 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, embedding: dict, project_id: str, dbschema_retriever: Any + table_retrieval: dict, + project_id: str, + dbschema_retriever: Any, + embedding: dict | None = None, + mdl_hash: str | None = None, ) -> list[Document]: tables = table_retrieval.get("documents", []) table_names = [] @@ -220,10 +237,7 @@ async def dbschema_retrieval( ], } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) results = await dbschema_retriever.run( query_embedding=embedding.get("embedding"), filters=filters @@ -297,18 +311,20 @@ async def classify_intent(prompt: dict, generator: Any, generator_name: str) -> @observe(capture_input=False) -def post_process(classify_intent: dict, construct_db_schemas: list[str]) -> dict: +def post_process( + classify_intent: dict, construct_db_schemas: list[str], query: str +) -> dict: try: results = orjson.loads(classify_intent.get("replies")[0]) return { - "rephrased_question": results["rephrased_question"], + "rephrased_question": results.get("rephrased_question") or query, "intent": results["results"], "reasoning": results["reasoning"], "db_schemas": construct_db_schemas, } except Exception: return { - "rephrased_question": "", + "rephrased_question": query, "intent": "TEXT_TO_SQL", "reasoning": "", "db_schemas": construct_db_schemas, @@ -379,6 +395,7 @@ async def run( self, query: str, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, histories: Optional[list[AskHistory]] = None, sql_samples: Optional[list[dict]] = None, instructions: Optional[list[dict]] = None, @@ -390,6 +407,7 @@ async def run( inputs={ "query": query, "project_id": project_id or "", + "mdl_hash": mdl_hash, "histories": histories or [], "sql_samples": sql_samples or [], "instructions": instructions or [], diff --git a/wren-ai-service/src/pipelines/generation/question_recommendation.py b/wren-ai-service/src/pipelines/generation/question_recommendation.py index a6e7c17b02..cccb1fc2e8 100644 --- a/wren-ai-service/src/pipelines/generation/question_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/question_recommendation.py @@ -18,7 +18,16 @@ system_prompt = """ -You are an expert in data analysis and SQL query generation. Given a data model specification, optionally a user's question, and a list of categories, your task is to generate insightful, specific questions that can be answered using the provided data model. Each question should be accompanied by a brief explanation of its relevance or importance. +You are an expert in data analysis and SQL query generation. Given a data model specification, optionally a user's question, and a list of categories, your task is to generate insightful, specific questions that can be answered using the provided data model. + +### Grounding Rules + +- The DATABASE SCHEMA is the only source for answerable business concepts. +- Generate questions only from tables, views, metrics, columns, measures, dimensions, calculated fields, and relationships that are present in DATABASE SCHEMA. +- Use aliases, descriptions, and comments only to understand meaning. Do not introduce nouns, measures, dimensions, periods, or entities that are not supported by the schema. +- If the same business concept appears in multiple modeled datasets, generate questions that can use each relevant modeled dataset, provided the schema exposes the needed fields. +- If a question would require fields from multiple datasets, generate it only when DATABASE SCHEMA provides either a relationship path, a view, a metric, or compatible fields that can be combined as separate rows. +- Do not use generic analytics examples, common business templates, or prior wording as a source of answerable concepts unless the concept is represented in DATABASE SCHEMA. ### JSON Output Structure @@ -28,8 +37,8 @@ { "questions": [ { - "question": "", - "category": "" + "question": "schema-grounded question text", + "category": "question category" }, ... ] @@ -58,6 +67,7 @@ - Generate questions that are closely related to the user's previous question, ensuring that the new questions build upon or provide deeper insights into the original query. - Use **random category selection** to introduce diverse perspectives while maintaining a focus on the context of the previous question. - Apply the analysis techniques above to enhance the relevance and depth of the generated questions. + - Keep only the parts of the previous question that are supported by DATABASE SCHEMA. 4. **If No User Question is Provided:** @@ -68,58 +78,23 @@ - Ensure questions can be answered using the data model. - Mix simple and complex questions. - Avoid open-ended questions - each should have a definite answer. - - Incorporate time-based analysis where relevant. - - Combine multiple analysis techniques when appropriate for deeper insights. + - Incorporate time-based analysis only when DATABASE SCHEMA exposes relevant time fields. + - Combine multiple analysis techniques when DATABASE SCHEMA supports the required fields and relationships. ### Categories of Questions 1. **Descriptive Questions** Summarize historical data. - - Example: _"What was the total sales volume for each product last quarter?"_ - 2. **Segmentation Questions** Identify meaningful data segments. - - Example: _"Which customer segments contributed most to revenue growth?"_ - 3. **Comparative Questions** Compare data across segments or periods. - - Example: _"How did Product A perform compared to Product B last year?"_ - 4. **Data Quality/Accuracy Questions** Assess data reliability and completeness. - - Example: _"Are there inconsistencies in the sales records for Q1?"_ - ---- - -### Example JSON Output - -```json -{ - "questions": [ - { - "question": "What was the total revenue generated by each region in the last year?", - "category": "Descriptive Questions" - }, - { - "question": "How do customer preferences differ between age groups?", - "category": "Segmentation Questions" - }, - { - "question": "How does the conversion rate vary across different lead sources?", - "category": "Comparative Questions" - }, - { - "question": "What percentage of contacts have incomplete or missing key properties (e.g., email, lifecycle stage, or deal association)", - "category": "Data Quality/Accuracy Questions" - } - ] -} -``` - --- ### Additional Instructions for Randomization @@ -156,7 +131,7 @@ {% endfor %} {% endif %} -Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based on the provided data model. Both the questions and category names should be translated into {{language}}{% if user_question %} and be related to the user's question{% endif %}. The output format should maintain the structure but with localized text. +Please generate {{max_questions}} insightful questions for each of the {{max_categories}} categories based only on the provided data model. Both the questions and category names should be translated into {{language}}{% if user_question %} and be related to the user's question{% endif %}. The output format should maintain the structure but with localized text. """ diff --git a/wren-ai-service/src/pipelines/generation/relationship_recommendation.py b/wren-ai-service/src/pipelines/generation/relationship_recommendation.py index e0d0bed675..b92ec99c54 100644 --- a/wren-ai-service/src/pipelines/generation/relationship_recommendation.py +++ b/wren-ai-service/src/pipelines/generation/relationship_recommendation.py @@ -27,13 +27,16 @@ - **type**: The type of relationship, which can be "MANY_TO_ONE", "ONE_TO_MANY" or "ONE_TO_ONE" only. - **toModel**: The name of the target model. - **toColumn**: The column in the target model that forms the relationship. -- **reason**: The reason for recommending this relationship. +- **reason**: A clear natural language business description of what the relationship means and why it is useful. Important guidelines: 1. Do not recommend relationships within the same model (fromModel and toModel must be different). 2. Only suggest relationships if there is a clear and beneficial reason to do so. 3. If there are no good relationships to recommend or if there are fewer than two models, return an empty list of relationships. 4. Use "MANY_TO_ONE" and "ONE_TO_MANY" instead of "MANY_TO_MANY" relationships. +5. Write the reason for business users. Do not merely repeat raw table names, model names, or column names. +6. Use available model names, display names, column names, descriptions, primary keys, and table context to explain how the entities are related. +7. Prefer descriptions like "Each order belongs to one customer, so revenue can be analyzed by customer." over descriptions like "orders.customer_id references customers.id." Output all relationships in the following JSON structure: @@ -46,7 +49,7 @@ "type": "", "toModel": "", "toColumn": "", - "reason": "" + "reason": "" } ... ] @@ -65,7 +68,7 @@ {{models}} **Please analyze these models and suggest optimizations for their relationships.** -Take into account best practices in database design, opportunities for normalization, indexing strategies, and any additional relationships that could improve data integrity and enhance query performance. +Suggest only relationship definitions between the provided models. Do not suggest indexes, normalization changes, new models, new columns, or relationships that require fields not present in the provided model payload. Use this for the relationship name and reason based on the localization language: {{language}} """ @@ -74,26 +77,45 @@ ## Start of Pipeline @observe(capture_input=False) def cleaned_models(mdl: dict) -> dict: - def remove_display_name(d: dict) -> dict: - if "properties" in d and isinstance(d["properties"], dict): - d["properties"] = d["properties"].copy() - d["properties"].pop("displayName", None) - return d + def properties(payload: dict) -> dict: + value = payload.get("properties") + return value if isinstance(value, dict) else {} + + def compact_properties(payload: dict) -> dict: + source = properties(payload) + return { + key: source.get(key) + for key in ["displayName", "description"] + if source.get(key) + } + + def compact_column(column: dict) -> dict: + compacted = { + "name": column.get("name", ""), + "type": column.get("type", ""), + } + if compacted_properties := compact_properties(column): + compacted["properties"] = compacted_properties + return compacted def column_filter(columns: list[dict]) -> list[dict]: filtered_columns = [] for column in columns: if "relationship" not in column: - # Create a copy of the column to avoid modifying the original - filtered_column = column.copy() - filtered_column = remove_display_name(filtered_column) - filtered_columns.append(filtered_column) + filtered_columns.append(compact_column(column)) return filtered_columns return [ - remove_display_name( - {**model, "columns": column_filter(model.get("columns", []))} - ) + { + "name": model.get("name", ""), + "primaryKey": model.get("primaryKey", ""), + "columns": column_filter(model.get("columns", [])), + **( + {"properties": compacted_properties} + if (compacted_properties := compact_properties(model)) + else {} + ), + } for model in mdl.get("models", []) ] diff --git a/wren-ai-service/src/pipelines/generation/semantics_description.py b/wren-ai-service/src/pipelines/generation/semantics_description.py index acc5fc8594..f26206ea72 100644 --- a/wren-ai-service/src/pipelines/generation/semantics_description.py +++ b/wren-ai-service/src/pipelines/generation/semantics_description.py @@ -166,15 +166,23 @@ def wrapper(text: str) -> str: reply = generate.get("replies")[0] # Expecting only one reply normalized = wrapper(reply) - return {model["name"]: model for model in normalized["models"]} + return {model["name"]: model for model in normalized.get("models", [])} @observe(capture_input=False) def output(normalize: dict, picked_models: list[dict]) -> dict: def _filter(enriched: list[dict], columns: list[dict]) -> list[dict]: - valid_columns = [col["name"] for col in columns] + valid_columns = { + col.get("name") + for col in columns + if isinstance(col, dict) and col.get("name") + } - return [col for col in enriched if col["name"] in valid_columns] + return [ + col + for col in enriched or [] + if isinstance(col, dict) and col.get("name") in valid_columns + ] models = {model["name"]: model for model in picked_models} @@ -206,6 +214,7 @@ class SemanticResult(BaseModel): SEMANTICS_DESCRIPTION_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { diff --git a/wren-ai-service/src/pipelines/generation/sql_answer.py b/wren-ai-service/src/pipelines/generation/sql_answer.py index 81289081b5..e6faf4de21 100644 --- a/wren-ai-service/src/pipelines/generation/sql_answer.py +++ b/wren-ai-service/src/pipelines/generation/sql_answer.py @@ -1,6 +1,7 @@ import asyncio import logging import sys +import time from typing import Any, Optional from hamilton import base @@ -19,23 +20,24 @@ sql_to_answer_system_prompt = """ ### TASK -You are a data analyst that great at answering non-technical user's questions based on the data, sql so that even non technical users can easily understand. -Please answer the user's question in concise and clear manner in Markdown format. +You are a data analyst answering a user's question using only the provided SQL result data. +Answer clearly for a non-technical user in Markdown. ### INSTRUCTIONS -1. Read the user's question and understand the user's intention. -2. Read the sql and understand the data. -3. Make sure the answer is aimed for non-technical users, so don't mention any technical terms such as SQL syntax. -4. Generate a concise and clear answer in string format to answerthe user's question based on the data and sql. -5. If answer is in list format, only list top few examples, and tell users there are more results omitted. -6. Answer must be in the same language user specified. -7. Do not include ```markdown or ``` in the answer. -8. If the user provides a custom instruction, it should be followed strictly and you should use it to change the style of response. +1. Use only the provided Data columns, row records, and rows. Do not invent values, totals, categories, dates, examples, or analysis outputs. +2. Treat row records as the clearest representation of the result because each value is paired with its column name. +3. Do not write code, Python, pseudo-code, SQL, code fences, implementation steps, or phrases such as "running the above code". +4. Do not mention SQL syntax, table names, or database internals unless the user explicitly asks for them. +5. If no rows are provided, say that no matching rows were returned. +6. If rows are detailed records, summarize the visible records directly. If rows are aggregates, answer using the aggregate values. +7. If the answer is a list, keep it concise and use only examples present in the provided data. +8. Answer must be in the same language user specified. +9. If the user provides a custom instruction, follow it strictly for the response style unless it conflicts with these data-grounding rules. ### OUTPUT FORMAT -Please provide your response in proper Markdown stringformat. +Return only the user-facing answer as a Markdown string. """ sql_to_answer_user_prompt_template = """ @@ -44,13 +46,14 @@ SQL: {{ sql }} Data: columns: {{ sql_data.columns }} +row records: {{ sql_data.row_records }} rows: {{ sql_data.data }} Language: {{ language }} Current Time: {{ current_time }} Custom Instruction: {{ custom_instruction }} -Please think step by step and answer the user's question. +Answer directly from the provided row records and rows. """ @@ -158,7 +161,8 @@ async def run( custom_instruction: Optional[str] = None, ) -> dict: logger.info("Sql_Answer Generation pipeline is running...") - return await self._pipe.execute( + started_at = time.perf_counter() + result = await self._pipe.execute( ["generate_answer"], inputs={ "query": query, @@ -171,3 +175,9 @@ async def run( **self._components, }, ) + logger.info( + "Ask timing query_id=%s stage=answer_formatting elapsed_ms=%.1f", + query_id or "", + (time.perf_counter() - started_at) * 1000, + ) + return result diff --git a/wren-ai-service/src/pipelines/generation/sql_correction.py b/wren-ai-service/src/pipelines/generation/sql_correction.py index 973b8c69a7..7da913ac28 100644 --- a/wren-ai-service/src/pipelines/generation/sql_correction.py +++ b/wren-ai-service/src/pipelines/generation/sql_correction.py @@ -1,13 +1,12 @@ import logging import sys -from typing import Any, Dict, List +from typing import Any, Dict from hamilton import base from hamilton.async_driver import AsyncDriver -from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider @@ -16,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, get_text_to_sql_rules, ) from src.pipelines.retrieval.sql_functions import SqlFunction @@ -30,12 +30,20 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) return f""" ### TASK ### -You are an ANSI SQL expert with exceptional logical thinking skills and debugging skills, you need to fix the syntactically incorrect ANSI SQL query. +You are a Wren SQL expert with exceptional logical thinking skills and debugging skills. Regenerate a grounded Wren SQL query from the user's question and the current DATABASE SCHEMA after a previous SQL attempt failed. ### SQL CORRECTION INSTRUCTIONS ### 1. First, think hard about the error message, and figure out the root cause first(please use the DATABASE SCHEMA, SQL FUNCTIONS and USER INSTRUCTIONS to help you figure out the root cause). 2. Then, generate the syntactically correct ANSI SQL query to correct the error. +3. If the failed SQL references a table, view, column, function, alias, or placeholder that is not present in DATABASE SCHEMA or SQL FUNCTIONS, do not preserve it. Regenerate from the USER QUESTION and DATABASE SCHEMA. +4. Treat invalid object name, dataset not found, table not found, invalid column name, and invalid identifier errors as schema-grounding failures. Use exact declared identifiers from DATABASE SCHEMA only. +5. If the error reports an unknown table or field, replace it only with an exact executable identifier declared in DATABASE SCHEMA or SQL FUNCTIONS. Do not retry the same unknown identifier. +6. Do not create dummy CTEs, placeholder tables, table-existence checks, or generic replacement names to make the query executable. If the requested intent is supported by retrieved schema objects, use those exact objects; otherwise return null for sql. +7. For grouped queries, repair SQL Server errors about ORDER BY columns not appearing in GROUP BY by ordering with selected grouping columns or selected aggregate aliases, or by adding the exact ordering key to both SELECT and GROUP BY when that key is declared in DATABASE SCHEMA. +8. Do not preserve generic log, file, JSON, payload, text, or app-metric scans when DATABASE SCHEMA contains exact modeled business columns for the user's requested entity, measure, status, date, or dimension. +9. If the failed SQL invented component fields for a metric that exists directly in DATABASE SCHEMA, replace the calculation with the exact declared metric column. +10. Do not route a question to a different business domain because of generic keyword overlap. Use only retrieved schema metadata that directly represents the requested entities, measures, filters, dates, and dimensions. ### SQL RULES ### Make sure you follow the SQL Rules strictly. @@ -43,16 +51,21 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and answers the user's requested intent. Do not create table or column identifiers from the user's wording. If no fully grounded SQL can be generated, return null for sql. {{ - "sql": + "sql": "corrected SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ sql_correction_user_prompt_template = """ {% if documents %} +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -74,29 +87,59 @@ def get_sql_correction_system_prompt(sql_knowledge: SqlKnowledge | None = None) {% endif %} ### QUESTION ### -SQL: {{ invalid_generation_result.sql }} -Error Message: {{ invalid_generation_result.error }} +{% if query %} +User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from the WREN SQL IDENTIFIER CONTRACT and DATABASE SCHEMA. Treat source metadata, physical names, lineage names, semantic labels, diagnostic text, and user question words as non-executable background unless the exact same identifier is declared in the contract. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, failed SQL, dry-run diagnostic, or reasoning plan. +{% endif %} +### FAILED SQL ### +The failed SQL below is diagnostic context only. It is not an executable schema source. +Only preserve an identifier, function, literal filter, grouping, ordering, or join from this SQL when it is also declared exactly in the WREN SQL IDENTIFIER CONTRACT, DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. +If it contains placeholders, assumed business names, connector-specific syntax, source/physical names, or unsupported objects, discard those parts and regenerate from the QUESTION plus DATABASE SCHEMA. + +{% if invalid_generation_result and invalid_generation_result.sql %} +{{ invalid_generation_result.sql }} +{% else %} +No failed SQL was provided. +{% endif %} + +### DRY-RUN DIAGNOSTIC ### +The diagnostic below explains why the previous SQL failed. Use it to understand the failure only. +Do not copy identifiers, source names, physical names, SQL fragments, or replacement candidates from the diagnostic unless they appear exactly in the WREN SQL IDENTIFIER CONTRACT, DATABASE SCHEMA, SQL FUNCTIONS, or USER INSTRUCTIONS. + +{% if invalid_generation_result and invalid_generation_result.error %} +{{ invalid_generation_result.error }} +{% else %} +No dry-run diagnostic was provided. +{% endif %} + +Regenerate from the user question, current DATABASE SCHEMA, and the diagnostic failure. Keep DATABASE SCHEMA as the only executable identifier source. -Let's think step by step. +Return only the final JSON SQL response. """ ## Start of Pipeline @observe(capture_input=False) def prompt( - documents: List[Document], + documents: list[str], invalid_generation_result: Dict, prompt_builder: PromptBuilder, + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, ) -> dict: _prompt = prompt_builder.run( + query=query, documents=documents, invalid_generation_result=invalid_generation_result, + sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -120,13 +163,20 @@ async def post_process( generate_sql_correction: dict, post_processor: SQLGenPostProcessor, data_source: str, + documents: list[str] | None = None, + query: str | None = None, project_id: str | None = None, + mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, ) -> dict: return await post_processor.run( generate_sql_correction.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, + contexts=documents, + fallback_query=query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -167,19 +217,25 @@ def __init__( @observe(name="SQL Correction") async def run( self, - contexts: List[Document], + contexts: list[str], invalid_generation_result: Dict[str, str], + query: str | None = None, + sql_generation_reasoning: str | None = None, instructions: list[dict] | None = None, sql_functions: list[SqlFunction] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, sql_knowledge: SqlKnowledge | None = None, + validation_contexts: list[str] | None = None, ): logger.info("SQLCorrection pipeline is running...") - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) + if project_id or use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) else: metadata = {} @@ -187,14 +243,18 @@ async def run( ["post_process"], inputs={ "invalid_generation_result": invalid_generation_result, + "query": query, "documents": contexts, + "sql_generation_reasoning": sql_generation_reasoning, "instructions": instructions, "sql_functions": sql_functions, "project_id": project_id, + "mdl_hash": mdl_hash, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, "data_source": metadata.get("data_source", "local_file"), "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py index 3f22b9d512..5ca4c552a9 100644 --- a/wren-ai-service/src/pipelines/generation/sql_diagnosis.py +++ b/wren-ai-service/src/pipelines/generation/sql_diagnosis.py @@ -142,6 +142,8 @@ async def run( invalid_sql: str, error_message: str, language: str, + project_id: str | None = None, + mdl_hash: str | None = None, ): logger.info("SQLDiagnosis pipeline is running...") @@ -153,6 +155,8 @@ async def run( "invalid_sql": invalid_sql, "error_message": error_message, "language": language, + "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_generation.py b/wren-ai-service/src/pipelines/generation/sql_generation.py index 1ee4952b3e..cb1bd15646 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation.py @@ -5,8 +5,8 @@ from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, LLMProvider @@ -15,10 +15,13 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, + generate_simple_analytics_sql, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, get_sql_generation_system_prompt, + unsupported_schema_generation_result, ) from src.pipelines.retrieval.sql_functions import SqlFunction from src.pipelines.retrieval.sql_knowledge import SqlKnowledge @@ -28,6 +31,11 @@ sql_generation_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -54,11 +62,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -71,13 +78,10 @@ ### QUESTION ### User's Question: {{ query }} +Answer the user's intent using the current DATABASE SCHEMA. Use comments, aliases, descriptions, calculated fields, metrics, and relationships only to understand meaning; the SQL must use exact declared table and column names from the WREN SQL IDENTIFIER CONTRACT and DATABASE SCHEMA. Treat source metadata, physical names, lineage names, semantic labels, and user question words as non-executable background unless the exact same identifier is declared in the contract. If a needed table, output column, filter column, grouping column, relation, date field, measure, or function is not declared in DATABASE SCHEMA or SQL FUNCTIONS, return null for sql instead of inventing, substituting, or approximating a similar name. If the retrieved schema does not ground the user's primary requested intent, return null for sql instead of querying an unrelated object. +If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, stop and return null for sql. Never create a table or column from the user's wording, even when the wording looks like a business term or object name. -{% if sql_generation_reasoning %} -### REASONING PLAN ### -{{ sql_generation_reasoning }} -{% endif %} - -Let's think step by step. +Return only the final JSON SQL response. """ @@ -99,7 +103,6 @@ def prompt( _prompt = prompt_builder.run( query=query, documents=documents, - sql_generation_reasoning=sql_generation_reasoning, instructions=construct_instructions( instructions=instructions, ), @@ -116,6 +119,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -139,7 +143,12 @@ async def post_process( generate_sql: dict, post_processor: SQLGenPostProcessor, data_source: str, + query: str | None = None, + grounding_query: str | None = None, + documents: list[str] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, @@ -147,6 +156,9 @@ async def post_process( return await post_processor.run( generate_sql.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, + contexts=validation_contexts or documents, + fallback_query=grounding_query or query, use_dry_plan=use_dry_plan, data_source=data_source, allow_dry_plan_fallback=allow_dry_plan_fallback, @@ -194,6 +206,7 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, @@ -202,32 +215,117 @@ async def run( allow_dry_plan_fallback: bool = True, allow_data_preview: bool = False, sql_knowledge: SqlKnowledge | None = None, + grounding_query: str | None = None, + validation_contexts: list[str] | None = None, ): - logger.info("SQL Generation pipeline is running...") + logger.info( + "SQL Generation pipeline is running for project_id=%s mdl_hash=%s", + project_id or "", + mdl_hash or "", + ) - if use_dry_plan: - metadata = await retrieve_metadata(project_id or "", self._retriever) + if project_id or use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) else: metadata = {} + data_source = metadata.get("data_source", "local_file") + + effective_grounding_query = grounding_query or query + unsupported_result = unsupported_schema_generation_result( + effective_grounding_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "SQL generation skipped before LLM because selected schema does not cover requested concepts: %s", + unsupported_result["invalid_generation_result"]["error"], + ) + return {"post_process": unsupported_result} return await self._pipe.execute( ["post_process"], inputs={ "query": query, + "grounding_query": effective_grounding_query, "documents": contexts, "sql_generation_reasoning": sql_generation_reasoning, "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, + "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, "sql_functions": sql_functions, "use_dry_plan": use_dry_plan, "allow_dry_plan_fallback": allow_dry_plan_fallback, - "data_source": metadata.get("data_source", "local_file"), + "data_source": data_source, "allow_data_preview": allow_data_preview, "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) + + async def run_deterministic_fast_path( + self, + query: str, + contexts: list[str], + project_id: str | None = None, + mdl_hash: str | None = None, + use_dry_plan: bool = False, + allow_dry_plan_fallback: bool = True, + allow_data_preview: bool = False, + grounding_query: str | None = None, + ) -> dict | None: + if use_dry_plan: + metadata = await retrieve_metadata( + project_id or "", self._retriever, mdl_hash + ) + else: + metadata = {} + data_source = metadata.get("data_source", "local_file") + effective_grounding_query = grounding_query or query + + unsupported_result = unsupported_schema_generation_result( + effective_grounding_query, + contexts=contexts, + data_source=data_source, + ) + if unsupported_result: + logger.info( + "SQL generation deterministic fast path returned unsupported schema before LLM." + ) + return {"post_process": unsupported_result, "fast_path": "unsupported"} + + deterministic_sql = generate_simple_analytics_sql( + effective_grounding_query, + contexts, + ) + if not deterministic_sql: + return None + + logger.info("SQL generation deterministic fast path produced a candidate.") + post_process = await self._components["post_processor"].run( + [deterministic_sql], + project_id=project_id, + mdl_hash=mdl_hash, + contexts=contexts, + fallback_query=effective_grounding_query, + use_dry_plan=use_dry_plan, + data_source=data_source, + allow_dry_plan_fallback=allow_dry_plan_fallback, + allow_data_preview=allow_data_preview, + ) + if post_process.get("valid_generation_result"): + logger.info("SQL generation deterministic fast path accepted candidate.") + return {"post_process": post_process, "fast_path": "deterministic"} + + logger.info( + "SQL generation deterministic fast path rejected candidate; continuing to LLM. reason=%s", + post_process.get("invalid_generation_result", {}).get("error"), + ) + return None diff --git a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py index 00b731cb2c..2178880caa 100644 --- a/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py +++ b/wren-ai-service/src/pipelines/generation/sql_generation_reasoning.py @@ -6,13 +6,15 @@ from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider from src.pipelines.common import clean_up_new_lines from src.pipelines.generation.utils.sql import ( construct_instructions, + construct_schema_identifier_catalog, + sanitize_sql_generation_reasoning, sql_generation_reasoning_system_prompt, ) from src.utils import trace_cost @@ -22,6 +24,13 @@ sql_generation_reasoning_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +Use exact model, column, and relationship names from that contract whenever reasoning names schema objects. +Treat source metadata, physical names, lineage names, semantic labels, and user wording as non-executable background unless the exact same identifier is declared in the contract. +If the contract does not contain a table, column, or relationship required by the user's intent, say the current schema context is insufficient instead of naming an assumed object. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -29,11 +38,10 @@ {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sql_sample in sql_samples %} Question: {{sql_sample.question}} -SQL: -{{sql_sample.sql}} {% endfor %} {% endif %} @@ -61,15 +69,18 @@ def prompt( sql_samples: list[dict], instructions: list[dict], prompt_builder: PromptBuilder, + validation_contexts: Optional[list[str]] = None, configuration: Configuration | None = Configuration(), ) -> dict: + schema_documents = validation_contexts or documents _prompt = prompt_builder.run( query=query, - documents=documents, + documents=schema_documents, sql_samples=sql_samples, instructions=construct_instructions( instructions=instructions, ), + schema_identifier_catalog=construct_schema_identifier_catalog(schema_documents), language=configuration.language, current_time=configuration.show_current_time(), ) @@ -90,7 +101,7 @@ async def generate_sql_reasoning( def post_process( generate_sql_reasoning: dict, ) -> dict: - return generate_sql_reasoning.get("replies")[0] + return sanitize_sql_generation_reasoning(generate_sql_reasoning.get("replies")[0]) ## End of Pipeline @@ -158,6 +169,9 @@ async def run( contexts: list[str], sql_samples: Optional[list[dict]] = None, instructions: Optional[list[str]] = None, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + validation_contexts: Optional[list[str]] = None, configuration: Configuration = Configuration(), query_id: Optional[str] = None, ): @@ -169,6 +183,9 @@ async def run( "documents": contexts, "sql_samples": sql_samples or [], "instructions": instructions or [], + "project_id": project_id, + "mdl_hash": mdl_hash, + "validation_contexts": validation_contexts, "configuration": configuration, "query_id": query_id, **self._components, diff --git a/wren-ai-service/src/pipelines/generation/sql_regeneration.py b/wren-ai-service/src/pipelines/generation/sql_regeneration.py index 4b7284aa26..604d63f40e 100644 --- a/wren-ai-service/src/pipelines/generation/sql_regeneration.py +++ b/wren-ai-service/src/pipelines/generation/sql_regeneration.py @@ -5,8 +5,8 @@ from hamilton import base from hamilton.async_driver import AsyncDriver from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe +from langfuse.decorators import observe from src.core.engine import Engine from src.core.pipeline import BasicPipeline from src.core.provider import LLMProvider @@ -15,6 +15,7 @@ SQL_GENERATION_MODEL_KWARGS, SQLGenPostProcessor, construct_instructions, + construct_schema_identifier_catalog, get_calculated_field_instructions, get_json_field_instructions, get_metric_instructions, @@ -34,23 +35,26 @@ def get_sql_regeneration_system_prompt( return f""" ### TASK ### -You are a great ANSI SQL expert. Now you are given database schema, SQL generation reasoning and an original SQL query, -please carefully review the reasoning, and then generate a new SQL query that matches the reasoning. -While generating the new SQL query, you should use the original SQL query as a reference. -While generating the new SQL query, make sure to use the database schema to generate the SQL query. +You are a Wren SQL expert. Generate a grounded Wren SQL query from the current DATABASE SCHEMA and the requested adjustment intent. +The DATABASE SCHEMA is the only source of executable table and column identifiers. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be a JSON object. Return null for sql if the requested adjustment cannot be grounded in DATABASE SCHEMA. {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ sql_regeneration_user_prompt_template = """ +{{ schema_identifier_catalog }} + +The WREN SQL IDENTIFIER CONTRACT above is the authoritative executable schema. +The DATABASE SCHEMA below provides type, semantic, and relationship details for those exact identifiers. + ### DATABASE SCHEMA ### {% for document in documents %} {{ document }} @@ -77,11 +81,10 @@ def get_sql_regeneration_system_prompt( {% if sql_samples %} ### SQL SAMPLES ### +These samples are examples of intent and style only. Their SQL bodies are intentionally omitted so they cannot provide executable identifiers, literal values, placeholders, functions, or SQL patterns. {% for sample in sql_samples %} Question: {{sample.question}} -SQL: -{{sample.sql}} {% endfor %} {% endif %} @@ -93,10 +96,10 @@ def get_sql_regeneration_system_prompt( {% endif %} ### QUESTION ### -SQL generation reasoning: {{ sql_generation_reasoning }} -Original SQL query: {{ sql }} +Adjustment intent: {{ sql_generation_reasoning }} +The previous SQL is intentionally omitted so it cannot provide executable identifiers, literal values, placeholders, functions, SQL patterns, or unsupported object names. -Let's think step by step. +Regenerate from the adjustment intent and current DATABASE SCHEMA only. Return only the final JSON SQL response. """ @@ -135,6 +138,7 @@ def prompt( ), sql_samples=sql_samples, sql_functions=sql_functions, + schema_identifier_catalog=construct_schema_identifier_catalog(documents), ) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} @@ -157,11 +161,16 @@ async def regenerate_sql( async def post_process( regenerate_sql: dict, post_processor: SQLGenPostProcessor, + documents: list[str], project_id: str | None = None, + mdl_hash: str | None = None, + validation_contexts: list[str] | None = None, ) -> dict: return await post_processor.run( regenerate_sql.get("replies"), project_id=project_id, + mdl_hash=mdl_hash, + contexts=validation_contexts or documents, ) @@ -200,11 +209,13 @@ async def run( sql_samples: list[dict] | None = None, instructions: list[dict] | None = None, project_id: str | None = None, + mdl_hash: str | None = None, has_calculated_field: bool = False, has_metric: bool = False, has_json_field: bool = False, sql_functions: list[SqlFunction] | None = None, sql_knowledge: SqlKnowledge | None = None, + validation_contexts: list[str] | None = None, ): logger.info("SQL Regeneration pipeline is running...") @@ -217,11 +228,13 @@ async def run( "sql_samples": sql_samples, "instructions": instructions, "project_id": project_id, + "mdl_hash": mdl_hash, "has_calculated_field": has_calculated_field, "has_metric": has_metric, "has_json_field": has_json_field, "sql_functions": sql_functions, "sql_knowledge": sql_knowledge, + "validation_contexts": validation_contexts, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/generation/sql_tables_extraction.py b/wren-ai-service/src/pipelines/generation/sql_tables_extraction.py index 7060e750b3..418c557ec7 100644 --- a/wren-ai-service/src/pipelines/generation/sql_tables_extraction.py +++ b/wren-ai-service/src/pipelines/generation/sql_tables_extraction.py @@ -75,7 +75,16 @@ async def extract_sql_tables(prompt: dict, generator: Any, generator_name: str) async def post_process( extract_sql_tables: dict, ) -> list[str]: - return orjson.loads(extract_sql_tables.get("replies")[0])["tables"] + try: + tables = orjson.loads(extract_sql_tables.get("replies")[0]).get("tables", []) + except Exception as e: + logger.warning("Failed to parse extracted SQL tables: %s", e) + return [] + + if not isinstance(tables, list): + return [] + + return [table for table in tables if isinstance(table, str)] ## End of Pipeline diff --git a/wren-ai-service/src/pipelines/generation/utils/chart.py b/wren-ai-service/src/pipelines/generation/utils/chart.py index 5d06b949a9..a725609a7e 100644 --- a/wren-ai-service/src/pipelines/generation/utils/chart.py +++ b/wren-ai-service/src/pipelines/generation/utils/chart.py @@ -1,4 +1,5 @@ import logging +import re from typing import Any, Dict, Literal, Optional import orjson @@ -10,6 +11,16 @@ logger = logging.getLogger("wren-ai-service") +_TEMPORAL_NAME_PATTERN = re.compile( + r"(date|time|timestamp|month|year|week|day|period)", re.I +) +_COUNT_QUERY_PATTERN = re.compile(r"\b(count|counts|number|volume|how many)\b", re.I) +_TREND_QUERY_PATTERN = re.compile( + r"\b(trend|over time|monthly|weekly|daily|yearly|by month|by week|by year)\b", + re.I, +) +_GROUPED_QUERY_PATTERN = re.compile(r"\b(each|across|within|by .+ by|per .+ per)\b", re.I) +_TOP_QUERY_PATTERN = re.compile(r"\b(top|highest|lowest|most|least|rank)\b", re.I) chart_generation_instructions = """ ### INSTRUCTIONS ### @@ -44,7 +55,7 @@ 1. Understanding Your Data Types - Nominal (Categorical): Names or labels without a specific order (e.g., types of fruits, countries). - Ordinal: Categorical data with a meaningful order but no fixed intervals (e.g., rankings, satisfaction levels). -- Quantitative: Numerical values representing counts or measurements (e.g., sales figures, temperatures). +- Quantitative: Numerical values representing counts or measurements. - Temporal: Date or time data (e.g., timestamps, dates). 2. Chart Types and When to Use Them - Bar Chart @@ -52,19 +63,19 @@ - Data Requirements: - One categorical variable (x-axis). - One quantitative variable (y-axis). - - Example: Comparing sales numbers for different product categories. + - Example: Comparing metric values for different categories. - Grouped Bar Chart - Use When: Comparing sub-categories within main categories. - Data Requirements: - Two categorical variables (x-axis grouped by one, color-coded by another). - One quantitative variable (y-axis). - - Example: Sales numbers for different products across various regions. + - Example: Comparing metric values for subcategories across categories. - Line Chart - Use When: Displaying trends over continuous data, especially time. - Data Requirements: - One temporal or ordinal variable (x-axis). - One quantitative variable (y-axis). - - Example: Tracking monthly revenue over a year. + - Example: Tracking a monthly metric over a year. - Multi Line Chart - Use When: Displaying trends over continuous data, especially time. - Data Requirements: @@ -88,7 +99,7 @@ - Stacked Bar Chart - Use When: Showing composition and comparison across categories. - Data Requirements: Same as grouped bar chart. - - Example: Sales by region and product type. + - Example: Metric by category and subcategory. - Guidelines for Selecting Chart Types - Comparing Categories: - Bar Chart: Best for simple comparisons across categories. @@ -106,28 +117,28 @@ 1. Bar Chart - Sample Data: [ - {"Region": "North", "Sales": 100}, - {"Region": "South", "Sales": 200}, - {"Region": "East", "Sales": 300}, - {"Region": "West", "Sales": 400} + {"Category": "A", "Metric": 100}, + {"Category": "B", "Metric": 200}, + {"Category": "C", "Metric": 300}, + {"Category": "D", "Metric": 400} ] - Chart Schema: { "title": , "mark": {"type": "bar"}, "encoding": { - "x": {"field": "Region", "type": "nominal", "title": }, - "y": {"field": "Sales", "type": "quantitative", "title": }, - "color": {"field": "Region", "type": "nominal", "title": ""} + "x": {"field": "Category", "type": "nominal", "title": }, + "y": {"field": "Metric", "type": "quantitative", "title": }, + "color": {"field": "Category", "type": "nominal", "title": ""} } } 2. Line Chart - Sample Data: [ - {"Date": "2022-01-01", "Sales": 100}, - {"Date": "2022-01-02", "Sales": 200}, - {"Date": "2022-01-03", "Sales": 300}, - {"Date": "2022-01-04", "Sales": 400} + {"Date": "2022-01-01", "Metric": 100}, + {"Date": "2022-01-02", "Metric": 200}, + {"Date": "2022-01-03", "Metric": 300}, + {"Date": "2022-01-04", "Metric": 400} ] - Chart Schema: { @@ -135,7 +146,7 @@ "mark": {"type": "line"}, "encoding": { "x": {"field": "Date", "type": "temporal", "title": }, - "y": {"field": "Sales", "type": "quantitative", "title": } + "y": {"field": "Metric", "type": "quantitative", "title": } } } 3. Pie Chart @@ -158,10 +169,10 @@ 4. Area Chart - Sample Data: [ - {"Date": "2022-01-01", "Sales": 100}, - {"Date": "2022-01-02", "Sales": 200}, - {"Date": "2022-01-03", "Sales": 300}, - {"Date": "2022-01-04", "Sales": 400} + {"Date": "2022-01-01", "Metric": 100}, + {"Date": "2022-01-02", "Metric": 200}, + {"Date": "2022-01-03", "Metric": 300}, + {"Date": "2022-01-04", "Metric": 400} ] - Chart Schema: { @@ -169,52 +180,52 @@ "mark": {"type": "area"}, "encoding": { "x": {"field": "Date", "type": "temporal", "title": ""}, - "y": {"field": "Sales", "type": "quantitative", "title": ""} + "y": {"field": "Metric", "type": "quantitative", "title": ""} } } 5. Stacked Bar Chart - Sample Data: [ - {"Region": "North", "Product": "A", "Sales": 100}, - {"Region": "North", "Product": "B", "Sales": 150}, - {"Region": "South", "Product": "A", "Sales": 200}, - {"Region": "South", "Product": "B", "Sales": 250}, - {"Region": "East", "Product": "A", "Sales": 300}, - {"Region": "East", "Product": "B", "Sales": 350}, - {"Region": "West", "Product": "A", "Sales": 400}, - {"Region": "West", "Product": "B", "Sales": 450} + {"Category": "A", "Subcategory": "One", "Metric": 100}, + {"Category": "A", "Subcategory": "Two", "Metric": 150}, + {"Category": "B", "Subcategory": "One", "Metric": 200}, + {"Category": "B", "Subcategory": "Two", "Metric": 250}, + {"Category": "C", "Subcategory": "One", "Metric": 300}, + {"Category": "C", "Subcategory": "Two", "Metric": 350}, + {"Category": "D", "Subcategory": "One", "Metric": 400}, + {"Category": "D", "Subcategory": "Two", "Metric": 450} ] - Chart Schema: { "title": "", "mark": {"type": "bar"}, "encoding": { - "x": {"field": "Region", "type": "nominal", "title": ""}, - "y": {"field": "Sales", "type": "quantitative", "title": "", "stack": "zero"}, - "color": {"field": "Product", "type": "nominal", "title": ""} + "x": {"field": "Category", "type": "nominal", "title": ""}, + "y": {"field": "Metric", "type": "quantitative", "title": "", "stack": "zero"}, + "color": {"field": "Subcategory", "type": "nominal", "title": ""} } } 6. Grouped Bar Chart - Sample Data: [ - {"Region": "North", "Product": "A", "Sales": 100}, - {"Region": "North", "Product": "B", "Sales": 150}, - {"Region": "South", "Product": "A", "Sales": 200}, - {"Region": "South", "Product": "B", "Sales": 250}, - {"Region": "East", "Product": "A", "Sales": 300}, - {"Region": "East", "Product": "B", "Sales": 350}, - {"Region": "West", "Product": "A", "Sales": 400}, - {"Region": "West", "Product": "B", "Sales": 450} + {"Category": "A", "Subcategory": "One", "Metric": 100}, + {"Category": "A", "Subcategory": "Two", "Metric": 150}, + {"Category": "B", "Subcategory": "One", "Metric": 200}, + {"Category": "B", "Subcategory": "Two", "Metric": 250}, + {"Category": "C", "Subcategory": "One", "Metric": 300}, + {"Category": "C", "Subcategory": "Two", "Metric": 350}, + {"Category": "D", "Subcategory": "One", "Metric": 400}, + {"Category": "D", "Subcategory": "Two", "Metric": 450} ] - Chart Schema: { "title": "", "mark": {"type": "bar"}, "encoding": { - "x": {"field": "Region", "type": "nominal", "title": ""}, - "y": {"field": "Sales", "type": "quantitative", "title": ""}, - "xOffset": {"field": "Product", "type": "nominal", "title": ""}, - "color": {"field": "Product", "type": "nominal", "title": ""} + "x": {"field": "Category", "type": "nominal", "title": ""}, + "y": {"field": "Metric", "type": "quantitative", "title": ""}, + "xOffset": {"field": "Subcategory", "type": "nominal", "title": ""}, + "color": {"field": "Subcategory", "type": "nominal", "title": ""} } } 7. Multi Line Chart @@ -267,10 +278,7 @@ def run( col: list(df[col].unique())[:sample_column_size] for col in df.columns } - if len(df) > sample_data_count: - sample_data = df.sample(n=sample_data_count).to_dict(orient="records") - else: - sample_data = df.to_dict(orient="records") + sample_data = df.head(sample_data_count).to_dict(orient="records") return { "sample_data": sample_data, @@ -278,6 +286,324 @@ def run( } +def build_fallback_chart_result(query: str, data: Dict[str, Any]) -> Dict[str, Any]: + columns = _column_names(data) + rows = data.get("data", []) or [] + records = _row_records(columns, rows) + if not columns or not records: + return _empty_chart_result("No chart was generated because the result is empty.") + + numeric_columns = [ + column for column in columns if _is_numeric_column(column, records) + ] + temporal_columns = [ + column for column in columns if _is_temporal_column(column, records) + ] + categorical_columns = [ + column + for column in columns + if column not in numeric_columns and column not in temporal_columns + ] + + if temporal_columns and not numeric_columns and _TREND_QUERY_PATTERN.search( + query or "" + ): + x_field = _best_named_column(query, temporal_columns) or temporal_columns[0] + return _chart_result( + chart_type="line", + chart_schema={ + "title": _title(query), + "mark": {"type": "line"}, + "encoding": { + "x": _encoding( + x_field, + "temporal", + title=x_field, + timeUnit=_time_unit(query), + ), + "y": _count_encoding(query), + }, + }, + ) + + if not numeric_columns: + count_chart_columns = categorical_columns or temporal_columns + if not count_chart_columns: + return _empty_chart_result( + "No chart was generated because the result is not suitable for visualization." + ) + + x_field = _best_named_column(query, count_chart_columns) or count_chart_columns[0] + secondary_categories = [ + column for column in categorical_columns if column != x_field + ] + if len(secondary_categories) >= 1 and ( + _GROUPED_QUERY_PATTERN.search(query or "") or len(categorical_columns) >= 2 + ): + group_field = ( + _group_field_from_each_clause(query, secondary_categories) + or secondary_categories[0] + ) + return _chart_result( + chart_type="grouped_bar", + chart_schema={ + "title": _title(query), + "mark": {"type": "bar"}, + "encoding": { + "x": _encoding(x_field, "nominal", title=x_field), + "y": _count_encoding(query), + "xOffset": _encoding(group_field, "nominal", title=group_field), + "color": _encoding(group_field, "nominal", title=group_field), + }, + }, + ) + + return _chart_result( + chart_type="bar", + chart_schema={ + "title": _title(query), + "mark": {"type": "bar"}, + "encoding": { + "x": _encoding(x_field, "nominal", title=x_field), + "y": _count_encoding(query), + "color": _encoding(x_field, "nominal", title=x_field), + }, + }, + ) + + y_field = _best_numeric_column(query, numeric_columns) + + if temporal_columns and (_TREND_QUERY_PATTERN.search(query or "") or len(columns) <= 3): + x_field = _best_named_column(query, temporal_columns) or temporal_columns[0] + chart_type = "line" + return _chart_result( + chart_type=chart_type, + chart_schema={ + "title": _title(query), + "mark": {"type": "line"}, + "encoding": { + "x": _encoding( + x_field, + "temporal", + title=x_field, + timeUnit=_time_unit(query), + ), + "y": _encoding(y_field, "quantitative", title=y_field), + }, + }, + ) + + category_candidates = categorical_columns or [ + column for column in columns if column != y_field + ] + if not category_candidates: + return {} + + x_field = _best_named_column(query, category_candidates) or category_candidates[0] + + secondary_categories = [ + column for column in categorical_columns if column != x_field + ] + should_group = ( + len(secondary_categories) >= 1 + and (_GROUPED_QUERY_PATTERN.search(query or "") or len(categorical_columns) >= 2) + ) + if should_group: + group_field = ( + _group_field_from_each_clause(query, secondary_categories) + or secondary_categories[0] + ) + return _chart_result( + chart_type="grouped_bar", + chart_schema={ + "title": _title(query), + "mark": {"type": "bar"}, + "encoding": { + "x": _encoding(x_field, "nominal", title=x_field), + "y": _encoding(y_field, "quantitative", title=y_field), + "xOffset": _encoding(group_field, "nominal", title=group_field), + "color": _encoding(group_field, "nominal", title=group_field), + }, + }, + ) + + return _chart_result( + chart_type="bar", + chart_schema={ + "title": _title(query), + "mark": {"type": "bar"}, + "encoding": { + "x": _encoding(x_field, "nominal", title=x_field), + "y": _encoding(y_field, "quantitative", title=y_field), + "color": _encoding(x_field, "nominal", title=x_field), + }, + }, + ) + + +def _is_schema_compatible_with_sample_data( + chart_schema: Dict[str, Any], sample_data: list[dict] +) -> bool: + if not chart_schema or not sample_data: + return False + + fields = _schema_fields(chart_schema) + if not fields: + return False + + sample_fields = set(sample_data[0].keys()) + if not fields.issubset(sample_fields): + return False + + for encoding in (chart_schema.get("encoding") or {}).values(): + if not isinstance(encoding, dict): + continue + if encoding.get("aggregate") == "count" and not encoding.get("field"): + return False + + return True + + +def _column_names(data: Dict[str, Any]) -> list[str]: + return [ + column.get("name", "") if isinstance(column, dict) else str(column) + for column in data.get("columns", []) or [] + if column + ] + + +def _row_records(columns: list[str], rows: list[Any]) -> list[dict]: + records = [] + for row in rows: + if isinstance(row, dict): + records.append({column: row.get(column) for column in columns}) + continue + if not isinstance(row, (list, tuple)): + row = [row] + records.append( + {column: row[index] if index < len(row) else None for index, column in enumerate(columns)} + ) + return records + + +def _is_numeric_column(column: str, records: list[dict]) -> bool: + values = [record.get(column) for record in records if record.get(column) is not None] + return bool(values) and all(isinstance(value, (int, float)) for value in values) + + +def _is_temporal_column(column: str, records: list[dict]) -> bool: + if _TEMPORAL_NAME_PATTERN.search(column): + return True + values = [str(record.get(column)) for record in records[:5] if record.get(column)] + return bool(values) and all( + re.match(r"^\d{4}(-\d{2}){0,2}(?:[ T]\d{2}:\d{2}:\d{2})?", value) + for value in values + ) + + +def _best_named_column(query: str, columns: list[str]) -> str | None: + query_terms = _terms(query) + scored = [] + for index, column in enumerate(columns): + score = len(query_terms & _terms(column)) + mention_position = _column_mention_position(query, column) + scored.append((score, -mention_position, -index, column)) + scored.sort(reverse=True) + return scored[0][3] if scored else None + + +def _best_numeric_column(query: str, columns: list[str]) -> str: + named = _best_named_column(query, columns) + if named and len(_terms(named) & _terms(query)) > 0: + return named + return columns[0] + + +def _terms(value: str) -> set[str]: + return { + term.lower().rstrip("s") + for term in re.findall(r"[A-Za-z][A-Za-z0-9]*", value or "") + if len(term) >= 3 + } + + +def _encoding(field: str, field_type: str, title: str, **extra) -> dict: + return {"field": field, "type": field_type, "title": title, **extra} + + +def _time_unit(query: str) -> str: + normalized = (query or "").lower() + if "year" in normalized: + return "year" + if "week" in normalized or "day" in normalized or "daily" in normalized: + return "yearmonthdate" + return "yearmonth" + + +def _count_title(query: str) -> str: + volume_match = re.search(r"\b([A-Za-z][A-Za-z0-9]*)\s+volume\b", query or "", re.I) + if volume_match: + return f"{volume_match.group(1).title()} Count" + return "Count" + + +def _count_encoding(query: str) -> dict: + return { + "aggregate": "count", + "type": "quantitative", + "title": _count_title(query), + } + + +def _title(query: str) -> str: + normalized = " ".join((query or "Chart").strip().split()) + return normalized[:120].rstrip(".") + + +def _chart_result(chart_type: str, chart_schema: Dict[str, Any]) -> Dict[str, Any]: + return { + "reasoning": "Generated from the SQL result columns and requested chart type.", + "chart_type": chart_type, + "chart_schema": chart_schema, + } + + +def _empty_chart_result(reasoning: str) -> Dict[str, Any]: + return {"reasoning": reasoning, "chart_type": "", "chart_schema": {}} + + +def _schema_fields(chart_schema: Dict[str, Any]) -> set[str]: + fields = set() + for encoding in (chart_schema.get("encoding") or {}).values(): + if isinstance(encoding, dict) and encoding.get("field"): + fields.add(encoding["field"]) + for transform in chart_schema.get("transform") or []: + if isinstance(transform, dict): + fields.update(transform.get("fold") or []) + return fields + + +def _column_mention_position(query: str, column: str) -> int: + query_terms = list(re.findall(r"[A-Za-z][A-Za-z0-9]*", query or "")) + column_terms = _terms(column) + for index, term in enumerate(query_terms): + if term.lower().rstrip("s") in column_terms: + return index + return 10_000 + + +def _group_field_from_each_clause(query: str, columns: list[str]) -> str | None: + match = re.search(r"\beach\s+([A-Za-z][A-Za-z0-9]*)", query or "", re.I) + if not match: + return None + + each_term = match.group(1).lower().rstrip("s") + for column in columns: + if each_term in _terms(column): + return column + return None + + @component class ChartGenerationPostProcessor: @component.output_types( diff --git a/wren-ai-service/src/pipelines/generation/utils/sql.py b/wren-ai-service/src/pipelines/generation/utils/sql.py index 088282574e..e01f274c6a 100644 --- a/wren-ai-service/src/pipelines/generation/utils/sql.py +++ b/wren-ai-service/src/pipelines/generation/utils/sql.py @@ -1,27 +1,5037 @@ import logging +import re +import time +from datetime import datetime, timezone +from functools import lru_cache from typing import Any, Dict, List import aiohttp import orjson +import sqlparse from haystack import component -from haystack.dataclasses import ChatMessage -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict +from sqlparse.sql import Function, Identifier, IdentifierList, TokenList +from sqlparse.tokens import Comment, Keyword from src.core.engine import ( Engine, clean_generation_result, ) from src.pipelines.retrieval.sql_knowledge import SqlKnowledge +from src.providers.llm import ChatMessage from src.web.v1.services.ask import AskHistory logger = logging.getLogger("wren-ai-service") +_DDL_CREATE_PATTERN = re.compile( + r"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P[^\s(]+)\s*\(", + re.IGNORECASE, +) +_DDL_COLUMN_KEYWORDS = { + "CHECK", + "CONSTRAINT", + "FOREIGN", + "INDEX", + "KEY", + "PRIMARY", + "UNIQUE", +} + +def _timing_ms(started_at: float) -> float: + return (time.perf_counter() - started_at) * 1000 + +_DDL_CREATE_PATTERN = re.compile( + r"\bCREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P
[^\s(]+)\s*\(", + re.IGNORECASE, +) +_DDL_COLUMN_KEYWORDS = { + "CHECK", + "CONSTRAINT", + "FOREIGN", + "INDEX", + "KEY", + "PRIMARY", + "UNIQUE", +} + + +_IDENTIFIER_TOKEN = r'"[^"]+"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*' +_QUALIFIED_IDENTIFIER = rf"(?:{_IDENTIFIER_TOKEN})(?:\s*\.\s*(?:{_IDENTIFIER_TOKEN}))*" +_SIMPLE_IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_$]*$") +_DDL_RELATION = re.compile( + rf"\bCREATE\s+(?:TABLE|VIEW)\s+(?P{_QUALIFIED_IDENTIFIER})", + re.IGNORECASE, +) +_RELATION_REFERENCE = re.compile( + rf"\b(?:FROM|JOIN)\s+(?P{_QUALIFIED_IDENTIFIER})" + rf"(?:\s+(?:AS\s+)?(?P{_IDENTIFIER_TOKEN}))?", + re.IGNORECASE, +) +_TSQL_TOP_LIMIT = re.compile( + r"(?is)^(\s*)SELECT\s+TOP\s*\(?\s*(\d+)\s*\)?\s+(?!PERCENT\b)(.+?)\s*;?\s*$" +) +_TO_DATE_SIMPLE = re.compile( + r"(?is)\bTO_DATE\s*\(\s*" + r"(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*" + r"(?:\s*\.\s*(?:\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*))?)" + r"\s*,\s*'[^']+'\s*\)" +) +_ORDER_BY_ALIAS_ITEM = re.compile( + r"(?is)^(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)" + r"(?P\s+(?:ASC|DESC))?$" +) +_JOINED_DESC_LIMIT = re.compile(r"(?i)\bDESC\s*LIMIT\b|DESCLIMIT") +_JOINED_ASC_LIMIT = re.compile(r"(?i)\bASC\s*LIMIT\b|ASCLIMIT") +_CTE_REFERENCE = re.compile( + rf"(?:\bWITH|,)\s+(?P{_IDENTIFIER_TOKEN})\s+AS\s*\(", + re.IGNORECASE, +) +_QUALIFIED_COLUMN = re.compile( + rf"(?P{_QUALIFIED_IDENTIFIER})\s*\.\s*(?P{_IDENTIFIER_TOKEN})" +) +_SQL_START = re.compile(r"^\s*(?:WITH|SELECT)\b", re.IGNORECASE | re.DOTALL) +_SQL_REASONING_DISALLOWED = re.compile( + r"(?is)```|\b(?:SELECT|WITH|FROM|JOIN|WHERE|GROUP\s+BY|ORDER\s+BY|" + r"HAVING|LIMIT|UNION)\b" +) +_SQL_REASONING_ASSUMPTION_WORDS = re.compile( + r"(?i)\b(?:assume|assuming|likely|possible|might|example)\b" +) +_SAFE_SQL_REASONING_PLAN = """1. **Identify Supported Data**: Use only the retrieved schema metadata for the active project to decide whether the question is supported. +2. **Ground Requested Filters**: Match user-entered names, statuses, products, and similar values only against verified schema metadata or sample values; ask for clarification when the match is not clear. +3. **Prepare The Result Shape**: Build the requested rows, totals, time breakdowns, ordering, or limits using only validated retrieved tables and columns.""" +_SQL_OBJECT_ALIAS_STOP_WORDS = { + "CROSS", + "EXCEPT", + "FETCH", + "FULL", + "GROUP", + "HAVING", + "INNER", + "INTERSECT", + "JOIN", + "LEFT", + "LIMIT", + "MATCH_RECOGNIZE", + "NATURAL", + "OFFSET", + "ORDER", + "RIGHT", + "TABLESAMPLE", + "UNION", + "WHERE", +} + + +def sanitize_sql_generation_reasoning(reasoning: Any) -> str: + text = str(reasoning or "").strip() + if not text: + return _SAFE_SQL_REASONING_PLAN + + if _SQL_REASONING_DISALLOWED.search(text) or _SQL_REASONING_ASSUMPTION_WORDS.search( + text + ): + logger.warning( + "SQL generation reasoning violated schema-grounding display contract; " + "using safe non-executable reasoning plan." + ) + return _SAFE_SQL_REASONING_PLAN + + return text +_KEYED_SEMANTIC_IDENTIFIER_ALIAS = re.compile( + r"""(?ix) + \b(?:alias|displayName|display_name|sourceColumnName|source_column_name| + sourceTableName|source_table_name|source_table|source_name| + physicalName|physical_name|referenceName|reference_name| + lineageName|lineage_name) + ['"]?\s*[:=]\s*['"](?P[^'"]+)['"] + """ +) +_UNQUALIFIED_QUOTED_IDENTIFIER = re.compile(r'(?(?:[^"]|"")*)"') +_UNQUALIFIED_BARE_IDENTIFIER = re.compile( + r"(?[A-Za-z_][A-Za-z0-9_$]*)\b(?!\s*\.)" +) +_SINGLE_QUOTED_LITERAL = re.compile(r"'(?:''|[^'])*'") +_SQL_RESERVED_WORDS = { + "ALL", + "ALTER", + "AND", + "ASC", + "AS", + "BETWEEN", + "BY", + "CASE", + "CAST", + "COUNT", + "CREATE", + "CROSS", + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "DELETE", + "DESC", + "DISTINCT", + "ELSE", + "END", + "EXCEPT", + "FALSE", + "FETCH", + "FOR", + "FROM", + "FULL", + "GROUP", + "HAVING", + "IN", + "ISNULL", + "INNER", + "INSERT", + "INTERVAL", + "INTERSECT", + "IS", + "JOIN", + "LEFT", + "LIKE", + "LIMIT", + "NATURAL", + "NO", + "NOT", + "NULL", + "OFFSET", + "ON", + "OR", + "ORDER", + "OUTER", + "RIGHT", + "SELECT", + "TABLE", + "TABLESAMPLE", + "THEN", + "TRUE", + "UNION", + "UPDATE", + "VALUES", + "WHEN", + "WHERE", + "WINDOW", + "WITH", +} +_SQL_FUNCTION_WORDS = { + "ABS", + "AVG", + "CAST", + "CEIL", + "CEILING", + "COALESCE", + "CONCAT", + "COUNT", + "COUNT_BIG", + "CURRENT_DATE", + "CURRENT_TIME", + "CURRENT_TIMESTAMP", + "DATE_TRUNC", + "DAY", + "EXTRACT", + "FLOOR", + "LOWER", + "MAX", + "MIN", + "MONTH", + "NULLIF", + "ROUND", + "SUM", + "TRIM", + "UPPER", + "YEAR", +} +_SQL_TYPE_WORDS = { + "BIGINT", + "BOOLEAN", + "CHAR", + "DATE", + "DATETIME", + "DECIMAL", + "DOUBLE", + "FLOAT", + "FLOAT4", + "FLOAT8", + "INT", + "INT2", + "INT4", + "INT8", + "INTEGER", + "NUMERIC", + "REAL", + "SMALLINT", + "TEXT", + "TIME", + "TIMESTAMP", + "VARCHAR", +} +_DATE_PART_WORDS = { + "DAY", + "DOW", + "DOY", + "HOUR", + "MICROSECOND", + "MILLISECOND", + "MINUTE", + "MONTH", + "QUARTER", + "SECOND", + "WEEK", + "YEAR", +} +_FALLBACK_TOKEN = re.compile(r"[a-z0-9]+") +_QUERY_VALUE_TOKEN = re.compile(r"[A-Za-z0-9]+(?:[-_][A-Za-z0-9]+)*") +_FILTER_VALUE_PHRASE = re.compile( + r"""(?ix) + \b(?Pcalled|contains|containing|for|from|like|matching|named|where|with)\s+ + (?P[A-Za-z0-9][A-Za-z0-9_ ./&'()-]*?) + (?= + \s+\b(?:across|after|before|by|during|group|grouped|having|in|limit|on| + order|ordered|per|since|sort|sorted|this|last|next|where|with|from|for)\b + |[?,!;:] + |$ + ) + """ +) +_OPEN_FILTER_VALUE_INTRODUCERS = {"for", "from"} +_FALLBACK_STOPWORDS = { + "a", + "an", + "and", + "are", + "by", + "from", + "in", + "of", + "show", + "that", + "the", + "to", + "with", +} +_MONTH_NAME_TO_NUMBER = { + "january": 1, + "february": 2, + "march": 3, + "april": 4, + "may": 5, + "june": 6, + "july": 7, + "august": 8, + "september": 9, + "october": 10, + "november": 11, + "december": 12, +} + + +def _fallback_token_variants(token: str) -> set[str]: + token = token.lower() + variants = {token} + + generic_tokens = globals().get("_GENERIC_SCHEMA_INTENT_TOKENS", set()) + if token in generic_tokens: + return variants + + if len(token) > 4 and token.endswith("ies"): + variants.add(token[:-3] + "y") + elif len(token) > 4 and token.endswith("es"): + if token.endswith(("ches", "shes", "sses", "uses", "xes", "zes")): + variants.add(token[:-2]) + else: + variants.add(token[:-1]) + elif len(token) > 3 and token.endswith("s") and not token.endswith(("ss", "us")): + variants.add(token[:-1]) + if len(token) > 4 and token.endswith("ed"): + stem = token[:-2] + variants.add(stem) + if not stem.endswith("e"): + variants.add(stem + "e") + if len(token) > 5 and token[-3] in {"d", "s", "t", "v", "z"}: + variants.add(token[:-1]) + if len(stem) > 2 and stem[-1] == stem[-2]: + variants.add(stem[:-1]) + if len(token) > 5 and token.endswith("ing"): + stem = token[:-3] + variants.add(stem) + if not stem.endswith("d"): + variants.add(stem + "e") + if len(stem) > 2 and stem[-1] == stem[-2]: + variants.add(stem[:-1]) + + return {variant for variant in variants if variant} + + +def _expand_fallback_token_variants(tokens: set[str]) -> set[str]: + expanded = set() + for token in tokens: + expanded.update(_fallback_token_variants(token)) + return expanded + + +_COMPOUND_IDENTIFIER_PART_TOKENS = { + "account", + "amount", + "balance", + "business", + "buyer", + "category", + "client", + "company", + "count", + "currency", + "customer", + "date", + "division", + "end", + "exchange", + "failure", + "gross", + "group", + "invoice", + "market", + "material", + "month", + "name", + "order", + "person", + "priority", + "product", + "quantity", + "rate", + "record", + "repair", + "sales", + "salesperson", + "severity", + "status", + "supplier", + "task", + "ticket", + "type", + "unit", + "value", + "vendor", + "year", +} +_COMPOUND_IDENTIFIER_ALIASES = { + "acct": {"account"}, + "amt": {"amount"}, + "bu": {"business", "unit"}, + "curr": {"currency"}, + "cust": {"customer"}, + "gl": {"general", "ledger"}, + "ord": {"order"}, + "prod": {"product"}, + "qty": {"quantity"}, + "vend": {"vendor"}, +} + + +def _compound_identifier_tokens(token: str) -> set[str]: + if len(token) < 5: + return set() + + tokens: set[str] = set() + for part in _COMPOUND_IDENTIFIER_PART_TOKENS: + if part != token and len(part) >= 3 and part in token: + tokens.add(part) + + if token not in _COMPOUND_IDENTIFIER_PART_TOKENS: + for alias, expansions in _COMPOUND_IDENTIFIER_ALIASES.items(): + if alias != token and (token.startswith(alias) or token.endswith(alias)): + tokens.add(alias) + tokens.update(expansions) + + return tokens + + +def normalize_wren_sql_dialect(sql: str) -> str: + if not sql: + return sql + + sql = _TO_DATE_SIMPLE.sub( + lambda match: f"CAST({match.group('expr')} AS DATE)", + sql, + ) + sql = _JOINED_DESC_LIMIT.sub("DESC LIMIT", sql) + sql = _JOINED_ASC_LIMIT.sub("ASC LIMIT", sql) + sql = _replace_order_by_aliases_with_select_expressions(sql) + sql = _JOINED_DESC_LIMIT.sub("DESC LIMIT", sql) + sql = _JOINED_ASC_LIMIT.sub("ASC LIMIT", sql) + + if re.search(r"(?i)\bLIMIT\s+\d+\b", sql): + return sql + + match = _TSQL_TOP_LIMIT.match(sql) + if not match: + return sql + + leading_space, limit, select_body = match.groups() + if re.search(r"(?i)\bWITH\s+TIES\b", select_body): + return sql + + return f"{leading_space}SELECT {select_body.strip()}\nLIMIT {limit}" + + +def _replace_order_by_aliases_with_select_expressions(sql: str) -> str: + select_clause = _extract_select_clause(sql) + order_by_clause = _extract_clause(sql, "ORDER BY", ("LIMIT", "OFFSET")) + if not select_clause or not order_by_clause: + return sql + + alias_to_expression = {} + for select_item in _split_sql_tokens(select_clause): + expression, alias = _split_select_expression_alias(select_item) + if alias and expression: + alias_to_expression[alias] = expression + + if not alias_to_expression: + return sql + + replaced_any = False + rewritten_order_items = [] + for order_item in _split_sql_tokens(order_by_clause): + match = _ORDER_BY_ALIAS_ITEM.match(order_item.strip()) + if not match: + rewritten_order_items.append(order_item) + continue + + alias = _unquote_identifier(match.group("identifier")) + expression = alias_to_expression.get(alias) + if not expression: + rewritten_order_items.append(order_item) + continue + + rewritten_order_items.append(f"{expression}{match.group('suffix') or ''}") + replaced_any = True + + if not replaced_any: + return sql + + order_by_match = re.search( + r"(?is)\bORDER\s+BY\b\s+.*?(?=\b(?:LIMIT|OFFSET)\b|$)", + sql, + ) + if not order_by_match: + return sql + + suffix = sql[order_by_match.end() :] + if suffix and not suffix[0].isspace(): + suffix = "\n" + suffix + replacement = "ORDER BY " + ", ".join(rewritten_order_items) + return sql[: order_by_match.start()] + replacement + suffix + + +def _unquote_identifier(identifier: str) -> str: + identifier = identifier.strip() + if len(identifier) >= 2 and identifier[0] == "[" and identifier[-1] == "]": + return identifier[1:-1].replace("]]", "]") + if len(identifier) >= 2 and identifier[0] == '"' and identifier[-1] == '"': + return identifier[1:-1].replace('""', '"') + return identifier + + +def _split_qualified_identifier(identifier: str) -> list[str]: + parts = [] + current = [] + in_double_quote = False + in_bracket = False + index = 0 + + while index < len(identifier): + char = identifier[index] + nxt = identifier[index + 1] if index + 1 < len(identifier) else None + + if in_double_quote: + current.append(char) + if char == '"' and nxt == '"': + current.append(nxt) + index += 2 + continue + if char == '"': + in_double_quote = False + index += 1 + continue + + if in_bracket: + current.append(char) + if char == "]": + in_bracket = False + index += 1 + continue + + if char == '"': + current.append(char) + in_double_quote = True + elif char == "[": + current.append(char) + in_bracket = True + elif char == ".": + part = "".join(current).strip() + if part: + parts.append(part) + current = [] + else: + current.append(char) + index += 1 + + part = "".join(current).strip() + if part: + parts.append(part) + + return parts + + +def _normalize_identifier(identifier: str) -> str: + return ".".join(_unquote_identifier(part) for part in _split_qualified_identifier(identifier)) + + +def _quote_identifier(identifier: str) -> str: + return f'"{identifier.replace(chr(34), chr(34) * 2)}"' + + +def _split_sql_tokens(sql: str) -> list[str]: + tokens = [] + current = [] + in_double_quote = False + + for char in sql: + if char == '"': + current.append(char) + in_double_quote = not in_double_quote + continue + if char == "," and not in_double_quote: + token = "".join(current).strip() + if token: + tokens.append(token) + current = [] + continue + current.append(char) + + token = "".join(current).strip() + if token: + tokens.append(token) + return tokens + + +def _identifier_needs_quotes(identifier: str) -> bool: + return ( + not _SIMPLE_IDENTIFIER.fullmatch(identifier) + or identifier.upper() in _SQL_RESERVED_WORDS + ) + + +def _iter_context_texts(contexts: list[Any] | None): + if not contexts: + return + for context in contexts: + yield getattr(context, "content", context) + + +def _context_cache_key(contexts: list[Any] | None) -> tuple[str, ...]: + if not contexts: + return tuple() + return tuple(str(context) for context in _iter_context_texts(contexts)) + + +def _clean_contract_value(value: str) -> str: + value = value.strip().strip(",") + if not value: + return "" + return _normalize_identifier(value) + + +def _parse_contract_values(value: str) -> list[str]: + value = value.strip().strip(",") + if not value: + return [] + + try: + loaded = orjson.loads(value) + except orjson.JSONDecodeError: + loaded = None + + if isinstance(loaded, list): + return [ + _clean_contract_value(str(item)) + for item in loaded + if _clean_contract_value(str(item)) + ] + + parsed = _clean_contract_value(value) + return [parsed] if parsed else [] + + +def _extract_contract_schema_index( + contexts: list[Any] | None, +) -> dict[str, set[str] | None]: + schema_index: dict[str, set[str] | None] = {} + if not contexts: + return schema_index + + for context in _iter_context_texts(contexts): + current_relation = None + reading_columns = False + + for raw_line in str(context).splitlines(): + line = raw_line.strip() + if not line: + continue + + if line.startswith("sql_table_name_use_exactly:"): + current_relation = _clean_contract_value(line.split(":", 1)[1]) + if current_relation: + schema_index.setdefault(current_relation, set()) + reading_columns = False + continue + + if line.startswith("sql_column_names_use_exactly:"): + reading_columns = True + if current_relation: + columns = schema_index.setdefault(current_relation, set()) + if columns is not None: + for value in _parse_contract_values(line.split(":", 1)[1]): + columns.add(value) + continue + + if line.startswith("relationship_constraints_use_exactly:"): + reading_columns = False + continue + + if line.startswith("sql_column_name_use_exactly:"): + if current_relation: + columns = schema_index.setdefault(current_relation, set()) + if columns is not None: + for value in _parse_contract_values(line.split(":", 1)[1]): + columns.add(value) + continue + + if reading_columns and line.startswith("-") and current_relation: + columns = schema_index.setdefault(current_relation, set()) + if columns is not None: + value = _clean_contract_value(line[1:]) + if value: + columns.add(value) + continue + + if line.startswith(("END WREN SQL IDENTIFIER CONTRACT", "Only ")): + reading_columns = False + + return schema_index + + +def _extract_schema_identifiers(contexts: list[Any] | None) -> list[str]: + if not contexts: + return [] + + identifiers: list[str] = [] + seen = set() + + def add(identifier: str) -> None: + identifier = _unquote_identifier(identifier.strip().rstrip(",")) + if not identifier or identifier.upper() in {"FOREIGN", "PRIMARY", "KEY"}: + return + if identifier not in seen: + seen.add(identifier) + identifiers.append(identifier) + + for relation, columns in _extract_schema_index(contexts).items(): + add(relation) + if columns: + for column in columns: + add(column) + + for context in _iter_context_texts(contexts): + for match in _DDL_RELATION.finditer(context): + add(match.group("name")) + + in_table = False + for raw_line in context.splitlines(): + line = raw_line.strip() + if not line or line.startswith("--") or line.startswith("/*"): + continue + if re.search(r"\bCREATE\s+TABLE\b", line, re.IGNORECASE): + in_table = True + remainder = line.split("(", 1) + if len(remainder) == 1: + continue + line = remainder[1].strip() + if not in_table: + continue + if line.startswith(");") or line == ")": + in_table = False + continue + line = line.split("--", 1)[0].strip().rstrip(",") + if not line or line.upper().startswith(("FOREIGN KEY", "PRIMARY KEY")): + continue + if line.startswith('"'): + end = line.find('"', 1) + while end != -1 and end + 1 < len(line) and line[end + 1] == '"': + end = line.find('"', end + 2) + if end > 0: + add(line[: end + 1]) + else: + add(line.split(None, 1)[0]) + + return identifiers + + +def _thaw_schema_index( + frozen_schema_index: tuple[tuple[str, tuple[str, ...] | None], ...], +) -> dict[str, set[str] | None]: + return { + relation: None if columns is None else set(columns) + for relation, columns in frozen_schema_index + } + + +def _freeze_schema_index( + schema_index: dict[str, set[str] | None], +) -> tuple[tuple[str, tuple[str, ...] | None], ...]: + return tuple( + sorted( + ( + relation, + None if columns is None else tuple(sorted(columns)), + ) + for relation, columns in schema_index.items() + ) + ) + + +def _extract_schema_index(contexts: list[Any] | None) -> dict[str, set[str] | None]: + return _thaw_schema_index(_cached_extract_schema_index(_context_cache_key(contexts))) + + +@lru_cache(maxsize=256) +def _cached_extract_schema_index( + context_texts: tuple[str, ...], +) -> tuple[tuple[str, tuple[str, ...] | None], ...]: + if not context_texts: + return tuple() + + schema_index = _extract_contract_schema_index(list(context_texts)) + + for context in context_texts: + relation_match = _DDL_RELATION.search(context) + if not relation_match: + continue + + relation_name = _normalize_identifier(relation_match.group("name")) + if re.search(r"\bCREATE\s+VIEW\b", context, re.IGNORECASE): + schema_index.setdefault(relation_name, None) + continue + + column_block_match = re.search( + r"\bCREATE\s+TABLE\b[^(]*\((?P.*)\)\s*;?", + context, + re.IGNORECASE | re.DOTALL, + ) + if not column_block_match: + schema_index.setdefault(relation_name, None) + continue + + columns = set() + for raw_column in _split_sql_tokens(column_block_match.group("columns")): + line = "\n".join( + line.strip() + for line in raw_column.splitlines() + if line.strip() + and not line.strip().startswith("--") + and not line.strip().startswith("/*") + ).strip() + if not line: + continue + line = line.split("--", 1)[0].strip() + if not line or line.startswith("/*"): + continue + if line.upper().startswith(("FOREIGN KEY", "PRIMARY KEY")): + continue + if line.startswith('"'): + end = line.find('"', 1) + while end != -1 and end + 1 < len(line) and line[end + 1] == '"': + end = line.find('"', end + 2) + if end > 0: + columns.add(_unquote_identifier(line[: end + 1])) + else: + columns.add(_unquote_identifier(line.split(None, 1)[0])) + + existing_columns = schema_index.get(relation_name) + if existing_columns is None: + schema_index[relation_name] = columns + else: + existing_columns.update(columns) + + return _freeze_schema_index(schema_index) + + +def _semantic_tokens_from_value(value: Any) -> set[str]: + tokens: set[str] = set() + if isinstance(value, str): + tokens.update(_fallback_tokens(value)) + elif isinstance(value, dict): + for nested_value in value.values(): + tokens.update(_semantic_tokens_from_value(nested_value)) + elif isinstance(value, list): + for nested_value in value: + tokens.update(_semantic_tokens_from_value(nested_value)) + return tokens + + +def _extract_semantic_context_payload(context: str) -> dict[str, Any]: + start_marker = "WREN RETRIEVED SEMANTIC CONTEXT" + end_marker = "WREN SQL IDENTIFIER CONTRACT" + start_index = context.upper().find(start_marker) + end_index = context.upper().find(end_marker, start_index) + if start_index < 0 or end_index < 0: + return {} + payload_text = context[start_index + len(start_marker) : end_index].strip() + try: + payload = orjson.loads(payload_text) + except orjson.JSONDecodeError: + return {} + return payload if isinstance(payload, dict) else {} + + +def _extract_semantic_tokens_by_column( + context: str, +) -> tuple[set[str], dict[str, set[str]], dict[str, list[str]]]: + payload = _extract_semantic_context_payload(context) + if not payload: + return set(), {}, {} + + table_tokens = _semantic_tokens_from_value( + payload.get("semantic_context_not_sql_identifiers") + ) + table_tokens.update(_semantic_tokens_from_value(payload.get("object_type"))) + + column_tokens: dict[str, set[str]] = {} + column_sample_values: dict[str, list[str]] = {} + for column in payload.get("columns", []) or []: + if not isinstance(column, dict): + continue + column_name = column.get("sql_column_name_use_exactly") + if not isinstance(column_name, str) or not column_name: + continue + tokens = _semantic_tokens_from_value( + column.get("semantic_context_not_sql_identifier") + ) + tokens.update(_semantic_tokens_from_value(column.get("display_name"))) + tokens.update(_semantic_tokens_from_value(column.get("source_column_name"))) + if tokens: + column_tokens[column_name] = tokens + sample_values = _extract_column_sample_values(column) + if sample_values: + column_sample_values[column_name] = sample_values + + return table_tokens, column_tokens, column_sample_values + + +def _extract_column_sample_values(column: dict[str, Any]) -> list[str]: + values: list[str] = [] + + def add(value: Any) -> None: + if value is None: + return + if isinstance(value, (list, tuple, set)): + for item in value: + add(item) + return + if isinstance(value, dict): + for item in value.values(): + add(item) + return + text = str(value).strip() + if text and text.lower() not in {item.lower() for item in values}: + values.append(text) + + for key in ( + "sample_values", + "sample_value", + "samples", + "values", + "example_values", + "examples", + "distinct_values", + ): + add(column.get(key)) + return values + + +def _thaw_schema_details( + frozen_schema_details: tuple[ + tuple[str, tuple[tuple[str, str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], ...]], + ..., + ], +) -> dict[str, list[dict[str, Any]]]: + return { + relation: [ + { + "name": name, + "data_type": data_type, + "semantic_tokens": set(semantic_tokens), + "sample_values": list(sample_values), + "_table_semantic_tokens": set(table_semantic_tokens), + } + for ( + name, + data_type, + semantic_tokens, + sample_values, + table_semantic_tokens, + ) in columns + ] + for relation, columns in frozen_schema_details + } + + +def _freeze_schema_details( + schema_details: dict[str, list[dict[str, Any]]], +) -> tuple[ + tuple[str, tuple[tuple[str, str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], ...]], + ..., +]: + return tuple( + ( + relation, + tuple( + ( + column["name"], + column["data_type"], + tuple(sorted(column.get("semantic_tokens") or set())), + tuple(column.get("sample_values") or []), + tuple(sorted(column.get("_table_semantic_tokens") or set())), + ) + for column in columns + ), + ) + for relation, columns in sorted(schema_details.items()) + ) + + +def _extract_schema_details( + contexts: list[Any] | None, +) -> dict[str, list[dict[str, Any]]]: + return _thaw_schema_details( + _cached_extract_schema_details(_context_cache_key(contexts)) + ) + + +@lru_cache(maxsize=256) +def _cached_extract_schema_details( + context_texts: tuple[str, ...], +) -> tuple[ + tuple[str, tuple[tuple[str, str, tuple[str, ...], tuple[str, ...], tuple[str, ...]], ...]], + ..., +]: + if not context_texts: + return tuple() + + schema_details: dict[str, list[dict[str, str]]] = {} + + for context in context_texts: + relation_match = _DDL_RELATION.search(context) + if not relation_match: + continue + + relation_name = _unquote_identifier(relation_match.group("name")) + table_semantic_tokens, column_semantic_tokens, column_sample_values = ( + _extract_semantic_tokens_by_column(context) + ) + column_block_match = re.search( + r"\bCREATE\s+TABLE\b[^(]*\((?P.*)\)\s*;?", + context, + re.IGNORECASE | re.DOTALL, + ) + if not column_block_match: + continue + + columns = [] + for raw_column in _split_sql_tokens(column_block_match.group("columns")): + line = "\n".join( + line.strip() + for line in raw_column.splitlines() + if line.strip() + and not line.strip().startswith("--") + and not line.strip().startswith("/*") + ).strip() + if not line: + continue + line = line.split("--", 1)[0].strip() + if not line or line.upper().startswith(("FOREIGN KEY", "PRIMARY KEY")): + continue + + if line.startswith('"'): + end = line.find('"', 1) + while end != -1 and end + 1 < len(line) and line[end + 1] == '"': + end = line.find('"', end + 2) + if end <= 0: + continue + name = _unquote_identifier(line[: end + 1]) + remainder = line[end + 1 :].strip() + else: + parts = line.split(None, 1) + if not parts: + continue + name = _unquote_identifier(parts[0]) + remainder = parts[1].strip() if len(parts) > 1 else "" + + data_type = remainder.split(None, 1)[0].upper() if remainder else "" + columns.append( + { + "name": name, + "data_type": data_type, + "semantic_tokens": column_semantic_tokens.get(name, set()), + "sample_values": column_sample_values.get(name, []), + "_table_semantic_tokens": table_semantic_tokens, + } + ) + + schema_details[relation_name] = columns + + return _freeze_schema_details(schema_details) + + +def _is_identifier_boundary(char: str | None) -> bool: + return char is None or not (char.isalnum() or char in {"_", "$", '"'}) + + +def _replace_identifier_outside_literals(sql: str, identifier: str) -> str: + quoted = _quote_identifier(identifier) + result = [] + index = 0 + in_single_quote = False + in_double_quote = False + in_line_comment = False + in_block_comment = False + length = len(sql) + identifier_length = len(identifier) + + while index < length: + current = sql[index] + nxt = sql[index + 1] if index + 1 < length else None + + if in_line_comment: + result.append(current) + if current == "\n": + in_line_comment = False + index += 1 + continue + if in_block_comment: + result.append(current) + if current == "*" and nxt == "/": + result.append(nxt) + index += 2 + in_block_comment = False + else: + index += 1 + continue + if in_single_quote: + result.append(current) + if current == "'" and nxt == "'": + result.append(nxt) + index += 2 + elif current == "'": + in_single_quote = False + index += 1 + else: + index += 1 + continue + if in_double_quote: + result.append(current) + if current == '"' and nxt == '"': + result.append(nxt) + index += 2 + elif current == '"': + in_double_quote = False + index += 1 + else: + index += 1 + continue + + if current == "-" and nxt == "-": + result.append(current) + result.append(nxt) + index += 2 + in_line_comment = True + continue + if current == "/" and nxt == "*": + result.append(current) + result.append(nxt) + index += 2 + in_block_comment = True + continue + if current == "'": + result.append(current) + index += 1 + in_single_quote = True + continue + if current == '"': + result.append(current) + index += 1 + in_double_quote = True + continue + + if sql.startswith(identifier, index): + before = sql[index - 1] if index > 0 else None + after_index = index + identifier_length + after = sql[after_index] if after_index < length else None + if _is_identifier_boundary(before) and _is_identifier_boundary(after): + result.append(quoted) + index = after_index + continue + + result.append(current) + index += 1 + + return "".join(result) + + +def _replace_bracket_identifiers(sql: str, valid_identifiers: set[str]) -> str: + result = [] + index = 0 + in_single_quote = False + in_double_quote = False + length = len(sql) + + while index < length: + current = sql[index] + nxt = sql[index + 1] if index + 1 < length else None + + if in_single_quote: + result.append(current) + if current == "'" and nxt == "'": + result.append(nxt) + index += 2 + elif current == "'": + in_single_quote = False + index += 1 + else: + index += 1 + continue + if in_double_quote: + result.append(current) + if current == '"' and nxt == '"': + result.append(nxt) + index += 2 + elif current == '"': + in_double_quote = False + index += 1 + else: + index += 1 + continue + if current == "'": + result.append(current) + in_single_quote = True + index += 1 + continue + if current == '"': + result.append(current) + in_double_quote = True + index += 1 + continue + if current == "[": + end = sql.find("]", index + 1) + if end > index: + identifier = sql[index + 1 : end] + if identifier in valid_identifiers: + result.append(_quote_identifier(identifier)) + index = end + 1 + continue + result.append(current) + index += 1 + + return "".join(result) + + +def _semantic_key_can_be_identifier(key_hint: str) -> bool: + key = key_hint.rsplit(".", 1)[-1].lower() + return key in { + "alias", + "displayname", + "display_name", + "lineagename", + "lineage_name", + "physicalname", + "physical_name", + "referencename", + "reference_name", + "sourcecolumnname", + "source_column_name", + "source_name", + "source_table", + "source_table_name", + } + + +def _clean_semantic_identifier_alias(value: Any) -> str | None: + if value is None: + return None + + text = str(value).strip() + if not text or "\n" in text or len(text) > 256: + return None + + text = re.sub(r"^\s*/\*+", "", text) + text = re.sub(r"\*+/\s*$", "", text) + text = text.strip(" \t\r\n'\"`.,;:(){}") + if not text or len(text) > 256: + return None + if re.search(r"[!?;]", text): + return None + + normalized = _normalize_identifier(text) + if not normalized: + return None + if normalized.upper() in _SQL_RESERVED_WORDS: + return None + if "." not in normalized and normalized.lower() in _GENERIC_SCHEMA_INTENT_TOKENS: + return None + + if re.search(r"\s", normalized): + tokens = _fallback_tokens(normalized) + if len(tokens) > 5: + return None + + return normalized + + +def _table_reference_aliases(value: Any) -> set[str]: + if not isinstance(value, dict): + return set() + + table_reference = value.get("source_table_reference") or value.get( + "tableReference" + ) + if not isinstance(table_reference, dict): + table_reference = value + + table = table_reference.get("table") + if not table: + return set() + + aliases: set[str] = set() + schema = table_reference.get("schema") + catalog = table_reference.get("catalog") + if schema: + aliases.add(f"{schema}.{table}") + if catalog and schema: + aliases.add(f"{catalog}.{schema}.{table}") + aliases.add(str(table)) + return { + alias + for alias in (_clean_semantic_identifier_alias(alias) for alias in aliases) + if alias + } + + +def _keyed_semantic_identifier_aliases(text: str) -> set[str]: + aliases: set[str] = set() + for match in _KEYED_SEMANTIC_IDENTIFIER_ALIAS.finditer(text): + alias = _clean_semantic_identifier_alias(match.group("value")) + if alias: + aliases.add(alias) + return aliases + + +def _semantic_identifier_aliases(value: Any, key_hint: str = "") -> set[str]: + aliases: set[str] = set() + if value is None: + return aliases + + if isinstance(value, dict): + aliases.update(_table_reference_aliases(value)) + for key, nested_value in value.items(): + nested_key = f"{key_hint}.{key}" if key_hint else str(key) + aliases.update(_semantic_identifier_aliases(nested_value, nested_key)) + return aliases + + if isinstance(value, (list, tuple, set)): + for nested_value in value: + aliases.update(_semantic_identifier_aliases(nested_value, key_hint)) + return aliases + + text = str(value) + if _semantic_key_can_be_identifier(key_hint): + alias = _clean_semantic_identifier_alias(text) + if alias: + aliases.add(alias) + aliases.update(_keyed_semantic_identifier_aliases(text)) + return aliases + + +def _identifier_alias_key(identifier: str) -> str: + return _normalize_identifier(identifier).lower() + + +def _identifier_style_aliases(identifier: str) -> set[str]: + cleaned = _clean_semantic_identifier_alias(identifier) + if not cleaned: + return set() + + aliases = {cleaned} + parts = re.findall(r"[A-Za-z0-9]+", cleaned) + if not 1 < len(parts) <= 6: + return aliases + + aliases.add("_".join(parts)) + aliases.add("".join(parts)) + aliases.add(parts[0] + "".join(part[:1].upper() + part[1:] for part in parts[1:])) + aliases.add("".join(part[:1].upper() + part[1:] for part in parts)) + return { + alias + for alias in aliases + if _clean_semantic_identifier_alias(alias) + } + + +def _add_unique_identifier_alias( + aliases: dict[str, str], + alias_values: dict[str, str], + alias: Any, + target: str, +) -> None: + if not target: + return + + for cleaned_alias in _identifier_style_aliases(alias): + if _identifier_alias_key(cleaned_alias) == _identifier_alias_key(target): + continue + + key = _identifier_alias_key(cleaned_alias) + existing = aliases.get(key) + if existing is None and key not in aliases: + aliases[key] = target + alias_values[key] = cleaned_alias + elif existing != target: + aliases[key] = "" + + +def _identifier_reference_variants(identifier: str) -> set[str]: + cleaned = _clean_semantic_identifier_alias(identifier) + if not cleaned: + return set() + + parts = cleaned.split(".") + variants = { + _quote_identifier(cleaned), + f"[{cleaned}]", + f"`{cleaned}`", + } + if not re.search(r"\s", cleaned): + variants.add(cleaned) + if len(parts) > 1: + variants.add(".".join(parts)) + variants.add(".".join(_quote_identifier(part) for part in parts)) + variants.add(".".join(f"[{part}]" for part in parts)) + variants.add(".".join(f"`{part}`" for part in parts)) + return variants + + +def _identifier_dot_aliases(identifier: str) -> set[str]: + if "." in identifier or "_" not in identifier: + return set() + + parts = [part for part in identifier.split("_") if part] + aliases = set() + if len(parts) >= 2: + aliases.add(f"{parts[0]}.{'_'.join(parts[1:])}") + if len(parts) >= 3: + aliases.add(f"{parts[0]}.{parts[1]}.{'_'.join(parts[2:])}") + return aliases + + +def _render_table_identifier(identifier: str) -> str: + return identifier if "." in identifier else _quote_identifier(identifier) + + +def _render_qualifier_identifier(identifier: str, table_names: set[str]) -> str: + if identifier in table_names: + return _render_table_identifier(identifier) + if "." in identifier or _identifier_needs_quotes(identifier): + return _quote_identifier(identifier) + return identifier + + +def _replace_sql_text_outside_literals( + sql: str, + target: str, + replacement: str, + case_insensitive: bool = False, +) -> str: + if not target: + return sql + + result = [] + index = 0 + in_single_quote = False + in_double_quote = False + in_line_comment = False + in_block_comment = False + length = len(sql) + target_length = len(target) + comparable_target = target.lower() if case_insensitive else target + + while index < length: + current = sql[index] + nxt = sql[index + 1] if index + 1 < length else None + + if in_line_comment: + result.append(current) + if current == "\n": + in_line_comment = False + index += 1 + continue + if in_block_comment: + result.append(current) + if current == "*" and nxt == "/": + result.append(nxt) + index += 2 + in_block_comment = False + else: + index += 1 + continue + if in_single_quote: + result.append(current) + if current == "'" and nxt == "'": + result.append(nxt) + index += 2 + elif current == "'": + in_single_quote = False + index += 1 + else: + index += 1 + continue + + candidate = sql[index : index + target_length] + comparable_candidate = candidate.lower() if case_insensitive else candidate + if comparable_candidate == comparable_target: + before = sql[index - 1] if index > 0 else None + after_index = index + target_length + after = sql[after_index] if after_index < length else None + if _is_identifier_boundary(before) and _is_identifier_boundary(after): + result.append(replacement) + index = after_index + continue + + if in_double_quote: + result.append(current) + if current == '"' and nxt == '"': + result.append(nxt) + index += 2 + elif current == '"': + in_double_quote = False + index += 1 + else: + index += 1 + continue + + if current == "-" and nxt == "-": + result.append(current) + result.append(nxt) + index += 2 + in_line_comment = True + continue + if current == "/" and nxt == "*": + result.append(current) + result.append(nxt) + index += 2 + in_block_comment = True + continue + if current == "'": + result.append(current) + index += 1 + in_single_quote = True + continue + if current == '"': + result.append(current) + index += 1 + in_double_quote = True + continue + + result.append(current) + index += 1 + + return "".join(result) + + +def _extract_sql_grounding(sql: str) -> dict[str, Any]: + cte_names = { + _normalize_identifier(match.group("name")) + for match in _CTE_REFERENCE.finditer(sql) + } + relation_references = [] + alias_to_relation = {} + + for match in _RELATION_REFERENCE.finditer(sql): + if _is_extract_from_clause(sql, match.start()): + continue + relation = _unquote_identifier(match.group("name")) + if relation.upper() in {"UNNEST", "LATERAL"}: + continue + alias = match.group("alias") + alias = _normalize_identifier(alias) if alias else relation + if alias.upper() in _SQL_OBJECT_ALIAS_STOP_WORDS: + alias = relation + relation_references.append(relation) + alias_to_relation[alias] = relation + alias_to_relation[relation] = relation + + qualified_columns = [ + ( + _normalize_identifier(match.group("qualifier")), + _normalize_identifier(match.group("column")), + ) + for match in _QUALIFIED_COLUMN.finditer(sql) + ] + + return { + "cte_names": cte_names, + "relation_references": relation_references, + "alias_to_relation": alias_to_relation, + "qualified_columns": qualified_columns, + } + + +def _is_extract_from_clause(sql: str, from_start: int) -> bool: + prefix = sql[:from_start] + last_open = prefix.rfind("(") + if last_open == -1 or prefix.rfind(")") > last_open: + return False + + before_open = prefix[:last_open].rstrip() + return before_open.upper().endswith("EXTRACT") + + +def _strip_string_literals(sql: str) -> str: + return _SINGLE_QUOTED_LITERAL.sub( + lambda match: " " * len(match.group(0)), + sql, + ) + + +def _extract_clause(sql: str, clause: str, end_clauses: tuple[str, ...]) -> str: + end_pattern = "|".join(re.escape(end_clause) for end_clause in end_clauses) + pattern = re.compile( + rf"(?is)\b{re.escape(clause)}\b\s+(?P.*?)(?=\b(?:{end_pattern})\b|$)" + ) + match = pattern.search(sql) + return match.group("body").strip() if match else "" + + +def _extract_select_clause(sql: str) -> str: + match = re.search(r"(?is)\bSELECT\b\s+(?P.*?)(?=\bFROM\b)", sql) + return match.group("body").strip() if match else "" + + +def _split_select_expression_alias(expression: str) -> tuple[str, str | None]: + as_match = re.search( + r"(?is)\s+AS\s+(?P\"[^\"]+\"|\[[^\]]+\]|[A-Za-z_][A-Za-z0-9_$]*)\s*$", + expression, + ) + if as_match: + return expression[: as_match.start()].strip(), _unquote_identifier( + as_match.group("alias") + ) + + return expression, None + + +def _extract_output_aliases(select_clause: str) -> set[str]: + aliases = set() + for expression in _split_sql_tokens(select_clause): + _, alias = _split_select_expression_alias(expression) + if alias: + aliases.add(alias) + return aliases + + +def _iter_unqualified_identifier_candidates(expression: str): + stripped = _strip_string_literals(expression) + + for match in _UNQUALIFIED_QUOTED_IDENTIFIER.finditer(stripped): + yield _unquote_identifier(f'"{match.group("name")}"') + + without_quoted = _UNQUALIFIED_QUOTED_IDENTIFIER.sub(" ", stripped) + for match in _UNQUALIFIED_BARE_IDENTIFIER.finditer(without_quoted): + name = match.group("name") + following = without_quoted[match.end() :].lstrip() + if following.startswith("("): + continue + yield name + + +def _sql_mentions_identifier(sql: str, identifier: str) -> bool: + stripped = _strip_string_literals(sql) + quoted_identifier = re.escape(_quote_identifier(identifier)) + bracket_identifier = re.escape(f"[{identifier}]") + bare_identifier = re.escape(identifier) + return bool( + re.search(rf'(? bool: + lowered_sql = sql.lower() + for value in values: + cleaned = _clean_filter_value(value) + if cleaned and _quote_literal(cleaned.lower()) in lowered_sql: + return True + return False + + +def _extract_sql_string_literals(sql: str) -> list[str]: + literals = [] + for match in _SINGLE_QUOTED_LITERAL.finditer(sql): + literal = match.group(0)[1:-1].replace("''", "'") + if literal: + literals.append(literal) + return literals + + +def _extract_column_filter_literals(sql: str, column_name: str) -> list[str]: + stripped = _strip_string_literals(sql) + quoted_column = re.escape(_quote_identifier(column_name)) + bare_column = re.escape(column_name) + column_pattern = rf"(?:{quoted_column}|(? set[str]: + cleaned = _clean_filter_value(str(literal).replace("%", " ")) + if not cleaned: + return set() + if re.fullmatch(r"-?\d+(?:\.\d+)?", cleaned): + return set() + if re.fullmatch( + r"\d{4}-\d{2}-\d{2}(?:[ t]\d{2}:\d{2}(?::\d{2}(?:\.\d+)?)?)?", + cleaned, + flags=re.IGNORECASE, + ): + return set() + return { + token + for token in _fallback_tokens(cleaned) + if token not in _GENERIC_SCHEMA_INTENT_TOKENS and not token.isdigit() + } + + +def _validate_literal_values_against_samples( + sql: str, + schema_details: dict[str, list[dict[str, str]]], + grounding: dict[str, Any], + query: str | None = None, +) -> str | None: + query_tokens = _fallback_tokens(query) if query else set() + referenced_relations = { + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + } + for relation in referenced_relations: + columns = schema_details.get(relation, []) + schema_tokens = _schema_tokens_for_table(relation, columns) + allowed_query_literal_tokens = _filter_value_tokens( + _schema_driven_value_terms(query, schema_tokens, columns) + ) + for column in columns: + sample_values = column.get("sample_values") or [] + literals = _extract_column_filter_literals(sql, column["name"]) + if not literals: + continue + sample_tokens = _sample_value_tokens(column) if sample_values else set() + sample_lowers = {str(value).lower() for value in sample_values} + for literal in literals: + literal_tokens = _literal_grounding_tokens(literal) + if not literal_tokens: + continue + if literal.lower() in sample_lowers or literal_tokens & sample_tokens: + continue + if ( + allowed_query_literal_tokens + and literal_tokens <= allowed_query_literal_tokens + ): + logger.info( + "Literal value accepted from grounded user filter column=%s literal=%s", + column["name"], + literal, + ) + continue + if sample_values: + return ( + "Schema grounding failed. The generated SQL filters column " + f"{column['name']} with a literal value not found in that " + "column's verified sample values or grounded filter terms." + ) + return ( + "Schema grounding failed. The generated SQL filters column " + f"{column['name']} with a literal value that is not grounded " + "as a filter value in the user question or verified sample " + "values." + ) + return None + + +def _validate_unqualified_columns_for_single_relation( + sql: str, + schema_index: dict[str, set[str] | None], + grounding: dict[str, Any], +) -> str | None: + real_relations = [ + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + ] + if grounding["cte_names"]: + return None + + unique_real_relations = list(dict.fromkeys(real_relations)) + if len(unique_real_relations) != 1: + return None + + relation = unique_real_relations[0] + valid_columns = schema_index.get(relation) + if valid_columns is None: + return None + + select_clause = _extract_select_clause(sql) + output_aliases = _extract_output_aliases(select_clause) + ignored_identifiers = ( + set(schema_index) + | set(grounding["alias_to_relation"]) + | set(grounding["cte_names"]) + | output_aliases + ) + + clause_expressions = [] + for expression in _split_sql_tokens(select_clause): + expression, _ = _split_select_expression_alias(expression) + clause_expressions.append(expression) + clause_expressions.extend( + filter( + None, + [ + _extract_clause( + sql, + "WHERE", + ("GROUP BY", "HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause( + sql, + "GROUP BY", + ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause( + sql, + "HAVING", + ("ORDER BY", "LIMIT", "OFFSET"), + ), + _extract_clause(sql, "ORDER BY", ("LIMIT", "OFFSET")), + ], + ) + ) + + invalid_columns = set() + for expression in clause_expressions: + for identifier in _iter_unqualified_identifier_candidates(expression): + upper_identifier = identifier.upper() + if ( + upper_identifier in _SQL_RESERVED_WORDS + or upper_identifier in _SQL_FUNCTION_WORDS + or upper_identifier in _SQL_TYPE_WORDS + or upper_identifier in _DATE_PART_WORDS + or identifier in ignored_identifiers + or identifier in valid_columns + ): + continue + invalid_columns.add(identifier) + + if not invalid_columns: + return None + + return ( + "Schema grounding failed. The SQL references unqualified columns that " + f"are not present in verified table or view {relation}: " + f"{', '.join(sorted(invalid_columns))}. Use only verified columns: " + f"{', '.join(sorted(valid_columns))}." + ) + + +def validate_sql_against_contexts( + sql: str, + contexts: list[Any] | None = None, +) -> str | None: + schema_index = _extract_schema_index(contexts) + if not schema_index: + return None + + valid_relations = set(schema_index) + grounding = _extract_sql_grounding(sql) + cte_names = grounding["cte_names"] + + shadowed_relations = sorted(cte_names & valid_relations) + if shadowed_relations: + return ( + "Schema grounding failed. The SQL creates CTEs with names that already " + f"belong to verified schema objects: {', '.join(shadowed_relations)}. " + "Do not create dummy CTEs for schema objects; use the verified tables or views directly." + ) + + invalid_relations = sorted( + { + relation + for relation in grounding["relation_references"] + if relation not in valid_relations and relation not in cte_names + } + ) + if invalid_relations: + return ( + "Schema grounding failed. The SQL references tables or views that are not " + f"in the retrieved schema for the active question: {', '.join(invalid_relations)}. " + f"Use only verified tables or views: {', '.join(sorted(valid_relations))}." + ) + + alias_to_relation = grounding["alias_to_relation"] + for qualifier, column in grounding["qualified_columns"]: + relation = alias_to_relation.get(qualifier) + if not relation or relation in cte_names: + continue + valid_columns = schema_index.get(relation) + if valid_columns is None: + continue + if column not in valid_columns: + return ( + "Schema grounding failed. The SQL references column " + f"{qualifier}.{column}, but column {column} is not present in verified " + f"table or view {relation}. Use only verified columns: " + f"{', '.join(sorted(valid_columns))}." + ) + + unqualified_column_error = _validate_unqualified_columns_for_single_relation( + sql, + schema_index, + grounding, + ) + if unqualified_column_error: + return unqualified_column_error + + return None + + +def _extract_sql_from_value(value: Any) -> str | None: + if value is None: + return None + + if isinstance(value, str): + text = value.strip() + if not text: + return None + + try: + parsed = orjson.loads(text) + except orjson.JSONDecodeError: + return text if _SQL_START.search(text) else None + + return _extract_sql_from_value(parsed) + + if isinstance(value, dict): + for key in ("sql", "query", "code"): + extracted = _extract_sql_from_value(value.get(key)) + if extracted: + return extracted + + extracted = _extract_sql_from_value(value.get("arguments")) + if extracted: + return extracted + + return None + + if isinstance(value, list): + for item in value: + extracted = _extract_sql_from_value(item) + if extracted: + return extracted + + return None + + +def _extract_generation_sql(generation_result: str | None) -> str | None: + if not generation_result: + return None + + extracted = _extract_sql_from_value(generation_result) + if extracted: + return extracted + + text = generation_result.strip() + return text if _SQL_START.search(text) else None + + +def _sql_has_aggregate_function(sql: str) -> bool: + return bool(re.search(r"(?is)\b(?:AVG|COUNT|MAX|MIN|SUM)\s*\(", sql)) + + +def _group_by_source_columns(sql: str) -> set[str]: + group_by_clause = _extract_clause( + sql, + "GROUP BY", + ("HAVING", "ORDER BY", "LIMIT", "OFFSET"), + ) + if not group_by_clause: + return set() + + columns: set[str] = set() + for expression in _split_sql_tokens(group_by_clause): + for identifier in _iter_unqualified_identifier_candidates(expression): + upper_identifier = identifier.upper() + if ( + upper_identifier in _SQL_RESERVED_WORDS + or upper_identifier in _SQL_FUNCTION_WORDS + or upper_identifier in _SQL_TYPE_WORDS + or upper_identifier in _DATE_PART_WORDS + ): + continue + columns.add(identifier) + return columns + + +def _query_allows_grouped_aggregate(query: str, query_tokens: set[str]) -> bool: + return bool( + _has_grouping_intent(query, query_tokens) + or _has_count_intent(query_tokens) + or _has_sum_intent(query_tokens) + or _is_rate_metric_intent(query_tokens) + or _is_distribution_metric_intent(query_tokens) + ) + + +def _missing_value_target_tokens(query: str, query_tokens: set[str]) -> set[str]: + match = re.search( + r"(?is)\b(?:blank|empty|missing|null)\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if match: + tokens = _fallback_tokens(match.group("value")) + else: + tokens = set(query_tokens) + return tokens - _GENERIC_SCHEMA_INTENT_TOKENS - _NULL_CHECK_TOKENS + + +def _sql_null_checked_columns(sql: str, columns: list[dict[str, str]]) -> set[str]: + checked_columns: set[str] = set() + stripped = _strip_string_literals(sql) + for column in columns: + quoted_column = re.escape(_quote_identifier(column["name"])) + bare_column = re.escape(column["name"]) + column_pattern = rf"(?:{quoted_column}|(? str | None: + grouping_tokens = _grouping_phrase_tokens(query) - _GENERIC_SCHEMA_INTENT_TOKENS + if not grouping_tokens or _grouping_phrase_has_multiple_dimensions(query): + return None + + group_columns = _group_by_source_columns(sql) + if len(group_columns) > 1: + return ( + "Schema grounding failed. The question asks for one grouping " + "dimension, but the generated SQL groups by multiple source " + "columns." + ) + if not group_columns: + return None + + grouped_column_names = {_normalize_identifier(name) for name in group_columns} + scored_columns: list[tuple[int, str]] = [] + for relation in referenced_relations: + for column in schema_details.get(relation) or []: + score = _column_score_for_tokens(column, grouping_tokens) + if score > 0: + scored_columns.append((score, column["name"])) + + if not scored_columns: + return None + + best_score = max(score for score, _ in scored_columns) + grouped_best_score = max( + ( + score + for score, name in scored_columns + if _normalize_identifier(name) in grouped_column_names + ), + default=0, + ) + if best_score > grouped_best_score: + logger.info( + "Generated SQL explicit grouping validation rejected grouping_tokens=%s grouped_columns=%s best_columns=%s", + sorted(grouping_tokens), + sorted(group_columns), + [ + name + for score, name in scored_columns + if score == best_score + ], + ) + return ( + "Schema grounding failed. The question asks for a specific grouping " + "dimension, but the generated SQL groups by a weaker matching " + "column." + ) + + return None + + +def validate_sql_semantic_coverage( + sql: str, + query: str | None, + contexts: list[Any] | None = None, +) -> str | None: + if not sql or not query: + return None + + raw_query_tokens = _fallback_tokens(query) + + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + + grounding = _extract_sql_grounding(sql) + referenced_relations = { + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] + } + if not referenced_relations: + return None + + schema_tokens = set() + for relation in referenced_relations: + columns = schema_details.get(relation) + if columns is not None: + schema_tokens.update(_schema_tokens_for_table(relation, columns)) + + required_tokens = _schema_required_query_tokens( + query, + raw_query_tokens, + schema_details, + ) + unsupported_tokens = _unsupported_query_tokens( + raw_query_tokens, + schema_details, + query=query, + ) + if unsupported_tokens: + return ( + "Schema grounding failed. The retrieved schema metadata does not " + "support these non-operational question term(s): " + f"{', '.join(sorted(unsupported_tokens))}. " + "Select a project with matching schema metadata or ask a supported " + "question." + ) + + missing_concepts = sorted( + token + for token in required_tokens + if not _schema_token_covered(token, schema_tokens) + ) + if missing_concepts: + return ( + "Schema grounding failed. The generated SQL uses verified identifiers, " + "but the selected table or view does not cover these schema-backed " + f"question tokens: {', '.join(missing_concepts)}. Use only schema " + "objects whose metadata supports the requested terms, or return no " + "SQL if the active project does not contain them." + ) + + subject_tokens = _query_subject_schema_tokens( + query, + raw_query_tokens, + schema_details, + ) + if subject_tokens: + identifier_tokens_by_table = { + table_name: _schema_identifier_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + selected_identifier_tokens = set() + for relation in referenced_relations: + selected_identifier_tokens.update( + identifier_tokens_by_table.get(relation, set()) + ) + if any( + _schema_tokens_cover(subject_tokens, tokens) + for tokens in identifier_tokens_by_table.values() + ) and not _schema_tokens_cover(subject_tokens, selected_identifier_tokens): + return ( + "Schema grounding failed. The generated SQL does not use a " + "verified table or column identifier covering the requested " + "primary subject. Choose a retrieved schema object whose declared " + "identifiers support the subject, or return no SQL if none is " + "available." + ) + if ( + not any( + _schema_tokens_cover(subject_tokens, tokens) + for tokens in identifier_tokens_by_table.values() + ) + and any( + _schema_token_match_count(subject_tokens, tokens) > 0 + for tokens in identifier_tokens_by_table.values() + ) + and _schema_token_match_count(subject_tokens, selected_identifier_tokens) + == 0 + ): + return ( + "Schema grounding failed. The generated SQL does not use a " + "verified table or column identifier covering the requested " + "primary subject. Choose a retrieved schema object whose declared " + "identifiers support the subject, or return no SQL if none is " + "available." + ) + + explicit_grouping_tokens = _grouping_phrase_tokens(query) + group_columns = _group_by_source_columns(sql) + if explicit_grouping_tokens and group_columns: + grouping_required_tokens = _schema_derived_query_tokens( + _query_content_tokens(explicit_grouping_tokens), + schema_details, + ) + if grouping_required_tokens: + grouped_column_tokens: set[str] = set() + grouped_column_names = { + _normalize_identifier(name) for name in group_columns + } + for relation in referenced_relations: + for column in schema_details.get(relation, []): + if _normalize_identifier(column["name"]) in grouped_column_names: + grouped_column_tokens.update(_column_business_tokens(column)) + if not _schema_tokens_cover( + grouping_required_tokens, + grouped_column_tokens, + ): + return ( + "Schema grounding failed. The generated SQL groups by a " + "column that does not cover the requested grouping " + "dimension. Choose a verified grouping column matching the " + "question, or return no SQL if none is available." + ) + + if _is_average_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bAVG\s*\(", sql): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL does not compute AVG over a " + "verified measure." + ) + if re.search(r"(?is)\bCOUNT\s*\(", sql) and not re.search( + r"(?is)\bAVG\s*\(", + sql, + ): + return ( + "Schema grounding failed. The question asks for an average " + "metric, but the generated SQL computes a count." + ) + + if _is_distribution_metric_intent(raw_query_tokens): + if not re.search(r"(?is)\bCOUNT\s*\(", sql) or not re.search( + r"(?is)\bGROUP\s+BY\b", + sql, + ): + return ( + "Schema grounding failed. The question asks for a distribution " + "or breakdown, but the generated SQL does not compute grouped " + "counts." + ) + + if ( + _has_extreme_intent(raw_query_tokens) + and re.search(r"(?is)\bGROUP\s+BY\b", sql) + and _sql_has_aggregate_function(sql) + and not _query_allows_grouped_aggregate(query, raw_query_tokens) + ): + return ( + "Schema grounding failed. The question asks for top or bottom " + "records, but the generated SQL returns a grouped aggregate. Use " + "a row-level ordering unless the question asks for grouping, count, " + "sum, rate, or distribution." + ) + + if ( + _grouping_phrase_tokens(query) + and not _grouping_phrase_has_multiple_dimensions(query) + ): + grouping_error = _validate_explicit_grouping_columns( + sql, + query, + raw_query_tokens, + schema_details, + referenced_relations, + ) + if grouping_error: + return grouping_error + + if _has_missing_value_intent(raw_query_tokens): + target_tokens = _missing_value_target_tokens(query, raw_query_tokens) + for relation in referenced_relations: + columns = schema_details.get(relation) or [] + if not columns or not target_tokens: + continue + checked_columns = _sql_null_checked_columns(sql, columns) + if not checked_columns: + continue + scored_columns = [ + (_column_score_for_tokens(column, target_tokens), column["name"]) + for column in columns + ] + best_score = max((score for score, _ in scored_columns), default=0) + checked_best_score = max( + ( + score + for score, name in scored_columns + if name in checked_columns + ), + default=0, + ) + if best_score > checked_best_score: + return ( + "Schema grounding failed. The question asks for missing " + "values on a specific schema concept, but the generated SQL " + "checks a weaker matching column for null or blank values." + ) + + return _validate_literal_values_against_samples( + sql, + schema_details, + grounding, + query=query, + ) + + +def unsupported_schema_message( + query: str | None, + contexts: list[Any] | None = None, +) -> str | None: + if not query: + return None + query_tokens = _fallback_tokens(query) + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + required_tokens = _schema_required_query_tokens( + query, + query_tokens, + schema_details, + ) + unsupported_tokens = _unsupported_query_tokens( + query_tokens, + schema_details, + query=query, + ) + if unsupported_tokens: + return ( + "No retrieved table or view in the active project contains verified " + "schema metadata for all requested non-operational term(s): " + f"{', '.join(sorted(unsupported_tokens))}. Select a project with " + "matching fields, add schema descriptions/sample values, or ask a " + "question supported by the selected project's schema." + ) + + table_tokens = _schema_tokens_by_table(schema_details) + if required_tokens and any( + _schema_tokens_cover(required_tokens, tokens) + for tokens in table_tokens.values() + ): + return None + + if not required_tokens and not unsupported_tokens: + return None + detail_tokens = sorted(required_tokens or unsupported_tokens) + return ( + "No retrieved table or view in the active project contains verified " + "schema metadata for all requested non-operational term(s): " + f"{', '.join(detail_tokens)}. Select a project with matching fields, " + "add schema descriptions/sample values, or ask a question supported by " + "the selected project's schema." + ) + + +def unsupported_schema_generation_result( + query: str | None, + contexts: list[Any] | None = None, + data_source: str = "", +) -> dict[str, Any] | None: + message = unsupported_schema_message(query, contexts=contexts) + if not message: + return None + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": message, + "correlation_id": "", + "data_source": data_source, + }, + } + + +def schema_grounding_failure_message( + query: str | None, + contexts: list[Any] | None = None, +) -> str: + unsupported_message = unsupported_schema_message(query, contexts=contexts) + if unsupported_message: + return unsupported_message + + schema_details = _extract_schema_details(contexts) + query_tokens = _fallback_tokens(query) + schema_backed_tokens = _schema_required_query_tokens( + query, + query_tokens, + schema_details, + ) + subject_tokens = _query_subject_schema_tokens( + query, + query_tokens, + schema_details, + ) + if schema_details and schema_backed_tokens: + table_tokens = _schema_tokens_by_table(schema_details) + if not any( + all(_fallback_token_variants(token) & tokens for token in schema_backed_tokens) + for tokens in table_tokens.values() + ): + return ( + "No retrieved table or view in the active project contains one " + "verified schema object covering all requested schema-backed " + f"term(s): {', '.join(sorted(schema_backed_tokens))}. The " + "retrieved schema may contain some terms only on different " + "tables or views; add schema metadata/relationships or ask a " + "question supported by one verified schema object." + ) + + content_tokens = sorted(_query_content_tokens(query_tokens)) + if content_tokens: + return ( + "No grounded SQL could be generated using only verified retrieved " + "tables and columns for the active project and requested term(s): " + f"{', '.join(content_tokens)}. Add matching schema metadata/sample " + "values or ask a question supported by the selected project's schema." + ) + + return ( + "No grounded SQL could be generated using only verified retrieved tables " + "and columns for the active project. Ask a question supported by the " + "selected project's schema." + ) + + +def normalize_sql_with_schema_identifiers( + sql: str, + contexts: list[Any] | None = None, +) -> str: + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) + sql = schema_catalog.normalize_sql(sql) + schema_identifiers = set(_extract_schema_identifiers(contexts)) + identifiers = [ + identifier + for identifier in schema_identifiers + if "." not in identifier and _identifier_needs_quotes(identifier) + ] + sql = _replace_bracket_identifiers(sql, schema_identifiers) + for identifier in sorted(identifiers, key=len, reverse=True): + sql = _replace_identifier_outside_literals(sql, identifier) + return sql + + +def _fallback_tokens(value: Any) -> set[str]: + if value is None: + return set() + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + tokens = { + token + for token in _FALLBACK_TOKEN.findall(text.lower()) + if token not in _FALLBACK_STOPWORDS + } + return _expand_fallback_token_variants(tokens) + + +def _identifier_tokens(value: Any) -> set[str]: + tokens = _fallback_tokens(value) + for token in list(tokens): + tokens.update(_compound_identifier_tokens(token)) + return _expand_fallback_token_variants(tokens) + + +def _column_business_tokens(column: dict[str, Any]) -> set[str]: + tokens = _identifier_tokens(column["name"]) + tokens.update(column.get("semantic_tokens") or set()) + return tokens + + +def _table_business_tokens( + table_name: str, + columns: list[dict[str, Any]], +) -> set[str]: + tokens = _identifier_tokens(table_name) + for column in columns: + tokens.update(column.get("_table_semantic_tokens") or set()) + return tokens + + +def _data_type_base(data_type: str) -> str: + return data_type.upper().split("(", 1)[0].strip() + + +def _is_numeric_type(data_type: str) -> bool: + return _data_type_base(data_type) in { + "BIGINT", + "DECIMAL", + "DOUBLE", + "FLOAT", + "FLOAT4", + "FLOAT8", + "INT", + "INT2", + "INT4", + "INT8", + "INTEGER", + "NUMERIC", + "REAL", + "SMALLINT", + } + + +def _is_date_type(data_type: str) -> bool: + return _data_type_base(data_type) in { + "DATE", + "DATETIME", + "DATETIME2", + "SMALLDATETIME", + "TIME", + "TIMESTAMP", + "TIMESTAMPTZ", + "TIMESTAMP_LTZ", + "TIMESTAMP_NTZ", + "TIMESTAMP_TZ", + } + + +_RATE_METRIC_TOKENS = {"rate", "ratio", "percent", "percentage"} +_COUNT_METRIC_TOKENS = {"count", "counts", "many", "most", "number", "total"} +_AVERAGE_METRIC_TOKENS = {"average", "avg", "mean"} +_DISTRIBUTION_METRIC_TOKENS = {"distribution", "breakdown"} +_SUM_METRIC_TOKENS = {"sum", "total"} +_MIN_METRIC_TOKENS = {"bottom", "least", "lowest", "min", "minimum", "smallest"} +_MAX_METRIC_TOKENS = {"greatest", "highest", "largest", "max", "maximum", "most", "top"} +_LATEST_METRIC_TOKENS = {"latest", "newest", "recent"} +_NULL_CHECK_TOKENS = {"blank", "empty", "missing", "null"} +_IMPLICIT_TEXT_VALUE_COLUMN_ROLE_TOKENS = { + "account", + "buyer", + "city", + "client", + "company", + "country", + "cust", + "customer", + "entity", + "label", + "market", + "name", + "org", + "organisation", + "organization", + "party", + "person", + "region", + "seller", + "state", + "supplier", + "title", + "vendor", +} +_IMPLICIT_TEXT_VALUE_IDENTIFIER_TOKENS = { + "code", + "id", + "identifier", + "key", + "no", + "num", + "number", + "po", + "ref", + "reference", + "uuid", +} +_GENERIC_SCHEMA_INTENT_TOKENS = { + "a", + "across", + "all", + "an", + "and", + "as", + "ascending", + "associated", + "association", + "average", + "avg", + "between", + "bottom", + "breakdown", + "bucket", + "buckets", + "by", + "compare", + "count", + "counts", + "date", + "day", + "descending", + "distribution", + "each", + "eight", + "eighteen", + "eleven", + "for", + "five", + "from", + "four", + "fourteen", + "group", + "grouped", + "groups", + "has", + "have", + "highest", + "how", + "in", + "is", + "latest", + "least", + "list", + "lowest", + "many", + "max", + "maximum", + "me", + "mean", + "min", + "minimum", + "month", + "monthly", + "most", + "newest", + "nine", + "nineteen", + "number", + "of", + "one", + "ordered", + "pair", + "pairs", + "per", + "please", + "quarter", + "recent", + "record", + "records", + "result", + "results", + "row", + "rows", + "seven", + "seventeen", + "show", + "six", + "sixteen", + "smallest", + "sort", + "sorted", + "sum", + "ten", + "the", + "there", + "thirteen", + "this", + "three", + "to", + "top", + "total", + "twelve", + "twenty", + "two", + "using", + "was", + "week", + "were", + "what", + "when", + "where", + "which", + "who", + "why", + "with", + "without", + "year", +} +_GENERIC_SCHEMA_INTENT_TOKENS.update(_NULL_CHECK_TOKENS) +_GENERIC_SCHEMA_INTENT_TOKENS.update(_MONTH_NAME_TO_NUMBER.keys()) +_FILTER_VALUE_BOUNDARY_TOKENS = _GENERIC_SCHEMA_INTENT_TOKENS | { + "after", + "before", + "column", + "columns", + "during", + "field", + "fields", + "having", + "limit", + "on", + "since", + "until", + "value", + "values", + "where", +} +_COLUMN_MENTION_STOP_TOKENS = _GENERIC_SCHEMA_INTENT_TOKENS | { + "filter", + "filters", + "or", + "use", + "uses", +} + + +def _is_rate_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _RATE_METRIC_TOKENS) + + +def _is_average_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _AVERAGE_METRIC_TOKENS) + + +def _is_distribution_metric_intent(raw_query_tokens: set[str]) -> bool: + return bool(raw_query_tokens & _DISTRIBUTION_METRIC_TOKENS) + + +def _is_rate_like_column(column: dict[str, str]) -> bool: + return bool(_identifier_tokens(column["name"]) & (_RATE_METRIC_TOKENS | {"score"})) + + +def _is_identifier_like_column(column: dict[str, str]) -> bool: + tokens = _identifier_tokens(column["name"]) + return bool(tokens & {"id", "identifier", "key", "uuid"}) + + +def _quote_joined(identifiers: list[str]) -> str: + return ", ".join(_quote_identifier(identifier) for identifier in identifiers) + + +def _schema_tokens_for_table(table_name: str, columns: list[dict[str, str]]) -> set[str]: + tokens = _table_business_tokens(table_name, columns) + for column in columns: + tokens.update(_column_business_tokens(column)) + tokens.update(_sample_value_tokens(column)) + return tokens + + +def _schema_subject_tokens_for_table( + table_name: str, + columns: list[dict[str, str]], +) -> set[str]: + tokens = _table_business_tokens(table_name, columns) + for column in columns: + tokens.update(_identifier_tokens(column["name"])) + tokens.update(column.get("semantic_tokens") or set()) + return tokens + + +def _schema_identifier_tokens_for_table( + table_name: str, + columns: list[dict[str, str]], +) -> set[str]: + tokens = _identifier_tokens(table_name) + for column in columns: + tokens.update(_identifier_tokens(column["name"])) + return tokens + + +def _schema_tokens_by_table( + schema_details: dict[str, list[dict[str, str]]], +) -> dict[str, set[str]]: + return { + table_name: _schema_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + + +def _schema_subject_tokens_by_table( + schema_details: dict[str, list[dict[str, str]]], +) -> dict[str, set[str]]: + return { + table_name: _schema_subject_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + + +def _schema_token_covered(token: str, schema_tokens: set[str]) -> bool: + return bool(_fallback_token_variants(token) & schema_tokens) + + +def _schema_tokens_cover(required_tokens: set[str], schema_tokens: set[str]) -> bool: + return all(_schema_token_covered(token, schema_tokens) for token in required_tokens) + + +def _schema_token_match_count( + required_tokens: set[str], + schema_tokens: set[str], +) -> int: + return sum(1 for token in required_tokens if _schema_token_covered(token, schema_tokens)) + + +def _schema_derived_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + supported_tokens: set[str] = set() + for token in query_tokens: + if token in _GENERIC_SCHEMA_INTENT_TOKENS or token.isdigit(): + continue + supported_tokens.update(_fallback_token_variants(token) & schema_tokens) + return supported_tokens + + +def _schema_derived_subject_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + schema_tokens_by_table = _schema_subject_tokens_by_table(schema_details) + schema_tokens = set().union(*schema_tokens_by_table.values()) if schema_tokens_by_table else set() + supported_tokens: set[str] = set() + for token in query_tokens: + if token in _GENERIC_SCHEMA_INTENT_TOKENS or token.isdigit(): + continue + supported_tokens.update(_fallback_token_variants(token) & schema_tokens) + return supported_tokens + + +def _schema_required_query_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + required_tokens = _schema_derived_query_tokens(query_tokens, schema_details) + if not required_tokens: + return required_tokens + + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + descriptor_tokens = _schema_adjacent_dimension_descriptor_tokens( + query, + query_tokens, + schema_tokens, + ) + user_value_tokens = _schema_driven_user_value_tokens( + query, + query_tokens, + schema_details, + ) + subject_tokens = _query_subject_schema_tokens( + query, + query_tokens, + schema_details, + ) + return (required_tokens - descriptor_tokens - user_value_tokens) | subject_tokens + + +def _query_subject_schema_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + subject_tokens = _query_subject_content_tokens(query) + if not subject_tokens: + return set() + + return _schema_derived_subject_query_tokens(subject_tokens, schema_details) + + +def _query_subject_unsupported_tokens( + query: str | None, + schema_details: dict[str, list[dict[str, str]]], +) -> set[str]: + subject_tokens = _query_subject_content_tokens(query) + if not subject_tokens: + return set() + + schema_tokens_by_table = _schema_subject_tokens_by_table(schema_details) + schema_tokens = set().union(*schema_tokens_by_table.values()) if schema_tokens_by_table else set() + all_columns = [ + column + for columns in schema_details.values() + for column in columns + ] + subject_value_tokens = _filter_value_tokens( + _schema_driven_column_value_terms( + _query_subject_text(query), + all_columns, + schema_tokens, + ) + ) + subject_tokens = subject_tokens - subject_value_tokens + return { + token + for token in subject_tokens + if not _schema_token_covered(token, schema_tokens) + } + + +def _schema_adjacent_dimension_descriptor_tokens( + query: str | None, + query_tokens: set[str], + schema_tokens: set[str], +) -> set[str]: + if not query: + return set() + dimension_intent_tokens = _DISTRIBUTION_METRIC_TOKENS | { + "across", + "by", + "each", + "group", + "grouped", + "per", + } + if not ( + query_tokens + & dimension_intent_tokens + ): + return set() + + descriptor_tokens: set[str] = set() + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + for index, match in enumerate(matches): + if index == 0: + continue + previous_tokens = _fallback_tokens( + matches[index - 1].group(0).replace("_", " ") + ) + value_tokens = _fallback_tokens(match.group(0).replace("_", " ")) + explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { + "'", + '"', + } + if ( + previous_tokens & schema_tokens + and value_tokens + and not (value_tokens & schema_tokens) + and not explicitly_quoted + ): + descriptor_tokens.update(value_tokens) + return descriptor_tokens + + +def _unsupported_query_tokens( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], + query: str | None = None, +) -> set[str]: + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + descriptor_tokens = _schema_adjacent_dimension_descriptor_tokens( + query, + query_tokens, + schema_tokens, + ) + user_value_tokens = _schema_driven_user_value_tokens( + query, + query_tokens, + schema_details, + ) + unsupported_tokens = { + token + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and not token.isdigit() + and not (_fallback_token_variants(token) & schema_tokens) + and token not in descriptor_tokens + and token not in user_value_tokens + } + unsupported_tokens.update(_query_subject_unsupported_tokens(query, schema_details)) + return unsupported_tokens + + +def _table_covers_requested_concepts( + table_name: str, + columns: list[dict[str, str]], + concept_tokens: set[str], +) -> bool: + required_tokens = { + token + for token in concept_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + } + if not required_tokens: + return True + schema_tokens = _schema_tokens_for_table(table_name, columns) + return _schema_tokens_cover(required_tokens, schema_tokens) + + +def _choose_fallback_table( + query_tokens: set[str], + schema_details: dict[str, list[dict[str, str]]], + concept_tokens: set[str] | None = None, + subject_tokens: set[str] | None = None, +) -> tuple[str, list[dict[str, str]]] | None: + concept_tokens = concept_tokens or query_tokens + subject_tokens = subject_tokens or set() + required_tokens = _schema_derived_query_tokens(concept_tokens, schema_details) + identifier_tokens_by_table = { + table_name: _schema_identifier_tokens_for_table(table_name, columns) + for table_name, columns in schema_details.items() + } + subject_identifier_requires_all = bool(subject_tokens) and any( + _schema_tokens_cover(subject_tokens, tokens) + for tokens in identifier_tokens_by_table.values() + ) + subject_identifier_requires_any = bool(subject_tokens) and not subject_identifier_requires_all and any( + _schema_token_match_count(subject_tokens, tokens) > 0 + for tokens in identifier_tokens_by_table.values() + ) + scored_tables = [] + for table_name, columns in schema_details.items(): + table_tokens = _table_business_tokens(table_name, columns) + table_name_tokens = _identifier_tokens(table_name) + identifier_tokens = identifier_tokens_by_table.get(table_name, set()) + column_token_union: set[str] = set() + sample_token_union: set[str] = set() + numeric_token_matches: set[str] = set() + date_bonus = 0 + table_name_matches = { + variant + for token in query_tokens + for variant in (_fallback_token_variants(token) & table_name_tokens) + } + required_table_name_matches = { + variant + for token in required_tokens + for variant in (_fallback_token_variants(token) & table_name_tokens) + } + score = len(table_name_matches) * 32 + score += len(required_table_name_matches) * 64 + score += len(query_tokens & table_tokens) * 8 + for column in columns: + column_tokens = _column_business_tokens(column) + sample_tokens = _sample_value_tokens(column) + column_token_union.update(column_tokens) + sample_token_union.update(sample_tokens) + if _is_numeric_type(column["data_type"]): + numeric_token_matches.update(query_tokens & column_tokens) + if _is_date_type(column["data_type"]): + date_bonus = 2 + + score += len(query_tokens & column_token_union) * 10 + score += len(query_tokens & sample_token_union) * 6 + score += len(numeric_token_matches) * 2 + score += date_bonus + + table_schema_tokens = _schema_tokens_for_table(table_name, columns) + if required_tokens and not _schema_tokens_cover( + required_tokens, + table_schema_tokens, + ): + continue + if subject_identifier_requires_all and not _schema_tokens_cover( + subject_tokens, + identifier_tokens, + ): + continue + if subject_identifier_requires_any and _schema_token_match_count( + subject_tokens, + identifier_tokens, + ) == 0: + continue + score += _schema_token_match_count(required_tokens, table_schema_tokens) * 20 + score += _schema_token_match_count(subject_tokens, identifier_tokens) * 80 + + if score > 0: + scored_tables.append((score, table_name, columns)) + + if not scored_tables: + return None + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return scored_tables[0][1], scored_tables[0][2] + + +def _choose_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, + date: bool | None = None, +) -> str | None: + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + if numeric is True and not _is_numeric_type(column["data_type"]): + continue + if date is True and not ( + _is_date_type(column["data_type"]) + or column_tokens & {"date", "day", "month", "time", "year"} + ): + continue + score = len(required_tokens & column_tokens) * 10 + if date is True and _is_date_type(column["data_type"]): + score += 20 + if required_tokens and required_tokens.issubset(column_tokens): + score += 30 + if score > 0: + candidates.append((score, column["name"])) + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1])) + return candidates[0][1] + + +def _column_score_for_tokens( + column: dict[str, str], + required_tokens: set[str], + numeric: bool | None = None, +) -> int: + column_tokens = _identifier_tokens(column["name"]) + column_tokens.update(column.get("semantic_tokens") or set()) + if numeric is True and not _is_numeric_type(column["data_type"]): + return 0 + score = len(required_tokens & column_tokens) * 10 + if required_tokens and required_tokens.issubset(column_tokens): + score += 30 + return score + + +def _choose_ranked_column_by_tokens( + columns: list[dict[str, str]], + required_tokens: set[str], + numeric: bool | None = None, +) -> dict[str, str] | None: + candidates = [ + (_column_score_for_tokens(column, required_tokens, numeric=numeric), column) + for column in columns + ] + candidates = [(score, column) for score, column in candidates if score > 0] + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _sample_value_tokens(column: dict[str, Any]) -> set[str]: + tokens: set[str] = set() + for value in column.get("sample_values") or []: + tokens.update(_fallback_tokens(value)) + return tokens + + +def _column_supports_filter_values( + column: dict[str, Any], + values: list[str], +) -> bool: + cleaned_values = [value for value in (_clean_filter_value(value) for value in values) if value] + if not cleaned_values: + return True + + column_tokens = _column_business_tokens(column) + sample_tokens = _sample_value_tokens(column) + value_tokens: set[str] = set() + for value in cleaned_values: + value_tokens.update(_fallback_tokens(value)) + + if value_tokens & sample_tokens: + return True + if value_tokens & column_tokens: + return True + return False + + +def _filter_value_tokens(values: list[str]) -> set[str]: + tokens: set[str] = set() + for value in values: + tokens.update(_fallback_tokens(value)) + return tokens + + +def _first_filter_value_search_token(values: list[str]) -> str | None: + tokens: list[str] = [] + seen: set[str] = set() + for value in values: + for token in _FALLBACK_TOKEN.findall(str(value).lower()): + if ( + token in seen + or token in _FALLBACK_STOPWORDS + or token in _GENERIC_SCHEMA_INTENT_TOKENS + or token.isdigit() + ): + continue + tokens.append(token) + seen.add(token) + + longer_tokens = [token for token in tokens if len(token) >= 4] + if longer_tokens: + return longer_tokens[0] + return tokens[0] if tokens else None + + +def _filter_value_search_pattern(values: list[str]) -> str | None: + cleaned_values = [ + value for value in (_clean_filter_value(value) for value in values) if value + ] + if not cleaned_values: + return None + + primary_value = cleaned_values[0] + primary_tokens = [ + token + for token in _FALLBACK_TOKEN.findall(primary_value.lower()) + if token not in _FALLBACK_STOPWORDS + and token not in _GENERIC_SCHEMA_INTENT_TOKENS + and not token.isdigit() + ] + if re.search(r"[.&/()'-]", primary_value) or len(primary_tokens) >= 2: + phrase = re.sub(r"\s+", " ", primary_value.lower()).strip() + if phrase: + return f"%{phrase}%" + + search_token = _first_filter_value_search_token(values) + return f"%{search_token.lower()}%" if search_token else None + + +def _implicit_text_value_predicate( + column: dict[str, str], + values: list[str], +) -> str: + if column.get("sample_values") and ( + _filter_value_tokens(values) & _sample_value_tokens(column) + ): + return _filter_predicate_for_values(column, values) + + search_pattern = _filter_value_search_pattern(values) + if not search_pattern: + return _filter_predicate_for_values(column, values) + + quoted_column = _quote_identifier(column["name"]) + return f"LOWER({quoted_column}) LIKE {_quote_literal(search_pattern)}" + + +def _filter_predicate_for_values( + column: dict[str, str], + values: list[str], +) -> str: + cleaned_values = [ + value for value in (_clean_filter_value(value) for value in values) if value + ] + if not cleaned_values: + return _non_missing_value_predicate(column) + + column_tokens = _column_business_tokens(column) + value_tokens: set[str] = set() + for value in cleaned_values: + value_tokens.update(_fallback_tokens(value)) + + return _value_match_predicate(column, cleaned_values[0], cleaned_values[1:]) + + +def _choose_filter_column_for_values( + columns: list[dict[str, str]], + values: list[str], +) -> dict[str, str] | None: + candidates: list[tuple[int, dict[str, str]]] = [] + + for column in columns: + column_tokens = _column_business_tokens(column) + if not _column_supports_filter_values(column, values): + continue + value_tokens = { + token + for value in values + for token in _fallback_tokens(value) + } + score = len(column_tokens & value_tokens) * 10 + score += len(_sample_value_tokens(column) & value_tokens) * 20 + if score > 0: + candidates.append((score, column)) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1]["name"])) + return candidates[0][1] + + +def _choose_temporal_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + candidates = [] + for column in columns: + column_tokens = _column_business_tokens(column) + if not _is_date_type(column["data_type"]): + continue + score = len(query_tokens & column_tokens) * 12 + score += 20 + if score > 0: + candidates.append((score, column["name"])) + + if not candidates: + return None + candidates.sort(key=lambda item: (-item[0], item[1])) + return candidates[0][1] + + +def _order_by_phrase_tokens(query: str) -> set[str]: + match = re.search( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+(?P[A-Za-z0-9_ /-]+)", + query, + ) + if not match: + return set() + return _fallback_tokens(match.group("value")) + + +def _choose_order_by_column( + query: str, + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + order_tokens = _order_by_phrase_tokens(query) + if not order_tokens: + return None + + column = _choose_ranked_column_by_tokens(columns, order_tokens) + return column["name"] if column else None + + +def _choose_dimension_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> str | None: + columns = _choose_dimension_columns(query_tokens, columns, max_columns=1) + return columns[0] if columns else None + + +def _choose_dimension_columns( + query_tokens: set[str], + columns: list[dict[str, str]], + max_columns: int = 3, +) -> list[str]: + name_candidates = [] + semantic_candidates = [] + filtered_query_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + identifier_requested = bool(filtered_query_tokens & {"id", "identifier", "key", "uuid"}) + for index, column in enumerate(columns): + if _is_identifier_like_column(column) and not identifier_requested: + continue + name_tokens = _identifier_tokens(column["name"]) + numeric_dimension_tokens = {"month", "period", "quarter", "year"} + if _is_numeric_type(column["data_type"]) and not ( + query_tokens & name_tokens & numeric_dimension_tokens + ): + continue + semantic_tokens = set(column.get("semantic_tokens") or set()) + name_score = len(filtered_query_tokens & name_tokens) * 10 + if query_tokens & name_tokens & numeric_dimension_tokens: + name_score += 10 + semantic_score = len(filtered_query_tokens & semantic_tokens) * 3 + if name_score > 0: + name_candidates.append((name_score + semantic_score, index, column["name"])) + elif semantic_score > 0: + semantic_candidates.append((semantic_score, index, column["name"])) + candidates = name_candidates or semantic_candidates + candidates.sort(key=lambda item: (-item[0], item[1])) + return [name for _, _, name in candidates[:max_columns]] + + +def _choose_missing_value_column( + query: str, + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + target_tokens = _missing_value_target_tokens(query, query_tokens) + column = _choose_ranked_column_by_tokens(columns, target_tokens) + if column: + return column + return _choose_ranked_column_by_tokens( + columns, + query_tokens - (_GENERIC_SCHEMA_INTENT_TOKENS | _NULL_CHECK_TOKENS), + ) + + +def _choose_count_subject_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + column = _choose_ranked_column_by_tokens( + columns, + query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS, + ) + if column: + return column + for column in columns: + if not _is_numeric_type(column["data_type"]) and not _is_rate_like_column(column): + return column + return columns[0] if columns else None + + +def _choose_average_measure_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + filtered_tokens = query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS + measure_candidates = [ + column for column in columns if not _is_identifier_like_column(column) + ] + column = _choose_ranked_column_by_tokens( + measure_candidates, + filtered_tokens, + numeric=True, + ) + if column: + return column + numeric_columns = [ + column + for column in columns + if _is_numeric_type(column["data_type"]) + and not _is_identifier_like_column(column) + ] + return numeric_columns[0] if len(numeric_columns) == 1 else None + + +def _is_text_type(data_type: str) -> bool: + return _data_type_base(data_type) in { + "CHAR", + "CHARACTER", + "NCHAR", + "NTEXT", + "NVARCHAR", + "STRING", + "TEXT", + "VARCHAR", + } + + +def _is_boolean_type(data_type: str) -> bool: + return _data_type_base(data_type) in {"BIT", "BOOL", "BOOLEAN"} + + +def _is_categorical_value_column(column: dict[str, Any]) -> bool: + return not ( + _is_identifier_like_column(column) + or _is_numeric_type(column["data_type"]) + or _is_date_type(column["data_type"]) + or _is_boolean_type(column["data_type"]) + ) + + +def _missing_value_predicate(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_text_type(column["data_type"]): + return f"({quoted_column} IS NULL OR {quoted_column} = '')" + return f"{quoted_column} IS NULL" + + +def _non_missing_value_predicate(column: dict[str, str]) -> str: + quoted_column = _quote_identifier(column["name"]) + if _is_text_type(column["data_type"]): + return f"({quoted_column} IS NOT NULL AND {quoted_column} <> '')" + return f"{quoted_column} IS NOT NULL" + + +def _aggregate_for_measure(measure_column: str) -> tuple[str, str]: + tokens = _identifier_tokens(measure_column) + if tokens & {"rate", "score", "percent", "percentage"}: + return "AVG", "average_value" + return "SUM", "total_value" + + +def _select_listing_columns( + query_tokens: set[str], + columns: list[dict[str, str]], + measure_column: str | None = None, + date_column: str | None = None, + max_columns: int = 6, +) -> list[str]: + scored_columns: list[tuple[int, int, str]] = [] + for index, column in enumerate(columns): + name = column["name"] + tokens = _column_business_tokens(column) + score = len(query_tokens & tokens) * 10 + if name == date_column: + score += 35 + if name == measure_column: + score += 12 + if score > 0: + scored_columns.append((score, index, name)) + + scored_columns.sort(key=lambda item: (-item[0], item[1])) + selected = [] + for _, _, name in scored_columns: + if name not in selected: + selected.append(name) + if len(selected) >= max_columns: + break + if not selected and columns: + selected = [column["name"] for column in columns[:max_columns]] + return selected + + +def _include_required_listing_columns( + selected_columns: list[str], + required_columns: list[str], + max_columns: int, +) -> list[str]: + selected = list(selected_columns) + for column_name in reversed(required_columns): + if column_name in selected: + continue + selected.insert(0, column_name) + return selected[:max_columns] + + +def _fallback_limit(query: str) -> int | None: + match = re.search(r"(?i)\btop\s+(\d+)\b", query) + if match: + return int(match.group(1)) + + word_numbers = { + "one": 1, + "two": 2, + "three": 3, + "four": 4, + "five": 5, + "six": 6, + "seven": 7, + "eight": 8, + "nine": 9, + "ten": 10, + "eleven": 11, + "twelve": 12, + "thirteen": 13, + "fourteen": 14, + "fifteen": 15, + "sixteen": 16, + "seventeen": 17, + "eighteen": 18, + "nineteen": 19, + "twenty": 20, + } + word_match = re.search( + r"(?i)\btop\s+(" + + "|".join(re.escape(word) for word in word_numbers) + + r")\b", + query, + ) + return word_numbers[word_match.group(1).lower()] if word_match else None + + +def _fallback_month_filter(query: str) -> tuple[int, int] | None: + tokens = _fallback_tokens(query) + for month_name, month_number in _MONTH_NAME_TO_NUMBER.items(): + if month_name in tokens: + return datetime.now(timezone.utc).year, month_number + return None + + +def _grouping_phrase_tokens(query: str) -> set[str]: + phrase = _grouping_phrase_text(query) + return _fallback_tokens(phrase) if phrase else set() + + +def _selected_dimensions_cover_grouping_tokens( + dimension_columns: list[str], + grouping_tokens: set[str], + columns: list[dict[str, Any]], + schema_details: dict[str, list[dict[str, str]]], +) -> bool: + grouping_required_tokens = _schema_derived_query_tokens( + _query_content_tokens(grouping_tokens), + schema_details, + ) + if not grouping_required_tokens: + return True + + selected_column_names = { + _normalize_identifier(column_name) for column_name in dimension_columns + } + selected_tokens: set[str] = set() + for column in columns: + if _normalize_identifier(column["name"]) in selected_column_names: + selected_tokens.update(_column_business_tokens(column)) + + return _schema_tokens_cover(grouping_required_tokens, selected_tokens) + + +def _grouping_phrase_has_multiple_dimensions(query: str) -> bool: + phrase = _grouping_phrase_text(query) + return bool(phrase and re.search(r"(?i)(?:,|/|\band\b|\bor\b|\bpairs?\b)", phrase)) + + +def _grouping_phrase_text(query: str) -> str | None: + query = re.sub( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", + "", + query, + ) + match = re.search( + r"""(?ixs) + \b(?:grouped\s+by|group\s+by|by|across|per)\s+ + (?P[A-Za-z0-9_ /-]+?) + (?= + \s+\b(?:after|before|during|for|from|having|in|limit|on|order(?:ed)?\s+by|since|sort(?:ed)?\s+by|where|with)\b(?!-) + |[?.!,;:] + |$ + ) + """, + query, + ) + if not match: + return None + return match.group("value") + + +def _current_year_where_clause( + date_column: str | None, + columns: list[dict[str, str]], + query_tokens: set[str], +) -> str: + if not date_column or not {"this", "year"}.issubset(query_tokens): + return "" + + column = next((column for column in columns if column["name"] == date_column), None) + if not column: + return "" + + current_year = datetime.now(timezone.utc).year + quoted_column = _quote_identifier(date_column) + if _is_numeric_type(column["data_type"]) or _identifier_tokens(date_column) & {"year"}: + return f"\nWHERE {quoted_column} = {current_year}" + if _is_date_type(column["data_type"]) or _identifier_tokens(date_column) & { + "date", + "day", + "month", + "time", + }: + return f"\nWHERE {_date_part_expression(date_column, 'YEAR')} = {current_year}" + return "" + + +def _quote_literal(value: str) -> str: + return "'" + value.replace("'", "''") + "'" + + +def _value_match_predicate( + column: dict[str, str] | str, + value: str, + alternate_values: list[str] | None = None, +) -> str: + column_name = column["name"] if isinstance(column, dict) else column + quoted_column = _quote_identifier(column_name) + values = [] + for candidate in [value] + (alternate_values or []): + cleaned = _clean_filter_value(candidate) + if cleaned and cleaned.lower() not in {item.lower() for item in values}: + values.append(cleaned) + + if not values: + return f"{quoted_column} IS NOT NULL" + + if isinstance(column, dict) and ( + _is_text_type(column["data_type"]) or _is_categorical_value_column(column) + ): + lowered_values = [_quote_literal(candidate.lower()) for candidate in values] + if len(lowered_values) == 1: + return f"LOWER({quoted_column}) = {lowered_values[0]}" + return f"LOWER({quoted_column}) IN ({', '.join(lowered_values)})" + + return f"{quoted_column} = {_quote_literal(values[0])}" + + +def _clean_filter_value(value: str | None) -> str | None: + if value is None: + return None + value = value.strip(" \t\r\n'\"`.,;:()[]{}") + return value or None + + +def _query_content_tokens(query_tokens: set[str]) -> set[str]: + return { + token + for token in query_tokens + if token not in _GENERIC_SCHEMA_INTENT_TOKENS + and token not in _MONTH_NAME_TO_NUMBER + and not token.isdigit() + } + + +def _query_subject_text(query: str | None) -> str: + if not query: + return "" + + return re.split( + r"""(?ix) + \b(?: + group(?:ed)?\s+by + |break(?:down)?\s+by + |for\s+each + |across + |after + |before + |by + |called + |contains + |containing + |during + |for + |from + |having + |in + |like + |matching + |named + |on + |per + |since + |until + |where + |with + )\b + """, + query, + maxsplit=1, + )[0] + + +def _query_subject_content_tokens(query: str | None) -> set[str]: + subject_text = _query_subject_text(query) + return _query_content_tokens(_fallback_tokens(subject_text)) + + +def _has_grouping_intent(query: str, query_tokens: set[str]) -> bool: + query_without_ordering = re.sub( + r"(?is)\b(?:order(?:ed)?|sort(?:ed)?)\s+by\s+[A-Za-z0-9_ /-]+", + "", + query, + ) + return bool( + _is_distribution_metric_intent(query_tokens) + or query_tokens & {"group", "grouped", "per"} + or re.search(r"(?i)\bby\s+[A-Za-z0-9_ -]+\b", query_without_ordering) + ) + + +def _has_count_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _COUNT_METRIC_TOKENS) + + +def _has_sum_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _SUM_METRIC_TOKENS) + + +def _has_extreme_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & (_MAX_METRIC_TOKENS | _MIN_METRIC_TOKENS)) + + +def _has_latest_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _LATEST_METRIC_TOKENS) + + +def _has_missing_value_intent(query_tokens: set[str]) -> bool: + return bool(query_tokens & _NULL_CHECK_TOKENS) + + +def _sort_direction_for_query(query_tokens: set[str]) -> str: + return "ASC" if query_tokens & _MIN_METRIC_TOKENS else "DESC" + + +def _choose_numeric_measure_column( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> dict[str, str] | None: + content_tokens = _query_content_tokens(query_tokens) + measure_candidates = [ + column for column in columns if not _is_identifier_like_column(column) + ] + column = _choose_ranked_column_by_tokens( + measure_candidates, + content_tokens, + numeric=True, + ) + if column: + return column + + numeric_columns = [ + column + for column in columns + if _is_numeric_type(column["data_type"]) + and not _is_identifier_like_column(column) + ] + if len(numeric_columns) == 1: + return numeric_columns[0] + return None + + +def _sample_value_filters( + query_tokens: set[str], + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + filters: list[tuple[dict[str, Any], list[str]]] = [] + consumed_tokens: set[str] = set() + content_tokens = _query_content_tokens(query_tokens) + if not content_tokens: + return filters + + for column in columns: + matches: list[str] = [] + for value in column.get("sample_values") or []: + cleaned_value = _clean_filter_value(str(value)) + if not cleaned_value: + continue + value_tokens = _fallback_tokens(cleaned_value) + if not value_tokens or not value_tokens <= content_tokens: + continue + if value_tokens <= consumed_tokens: + continue + if cleaned_value.lower() not in {item.lower() for item in matches}: + matches.append(cleaned_value) + consumed_tokens.update(value_tokens) + if matches: + filters.append((column, matches)) + + return filters + + +def _mentioned_text_columns_for_query( + query_tokens: set[str], + columns: list[dict[str, Any]], +) -> list[dict[str, Any]]: + mentioned_columns = [] + for column in columns: + if not _is_categorical_value_column(column): + continue + column_tokens = _column_business_tokens(column) + if column_tokens & _IMPLICIT_TEXT_VALUE_IDENTIFIER_TOKENS: + continue + if any(_fallback_token_variants(token) & column_tokens for token in query_tokens): + mentioned_columns.append(column) + return mentioned_columns + + +def _matched_column_query_tokens( + query_tokens: set[str], + column: dict[str, Any], +) -> set[str]: + column_tokens = _column_business_tokens(column) + matched_tokens: set[str] = set() + for token in query_tokens: + matched_tokens.update(_fallback_token_variants(token) & column_tokens) + return matched_tokens + + +def _matched_column_identifier_query_tokens( + query_tokens: set[str], + column: dict[str, Any], +) -> set[str]: + column_tokens = _identifier_tokens(column["name"]) + matched_tokens: set[str] = set() + for token in query_tokens: + matched_tokens.update(_fallback_token_variants(token) & column_tokens) + return matched_tokens + + +def _dedupe_filter_values(values: list[str]) -> list[str]: + deduped: list[str] = [] + seen: set[str] = set() + for value in values: + cleaned = _clean_filter_value(value) + if not cleaned: + continue + key = cleaned.lower() + if key in seen: + continue + deduped.append(cleaned) + seen.add(key) + return deduped + + +def _token_matches_column(token: str, column_tokens: set[str]) -> bool: + meaningful_column_tokens = { + token + for token in column_tokens + if len(token) > 1 and token not in _COLUMN_MENTION_STOP_TOKENS + } + return bool(_fallback_token_variants(token) & meaningful_column_tokens) + + +def _is_filter_value_boundary_token( + token: str, + *, + schema_tokens: set[str], + current_column_tokens: set[str], + allow_schema_token_overlap: bool = False, +) -> bool: + value_tokens = _fallback_tokens(token) + if not value_tokens: + return True + if value_tokens & _FILTER_VALUE_BOUNDARY_TOKENS: + return True + if ( + not allow_schema_token_overlap + and value_tokens & schema_tokens + and not (value_tokens & current_column_tokens) + ): + return True + return False + + +def _schema_driven_column_value_filters( + query: str | None, + columns: list[dict[str, Any]], + schema_tokens: set[str], +) -> list[tuple[dict[str, Any], list[str]]]: + if not query: + return [] + + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + if not matches: + return [] + + filters: list[tuple[dict[str, Any], list[str]]] = [] + + def add_value(column: dict[str, Any], parts: list[str]) -> bool: + if not parts: + return False + cleaned = _clean_filter_value(" ".join(parts)) + if not cleaned: + return False + value_tokens = _fallback_tokens(cleaned) + if not value_tokens or value_tokens & _FILTER_VALUE_BOUNDARY_TOKENS: + return False + if value_tokens <= _column_business_tokens(column): + return False + if all(token.isdigit() for token in value_tokens): + return False + + for existing_column, existing_values in filters: + if existing_column["name"] != column["name"]: + continue + if cleaned.lower() not in {value.lower() for value in existing_values}: + existing_values.append(cleaned) + return True + + filters.append((column, [cleaned])) + return True + + def add_preceding_values( + column: dict[str, Any], + start_index: int, + column_tokens: set[str], + ) -> bool: + added = False + preceding_parts: list[str] = [] + cursor = start_index + while cursor >= 0 and len(preceding_parts) < 4: + candidate = matches[cursor].group(0).replace("_", " ") + candidate_tokens = _fallback_tokens(candidate) + if ( + candidate_tokens + and candidate_tokens <= schema_tokens + and not (candidate_tokens & column_tokens) + ): + break + if _is_filter_value_boundary_token( + candidate, + schema_tokens=schema_tokens, + current_column_tokens=column_tokens, + allow_schema_token_overlap=not preceding_parts, + ): + break + preceding_parts.insert(0, candidate) + cursor -= 1 + + boundary = matches[cursor].group(0).lower() if cursor >= 0 else "" + + if boundary in {"and", "or"}: + alternate_parts: list[str] = [] + cursor -= 1 + while cursor >= 0 and len(alternate_parts) < 4: + candidate = matches[cursor].group(0).replace("_", " ") + candidate_tokens = _fallback_tokens(candidate) + if ( + candidate_tokens + and candidate_tokens <= schema_tokens + and not (candidate_tokens & column_tokens) + ): + break + if _is_filter_value_boundary_token( + candidate, + schema_tokens=schema_tokens, + current_column_tokens=column_tokens, + allow_schema_token_overlap=not alternate_parts, + ): + break + alternate_parts.insert(0, candidate) + cursor -= 1 + added = add_value(column, alternate_parts) or added + + if boundary not in _OPEN_FILTER_VALUE_INTRODUCERS: + added = add_value(column, preceding_parts) or added + + return added + + def add_following_values( + column: dict[str, Any], + start_index: int, + column_tokens: set[str], + *, + allow_initial_schema_overlap: bool = True, + stop_schema_overlap: bool = False, + ) -> None: + following_parts: list[str] = [] + cursor = start_index + while cursor < len(matches) and len(following_parts) < 6: + candidate = matches[cursor].group(0).replace("_", " ") + candidate_tokens = _fallback_tokens(candidate) + if ( + stop_schema_overlap + and + candidate_tokens + and candidate_tokens <= schema_tokens + and not (candidate_tokens & column_tokens) + ): + break + if _is_filter_value_boundary_token( + candidate, + schema_tokens=schema_tokens, + current_column_tokens=column_tokens, + allow_schema_token_overlap=allow_initial_schema_overlap + and not following_parts, + ): + if candidate.lower() in {"and", "or"} and following_parts: + add_value(column, following_parts) + following_parts = [] + cursor += 1 + continue + break + following_parts.append(candidate) + cursor += 1 + add_value(column, following_parts) + + for column in columns: + if ( + _is_numeric_type(column["data_type"]) + or _is_date_type(column["data_type"]) + or _is_boolean_type(column["data_type"]) + or _is_identifier_like_column(column) + ): + continue + + column_tokens = _identifier_tokens(column["name"]) + if not column_tokens: + continue + + for index, match in enumerate(matches): + raw_token = match.group(0).replace("_", " ") + if not _token_matches_column(raw_token, column_tokens): + continue + + next_index = index + 1 + while next_index < len(matches) and _token_matches_column( + matches[next_index].group(0).replace("_", " "), + column_tokens, + ): + next_index += 1 + + previous_index = index - 1 + while previous_index >= 0 and _token_matches_column( + matches[previous_index].group(0).replace("_", " "), + column_tokens, + ): + previous_index -= 1 + + preceding_added = add_preceding_values( + column, + previous_index, + column_tokens, + ) + add_following_values( + column, + next_index, + column_tokens, + allow_initial_schema_overlap=not preceding_added, + stop_schema_overlap=( + previous_index >= 0 + and matches[previous_index].group(0).lower() + in {"across", "by", "per"} + ), + ) + + return filters + + +def _schema_driven_column_value_terms( + query: str | None, + columns: list[dict[str, Any]], + schema_tokens: set[str], +) -> list[str]: + values: list[str] = [] + for _, filter_values in _schema_driven_column_value_filters( + query, + columns, + schema_tokens, + ): + values.extend(filter_values) + return _dedupe_filter_values(values) + + +def _schema_driven_value_terms( + query: str | None, + schema_tokens: set[str], + columns: list[dict[str, Any]], + introducers: set[str] | None = None, +) -> list[str]: + value_terms = _query_schema_value_terms( + query, + schema_tokens, + introducers=introducers, + ) + if introducers is not None: + return _dedupe_filter_values(value_terms) + + column_value_terms = _schema_driven_column_value_terms( + query, + columns, + schema_tokens, + ) + if not column_value_terms: + return _dedupe_filter_values(value_terms) + + column_value_token_sets = [_fallback_tokens(value) for value in column_value_terms] + filtered_terms = [] + for value in value_terms: + value_tokens = _fallback_tokens(value) + if any( + column_tokens + and column_tokens < value_tokens + and value_tokens & schema_tokens + for column_tokens in column_value_token_sets + ): + continue + filtered_terms.append(value) + + return _dedupe_filter_values([*filtered_terms, *column_value_terms]) + + +def _query_schema_value_terms( + query: str | None, + schema_tokens: set[str], + introducers: set[str] | None = None, +) -> list[str]: + if not query: + return [] + + values: list[str] = [] + matches = list(_QUERY_VALUE_TOKEN.finditer(query)) + + def add_value( + raw_value: str, + *, + allow_schema_token_overlap: bool = False, + ) -> None: + value_tokens = _fallback_tokens(raw_value) + if not value_tokens: + return + if value_tokens & _GENERIC_SCHEMA_INTENT_TOKENS: + return + if value_tokens & schema_tokens and ( + not allow_schema_token_overlap or value_tokens <= schema_tokens + ): + return + if all(token.isdigit() for token in value_tokens): + return + cleaned_value = _clean_filter_value(raw_value) + if cleaned_value and cleaned_value.lower() not in { + value.lower() for value in values + }: + values.append(cleaned_value) + + for phrase_match in _FILTER_VALUE_PHRASE.finditer(query): + introducer = phrase_match.group("introducer").lower() + if introducers is not None and introducer not in introducers: + continue + add_value( + phrase_match.group("value"), + allow_schema_token_overlap=introducer in _OPEN_FILTER_VALUE_INTRODUCERS, + ) + + phrase_value_tokens = [_fallback_tokens(value) for value in values] + + def has_filter_value_context(index: int, explicitly_quoted: bool) -> bool: + if explicitly_quoted: + return True + previous_raw = matches[index - 1].group(0).lower() if index > 0 else "" + if introducers is not None and previous_raw not in introducers: + return False + if previous_raw in { + "called", + "contains", + "containing", + "equal", + "equals", + "for", + "from", + "is", + "like", + "matching", + "named", + "where", + "with", + }: + return True + + return False + + for index, match in enumerate(matches): + raw_value = match.group(0).replace("_", " ") + value_tokens = _fallback_tokens(raw_value) + if not value_tokens: + continue + previous_tokens = ( + _fallback_tokens(matches[index - 1].group(0).replace("_", " ")) + if index > 0 + else set() + ) + previous_raw = matches[index - 1].group(0).lower() if index > 0 else "" + explicitly_quoted = match.start() > 0 and query[match.start() - 1] in { + "'", + '"', + } + if not has_filter_value_context(index, explicitly_quoted): + continue + if previous_tokens & schema_tokens and not explicitly_quoted: + continue + if any(value_tokens <= phrase_tokens for phrase_tokens in phrase_value_tokens): + continue + add_value( + raw_value, + allow_schema_token_overlap=previous_raw + in _OPEN_FILTER_VALUE_INTRODUCERS, + ) + + return values + + +def _schema_driven_user_value_tokens( + query: str | None, + query_tokens: set[str], + schema_details: dict[str, list[dict[str, Any]]], +) -> set[str]: + if not query: + return set() + + mentioned_columns: list[dict[str, Any]] = [] + all_columns: list[dict[str, Any]] = [] + for columns in schema_details.values(): + all_columns.extend(columns) + mentioned_columns.extend(_mentioned_text_columns_for_query(query_tokens, columns)) + + matched_concepts = [ + _matched_column_query_tokens(query_tokens, column) + for column in mentioned_columns + ] + matched_concepts = [concepts for concepts in matched_concepts if concepts] + schema_tokens = set().union(*_schema_tokens_by_table(schema_details).values()) + column_value_terms = _schema_driven_column_value_terms( + query, + all_columns, + schema_tokens, + ) + value_terms = _schema_driven_value_terms( + query, + schema_tokens, + all_columns, + ) + open_value_terms = _schema_driven_value_terms( + query, + schema_tokens, + all_columns, + introducers=_OPEN_FILTER_VALUE_INTRODUCERS, + ) + if not matched_concepts: + if value_terms: + value_tokens: set[str] = set() + for value in value_terms: + value_tokens.update(_fallback_tokens(value)) + open_value_tokens: set[str] = set() + for value in open_value_terms: + open_value_tokens.update(_fallback_tokens(value)) + sample_tokens: set[str] = set() + for columns in schema_details.values(): + for column in columns: + sample_tokens.update(_sample_value_tokens(column)) + if sample_tokens and not (value_tokens & sample_tokens) and not open_value_tokens: + return set() + return value_tokens + logger.info( + "Schema-derived user value grounding skipped: no mentioned categorical columns query=%s", + query, + ) + return set() + shared_concepts = set.intersection(*matched_concepts) + if not shared_concepts: + if column_value_terms: + value_tokens: set[str] = set() + for value in column_value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding accepted explicit column values despite ambiguous categorical concepts query=%s values=%s value_tokens=%s columns=%s concepts=%s", + query, + column_value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return value_tokens + if open_value_terms: + value_tokens: set[str] = set() + for value in open_value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding accepted open filter values despite ambiguous categorical concepts query=%s values=%s value_tokens=%s columns=%s concepts=%s", + query, + open_value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return value_tokens + logger.info( + "Schema-derived user value grounding skipped: ambiguous categorical concepts query=%s columns=%s concepts=%s", + query, + [column["name"] for column in mentioned_columns], + [sorted(concepts) for concepts in matched_concepts], + ) + return set() + + value_tokens: set[str] = set() + for value in value_terms: + value_tokens.update(_fallback_tokens(value)) + logger.info( + "Schema-derived user value grounding query=%s shared_concepts=%s values=%s value_tokens=%s columns=%s", + query, + sorted(shared_concepts), + value_terms, + sorted(value_tokens), + [column["name"] for column in mentioned_columns], + ) + return value_tokens + + +def _schema_driven_user_value_filters( + query: str | None, + query_tokens: set[str], + table_name: str, + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + schema_tokens = _schema_tokens_for_table(table_name, columns) + column_value_filters = _schema_driven_column_value_filters( + query, + columns, + schema_tokens, + ) + if column_value_filters: + logger.info( + "Schema-derived explicit column value filters selected query=%s table=%s filters=%s", + query, + table_name, + [ + {"column": column["name"], "values": values} + for column, values in column_value_filters + ], + ) + return column_value_filters + + open_value_terms = _schema_driven_value_terms( + query, + schema_tokens, + columns, + introducers=_OPEN_FILTER_VALUE_INTRODUCERS, + ) + values = _schema_driven_value_terms(query, schema_tokens, columns) + if not values: + return [] + + filter_query_tokens = query_tokens - _filter_value_tokens(values) + if open_value_terms: + filter_query_tokens = filter_query_tokens - ( + _grouping_phrase_tokens(query) if query else set() + ) + mentioned_columns = _mentioned_text_columns_for_query(filter_query_tokens, columns) + if len(mentioned_columns) != 1: + return [] + + return [(mentioned_columns[0], values)] + + +def _implicit_text_value_filter_score( + column: dict[str, Any], + query_tokens: set[str], + table_tokens: set[str], + value_tokens: set[str], + allow_unmatched_samples: bool = False, +) -> int: + if ( + not _is_categorical_value_column(column) + or _is_identifier_like_column(column) + or not value_tokens + ): + return 0 + + column_tokens = _column_business_tokens(column) + if column_tokens & _IMPLICIT_TEXT_VALUE_IDENTIFIER_TOKENS: + return 0 + + sample_tokens = _sample_value_tokens(column) + if sample_tokens and not (value_tokens & sample_tokens) and not allow_unmatched_samples: + return 0 + + filter_context_tokens = ( + query_tokens - _GENERIC_SCHEMA_INTENT_TOKENS - table_tokens - value_tokens + ) + score = len(column_tokens & _IMPLICIT_TEXT_VALUE_COLUMN_ROLE_TOKENS) * 12 + score += len(column_tokens & filter_context_tokens) * 20 + if sample_tokens and value_tokens & sample_tokens: + score += 40 + return score + + +def _schema_driven_implicit_text_value_filters( + query: str | None, + query_tokens: set[str], + table_name: str, + columns: list[dict[str, Any]], +) -> list[tuple[dict[str, Any], list[str]]]: + schema_tokens = _schema_tokens_for_table(table_name, columns) + values = _schema_driven_value_terms(query, schema_tokens, columns) + if not values: + return [] + + value_tokens = _filter_value_tokens(values) + open_value_terms = _schema_driven_value_terms( + query, + schema_tokens, + columns, + introducers=_OPEN_FILTER_VALUE_INTRODUCERS, + ) + allow_unmatched_samples = bool(open_value_terms) + table_tokens = _table_business_tokens(table_name, columns) + scoring_query_tokens = query_tokens - (_grouping_phrase_tokens(query) if query else set()) + candidates: list[tuple[int, str, dict[str, Any]]] = [] + for column in columns: + score = _implicit_text_value_filter_score( + column, + scoring_query_tokens, + table_tokens, + value_tokens, + allow_unmatched_samples=allow_unmatched_samples, + ) + if score > 0: + candidates.append((score, column["name"], column)) + + if not candidates: + logger.info( + "Schema-derived implicit text filter skipped no_candidate query=%s table=%s values=%s", + query, + table_name, + values, + ) + return [] + + candidates.sort(key=lambda item: (-item[0], item[1])) + if len(candidates) > 1 and candidates[0][0] == candidates[1][0]: + top_score = candidates[0][0] + logger.info( + "Schema-derived implicit text filter skipped ambiguous columns query=%s columns=%s", + query, + [column["name"] for score, _, column in candidates if score == top_score], + ) + return [] + logger.info( + "Schema-derived implicit text filter selected query=%s table=%s column=%s score=%s values=%s sample_values_available=%s", + query, + table_name, + candidates[0][2]["name"], + candidates[0][0], + values, + bool(candidates[0][2].get("sample_values")), + ) + return [(candidates[0][2], values)] + + +def _where_clause(predicates: list[str]) -> str: + return f"\nWHERE {' AND '.join(predicates)}" if predicates else "" + + +def _count_expression_for_query( + query_tokens: set[str], + columns: list[dict[str, str]], +) -> tuple[str, list[str]]: + subject_column = _choose_count_subject_column(query_tokens, columns) + if not subject_column: + return "COUNT(*)", [] + return ( + f"COUNT({_quote_identifier(subject_column['name'])})", + [_non_missing_value_predicate(subject_column)], + ) + + +def _date_bucket_expressions(date_column: str) -> tuple[str, str]: + return ( + _date_part_expression(date_column, "YEAR"), + _date_part_expression(date_column, "MONTH"), + ) + + +def _date_part_expression(date_column: str, part: str) -> str: + return f"CAST(EXTRACT({part} FROM {_quote_identifier(date_column)}) AS BIGINT)" + + +def generate_simple_analytics_sql( + query: str | None, + contexts: list[Any] | None, +) -> str | None: + if not query: + return None + + query_tokens = _fallback_tokens(query) + if not query_tokens: + return None + + schema_details = _extract_schema_details(contexts) + if not schema_details: + return None + + content_tokens = _query_content_tokens(query_tokens) + schema_backed_tokens = _schema_required_query_tokens( + query, + query_tokens, + schema_details, + ) + subject_tokens = _query_subject_schema_tokens( + query, + query_tokens, + schema_details, + ) + unsupported_tokens = _unsupported_query_tokens( + query_tokens, + schema_details, + query=query, + ) + if unsupported_tokens: + logger.info( + "Schema-derived SQL fallback skipped unsupported_tokens=%s", + sorted(unsupported_tokens), + ) + return None + if content_tokens and not schema_backed_tokens: + logger.info( + "Schema-derived SQL fallback skipped no_schema_backed_tokens=%s", + sorted(content_tokens), + ) + return None + if not content_tokens and len(schema_details) != 1: + logger.info( + "Schema-derived SQL fallback skipped ambiguous_schema_only_request tables=%s", + sorted(schema_details), + ) + return None + + chosen = _choose_fallback_table( + query_tokens, + schema_details, + concept_tokens=schema_backed_tokens or query_tokens, + subject_tokens=subject_tokens, + ) + if not chosen: + return None + + table_name, columns = chosen + column_names = [column["name"] for column in columns] + quoted_table = _quote_identifier(table_name) + limit = _fallback_limit(query) + date_column = _choose_temporal_column(query_tokens, columns) + order_column = _choose_order_by_column(query, query_tokens, columns) + sample_filters = _sample_value_filters(query_tokens, columns) + sample_filter_column_names = {column["name"] for column, _ in sample_filters} + user_value_filters = [ + (column, values) + for column, values in _schema_driven_user_value_filters( + query, + query_tokens, + table_name, + columns, + ) + if column["name"] not in sample_filter_column_names + ] + implicit_value_filters = [] + if not sample_filters and not user_value_filters: + implicit_value_filters = _schema_driven_implicit_text_value_filters( + query, + query_tokens, + table_name, + columns, + ) + value_filters = [*sample_filters, *user_value_filters] + sample_predicates = [ + _filter_predicate_for_values(column, values) + for column, values in value_filters + ] + sample_predicates.extend( + _implicit_text_value_predicate(column, values) + for column, values in implicit_value_filters + ) + selected_sample_filter_columns = [ + column["name"] for column, _ in [*value_filters, *implicit_value_filters] + ] + filter_column_names = selected_sample_filter_columns + month_filter = _fallback_month_filter(query) + metric_intent = { + "average": _is_average_metric_intent(query_tokens), + "count": _has_count_intent(query_tokens), + "distribution": _is_distribution_metric_intent(query_tokens), + "extreme": _has_extreme_intent(query_tokens), + "latest": _has_latest_intent(query_tokens), + "missing": _has_missing_value_intent(query_tokens), + "rate": _is_rate_metric_intent(query_tokens), + "sum": _has_sum_intent(query_tokens), + } + logger.info( + "Schema-derived SQL fallback selected table=%s schema_tokens=%s verified_columns=%s sample_filter_columns=%s metric_intent=%s", + table_name, + sorted(schema_backed_tokens), + column_names, + selected_sample_filter_columns, + metric_intent, + ) + + if _has_missing_value_intent(query_tokens): + missing_column = _choose_missing_value_column(query, query_tokens, columns) + if not missing_column: + return None + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + selected_columns = _include_required_listing_columns( + selected_columns, + [missing_column["name"], *filter_column_names], + max_columns=8, + ) + predicates = [*sample_predicates, _missing_value_predicate(missing_column)] + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" + ) + + month_predicate = "" + if date_column and month_filter: + year, month = month_filter + start_date = f"{year:04d}-{month:02d}-01" + end_year = year + 1 if month == 12 else year + end_month = 1 if month == 12 else month + 1 + end_date = f"{end_year:04d}-{end_month:02d}-01" + quoted_date = _quote_identifier(date_column) + month_predicate = ( + f"{quoted_date} >= '{start_date}' AND {quoted_date} < '{end_date}'" + ) + + predicates = list(sample_predicates) + if month_predicate: + predicates.append(month_predicate) + + if _is_average_metric_intent(query_tokens): + measure_column = _choose_average_measure_column(query_tokens, columns) + if not measure_column: + return None + grouping_tokens = _grouping_phrase_tokens(query) or query_tokens + dimension_columns = _choose_dimension_columns( + grouping_tokens, + columns, + max_columns=2, + ) + aggregate_expr = f"AVG({_quote_identifier(measure_column['name'])})" + if dimension_columns: + quoted_dimensions = _quote_joined(dimension_columns) + return ( + f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier('average_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('average_value')} DESC" + ) + return ( + f"SELECT {aggregate_expr} AS {_quote_identifier('average_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" + ) + + measure_column = _choose_numeric_measure_column(query_tokens, columns) + if ( + date_column + and query_tokens & {"month", "monthly"} + and not ( + _has_sum_intent(query_tokens) + or _is_average_metric_intent(query_tokens) + or _is_rate_metric_intent(query_tokens) + or _has_extreme_intent(query_tokens) + ) + ): + year_expr, month_expr = _date_bucket_expressions(date_column) + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) + + if ( + measure_column + and date_column + and query_tokens & {"month", "monthly"} + and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) + ): + year_expr, month_expr = _date_bucket_expressions(date_column) + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{month_expr} AS {_quote_identifier('month')}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}, {month_expr}\n" + f"ORDER BY {year_expr}, {month_expr}" + ) + + if ( + measure_column + and date_column + and "year" in query_tokens + and "month" not in query_tokens + and "monthly" not in query_tokens + and (_has_sum_intent(query_tokens) or _has_grouping_intent(query, query_tokens)) + ): + year_expr = _date_part_expression(date_column, "YEAR") + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {year_expr} AS {_quote_identifier('year')}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(sample_predicates)}\n" + f"GROUP BY {year_expr}\nORDER BY {year_expr}" + ) + + explicit_grouping_tokens = _grouping_phrase_tokens(query) + if ( + _has_extreme_intent(query_tokens) + and measure_column + and explicit_grouping_tokens + and _matched_column_query_tokens(explicit_grouping_tokens, measure_column) + ): + direction = _sort_direction_for_query(query_tokens) + selected_columns = _select_listing_columns( + query_tokens, + columns, + measure_column=measure_column["name"], + date_column=date_column, + max_columns=8, + ) + if measure_column["name"] not in selected_columns: + selected_columns.insert(0, measure_column["name"]) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(measure_column['name'])} {direction}{limit_clause}" + ) + + if _has_count_intent(query_tokens) and not _has_grouping_intent(query, query_tokens): + return ( + f"SELECT COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" + ) + + if _has_latest_intent(query_tokens): + if not date_column: + return None + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if date_column not in selected_columns: + selected_columns.insert(0, date_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + ) + + if _has_grouping_intent(query, query_tokens): + grouping_tokens = explicit_grouping_tokens or query_tokens + max_dimensions = 1 if _has_extreme_intent(query_tokens) else 3 + if explicit_grouping_tokens and not _grouping_phrase_has_multiple_dimensions(query): + max_dimensions = 1 + dimension_columns = _choose_dimension_columns( + grouping_tokens, + columns, + max_columns=max_dimensions, + ) + if explicit_grouping_tokens and dimension_columns: + if not _selected_dimensions_cover_grouping_tokens( + dimension_columns, + explicit_grouping_tokens, + columns, + schema_details, + ): + return None + if dimension_columns: + quoted_dimensions = _quote_joined(dimension_columns) + if measure_column and (_has_sum_intent(query_tokens) or _is_rate_metric_intent(query_tokens)): + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + aggregate_expr = f"{aggregate}({_quote_identifier(measure_column['name'])})" + direction = _sort_direction_for_query(query_tokens) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimensions}, {aggregate_expr} AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {quoted_dimensions}, COUNT(*) AS {_quote_identifier('record_count')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {quoted_dimensions}\n" + f"ORDER BY {_quote_identifier('record_count')} DESC{limit_clause}" + ) + return None + + if measure_column and _has_sum_intent(query_tokens): + return ( + f"SELECT SUM({_quote_identifier(measure_column['name'])}) AS {_quote_identifier('total_value')}\n" + f"FROM {quoted_table}{_where_clause(predicates)}" + ) + + if _has_extreme_intent(query_tokens): + direction = _sort_direction_for_query(query_tokens) + measure_column_is_grounded = bool( + measure_column + and _matched_column_identifier_query_tokens( + content_tokens | explicit_grouping_tokens, + measure_column, + ) + ) + if measure_column and measure_column_is_grounded: + limit_clause = f"\nLIMIT {limit}" if limit else "" + if _query_allows_grouped_aggregate(query, query_tokens): + dimension_column = _choose_dimension_column(query_tokens, columns) + else: + dimension_column = None + if dimension_column: + aggregate, alias = _aggregate_for_measure(measure_column["name"]) + return ( + f"SELECT {_quote_identifier(dimension_column)}, " + f"{aggregate}({_quote_identifier(measure_column['name'])}) AS {_quote_identifier(alias)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"GROUP BY {_quote_identifier(dimension_column)}\n" + f"ORDER BY {_quote_identifier(alias)} {direction}{limit_clause}" + ) + selected_columns = _select_listing_columns( + query_tokens, + columns, + measure_column=measure_column["name"], + date_column=date_column, + max_columns=8, + ) + if measure_column["name"] not in selected_columns: + selected_columns.insert(0, measure_column["name"]) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(measure_column['name'])} {direction}{limit_clause}" + ) + if date_column and (month_predicate or query_tokens & {"record", "records", "row", "rows"}): + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if date_column not in selected_columns: + selected_columns.insert(0, date_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + ) + if order_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if order_column not in selected_columns: + selected_columns.insert(0, order_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(order_column)} {direction}{limit_clause}" + ) + return None + + if month_predicate and date_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + measure_column=measure_column["name"] if measure_column else None, + date_column=date_column, + max_columns=8, + ) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(date_column)} DESC{limit_clause}" + ) + + if order_column: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + if order_column not in selected_columns: + selected_columns.insert(0, order_column) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}\n" + f"ORDER BY {_quote_identifier(order_column)} ASC{limit_clause}" + ) + + if sample_predicates or query_tokens & {"all", "list"}: + selected_columns = _select_listing_columns( + query_tokens, + columns, + date_column=date_column, + max_columns=8, + ) + selected_columns = _include_required_listing_columns( + selected_columns, + filter_column_names, + max_columns=8, + ) + limit_clause = f"\nLIMIT {limit}" if limit else "" + return ( + f"SELECT {_quote_joined(selected_columns)}\n" + f"FROM {quoted_table}{_where_clause(predicates)}{limit_clause}" + ) + + return None + @component class SQLGenPostProcessor: def __init__(self, engine: Engine): self._engine = engine + @staticmethod + def _log_timing( + stage: str, + started_at: float, + project_id: str | None = None, + **fields: Any, + ) -> None: + suffix = " ".join( + f"{key}={value}" for key, value in fields.items() if value is not None + ) + logger.info( + "Ask timing stage=%s project_id=%s elapsed_ms=%.1f%s%s", + stage, + project_id or "", + _timing_ms(started_at), + " " if suffix else "", + suffix, + ) + @component.output_types( valid_generation_result=Dict[str, Any], invalid_generation_result=Dict[str, Any], @@ -30,31 +5040,248 @@ async def run( self, replies: List[str] | List[List[str]], project_id: str | None = None, + mdl_hash: str | None = None, + contexts: list[Any] | None = None, + fallback_query: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", allow_data_preview: bool = False, ) -> dict: try: - cleaned_generation_result = clean_generation_result(replies[0]) + total_started_at = time.perf_counter() + extraction_started_at = time.perf_counter() + cleaned_generation_result, extraction_error = _extract_sql_response( + clean_generation_result(replies[0]) + ) + self._log_timing( + "sql_response_extraction", + extraction_started_at, + project_id, + sql_present=bool(cleaned_generation_result), + ) + grounding_invalid_generation_result = None + + def validate_candidate_sql(candidate_sql: str) -> str | None: + validation_started_at = time.perf_counter() + schema_catalog = _SchemaCatalog.from_contexts(contexts or []) + grounding_error = schema_catalog.validate_sql(candidate_sql) + if not grounding_error: + grounding_error = validate_sql_against_contexts( + candidate_sql, + contexts=contexts, + ) + if not grounding_error: + grounding_error = validate_sql_semantic_coverage( + candidate_sql, + fallback_query, + contexts=contexts, + ) + self._log_timing( + "sql_validation", + validation_started_at, + project_id, + status="rejected" if grounding_error else "grounded", + ) + return grounding_error + + if cleaned_generation_result: + cleaned_generation_result = normalize_sql_with_schema_identifiers( + cleaned_generation_result, + contexts=contexts, + ) + cleaned_generation_result = normalize_wren_sql_dialect( + cleaned_generation_result + ) + grounding_error = validate_candidate_sql(cleaned_generation_result) + if grounding_error: + logger.info( + "Generated SQL validation result project_id=%s status=rejected reason=%s sql=%s", + project_id or "", + grounding_error, + cleaned_generation_result, + ) + grounding_invalid_generation_result = { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": grounding_error, + "correlation_id": "", + "data_source": data_source, + } + else: + logger.info( + "Generated SQL validation result project_id=%s status=grounded sql=%s", + project_id or "", + cleaned_generation_result, + ) + elif extraction_error: + logger.info( + "Generated SQL extraction result project_id=%s status=rejected reason=%s", + project_id or "", + extraction_error, + ) + + fallback_started_at = time.perf_counter() + fallback_generation_result = generate_simple_analytics_sql( + fallback_query, + contexts, + ) + self._log_timing( + "schema_fast_path_generation", + fallback_started_at, + project_id, + sql_present=bool(fallback_generation_result), + ) + if fallback_generation_result: + logger.info( + "Deterministic SQL fallback generated project_id=%s sql=%s", + project_id or "", + fallback_generation_result, + ) + fallback_generation_result = normalize_sql_with_schema_identifiers( + fallback_generation_result, + contexts=contexts, + ) + fallback_generation_result = normalize_wren_sql_dialect( + fallback_generation_result + ) + fallback_grounding_error = validate_candidate_sql( + fallback_generation_result + ) + logger.info( + "Deterministic SQL fallback validation result project_id=%s status=%s%s", + project_id or "", + "grounded" if not fallback_grounding_error else "rejected", + "" + if not fallback_grounding_error + else f" reason={fallback_grounding_error}", + ) + if not fallback_grounding_error: + engine_validation_started_at = time.perf_counter() + ( + fallback_valid_generation_result, + fallback_invalid_generation_result, + ) = await self._classify_generation_result( + fallback_generation_result, + project_id=project_id, + mdl_hash=mdl_hash, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + data_source=data_source, + allow_data_preview=allow_data_preview, + ) + self._log_timing( + "sql_engine_validation", + engine_validation_started_at, + project_id, + status="valid" + if fallback_valid_generation_result + else "invalid", + ) + if fallback_valid_generation_result: + logger.info( + "Using deterministic schema-grounded SQL fallback for query." + ) + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="valid", + ) + return { + "valid_generation_result": fallback_valid_generation_result, + "invalid_generation_result": {}, + } + logger.info( + "Deterministic SQL fallback did not validate: %s", + fallback_invalid_generation_result.get("error"), + ) + + if grounding_invalid_generation_result: + unsupported_message = schema_grounding_failure_message( + fallback_query, + contexts=contexts, + ) + logger.info( + "Generated SQL grounding rejection converted to unsupported schema project_id=%s reason=%s original_reason=%s", + project_id or "", + unsupported_message, + grounding_invalid_generation_result.get("error"), + ) + grounding_invalid_generation_result["type"] = "NO_RELEVANT_SQL" + grounding_invalid_generation_result["error"] = unsupported_message + grounding_invalid_generation_result["sql"] = "" + grounding_invalid_generation_result["original_sql"] = "" + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="invalid", + ) + return { + "valid_generation_result": {}, + "invalid_generation_result": grounding_invalid_generation_result, + } + + unsupported_result = unsupported_schema_generation_result( + fallback_query, + contexts=contexts, + data_source=data_source, + ) + if not cleaned_generation_result and unsupported_result: + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="unsupported", + ) + return unsupported_result - # test if cleaned_generation_result in string format is actually a dictionary with key 'sql' - if cleaned_generation_result.startswith("{"): - cleaned_generation_result = orjson.loads(cleaned_generation_result)[ - "sql" - ] + if not cleaned_generation_result and extraction_error: + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="invalid", + ) + return { + "valid_generation_result": {}, + "invalid_generation_result": { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": extraction_error, + "correlation_id": "", + "data_source": data_source, + }, + } + engine_validation_started_at = time.perf_counter() ( valid_generation_result, invalid_generation_result, ) = await self._classify_generation_result( cleaned_generation_result, project_id=project_id, + mdl_hash=mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, data_source=data_source, allow_data_preview=allow_data_preview, ) + self._log_timing( + "sql_engine_validation", + engine_validation_started_at, + project_id, + status="valid" if valid_generation_result else "invalid", + ) + self._log_timing( + "sql_post_process_total", + total_started_at, + project_id, + status="valid" if valid_generation_result else "invalid", + ) return { "valid_generation_result": valid_generation_result, @@ -70,8 +5297,9 @@ async def run( async def _classify_generation_result( self, - generation_result: str, + generation_result: str | None, project_id: str | None = None, + mdl_hash: str | None = None, use_dry_plan: bool = False, allow_dry_plan_fallback: bool = True, data_source: str = "", @@ -81,12 +5309,23 @@ async def _classify_generation_result( invalid_generation_result = {} use_dry_run = not allow_data_preview + if not generation_result: + return valid_generation_result, { + "sql": "", + "original_sql": "", + "type": "NO_RELEVANT_SQL", + "error": "No grounded SQL was generated from the current schema.", + "correlation_id": "", + } + async with aiohttp.ClientSession() as session: if use_dry_plan: dry_plan_result, error_message = await self._engine.dry_plan( session, generation_result, data_source, + project_id=project_id, + mdl_hash=mdl_hash, allow_fallback=allow_dry_plan_fallback, ) @@ -98,6 +5337,7 @@ async def _classify_generation_result( else: invalid_generation_result = { "sql": generation_result, + "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") else "DRY_PLAN", @@ -105,10 +5345,32 @@ async def _classify_generation_result( "correlation_id": "", } elif use_dry_run: + dry_plan_result, error_message = await self._engine.dry_plan( + session, + generation_result, + data_source, + project_id=project_id, + mdl_hash=mdl_hash, + allow_fallback=allow_dry_plan_fallback, + ) + + if not dry_plan_result: + invalid_generation_result = { + "sql": generation_result, + "original_sql": generation_result, + "type": "TIME_OUT" + if error_message.startswith("Request timed out") + else "DRY_PLAN", + "error": error_message, + "correlation_id": "", + } + return valid_generation_result, invalid_generation_result + success, _, addition = await self._engine.execute_sql( generation_result, session, project_id=project_id, + mdl_hash=mdl_hash, limit=1, dry_run=True, ) @@ -121,7 +5383,7 @@ async def _classify_generation_result( else: error_message = addition.get("error_message", "") invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": generation_result, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -134,6 +5396,7 @@ async def _classify_generation_result( generation_result, session, project_id=project_id, + mdl_hash=mdl_hash, limit=1, dry_run=False, ) @@ -151,7 +5414,7 @@ async def _classify_generation_result( else "PREVIEW_FAILED" ) invalid_generation_result = { - "sql": addition.get("error_sql", generation_result), + "sql": generation_result, "original_sql": generation_result, "type": "TIME_OUT" if error_message.startswith("Request timed out") @@ -163,6 +5426,820 @@ async def _classify_generation_result( return valid_generation_result, invalid_generation_result +_SCHEMA_CATALOG_CACHE: dict[tuple[str, ...], "_SchemaCatalog"] = {} + + +class _SchemaCatalog: + def __init__( + self, + tables: dict[str, set[str]], + relationships: dict[str, set[str]] | None = None, + table_aliases: dict[str, str] | None = None, + table_alias_values: dict[str, str] | None = None, + column_aliases: dict[str, dict[str, str]] | None = None, + column_alias_values: dict[str, dict[str, str]] | None = None, + ): + self._tables = tables + self._relationships = relationships or {} + self._table_aliases = table_aliases or {} + self._table_alias_values = table_alias_values or {} + self._column_aliases = column_aliases or {} + self._column_alias_values = column_alias_values or {} + + @classmethod + def from_contexts(cls, contexts: list[str]) -> "_SchemaCatalog": + context_cache_key = _context_cache_key(contexts) + cached_catalog = _SCHEMA_CATALOG_CACHE.get(context_cache_key) + if cached_catalog is not None: + return cached_catalog + + tables: dict[str, set[str]] = {} + relationships: dict[str, set[str]] = {} + table_aliases: dict[str, str] = {} + table_alias_values: dict[str, str] = {} + column_aliases: dict[str, dict[str, str]] = {} + column_alias_values: dict[str, dict[str, str]] = {} + + for context in contexts: + context = str(context) + cls._add_contract_identifiers(context, tables, relationships) + cls._add_ddl_identifiers(context, tables) + cls._add_semantic_aliases( + context, + table_aliases, + table_alias_values, + column_aliases, + column_alias_values, + ) + + for table_name, column_names in tables.items(): + for alias in _identifier_dot_aliases(table_name): + _add_unique_identifier_alias( + table_aliases, + table_alias_values, + alias, + table_name, + ) + for column_name in column_names: + column_aliases.setdefault(table_name, {}) + column_alias_values.setdefault(table_name, {}) + + catalog = cls( + tables, + relationships, + table_aliases, + table_alias_values, + column_aliases, + column_alias_values, + ) + if len(_SCHEMA_CATALOG_CACHE) >= 256: + _SCHEMA_CATALOG_CACHE.pop(next(iter(_SCHEMA_CATALOG_CACHE))) + _SCHEMA_CATALOG_CACHE[context_cache_key] = catalog + return catalog + + @staticmethod + def _context_table_name(context: str) -> str | None: + payload = _extract_semantic_context_payload(context) + contract = payload.get("sql_identifier_contract") if payload else {} + if isinstance(contract, dict): + table_name = contract.get("sql_table_name_use_exactly") + if isinstance(table_name, str) and table_name: + return _clean_contract_value(table_name) + + for raw_line in context.splitlines(): + line = raw_line.strip() + if line.startswith("sql_table_name_use_exactly:"): + table_name = _clean_contract_value(line.split(":", 1)[1]) + if table_name: + return table_name + + match = _DDL_RELATION.search(context) + if match: + return _normalize_identifier(match.group("name")) + return None + + @staticmethod + def _add_semantic_aliases( + context: str, + table_aliases: dict[str, str], + table_alias_values: dict[str, str], + column_aliases: dict[str, dict[str, str]], + column_alias_values: dict[str, dict[str, str]], + ) -> None: + table_name = _SchemaCatalog._context_table_name(context) + payload = _extract_semantic_context_payload(context) + if not table_name or not payload: + return + + table_semantic = payload.get("semantic_context_not_sql_identifiers") + for alias in _semantic_identifier_aliases(table_semantic): + _add_unique_identifier_alias( + table_aliases, + table_alias_values, + alias, + table_name, + ) + + column_aliases.setdefault(table_name, {}) + column_alias_values.setdefault(table_name, {}) + for column in payload.get("columns", []) or []: + if not isinstance(column, dict): + continue + column_name = column.get("sql_column_name_use_exactly") + if not isinstance(column_name, str) or not column_name: + continue + for alias in _semantic_identifier_aliases(column): + _add_unique_identifier_alias( + column_aliases[table_name], + column_alias_values[table_name], + alias, + column_name, + ) + + @staticmethod + def _add_contract_identifiers( + context: str, + tables: dict[str, set[str]], + relationships: dict[str, set[str]], + ) -> None: + current_table: str | None = None + current_section: str | None = None + + for raw_line in context.splitlines(): + line = raw_line.strip() + if line.startswith("table: "): + current_table = _clean_identifier(line.removeprefix("table: ")) + if current_table: + tables.setdefault(current_table, set()) + relationships.setdefault(current_table, set()) + current_section = None + continue + + if line.startswith("sql_table_name_use_exactly:"): + current_table = _clean_identifier( + line.removeprefix("sql_table_name_use_exactly:") + ) + if current_table: + tables.setdefault(current_table, set()) + relationships.setdefault(current_table, set()) + current_section = None + continue + + if current_table and line in {"columns:", "sql_column_names_use_exactly:"}: + current_section = "columns" + continue + + if current_table and line in { + "relationships:", + "relationship_constraints_use_exactly:", + }: + current_section = "relationships" + continue + + if current_section and current_table and line.startswith("- "): + value = line.removeprefix("- ").strip() + if not value: + continue + if current_section == "columns": + column_name = _clean_identifier(value) + if column_name: + tables.setdefault(current_table, set()).add(column_name) + elif current_section == "relationships": + relationships.setdefault(current_table, set()).add(value) + continue + + if current_section and line and not line.startswith("- "): + current_section = None + + @staticmethod + def _add_ddl_identifiers( + context: str, + tables: dict[str, set[str]], + ) -> None: + for match in _DDL_CREATE_PATTERN.finditer(context): + table_name = _clean_identifier(match.group("table")) + if not table_name: + continue + + body_start = match.end() + body_end = _find_matching_parenthesis(context, body_start) + if body_end is None: + tables.setdefault(table_name, set()) + continue + + tables.setdefault(table_name, set()).update( + _extract_ddl_column_names(context[body_start:body_end]) + ) + + def normalize_sql(self, sql: str) -> str: + if not sql: + return sql + sql = self._normalize_table_aliases(sql) + return self._normalize_column_aliases(sql) + + def _normalize_table_aliases(self, sql: str) -> str: + sorted_aliases = sorted( + self._table_aliases.items(), + key=lambda item: len(self._table_alias_values.get(item[0], item[0])), + reverse=True, + ) + for alias_key, table_name in sorted_aliases: + if not table_name: + continue + alias = self._table_alias_values.get(alias_key, alias_key) + replacement = _render_table_identifier(table_name) + for variant in sorted( + _identifier_reference_variants(alias), + key=len, + reverse=True, + ): + sql = _replace_sql_text_outside_literals( + sql, + variant, + replacement, + case_insensitive=True, + ) + return sql + + def _normalize_column_aliases(self, sql: str) -> str: + grounding = _extract_sql_grounding(sql) + real_relations = [ + relation + for relation in grounding["relation_references"] + if relation not in grounding["cte_names"] and relation in self._tables + ] + if not real_relations: + return sql + + qualifiers_by_relation: dict[str, set[str]] = {} + for qualifier, relation in grounding["alias_to_relation"].items(): + if relation in self._tables: + qualifiers_by_relation.setdefault(relation, set()).add(qualifier) + + for relation in real_relations: + sql = self._normalize_qualified_column_aliases( + sql, + relation, + qualifiers_by_relation.get(relation, set()), + ) + + unique_real_relations = list(dict.fromkeys(real_relations)) + if len(unique_real_relations) == 1: + sql = self._normalize_unqualified_column_aliases( + sql, + unique_real_relations[0], + ) + return sql + + def _normalize_qualified_column_aliases( + self, + sql: str, + relation: str, + qualifiers: set[str], + ) -> str: + aliases = self._column_aliases.get(relation, {}) + alias_values = self._column_alias_values.get(relation, {}) + for alias_key, column_name in aliases.items(): + if not column_name: + continue + alias = alias_values.get(alias_key, alias_key) + column_variants = sorted( + _identifier_reference_variants(alias), + key=len, + reverse=True, + ) + for qualifier in qualifiers: + qualifier_variants = sorted( + _identifier_reference_variants(qualifier), + key=len, + reverse=True, + ) + qualifier_replacement = _render_qualifier_identifier( + qualifier, + set(self._tables), + ) + replacement = f"{qualifier_replacement}.{_quote_identifier(column_name)}" + for qualifier_variant in qualifier_variants: + for column_variant in column_variants: + sql = _replace_sql_text_outside_literals( + sql, + f"{qualifier_variant}.{column_variant}", + replacement, + case_insensitive=True, + ) + return sql + + def _normalize_unqualified_column_aliases(self, sql: str, relation: str) -> str: + aliases = self._column_aliases.get(relation, {}) + alias_values = self._column_alias_values.get(relation, {}) + sorted_aliases = sorted( + aliases.items(), + key=lambda item: len(alias_values.get(item[0], item[0])), + reverse=True, + ) + for alias_key, column_name in sorted_aliases: + if not column_name: + continue + alias = alias_values.get(alias_key, alias_key) + replacement = _quote_identifier(column_name) + for variant in sorted( + _identifier_reference_variants(alias), + key=len, + reverse=True, + ): + sql = _replace_sql_text_outside_literals( + sql, + variant, + replacement, + case_insensitive=True, + ) + return sql + + def to_prompt(self) -> str: + if not self._tables: + return "" + + lines = [ + "### VALIDATED RETRIEVED SCHEMA IDENTIFIERS ###", + "The SQL must use only these exact deployed Wren identifiers.", + "Each table value below is one indivisible Wren model identifier; never split it into database, schema, or table parts.", + "Use a multipart table reference only when that exact multipart identifier is listed below as a table value.", + "Do not derive table or column names from the user's wording, source SQL, physical names, comments, aliases, or descriptions.", + ] + for table_name, column_names in self._tables.items(): + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in sorted(column_names)) + table_relationships = self._relationships.get(table_name) + if table_relationships: + lines.append("relationships:") + lines.extend( + f"- {relationship}" + for relationship in sorted(table_relationships) + ) + lines.extend( + [ + "If the requested intent cannot be expressed with these exact identifiers, return null for sql.", + "### END VALIDATED RETRIEVED SCHEMA IDENTIFIERS ###", + ] + ) + return "\n".join(lines) + + def validate_sql(self, sql: str | None) -> str | None: + if not sql or not self._tables: + return None + + parsed_statements = sqlparse.parse(sql) + if not parsed_statements: + return "Generated SQL could not be parsed for schema grounding." + + referenced_tables: set[str] = set() + table_aliases: dict[str, str] = {} + qualified_columns: list[tuple[str, str]] = [] + cte_names: set[str] = set() + + for statement in parsed_statements: + cte_names.update(_extract_cte_names(statement)) + statement_tables, statement_aliases = _extract_table_references(statement) + referenced_tables.update(statement_tables) + table_aliases.update(statement_aliases) + qualified_columns.extend(_extract_qualified_columns(statement)) + + executable_tables = referenced_tables - cte_names + unknown_tables = sorted( + table_name + for table_name in executable_tables + if table_name not in self._tables + ) + if unknown_tables: + return ( + "Generated SQL referenced table(s) not present in the retrieved " + f"schema context: {', '.join(unknown_tables)}." + ) + + unknown_columns = [] + for qualifier, column_name in qualified_columns: + table_name = table_aliases.get(qualifier, qualifier) + if table_name in cte_names: + continue + if table_name in self._tables and self._tables[table_name]: + if column_name not in self._tables[table_name]: + unknown_columns.append(f"{qualifier}.{column_name}") + + if unknown_columns: + return ( + "Generated SQL referenced column(s) not present in the retrieved " + f"schema context: {', '.join(sorted(set(unknown_columns)))}." + ) + + return None + + +def construct_schema_identifier_catalog(contexts: list[str] | None) -> str: + return _SchemaCatalog.from_contexts(contexts or []).to_prompt() + + +def _extract_cte_names(token_list: TokenList) -> set[str]: + cte_names: set[str] = set() + with_seen = False + + for token in token_list.tokens: + if token.is_whitespace or token.ttype in Comment: + continue + + if token.normalized == "WITH": + with_seen = True + continue + + if not with_seen: + continue + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + name = _clean_identifier(identifier.get_name()) + if name: + cte_names.add(name) + break + + if isinstance(token, Identifier): + name = _clean_identifier(token.get_name()) + if name: + cte_names.add(name) + break + + if token.ttype is Keyword: + break + + return cte_names + + +def _extract_table_references(token_list: TokenList) -> tuple[set[str], dict[str, str]]: + table_names: set[str] = set() + aliases: dict[str, str] = {} + expect_table = False + + for token in token_list.tokens: + if token.is_whitespace or token.ttype in Comment: + continue + + if isinstance(token, Function): + continue + + if isinstance(token, TokenList): + nested_tables, nested_aliases = _extract_table_references(token) + table_names.update(nested_tables) + aliases.update(nested_aliases) + + if token.ttype is Keyword and token.normalized in { + "FROM", + "JOIN", + "INNER JOIN", + "LEFT JOIN", + "LEFT OUTER JOIN", + "RIGHT JOIN", + "RIGHT OUTER JOIN", + "FULL JOIN", + "FULL OUTER JOIN", + "CROSS JOIN", + }: + expect_table = True + continue + + if not expect_table: + continue + + if isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + _add_table_reference(identifier, table_names, aliases) + expect_table = False + continue + + if isinstance(token, Identifier): + _add_table_reference(token, table_names, aliases) + expect_table = False + continue + + if token.ttype is Keyword: + expect_table = False + + return table_names, aliases + + +def _add_table_reference( + identifier: Identifier, table_names: set[str], aliases: dict[str, str] +) -> None: + if not isinstance(identifier, Identifier): + return + + table_name = _table_reference_name(identifier) + alias = _clean_identifier(identifier.get_alias()) + if not table_name: + return + + table_names.add(table_name) + aliases[table_name] = table_name + if alias: + aliases[alias] = table_name + + +def _table_reference_name(identifier: Identifier) -> str | None: + if not isinstance(identifier, Identifier): + return None + + parent_getter = getattr(identifier, "get_parent_name", None) + real_getter = getattr(identifier, "get_real_name", None) + parent_name = _clean_identifier(parent_getter() if parent_getter else None) + real_name = _clean_identifier(real_getter() if real_getter else None) + if parent_name and real_name: + return f"{parent_name}.{real_name}" + return real_name + + +def _extract_qualified_columns(token_list: TokenList) -> list[tuple[str, str]]: + columns: list[tuple[str, str]] = [] + + for token in token_list.tokens: + if isinstance(token, Identifier): + _add_qualified_column(token, columns) + elif isinstance(token, IdentifierList): + for identifier in token.get_identifiers(): + if isinstance(identifier, Identifier): + _add_qualified_column(identifier, columns) + elif isinstance(identifier, TokenList): + columns.extend(_extract_qualified_columns(identifier)) + elif isinstance(token, TokenList): + columns.extend(_extract_qualified_columns(token)) + + return columns + + +def _add_qualified_column( + identifier: Identifier, columns: list[tuple[str, str]] +) -> None: + if not isinstance(identifier, Identifier): + return + + parent_getter = getattr(identifier, "get_parent_name", None) + real_getter = getattr(identifier, "get_real_name", None) + parent_name = _clean_identifier(parent_getter() if parent_getter else None) + column_name = _clean_identifier(real_getter() if real_getter else None) + if parent_name and column_name and column_name != "*": + columns.append((parent_name, column_name)) + + +def _find_matching_parenthesis(text: str, body_start: int) -> int | None: + depth = 1 + quote: str | None = None + i = body_start + + while i < len(text): + char = text[i] + if quote: + if char == quote: + if quote == "'" and i + 1 < len(text) and text[i + 1] == "'": + i += 2 + continue + quote = None + i += 1 + continue + + if char in {"'", '"', "`"}: + quote = char + elif char == "[": + closing = text.find("]", i + 1) + if closing == -1: + return None + i = closing + elif char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + return i + i += 1 + + return None + + +def _extract_ddl_column_names(ddl_body: str) -> set[str]: + column_names: set[str] = set() + for column_definition in _split_top_level_commas(ddl_body): + column_name = _extract_ddl_column_name(column_definition) + if column_name: + column_names.add(column_name) + return column_names + + +def _split_top_level_commas(value: str) -> list[str]: + parts: list[str] = [] + depth = 0 + quote: str | None = None + start = 0 + i = 0 + + while i < len(value): + char = value[i] + if quote: + if char == quote: + if quote == "'" and i + 1 < len(value) and value[i + 1] == "'": + i += 2 + continue + quote = None + i += 1 + continue + + if char in {"'", '"', "`"}: + quote = char + elif char == "[": + closing = value.find("]", i + 1) + if closing == -1: + break + i = closing + elif char == "(": + depth += 1 + elif char == ")": + depth = max(depth - 1, 0) + elif char == "," and depth == 0: + parts.append(value[start:i].strip()) + start = i + 1 + i += 1 + + tail = value[start:].strip() + if tail: + parts.append(tail) + return parts + + +def _extract_ddl_column_name(column_definition: str) -> str | None: + definition = _strip_leading_sql_comments(column_definition.strip()) + if not definition: + return None + + first_word = definition.split(maxsplit=1)[0].strip().strip('"`[]').upper() + if first_word in _DDL_COLUMN_KEYWORDS: + return None + + if definition.startswith("["): + closing = definition.find("]") + if closing > 0: + return _clean_identifier(definition[: closing + 1]) + + if definition.startswith('"'): + closing = definition.find('"', 1) + if closing > 0: + return _clean_identifier(definition[: closing + 1]) + + if definition.startswith("`"): + closing = definition.find("`", 1) + if closing > 0: + return _clean_identifier(definition[: closing + 1]) + + return _clean_identifier(definition.split(maxsplit=1)[0]) + + +def _strip_leading_sql_comments(value: str) -> str: + stripped = value.strip() + while stripped: + if stripped.startswith("--"): + lines = stripped.splitlines() + stripped = "\n".join(lines[1:]).strip() + continue + if stripped.startswith("/*"): + closing = stripped.find("*/") + if closing == -1: + return "" + stripped = stripped[closing + 2 :].strip() + continue + break + return stripped + + +def _clean_identifier(identifier: str | None) -> str | None: + if identifier is None: + return None + cleaned = identifier.strip().strip('"`[]') + return cleaned or None + + +def _extract_sql_from_json_value(value: Any) -> str | None: + if isinstance(value, str): + candidate = value.strip() + if not candidate: + return None + if candidate.upper().startswith(("SELECT", "WITH")): + return candidate + if candidate.startswith(("{", "[")): + try: + return _extract_sql_from_json_value(orjson.loads(candidate)) + except orjson.JSONDecodeError: + return None + return None + + if isinstance(value, dict): + for key in ("sql", "query"): + sql = _extract_sql_from_json_value(value.get(key)) + if sql: + return sql + + for key in ("arguments", "content", "tool_calls", "function_call", "message"): + sql = _extract_sql_from_json_value(value.get(key)) + if sql: + return sql + + for nested_value in value.values(): + sql = _extract_sql_from_json_value(nested_value) + if sql: + return sql + + if isinstance(value, list): + for item in value: + sql = _extract_sql_from_json_value(item) + if sql: + return sql + + return None + + +def _extract_sql_response(generation_result: str) -> tuple[str | None, str | None]: + cleaned_generation_result = generation_result.strip() + if not cleaned_generation_result: + return None, "No grounded SQL was generated from the current schema." + + if cleaned_generation_result.startswith(("{", "[")): + try: + payload = orjson.loads(cleaned_generation_result) + except orjson.JSONDecodeError: + return ( + None, + "SQL generation response did not include a supported SQL JSON payload.", + ) + + sql = _extract_sql_from_json_value(payload) + if sql: + return sql, None + + return ( + None, + f"SQL generation response did not include a supported SQL field: {payload}", + ) + + if cleaned_generation_result.upper().startswith(("SELECT", "WITH")): + return cleaned_generation_result, None + + return None, "SQL generation response did not include a supported SQL JSON payload." + + +_MANDATORY_SQL_GROUNDING_RULES = """ +### MANDATORY SQL GROUNDING RULES ### +- Treat the DATABASE SCHEMA section as the only source of executable table and column identifiers. +- Every table and column referenced in SELECT, FROM, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY must appear exactly in the CREATE TABLE, CREATE VIEW, or metric schema text provided in DATABASE SCHEMA. +- Comments, aliases, display labels, descriptions, reasoning text, SQL samples, and user wording are semantic hints only. They are never source table or source column identifiers. +- Physical datasource names, source database names, source schema names, source table names, source column names, lineage names, and names embedded inside descriptions or comments are semantic context only. Never use them as executable Wren table or column identifiers unless the exact same identifier is declared in DATABASE SCHEMA. +- Interpret the user's intent from the question wording, schema descriptions, aliases, display labels, calculated fields, metrics, and relationships, then express that intent with exact executable identifiers from DATABASE SCHEMA. +- When DATABASE SCHEMA contains WREN RETRIEVED SEMANTIC CONTEXT blocks, first use those blocks to understand each retrieved object's exact SQL identifier contract, semantic meaning, relationships, views, metrics, and calculated fields. +- When DATABASE SCHEMA contains WREN SQL IDENTIFIER CONTRACT sections, treat them as the compact authoritative list of executable identifiers for each retrieved object before reading semantic descriptions. +- In WREN RETRIEVED SEMANTIC CONTEXT, copy executable identifiers only from sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, or the following DDL declarations. +- Treat every retrieved Wren table/model name as one indivisible executable identifier. Prefixes, suffixes, underscores, source schema names, connector names, or words that look like database/schema parts are still part of that single Wren identifier. +- Never convert an exact Wren table/model name into a multipart native database reference. If DATABASE SCHEMA declares a table named abc_def, use "abc_def"; do not write abc.def, "abc"."def", or any other split form. +- Use multipart table references such as schema.table or "schema"."table" only when DATABASE SCHEMA declares that exact multipart Wren identifier as the executable table/model name. +- Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier are meaning only. Do not combine words, labels, ordinals, prefixes, suffixes, abbreviations, comments, or descriptions from those values into a table or column identifier. +- When a business term is represented by a column alias, display label, or description, use the corresponding real table and column name from DATABASE SCHEMA in the SQL, not the display text. +- The executable identifier is the name in the CREATE TABLE, CREATE VIEW, or metric field declaration. Do not derive executable identifiers by rewriting, translating, singularizing, pluralizing, spacing, casing, or abbreviating natural language, comments, aliases, display labels, or descriptions. +- Never generate SQL from assumptions such as "assuming the table contains", "assuming this column exists", or "a possible table/column". Use only schema-confirmed identifiers. +- Never generate placeholder identifiers, placeholder table names, or template markers in the SQL. If the retrieved metadata does not contain an executable object or column for a requested concept, omit that unsupported concept. +- Never create an identifier from user question wording by changing spaces, casing, punctuation, singular/plural form, abbreviations, prefixes, or suffixes. If the exact requested table or column concept is not represented by a retrieved schema identifier, return null for sql. +- If a requested concept, output column, filter, sort, join, grouping, measure, or time field is not represented by an exact table or column in DATABASE SCHEMA, do not invent a field for it. If that field is required to answer the request, return null for sql. +- If the DATABASE SCHEMA does not contain an identifier needed to answer part of the request, return null for sql or omit that unsupported part instead of naming a substitute table or column. +- When a dry run error reports an invalid object name or invalid column name, remove that identifier unless it appears exactly in DATABASE SCHEMA. Correct it only to an exact schema identifier. +- Do not replace an invalid identifier with a similar-looking physical, source, lineage, alias, display, description, sample, or error-message name. Regenerate from the user's intent and the current DATABASE SCHEMA, and omit unsupported parts instead of substituting non-schema identifiers. +- When using multiple tables to combine fields into the same output row, join only through the FOREIGN KEY relationships shown in DATABASE SCHEMA. If no relationship is shown for the needed tables, prefer a single table, view, or metric that already contains the requested fields. +- When the same requested result can be answered from multiple schema objects with compatible fields or metrics, include all relevant objects by combining separate result rows with UNION ALL instead of choosing only one object. +- Use UNION ALL only when each SELECT branch is independently valid from DATABASE SCHEMA and returns the same result shape. Do not use UNION ALL to combine unrelated concepts or to compensate for missing columns. +- If the question requires fields that are spread across multiple schema objects, use all required related tables, views, or metrics only when the DATABASE SCHEMA provides the needed columns and relationship path. +- Do not query INFORMATION_SCHEMA, system catalogs, metadata tables, or table-existence checks to answer the user. Query only the business tables, views, and metrics in DATABASE SCHEMA. +- SQL samples and query history are examples of intent and style only. Never copy a table name, column name, alias, literal value, or function from them unless it is also valid for the current DATABASE SCHEMA and SQL FUNCTIONS. +- Identifiers shown in prompt examples are illustrative only unless the same identifier appears exactly in DATABASE SCHEMA for this request. +- Generate Wren SQL only, not the native SQL dialect of the connected warehouse. Do not use SQL Server TOP, square-bracket quoting, backtick quoting, FETCH FIRST, OFFSET/FETCH pagination, or warehouse-specific functions unless they are explicitly listed in SQL FUNCTIONS for this request. +- For top, first, highest, lowest, largest, smallest, or other limited result requests, express the ranking/order with ORDER BY and apply a final LIMIT clause in Wren SQL. Never use SELECT TOP n. +- Apply relative date or time filters only when DATABASE SCHEMA contains an exact date/time field for the requested time concept and SQL FUNCTIONS contains the exact date/time operation needed. Do not compare text fields to date functions. +- Treat reasoning plans, correction notes, and error messages as non-executable context. Never copy SQL fragments, inferred identifiers, placeholder names, template markers, literal values, or unsupported functions from them. +- If a column comment, alias, display label, or description names a business concept, first locate the exact declared source column for that concept in DATABASE SCHEMA. If no exact declared source column exists, omit that concept. +- Use `display_label` and `description` only to understand business meaning; when a concept maps to a schema identifier, generated SQL must use that exact identifier. +- For aggregate sorting, select the aggregate with an alias and order by that alias instead of ordering directly by an aggregate expression. +- In grouped queries, every non-aggregate ORDER BY expression must be a selected grouping column, a selected ordering helper column that is also present in GROUP BY, or a selected aggregate alias. Do not order grouped SQL by a hidden column. +- Before returning the final SQL, silently check that each identifier and function in the SQL is grounded in DATABASE SCHEMA or SQL FUNCTIONS. If any identifier or function is ungrounded, remove that part. If the ungrounded part is needed to answer the user's requested intent, return null for sql. +- If the retrieved DATABASE SCHEMA does not contain a table, column, relationship, or supported function needed for part of the user's request, leave that part out instead of inventing a replacement. +- If a requested noun, output column, grouping, filter, or measure appears only in the user's wording and not in DATABASE SCHEMA, do not translate it into a generic object name. Use only schema-supported concepts and omit unsupported parts. +- If the user's primary requested subject, output column, grouping, filter, timeframe, measure, or required relationship cannot be grounded by the retrieved DATABASE SCHEMA, return null for sql instead of producing an approximate query. +- Do not answer by selecting a nearby table only because it was retrieved. A retrieved object is usable only when its declared table, columns, relationships, or metric fields support the user's requested intent. +- Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata contains specific modeled columns for the requested entity, measure, filter, date, or dimension. +- Prefer exact modeled fields over generic text search. If a requested concept is represented by an explicit declared column, use that column rather than searching a generic payload field with LIKE. +- If the schema already exposes a measure that directly matches the requested metric, use that exact measure column instead of recomputing it from invented component fields. +- Do not prefer or exclude any business domain by built-in rules. Ground every choice in the DATABASE SCHEMA supplied for this request. +""" + + _DEFAULT_TEXT_TO_SQL_RULES = """ ### SQL RULES ### - ONLY USE SELECT statements, NO DELETE, UPDATE OR INSERT etc. statements that might change the data in the database. @@ -172,11 +6249,13 @@ async def _classify_generation_result( - DON'T INCLUDE comments in the generated SQL query. - YOU MUST USE "JOIN" if you choose columns from multiple tables! - PREFER USING CTEs over subqueries. +- Copy table names exactly as one Wren identifier from DATABASE SCHEMA. Do not split underscores or source-schema-like prefixes into dot-qualified database/schema/table references. +- Use table aliases only as SQL aliases for already-declared Wren table names; never use aliases or source schema names as replacements for Wren table names. +- Qualify every source column reference with its exact Wren table name or SQL table alias in SELECT, JOIN, WHERE, GROUP BY, HAVING, and ORDER BY. Output aliases in the final SELECT may be unqualified. - When generating SQL query, always: - - Put double quotes around column and table names. + - Put double quotes around exact column and table names copied from DATABASE SCHEMA. - Put single quotes around string literals. - Never quote numeric literals. - For example: SELECT "customers"."customer_name" FROM "customers" WHERE "customers"."city" = 'Taipei' and "customers"."year" = 1992; - YOU MUST USE "lower(.) like lower()" function or "lower(.) = lower()" function for case-insensitive comparison! - Use "lower(.) LIKE lower()" when: - The user requests a pattern or partial match. @@ -193,23 +6272,12 @@ async def _classify_generation_result( - example 1: CAST(properties_closedate AS TIMESTAMP WITH TIME ZONE) - example 2: CAST('2024-11-09 00:00:00' AS TIMESTAMP WITH TIME ZONE) - example 3: CAST(DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AS TIMESTAMP WITH TIME ZONE) -- If the user asks for a specific date, please give the date range in SQL query - - example: "What is the total revenue for the month of 2024-11-01?" - - answer: "SELECT SUM(r.PriceSum) FROM Revenue r WHERE CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) >= CAST('2024-11-01 00:00:00' AS TIMESTAMP WITH TIME ZONE) AND CAST(r.PurchaseTimestamp AS TIMESTAMP WITH TIME ZONE) < CAST('2024-11-02 00:00:00' AS TIMESTAMP WITH TIME ZONE)" +- If the user asks for a specific date, use a date range over an exact date/time column from DATABASE SCHEMA. + - example: filter an exact date/time column with a start timestamp and the next boundary timestamp for the requested period. - USE THE VIEW TO SIMPLIFY THE QUERY. - DON'T MISUSE THE VIEW NAME. THE ACTUAL NAME IS FOLLOWING THE CREATE VIEW STATEMENT. - ONLY USE table/column alias in the final SELECT clause; don't use table/columnalias in the other clauses. -- Refer to the value of alias from the comment section of the corresponding table or column in the DATABASE SCHEMA section for reference when using alias in the final SELECT clause. - - EXAMPLE - DATABASE SCHEMA - /* {"alias":"_orders","description":"A model representing the orders data."} */ - CREATE TABLE orders ( - -- {"description":"A column that represents the timestamp when the order was approved.","alias":"_timestamp"} - ApprovedTimestamp TIMESTAMP - } - - SQL - SELECT "_orders"."ApprovedTimestamp" AS "_timestamp" FROM "orders" AS "_orders"; +- Refer to the alias value in the DATABASE SCHEMA comment for the corresponding table or column only as the output label in the final SELECT clause. Do not use alias values as source table or source column identifiers unless they are also exact executable identifiers declared in DATABASE SCHEMA. - DON'T USE '.' in column/table alias, replace '.' with '_' in column/table alias. - DON'T USE "FILTER(WHERE )" clause in the generated SQL query. - DON'T USE "EXTRACT(EPOCH FROM )" clause in the generated SQL query. @@ -218,8 +6286,11 @@ async def _classify_generation_result( - DON'T USE "TO_CHAR" function in the generated SQL query. - Aggregate functions are not allowed in the WHERE clause. Instead, they belong in the HAVING clause, which is used to filter after aggregation. - You can only add "ORDER BY" and "LIMIT" to the final "UNION" result. -- For the ranking problem, you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -- For the ranking problem, you must add the ranking column to the final SELECT clause. +- Do not use SELECT TOP n, FETCH FIRST, OFFSET/FETCH, square-bracket quoting, or backtick quoting. Use Wren SQL syntax with ORDER BY and a final LIMIT n clause for limited or top-N results. +- For top, bottom, highest, lowest, first, or last requests, sort by an exact selected column or aggregate alias and use LIMIT unless the user explicitly asks for rank values. +- For explicit ranking requests, use the ranking function `DENSE_RANK()`, add the ranking column to the final SELECT clause, and filter rank values with WHERE. +- For grouped trend queries, include any non-aggregate ordering key in both SELECT and GROUP BY, or order by selected grouping columns/aggregate aliases only. +- Reuse exact metric/measure columns when present. Do not invent component columns in order to calculate a requested metric that already exists in DATABASE SCHEMA. """ @@ -230,44 +6301,7 @@ async def _classify_generation_result( First, provide a brief explanation of what each field represents in the context of the schema, including how each field is computed using the relationships between models. Then, during the following tasks, if the user queries pertain to any calculated fields defined in the database schema, ensure to utilize those calculated fields appropriately in the output SQL queries. The goal is to accurately reflect the intent of the question in the SQL syntax, leveraging the pre-computed logic embedded within the calculated fields. - -EXAMPLES: -The given schema is created by the SQL command: - -CREATE TABLE orders ( - OrderId VARCHAR PRIMARY KEY, - CustomerId VARCHAR, - -- This column is a Calculated Field - -- column expression: avg(reviews.Score) - Rating DOUBLE, - -- This column is a Calculated Field - -- column expression: count(reviews.Id) - ReviewCount BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) - Size BIGINT, - -- This column is a Calculated Field - -- column expression: count(order_items.ItemNumber) > 1 - Large BOOLEAN, - FOREIGN KEY (CustomerId) REFERENCES customers(Id) -); - -Interpret the columns that are marked as Calculated Fields in the schema: -Rating (DOUBLE) - Calculated as the average score (avg) of the Score field from the reviews table where the reviews are associated with the order. This field represents the overall customer satisfaction rating for the order based on review scores. -ReviewCount (BIGINT) - Calculated by counting (count) the number of entries in the reviews table associated with this order. It measures the volume of customer feedback received for the order. -Size (BIGINT) - Represents the total number of items in the order, calculated by counting the number of item entries (ItemNumber) in the order_items table linked to this order. This field is useful for understanding the scale or size of an order. -Large (BOOLEAN) - A boolean value calculated to check if the number of items in the order exceeds one (count(order_items.ItemNumber) > 1). It indicates whether the order is considered large in terms of item quantity. - -And if the user input queries like these: -1. "How many large orders have been placed by customer with ID 'C1234'?" -2. "What is the average customer rating for orders that were rated by more than 10 reviewers?" - -For the first query: -First try to intepret the user query, the user wants to know the average rating for orders which have attracted significant review activity, specifically those with more than 10 reviews. -Then, according to the above intepretation about the given schema, the term 'Rating' is predefined in the Calculated Field of the 'orders' model. And, the number of reviews is also predefined in the 'ReviewCount' Calculated Field. -So utilize those Calculated Fields in the SQL generation process to give an answer like this: - -SQL Query: SELECT AVG(Rating) FROM orders WHERE ReviewCount > 10 +Use calculated fields only when their exact field names are declared in DATABASE SCHEMA and their descriptions or expressions match the user's intent. Do not recreate a calculated field expression with undeclared source tables or columns, and do not invent relationships that are not declared in DATABASE SCHEMA. """ _DEFAULT_METRIC_INSTRUCTIONS = """ @@ -297,68 +6331,7 @@ async def _classify_generation_result( If the given schema contains the structures marked as 'metric', you should first interpret the metric schema based on the above definition. Then, during the following tasks, if the user queries pertain to any metrics defined in the database schema, ensure to utilize those metrics appropriately in the output SQL queries. The target is making complex data analysis more accessible and manageable by pre-aggregating data and structuring it using the metric structure, and supporting direct querying for business insights. - -EXAMPLES: -The given schema is created by the SQL command: - -/* This table is a metric */ -/* Metric Base Object: orders */ -CREATE TABLE Revenue ( - -- This column is a dimension - PurchaseTimestamp TIMESTAMP, - -- This column is a dimension - CustomerId VARCHAR, - -- This column is a dimension - Status VARCHAR, - -- This column is a measure - -- expression: sum(order_items.Price) - PriceSum DOUBLE, - -- This column is a measure - -- expression: count(OrderId) - NumberOfOrders BIGINT -); - -Interpret the metric with the understanding of the metric structure: -1. Base Object: orders -This is the primary data source for the metric. -The orders table provides the underlying data from which dimensions and measures are derived. -It is the foundation upon which the metric is built, though it itself is not directly used in queries against the Revenue table. -It shows the reference between the 'Revenue' metric and the 'orders' model. For the user queries pretain to the 'Revenue' of 'orders', the metric should be utilize in the sql generation process. -2. Dimensions -The metric contains the columns marked as 'dimension'. They can be interpreted as below: -- PurchaseTimestamp (TIMESTAMP) - Acts as a temporal dimension, allowing analysis of revenue over time. This can be used to observe trends, seasonal variations, or performance over specific periods. -- CustomerId (VARCHAR) - A key dimension for customer segmentation, it enables the analysis of revenue generated from individual customers or customer groups. -- Status (VARCHAR) - Reflects the current state of an order (e.g., pending, completed, cancelled). This dimension is crucial for analyses that differentiate performance based on order status. -3. Measures -The metric contains the columns marked as 'measure'. They can be interpreted as below: -- PriceSum (DOUBLE) - A financial measure calculated as sum(order_items.Price), representing the total revenue generated from orders. This measure is vital for tracking overall sales performance and is the primary output of interest in many financial and business analyses. -- NumberOfOrders (BIGINT) - A count measure that provides the total number of orders. This is essential for operational metrics, such as assessing the volume of business activity and evaluating the efficiency of sales processes. - -Now, if the user input queries like this: -Question: "What was the total revenue from each customer last month?" - -First try to intepret the user query, the user asks for a breakdown of the total revenue generated by each customer in the previous calendar month. -The user is specifically interested in understanding how much each customer contributed to the total sales during this period. -To answer this question, it is suitable to use the following components from the metric: -1. CustomerId (Dimension): This will be used to group the revenue data by each unique customer, allowing us to segment the total revenue by customer. -2. PurchaseTimestamp (Dimension): This timestamp field will be used to filter the data to only include orders from the last month. -3. PriceSum (Measure): Since PriceSum is a pre-aggregated measure of total revenue (sum of order_items.Price), it can be directly used to sum up the revenue without needing further aggregation in the SQL query. -So utilize those metric components in the SQL generation process to give an answer like this: - -SQL Query: -SELECT - CustomerId, - PriceSum AS TotalRevenue -FROM - Revenue -WHERE - PurchaseTimestamp >= DATE_TRUNC('month', CURRENT_DATE - INTERVAL '1 month') AND - PurchaseTimestamp < DATE_TRUNC('month', CURRENT_DATE) +Use metric dimensions and measures only when their exact metric field names are declared in DATABASE SCHEMA and match the user's requested grouping, filtering, or aggregation. Treat the metric base object as semantic context; do not query the base object unless it is also declared as a retrieved executable table or metric in DATABASE SCHEMA. """ _DEFAULT_JSON_FIELD_INSTRUCTIONS = """ @@ -369,31 +6342,12 @@ async def _classify_generation_result( - LAX_FLOAT64 for double and float fields - LAX_INT64 for bigint fields - LAX_STRING for varchar fields - - For Example: - DATA SCHEMA: - `/* {"alias":"users","description":"A model representing the users data."} */ - CREATE TABLE users ( - -- {"alias":"address","description":"A JSON object that represents address information of this user.","json_type":"JSON","json_fields":{"json_type":"JSON","address.json.city":{"name":"city","type":"varchar","path":"$.city","properties":{"alias":"city","description":"City Name."}},"address.json.state":{"name":"state","type":"varchar","path":"$.state","properties":{"alias":"state","description":"ISO code or name of the state, province or district."}},"address.json.postcode":{"name":"postcode","type":"varchar","path":"$.postcode","properties":{"alias":"postcode","description":"Postal code."}},"address.json.country":{"name":"country","type":"varchar","path":"$.country","properties":{"alias":"country","description":"ISO code of the country."}}}} - address JSON - )` - To get the city of address in user table use SQL: - `SELECT LAX_STRING(JSON_QUERY(u.address, '$.city')) FROM user as u` + - Use only the exact JSON column and JSON field paths declared in DATABASE SCHEMA. - ONLY USE JSON_QUERY_ARRAY for querying "json_type":"JSON_ARRAY" is identified in the comment of the column, NOT the deprecated JSON_EXTRACT_ARRAY. - USE UNNEST to analysis each item individually in the ARRAY. YOU MUST SELECT FROM the parent table ahead of the UNNEST ARRAY. - The alias of the UNNEST(ARRAY) should be in the format `unnest_table_alias(individual_item_alias)` - - For Example: `SELECT item FROM UNNEST(ARRAY[1,2,3]) as my_unnested_table(item)` - If the items in the ARRAY are JSON objects, use JSON_QUERY to query the fields inside each JSON item. - - For Example: - DATA SCHEMA - `/* {"alias":"my_table","description":"A test my_table"} */ - CREATE TABLE my_table ( - -- {"alias":"elements","description":"elements column","json_type":"JSON_ARRAY","json_fields":{"json_type":"JSON_ARRAY","elements.json_array.id":{"name":"id","type":"bigint","path":"$.id","properties":{"alias":"id","description":"data ID."}},"elements.json_array.key":{"name":"key","type":"varchar","path":"$.key","properties":{"alias":"key","description":"data Key."}},"elements.json_array.value":{"name":"value","type":"varchar","path":"$.value","properties":{"alias":"value","description":"data Value."}}}} - elements JSON - )` - To get the number of elements in my_table table use SQL: - `SELECT LAX_INT64(JSON_QUERY(element, '$.number')) FROM my_table as t, UNNEST(JSON_QUERY_ARRAY(elements)) AS my_unnested_table(element) WHERE LAX_FLOAT64(JSON_QUERY(element, '$.value')) > 3.5` - To JOIN ON the fields inside UNNEST(ARRAY), YOU MUST SELECT FROM the parent table ahead of the UNNEST syntax, and the alias of the UNNEST(ARRAY) SHOULD BE IN THE FORMAT unnest_table_alias(individual_item_alias) - - For Example: `SELECT p.column_1, j.column_2 FROM parent_table AS p, join_table AS j JOIN UNNEST(p.array_column) AS unnested(array_item) ON j.id = array_item.id` - DON'T USE JSON_QUERY and JSON_QUERY_ARRAY when "json_type":"". - DON'T USE LAX_BOOL, LAX_FLOAT64, LAX_INT64, LAX_STRING when "json_type":"". """ @@ -401,40 +6355,34 @@ async def _classify_generation_result( sql_samples_instructions = """ #### Instructions for SQL Samples #### -Finally, you will learn from the sample SQL queries provided in the input. These samples demonstrate best practices and common patterns for querying this specific database. +Finally, you will learn from the sample questions provided in the input. These samples demonstrate intent and response style for this specific database. For each sample, you should: 1. Study the question that explains what the query aims to accomplish -2. Analyze the SQL implementation to understand: - - Table structures and relationships used - - Specific functions and operators employed - - Query patterns and techniques demonstrated -3. Use these samples as reference patterns when generating similar queries -4. Adapt the techniques shown in the samples to match new query requirements while maintaining consistent style and approach +2. Use these samples as intent and style context only, but treat the DATABASE SCHEMA as the only valid source of executable table and column names +3. Adapt the intent patterns to match new query requirements while maintaining consistent style and approach +4. Never copy table names, column names, aliases, literal values, placeholders, or functions from samples The samples will help you understand: -- Preferred table join patterns -- Common aggregation methods -- Specific function usage -- Query structure and formatting conventions +- Common analytical intents +- Common aggregation requests +- Preferred answer style -When generating new queries, try to follow similar patterns when applicable, while adapting them to the specific requirements of each new query. +When generating new queries, follow similar intent patterns when applicable, while adapting them to the specific requirements of each new query. -Learn about the usage of the schema structures and generate SQL based on them. +Learn about the user's intent from the samples and generate SQL from the current DATABASE SCHEMA and SQL FUNCTIONS only. """ sql_generation_reasoning_system_prompt = """ ### TASK ### -You are a helpful data analyst who is great at thinking deeply and reasoning about the user's question and the database schema, and you provide a step-by-step reasoning plan in order to answer the user's question. +You are a helpful data analyst who explains the user's analytical intent and provides a concise, non-executable reasoning plan for answering the user's question. ### INSTRUCTIONS ### 1. Think deeply and reason about the user's question, the database schema, and the user's query history if provided. -2. Explicitly state the following information in the reasoning plan: -if the user puts any specific timeframe(e.g. YYYY-MM-DD) in the user's question(excluding the value of the current time), you will put the absolute time frame in the SQL query; -otherwise, you will put the relative timeframe in the SQL query. -3. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must use the ranking function, `DENSE_RANK()` to rank the results and then use `WHERE` clause to filter the results. -4. For the ranking problem(e.g. "top x", "bottom x", "first x", "last x"), you must add the ranking column to the final SELECT clause. +2. Explicitly state requested timeframes in natural language only. Mention exact date/time columns only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +3. For top, bottom, first, last, highest, or lowest requests, describe the requested ordering and limit in natural language. Mention exact ordering columns or measures only when they are declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +4. Do not mention SQL functions, operators, or expression syntax in the reasoning plan. 5. If USER INSTRUCTIONS section is provided, make sure to consider them in the reasoning plan. 6. If SQL SAMPLES section is provided, make sure to consider them in the reasoning plan. 7. Give a step by step reasoning plan in order to answer user's question. @@ -442,9 +6390,22 @@ async def _classify_generation_result( 9. Don't include SQL in the reasoning plan. 10. Each step in the reasoning plan must start with a number, a title(in bold format in markdown), and a reasoning for the step. 11. Do not include ```markdown or ``` in the answer. -12. A table name in the reasoning plan must be in this format: `table: `. -13. A column name in the reasoning plan must be in this format: `column: .`. -14. ONLY SHOWING the reasoning plan in bullet points. +12. Mention table names only by writing the literal prefix `table:` followed by an exact table name declared in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +13. Mention column names only by writing the literal prefix `column:` followed by an exact declared table name, a dot, and an exact column name declared for that table in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +14. Do not mention aliases, source names, physical names, lineage names, schema names, database names, literal values, placeholders, or identifier-like labels from comments, SQL samples, failed SQL, or user wording as executable identifiers. +15. Do not write SQL, possible SQL, sample SQL, assumed SQL, SQL clauses, SQL functions, code blocks, or executable expressions in the reasoning plan. Do not write date/time expressions in the reasoning plan. +16. Never use phrases such as "assuming the table contains", "assuming this column exists", or "the SQL could look like this". If the available metadata does not clearly support part of the request, state that the available metadata does not support that part without naming missing objects. +17. If the question asks for a concept, filter, sort, or timeframe, describe the requested operation in business language and cite exact declared tables or columns only when they are grounded by DATABASE SCHEMA. +18. Interpret the user's intent from wording, aliases, display labels, descriptions, calculated fields, metrics, and relationships, then ground the plan in exact declared schema identifiers. +19. If multiple schema objects are required, identify the exact declared relationship path from DATABASE SCHEMA. If no relationship path is declared, say that the retrieved metadata does not provide a join path. +20. Treat SQL samples and query history as examples only. Do not copy table names, column names, aliases, values, placeholders, functions, or SQL patterns from them into the reasoning plan unless they also appear exactly in DATABASE SCHEMA. +21. Do not mention placeholder SQL, metadata-table checks, INFORMATION_SCHEMA, or replacement instructions to the user. +22. The reasoning plan is semantic context for intent only, not a source of executable identifiers. SQL generation must re-read DATABASE SCHEMA and WREN SQL IDENTIFIER CONTRACT before using any identifier. +23. ONLY SHOWING the reasoning plan in bullet points. +24. Do not use the words "assume", "assuming", "likely", "possible", "might", or "example" when describing tables, columns, filters, or SQL. +25. If exact deployed table and column identifiers are not available for a requested part, say only that the retrieved metadata does not support that part. Do not propose a replacement name. +26. Do not write table names or column names from the user's wording unless the same identifier appears exactly in DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT. +27. Do not include code blocks, inline SQL fragments, SELECT statements, WHERE clauses, join clauses, or any query-shaped text in the reasoning plan. ### FINAL ANSWER FORMAT ### The final answer must be a reasoning plan in plain Markdown string format @@ -462,12 +6423,19 @@ def _extract_from_sql_knowledge( def get_text_to_sql_rules(sql_knowledge: SqlKnowledge | None = None) -> str: + rules = _DEFAULT_TEXT_TO_SQL_RULES if sql_knowledge is not None: - return _extract_from_sql_knowledge( - sql_knowledge, "text_to_sql_rule", _DEFAULT_TEXT_TO_SQL_RULES + connector_rules = _extract_from_sql_knowledge( + sql_knowledge, "text_to_sql_rule", "" ) + if connector_rules: + rules = f"""{rules} + +### CONNECTOR SQL KNOWLEDGE ### +Use the following connector-specific knowledge only when it does not conflict with Wren SQL syntax, DATABASE SCHEMA identifiers, SQL FUNCTIONS, or WREN SQL IDENTIFIER CONTRACT. +{connector_rules}""" - return _DEFAULT_TEXT_TO_SQL_RULES + return f"{rules}\n\n{_MANDATORY_SQL_GROUNDING_RULES}" def get_calculated_field_instructions(sql_knowledge: SqlKnowledge | None = None) -> str: @@ -503,41 +6471,54 @@ def get_sql_generation_system_prompt(sql_knowledge: SqlKnowledge | None = None) text_to_sql_rules = get_text_to_sql_rules(sql_knowledge) return f""" -You are a helpful assistant that converts natural language queries into ANSI SQL queries. +You are a helpful assistant that converts natural language queries into Wren SQL queries. -Given user's question, database schema, etc., you should think deeply and carefully and generate the SQL query based on the given reasoning plan step by step. +Given the user's question and database schema, generate one grounded Wren SQL query. The DATABASE SCHEMA is the only source of executable identifiers. ### GENERAL RULES ### 1. YOU MUST FOLLOW the instructions strictly to generate the SQL query if the section of USER INSTRUCTIONS is available in user's input. -2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. -3. YOU MUST REFER to the sql samples and learn the usage of the schema structures and how SQL is written based on them if the section of SQL SAMPLES is available in user's input. -4. YOU MUST FOLLOW the reasoning plan step by step strictly to generate the SQL query if the section of REASONING PLAN is available in user's input. -5. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. +2. YOU MUST ONLY CHOOSE the appropriate functions from the sql functions list and use them in the SQL query if the section of SQL FUNCTIONS is available in user's input. Use the exact supported syntax shown there; otherwise omit the function-dependent part of the request. +3. YOU MUST REFER to the sql samples only as examples of intent and style if the section of SQL SAMPLES is available in user's input. Do not copy identifiers, literals, placeholders, SQL patterns, or functions from samples. +4. YOU MUST treat the reasoning plan as semantic context for intent only. Do not copy identifiers, functions, literal values, SQL fragments, template markers, or placeholders from the reasoning plan. Choose every executable identifier only from DATABASE SCHEMA or WREN SQL IDENTIFIER CONTRACT, and every function only from SQL FUNCTIONS. +5. YOU MUST answer the user's intent, not just exact wording. Use schema aliases, descriptions, calculated fields, metrics, and relationships to understand intent, then generate SQL with exact DATABASE SCHEMA identifiers only. +6. YOU MUST first read any WREN SQL IDENTIFIER CONTRACT and WREN RETRIEVED SEMANTIC CONTEXT block attached to each schema object. Use sql_table_name_use_exactly, sql_column_name_use_exactly, sql_column_names_use_exactly, relationship_constraints_use_exactly, and the following DDL declarations as executable grounding. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand business meaning. +7. When DATABASE SCHEMA contains EXECUTABLE WREN IDENTIFIER CATALOG sections, treat those sections as the first and clearest list of allowed executable identifiers. +8. If the user asks for fields that exist across multiple related schema objects, include those objects only when DATABASE SCHEMA shows the exact columns and relationship path needed to join them. +9. If the user asks for a result that is represented in multiple schema objects with compatible fields, include all relevant objects using independently valid SELECT branches combined with UNION ALL. Use joins only for relationship-backed row-level combinations. +10. Before finalizing the JSON response, YOU MUST perform a silent grounding check: every table, column, join key, filter field, grouping field, ordering field, and function in the SQL must be present in DATABASE SCHEMA or SQL FUNCTIONS. If a planned element is not grounded, omit that element. If the element is needed to answer the user's requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. +11. YOU MUST treat source database/schema/table names, physical datasource names, lineage names, comments, aliases, and display labels as semantic context only. Never use them as executable identifiers unless the exact same identifier appears in DATABASE SCHEMA. +12. If an identifier, literal value, placeholder, template marker, or function appears only in SQL samples, failed SQL, descriptions, lineage, reasoning text, or error messages, it is not executable for this request; ignore those parts when generating executable SQL. +13. If any planned SQL identifier cannot be copied exactly from DATABASE SCHEMA, EXECUTABLE WREN IDENTIFIER CATALOG, or WREN SQL IDENTIFIER CONTRACT, return null for sql. Never create a table or column from the user's wording. +14. YOU MUST FOLLOW SQL Rules if they are not contradicted with instructions. {text_to_sql_rules} ### FINAL ANSWER FORMAT ### -The final answer must be a ANSI SQL query in JSON format: +The final answer must be JSON. Return a SQL string only when it is fully grounded in DATABASE SCHEMA and SQL FUNCTIONS and it answers the user's requested intent. Do not create table or column identifiers from the user's wording. If the retrieved schema does not ground the requested subject, output column, filter, grouping, measure, timeframe, or relationship, return null for sql. {{ - "sql": + "sql": "SQL query string using only identifiers declared in DATABASE SCHEMA, or null" }} """ class SqlGenerationResult(BaseModel): - sql: str + model_config = ConfigDict(extra="forbid") + + sql: str | None SQL_GENERATION_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { "name": "sql_generation_result", + "strict": True, "schema": SqlGenerationResult.model_json_schema(), }, - } + }, } diff --git a/wren-ai-service/src/pipelines/indexing/db_schema.py b/wren-ai-service/src/pipelines/indexing/db_schema.py index 394d087b46..e78bad906f 100644 --- a/wren-ai-service/src/pipelines/indexing/db_schema.py +++ b/wren-ai-service/src/pipelines/indexing/db_schema.py @@ -19,8 +19,8 @@ AsyncDocumentWriter, DocumentCleaner, MDLValidator, - clean_display_name, ) +from src.pipelines.common import build_project_deploy_filter from src.pipelines.indexing.utils import helper logger = logging.getLogger("wren-ai-service") @@ -34,9 +34,15 @@ async def run( mdl: Dict[str, Any], column_batch_size: int, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ): def _additional_meta() -> Dict[str, Any]: - return {"project_id": project_id} if project_id else {} + metadata = {} + if project_id: + metadata["project_id"] = project_id + if mdl_hash: + metadata["mdl_hash"] = mdl_hash + return metadata chunks = [ { @@ -96,6 +102,8 @@ async def _preprocessor(model: Dict[str, Any], **kwargs) -> Dict[str, Any]: return { "name": model.get("name", ""), "properties": model.get("properties", {}), + "tableReference": model.get("tableReference"), + "refSql": model.get("refSql"), "columns": columns, "primaryKey": model.get("primaryKey", ""), } @@ -132,10 +140,7 @@ def _convert_models_and_relationships( def _model_command(model: Dict[str, Any]) -> dict: properties = model.get("properties", {}) - model_properties = { - "alias": clean_display_name(properties.get("displayName", "")), - "description": properties.get("description", ""), - } + model_properties = helper.normalize_semantic_properties(properties) comment = f"\n/* {str(model_properties)} */\n" table_name = model["name"] @@ -143,6 +148,9 @@ def _model_command(model: Dict[str, Any]) -> dict: "type": "TABLE", "comment": comment, "name": table_name, + "properties": properties, + "tableReference": model.get("tableReference"), + "refSql": model.get("refSql"), } return {"name": table_name, "payload": str(payload)} @@ -162,6 +170,7 @@ def _column_command(column: Dict[str, Any], model: Dict[str, Any]) -> dict: "name": column["name"], "data_type": column["type"], "is_primary_key": column["name"] == model["primaryKey"], + "properties": column.get("properties", {}), } def _relationship_command( @@ -305,11 +314,13 @@ async def chunk( chunker: DDLChunker, column_batch_size: int, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: return await chunker.run( mdl=mdl, column_batch_size=column_batch_size, project_id=project_id, + mdl_hash=mdl_hash, ) @@ -345,6 +356,7 @@ def __init__( **kwargs, ) -> None: dbschema_store = document_store_provider.get_store() + self._store = dbschema_store self._components = { "cleaner": DocumentCleaner([dbschema_store]), @@ -368,7 +380,10 @@ def __init__( @observe(name="DB Schema Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id}, DB Schema Indexing pipeline is running..." @@ -378,6 +393,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, **self._configs, }, @@ -390,3 +406,22 @@ async def clean(self, project_id: Optional[str] = None) -> None: cleaner=self._components["cleaner"], project_id=project_id, ) + + async def count_documents( + self, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ) -> int: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + if project_deploy_filter := build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ): + filters["conditions"] += project_deploy_filter["conditions"] + + return await self._store.count_documents(filters=filters) diff --git a/wren-ai-service/src/pipelines/indexing/historical_question.py b/wren-ai-service/src/pipelines/indexing/historical_question.py index ff30d91f0f..2873672292 100644 --- a/wren-ai-service/src/pipelines/indexing/historical_question.py +++ b/wren-ai-service/src/pipelines/indexing/historical_question.py @@ -53,7 +53,12 @@ class ViewChunker: """ @component.output_types(documents=List[Document]) - def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None) -> None: + def run( + self, + mdl: Dict[str, Any], + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ) -> None: def _get_content(view: Dict[str, Any]) -> str: properties = view.get("properties", {}) historical_queries = properties.get("historical_queries", []) @@ -70,7 +75,12 @@ def _get_meta(view: Dict[str, Any]) -> Dict[str, Any]: } def _additional_meta() -> Dict[str, Any]: - return {"project_id": project_id} if project_id else {} + metadata = {} + if project_id: + metadata["project_id"] = project_id + if mdl_hash: + metadata["mdl_hash"] = mdl_hash + return metadata chunks = [ { @@ -105,12 +115,16 @@ def chunk( mdl: Dict[str, Any], chunker: ViewChunker, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: - return chunker.run(mdl=mdl, project_id=project_id) + return chunker.run(mdl=mdl, project_id=project_id, mdl_hash=mdl_hash) @observe(capture_input=False, capture_output=False) async def embedding(chunk: Dict[str, Any], embedder: Any) -> Dict[str, Any]: + if not chunk["documents"]: + return chunk + return await embedder.run(documents=chunk["documents"]) @@ -161,7 +175,10 @@ def __init__( @observe(name="Historical Question Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id}, Historical Question Indexing pipeline is running..." @@ -171,6 +188,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/indexing/instructions.py b/wren-ai-service/src/pipelines/indexing/instructions.py index b23f3cf2ab..96b205b062 100644 --- a/wren-ai-service/src/pipelines/indexing/instructions.py +++ b/wren-ai-service/src/pipelines/indexing/instructions.py @@ -29,10 +29,19 @@ class Instruction(BaseModel): @component class InstructionsConverter: @component.output_types(documents=List[Document]) - def run(self, instructions: list[Instruction], project_id: str = ""): + def run( + self, + instructions: list[Instruction], + project_id: str = "", + mdl_hash: Optional[str] = None, + ): logger.info(f"Project ID: {project_id} Converting instructions to documents...") - addition = {"project_id": project_id} if project_id else {} + addition = {} + if project_id: + addition["project_id"] = project_id + if mdl_hash: + addition["mdl_hash"] = mdl_hash return { "documents": [ @@ -61,20 +70,25 @@ def __init__(self, instructions_store: DocumentStore) -> None: @component.output_types() async def run( - self, instruction_ids: List[str], project_id: Optional[str] = None + self, + instruction_ids: List[str], + project_id: Optional[str] = None, + delete_all: bool = False, ) -> None: - filter = { - "operator": "AND", - "conditions": [ - {"field": "instruction_id", "operator": "in", "value": instruction_ids}, - ], - } + conditions = [] + + if not delete_all: + conditions.append( + {"field": "instruction_id", "operator": "in", "value": instruction_ids} + ) if project_id: - filter["conditions"].append( + conditions.append( {"field": "project_id", "operator": "==", "value": project_id} ) + filter = {"operator": "AND", "conditions": conditions} if conditions else None + return await self.store.delete_documents(filter) @@ -108,7 +122,11 @@ async def clean( ) -> Dict[str, Any]: instruction_ids = [instruction.id for instruction in instructions] if instruction_ids or delete_all: - await cleaner.run(instruction_ids=instruction_ids, project_id=project_id) + await cleaner.run( + instruction_ids=instruction_ids, + project_id=project_id, + delete_all=delete_all, + ) return embedding @@ -152,6 +170,7 @@ async def run( self, instructions: list[Instruction], project_id: str = "", + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} Instructions Indexing pipeline is running..." @@ -159,6 +178,7 @@ async def run( input = { "project_id": project_id, + "mdl_hash": mdl_hash, "instructions": instructions, **self._components, } diff --git a/wren-ai-service/src/pipelines/indexing/project_meta.py b/wren-ai-service/src/pipelines/indexing/project_meta.py index 2ee4a08074..566d0b52ac 100644 --- a/wren-ai-service/src/pipelines/indexing/project_meta.py +++ b/wren-ai-service/src/pipelines/indexing/project_meta.py @@ -30,9 +30,14 @@ def validate_mdl(mdl_str: str, validator: MDLValidator) -> dict[str, Any]: def chunk( mdl: dict[str, Any], project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> dict[str, Any]: - addition = {"project_id": project_id} if project_id else {} - data_source = mdl.get("dataSource", "local_file").lower() + addition = {} + if project_id: + addition["project_id"] = project_id + if mdl_hash: + addition["mdl_hash"] = mdl_hash + data_source = str(mdl.get("dataSource") or "local_file").lower() if data_source == "duckdb": # fix duckdb to local_file due to wren-ibis implementation at the moment @@ -87,7 +92,10 @@ def __init__( @observe(name="Project Meta Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> dict[str, Any]: logger.info( f"Project ID: {project_id}, Project Meta Indexing pipeline is running..." @@ -97,6 +105,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, }, ) diff --git a/wren-ai-service/src/pipelines/indexing/sql_pairs.py b/wren-ai-service/src/pipelines/indexing/sql_pairs.py index a92fb36df1..65eec14266 100644 --- a/wren-ai-service/src/pipelines/indexing/sql_pairs.py +++ b/wren-ai-service/src/pipelines/indexing/sql_pairs.py @@ -28,10 +28,19 @@ class SqlPair(BaseModel): @component class SqlPairsConverter: @component.output_types(documents=List[Document]) - def run(self, sql_pairs: List[SqlPair], project_id: str = ""): + def run( + self, + sql_pairs: List[SqlPair], + project_id: str = "", + mdl_hash: Optional[str] = None, + ): logger.info(f"Project ID: {project_id} Converting SQL pairs to documents...") - addition = {"project_id": project_id} if project_id else {} + addition = {} + if project_id: + addition["project_id"] = project_id + if mdl_hash: + addition["mdl_hash"] = mdl_hash return { "documents": [ @@ -56,20 +65,25 @@ def __init__(self, sql_pairs_store: DocumentStore) -> None: @component.output_types() async def run( - self, sql_pair_ids: List[str], project_id: Optional[str] = None + self, + sql_pair_ids: List[str], + project_id: Optional[str] = None, + delete_all: bool = False, ) -> None: - filter = { - "operator": "AND", - "conditions": [ - {"field": "sql_pair_id", "operator": "in", "value": sql_pair_ids}, - ], - } + conditions = [] + + if not delete_all: + conditions.append( + {"field": "sql_pair_id", "operator": "in", "value": sql_pair_ids} + ) if project_id: - filter["conditions"].append( + conditions.append( {"field": "project_id", "operator": "==", "value": project_id} ) + filter = {"operator": "AND", "conditions": conditions} if conditions else None + return await self.store.delete_documents(filter) @@ -81,7 +95,7 @@ def boilerplates( mdl = orjson.loads(mdl_str) return { - boilerplate.lower() + str(boilerplate).lower() for model in mdl.get("models", []) if (boilerplate := model.get("properties", {}).get("boilerplate")) } @@ -91,7 +105,11 @@ def boilerplates( def sql_pairs( boilerplates: Set[str], external_pairs: Dict[str, Any], + include_default_pairs: bool = True, ) -> List[SqlPair]: + if not include_default_pairs and not external_pairs: + return [] + return [ SqlPair( id=pair.get("id"), @@ -109,8 +127,13 @@ def to_documents( sql_pairs: List[SqlPair], document_converter: SqlPairsConverter, project_id: str = "", + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: - return document_converter.run(sql_pairs=sql_pairs, project_id=project_id) + return document_converter.run( + sql_pairs=sql_pairs, + project_id=project_id, + mdl_hash=mdl_hash, + ) @observe(capture_input=False, capture_output=False) @@ -118,6 +141,9 @@ async def embedding( to_documents: Dict[str, Any], embedder: Any, ) -> Dict[str, Any]: + if not to_documents["documents"]: + return to_documents + return await embedder.run(documents=to_documents["documents"]) @@ -131,7 +157,11 @@ async def clean( ) -> Dict[str, Any]: sql_pair_ids = [sql_pair.id for sql_pair in sql_pairs] if sql_pair_ids or delete_all: - await cleaner.run(sql_pair_ids=sql_pair_ids, project_id=project_id) + await cleaner.run( + sql_pair_ids=sql_pair_ids, + project_id=project_id, + delete_all=delete_all, + ) return embedding @@ -141,6 +171,9 @@ async def write( clean: Dict[str, Any], writer: AsyncDocumentWriter, ) -> None: + if not clean["documents"]: + return None + return await writer.run(documents=clean["documents"]) @@ -194,19 +227,27 @@ async def run( self, mdl_str: str, project_id: str = "", + mdl_hash: Optional[str] = None, external_pairs: Optional[Dict[str, Any]] = None, + delete_all: bool = False, + include_default_pairs: bool = True, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id} SQL Pairs Indexing pipeline is running..." ) + pairs = { + **(self._external_pairs if include_default_pairs else {}), + **(external_pairs or {}), + } + input = { "mdl_str": mdl_str, "project_id": project_id, - "external_pairs": { - **self._external_pairs, - **(external_pairs or {}), - }, + "mdl_hash": mdl_hash, + "external_pairs": pairs, + "delete_all": delete_all, + "include_default_pairs": include_default_pairs, **self._components, } diff --git a/wren-ai-service/src/pipelines/indexing/table_description.py b/wren-ai-service/src/pipelines/indexing/table_description.py index 6da100868f..3d61238d79 100644 --- a/wren-ai-service/src/pipelines/indexing/table_description.py +++ b/wren-ai-service/src/pipelines/indexing/table_description.py @@ -14,6 +14,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider +from src.pipelines.common import build_project_deploy_filter from src.pipelines.indexing import AsyncDocumentWriter, DocumentCleaner, MDLValidator logger = logging.getLogger("wren-ai-service") @@ -22,9 +23,19 @@ @component class TableDescriptionChunker: @component.output_types(documents=List[Document]) - def run(self, mdl: Dict[str, Any], project_id: Optional[str] = None): + def run( + self, + mdl: Dict[str, Any], + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): def _additional_meta() -> Dict[str, Any]: - return {"project_id": project_id} if project_id else {} + metadata = {} + if project_id: + metadata["project_id"] = project_id + if mdl_hash: + metadata["mdl_hash"] = mdl_hash + return metadata chunks = [ { @@ -88,12 +99,16 @@ def chunk( mdl: Dict[str, Any], chunker: TableDescriptionChunker, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: - return chunker.run(mdl=mdl, project_id=project_id) + return chunker.run(mdl=mdl, project_id=project_id, mdl_hash=mdl_hash) @observe(capture_input=False, capture_output=False) async def embedding(chunk: Dict[str, Any], embedder: Any) -> Dict[str, Any]: + if not chunk["documents"]: + return chunk + return await embedder.run(documents=chunk["documents"]) @@ -125,6 +140,7 @@ def __init__( table_description_store = document_store_provider.get_store( dataset_name="table_descriptions" ) + self._store = table_description_store self._components = { "cleaner": DocumentCleaner([table_description_store]), @@ -145,7 +161,10 @@ def __init__( @observe(name="Table Description Indexing") async def run( - self, mdl_str: str, project_id: Optional[str] = None + self, + mdl_str: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Dict[str, Any]: logger.info( f"Project ID: {project_id}, Table Description Indexing pipeline is running..." @@ -155,6 +174,7 @@ async def run( inputs={ "mdl_str": mdl_str, "project_id": project_id, + "mdl_hash": mdl_hash, **self._components, **self._configs, }, @@ -167,3 +187,22 @@ async def clean(self, project_id: Optional[str] = None) -> None: cleaner=self._components["cleaner"], project_id=project_id, ) + + async def count_documents( + self, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ) -> int: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + ], + } + if project_deploy_filter := build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ): + filters["conditions"] += project_deploy_filter["conditions"] + + return await self._store.count_documents(filters=filters) diff --git a/wren-ai-service/src/pipelines/indexing/utils/helper.py b/wren-ai-service/src/pipelines/indexing/utils/helper.py index 3829324a0a..b640de26be 100644 --- a/wren-ai-service/src/pipelines/indexing/utils/helper.py +++ b/wren-ai-service/src/pipelines/indexing/utils/helper.py @@ -11,6 +11,18 @@ logger = logging.getLogger("wren-ai-service") +SEMANTIC_METADATA_KEYS = ( + "aliases", + "synonyms", + "businessContext", + "dataMeaning", + "semanticType", + "useCases", + "aggregationDefault", + "format", + "examples", +) + class Helper: def __init__( @@ -28,12 +40,34 @@ def __call__(self, column: Dict[str, Any], **kwargs) -> Any: return self.helper(column, **kwargs) -def _properties_comment(column: Dict[str, Any], **_) -> str: - props = column["properties"] - column_properties = { +def normalize_semantic_properties( + props: Dict[str, Any], + include_source_column_name: bool = False, +) -> Dict[str, Any]: + if not isinstance(props, dict): + props = {} + + semantic_properties = { "alias": clean_display_name(props.get("displayName", "")), "description": props.get("description", ""), } + if include_source_column_name: + semantic_properties["sourceColumnName"] = props.get("sourceColumnName", "") + + for key in SEMANTIC_METADATA_KEYS: + value = props.get(key) + if value not in ("", None, [], {}): + semantic_properties[key] = value + + return semantic_properties + + +def _properties_comment(column: Dict[str, Any], **_) -> str: + props = column["properties"] + column_properties = normalize_semantic_properties( + props, + include_source_column_name=True, + ) # Add any nested columns if they exist nested = {k: v for k, v in props.items() if k.startswith("nested")} diff --git a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py index 6c8dd7bbe3..5bf9d48382 100644 --- a/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/db_schema_retrieval.py @@ -1,20 +1,28 @@ import ast +import asyncio import logging +import re import sys +import time +from functools import lru_cache from typing import Any, Optional import orjson +import sqlparse import tiktoken from hamilton import base from hamilton.async_driver import AsyncDriver from haystack import Document from haystack.components.builders.prompt_builder import PromptBuilder -from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict +from sqlparse.sql import Identifier, IdentifierList +from sqlparse.tokens import DML, Comment, Keyword +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider, LLMProvider from src.pipelines.common import ( + build_project_deploy_filter, build_table_ddl, clean_up_new_lines, get_engine_supported_data_type, @@ -24,12 +32,193 @@ logger = logging.getLogger("wren-ai-service") +_SEMANTIC_TABLE_NAME_MERGE_LIMIT = 8 +_LEXICAL_SCHEMA_TABLE_NAME_MERGE_LIMIT = 8 +_MAX_RETRIEVED_TABLE_NAMES = 24 +_MAX_LLM_SCHEMA_CONTEXT_TABLES = 8 +_MAX_LLM_SCHEMA_CONTEXT_TOKENS = 12_000 +_MAX_RELATED_TABLE_EXPANSION_DEPTH = 1 +_RANK_TOKEN = re.compile(r"[a-z0-9]+") +_RANK_GENERIC_QUERY_TOKENS = { + "a", + "across", + "all", + "an", + "and", + "as", + "average", + "avg", + "between", + "bottom", + "breakdown", + "bucket", + "buckets", + "by", + "count", + "counts", + "date", + "day", + "descending", + "distribution", + "each", + "for", + "from", + "group", + "grouped", + "groups", + "has", + "have", + "highest", + "how", + "in", + "is", + "latest", + "least", + "list", + "lowest", + "many", + "max", + "maximum", + "me", + "mean", + "min", + "minimum", + "blank", + "empty", + "missing", + "null", + "month", + "monthly", + "most", + "newest", + "number", + "of", + "ordered", + "per", + "quarter", + "recent", + "record", + "records", + "result", + "results", + "row", + "rows", + "show", + "sort", + "sorted", + "sum", + "the", + "there", + "this", + "to", + "top", + "total", + "using", + "was", + "week", + "were", + "what", + "where", + "which", + "with", + "year", +} +_COMPOUND_IDENTIFIER_PART_TOKENS = { + "account", + "amount", + "balance", + "business", + "buyer", + "category", + "client", + "company", + "count", + "customer", + "date", + "division", + "failure", + "gross", + "group", + "invoice", + "market", + "material", + "month", + "name", + "order", + "person", + "priority", + "product", + "quantity", + "record", + "repair", + "sales", + "salesperson", + "severity", + "status", + "supplier", + "ticket", + "type", + "unit", + "value", + "vendor", + "year", +} +_COMPOUND_IDENTIFIER_ALIASES = { + "acct": {"account"}, + "amt": {"amount"}, + "bu": {"business", "unit"}, + "cust": {"customer"}, + "gl": {"general", "ledger"}, + "ord": {"order"}, + "prod": {"product"}, + "qty": {"quantity"}, + "vend": {"vendor"}, +} + +_ALL_SCHEMA_DOCUMENTS_CACHE: dict[tuple[int, str, str], list[Document]] = {} +_ALL_SCHEMA_DOCUMENTS_CACHE_LOCKS: dict[tuple[int, str, str], asyncio.Lock] = {} +_SCHEMA_DOCUMENTS_CACHE: dict[ + tuple[int, str, str, tuple[str, ...]], list[Document] +] = {} +_SCHEMA_DOCUMENTS_CACHE_LOCKS: dict[ + tuple[int, str, str, tuple[str, ...]], asyncio.Lock +] = {} + + +def _elapsed_ms(started_at: float) -> float: + return (time.perf_counter() - started_at) * 1000 + + +def _log_retrieval_timing( + stage: str, + started_at: float, + project_id: str | None = None, + **fields: Any, +) -> None: + suffix = " ".join( + f"{key}={value}" for key, value in fields.items() if value is not None + ) + logger.info( + "Ask timing project_id=%s stage=%s elapsed_ms=%.1f%s%s", + project_id or "", + stage, + _elapsed_ms(started_at), + " " if suffix else "", + suffix, + ) + table_columns_selection_system_prompt = """ ### TASK ### You are a highly skilled data analyst. Your goal is to examine the provided database schema, interpret the posed question, and identify the specific columns from the relevant tables required to construct an accurate SQL query. -The database schema includes tables, columns, primary keys, foreign keys, relationships, and any relevant constraints. +The database schema includes structural, semantic, and business modeling metadata: +- Models are logical datasets backed by physical tables or SQL definitions. +- Columns are exposed fields, including renamed fields, expressions, primary keys, and calculated fields. +- Relationships are reusable join logic between models. +- Calculated fields are business logic defined once and reused across queries. +- Views are named SQL statements that behave like stable virtual tables. +- Metrics are structured aggregation objects with measures and dimensions. ### INSTRUCTIONS ### 1. Carefully analyze the schema and identify the essential tables and columns needed to answer the question. @@ -39,6 +228,23 @@ 5. The number of columns chosen must match the number of reasoning. 6. Final chosen columns must be only column names, don't prefix it with table names. 7. If the chosen column is a child column of a STRUCT type column, choose the parent column instead of the child column. +8. Map the business question to the modeled datasets whose descriptions, aliases, columns, calculated fields, views, metrics, and relationships support the intent. +9. Prefer modeled analytical interfaces such as views and metrics when they expose the fields needed to answer the question. +10. If the answer needs fields, filters, time dimensions, ordering, aggregations, or relationship keys from multiple related datasets, include every required related dataset and the columns needed from each one. +11. Reuse calculated fields and metric measures or dimensions when they already represent the requested business concept. +12. Follow only the relationships shown in the provided schema when selecting columns across datasets. +13. Do not stop at a single top candidate when the question needs multiple related datasets. +14. If the same business concept is represented by multiple modeled datasets, select each relevant dataset and the fields needed to answer the shared intent. +15. If multiple modeled datasets expose compatible fields for the same requested result shape, keep each relevant dataset available so SQL generation can combine them as separate result rows instead of discarding all but one. +16. Prefer the set of deployed models, views, metrics, columns, and relationships that best support the current question. +17. If WREN RETRIEVED SEMANTIC CONTEXT is present, use sql_table_name_use_exactly and sql_column_name_use_exactly values as the exact names to return. +18. Use semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier only to understand meaning. Do not return descriptions, labels, source metadata, or rewritten variants as table or column names. +19. Prefer tables and columns whose supplied names, descriptions, relationships, metrics, or sample values directly support the requested entities, measures, filters, dates, identifiers, and dimensions. Do not answer from generic log, file, JSON, payload, text, or app-metric columns when retrieved schema metadata provides specific modeled columns for the same requested concept. +20. Compare the user's requested entities, measures, filters, dates, and dimensions only with schema metadata supplied for the active project. Do not use built-in business synonym lists. +21. If a table only contains generic data/payload/text fields and another table exposes exact business columns that match the request, choose the business table instead of searching the generic field with LIKE. +22. Never return placeholder table or column names or any user-worded identifier unless the exact same identifier appears in the provided CREATE TABLE statement or identifier contract. +23. If a requested measure, dimension, filter, or time field is not represented by retrieved schema metadata, leave it unsupported instead of substituting a similar-looking field. +24. Metric intent such as count, sum, average, minimum, maximum, ranking, date bucketing, and grouping must be satisfied by declared columns or metric fields from the retrieved schema. ### FINAL ANSWER FORMAT ### Please provide your response as a JSON object, structured as follows: @@ -99,46 +305,504 @@ """ +def _project_filter_conditions( + project_id: str | None, + mdl_hash: str | None = None, +) -> list[dict[str, Any]]: + project_deploy_filter = build_project_deploy_filter( + project_id=project_id, + mdl_hash=mdl_hash, + ) + return project_deploy_filter["conditions"] if project_deploy_filter else [] + + def _build_metric_ddl(content: dict) -> str: - columns_ddl = [ - f"{column['comment']}{column['name']} {get_engine_supported_data_type(column['data_type'])}" + columns = [ + column for column in content["columns"] - if column["data_type"].lower() - != "unknown" # quick fix: filtering out UNKNOWN column type + if column["data_type"].lower() != "unknown" + ] + context = _format_semantic_context( + { + "object_type": "metric", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable analytical aggregation interface", + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "semantic_context_not_sql_identifier": column["comment"], + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column['data_type'])}" + for column in columns ] return ( - f"{content['comment']}CREATE TABLE {content['name']} (\n " + f"{context}CREATE TABLE {content['name']} (\n " + ",\n ".join(columns_ddl) + "\n);" ) +def _strip_identifier_quotes(identifier: str | None) -> str | None: + if not identifier: + return identifier + + return identifier.strip().strip('"`[]') + + +def _view_columns_from_statement(statement: str) -> list[dict]: + if not statement: + return [] + + parsed = sqlparse.parse(statement) + if not parsed: + return [] + + statement_tokens = parsed[0].tokens + select_seen = False + output_columns: list[str] = [] + + for token in statement_tokens: + if token.is_whitespace or token.ttype in Comment: + continue + + if token.ttype is DML and token.normalized == "SELECT": + select_seen = True + continue + + if not select_seen: + continue + + if token.ttype is Keyword and token.normalized == "FROM": + break + + identifiers: list[Identifier] = [] + if isinstance(token, IdentifierList): + identifiers.extend( + identifier + for identifier in token.get_identifiers() + if isinstance(identifier, Identifier) + ) + elif isinstance(token, Identifier): + identifiers.append(token) + + for identifier in identifiers: + column_name = _strip_identifier_quotes( + identifier.get_alias() or identifier.get_real_name() + ) + if column_name and column_name != "*": + output_columns.append(column_name) + + deduplicated_columns = list(dict.fromkeys(output_columns)) + return [ + { + "name": column_name, + "data_type": "VARCHAR", + "comment": "Output column declared by the view statement.", + } + for column_name in deduplicated_columns + ] + + +def _source_table_identifier(table_reference: dict | None) -> str: + if not isinstance(table_reference, dict): + return "" + + parts = [ + str(table_reference.get(part, "")).strip() + for part in ("schema", "table") + ] + parts = [part for part in parts if part] + return ".".join(parts) + + def _build_view_ddl(content: dict) -> str: + columns = [ + column + for column in content.get("columns", []) + if column.get("name") and column.get("data_type", "").lower() != "unknown" + ] + statement = content.get("statement", "") + if not columns: + columns = _view_columns_from_statement(statement) + + context = _format_semantic_context( + { + "object_type": "view", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [column["name"] for column in columns], + }, + "semantic_context_not_sql_identifiers": { + "role": "stable virtual table interface", + "description": content["comment"], + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type( + column.get("data_type") + ), + "semantic_context_not_sql_identifier": column.get("comment", ""), + } + for column in columns + ], + } + ) + columns_ddl = [ + f"{column['name']} {get_engine_supported_data_type(column.get('data_type'))}" + for column in columns + ] + + return ( + f"{context}CREATE TABLE {content['name']} (\n " + + ",\n ".join(columns_ddl) + + "\n);" + ) + + +def _format_semantic_context(context: dict) -> str: return ( - f"{content['comment']}CREATE VIEW {content['name']}\nAS {content['statement']}" + "/*\n" + "WREN RETRIEVED SEMANTIC CONTEXT\n" + f"{orjson.dumps(context).decode('utf-8')}\n" + f"{_format_identifier_contract(context)}" + "Only values in sql_identifier_contract, sql_column_name_use_exactly, and identifiers declared in the following DDL are executable in Wren SQL.\n" + "Values under semantic_context_not_sql_identifiers and semantic_context_not_sql_identifier explain meaning only and must not be copied, combined, or rewritten as executable SQL identifiers.\n" + "*/\n" + f"{_format_executable_identifier_catalog(context)}" + ) + + +def _format_executable_identifier_catalog(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "### EXECUTABLE WREN IDENTIFIER CATALOG ###", + "Copy SQL identifiers only from this catalog or the following DDL.", + "Do not create identifiers from user wording, semantic descriptions, display labels, source names, physical names, failed SQL, or reasoning text.", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"table: {table_name}") + if column_names: + lines.append("columns:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationships:") + lines.extend(f"- {constraint}" for constraint in relationship_constraints) + lines.extend( + [ + "If a needed table, column, or relationship is not listed here or declared in the following DDL, return null for sql.", + "### END EXECUTABLE WREN IDENTIFIER CATALOG ###", + "", + ] + ) + return "\n".join(lines) + + +def _format_identifier_contract(context: dict) -> str: + contract = context.get("sql_identifier_contract", {}) + table_name = contract.get("sql_table_name_use_exactly") + column_names = contract.get("sql_column_names_use_exactly") or [ + column["sql_column_name_use_exactly"] + for column in context.get("columns", []) + if column.get("sql_column_name_use_exactly") + ] + relationship_constraints = contract.get("relationship_constraints_use_exactly") or [ + relationship["sql_relationship_constraint_use_exactly"] + for relationship in context.get("relationships", []) + if relationship.get("sql_relationship_constraint_use_exactly") + ] + + lines = [ + "WREN SQL IDENTIFIER CONTRACT", + f"object_type: {context.get('object_type', '')}", + ] + if table_name: + lines.append(f"sql_table_name_use_exactly: {table_name}") + if column_names: + lines.append("sql_column_names_use_exactly:") + lines.extend(f"- {column_name}" for column_name in column_names) + if relationship_constraints: + lines.append("relationship_constraints_use_exactly:") + lines.extend( + f"- {relationship_constraint}" + for relationship_constraint in relationship_constraints + ) + lines.extend( + [ + "Only the identifiers listed in this contract and the identifiers declared in the following DDL are executable.", + "Semantic descriptions, source names, aliases, examples, and user wording are not executable identifiers.", + "END WREN SQL IDENTIFIER CONTRACT", + "", + ] + ) + return "\n".join(lines) + + +def _included_relationship_columns(content: dict, tables: Optional[set[str]]) -> set: + relationship_columns = { + column.get("column") + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + } + relationship_columns.discard(None) + return relationship_columns + + +def _included_columns( + content: dict, columns: Optional[set[str]], tables: Optional[set[str]] +) -> list[dict]: + relationship_columns = _included_relationship_columns(content, tables) + return [ + column + for column in content["columns"] + if column["type"] == "COLUMN" + and ( + not columns + or column["name"] in columns + or column["name"] in relationship_columns + or column["is_primary_key"] + ) + and column["data_type"].lower() != "unknown" + ] + + +def _included_relationships(content: dict, tables: Optional[set[str]]) -> list[dict]: + return [ + column + for column in content["columns"] + if column["type"] == "FOREIGN_KEY" + and (not tables or set(column.get("tables", [])).issubset(tables)) + ] + + +def _build_table_retrieval_context( + content: dict, columns: Optional[set[str]] = None, tables: Optional[set[str]] = None +) -> tuple[str, bool, bool]: + ddl, has_calculated_field, has_json_field = build_table_ddl( + content, + columns=columns, + tables=tables, + include_semantic_comments=False, ) + included_columns = _included_columns(content, columns, tables) + included_relationships = _included_relationships(content, tables) + table_properties = content.get("properties") or {} + table_reference = content.get("tableReference") or {} + context = _format_semantic_context( + { + "object_type": "model", + "sql_identifier_contract": { + "sql_table_name_use_exactly": content["name"], + "sql_column_names_use_exactly": [ + column["name"] for column in included_columns + ], + "relationship_constraints_use_exactly": [ + relationship["constraint"] + for relationship in included_relationships + ], + }, + "semantic_context_not_sql_identifiers": { + "description": content["comment"], + "display_name": table_properties.get("displayName"), + "source_table_name": _source_table_identifier(table_reference), + "source_table_reference": table_reference, + }, + "columns": [ + { + "sql_column_name_use_exactly": column["name"], + "data_type": get_engine_supported_data_type(column["data_type"]), + "is_primary_key": column["is_primary_key"], + "semantic_context_not_sql_identifier": column["comment"], + "display_name": (column.get("properties") or {}).get( + "displayName" + ), + "source_column_name": (column.get("properties") or {}).get( + "sourceColumnName" + ), + } + for column in included_columns + ], + "relationships": [ + { + "semantic_context_not_sql_identifier": relationship["comment"], + "sql_relationship_constraint_use_exactly": relationship[ + "constraint" + ], + "related_models_use_exactly": relationship.get("tables", []), + } + for relationship in included_relationships + ], + } + ) + return f"{context}{ddl}", has_calculated_field, has_json_field + + +def _identifier_context(table_name: str, column_names: list[str]) -> str: + return "\n".join( + [f"table: {table_name}", "columns:", *[f"- {name}" for name in column_names]] + ) + + +def _build_retrieval_item(table_schema: dict) -> tuple[dict[str, str], bool, bool]: + ddl, has_calculated_field, has_json_field = _build_table_retrieval_context( + table_schema + ) + return ( + { + "table_name": table_schema["name"], + "table_ddl": ddl, + "identifier_context": _identifier_context( + table_schema["name"], + [ + column["name"] + for column in _included_columns(table_schema, None, None) + ], + ), + }, + has_calculated_field, + has_json_field, + ) + + +def _build_pruning_context(content: dict) -> str: + if content["type"] == "TABLE": + return _build_table_retrieval_context(content)[0] + if content["type"] == "METRIC": + return _build_metric_ddl(content) + if content["type"] == "VIEW": + return _build_view_ddl(content) + return "" + + +def _fallback_retrieval_results( + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], +) -> dict[str, Any]: + retrieval_results = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE": + retrieval_item, _has_calculated_field, _has_json_field = ( + _build_retrieval_item(table_schema) + ) + retrieval_results.append(retrieval_item) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + for document in dbschema_retrieval: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + ), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), + } + ) + + return { + "db_schemas": retrieval_results, + "tokens": _token_count, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + + +def _empty_retrieval_results() -> dict[str, Any]: + return { + "retrieval_results": [], + "has_calculated_field": False, + "has_metric": False, + "has_json_field": False, + } + + +def _merge_names(*name_lists: list[str]) -> list[str]: + merged: list[str] = [] + seen: set[str] = set() + for names in name_lists: + for name in names: + if name in seen: + continue + merged.append(name) + seen.add(name) + return merged ## Start of Pipeline @observe(capture_input=False, capture_output=False) async def embedding(query: str, embedder: Any, histories: list[AskHistory]) -> dict: if query: - if histories: - previous_query_summaries = [history.question for history in histories] - else: - previous_query_summaries = [] - - query = "\n".join(previous_query_summaries) + "\n" + query - - return await embedder.run(query) + return await embedder.run(_augment_retrieval_query(query)) else: return {} @observe(capture_input=False) async def table_retrieval( - embedding: dict, project_id: str, tables: list[str], table_retriever: Any + embedding: dict, + project_id: str, + tables: list[str], + table_retriever: Any, + mdl_hash: str | None = None, ) -> dict: filters = { "operator": "AND", @@ -147,10 +811,7 @@ async def table_retrieval( ], } - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} - ) + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) if embedding: return await table_retriever.run( @@ -170,37 +831,783 @@ async def table_retrieval( @observe(capture_input=False) async def dbschema_retrieval( - table_retrieval: dict, project_id: str, dbschema_retriever: Any + table_retrieval: dict, + project_id: str, + dbschema_retriever: Any, + query: str | None = None, + embedding: dict | None = None, + mdl_hash: str | None = None, + include_related_models: bool = True, +) -> list[Document]: + table_names = _table_names_from_description_documents( + table_retrieval.get("documents", []) + ) + documents = [] + if embedding: + semantic_started_at = time.perf_counter() + semantic_documents = await _retrieve_semantic_schema_documents( + embedding, project_id, mdl_hash, dbschema_retriever + ) + _log_retrieval_timing( + "schema_retrieval_semantic", + semantic_started_at, + project_id, + document_count=len(semantic_documents), + ) + semantic_table_names = _table_names_from_schema_documents(semantic_documents)[ + :_SEMANTIC_TABLE_NAME_MERGE_LIMIT + ] + table_names = _merge_names(table_names, semantic_table_names)[ + :_MAX_RETRIEVED_TABLE_NAMES + ] + lexical_started_at = time.perf_counter() + lexical_documents, lexical_table_names = await _retrieve_lexical_schema_hits( + query=query, + project_id=project_id, + mdl_hash=mdl_hash, + dbschema_retriever=dbschema_retriever, + existing_table_names=set(table_names), + ) + _log_retrieval_timing( + "schema_retrieval_lexical_scan", + lexical_started_at, + project_id, + document_count=len(lexical_documents), + table_count=len(lexical_table_names), + ) + table_names = _merge_names(table_names, lexical_table_names)[ + :_MAX_RETRIEVED_TABLE_NAMES + ] + ranking_started_at = time.perf_counter() + table_names = _rank_table_names_by_query( + table_names, + _dedupe_documents(semantic_documents + lexical_documents), + query, + ) + _log_retrieval_timing( + "candidate_ranking", + ranking_started_at, + project_id, + candidate_count=len(table_names), + ) + selected_semantic_table_names = set(semantic_table_names) + documents = [ + document + for document in semantic_documents + if document.meta.get("name") in selected_semantic_table_names + ] + documents = _dedupe_documents(documents + lexical_documents) + + if table_names: + if include_related_models: + expansion_started_at = time.perf_counter() + retrieved_table_names = set() + pending_table_names = table_names + remaining_expansion_depth = _MAX_RELATED_TABLE_EXPANSION_DEPTH + + while pending_table_names: + retrieved_table_names.update(pending_table_names) + retrieved_documents = await _retrieve_schema_documents( + pending_table_names, project_id, mdl_hash, dbschema_retriever + ) + documents = _dedupe_documents(documents + retrieved_documents) + if remaining_expansion_depth <= 0: + break + remaining_expansion_depth -= 1 + remaining_slots = _MAX_RETRIEVED_TABLE_NAMES - len( + retrieved_table_names + ) + if remaining_slots <= 0: + break + pending_table_names = [ + table_name + for table_name in _related_table_names(documents) + if table_name not in retrieved_table_names + ][:remaining_slots] + + ranking_started_at = time.perf_counter() + ranked_documents = _rank_documents_for_query(documents, table_names, query) + _log_retrieval_timing( + "schema_retrieval_related_expansion", + expansion_started_at, + project_id, + table_count=len(retrieved_table_names), + document_count=len(documents), + ) + _log_retrieval_timing( + "candidate_ranking", + ranking_started_at, + project_id, + candidate_count=len(ranked_documents), + ) + logger.info( + "Ask schema retrieval project_id=%s retrieved_tables=%s", + project_id, + [ + { + "table": document.meta.get("name"), + "score": getattr(document, "score", None), + } + for document in ranked_documents + ], + ) + return ranked_documents + + named_started_at = time.perf_counter() + retrieved_documents = await _retrieve_schema_documents( + table_names, project_id, mdl_hash, dbschema_retriever + ) + _log_retrieval_timing( + "schema_retrieval_named_fetch", + named_started_at, + project_id, + table_count=len(table_names), + document_count=len(retrieved_documents), + ) + documents = _dedupe_documents(documents + retrieved_documents) + ranking_started_at = time.perf_counter() + ranked_documents = _rank_documents_for_query(documents, table_names, query) + _log_retrieval_timing( + "candidate_ranking", + ranking_started_at, + project_id, + candidate_count=len(ranked_documents), + ) + logger.info( + "Ask schema retrieval project_id=%s retrieved_tables=%s", + project_id, + [ + { + "table": document.meta.get("name"), + "score": getattr(document, "score", None), + } + for document in ranked_documents + ], + ) + return ranked_documents + + logger.info("Ask schema retrieval project_id=%s retrieved_tables=[]", project_id) + return [] + + +def _tokenize_schema_text(value: Any) -> set[str]: + if value is None: + return set() + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + return _expand_schema_token_variants(set(_RANK_TOKEN.findall(text.lower()))) + + +def _normalized_schema_mention_text(value: Any) -> str: + if value is None: + return "" + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(value)) + tokens = _RANK_TOKEN.findall(text.lower()) + return f" {' '.join(tokens)} " if tokens else "" + + +def _schema_identifier_mention_variants(identifier: str) -> set[str]: + text = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", str(identifier)) + tokens = _RANK_TOKEN.findall(text.lower()) + variants = {" ".join(tokens)} if tokens else set() + if len(tokens) > 1: + variants.add(" ".join(tokens[1:])) + return {variant for variant in variants if variant} + + +def _query_mentions_schema_identifier(normalized_query: str, identifier: str) -> bool: + if not normalized_query: + return False + return any( + f" {variant} " in normalized_query + for variant in _schema_identifier_mention_variants(identifier) + ) + + +def _rank_content_tokens(query_tokens: set[str]) -> set[str]: + return { + token + for token in query_tokens + if token not in _RANK_GENERIC_QUERY_TOKENS and not token.isdigit() + } + + +def _schema_token_variants(token: str) -> set[str]: + token = token.lower() + variants = {token} + if len(token) > 4 and token.endswith("ies"): + variants.add(token[:-3] + "y") + elif len(token) > 4 and token.endswith("es"): + if token.endswith(("ches", "shes", "sses", "uses", "xes", "zes")): + variants.add(token[:-2]) + else: + variants.add(token[:-1]) + elif len(token) > 3 and token.endswith("s") and not token.endswith(("ss", "us")): + variants.add(token[:-1]) + return {variant for variant in variants if variant} + + +def _expand_schema_token_variants(tokens: set[str]) -> set[str]: + expanded: set[str] = set() + for token in tokens: + expanded.update(_schema_token_variants(token)) + return expanded + + +def _compound_identifier_tokens(token: str) -> set[str]: + if len(token) < 5: + return set() + + tokens: set[str] = set() + for part in _COMPOUND_IDENTIFIER_PART_TOKENS: + if part != token and len(part) >= 4 and part in token: + tokens.add(part) + + for alias, expansions in _COMPOUND_IDENTIFIER_ALIASES.items(): + if alias != token and (token.startswith(alias) or token.endswith(alias)): + tokens.add(alias) + tokens.update(expansions) + + return tokens + + +def _tokenize_schema_identifier_text(value: Any) -> set[str]: + tokens = _tokenize_schema_text(value) + for token in list(tokens): + tokens.update(_compound_identifier_tokens(token)) + return _expand_schema_token_variants(tokens) + + +def _tokenize_nested_schema_identifier_text(value: Any) -> set[str]: + if value is None: + return set() + if isinstance(value, dict): + tokens: set[str] = set() + for nested_key, nested_value in value.items(): + tokens.update(_tokenize_schema_identifier_text(nested_key)) + tokens.update(_tokenize_nested_schema_identifier_text(nested_value)) + return tokens + if isinstance(value, (list, tuple, set)): + tokens: set[str] = set() + for nested_value in value: + tokens.update(_tokenize_nested_schema_identifier_text(nested_value)) + return tokens + return _tokenize_schema_identifier_text(value) + + +def _tokenize_nested_schema_text(value: Any) -> set[str]: + if value is None: + return set() + if isinstance(value, dict): + tokens: set[str] = set() + for nested_key, nested_value in value.items(): + tokens.update(_tokenize_schema_text(nested_key)) + tokens.update(_tokenize_nested_schema_text(nested_value)) + return tokens + if isinstance(value, (list, tuple, set)): + tokens: set[str] = set() + for nested_value in value: + tokens.update(_tokenize_nested_schema_text(nested_value)) + return tokens + return _tokenize_schema_text(value) + + +def _schema_rank_document_key(documents: list[Document]) -> tuple[tuple[str, str, str], ...]: + return tuple( + ( + str(document.meta.get("type", "")), + str(document.meta.get("name", "")), + str(document.content), + ) + for document in documents + ) + + +def _schema_rank_text_by_table(documents: list[Document]) -> dict[str, dict[str, set[str]]]: + return _cached_schema_rank_text_by_table(_schema_rank_document_key(documents)) + + +@lru_cache(maxsize=128) +def _cached_schema_rank_text_by_table( + document_key: tuple[tuple[str, str, str], ...], +) -> dict[str, dict[str, set[str]]]: + table_text: dict[str, dict[str, set[str]]] = {} + + def ensure(table_name: str) -> dict[str, set[str]]: + if table_name not in table_text: + table_text[table_name] = { + "table": set(), + "columns": set(), + "comments": set(), + } + return table_text[table_name] + + for _, meta_name, content_text in document_key: + try: + content = ast.literal_eval(content_text) + except (SyntaxError, ValueError): + continue + + table_name = meta_name or content.get("name") + if not table_name: + continue + + bucket = ensure(table_name) + bucket["table"].update(_tokenize_schema_identifier_text(table_name)) + bucket["table"].update(_tokenize_schema_identifier_text(content.get("name"))) + bucket["table"].update( + _tokenize_nested_schema_identifier_text(content.get("properties")) + ) + bucket["table"].update( + _tokenize_nested_schema_identifier_text(content.get("tableReference")) + ) + bucket["comments"].update(_tokenize_schema_text(content.get("comment"))) + bucket["comments"].update(_tokenize_schema_text(content.get("description"))) + bucket["comments"].update(_tokenize_nested_schema_text(content.get("refSql"))) + + for column in content.get("columns", []) or []: + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("name")) + ) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("column")) + ) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("display_name")) + ) + bucket["columns"].update( + _tokenize_schema_identifier_text(column.get("displayName")) + ) + bucket["columns"].update( + _tokenize_nested_schema_identifier_text(column.get("properties")) + ) + bucket["comments"].update(_tokenize_schema_text(column.get("comment"))) + bucket["comments"].update(_tokenize_schema_text(column.get("description"))) + + return table_text + + +def _rank_table_names_by_query( + table_names: list[str], + semantic_documents: list[Document], + query: str | None, + require_positive_score: bool = False, +) -> list[str]: + if not query or not table_names: + return [] if require_positive_score else table_names + + query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) + if not query_tokens: + return [] if require_positive_score else table_names + + table_text = _schema_rank_text_by_table(semantic_documents) + normalized_query = _normalized_schema_mention_text(query) + content_tokens = _rank_content_tokens(query_tokens) + + def score(table_name: str) -> int: + bucket = table_text.get(table_name, {}) + table_tokens = set(bucket.get("table", set())) | _tokenize_schema_identifier_text( + table_name + ) + column_tokens = set(bucket.get("columns", set())) + comment_tokens = set(bucket.get("comments", set())) + direct_table_matches = query_tokens & table_tokens + direct_column_matches = query_tokens & column_tokens + direct_comment_matches = query_tokens & comment_tokens + covered_tokens = direct_table_matches | direct_column_matches | direct_comment_matches + direct_content_matches = content_tokens & covered_tokens + direct_content_column_matches = content_tokens & column_tokens + + value = ( + len(direct_table_matches) * 6 + + len(direct_column_matches) * 8 + + len(direct_comment_matches) + + len(direct_content_matches) * 32 + + len(direct_content_column_matches) * 24 + ) + if _query_mentions_schema_identifier(normalized_query, table_name): + value += 80 + value += len( + direct_table_matches | direct_column_matches | direct_comment_matches + ) ** 2 + if content_tokens and all( + _schema_token_variants(token) & covered_tokens + for token in content_tokens + ): + value += 80 + len(content_tokens) * 16 + + if len(direct_column_matches) >= 2: + value += 8 + if direct_table_matches and direct_column_matches: + value += 8 + return value + + scored = [ + (index, table_name, score(table_name)) + for index, table_name in enumerate(table_names) + ] + if require_positive_score: + scored = [item for item in scored if item[2] > 0] + + ranked = sorted(scored, key=lambda item: (-item[2], item[0])) + return [table_name for _, table_name, _ in ranked] + + +def _rank_documents_by_table_names( + documents: list[Document], + table_names: list[str], +) -> list[Document]: + table_rank = {table_name: index for index, table_name in enumerate(table_names)} + return sorted( + documents, + key=lambda document: ( + table_rank.get(document.meta.get("name"), len(table_rank)), + document.meta.get("type", ""), + ), + ) + + +def _rank_documents_for_query( + documents: list[Document], + table_names: list[str], + query: str | None, +) -> list[Document]: + ranked_table_names = _rank_table_names_by_query( + _merge_names(_table_names_from_schema_documents(documents), table_names), + documents, + query, + ) + return _rank_documents_by_table_names(documents, ranked_table_names) + + +def _augment_retrieval_query(query: str) -> str: + return query + + +async def _retrieve_semantic_schema_documents( + embedding: dict, + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + + results = await dbschema_retriever.run( + query_embedding=embedding.get("embedding"), + filters=filters, + ) + return results["documents"] + + +async def _retrieve_all_schema_documents( + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, ) -> list[Document]: - tables = table_retrieval.get("documents", []) + cache_key = (id(dbschema_retriever), project_id, mdl_hash or "") + if cache_key in _ALL_SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=all count=%s", + project_id, + mdl_hash or "", + len(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + lock = _ALL_SCHEMA_DOCUMENTS_CACHE_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + if cache_key in _ALL_SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=all count=%s", + project_id, + mdl_hash or "", + len(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_ALL_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + started_at = time.perf_counter() + documents = await _retrieve_all_schema_documents_uncached( + project_id, + mdl_hash, + dbschema_retriever, + ) + if len(_ALL_SCHEMA_DOCUMENTS_CACHE) >= 64: + _ALL_SCHEMA_DOCUMENTS_CACHE.pop(next(iter(_ALL_SCHEMA_DOCUMENTS_CACHE))) + _ALL_SCHEMA_DOCUMENTS_CACHE[cache_key] = documents + _log_retrieval_timing( + "schema_retrieval_all_documents", + started_at, + project_id, + cache_hit=False, + document_count=len(documents), + ) + return list(documents) + + +async def _retrieve_all_schema_documents_uncached( + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + ], + } + + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] + + +async def _retrieve_lexical_schema_hits( + query: str | None, + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, + existing_table_names: set[str], +) -> tuple[list[Document], list[str]]: + if not query: + return [], [] + + documents = await _retrieve_all_schema_documents( + project_id, + mdl_hash, + dbschema_retriever, + ) + candidate_table_names = [ + table_name + for table_name in _rank_table_names_by_query( + _table_names_from_schema_documents(documents), + documents, + query, + require_positive_score=True, + ) + if table_name not in existing_table_names + ][:_LEXICAL_SCHEMA_TABLE_NAME_MERGE_LIMIT] + if not candidate_table_names: + return [], [] + + logger.info( + "Ask schema lexical project scan project_id=%s mdl_hash=%s retrieved_tables=%s", + project_id, + mdl_hash or "", + candidate_table_names, + ) + candidate_table_name_set = set(candidate_table_names) + return [ + document + for document in documents + if document.meta.get("name") in candidate_table_name_set + ], candidate_table_names + + +def _table_names_from_schema_documents(documents: list[Document]) -> list[str]: table_names = [] - for table in tables: - content = ast.literal_eval(table.content) - table_names.append(content["name"]) + seen = set() + + for document in documents: + table_name = document.meta.get("name") + if not table_name: + content = ast.literal_eval(document.content) + table_name = content.get("name") + if table_name and table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + + +def _table_names_from_description_documents(documents: list[Document]) -> list[str]: + table_names = [] + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + table_name = content["name"] + if table_name not in seen: + table_names.append(table_name) + seen.add(table_name) + + return table_names + + +async def _retrieve_schema_documents( + table_names: list[str], + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, +) -> list[Document]: + table_names = list(dict.fromkeys(table_names)) table_name_conditions = [ {"field": "name", "operator": "==", "value": table_name} for table_name in table_names ] - if table_name_conditions: - filters = { - "operator": "AND", - "conditions": [ - {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, - {"operator": "OR", "conditions": table_name_conditions}, - ], - } + if not table_name_conditions: + return [] - if project_id: - filters["conditions"].append( - {"field": "project_id", "operator": "==", "value": project_id} + cache_key = ( + id(dbschema_retriever), + project_id, + mdl_hash or "", + tuple(table_names), + ) + if cache_key in _SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=named table_count=%s document_count=%s", + project_id, + mdl_hash or "", + len(table_names), + len(_SCHEMA_DOCUMENTS_CACHE[cache_key]), + ) + return list(_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + lock = _SCHEMA_DOCUMENTS_CACHE_LOCKS.setdefault(cache_key, asyncio.Lock()) + async with lock: + if cache_key in _SCHEMA_DOCUMENTS_CACHE: + logger.info( + "Ask schema document cache hit project_id=%s mdl_hash=%s scope=named table_count=%s document_count=%s", + project_id, + mdl_hash or "", + len(table_names), + len(_SCHEMA_DOCUMENTS_CACHE[cache_key]), ) + return list(_SCHEMA_DOCUMENTS_CACHE[cache_key]) + + started_at = time.perf_counter() + documents = await _retrieve_schema_documents_uncached( + table_names, + project_id, + mdl_hash, + dbschema_retriever, + table_name_conditions, + ) + if len(_SCHEMA_DOCUMENTS_CACHE) >= 128: + _SCHEMA_DOCUMENTS_CACHE.pop(next(iter(_SCHEMA_DOCUMENTS_CACHE))) + _SCHEMA_DOCUMENTS_CACHE[cache_key] = documents + _log_retrieval_timing( + "schema_retrieval_named_documents", + started_at, + project_id, + cache_hit=False, + table_count=len(table_names), + document_count=len(documents), + ) + return list(documents) - results = await dbschema_retriever.run(query_embedding=[], filters=filters) - return results["documents"] - return [] +async def _retrieve_schema_documents_uncached( + table_names: list[str], + project_id: str, + mdl_hash: str | None, + dbschema_retriever: Any, + table_name_conditions: list[dict[str, Any]], +) -> list[Document]: + filters = { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"operator": "OR", "conditions": table_name_conditions}, + ], + } + + filters["conditions"].extend(_project_filter_conditions(project_id, mdl_hash)) + + results = await dbschema_retriever.run(query_embedding=[], filters=filters) + return results["documents"] + + +def _related_table_names(documents: list[Document]) -> list[str]: + related_table_names = [] + seen = set() + + for document in documents: + content = ast.literal_eval(document.content) + if content.get("type") != "TABLE_COLUMNS": + continue + + for column in content.get("columns", []): + if column.get("type") != "FOREIGN_KEY": + continue + + for table_name in column.get("tables", []): + if table_name not in seen: + related_table_names.append(table_name) + seen.add(table_name) + + return related_table_names + + +def _dedupe_documents(documents: list[Document]) -> list[Document]: + deduped = [] + seen = set() + + for document in documents: + identity = ( + document.meta.get("type"), + document.meta.get("name"), + document.content, + ) + if identity in seen: + continue + deduped.append(document) + seen.add(identity) + + return deduped + + +def _limit_retrieval_results_for_generation( + retrieval_results: list[dict[str, Any]], + encoding: tiktoken.Encoding, +) -> tuple[list[dict[str, Any]], int, int, str | None]: + if not retrieval_results: + return retrieval_results, 0, 0, None + + original_tokens = len( + encoding.encode( + " ".join( + retrieval_result.get("table_ddl", "") + for retrieval_result in retrieval_results + ) + ) + ) + limited_results: list[dict[str, Any]] = [] + limited_tokens = 0 + skipped_for_token_budget = False + + for retrieval_result in retrieval_results: + if len(limited_results) >= _MAX_LLM_SCHEMA_CONTEXT_TABLES: + break + + table_tokens = len(encoding.encode(retrieval_result.get("table_ddl", ""))) + if table_tokens > _MAX_LLM_SCHEMA_CONTEXT_TOKENS: + skipped_for_token_budget = True + continue + if limited_tokens + table_tokens > _MAX_LLM_SCHEMA_CONTEXT_TOKENS: + skipped_for_token_budget = True + continue + + limited_results.append(retrieval_result) + limited_tokens += table_tokens + + if not limited_results: + return [], original_tokens, 0, "all_tables_exceed_token_budget" + + if len(limited_results) == len(retrieval_results): + return retrieval_results, original_tokens, original_tokens, None + + reason = "ranked_top_k_skipped_token_budget" if skipped_for_token_budget else "ranked_top_k" + return limited_results, original_tokens, limited_tokens, reason @observe() @@ -224,9 +1631,16 @@ def construct_db_schemas(dbschema_retrieval: list[Document]) -> list[dict]: db_schemas[document.meta["name"]]["columns"] = content["columns"] else: db_schemas[document.meta["name"]]["columns"] += content["columns"] + elif content["type"] in {"VIEW", "METRIC"}: + db_schemas[document.meta["name"]] = content # remove incomplete schemas - db_schemas = {k: v for k, v in db_schemas.items() if "type" in v and "columns" in v} + db_schemas = { + k: v + for k, v in db_schemas.items() + if v.get("type") in {"VIEW", "METRIC"} + or (v.get("type") == "TABLE" and "columns" in v) + } return list(db_schemas.values()) @@ -246,11 +1660,20 @@ def check_using_db_schemas_without_pruning( for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE": - ddl, _has_calculated_field, _has_json_field = build_table_ddl(table_schema) + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) retrieval_results.append( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_context": _identifier_context( + table_schema["name"], + [ + column["name"] + for column in _included_columns(table_schema, None, None) + ], + ), } ) if _has_calculated_field: @@ -266,6 +1689,14 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_metric_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content["columns"] + if column["data_type"].lower() != "unknown" + ], + ), } ) has_metric = True @@ -274,6 +1705,15 @@ def check_using_db_schemas_without_pruning( { "table_name": content["name"], "table_ddl": _build_view_ddl(content), + "identifier_context": _identifier_context( + content["name"], + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), } ) @@ -290,6 +1730,41 @@ def check_using_db_schemas_without_pruning( "has_json_field": has_json_field, } + ( + limited_retrieval_results, + original_token_count, + limited_token_count, + limit_reason, + ) = _limit_retrieval_results_for_generation(retrieval_results, encoding) + if limit_reason: + if not limited_retrieval_results: + logger.info( + "Ask retrieval selected schema context exceeded generation budget; using column pruning reason=%s original_tables=%s original_tokens=%s token_budget=%s", + limit_reason, + len(retrieval_results), + original_token_count, + _MAX_LLM_SCHEMA_CONTEXT_TOKENS, + ) + return { + "db_schemas": [], + "tokens": original_token_count, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + logger.info( + "Ask retrieval capped generation schema context reason=%s original_tables=%s selected_tables=%s original_tokens=%s selected_tokens=%s table_budget=%s token_budget=%s", + limit_reason, + len(retrieval_results), + len(limited_retrieval_results), + original_token_count, + limited_token_count, + _MAX_LLM_SCHEMA_CONTEXT_TABLES, + _MAX_LLM_SCHEMA_CONTEXT_TOKENS, + ) + retrieval_results = limited_retrieval_results + _token_count = limited_token_count + return { "db_schemas": retrieval_results, "tokens": _token_count, @@ -308,23 +1783,21 @@ def prompt( histories: list[AskHistory], ) -> dict: if not check_using_db_schemas_without_pruning["db_schemas"]: - db_schemas = [ - build_table_ddl(construct_db_schema)[0] - for construct_db_schema in construct_db_schemas - ] - - previous_query_summaries = ( - [history.question for history in histories] if histories else [] + db_schemas = list( + filter( + None, + [ + _build_pruning_context(construct_db_schema) + for construct_db_schema in construct_db_schemas + ], + ) ) - query = "\n".join(previous_query_summaries) + "\n" + query - _prompt = prompt_builder.run(question=query, db_schemas=db_schemas) return {"prompt": clean_up_new_lines(_prompt.get("prompt"))} else: return {} - @observe(as_type="generation", capture_input=False) @trace_cost async def filter_columns_in_tables( @@ -344,32 +1817,59 @@ def construct_retrieval_results( filter_columns_in_tables: dict, construct_db_schemas: list[dict], dbschema_retrieval: list[Document], + query: str | None = None, ) -> dict[str, Any]: if filter_columns_in_tables: - columns_and_tables_needed = orjson.loads( - filter_columns_in_tables["replies"][0] - )["results"] - - # we need to change the below code to match the new schema of structured output - # the objective of this loop is to change the structure of JSON to match the needed format - reformated_json = {} - for table in columns_and_tables_needed: - reformated_json[table["table_name"]] = table["table_contents"] - columns_and_tables_needed = reformated_json + columns_and_tables_needed = _parse_column_selection_response( + filter_columns_in_tables + ) + lexical_columns_and_tables_needed = _lexical_columns_and_tables_needed( + construct_db_schemas, + query, + ) + columns_and_tables_needed = _merge_column_selection( + columns_and_tables_needed, + lexical_columns_and_tables_needed, + ) + if not columns_and_tables_needed: + logger.warning( + "Column pruning did not return grounded schema selections; " + "skipping broad schema fallback." + ) + return _empty_retrieval_results() + tables = set(columns_and_tables_needed.keys()) retrieval_results = [] + selected_schema_log = [] has_calculated_field = False has_metric = False has_json_field = False for table_schema in construct_db_schemas: if table_schema["type"] == "TABLE" and table_schema["name"] in tables: - ddl, _has_calculated_field, _has_json_field = build_table_ddl( - table_schema, - columns=set( - columns_and_tables_needed[table_schema["name"]]["columns"] - ), - tables=tables, + selected_columns = set( + columns_and_tables_needed[table_schema["name"]]["columns"] + ) + executable_columns = { + column["name"] + for column in table_schema["columns"] + if column["type"] == "COLUMN" + and column["data_type"].lower() != "unknown" + } + columns = selected_columns.intersection(executable_columns) + if selected_columns and not columns: + logger.warning( + "Column pruning selected no executable columns for %s; " + "including the full model schema to preserve grounding.", + table_schema["name"], + ) + columns = None + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context( + table_schema, + columns=columns, + tables=tables, + ) ) if _has_calculated_field: has_calculated_field = True @@ -380,29 +1880,97 @@ def construct_retrieval_results( { "table_name": table_schema["name"], "table_ddl": ddl, + "identifier_context": _identifier_context( + table_schema["name"], + [ + column["name"] + for column in _included_columns( + table_schema, columns, tables + ) + ], + ), } ) + selected_schema_log.append( + { + "table": table_schema["name"], + "columns": sorted(selected_columns), + } + ) + + if not retrieval_results: + logger.warning( + "Column-selection output did not match retrieved schemas; " + "falling back to unpruned retrieved schema context." + ) + return _build_unpruned_retrieval_results( + construct_db_schemas, dbschema_retrieval + ) for document in dbschema_retrieval: - if document.meta["name"] in columns_and_tables_needed: + try: content = ast.literal_eval(document.content) + except (ValueError, SyntaxError): + logger.warning( + "Skipping malformed retrieved schema document during schema pruning: %s", + document.meta, + ) + continue - if content["type"] == "METRIC": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_metric_ddl(content), - } - ) - has_metric = True - elif content["type"] == "VIEW": - retrieval_results.append( - { - "table_name": content["name"], - "table_ddl": _build_view_ddl(content), - } - ) + if not isinstance(content, dict): + logger.warning( + "Skipping non-object retrieved schema document during schema pruning: %s", + document.meta, + ) + continue + + content_name = content.get("name") + content_type = content.get("type") + if not content_name: + logger.warning( + "Skipping retrieved schema document without name during schema pruning: %s", + document.meta, + ) + continue + if content_name not in tables: + continue + + if content_type == "METRIC": + retrieval_results.append( + { + "table_name": content_name, + "table_ddl": _build_metric_ddl(content), + "identifier_context": _identifier_context( + content_name, + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), + } + ) + has_metric = True + elif content_type == "VIEW": + retrieval_results.append( + { + "table_name": content_name, + "table_ddl": _build_view_ddl(content), + "identifier_context": _identifier_context( + content_name, + [ + column["name"] + for column in content.get("columns", []) + if column.get("name") + and column.get("data_type", "").lower() != "unknown" + ], + ), + } + ) + + logger.info("Ask retrieval selected schema objects=%s", selected_schema_log) return { "retrieval_results": retrieval_results, "has_calculated_field": has_calculated_field, @@ -411,7 +1979,13 @@ def construct_retrieval_results( } else: retrieval_results = check_using_db_schemas_without_pruning["db_schemas"] - + logger.info( + "Ask retrieval selected schema objects=%s", + [ + {"table": retrieval_result.get("table_name"), "columns": "all"} + for retrieval_result in retrieval_results + ], + ) return { "retrieval_results": retrieval_results, "has_calculated_field": check_using_db_schemas_without_pruning[ @@ -422,27 +1996,293 @@ def construct_retrieval_results( } +def _normalize_column_selection_results(parsed_response: Any) -> list[dict]: + if isinstance(parsed_response, list): + return [item for item in parsed_response if isinstance(item, dict)] + + if not isinstance(parsed_response, dict): + return [] + + for key in ( + "results", + "tables", + "selected_tables", + "retrieval_results", + "matches", + "data", + "result", + "output", + ): + if key in parsed_response: + normalized = _normalize_column_selection_results(parsed_response[key]) + if normalized: + return normalized + + if "table_name" in parsed_response and ( + "table_contents" in parsed_response or "columns" in parsed_response + ): + return [parsed_response] + + keyed_tables = [] + for table_name, table_contents in parsed_response.items(): + if not isinstance(table_name, str) or not isinstance(table_contents, dict): + continue + if "table_contents" in table_contents: + keyed_tables.append( + { + "table_name": table_name, + "table_contents": table_contents["table_contents"], + } + ) + elif "columns" in table_contents: + keyed_tables.append( + {"table_name": table_name, "table_contents": table_contents} + ) + + return keyed_tables + + +def _parse_column_selection_response(filter_columns_in_tables: dict) -> dict: + raw_reply = (filter_columns_in_tables.get("replies") or [""])[0] + try: + parsed_response = orjson.loads(raw_reply) + except orjson.JSONDecodeError as exc: + logger.warning("Unable to parse column-selection JSON response: %s", exc) + return {} + + normalized_tables = _normalize_column_selection_results(parsed_response) + reformatted_json = {} + for table in normalized_tables: + table_name = table.get("table_name") or table.get("name") + table_contents = table.get("table_contents") or {} + if not table_contents and "columns" in table: + table_contents = table + + columns = ( + table_contents.get("columns") if isinstance(table_contents, dict) else None + ) + if not isinstance(table_name, str) or not isinstance(columns, list): + continue + + reformatted_json[table_name] = { + **table_contents, + "columns": [column for column in columns if isinstance(column, str)], + } + + if not reformatted_json: + response_shape = ( + f"keys={list(parsed_response.keys())[:8]}" + if isinstance(parsed_response, dict) + else type(parsed_response).__name__ + ) + logger.warning( + "Column-selection response did not include usable table columns (%s).", + response_shape, + ) + + return reformatted_json + + +def _build_unpruned_retrieval_results( + construct_db_schemas: list[dict], + dbschema_retrieval: list[Document], +) -> dict: + retrieval_results = [] + has_calculated_field = False + has_metric = False + has_json_field = False + + for table_schema in construct_db_schemas: + if table_schema["type"] == "TABLE": + ddl, _has_calculated_field, _has_json_field = ( + _build_table_retrieval_context(table_schema) + ) + retrieval_results.append( + { + "table_name": table_schema["name"], + "table_ddl": ddl, + } + ) + if _has_calculated_field: + has_calculated_field = True + if _has_json_field: + has_json_field = True + + for document in dbschema_retrieval: + content = ast.literal_eval(document.content) + + if content["type"] == "METRIC": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_metric_ddl(content), + } + ) + has_metric = True + elif content["type"] == "VIEW": + retrieval_results.append( + { + "table_name": content["name"], + "table_ddl": _build_view_ddl(content), + } + ) + + return { + "retrieval_results": retrieval_results, + "has_calculated_field": has_calculated_field, + "has_metric": has_metric, + "has_json_field": has_json_field, + } + + +def _merge_column_selection( + primary: dict[str, dict], + secondary: dict[str, dict], +) -> dict[str, dict]: + merged = { + table_name: { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + for table_name, table_contents in primary.items() + } + + for table_name, table_contents in secondary.items(): + if table_name not in merged: + merged[table_name] = { + **table_contents, + "columns": list(table_contents.get("columns", [])), + } + continue + + columns = list(merged[table_name].get("columns", [])) + for column in table_contents.get("columns", []): + if column not in columns: + columns.append(column) + merged[table_name]["columns"] = columns + + return merged + + +def _lexical_columns_and_tables_needed( + construct_db_schemas: list[dict], + query: str | None, + max_tables: int = 4, + max_columns_per_table: int = 12, +) -> dict[str, dict]: + if not query: + return {} + + query_tokens = _tokenize_schema_text(_augment_retrieval_query(query)) + if not query_tokens: + return {} + + scored_tables = [] + for table_schema in construct_db_schemas: + if table_schema.get("type") != "TABLE": + continue + + table_tokens = ( + _tokenize_schema_identifier_text(table_schema.get("name")) + | _tokenize_schema_text(table_schema.get("comment")) + | _tokenize_nested_schema_identifier_text(table_schema.get("properties")) + | _tokenize_nested_schema_identifier_text(table_schema.get("tableReference")) + ) + table_matches = query_tokens & table_tokens + table_score = len(table_matches) * 6 + column_scores = [] + column_match_union: set[str] = set() + + for column in table_schema.get("columns", []): + if ( + column.get("type") != "COLUMN" + or column.get("data_type", "").lower() == "unknown" + ): + continue + + column_tokens = _tokenize_schema_identifier_text(column.get("name")) + column_tokens.update( + _tokenize_schema_identifier_text(column.get("display_name")) + ) + column_tokens.update( + _tokenize_schema_identifier_text(column.get("displayName")) + ) + column_tokens.update( + _tokenize_nested_schema_identifier_text(column.get("properties")) + ) + comment_tokens = _tokenize_schema_text( + column.get("comment") + ) | _tokenize_schema_text(column.get("description")) + column_matches = query_tokens & column_tokens + comment_matches = query_tokens & comment_tokens + score = len(column_matches) * 10 + score += len(comment_matches) * 2 + if score > 0: + column_match_union.update(column_matches | comment_matches) + column_scores.append( + (score, column["name"], column.get("is_primary_key")) + ) + + if not column_scores and table_score <= 0: + continue + + total_score = table_score + sum(score for score, _, _ in column_scores) + total_score += len(table_matches | column_match_union) ** 2 + if total_score <= 0: + continue + + selected_columns = [] + for _, column_name, _ in sorted( + column_scores, + key=lambda item: (-item[0], item[1]), + ): + if column_name not in selected_columns: + selected_columns.append(column_name) + if len(selected_columns) >= max_columns_per_table: + break + for _, column_name, is_primary_key in column_scores: + if is_primary_key and column_name not in selected_columns: + selected_columns.append(column_name) + + scored_tables.append((total_score, table_schema["name"], selected_columns)) + + scored_tables.sort(key=lambda item: (-item[0], item[1])) + return { + table_name: {"columns": columns} + for _, table_name, columns in scored_tables[:max_tables] + if columns + } + + ## End of Pipeline class MatchingTableContents(BaseModel): + model_config = ConfigDict(extra="forbid") + chain_of_thought_reasoning: list[str] columns: list[str] class MatchingTable(BaseModel): + model_config = ConfigDict(extra="forbid") + table_name: str table_contents: MatchingTableContents table_selection_reason: str class RetrievalResults(BaseModel): + model_config = ConfigDict(extra="forbid") + results: list[MatchingTable] RETRIEVAL_MODEL_KWARGS = { + "preserve_json_schema": True, "response_format": { "type": "json_schema", "json_schema": { "name": "retrieval_schema", + "strict": True, "schema": RetrievalResults.model_json_schema(), }, } @@ -455,8 +2295,9 @@ def __init__( llm_provider: LLMProvider, embedder_provider: EmbedderProvider, document_store_provider: DocumentStoreProvider, - table_retrieval_size: int = 10, + table_retrieval_size: int = 50, table_column_retrieval_size: int = 100, + include_related_models: bool = False, **kwargs, ): self._components = { @@ -489,6 +2330,7 @@ def __init__( self._configs = { "encoding": _encoding, "context_window_size": llm_provider.get_context_window_size(), + "include_related_models": include_related_models, } super().__init__( @@ -501,19 +2343,37 @@ async def run( query: str = "", tables: Optional[list[str]] = None, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, histories: Optional[list[AskHistory]] = None, enable_column_pruning: bool = False, ): - logger.info("Ask Retrieval pipeline is running...") - return await self._pipe.execute( + logger.info( + "Ask Retrieval pipeline is running for project_id=%s mdl_hash=%s", + project_id or "", + mdl_hash or "", + ) + started_at = time.perf_counter() + result = await self._pipe.execute( ["construct_retrieval_results"], inputs={ "query": query, "tables": tables, "project_id": project_id or "", + "mdl_hash": mdl_hash, "histories": histories or [], "enable_column_pruning": enable_column_pruning, **self._components, **self._configs, }, ) + retrieval_results = result.get("construct_retrieval_results", {}).get( + "retrieval_results", + [], + ) + _log_retrieval_timing( + "schema_retrieval_total", + started_at, + project_id, + retrieval_result_count=len(retrieval_results), + ) + return result diff --git a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py index 0dcbc839ab..1d16ae1f99 100644 --- a/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/historical_question_retrieval.py @@ -10,7 +10,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.common import ScoreFilter +from src.pipelines.common import ScoreFilter, build_project_deploy_filter logger = logging.getLogger("wren-ai-service") @@ -39,19 +39,12 @@ def run(self, documents: List[Document]): async def count_documents( view_questions_store: QdrantDocumentStore, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> int: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) - return await view_questions_store.count_documents(filters=filters) + count = await view_questions_store.count_documents(filters=filters) + return count @observe(capture_input=False, capture_output=False) @@ -67,18 +60,10 @@ async def retrieval( embedding: dict, project_id: str, view_questions_retriever: Any, + mdl_hash: str = "", ) -> dict: if embedding: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) view_question_res = await view_questions_retriever.run( query_embedding=embedding.get("embedding"), @@ -148,13 +133,19 @@ def __init__( ) @observe(name="Historical Question") - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): logger.info("HistoricalQuestion Retrieval pipeline is running...") return await self._pipe.execute( ["formatted_output"], inputs={ "query": query, "project_id": project_id or "", + "mdl_hash": mdl_hash or "", **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/retrieval/instructions.py b/wren-ai-service/src/pipelines/retrieval/instructions.py index 86c17e93de..7753bfa819 100644 --- a/wren-ai-service/src/pipelines/retrieval/instructions.py +++ b/wren-ai-service/src/pipelines/retrieval/instructions.py @@ -10,7 +10,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.common import ScoreFilter +from src.pipelines.common import ScoreFilter, build_project_deploy_filter logger = logging.getLogger("wren-ai-service") @@ -57,18 +57,11 @@ def run( ## Start of Pipeline @observe(capture_input=False) async def count_documents( - store: QdrantDocumentStore, project_id: Optional[str] = None + store: QdrantDocumentStore, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> int: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) document_count = await store.count_documents(filters=filters) return document_count @@ -82,7 +75,12 @@ async def embedding(count_documents: int, query: str, embedder: Any) -> dict: @observe(capture_input=False) -async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: +async def retrieval( + embedding: dict, + project_id: str, + retriever: Any, + mdl_hash: str = "", +) -> dict: if not embedding: return {} @@ -98,6 +96,11 @@ async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: {"field": "project_id", "operator": "==", "value": project_id} ) + if mdl_hash: + filters["conditions"].append( + {"field": "mdl_hash", "operator": "==", "value": mdl_hash} + ) + res = await retriever.run( query_embedding=embedding.get("embedding"), filters=filters, @@ -136,6 +139,7 @@ async def default_instructions( project_id: str, scope_filter: ScopeFilter, scope: str, + mdl_hash: str = "", ) -> list[Document]: if not count_documents: return [] @@ -152,6 +156,11 @@ async def default_instructions( {"field": "project_id", "operator": "==", "value": project_id} ) + if mdl_hash: + filters["conditions"].append( + {"field": "mdl_hash", "operator": "==", "value": mdl_hash} + ) + _res = await retriever.run( query_embedding=None, filters=filters, @@ -213,7 +222,11 @@ def __init__( @observe(name="Instructions Retrieval") async def run( - self, query: str, project_id: Optional[str] = None, scope: str = "sql" + self, + query: str, + project_id: Optional[str] = None, + scope: str = "sql", + mdl_hash: Optional[str] = None, ): logger.info("Instructions Retrieval pipeline is running...") return await self._pipe.execute( @@ -222,6 +235,7 @@ async def run( "query": query, "project_id": project_id or "", "scope": scope, + "mdl_hash": mdl_hash or "", **self._components, **self._configs, }, diff --git a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py index e6dadd32d0..cf282b7899 100644 --- a/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py +++ b/wren-ai-service/src/pipelines/retrieval/preprocess_sql_data.py @@ -1,6 +1,7 @@ import logging import sys -from typing import Dict +from copy import deepcopy +from typing import Any, Dict import tiktoken from hamilton import base @@ -14,6 +15,46 @@ ## Start of Pipeline +def _get_column_name(column: Any) -> str: + if isinstance(column, dict): + return str(column.get("name", "")) + + return str(column) + + +def _build_row_records(sql_data: Dict) -> list[dict]: + column_names = [ + column_name + for column_name in ( + _get_column_name(column) for column in sql_data.get("columns", []) + ) + if column_name + ] + + if not column_names: + return [] + + row_records = [] + for row in sql_data.get("data", []): + if isinstance(row, dict): + row_records.append( + {column_name: row.get(column_name) for column_name in column_names} + ) + continue + + if not isinstance(row, (list, tuple)): + row = [row] + + row_records.append( + { + column_name: row[index] if index < len(row) else None + for index, column_name in enumerate(column_names) + } + ) + + return row_records + + @observe(capture_input=False, capture_output=False) def preprocess( sql_data: Dict, @@ -46,6 +87,9 @@ def reduce_data_size(data: list, reduction_step: int = 50) -> list: return returned_data + sql_data = deepcopy(sql_data) + sql_data["row_records"] = _build_row_records(sql_data) + _token_count = len(encoding.encode(str(sql_data))) num_rows_used_in_llm = len(sql_data.get("data", [])) iteration = 0 @@ -62,6 +106,7 @@ def reduce_data_size(data: list, reduction_step: int = 50) -> list: data = sql_data.get("data", []) sql_data["data"] = reduce_data_size(data) + sql_data["row_records"] = _build_row_records(sql_data) num_rows_used_in_llm = len(sql_data.get("data", [])) _token_count = len(encoding.encode(str(sql_data))) logger.info(f"Token count: {_token_count}") diff --git a/wren-ai-service/src/pipelines/retrieval/sql_executor.py b/wren-ai-service/src/pipelines/retrieval/sql_executor.py index b41151469f..74b5dd64ad 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_executor.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_executor.py @@ -16,8 +16,9 @@ @component class DataFetcher: - def __init__(self, engine: Engine): + def __init__(self, engine: Engine, data_source: str | None = None): self._engine = engine + self._data_source = data_source @component.output_types( results=Optional[Dict[str, Any]], @@ -26,6 +27,7 @@ async def run( self, sql: str, project_id: str | None = None, + mdl_hash: str | None = None, limit: int = 500, ): async with aiohttp.ClientSession() as session: @@ -33,6 +35,7 @@ async def run( sql, session, project_id=project_id, + mdl_hash=mdl_hash, dry_run=False, limit=limit, ) @@ -48,11 +51,13 @@ async def execute_sql( sql: str, data_fetcher: DataFetcher, project_id: str | None = None, + mdl_hash: str | None = None, limit: int = 500, ) -> dict: return await data_fetcher.run( sql=sql, project_id=project_id, + mdl_hash=mdl_hash, limit=limit, ) @@ -64,10 +69,11 @@ class SQLExecutor(BasicPipeline): def __init__( self, engine: Engine, + data_source: str | None = None, **kwargs, ): self._components = { - "data_fetcher": DataFetcher(engine=engine), + "data_fetcher": DataFetcher(engine=engine, data_source=data_source), } super().__init__( @@ -76,7 +82,11 @@ def __init__( @observe(name="SQL Execution") async def run( - self, sql: str, project_id: str | None = None, limit: int = 500 + self, + sql: str, + project_id: str | None = None, + mdl_hash: str | None = None, + limit: int = 500, ) -> dict: logger.info("SQL Execution pipeline is running...") return await self._pipe.execute( @@ -84,6 +94,7 @@ async def run( inputs={ "sql": sql, "project_id": project_id, + "mdl_hash": mdl_hash, "limit": limit, **self._components, }, diff --git a/wren-ai-service/src/pipelines/retrieval/sql_functions.py b/wren-ai-service/src/pipelines/retrieval/sql_functions.py index 016aa9b1e6..5d1bf98519 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_functions.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_functions.py @@ -21,7 +21,7 @@ class SqlFunction: _expr: str = None def __init__(self, definition: dict): - def _extract() -> tuple[str, list, str]: + def _extract() -> tuple[str, str, str]: return ( definition.get("name", "").upper(), definition.get("function_type", ""), @@ -104,13 +104,18 @@ def __init__( async def run( self, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> List[SqlFunction]: logger.info( f"Project ID: {project_id} SQL Functions Retrieval pipeline is running..." ) - metadata = await retrieve_metadata(project_id or "", self._retriever) - _data_source = metadata.get("data_source", "local_file") + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) + _data_source = metadata.get("data_source", "local_file").lower() if _data_source in self._cache: logger.info(f"Hit cache of SQL Functions for {_data_source}") diff --git a/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py b/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py index 167969047c..ad6dde9ad1 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_knowledge.py @@ -114,12 +114,17 @@ def __init__( async def run( self, project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> Optional[SqlKnowledge]: logger.info( f"Project ID: {project_id} SQL Knowledge Retrieval pipeline is running..." ) - metadata = await retrieve_metadata(project_id or "", self._retriever) + metadata = await retrieve_metadata( + project_id or "", + self._retriever, + mdl_hash=mdl_hash, + ) _data_source = metadata.get("data_source", "local_file") if _data_source in self._cache: diff --git a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py index 3fe44f32eb..d836cbe37a 100644 --- a/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py +++ b/wren-ai-service/src/pipelines/retrieval/sql_pairs_retrieval.py @@ -10,7 +10,7 @@ from src.core.pipeline import BasicPipeline from src.core.provider import DocumentStoreProvider, EmbedderProvider -from src.pipelines.common import ScoreFilter +from src.pipelines.common import ScoreFilter, build_project_deploy_filter logger = logging.getLogger("wren-ai-service") @@ -36,18 +36,11 @@ def run(self, documents: List[Document]): ## Start of Pipeline @observe(capture_input=False) async def count_documents( - store: QdrantDocumentStore, project_id: Optional[str] = None + store: QdrantDocumentStore, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, ) -> int: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) document_count = await store.count_documents(filters=filters) return document_count @@ -61,18 +54,14 @@ async def embedding(count_documents: int, query: str, embedder: Any) -> dict: @observe(capture_input=False) -async def retrieval(embedding: dict, project_id: str, retriever: Any) -> dict: +async def retrieval( + embedding: dict, + project_id: str, + retriever: Any, + mdl_hash: str = "", +) -> dict: if embedding: - filters = ( - { - "operator": "AND", - "conditions": [ - {"field": "project_id", "operator": "==", "value": project_id}, - ], - } - if project_id - else None - ) + filters = build_project_deploy_filter(project_id=project_id, mdl_hash=mdl_hash) res = await retriever.run( query_embedding=embedding.get("embedding"), @@ -143,13 +132,19 @@ def __init__( ) @observe(name="SqlPairs Retrieval") - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): logger.info("SqlPairs Retrieval pipeline is running...") return await self._pipe.execute( ["formatted_output"], inputs={ "query": query, "project_id": project_id or "", + "mdl_hash": mdl_hash or "", **self._components, **self._configs, }, diff --git a/wren-ai-service/src/providers/document_store/qdrant.py b/wren-ai-service/src/providers/document_store/qdrant.py index b90961c456..b92cd32d7d 100644 --- a/wren-ai-service/src/providers/document_store/qdrant.py +++ b/wren-ai-service/src/providers/document_store/qdrant.py @@ -1,3 +1,4 @@ +import inspect import logging import os from typing import Any, Dict, List, Optional @@ -21,6 +22,7 @@ from haystack_integrations.document_stores.qdrant.filters import ( convert_filters_to_qdrant, ) +from qdrant_client.http.exceptions import UnexpectedResponse from qdrant_client.http import models as rest from tqdm import tqdm @@ -30,6 +32,21 @@ logger = logging.getLogger("wren-ai-service") +def _env_flag(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return str(value).strip().lower() in {"1", "true", "yes", "on"} + + +def _is_missing_collection_error(err: Exception) -> bool: + if not isinstance(err, UnexpectedResponse): + return False + + status_code = getattr(err, "status_code", None) + return status_code == 404 and "doesn't exist" in str(err) + + def convert_haystack_documents_to_qdrant_points( documents: List[Document], *, @@ -161,6 +178,73 @@ def __init__( self.client.create_payload_index( collection_name=index, field_name="project_id", field_schema="keyword" ) + self.client.create_payload_index( + collection_name=index, field_name="mdl_hash", field_schema="keyword" + ) + + def recreate_collection( + self, + collection_name: str, + distance, + embedding_dim: int, + on_disk: Optional[bool] = None, + use_sparse_embeddings: Optional[bool] = None, + sparse_idf: bool = False, + ): + if on_disk is None: + on_disk = self.on_disk + + if use_sparse_embeddings is None: + use_sparse_embeddings = self.use_sparse_embeddings + + vectors_config = rest.VectorParams( + size=embedding_dim, + on_disk=on_disk, + distance=distance, + ) + sparse_vectors_config = None + + if use_sparse_embeddings: + vectors_config = {DENSE_VECTORS_NAME: vectors_config} + sparse_vectors_config = { + SPARSE_VECTORS_NAME: rest.SparseVectorParams( + index=rest.SparseIndexParams(on_disk=on_disk), + modifier=rest.Modifier.IDF if sparse_idf else None, + ), + } + + if self.client.collection_exists(collection_name): + self.client.delete_collection(collection_name) + + create_collection_kwargs = { + "collection_name": collection_name, + "vectors_config": vectors_config, + "sparse_vectors_config": sparse_vectors_config + if use_sparse_embeddings + else None, + "shard_number": self.shard_number, + "replication_factor": self.replication_factor, + "write_consistency_factor": self.write_consistency_factor, + "on_disk_payload": self.on_disk_payload, + "hnsw_config": self.hnsw_config, + "optimizers_config": self.optimizers_config, + "wal_config": self.wal_config, + "quantization_config": self.quantization_config, + } + if self.init_from is not None: + create_collection_signature = inspect.signature( + self.client.create_collection + ) + if "init_from" in create_collection_signature.parameters: + create_collection_kwargs["init_from"] = self.init_from + else: + logger.warning( + "Ignoring init_from for collection %s because the installed " + "qdrant-client does not support that argument", + collection_name, + ) + + self.client.create_collection(**create_collection_kwargs) async def _query_by_embedding( self, @@ -172,27 +256,36 @@ async def _query_by_embedding( ) -> List[Document]: qdrant_filters = convert_filters_to_qdrant(filters) - points = await self.async_client.search( - collection_name=self.index, - query_vector=rest.NamedVector( - name=DENSE_VECTORS_NAME if self.use_sparse_embeddings else "", - vector=query_embedding, - ), - search_params=( - rest.SearchParams( - quantization=rest.QuantizationSearchParams( - rescore=True, - oversampling=3.0, - ), + try: + points = await self.async_client.search( + collection_name=self.index, + query_vector=rest.NamedVector( + name=DENSE_VECTORS_NAME if self.use_sparse_embeddings else "", + vector=query_embedding, + ), + search_params=( + rest.SearchParams( + quantization=rest.QuantizationSearchParams( + rescore=True, + oversampling=3.0, + ), + ) + if len(query_embedding) + >= 1024 # reference: https://qdrant.tech/articles/binary-quantization/#when-should-you-not-use-bq + else None + ), + query_filter=qdrant_filters, + limit=top_k, + with_vectors=return_embedding, + ) + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, returning no documents", + self.index, ) - if len(query_embedding) - >= 1024 # reference: https://qdrant.tech/articles/binary-quantization/#when-should-you-not-use-bq - else None - ), - query_filter=qdrant_filters, - limit=top_k, - with_vectors=return_embedding, - ) + return [] + raise results = [ convert_qdrant_point_to_haystack_document( point, use_sparse_embeddings=self.use_sparse_embeddings @@ -218,12 +311,21 @@ async def _query_by_filters( points_list = [] offset = None while True: - points = await self.async_client.scroll( - collection_name=self.index, - offset=offset, - scroll_filter=qdrant_filters, - limit=top_k, - ) + try: + points = await self.async_client.scroll( + collection_name=self.index, + offset=offset, + scroll_filter=qdrant_filters, + limit=top_k, + ) + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, returning no documents", + self.index, + ) + return [] + raise points_list.extend(points[0]) if points[1] is None: break @@ -255,6 +357,14 @@ async def delete_documents(self, filters: Optional[Dict[str, Any]] = None): logger.warning( "Called QdrantDocumentStore.delete_documents() on a non-existing ID", ) + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, skipping delete", + self.index, + ) + return + raise async def count_documents(self, filters: Optional[Dict[str, Any]] = None) -> int: if not filters: @@ -262,11 +372,20 @@ async def count_documents(self, filters: Optional[Dict[str, Any]] = None) -> int else: qdrant_filters = convert_filters_to_qdrant(filters) - return ( - await self.async_client.count( - collection_name=self.index, count_filter=qdrant_filters - ) - ).count + try: + return ( + await self.async_client.count( + collection_name=self.index, count_filter=qdrant_filters + ) + ).count + except Exception as err: + if _is_missing_collection_error(err): + logger.warning( + "Qdrant collection %s does not exist yet, returning count 0", + self.index, + ) + return 0 + raise async def write_documents( self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.FAIL @@ -378,11 +497,7 @@ def __init__( if os.getenv("EMBEDDING_MODEL_DIMENSION") else 0 ), - recreate_index: bool = ( - bool(os.getenv("SHOULD_FORCE_DEPLOY")) - if os.getenv("SHOULD_FORCE_DEPLOY") - else False - ), + recreate_index: bool = _env_flag("SHOULD_FORCE_DEPLOY"), **_, ): self._location = location diff --git a/wren-ai-service/src/providers/embedder/litellm.py b/wren-ai-service/src/providers/embedder/litellm.py index 4d051e3284..2c5846ad44 100644 --- a/wren-ai-service/src/providers/embedder/litellm.py +++ b/wren-ai-service/src/providers/embedder/litellm.py @@ -1,12 +1,14 @@ import asyncio +import json import logging import os +from types import SimpleNamespace from typing import Any, Dict, List, Optional, Tuple +import aiohttp import backoff import openai from haystack import Document, component -from litellm import aembedding from src.core.provider import EmbedderProvider from src.providers.loader import provider @@ -14,8 +16,172 @@ logger = logging.getLogger("wren-ai-service") +MIN_EMBED_INPUT_CHARS = 128 -def _prepare_texts_to_embed(documents: List[Document]) -> List[str]: + +class EmbeddingRequestError(Exception): + pass + + +def _normalize_model_name(model: str, api_base_url: Optional[str]) -> str: + # OpenAI-compatible local servers often expect the raw model name and will + # reject litellm-style "openai/" prefixes. + model = str(model or "") + if api_base_url and model.startswith("openai/"): + return model.split("/", 1)[1] + return model + + +def _should_use_minimal_http_client(api_base_url: Optional[str]) -> bool: + if not api_base_url: + return False + + return "api.openai.com" not in str(api_base_url).lower() + + +def _build_embedding_meta(response: Any) -> Dict[str, Any]: + usage = getattr(response, "usage", {}) or {} + usage_dict = dict(usage) if isinstance(usage, dict) or hasattr(usage, "__iter__") else {} + + return { + "model": getattr(response, "model", ""), + "usage": usage_dict, + } + + +def _get_usage_value(usage: Any, key: str) -> int: + if isinstance(usage, dict): + return usage.get(key, 0) or 0 + + return getattr(usage, key, 0) or 0 + + +def _coerce_embedding_response(payload: Dict[str, Any]) -> Any: + data = payload.get("data") + if not data and payload.get("embedding") is not None: + data = [{"embedding": payload["embedding"]}] + + if not isinstance(data, list) or not data: + raise EmbeddingRequestError( + "Embedding provider returned an invalid response payload." + ) + + normalized_data = [] + for item in data: + embedding = item.get("embedding") if isinstance(item, dict) else None + if embedding is None: + raise EmbeddingRequestError( + "Embedding provider response did not include an embedding." + ) + normalized_data.append(SimpleNamespace(embedding=embedding)) + + return SimpleNamespace( + model=payload.get("model", ""), + data=normalized_data, + usage=payload.get("usage", {}) or {}, + ) + + +async def _create_embedding_via_http( + *, + model: str, + input_text: str, + api_key: Optional[str], + api_base_url: str, + timeout: Optional[float], + **kwargs, +): + endpoint = f"{remove_trailing_slash(api_base_url)}/embeddings" + headers = {"Content-Type": "application/json"} + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + payload = { + "model": _normalize_model_name(model, api_base_url), + "input": input_text, + } + payload.update({key: value for key, value in kwargs.items() if value is not None}) + + client_timeout = aiohttp.ClientTimeout(total=timeout) if timeout else None + async with aiohttp.ClientSession() as session: + async with session.post( + endpoint, + json=payload, + headers=headers, + timeout=client_timeout, + ) as response: + body = await response.text() + if response.status >= 400: + raise EmbeddingRequestError( + f"Embedding request failed with status {response.status}: {body}" + ) + + try: + payload = json.loads(body) + except json.JSONDecodeError as error: + raise EmbeddingRequestError( + "Embedding provider returned a non-JSON response." + ) from error + + return _coerce_embedding_response(payload) + + +async def _create_embedding( + *, + model: str, + input_text: str, + api_key: Optional[str], + api_base_url: Optional[str], + timeout: Optional[float], + **kwargs, +): + if _should_use_minimal_http_client(api_base_url): + return await _create_embedding_via_http( + model=model, + input_text=input_text, + api_key=api_key, + api_base_url=api_base_url, + timeout=timeout, + **kwargs, + ) + + client = openai.AsyncOpenAI( + api_key=api_key, + base_url=api_base_url, + timeout=timeout, + ) + return await client.embeddings.create( + model=_normalize_model_name(model, api_base_url), + input=input_text, + **kwargs, + ) + + +def _truncate_text_for_embedding( + text: str, max_input_chars: Optional[int] +) -> str: + if max_input_chars is None or len(text) <= max_input_chars: + return text + + return text[:max_input_chars].rstrip() + "..." + + +def _is_input_too_large_error(error: Exception) -> bool: + error_message = str(error).lower() + return any( + phrase in error_message + for phrase in [ + "too large to process", + "context size has been exceeded", + "physical batch size", + "input (", + ] + ) + + +def _prepare_texts_to_embed( + documents: List[Document], max_input_chars: Optional[int] +) -> List[str]: """ Prepare the texts to embed by concatenating the Document text with the metadata fields to embed. """ @@ -26,10 +192,19 @@ def _prepare_texts_to_embed(documents: List[Document]) -> List[str]: # copied from OpenAI embedding_utils (https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py) # replace newlines, which can negatively affect performance. text_to_embed = text_to_embed.replace("\n", " ") + text_to_embed = _truncate_text_for_embedding(text_to_embed, max_input_chars) texts_to_embed.append(text_to_embed) return texts_to_embed +def _iter_batches(items: List[str], batch_size: int) -> List[List[str]]: + effective_batch_size = max(batch_size, 1) + return [ + items[index : index + effective_batch_size] + for index in range(0, len(items), effective_batch_size) + ] + + @component class AsyncTextEmbedder: def __init__( @@ -38,16 +213,27 @@ def __init__( api_key: Optional[str] = None, api_base_url: Optional[str] = None, timeout: Optional[float] = None, + max_input_chars: Optional[int] = None, + query_prefix: str = "", **kwargs, ): self._api_key = api_key self._model = model self._api_base_url = api_base_url self._timeout = timeout + self._max_input_chars = ( + max(max_input_chars, 1) if max_input_chars is not None else None + ) + self._query_prefix = query_prefix self._kwargs = kwargs @component.output_types(embedding=List[float], meta=Dict[str, Any]) - @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) + @backoff.on_exception( + backoff.expo, + (aiohttp.ClientError, asyncio.TimeoutError, EmbeddingRequestError, openai.APIError), + max_time=60.0, + max_tries=3, + ) async def run(self, text: str): if not isinstance(text, str): raise TypeError( @@ -58,22 +244,50 @@ async def run(self, text: str): # copied from OpenAI embedding_utils (https://github.com/openai/openai-python/blob/main/openai/embeddings_utils.py) # replace newlines, which can negatively affect performance. text_to_embed = text.replace("\n", " ") - - response = await aembedding( - model=self._model, - input=[text_to_embed], - api_key=self._api_key, - api_base=self._api_base_url, - timeout=self._timeout, - **self._kwargs, + if self._query_prefix and not text_to_embed.startswith(self._query_prefix): + text_to_embed = f"{self._query_prefix}{text_to_embed}" + text_to_embed = _truncate_text_for_embedding( + text_to_embed, + self._max_input_chars, ) - meta = { - "model": response.model, - "usage": dict(response.usage) if hasattr(response, "usage") else {}, - } + candidate_text = text_to_embed + while True: + try: + response = await _create_embedding( + model=self._model, + input_text=candidate_text, + api_key=self._api_key, + api_base_url=self._api_base_url, + timeout=self._timeout, + **self._kwargs, + ) + break + except ( + aiohttp.ClientError, + asyncio.TimeoutError, + EmbeddingRequestError, + openai.APIError, + ) as error: + if ( + not _is_input_too_large_error(error) + or len(candidate_text) <= MIN_EMBED_INPUT_CHARS + ): + raise + + next_max_chars = max(len(candidate_text) // 2, MIN_EMBED_INPUT_CHARS) + logger.warning( + "Embedding input exceeded provider limits; retrying with %s characters", + next_max_chars, + ) + candidate_text = _truncate_text_for_embedding( + candidate_text, + next_max_chars, + ) - return {"embedding": response.data[0]["embedding"], "meta": meta} + meta = _build_embedding_meta(response) + + return {"embedding": response.data[0].embedding, "meta": meta} @component @@ -85,6 +299,8 @@ def __init__( api_key: Optional[str] = None, api_base_url: Optional[str] = None, timeout: Optional[float] = None, + max_input_chars: Optional[int] = None, + document_prefix: str = "", **kwargs, ): self._api_key = api_key @@ -92,52 +308,90 @@ def __init__( self._batch_size = batch_size self._api_base_url = api_base_url self._timeout = timeout + self._max_input_chars = ( + max(max_input_chars, 1) if max_input_chars is not None else None + ) + self._document_prefix = document_prefix self._kwargs = kwargs async def _embed_batch( self, texts_to_embed: List[str], batch_size: int ) -> Tuple[List[List[float]], Dict[str, Any]]: - async def embed_single_batch(batch: List[str]) -> Any: - return await aembedding( - model=self._model, - input=batch, - api_key=self._api_key, - api_base=self._api_base_url, - timeout=self._timeout, - **self._kwargs, - ) + # Some OpenAI-compatible local embedding servers accept scalar string input + # but fail on array input. Embed documents individually to avoid that path. + async def embed_single_text(text: str) -> Any: + candidate_text = text + while True: + try: + return await _create_embedding( + model=self._model, + input_text=candidate_text, + api_key=self._api_key, + api_base_url=self._api_base_url, + timeout=self._timeout, + **self._kwargs, + ) + except ( + aiohttp.ClientError, + asyncio.TimeoutError, + EmbeddingRequestError, + openai.APIError, + ) as error: + if ( + not _is_input_too_large_error(error) + or len(candidate_text) <= MIN_EMBED_INPUT_CHARS + ): + raise - batches = [ - texts_to_embed[i : i + batch_size] - for i in range(0, len(texts_to_embed), batch_size) - ] - - responses = await asyncio.gather( - *[embed_single_batch(batch) for batch in batches] - ) + next_max_chars = max(len(candidate_text) // 2, MIN_EMBED_INPUT_CHARS) + logger.warning( + "Embedding input exceeded provider limits; retrying with %s characters", + next_max_chars, + ) + candidate_text = _truncate_text_for_embedding( + candidate_text, + next_max_chars, + ) all_embeddings = [] meta: Dict[str, Any] = {} - for response in responses: - embeddings = [el["embedding"] for el in response.data] - all_embeddings.extend(embeddings) + for batch in _iter_batches(texts_to_embed, batch_size): + responses = await asyncio.gather( + *[embed_single_text(text) for text in batch] + ) - if "model" not in meta: - meta["model"] = response.model - if "usage" not in meta: - meta["usage"] = ( - dict(response.usage) if hasattr(response, "usage") else {} - ) - else: - if hasattr(response, "usage"): - meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens - meta["usage"]["total_tokens"] += response.usage.total_tokens + for response in responses: + embeddings = [ + el.embedding if hasattr(el, "embedding") else el["embedding"] + for el in response.data + ] + all_embeddings.extend(embeddings) + + if "model" not in meta: + meta["model"] = getattr(response, "model", "") + if "usage" not in meta: + meta["usage"] = _build_embedding_meta(response)["usage"] + else: + if hasattr(response, "usage"): + meta["usage"]["prompt_tokens"] += _get_usage_value( + response.usage, + "prompt_tokens", + ) + meta["usage"]["total_tokens"] += _get_usage_value( + response.usage, + "total_tokens", + ) return all_embeddings, meta @component.output_types(documents=List[Document], meta=Dict[str, Any]) - @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) + @backoff.on_exception( + backoff.expo, + (aiohttp.ClientError, asyncio.TimeoutError, EmbeddingRequestError, openai.APIError), + max_time=60.0, + max_tries=3, + ) async def run(self, documents: List[Document]): if ( not isinstance(documents, list) @@ -149,7 +403,20 @@ async def run(self, documents: List[Document]): "In case you want to embed a string, please use the AsyncTextEmbedder." ) - texts_to_embed = _prepare_texts_to_embed(documents=documents) + if not documents: + return {"documents": documents, "meta": {}} + + texts_to_embed = _prepare_texts_to_embed( + documents=documents, + max_input_chars=self._max_input_chars, + ) + if self._document_prefix: + texts_to_embed = [ + text + if text.startswith(self._document_prefix) + else f"{self._document_prefix}{text}" + for text in texts_to_embed + ] embeddings, meta = await self._embed_batch( texts_to_embed=texts_to_embed, @@ -172,12 +439,16 @@ def __init__( ] = None, # e.g. EMBEDDER_OPENAI_API_KEY, EMBEDDER_ANTHROPIC_API_KEY, etc. api_base: Optional[str] = None, timeout: float = 120.0, + query_prefix: Optional[str] = None, + document_prefix: Optional[str] = None, **kwargs, ): self._api_key = os.getenv(api_key_name) if api_key_name else None self._api_base = remove_trailing_slash(api_base) if api_base else None self._embedding_model = model self._timeout = timeout + self._query_prefix = query_prefix or "" + self._document_prefix = document_prefix or "" if "provider" in kwargs: del kwargs["provider"] self._kwargs = kwargs @@ -188,6 +459,7 @@ def get_text_embedder(self): api_base_url=self._api_base, model=self._embedding_model, timeout=self._timeout, + query_prefix=self._query_prefix, **self._kwargs, ) @@ -197,5 +469,6 @@ def get_document_embedder(self): api_base_url=self._api_base, model=self._embedding_model, timeout=self._timeout, + document_prefix=self._document_prefix, **self._kwargs, ) diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py index 3a92853e04..eaea5016c3 100644 --- a/wren-ai-service/src/providers/engine/wren.py +++ b/wren-ai-service/src/providers/engine/wren.py @@ -28,6 +28,7 @@ async def execute_sql( sql: str, session: aiohttp.ClientSession, project_id: str | None = None, + mdl_hash: str | None = None, dry_run: bool = True, timeout: float = settings.engine_timeout, limit: int = 500, @@ -37,6 +38,9 @@ async def execute_sql( "sql": remove_limit_statement(sql), "projectId": project_id, } + if mdl_hash: + data["hash"] = mdl_hash + if dry_run: data["dryRun"] = True data["limit"] = 1 @@ -138,6 +142,47 @@ async def execute_sql( {"error_message": f"Request timed out: {timeout} seconds"}, ) + async def dry_plan( + self, + session: aiohttp.ClientSession, + sql: str, + data_source: str, + project_id: str | None = None, + mdl_hash: str | None = None, + timeout: float = settings.engine_timeout, + allow_fallback: bool = True, + **kwargs, + ) -> Tuple[bool, str]: + data = { + "sql": sql, + "projectId": project_id, + "allowFallback": allow_fallback, + } + if mdl_hash: + data["hash"] = mdl_hash + + try: + async with session.post( + f"{self._endpoint}/api/graphql", + json={ + "query": "mutation DryPlanSql($data: DryPlanSQLDataInput) { dryPlanSql(data: $data) }", + "variables": {"data": data}, + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + res_json = await response.json() + if res_data := res_json.get("data"): + return bool(res_data.get("dryPlanSql")), "" + + error_message = res_json.get("errors", [{}])[0].get( + "message", "Unknown error" + ) + logger.error(f"Error dry planning SQL: {error_message}") + return False, error_message + except asyncio.TimeoutError: + logger.error(f"Request timed out: {timeout} seconds") + return False, f"Request timed out: {timeout} seconds" + @provider("wren_ibis") class WrenIbis(Engine): @@ -348,3 +393,33 @@ async def execute_sql( ) except asyncio.TimeoutError: return False, None, f"Request timed out: {timeout} seconds" + + async def dry_plan( + self, + session: aiohttp.ClientSession, + sql: str, + timeout: float = settings.engine_timeout, + **kwargs, + ) -> Tuple[bool, str]: + api_endpoint = f"{self._endpoint}/v1/mdl/dry-plan" + + try: + async with session.get( + api_endpoint, + json={ + "manifest": orjson.loads(base64.b64decode(self._manifest)) + if self._manifest + else {}, + "sql": sql, + }, + timeout=aiohttp.ClientTimeout(total=timeout), + ) as response: + res = await response.text() + + if response.status == 200: + return True, "" + + return False, res + except asyncio.TimeoutError: + logger.error(f"Request timed out: {timeout} seconds") + return False, f"Request timed out: {timeout} seconds" diff --git a/wren-ai-service/src/providers/llm/__init__.py b/wren-ai-service/src/providers/llm/__init__.py index c3f2a47b3f..0469adf7e8 100644 --- a/wren-ai-service/src/providers/llm/__init__.py +++ b/wren-ai-service/src/providers/llm/__init__.py @@ -221,7 +221,50 @@ def build_chunk(chunk: Any) -> StreamingChunk: return chunk_message -def convert_message_to_openai_format(message: ChatMessage) -> Dict[str, str]: +def _get_message_role_value(message: Any) -> str: + role = getattr(message, "role", ChatRole.USER) + return role.value if hasattr(role, "value") else str(role) + + +def _get_message_text_content(message: Any) -> Optional[str]: + try: + content = getattr(message, "content", None) + if isinstance(content, str) and content: + return content + except AttributeError: + # Haystack 2.x removed `.content` in favor of `.text`. + pass + + text = getattr(message, "text", None) + if isinstance(text, str) and text: + return text + + raw_content = getattr(message, "_content", None) + if isinstance(raw_content, str) and raw_content: + return raw_content + + if isinstance(raw_content, list): + text_parts = [] + for part in raw_content: + if isinstance(part, str) and part: + text_parts.append(part) + continue + + part_text = getattr(part, "text", None) + if isinstance(part_text, str) and part_text: + text_parts.append(part_text) + if text_parts: + return "\n".join(text_parts) + + return None + + +def _get_message_image_url(message: Any) -> Optional[str]: + image_url = getattr(message, "image_url", None) + return image_url if isinstance(image_url, str) and image_url else None + + +def convert_message_to_openai_format(message: Any) -> Dict[str, Any]: """ Convert a message to the format expected by OpenAI's Chat API. @@ -232,21 +275,24 @@ def convert_message_to_openai_format(message: ChatMessage) -> Dict[str, str]: - `content` - `name` (optional) """ - openai_msg = {"role": message.role.value} + openai_msg = {"role": _get_message_role_value(message)} + message_text = _get_message_text_content(message) + image_url = _get_message_image_url(message) - if message.content and hasattr(message, "image_url") and message.image_url: + if message_text and image_url: openai_msg["content"] = [ - {"type": "text", "text": message.content}, - {"type": "image_url", "image_url": {"url": message.image_url}}, + {"type": "text", "text": message_text}, + {"type": "image_url", "image_url": {"url": image_url}}, ] - elif message.content: - openai_msg["content"] = message.content - elif hasattr(message, "image_url") and message.image_url: + elif message_text: + openai_msg["content"] = message_text + elif image_url: openai_msg["content"] = [ - {"type": "image_url", "image_url": {"url": message.image_url}} + {"type": "image_url", "image_url": {"url": image_url}} ] - if hasattr(message, "name") and message.name: - openai_msg["name"] = message.name + name = getattr(message, "name", None) + if name: + openai_msg["name"] = name return openai_msg diff --git a/wren-ai-service/src/providers/llm/litellm.py b/wren-ai-service/src/providers/llm/litellm.py index e94918b8c5..fc9406b75b 100644 --- a/wren-ai-service/src/providers/llm/litellm.py +++ b/wren-ai-service/src/providers/llm/litellm.py @@ -19,6 +19,10 @@ from src.utils import extract_braces_content, remove_trailing_slash +def _is_openai_api_base(api_base: Optional[str]) -> bool: + return bool(api_base) and "api.openai.com" in api_base.lower() + + @provider("litellm_llm") class LitellmLLMProvider(LLMProvider): def __init__( @@ -65,10 +69,42 @@ def get_generator( generation_kwargs: Optional[Dict[str, Any]] = None, streaming_callback: Optional[Callable[[StreamingChunk], None]] = None, ): - combined_generation_kwargs = { - **(generation_kwargs or {}), - **(self._model_kwargs or {}), - } + component_generation_kwargs = generation_kwargs or {} + + def _normalize_generation_kwargs( + kwargs: Optional[Dict[str, Any]], + explicit_response_format: bool = False, + ) -> Dict[str, Any]: + normalized = dict(kwargs or {}) + preserve_json_schema = normalized.pop("preserve_json_schema", False) + response_format = normalized.get("response_format") + + # Plain text is the default chat-completions behavior. + # Some OpenAI-compatible endpoints reject an explicit + # {"type": "text"} payload or serialize it incorrectly. + if ( + isinstance(response_format, dict) + and response_format.get("type") == "text" + ): + normalized.pop("response_format", None) + + if ( + self._api_base + and not _is_openai_api_base(self._api_base) + and isinstance(response_format, dict) + and response_format.get("type") == "json_schema" + and not preserve_json_schema + ): + if explicit_response_format: + normalized["response_format"] = {"type": "json_object"} + else: + normalized.pop("response_format", None) + + if self._api_base and not _is_openai_api_base(self._api_base): + # Some local OpenAI-compatible servers reject non-OpenAI keys. + normalized.pop("speed", None) + + return normalized @backoff.on_exception(backoff.expo, openai.APIError, max_time=60.0, max_tries=3) async def _run( @@ -99,10 +135,42 @@ async def _run( convert_message_to_openai_format(message) for message in messages ] - generation_kwargs = { - **combined_generation_kwargs, - **(generation_kwargs or {}), + runtime_generation_kwargs = generation_kwargs or {} + model_generation_kwargs = self._model_kwargs or {} + explicit_response_format = ( + "response_format" in component_generation_kwargs + or "response_format" in model_generation_kwargs + or "response_format" in runtime_generation_kwargs + ) + merged_generation_kwargs = { + **component_generation_kwargs, + **model_generation_kwargs, + **runtime_generation_kwargs, } + if ( + component_generation_kwargs.get("preserve_json_schema") + and isinstance( + component_generation_kwargs.get("response_format"), dict + ) + and component_generation_kwargs["response_format"].get("type") + == "json_schema" + and "response_format" not in runtime_generation_kwargs + ): + merged_generation_kwargs["response_format"] = ( + component_generation_kwargs["response_format"] + ) + merged_generation_kwargs["preserve_json_schema"] = True + + generation_kwargs = _normalize_generation_kwargs( + merged_generation_kwargs, + explicit_response_format=explicit_response_format, + ) + completion_timeout = generation_kwargs.pop("timeout", self._timeout) + should_stream = ( + streaming_callback is not None + and query_id is not None + and generation_kwargs.pop("stream", True) + ) allowed_openai_params = generation_kwargs.get( "allowed_openai_params", [] @@ -112,9 +180,10 @@ async def _run( completion = await self._router.acompletion( model=self._model, messages=openai_formatted_messages, - stream=streaming_callback is not None, + stream=should_stream, allowed_openai_params=allowed_openai_params, mock_testing_fallbacks=self._enable_fallback_testing, + timeout=completion_timeout, **generation_kwargs, ) else: @@ -123,15 +192,15 @@ async def _run( api_key=self._api_key, api_base=self._api_base, api_version=self._api_version, - timeout=self._timeout, + timeout=completion_timeout, messages=openai_formatted_messages, - stream=streaming_callback is not None, + stream=should_stream, allowed_openai_params=allowed_openai_params, **generation_kwargs, ) completions: List[ChatMessage] = [] - if streaming_callback is not None: + if should_stream: num_responses = generation_kwargs.pop("n", 1) if num_responses > 1: raise ValueError( diff --git a/wren-ai-service/src/utils.py b/wren-ai-service/src/utils.py index d368080c3c..e6820b5efc 100644 --- a/wren-ai-service/src/utils.py +++ b/wren-ai-service/src/utils.py @@ -217,4 +217,40 @@ def extract_braces_content(resp: str) -> str: Returns the JSON string including braces, or the original string if no match is found. """ match = re.search(r"```json\s*(\{.*?\})\s*```", resp, re.DOTALL) - return match.group(1) if match else resp + if match: + return match.group(1) + + start = resp.find("{") + if start == -1: + return resp + + depth = 0 + in_string = False + escaped = False + + for index in range(start, len(resp)): + char = resp[index] + + if escaped: + escaped = False + continue + + if char == "\\": + escaped = True + continue + + if char == '"': + in_string = not in_string + continue + + if in_string: + continue + + if char == "{": + depth += 1 + elif char == "}": + depth -= 1 + if depth == 0: + return resp[start : index + 1] + + return resp diff --git a/wren-ai-service/src/web/v1/routers/ask.py b/wren-ai-service/src/web/v1/routers/ask.py index 84bc5dbb37..92b259dbad 100644 --- a/wren-ai-service/src/web/v1/routers/ask.py +++ b/wren-ai-service/src/web/v1/routers/ask.py @@ -1,7 +1,10 @@ +import asyncio +import logging +import time import uuid from dataclasses import asdict -from fastapi import APIRouter, BackgroundTasks, Depends +from fastapi import APIRouter, Depends from fastapi.responses import StreamingResponse from src.globals import ( @@ -20,25 +23,49 @@ ) router = APIRouter() +logger = logging.getLogger("wren-ai-service") @router.post("/asks") async def ask( ask_request: AskRequest, - background_tasks: BackgroundTasks, service_container: ServiceContainer = Depends(get_service_container), service_metadata: ServiceMetadata = Depends(get_service_metadata), ) -> AskResponse: + started_at = time.perf_counter() query_id = str(uuid.uuid4()) ask_request.query_id = query_id - service_container.ask_service._ask_results[query_id] = AskResultResponse( + ask_service = service_container.ask_service + ask_service._ask_results[query_id] = AskResultResponse( status="understanding", ) - background_tasks.add_task( - service_container.ask_service.ask, - ask_request, - service_metadata=asdict(service_metadata), + task = asyncio.create_task( + ask_service.ask( + ask_request, + service_metadata=asdict(service_metadata), + ) + ) + + def _handle_task_done(completed_task: asyncio.Task): + try: + completed_task.result() + except Exception: + # ask() already captures and records task failures, but we still + # log unexpected task-level exceptions instead of dropping them. + import logging + + logging.getLogger("wren-ai-service").exception( + "Unhandled exception in ask background task for query_id %s", + query_id, + ) + + task.add_done_callback(_handle_task_done) + logger.info( + "Ask timing query_id=%s project_id=%s stage=task_creation elapsed_ms=%.1f", + query_id, + ask_request.project_id or "", + (time.perf_counter() - started_at) * 1000, ) return AskResponse(query_id=query_id) @@ -47,14 +74,10 @@ async def ask( async def stop_ask( query_id: str, stop_ask_request: StopAskRequest, - background_tasks: BackgroundTasks, service_container: ServiceContainer = Depends(get_service_container), ) -> StopAskResponse: stop_ask_request.query_id = query_id - background_tasks.add_task( - service_container.ask_service.stop_ask, - stop_ask_request, - ) + service_container.ask_service.stop_ask(stop_ask_request) return StopAskResponse(query_id=query_id) diff --git a/wren-ai-service/src/web/v1/routers/question_recommendation.py b/wren-ai-service/src/web/v1/routers/question_recommendation.py index f029804996..8aa89e3f99 100644 --- a/wren-ai-service/src/web/v1/routers/question_recommendation.py +++ b/wren-ai-service/src/web/v1/routers/question_recommendation.py @@ -22,7 +22,9 @@ class PostRequest(BaseRequest): max_questions: int = 5 max_categories: int = 3 regenerate: bool = False - allow_data_preview: bool = True + allow_data_preview: bool = False + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False class PostResponse(BaseModel): diff --git a/wren-ai-service/src/web/v1/routers/relationship_recommendation.py b/wren-ai-service/src/web/v1/routers/relationship_recommendation.py index 384b856eb1..6f21071878 100644 --- a/wren-ai-service/src/web/v1/routers/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/routers/relationship_recommendation.py @@ -1,8 +1,10 @@ +import asyncio +import logging import uuid from dataclasses import asdict from typing import Literal, Optional -from fastapi import APIRouter, BackgroundTasks, Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel from src.globals import ( @@ -14,6 +16,23 @@ from src.web.v1.services import BaseRequest, RelationshipRecommendation router = APIRouter() +logger = logging.getLogger("wren-ai-service") + + +async def _run_recommendation_task( + service: RelationshipRecommendation, + input: RelationshipRecommendation.Input, + service_metadata: dict, +): + try: + await service.recommend(input, service_metadata=service_metadata) + except Exception as e: + logger.exception("Unexpected relationship recommendation task failure") + service._handle_exception( + input, + f"Unexpected relationship recommendation task failure: {str(e)}", + request_from=input.request_from, + ) class PostRequest(BaseRequest): @@ -30,7 +49,6 @@ class PostResponse(BaseModel): ) async def recommend( request: PostRequest, - background_tasks: BackgroundTasks, service_container: ServiceContainer = Depends(get_service_container), service_metadata: ServiceMetadata = Depends(get_service_metadata), ) -> PostResponse: @@ -45,8 +63,12 @@ async def recommend( configuration=request.configurations, ) - background_tasks.add_task( - service.recommend, input, service_metadata=asdict(service_metadata) + asyncio.create_task( + _run_recommendation_task( + service, + input, + asdict(service_metadata), + ) ) return PostResponse(id=id) diff --git a/wren-ai-service/src/web/v1/routers/semantics_description.py b/wren-ai-service/src/web/v1/routers/semantics_description.py index 16727bdd6c..d48dd202e5 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_description.py +++ b/wren-ai-service/src/web/v1/routers/semantics_description.py @@ -1,8 +1,10 @@ +import asyncio +import logging import uuid from dataclasses import asdict from typing import Literal, Optional -from fastapi import APIRouter, BackgroundTasks, Depends +from fastapi import APIRouter, Depends from pydantic import BaseModel from src.globals import ( @@ -14,6 +16,26 @@ from src.web.v1.services import BaseRequest, SemanticsDescription router = APIRouter() +logger = logging.getLogger("wren-ai-service") + + +async def _run_generate_task( + service: SemanticsDescription, + generate_request: SemanticsDescription.GenerateRequest, + service_metadata: dict, +): + try: + await service.generate( + generate_request, + service_metadata=service_metadata, + ) + except Exception as e: + logger.exception("Unexpected semantics description task failure") + service._handle_exception( + generate_request.id, + f"Unexpected semantics description task failure: {str(e)}", + request_from=generate_request.request_from, + ) class PostRequest(BaseRequest): @@ -32,7 +54,6 @@ class PostResponse(BaseModel): ) async def generate( request: PostRequest, - background_tasks: BackgroundTasks, service_container: ServiceContainer = Depends(get_service_container), service_metadata: ServiceMetadata = Depends(get_service_metadata), ) -> PostResponse: @@ -44,10 +65,12 @@ async def generate( id=id, **request.model_dump() ) - background_tasks.add_task( - service.generate, - generate_request, - service_metadata=asdict(service_metadata), + asyncio.create_task( + _run_generate_task( + service, + generate_request, + asdict(service_metadata), + ) ) return PostResponse(id=id) @@ -74,25 +97,32 @@ def _formatter(response: Optional[dict]) -> Optional[list[dict]]: if response is None: return None + def _properties(payload: dict) -> dict: + properties = payload.get("properties") + return properties if isinstance(properties, dict) else {} + return [ { "name": model_name, "columns": [ { - "name": column["name"], - "description": column["properties"].get("description", ""), + "name": column.get("name", ""), + "type": column.get("type", ""), + "description": _properties(column).get("description", ""), } - for column in model_data["columns"] + for column in model_data.get("columns", []) + if isinstance(column, dict) ], - "description": model_data["properties"].get("description", ""), + "description": _properties(model_data).get("description", ""), } for model_name, model_data in response.items() + if isinstance(model_data, dict) ] return GetResponse( id=resource.id, status=resource.status, - response=resource.response and _formatter(resource.response), + response=_formatter(resource.response), error=resource.error and resource.error.model_dump(), trace_id=resource.trace_id, ) diff --git a/wren-ai-service/src/web/v1/routers/semantics_preparation.py b/wren-ai-service/src/web/v1/routers/semantics_preparation.py index 1a8ed15dc7..8434cfaa7d 100644 --- a/wren-ai-service/src/web/v1/routers/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/routers/semantics_preparation.py @@ -42,10 +42,11 @@ async def prepare_semantics( @router.get("/semantics-preparations/{mdl_hash}/status") async def get_prepare_semantics_status( mdl_hash: str, + project_id: str | None = None, service_container: ServiceContainer = Depends(get_service_container), ) -> SemanticsPreparationStatusResponse: - return service_container.semantics_preparation_service.get_prepare_semantics_status( - SemanticsPreparationStatusRequest(mdl_hash=mdl_hash) + return await service_container.semantics_preparation_service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash=mdl_hash, project_id=project_id) ) diff --git a/wren-ai-service/src/web/v1/services/__init__.py b/wren-ai-service/src/web/v1/services/__init__.py index 5c76e4f3fb..c6d64026f8 100644 --- a/wren-ai-service/src/web/v1/services/__init__.py +++ b/wren-ai-service/src/web/v1/services/__init__.py @@ -56,8 +56,12 @@ def serialize(self): # for POST, PATCH, UPDATE, DELETE requests class BaseRequest(BaseModel): - _query_id: str | None = None - project_id: Optional[str] = None + query_id: Optional[str] = Field(default=None, exclude=True) + project_id: Optional[str] = Field( + default=None, + validation_alias=AliasChoices("project_id", "projectId"), + ) + mdl_hash: Optional[str] = None thread_id: Optional[str] = None configurations: Configuration = Field( default_factory=Configuration, @@ -65,14 +69,6 @@ class BaseRequest(BaseModel): ) request_from: Literal["ui", "api"] = "ui" - @property - def query_id(self) -> str: - return self._query_id - - @query_id.setter - def query_id(self, query_id: str): - self._query_id = query_id - # Put the services imports here to avoid circular imports and make them accessible directly to the rest of packages from .ask import AskService # noqa: E402 diff --git a/wren-ai-service/src/web/v1/services/ask.py b/wren-ai-service/src/web/v1/services/ask.py index aa26fa3f81..1d205b81c3 100644 --- a/wren-ai-service/src/web/v1/services/ask.py +++ b/wren-ai-service/src/web/v1/services/ask.py @@ -1,6 +1,8 @@ import asyncio import logging -from typing import Dict, List, Literal, Optional +import re +import time +from typing import Any, Dict, List, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe @@ -12,12 +14,119 @@ logger = logging.getLogger("wren-ai-service") +_SIMPLE_ANALYTICS_FAST_PATH_PATTERN = re.compile( + r"(?i)\b(" + r"how\s+many|count|counts|number\s+of|total|sum|average|avg|" + r"top\s+(?:\d+|one|two|three|four|five|six|seven|eight|nine|ten|" + r"eleven|twelve|thirteen|fourteen|fifteen|sixteen|seventeen|" + r"eighteen|nineteen|twenty)|" + r"latest|newest|recent|distribution|breakdown|" + r"group(?:ed)?\s+by|by\s+(?:day|week|month|quarter|year)|" + r"each\s+(?:day|week|month|quarter|year)|monthly|per|" + r"missing|blank|empty|null" + r")\b" +) +_SHOW_BY_FAST_PATH_PATTERN = re.compile( + r"(?is)\b(?:show|list)\b.+\bby\s+[A-Za-z0-9_ -]+\b" +) +_GROUP_RESULT_FAST_PATH_PATTERN = re.compile( + r"(?is)\b(?:group|groups)\b.+\b(?:result|results|those|that)\b" +) +_LISTING_FAST_PATH_PATTERN = re.compile( + r"(?is)\b(?:show|list)\b.+\b(?:record|records|row|rows)\b" +) +_GENERAL_HELP_PATTERN = re.compile( + r"(?i)\b(how\s+to|help|guide|docs|documentation|connect|configure|setting|settings)\b" +) +_DATA_SHAPE_PATTERN = re.compile( + r"(?i)\b(record|records|row|rows|table|tables|field|fields|column|columns)\b" +) +_HISTORY_SQL_IDENTIFIER_PATTERN = re.compile( + r'"([^"]+)"|\b(?:FROM|JOIN)\s+([A-Za-z_][A-Za-z0-9_$-]*)', + re.IGNORECASE, +) +_HISTORY_SQL_TABLE_PATTERN = re.compile( + r'\b(?:FROM|JOIN)\s+(?:"([^"]+)"|([A-Za-z_][A-Za-z0-9_$-]*))', + re.IGNORECASE, +) + + +def _looks_like_simple_analytics_request(query: str | None) -> bool: + if not query or _GENERAL_HELP_PATTERN.search(query): + return False + return bool( + _SIMPLE_ANALYTICS_FAST_PATH_PATTERN.search(query) + or _SHOW_BY_FAST_PATH_PATTERN.search(query) + or _GROUP_RESULT_FAST_PATH_PATTERN.search(query) + or _LISTING_FAST_PATH_PATTERN.search(query) + ) + + +def _can_return_pre_intent_schema_unsupported(query: str | None) -> bool: + if not query or _GENERAL_HELP_PATTERN.search(query): + return False + return bool( + _DATA_SHAPE_PATTERN.search(query) + or _SIMPLE_ANALYTICS_FAST_PATH_PATTERN.search(query) + or _SHOW_BY_FAST_PATH_PATTERN.search(query) + or _GROUP_RESULT_FAST_PATH_PATTERN.search(query) + or _LISTING_FAST_PATH_PATTERN.search(query) + ) + class AskHistory(BaseModel): sql: str question: str +def _history_sql_identifiers(sql: str | None) -> list[str]: + if not sql: + return [] + + identifiers = [] + seen = set() + for match in _HISTORY_SQL_IDENTIFIER_PATTERN.finditer(sql): + identifier = match.group(1) or match.group(2) + if not identifier or identifier in seen: + continue + identifiers.append(identifier) + seen.add(identifier) + return identifiers + + +def _history_sql_table_names(sql: str | None) -> list[str]: + if not sql: + return [] + + table_names = [] + seen = set() + for match in _HISTORY_SQL_TABLE_PATTERN.finditer(sql): + table_name = match.group(1) or match.group(2) + if not table_name or table_name in seen: + continue + table_names.append(table_name) + seen.add(table_name) + return table_names + + +def _build_fast_path_grounding_query( + query: str | None, + histories: list[AskHistory], +) -> str: + if not histories: + return query or "" + + latest_history = histories[0] + parts = [] + if latest_history.question: + parts.append(latest_history.question) + history_identifiers = _history_sql_identifiers(latest_history.sql) + if history_identifiers: + parts.append(" ".join(history_identifiers)) + parts.append(query or "") + return "\n".join(parts) + + # POST /v1/asks class AskRequest(BaseRequest): query: str @@ -25,10 +134,10 @@ class AskRequest(BaseRequest): # so we need to support as a choice, and will remove it in the future mdl_hash: Optional[str] = Field(validation_alias=AliasChoices("mdl_hash", "id")) histories: Optional[list[AskHistory]] = Field(default_factory=list) - ignore_sql_generation_reasoning: bool = False + ignore_sql_generation_reasoning: bool = True enable_column_pruning: bool = False - use_dry_plan: bool = False - allow_dry_plan_fallback: bool = True + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False custom_instruction: Optional[str] = None @@ -94,17 +203,62 @@ class AskResultResponse(_AskResultResponse): ] = Field(None, exclude=True) +def _build_sql_correction_error( + error_message: str | None, + sql_diagnosis_reasoning: str | None = None, +) -> str: + raw_error = error_message or "" + if sql_diagnosis_reasoning: + return ( + f"{sql_diagnosis_reasoning}\n\n" + f"Original Wren Engine validation error:\n{raw_error}" + ) + + return f"Original Wren Engine validation error:\n{raw_error}" + + +class _AskStageTimer: + def __init__(self, query_id: str, project_id: str | None): + self._query_id = query_id + self._project_id = project_id or "" + self._started_at = time.perf_counter() + self._last_at = self._started_at + + def mark( + self, + stage: str, + started_at: float | None = None, + **fields: Any, + ) -> None: + ended_at = time.perf_counter() + stage_started_at = started_at if started_at is not None else self._last_at + suffix = " ".join( + f"{key}={value}" for key, value in fields.items() if value is not None + ) + logger.info( + "Ask timing query_id=%s project_id=%s stage=%s elapsed_ms=%.1f total_ms=%.1f%s%s", + self._query_id, + self._project_id, + stage, + (ended_at - stage_started_at) * 1000, + (ended_at - self._started_at) * 1000, + " " if suffix else "", + suffix, + ) + self._last_at = ended_at + + class AskService: def __init__( self, pipelines: Dict[str, BasicPipeline], allow_intent_classification: bool = True, - allow_sql_generation_reasoning: bool = True, + allow_sql_generation_reasoning: bool = False, allow_sql_functions_retrieval: bool = True, allow_sql_diagnosis: bool = True, allow_sql_knowledge_retrieval: bool = True, - enable_column_pruning: bool = False, - max_sql_correction_retries: int = 3, + enable_column_pruning: bool = True, + max_sql_correction_retries: int = 0, max_histories: int = 5, maxsize: int = 1_000_000, ttl: int = 120, @@ -122,6 +276,79 @@ def __init__( self._max_histories = max_histories self._max_sql_correction_retries = max_sql_correction_retries + async def _retrieve_schema_context( + self, + ask_request: AskRequest, + user_query: str, + histories: list[AskHistory], + enable_column_pruning: bool, + timer: _AskStageTimer, + phase: str | None = None, + tables: list[str] | None = None, + ) -> tuple[dict, list[dict], list[str], list[str]]: + schema_retrieval_started_at = time.perf_counter() + retrieval_result = await self._pipelines["db_schema_retrieval"].run( + query=user_query, + tables=tables, + histories=histories, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + enable_column_pruning=enable_column_pruning, + ) + retrieval_payload = retrieval_result.get("construct_retrieval_results", {}) + documents = retrieval_payload.get("retrieval_results", []) + table_names = [document.get("table_name") for document in documents] + table_ddls = [document.get("table_ddl") for document in documents] + timer.mark( + "schema_retrieval", + schema_retrieval_started_at, + retrieved_table_count=len(table_names), + phase=phase, + ) + return retrieval_payload, documents, table_names, table_ddls + + async def _run_schema_fast_path( + self, + ask_request: AskRequest, + user_query: str, + table_ddls: list[str], + histories: list[AskHistory], + grounding_query: str, + use_dry_plan: bool, + allow_dry_plan_fallback: bool, + timer: _AskStageTimer, + phase: str | None = None, + ) -> dict | None: + fast_path_pipeline_name = ( + "followup_sql_generation" if histories else "sql_generation" + ) + fast_path_pipeline = self._pipelines[fast_path_pipeline_name] + fast_path_runner = getattr( + fast_path_pipeline, + "run_deterministic_fast_path", + None, + ) + if not fast_path_runner: + return None + + fast_path_started_at = time.perf_counter() + fast_path_result = await fast_path_runner( + query=user_query, + contexts=table_ddls, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + grounding_query=grounding_query, + ) + timer.mark( + "sql_generation_fast_path", + fast_path_started_at, + result=fast_path_result.get("fast_path") if fast_path_result else "miss", + phase=phase, + ) + return fast_path_result + def _is_stopped(self, query_id: str, container: dict): if ( result := container.get(query_id) @@ -159,8 +386,11 @@ async def ask( instructions = [] api_results = [] table_names = [] + table_ddls = [] + _retrieval_result = {} error_message = None invalid_sql = None + fast_path_terminal = False allow_sql_generation_reasoning = ( self._allow_sql_generation_reasoning and not ask_request.ignore_sql_generation_reasoning @@ -178,7 +408,14 @@ async def ask( sql_knowledge = None try: - user_query = ask_request.query + original_user_query = ask_request.query + user_query = original_user_query + timer = _AskStageTimer(query_id, ask_request.project_id) + timer.mark( + "frontend_request", + query_chars=len(original_user_query or ""), + history_count=len(histories), + ) # ask status can be understanding, searching, generating, finished, failed, stopped # we will need to handle business logic for each status @@ -189,51 +426,132 @@ async def ask( is_followup=True if histories else False, ) - historical_question = await self._pipelines["historical_question"].run( - query=user_query, - project_id=ask_request.project_id, - ) - - # we only return top 1 result - historical_question_result = historical_question.get( - "formatted_output", {} - ).get("documents", [])[:1] - - if historical_question_result: - api_results = [ - AskResult( - **{ - "sql": result.get("statement"), - "type": "view" if result.get("viewId") else "llm", - "viewId": result.get("viewId"), - } - ) - for result in historical_question_result - ] - sql_generation_reasoning = "" - else: + if not api_results: # Run both pipeline operations concurrently + support_context_started_at = time.perf_counter() sql_samples_task, instructions_task = await asyncio.gather( self._pipelines["sql_pairs_retrieval"].run( query=user_query, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ), self._pipelines["instructions_retrieval"].run( query=user_query, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, scope="sql", ), ) - # Extract results from completed tasks sql_samples = sql_samples_task["formatted_output"].get( "documents", [] ) instructions = instructions_task["formatted_output"].get( "documents", [] ) + timer.mark( + "schema_retrieval_support_context", + support_context_started_at, + sql_sample_count=len(sql_samples), + instruction_count=len(instructions), + ) - if self._allow_intent_classification: + if ( + _looks_like_simple_analytics_request(user_query) + and not self._is_stopped(query_id, self._ask_results) + ): + pre_intent_grounding_query = ( + _build_fast_path_grounding_query(user_query, histories) + if histories + else original_user_query + ) + pre_intent_tables = ( + _history_sql_table_names(histories[0].sql) + if histories + else None + ) + self._ask_results[query_id] = AskResultResponse( + status="searching", + type="TEXT_TO_SQL", + rephrased_question=rephrased_question, + intent_reasoning=intent_reasoning, + trace_id=trace_id, + is_followup=True if histories else False, + ) + ( + _retrieval_result, + documents, + table_names, + table_ddls, + ) = await self._retrieve_schema_context( + ask_request=ask_request, + user_query="" + if pre_intent_tables + else pre_intent_grounding_query, + histories=histories, + enable_column_pruning=enable_column_pruning, + timer=timer, + phase="pre_intent", + tables=pre_intent_tables, + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="schema_retrieval") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + if documents: + fast_path_result = await self._run_schema_fast_path( + ask_request=ask_request, + user_query=user_query, + table_ddls=table_ddls, + histories=histories, + grounding_query=pre_intent_grounding_query, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + timer=timer, + phase="pre_intent", + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark( + "cancelled", + at_stage="sql_generation_fast_path", + ) + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if fast_path_result: + post_process = fast_path_result["post_process"] + if sql_valid_result := post_process.get( + "valid_generation_result" + ): + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + fast_path_terminal = True + elif ( + failed_result := post_process.get( + "invalid_generation_result" + ) + ) and ( + failed_result.get("type") == "NO_RELEVANT_SQL" + and _can_return_pre_intent_schema_unsupported( + user_query + ) + ): + error_message = failed_result.get("error") + invalid_sql = "" + fast_path_terminal = True + + if ( + self._allow_intent_classification + and not api_results + and not fast_path_terminal + ): + intent_started_at = time.perf_counter() intent_classification_result = ( await self._pipelines["intent_classification"].run( query=user_query, @@ -241,6 +559,7 @@ async def ask( sql_samples=sql_samples, instructions=instructions, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, configuration=ask_request.configurations, ) ).get("post_process", {}) @@ -249,10 +568,20 @@ async def ask( "rephrased_question" ) intent_reasoning = intent_classification_result.get("reasoning") + timer.mark( + "llm_intent_generation", + intent_started_at, + intent=intent, + ) if rephrased_question: user_query = rephrased_question + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="llm_intent_generation") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if intent == "MISLEADING_QUERY": asyncio.create_task( self._pipelines["misleading_assistance"].run( @@ -333,7 +662,16 @@ async def ask( trace_id=trace_id, is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + grounding_query = ( + _build_fast_path_grounding_query(user_query, histories) + if histories + else original_user_query + ) + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not fast_path_terminal + ): self._ask_results[query_id] = AskResultResponse( status="searching", type="TEXT_TO_SQL", @@ -343,18 +681,23 @@ async def ask( is_followup=True if histories else False, ) - retrieval_result = await self._pipelines["db_schema_retrieval"].run( - query=user_query, + ( + _retrieval_result, + documents, + table_names, + table_ddls, + ) = await self._retrieve_schema_context( + ask_request=ask_request, + user_query=user_query, histories=histories, - project_id=ask_request.project_id, enable_column_pruning=enable_column_pruning, + timer=timer, ) - _retrieval_result = retrieval_result.get( - "construct_retrieval_results", {} - ) - documents = _retrieval_result.get("retrieval_results", []) - table_names = [document.get("table_name") for document in documents] - table_ddls = [document.get("table_ddl") for document in documents] + + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="schema_retrieval") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if not documents: logger.exception(f"ask pipeline - NO_RELEVANT_DATA: {user_query}") @@ -375,9 +718,44 @@ async def ask( results["metadata"]["type"] = "TEXT_TO_SQL" return results + fast_path_result = await self._run_schema_fast_path( + ask_request=ask_request, + user_query=user_query, + table_ddls=table_ddls, + histories=histories, + grounding_query=grounding_query, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, + timer=timer, + ) + if fast_path_result: + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_generation_fast_path") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + post_process = fast_path_result["post_process"] + if sql_valid_result := post_process.get( + "valid_generation_result" + ): + api_results = [ + AskResult( + **{ + "sql": sql_valid_result.get("sql"), + "type": "llm", + } + ) + ] + fast_path_terminal = True + elif failed_result := post_process.get("invalid_generation_result"): + if failed_result.get("type") == "NO_RELEVANT_SQL": + error_message = failed_result.get("error") + invalid_sql = "" + fast_path_terminal = True + if ( not self._is_stopped(query_id, self._ask_results) and not api_results + and not fast_path_terminal and allow_sql_generation_reasoning ): self._ask_results[query_id] = AskResultResponse( @@ -390,6 +768,7 @@ async def ask( is_followup=True if histories else False, ) + sql_reasoning_started_at = time.perf_counter() if histories: sql_generation_reasoning = ( await self._pipelines["followup_sql_generation_reasoning"].run( @@ -398,6 +777,8 @@ async def ask( histories=histories, sql_samples=sql_samples, instructions=instructions, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, configuration=ask_request.configurations, query_id=query_id, ) @@ -409,10 +790,21 @@ async def ask( contexts=table_ddls, sql_samples=sql_samples, instructions=instructions, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, configuration=ask_request.configurations, query_id=query_id, ) ).get("post_process", {}) + timer.mark( + "sql_generation_reasoning", + sql_reasoning_started_at, + is_followup=bool(histories), + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_generation_reasoning") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results self._ask_results[query_id] = AskResultResponse( status="planning", @@ -425,7 +817,11 @@ async def ask( is_followup=True if histories else False, ) - if not self._is_stopped(query_id, self._ask_results) and not api_results: + if ( + not self._is_stopped(query_id, self._ask_results) + and not api_results + and not fast_path_terminal + ): self._ask_results[query_id] = AskResultResponse( status="generating", type="TEXT_TO_SQL", @@ -437,11 +833,13 @@ async def ask( is_followup=True if histories else False, ) + auxiliary_retrieval_started_at = time.perf_counter() if allow_sql_functions_retrieval: sql_functions = await self._pipelines[ "sql_functions_retrieval" ].run( project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ) else: sql_functions = [] @@ -451,7 +849,18 @@ async def ask( "sql_knowledge_retrieval" ].run( project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ) + timer.mark( + "sql_auxiliary_retrieval", + auxiliary_retrieval_started_at, + functions_enabled=allow_sql_functions_retrieval, + knowledge_enabled=allow_sql_knowledge_retrieval, + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_auxiliary_retrieval") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results has_calculated_field = _retrieval_result.get( "has_calculated_field", False @@ -459,6 +868,7 @@ async def ask( has_metric = _retrieval_result.get("has_metric", False) has_json_field = _retrieval_result.get("has_json_field", False) + sql_generation_started_at = time.perf_counter() if histories: text_to_sql_generation_results = await self._pipelines[ "followup_sql_generation" @@ -468,6 +878,7 @@ async def ask( sql_generation_reasoning=sql_generation_reasoning, histories=histories, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -477,6 +888,7 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + grounding_query=grounding_query, ) else: text_to_sql_generation_results = await self._pipelines[ @@ -486,6 +898,7 @@ async def ask( contexts=table_ddls, sql_generation_reasoning=sql_generation_reasoning, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -495,7 +908,22 @@ async def ask( use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, + grounding_query=original_user_query, ) + timer.mark( + "sql_generation", + sql_generation_started_at, + is_followup=bool(histories), + status="valid" + if text_to_sql_generation_results["post_process"].get( + "valid_generation_result" + ) + else "invalid", + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_generation") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if sql_valid_result := text_to_sql_generation_results["post_process"][ "valid_generation_result" @@ -512,13 +940,27 @@ async def ask( "post_process" ]["invalid_generation_result"]: while current_sql_correction_retries < max_sql_correction_retries: - if failed_dry_run_result["type"] == "TIME_OUT": + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_correction") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + if failed_dry_run_result["type"] in ( + "TIME_OUT", + "NO_RELEVANT_SQL", + ): + error_message = failed_dry_run_result["error"] + invalid_sql = ( + "" + if failed_dry_run_result["type"] == "NO_RELEVANT_SQL" + else failed_dry_run_result["sql"] + ) break original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] current_sql_correction_retries += 1 + sql_diagnosis_reasoning = None self._ask_results[query_id] = AskResultResponse( status="correcting", @@ -532,6 +974,7 @@ async def ask( ) if allow_sql_diagnosis: + diagnosis_started_at = time.perf_counter() sql_diagnosis_results = await self._pipelines[ "sql_diagnosis" ].run( @@ -540,28 +983,60 @@ async def ask( invalid_sql=invalid_sql, error_message=error_message, language=ask_request.configurations.language, + project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") + timer.mark( + "sql_diagnosis", + diagnosis_started_at, + retry=current_sql_correction_retries, + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_diagnosis") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results + + correction_error = _build_sql_correction_error( + error_message=error_message, + sql_diagnosis_reasoning=sql_diagnosis_reasoning, + ) + correction_started_at = time.perf_counter() sql_correction_results = await self._pipelines[ "sql_correction" ].run( contexts=table_ddls, + query=user_query, + sql_generation_reasoning=sql_generation_reasoning, instructions=instructions, invalid_generation_result={ "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "error": correction_error, }, project_id=ask_request.project_id, + mdl_hash=ask_request.mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_functions=sql_functions, sql_knowledge=sql_knowledge, ) + timer.mark( + "sql_correction", + correction_started_at, + retry=current_sql_correction_retries, + status="valid" + if sql_correction_results["post_process"].get( + "valid_generation_result" + ) + else "invalid", + ) + if self._is_stopped(query_id, self._ask_results): + timer.mark("cancelled", at_stage="sql_correction") + results["metadata"]["type"] = "TEXT_TO_SQL" + return results if valid_generation_result := sql_correction_results[ "post_process" @@ -595,6 +1070,7 @@ async def ask( ) results["ask_result"] = api_results results["metadata"]["type"] = "TEXT_TO_SQL" + timer.mark("ask_total", status="finished") else: logger.exception(f"ask pipeline - NO_RELEVANT_SQL: {user_query}") if not self._is_stopped(query_id, self._ask_results): @@ -616,6 +1092,7 @@ async def ask( results["metadata"]["error_type"] = "NO_RELEVANT_SQL" results["metadata"]["error_message"] = error_message results["metadata"]["type"] = "TEXT_TO_SQL" + timer.mark("ask_total", status="failed", error_type="NO_RELEVANT_SQL") return results except Exception as e: @@ -641,9 +1118,16 @@ def stop_ask( self, stop_ask_request: StopAskRequest, ): + started_at = time.perf_counter() self._ask_results[stop_ask_request.query_id] = AskResultResponse( status="stopped", ) + logger.info( + "Ask timing query_id=%s project_id=%s stage=cancel_request elapsed_ms=%.1f status=stopped", + stop_ask_request.query_id, + stop_ask_request.project_id or "", + (time.perf_counter() - started_at) * 1000, + ) def get_ask_result( self, diff --git a/wren-ai-service/src/web/v1/services/ask_feedback.py b/wren-ai-service/src/web/v1/services/ask_feedback.py index 9c5b06a772..ba73369c01 100644 --- a/wren-ai-service/src/web/v1/services/ask_feedback.py +++ b/wren-ai-service/src/web/v1/services/ask_feedback.py @@ -4,7 +4,7 @@ from cachetools import TTLCache from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import AliasChoices, BaseModel, Field from src.core.pipeline import BasicPipeline from src.utils import trace_metadata @@ -20,6 +20,9 @@ class AskFeedbackRequest(BaseRequest): tables: List[str] sql_generation_reasoning: str sql: str + mdl_hash: Optional[str] = Field( + default=None, validation_alias=AliasChoices("mdl_hash", "id") + ) class AskFeedbackResponse(BaseModel): @@ -122,14 +125,17 @@ async def ask_feedback( self._pipelines["db_schema_retrieval"].run( tables=ask_feedback_request.tables, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ), self._pipelines["sql_pairs_retrieval"].run( query=ask_feedback_request.question, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ), self._pipelines["instructions_retrieval"].run( query=ask_feedback_request.question, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, scope="sql", ), ) @@ -139,6 +145,7 @@ async def ask_feedback( "sql_functions_retrieval" ].run( project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ) else: sql_functions = [] @@ -148,6 +155,7 @@ async def ask_feedback( "sql_knowledge_retrieval" ].run( project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ) # Extract results from completed tasks @@ -176,9 +184,11 @@ async def ask_feedback( "sql_regeneration" ].run( contexts=table_ddls, + query=ask_feedback_request.question, sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, sql=ask_feedback_request.sql, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, sql_samples=sql_samples, instructions=instructions, has_calculated_field=has_calculated_field, @@ -202,10 +212,14 @@ async def ask_feedback( elif failed_dry_run_result := text_to_sql_generation_results[ "post_process" ]["invalid_generation_result"]: - if failed_dry_run_result["type"] != "TIME_OUT": + if failed_dry_run_result["type"] == "NO_RELEVANT_SQL": + invalid_sql = "" + error_message = failed_dry_run_result["error"] + elif failed_dry_run_result["type"] != "TIME_OUT": original_sql = failed_dry_run_result["original_sql"] invalid_sql = failed_dry_run_result["sql"] error_message = failed_dry_run_result["error"] + sql_diagnosis_reasoning = None self._ask_feedback_results[ query_id @@ -222,23 +236,34 @@ async def ask_feedback( original_sql=original_sql, invalid_sql=invalid_sql, error_message=error_message, + language=ask_feedback_request.configurations.language, + project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, ) sql_diagnosis_reasoning = sql_diagnosis_results[ "post_process" ].get("reasoning") + correction_error_message = error_message + if sql_diagnosis_reasoning: + correction_error_message = ( + f"{error_message}\nDiagnosis: {sql_diagnosis_reasoning}" + ) + sql_correction_results = await self._pipelines[ "sql_correction" ].run( contexts=table_ddls, + query=ask_feedback_request.question, + sql_generation_reasoning=ask_feedback_request.sql_generation_reasoning, instructions=instructions, invalid_generation_result={ - "sql": original_sql, - "error": sql_diagnosis_reasoning - if allow_sql_diagnosis - else error_message, + "original_sql": original_sql, + "sql": invalid_sql, + "error": correction_error_message, }, project_id=ask_feedback_request.project_id, + mdl_hash=ask_feedback_request.mdl_hash, sql_functions=sql_functions, sql_knowledge=sql_knowledge, ) diff --git a/wren-ai-service/src/web/v1/services/chart.py b/wren-ai-service/src/web/v1/services/chart.py index ed3f47f1bc..46f28820e6 100644 --- a/wren-ai-service/src/web/v1/services/chart.py +++ b/wren-ai-service/src/web/v1/services/chart.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from src.core.pipeline import BasicPipeline +from src.pipelines.generation.utils.chart import build_fallback_chart_result from src.utils import trace_metadata from src.web.v1.services import BaseRequest @@ -139,6 +140,21 @@ async def chart( trace_id=trace_id, ) + fallback_chart_result = build_fallback_chart_result( + chart_request.query, + sql_data, + ) + if fallback_chart_result.get("chart_schema") or fallback_chart_result.get( + "reasoning" + ): + self._chart_results[query_id] = ChartResultResponse( + status="finished", + response=ChartResult(**fallback_chart_result), + trace_id=trace_id, + ) + results["chart_result"] = fallback_chart_result + return results + chart_generation_result = await self._pipelines["chart_generation"].run( query=chart_request.query, sql=chart_request.sql, diff --git a/wren-ai-service/src/web/v1/services/chart_adjustment.py b/wren-ai-service/src/web/v1/services/chart_adjustment.py index 92c1c590b9..594a74c20b 100644 --- a/wren-ai-service/src/web/v1/services/chart_adjustment.py +++ b/wren-ai-service/src/web/v1/services/chart_adjustment.py @@ -1,5 +1,5 @@ import logging -from typing import Dict, Literal, Optional +from typing import Any, Dict, Literal, Optional from cachetools import TTLCache from langfuse.decorators import observe @@ -27,6 +27,7 @@ class ChartAdjustmentOption(BaseModel): class ChartAdjustmentRequest(BaseRequest): query: str sql: str + data: Optional[Dict[str, Any]] = None adjustment_option: ChartAdjustmentOption chart_schema: dict @@ -117,15 +118,21 @@ async def chart_adjustment( trace_id=trace_id, ) - execute_sql_result = ( - await self._pipelines["sql_executor"].run( - sql=chart_adjustment_request.sql, - project_id=chart_adjustment_request.project_id, + if not chart_adjustment_request.data: + execute_sql_result = ( + await self._pipelines["sql_executor"].run( + sql=chart_adjustment_request.sql, + project_id=chart_adjustment_request.project_id, + ) + )["execute_sql"] + + sql_data = execute_sql_result["results"] + execute_sql_error_message = execute_sql_result.get( + "error_message", None ) - )["execute_sql"] - - sql_data = execute_sql_result["results"] - execute_sql_error_message = execute_sql_result.get("error_message", None) + else: + sql_data = chart_adjustment_request.data + execute_sql_error_message = None if execute_sql_error_message: self._chart_adjustment_results[ diff --git a/wren-ai-service/src/web/v1/services/instructions.py b/wren-ai-service/src/web/v1/services/instructions.py index 8a825938d2..1dae1649dd 100644 --- a/wren-ai-service/src/web/v1/services/instructions.py +++ b/wren-ai-service/src/web/v1/services/instructions.py @@ -103,6 +103,7 @@ async def index( await self._pipelines["instructions_indexing"].run( project_id=request.project_id, + mdl_hash=request.mdl_hash, instructions=instructions, ) diff --git a/wren-ai-service/src/web/v1/services/question_recommendation.py b/wren-ai-service/src/web/v1/services/question_recommendation.py index 6033237a45..f13107849d 100644 --- a/wren-ai-service/src/web/v1/services/question_recommendation.py +++ b/wren-ai-service/src/web/v1/services/question_recommendation.py @@ -4,9 +4,9 @@ import orjson from cachetools import TTLCache -from langfuse.decorators import observe from pydantic import BaseModel +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest, MetadataTraceable @@ -59,6 +59,16 @@ def _handle_exception( ) logger.error(error_message) + @staticmethod + def _is_valid_candidate(candidate: dict) -> bool: + return ( + isinstance(candidate, dict) + and isinstance(candidate.get("question"), str) + and bool(candidate["question"].strip()) + and isinstance(candidate.get("category"), str) + and bool(candidate["category"].strip()) + ) + @observe(name="Validate Question") async def _validate_question( self, @@ -67,8 +77,16 @@ async def _validate_question( max_questions: int, max_categories: int, project_id: Optional[str] = None, - allow_data_preview: bool = True, + allow_data_preview: bool = False, + use_dry_plan: bool = True, + allow_dry_plan_fallback: bool = False, ): + if not self._is_valid_candidate(candidate): + logger.warning( + f"Request {request_id}: Skipping invalid recommended question payload." + ) + return None + async def _document_retrieval() -> tuple[list[str], bool, bool, bool]: retrieval_result = await self._pipelines["db_schema_retrieval"].run( query=candidate["question"], @@ -131,6 +149,8 @@ async def _instructions_retrieval() -> list[dict]: has_metric=has_metric, has_json_field=has_json_field, sql_functions=sql_functions, + use_dry_plan=use_dry_plan, + allow_dry_plan_fallback=allow_dry_plan_fallback, allow_data_preview=allow_data_preview, sql_knowledge=sql_knowledge, ) @@ -172,11 +192,23 @@ class Request(BaseRequest): max_questions: int = 5 max_categories: int = 3 regenerate: bool = False - allow_data_preview: bool = True + allow_data_preview: bool = False + use_dry_plan: bool = True + allow_dry_plan_fallback: bool = False async def _recommend(self, request: dict): resp = await self._pipelines["question_recommendation"].run(**request) - questions = resp.get("normalized", {}).get("questions", []) + normalized = resp.get("normalized", {}) + if isinstance(normalized, list): + questions = normalized + elif isinstance(normalized, dict): + questions = normalized.get("questions", []) + else: + questions = [] + + questions = [ + question for question in questions if self._is_valid_candidate(question) + ] validation_tasks = [ self._validate_question( question, @@ -185,6 +217,8 @@ async def _recommend(self, request: dict): request["max_categories"], project_id=request["project_id"], allow_data_preview=request["allow_data_preview"], + use_dry_plan=request["use_dry_plan"], + allow_dry_plan_fallback=request["allow_dry_plan_fallback"], ) for question in questions ] @@ -218,6 +252,8 @@ async def recommend(self, input: Request, **kwargs) -> Event: "project_id": input.project_id, "event_id": input.event_id, "allow_data_preview": input.allow_data_preview, + "use_dry_plan": input.use_dry_plan, + "allow_dry_plan_fallback": input.allow_dry_plan_fallback, } await self._recommend(request) diff --git a/wren-ai-service/src/web/v1/services/relationship_recommendation.py b/wren-ai-service/src/web/v1/services/relationship_recommendation.py index 25c0bc5820..75ea10a1f3 100644 --- a/wren-ai-service/src/web/v1/services/relationship_recommendation.py +++ b/wren-ai-service/src/web/v1/services/relationship_recommendation.py @@ -1,3 +1,4 @@ +import asyncio import logging from typing import Dict, Literal, Optional @@ -35,11 +36,13 @@ def __init__( pipelines: Dict[str, BasicPipeline], maxsize: int = 1_000_000, ttl: int = 120, + generation_timeout_seconds: float = 180.0, ): self._pipelines = pipelines self._cache: Dict[str, RelationshipRecommendation.Resource] = TTLCache( maxsize=maxsize, ttl=ttl ) + self._generation_timeout_seconds = generation_timeout_seconds def _handle_exception( self, @@ -72,12 +75,32 @@ async def recommend(self, request: Input, **kwargs) -> Resource: "language": request.configurations.language, } - resp = await self._pipelines["relationship_recommendation"].run(**input) + logger.info( + "Calling configured LLM for relationship recommendations. " + "timeout_seconds=%s", + self._generation_timeout_seconds, + ) + try: + resp = await asyncio.wait_for( + self._pipelines["relationship_recommendation"].run(**input), + timeout=self._generation_timeout_seconds, + ) + except TimeoutError: + raise TimeoutError( + "Relationship recommendation generation timed out after " + f"{self._generation_timeout_seconds} seconds" + ) + response = resp.get("validated") + if response is None: + raise ValueError( + "Relationship recommendation pipeline returned no " + "validated response" + ) self._cache[request.id] = self.Resource( id=request.id, status="finished", - response=resp.get("validated"), + response=response, trace_id=trace_id, request_from=request.request_from, ) diff --git a/wren-ai-service/src/web/v1/services/semantics_description.py b/wren-ai-service/src/web/v1/services/semantics_description.py index 85bf14b188..c914a40e2b 100644 --- a/wren-ai-service/src/web/v1/services/semantics_description.py +++ b/wren-ai-service/src/web/v1/services/semantics_description.py @@ -86,10 +86,122 @@ def _chunking( for chunk in chunks ] - async def _generate_task(self, request_id: str, chunk: dict): + def _description(self, payload: dict) -> str: + properties = payload.get("properties") + if not isinstance(properties, dict): + properties = {} + value = payload.get("description") or properties.get("description", "") + return "" if value is None else str(value).strip() + + def _validate_chunk_output(self, chunk: dict, output: dict) -> dict: + if not isinstance(output, dict): + raise ValueError("Semantics description pipeline returned invalid output") + + selected_models = set(chunk.get("selected_models", [])) + models = { + model.get("name"): model + for model in chunk.get("mdl", {}).get("models", []) + if model.get("name") in selected_models + } + + for model_name, model in models.items(): + generated_model = output.get(model_name) + if not isinstance(generated_model, dict): + raise ValueError( + f"Semantics description output omitted selected model: {model_name}" + ) + + if not self._description(generated_model): + raise ValueError( + "Semantics description output omitted description for model: " + f"{model_name}" + ) + + generated_columns = { + column.get("name"): column + for column in generated_model.get("columns", []) + if isinstance(column, dict) and column.get("name") + } + + for column in model.get("columns", []): + if not isinstance(column, dict): + continue + + column_name = column.get("name", "") + generated_column = generated_columns.get(column_name) + if not generated_column: + raise ValueError( + "Semantics description output omitted selected column: " + f"{model_name}.{column_name}" + ) + + if not self._description(generated_column): + raise ValueError( + "Semantics description output omitted description for column: " + f"{model_name}.{column_name}" + ) + + return output + + def _chunk_columns(self, chunk: dict) -> list[dict]: + models = chunk.get("mdl", {}).get("models", []) + if not models: + return [] + return models[0].get("columns", []) or [] + + def _split_chunk(self, chunk: dict) -> list[dict]: + columns = self._chunk_columns(chunk) + if len(columns) <= 1: + return [] + + split_at = max(1, len(columns) // 2) + model = chunk["mdl"]["models"][0] + return [ + { + **chunk, + "mdl": {"models": [{**model, "columns": column_chunk}]}, + } + for column_chunk in (columns[:split_at], columns[split_at:]) + if column_chunk + ] + + def _is_retryable_chunk_error(self, error: Exception) -> bool: + message = str(error) + return ( + "malformed JSON" in message + or "omitted selected model" in message + or "omitted description for model" in message + or "omitted selected column" in message + or "omitted description for column" in message + ) + + async def _generate_chunk(self, chunk: dict) -> dict: resp = await self._pipelines["semantics_description"].run(**chunk) - output = resp.get("output") + output = resp.get("output") or {} + return self._validate_chunk_output(chunk, output) + + async def _generate_chunk_with_retry_splitting(self, chunk: dict) -> list[dict]: + try: + return [await self._generate_chunk(chunk)] + except ValueError as e: + split_chunks = self._split_chunk(chunk) + if not split_chunks or not self._is_retryable_chunk_error(e): + raise + model_name = chunk.get("selected_models", [""])[0] + logger.warning( + "Retrying semantics description for model %s with smaller " + "column chunks after incomplete or malformed response.", + model_name, + ) + outputs: list[dict] = [] + for split_chunk in split_chunks: + outputs.extend( + await self._generate_chunk_with_retry_splitting(split_chunk) + ) + return outputs + + def _merge_output(self, request_id: str, output: dict): current = self[request_id] current.response = current.response or {} @@ -100,6 +212,11 @@ async def _generate_task(self, request_id: str, chunk: dict): current.response[key]["columns"].extend(output[key]["columns"]) + async def _generate_task(self, request_id: str, chunk: dict): + outputs = await self._generate_chunk_with_retry_splitting(chunk) + for output in outputs: + self._merge_output(request_id, output) + @observe(name="Generate Semantics Description") @trace_metadata async def generate(self, request: GenerateRequest, **kwargs) -> Resource: diff --git a/wren-ai-service/src/web/v1/services/semantics_preparation.py b/wren-ai-service/src/web/v1/services/semantics_preparation.py index 2ff6215cbe..5de4d7847d 100644 --- a/wren-ai-service/src/web/v1/services/semantics_preparation.py +++ b/wren-ai-service/src/web/v1/services/semantics_preparation.py @@ -28,7 +28,7 @@ class SemanticsPreparationResponse(BaseModel): # GET /v1/semantics-preparations/{mdl_hash}/status -class SemanticsPreparationStatusRequest(BaseModel): +class SemanticsPreparationStatusRequest(BaseRequest): # don't recommend to use id as a field name, but it's used in the API spec # so we need to support as a choice, and will remove it in the future mdl_hash: str = Field(validation_alias=AliasChoices("mdl_hash", "id")) @@ -54,6 +54,9 @@ def __init__( self._prepare_semantics_statuses: Dict[ str, SemanticsPreparationStatusResponse ] = TTLCache(maxsize=maxsize, ttl=ttl) + self._prepare_semantics_project_ids: Dict[str, str] = TTLCache( + maxsize=maxsize, ttl=ttl + ) @observe(name="Prepare Semantics") @trace_metadata @@ -76,26 +79,53 @@ async def prepare_semantics( input = { "mdl_str": prepare_semantics_request.mdl, "project_id": prepare_semantics_request.project_id, + "mdl_hash": prepare_semantics_request.mdl_hash, } + project_scoped_index_names = [ + "db_schema", + "historical_question", + "table_description", + "sql_pairs", + "project_meta", + ] + + await asyncio.gather( + *[ + self._pipelines[name].clean( + project_id=prepare_semantics_request.project_id, + delete_all=True, + ) + if name == "sql_pairs" + else self._pipelines[name].clean( + project_id=prepare_semantics_request.project_id + ) + for name in project_scoped_index_names + ] + ) tasks = [ self._pipelines[name].run(**input) - for name in [ - "db_schema", - "historical_question", - "table_description", - "sql_pairs", - "project_meta", - ] + for name in project_scoped_index_names ] await asyncio.gather(*tasks) + if not await self._has_indexed_schema_documents( + prepare_semantics_request.project_id, + prepare_semantics_request.mdl_hash, + ): + raise RuntimeError( + "No indexed schema documents were found for the prepared deployment" + ) + self._prepare_semantics_statuses[ prepare_semantics_request.mdl_hash ] = SemanticsPreparationStatusResponse( status="finished", ) + self._prepare_semantics_project_ids[ + prepare_semantics_request.mdl_hash + ] = prepare_semantics_request.project_id or "" except Exception as e: logger.exception(f"Failed to prepare semantics: {e}") @@ -114,7 +144,31 @@ async def prepare_semantics( return results - def get_prepare_semantics_status( + async def _has_indexed_schema_documents( + self, + project_id: str, + mdl_hash: str, + ) -> bool: + dbschema_count, table_description_count = await asyncio.gather( + self._pipelines["db_schema"].count_documents( + project_id=project_id, + mdl_hash=mdl_hash, + ), + self._pipelines["table_description"].count_documents( + project_id=project_id, + mdl_hash=mdl_hash, + ), + ) + logger.info( + "Project ID: %s, MDL hash %s indexed schema document counts: db_schema=%s, table_description=%s", + project_id, + mdl_hash, + dbschema_count, + table_description_count, + ) + return dbschema_count > 0 and table_description_count > 0 + + async def get_prepare_semantics_status( self, prepare_semantics_status_request: SemanticsPreparationStatusRequest ) -> SemanticsPreparationStatusResponse: if ( @@ -122,6 +176,15 @@ def get_prepare_semantics_status( prepare_semantics_status_request.mdl_hash ) ) is None: + if ( + prepare_semantics_status_request.project_id + and await self._has_indexed_schema_documents( + prepare_semantics_status_request.project_id, + prepare_semantics_status_request.mdl_hash, + ) + ): + return SemanticsPreparationStatusResponse(status="finished") + logger.exception( f"id is not found for SemanticsPreparation: {prepare_semantics_status_request.mdl_hash}" ) @@ -129,11 +192,35 @@ def get_prepare_semantics_status( status="failed", error=SemanticsPreparationStatusResponse.SemanticsPreparationError( code="OTHERS", - message="{prepare_semantics_status_request.id} is not found", + message=f"{prepare_semantics_status_request.mdl_hash} is not found", ), ) - return result + if result.status != "finished": + return result + + project_id = ( + prepare_semantics_status_request.project_id + or self._prepare_semantics_project_ids.get( + prepare_semantics_status_request.mdl_hash + ) + ) + if not project_id: + return result + + if await self._has_indexed_schema_documents( + project_id, + prepare_semantics_status_request.mdl_hash, + ): + return result + + return SemanticsPreparationStatusResponse( + status="failed", + error=SemanticsPreparationStatusResponse.SemanticsPreparationError( + code="OTHERS", + message="Prepared schema documents are missing for this deployment", + ), + ) @observe(name="Delete Semantics Documents") @trace_metadata diff --git a/wren-ai-service/src/web/v1/services/sql_answer.py b/wren-ai-service/src/web/v1/services/sql_answer.py index c907387300..a87c17d203 100644 --- a/wren-ai-service/src/web/v1/services/sql_answer.py +++ b/wren-ai-service/src/web/v1/services/sql_answer.py @@ -1,5 +1,6 @@ import asyncio import logging +import time from typing import Dict, Literal, Optional from cachetools import TTLCache @@ -73,15 +74,23 @@ async def sql_answer( try: query_id = sql_answer_request.query_id + total_started_at = time.perf_counter() self._sql_answer_results[query_id] = SqlAnswerResultResponse( status="preprocessing", trace_id=trace_id, ) + preprocess_started_at = time.perf_counter() preprocessed_sql_data = self._pipelines["preprocess_sql_data"].run( sql_data=sql_answer_request.sql_data, )["preprocess"] + logger.info( + "Ask timing query_id=%s stage=answer_formatting_preprocess elapsed_ms=%.1f row_count=%s", + query_id, + (time.perf_counter() - preprocess_started_at) * 1000, + preprocessed_sql_data.get("num_rows_used_in_llm"), + ) if preprocessed_sql_data.get("num_rows_used_in_llm") == 0: results["metadata"]["error_type"] = "NO_DATA" @@ -93,6 +102,7 @@ async def sql_answer( trace_id=trace_id, ) + formatting_task_started_at = time.perf_counter() asyncio.create_task( self._pipelines["sql_answer"].run( query=sql_answer_request.query, @@ -104,6 +114,12 @@ async def sql_answer( custom_instruction=sql_answer_request.custom_instruction, ) ) + logger.info( + "Ask timing query_id=%s stage=answer_formatting_task_creation elapsed_ms=%.1f total_ms=%.1f", + query_id, + (time.perf_counter() - formatting_task_started_at) * 1000, + (time.perf_counter() - total_started_at) * 1000, + ) return results except Exception as e: diff --git a/wren-ai-service/src/web/v1/services/sql_corrections.py b/wren-ai-service/src/web/v1/services/sql_corrections.py index 86d0f55301..193973fd4a 100644 --- a/wren-ai-service/src/web/v1/services/sql_corrections.py +++ b/wren-ai-service/src/web/v1/services/sql_corrections.py @@ -2,9 +2,9 @@ from typing import List, Literal, Optional from cachetools import TTLCache -from langfuse.decorators import observe -from pydantic import BaseModel +from pydantic import AliasChoices, BaseModel, Field +from langfuse.decorators import observe from src.core.pipeline import BasicPipeline from src.utils import trace_metadata from src.web.v1.services import BaseRequest, MetadataTraceable @@ -60,6 +60,9 @@ class CorrectionRequest(BaseRequest): event_id: str sql: str error: str + mdl_hash: Optional[str] = Field( + default=None, validation_alias=AliasChoices("mdl_hash", "id") + ) retrieved_tables: Optional[List[str]] = None use_dry_plan: bool = False allow_dry_plan_fallback: bool = True @@ -98,12 +101,14 @@ async def correct( if self._allow_sql_knowledge_retrieval: sql_knowledge = await self._pipelines["sql_knowledge_retrieval"].run( project_id=project_id, + mdl_hash=request.mdl_hash, ) documents = ( ( await self._pipelines["db_schema_retrieval"].run( project_id=project_id, + mdl_hash=request.mdl_hash, tables=retrieved_tables, ) ) @@ -116,6 +121,7 @@ async def correct( contexts=table_ddls, invalid_generation_result=_invalid, project_id=project_id, + mdl_hash=request.mdl_hash, use_dry_plan=use_dry_plan, allow_dry_plan_fallback=allow_dry_plan_fallback, sql_knowledge=sql_knowledge, @@ -131,7 +137,6 @@ async def correct( event_id, f"An error occurred during SQL correction: {error_message}", trace_id=trace_id, - invalid_sql=invalid["sql"], request_from=request.request_from, ) else: diff --git a/wren-ai-service/src/web/v1/services/sql_pairs.py b/wren-ai-service/src/web/v1/services/sql_pairs.py index 84291baf8b..f6750b05cf 100644 --- a/wren-ai-service/src/web/v1/services/sql_pairs.py +++ b/wren-ai-service/src/web/v1/services/sql_pairs.py @@ -69,6 +69,7 @@ async def index( input = { "mdl_str": '{"models": [{"properties": {"boilerplate": "sql_pairs"}}]}', "project_id": request.project_id, + "mdl_hash": request.mdl_hash, "external_pairs": { "sql_pairs": [ sql_pair.model_dump() for sql_pair in request.sql_pairs diff --git a/wren-ai-service/tests/data/config.test.yaml b/wren-ai-service/tests/data/config.test.yaml index 65613d2c83..2c57e4d28a 100644 --- a/wren-ai-service/tests/data/config.test.yaml +++ b/wren-ai-service/tests/data/config.test.yaml @@ -80,7 +80,7 @@ settings: column_indexing_batch_size: 50 doc_endpoint: https://docs.getwren.ai is_oss: true - table_retrieval_size: 10 + table_retrieval_size: 50 table_column_retrieval_size: 1000 query_cache_maxsize: 1000 query_cache_ttl: 3600 diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py new file mode 100644 index 0000000000..9e01ab3bcc --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_chart_generation_utils.py @@ -0,0 +1,188 @@ +import pytest + +from src.pipelines.generation.utils.chart import ( + ChartDataPreprocessor, + _is_schema_compatible_with_sample_data, + build_fallback_chart_result, +) +from src.web.v1.services.chart import ChartRequest, ChartService + + +def test_chart_preprocessor_preserves_sql_result_order(): + data = { + "columns": [{"name": "Market"}, {"name": "Revenue"}], + "data": [[f"Market {index}", index] for index in range(20)], + } + + result = ChartDataPreprocessor().run(data, sample_data_count=3) + + assert result["sample_data"] == [ + {"Market": "Market 0", "Revenue": 0}, + {"Market": "Market 1", "Revenue": 1}, + {"Market": "Market 2", "Revenue": 2}, + ] + + +def test_fallback_chart_counts_categorical_only_results(): + result = build_fallback_chart_result( + "Create a chart comparing completed repairs across engineers.", + { + "columns": [{"name": "Status"}], + "data": [["completed"], ["completed"], ["in-progress"]], + }, + ) + + assert result["chart_type"] == "bar" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Status" + assert result["chart_schema"]["encoding"]["y"] == { + "aggregate": "count", + "type": "quantitative", + "title": "Count", + } + + +def test_fallback_chart_counts_multi_categorical_results(): + result = build_fallback_chart_result( + "Show matching records by status and owner.", + { + "columns": [{"name": "Status"}, {"name": "Owner"}], + "data": [ + ["active", "Asha"], + ["active", "Ben"], + ["inactive", "Asha"], + ], + }, + ) + + assert result["chart_type"] == "grouped_bar" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Status" + assert result["chart_schema"]["encoding"]["y"] == { + "aggregate": "count", + "type": "quantitative", + "title": "Count", + } + assert result["chart_schema"]["encoding"]["color"]["field"] == "Owner" + + +def test_fallback_chart_uses_grouped_bar_for_two_business_dimensions(): + result = build_fallback_chart_result( + "Which Customers have the highest New Orders in each Market?", + { + "columns": [ + {"name": "Market"}, + {"name": "Customer"}, + {"name": "OrderCount"}, + ], + "data": [["North", "Acme", 10], ["South", "Globex", 8]], + }, + ) + + assert result["chart_type"] == "grouped_bar" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Customer" + assert result["chart_schema"]["encoding"]["y"]["field"] == "OrderCount" + assert result["chart_schema"]["encoding"]["color"]["field"] == "Market" + assert result["chart_schema"]["encoding"]["xOffset"]["field"] == "Market" + + +def test_fallback_chart_selects_question_relevant_x_and_numeric_y(): + result = build_fallback_chart_result( + "Show Top 20 Sales Accounts by revenue.", + { + "columns": [ + {"name": "SalesAccount"}, + {"name": "Refund_Status"}, + {"name": "Revenue"}, + ], + "data": [["Acme", "TRANSMITTED", 100], ["Globex", "TRANSMITTED", 200]], + }, + ) + + assert result["chart_schema"]["encoding"]["x"]["field"] == "SalesAccount" + assert result["chart_schema"]["encoding"]["y"]["field"] == "Revenue" + assert result["chart_schema"]["encoding"]["y"]["type"] == "quantitative" + + +def test_fallback_chart_never_uses_same_numeric_field_for_x_and_y(): + result = build_fallback_chart_result( + "Show revenue by discount.", + { + "columns": [{"name": "Discount"}, {"name": "Revenue"}], + "data": [[5, 100], [10, 200]], + }, + ) + + assert result["chart_schema"]["encoding"]["x"]["field"] == "Discount" + assert result["chart_schema"]["encoding"]["y"]["field"] == "Revenue" + assert ( + result["chart_schema"]["encoding"]["x"]["field"] + != result["chart_schema"]["encoding"]["y"]["field"] + ) + + +def test_fallback_chart_defers_numeric_only_results_to_llm_chart_generation(): + assert ( + build_fallback_chart_result( + "Show total metric value.", + { + "columns": [{"name": "metric_value"}], + "data": [[100], [200]], + }, + ) + == {} + ) + + +def test_line_chart_uses_count_y_when_no_numeric_metric_exists(): + result = build_fallback_chart_result( + "Generate a line chart showing monthly repair volume.", + { + "columns": [{"name": "Month"}, {"name": "Status"}], + "data": [["2026-01", "completed"], ["2026-02", "in-progress"]], + }, + ) + + assert result["chart_type"] == "line" + assert result["chart_schema"]["encoding"]["x"]["field"] == "Month" + assert result["chart_schema"]["encoding"]["y"] == { + "aggregate": "count", + "type": "quantitative", + "title": "Repair Count", + } + + +def test_chart_schema_rejects_vega_aggregate_count_without_sql_metric(): + assert not _is_schema_compatible_with_sample_data( + { + "mark": {"type": "bar"}, + "encoding": { + "x": {"field": "Inv Date", "type": "temporal"}, + "y": {"aggregate": "count", "type": "quantitative"}, + }, + }, + [{"Inv Date": "2026-01-01"}, {"Inv Date": "2026-07-01"}], + ) + + +@pytest.mark.asyncio +async def test_chart_service_returns_deterministic_chart_without_llm_wait(): + class FakeChartGenerationPipeline: + async def run(self, **kwargs): + raise AssertionError("chart LLM pipeline should not be called") + + service = ChartService({"chart_generation": FakeChartGenerationPipeline()}) + request = ChartRequest( + query_id="chart-task", + query="Compare customer performance across markets.", + sql="SELECT Market, SUM(Revenue) AS Revenue FROM Sales GROUP BY Market", + data={ + "columns": [{"name": "Market"}, {"name": "Revenue"}], + "data": [["North", 100], ["South", 200]], + }, + ) + + result = await service.chart(request) + + assert result["chart_result"]["reasoning"] == ( + "Generated from the SQL result columns and requested chart type." + ) + assert result["chart_result"]["chart_type"] == "bar" diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py new file mode 100644 index 0000000000..bd11a6470d --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_prompt_grounding_contracts.py @@ -0,0 +1,56 @@ +from pathlib import Path + + +SERVICE_ROOT = Path(__file__).resolve().parents[4] + + +def _read_source(relative_path: str) -> str: + return (SERVICE_ROOT / relative_path).read_text(encoding="utf-8") + + +def test_intent_classification_does_not_require_exact_user_schema_names(): + source = _read_source("src/pipelines/generation/intent_classification.py") + + assert "schema-resolvable references" in source + assert "even if the user did not type exact table or column names" in source + assert "do not require the user to write exact schema identifiers" in source + assert ( + "Do not classify a data retrieval or analytics question as MISLEADING only " + "because the user did not write exact table or column names" + ) in source + + +def test_data_assistance_does_not_invent_hypothetical_schema(): + source = _read_source("src/pipelines/generation/data_assistance.py") + + assert "MUST NOT add SQL code" in source + assert "Use only the provided DATABASE SCHEMA as context" in source + assert "Do not invent, assume, or name tables or columns" in source + assert "do not provide hypothetical schema" in source + + +def test_sql_reasoning_contract_rejects_substitute_identifiers(): + source = _read_source("src/pipelines/generation/utils/sql.py") + + assert "If the DATABASE SCHEMA does not contain an identifier needed" in source + assert "instead of naming a substitute table or column" in source + assert "prompt examples" in source + assert "Identifiers shown in prompt examples are illustrative only" in source + assert "Use `display_label` and `description` only to understand" in source + assert "generated SQL must use that exact identifier" in source + + +def test_sql_correction_receives_raw_wren_engine_validation_error(): + source = _read_source("src/web/v1/services/ask.py") + + assert "_build_sql_correction_error" in source + assert "Original Wren Engine validation error" in source + assert "error_message" in source + + +def test_sql_correction_unknown_identifier_contract(): + source = _read_source("src/pipelines/generation/sql_correction.py") + + assert "If the error reports an unknown table or field" in source + assert "replace it only with an exact executable identifier" in source + assert "Do not retry the same unknown identifier" in source diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py index e7180fa418..6459d96afb 100644 --- a/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_semantics_enrichment.py @@ -1,4 +1,6 @@ -from src.pipelines.generation.semantics_description import output +import pytest + +from src.pipelines.generation.semantics_description import normalize, output def test_without_hallucination(): @@ -64,3 +66,190 @@ def test_with_hallucination_and_no_columns(): assert "model1" in result assert result["model1"]["name"] == "model1" assert len(result["model1"]["columns"]) == 0 + + +def test_single_model_output_is_bound_to_picked_model_name(): + test_normalize = { + "source table": { + "name": "source table", + "columns": [{"name": "entity_code"}], + "properties": { + "description": "Business records for source-table activity.", + "displayName": "source records, source activity", + }, + } + } + test_picked_models = [ + { + "name": "schema_source_table", + "columns": [{"name": "entity_code"}], + } + ] + + result = output(test_normalize, test_picked_models) + + assert list(result) == ["schema_source_table"] + assert result["schema_source_table"]["name"] == "schema_source_table" + assert result["schema_source_table"]["columns"][0]["name"] == "entity_code" + + +def test_single_column_output_is_bound_to_picked_column_name(): + test_normalize = { + "users": { + "name": "users", + "columns": [ + { + "name": "created at", + "properties": { + "description": "Timestamp when the user account was created.", + "displayName": "created date, signup date", + }, + } + ], + } + } + test_picked_models = [ + { + "name": "users", + "columns": [{"name": "created_at"}], + } + ] + + result = output(test_normalize, test_picked_models) + + assert result["users"]["columns"][0]["name"] == "created_at" + assert result["users"]["columns"][0]["properties"]["displayName"] == ( + "created date, signup date" + ) + + +def test_malformed_json_fails_instead_of_returning_empty_output(): + with pytest.raises(ValueError, match="malformed JSON"): + normalize({"replies": ['{"models": [']}) + + +def test_normalize_requires_generated_aliases(): + with pytest.raises(ValueError, match="incomplete semantic metadata"): + normalize( + { + "replies": [ + """ + { + "models": [ + { + "name": "orders", + "properties": { + "description": "Customer order transactions.", + "displayName": "orders" + }, + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Unique order identifier.", + "displayName": "" + } + } + ] + } + ] + } + """ + ] + } + ) + + +def test_normalize_preserves_generated_aliases(): + result = normalize( + { + "replies": [ + """ + { + "models": [ + { + "name": "orders", + "properties": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders" + }, + "columns": [ + { + "name": "order_id", + "type": "VARCHAR", + "properties": { + "description": "Unique order identifier.", + "displayName": "order id, order number" + } + } + ] + } + ] + } + """ + ] + } + ) + + assert result["orders"]["properties"]["displayName"] == "orders, sales orders" + assert result["orders"]["columns"][0]["properties"]["displayName"] == ( + "order id, order number" + ) + assert result["orders"]["columns"][0]["type"] == "VARCHAR" + + +def test_normalize_ignores_extra_llm_fields_but_keeps_semantics(): + result = normalize( + { + "replies": [ + """ + { + "models": [ + { + "name": "orders", + "entity": "transaction", + "properties": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders", + "businessUse": "reporting" + }, + "columns": [ + { + "name": "order_id", + "type": "VARCHAR", + "role": "identifier", + "nullable": false, + "properties": { + "description": "Unique order identifier.", + "displayName": "order id, order number", + "examples": ["1001"] + } + } + ] + } + ] + } + """ + ] + } + ) + + assert result == { + "orders": { + "name": "orders", + "columns": [ + { + "name": "order_id", + "type": "VARCHAR", + "properties": { + "description": "Unique order identifier.", + "displayName": "order id, order number", + }, + } + ], + "properties": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders", + }, + } + } diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py new file mode 100644 index 0000000000..a657582022 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_answer_prompt.py @@ -0,0 +1,86 @@ +from haystack.components.builders.prompt_builder import PromptBuilder + +import pytest + +from src.pipelines.generation.sql_answer import ( + SQLAnswer, + prompt, + sql_to_answer_system_prompt, + sql_to_answer_user_prompt_template, +) + + +def test_sql_answer_prompt_uses_sql_data_rows(): + result = prompt( + query="What is the cost per unit of production volumes for each supplier?", + sql="SELECT supplier_name, manufacturing_cost_per_unit FROM SupplierManufacturing", + sql_data={ + "columns": [ + {"name": "supplier_name", "type": "varchar"}, + {"name": "manufacturing_cost_per_unit", "type": "double"}, + ], + "data": [["Supplier 1", 0.06]], + "row_records": [ + { + "supplier_name": "Supplier 1", + "manufacturing_cost_per_unit": 0.06, + } + ], + }, + language="English", + current_time="2026-08-03T00:00:00", + custom_instruction="", + prompt_builder=PromptBuilder(template=sql_to_answer_user_prompt_template), + ) + + generated_prompt = result["prompt"] + + assert "rows:" in generated_prompt + assert "row records:" in generated_prompt + assert "Supplier 1" in generated_prompt + assert "supplier_name" in generated_prompt + assert "result rows:" not in generated_prompt + assert "Please think step by step" not in generated_prompt + + +def test_sql_answer_prompt_blocks_code_style_hallucinated_analysis(): + assert "Do not write code" in sql_to_answer_system_prompt + assert "Python" in sql_to_answer_system_prompt + assert "running the above code" in sql_to_answer_system_prompt + assert "Do not invent values" in sql_to_answer_system_prompt + + +@pytest.mark.asyncio +async def test_sql_answer_uses_llm_generation_with_query_id(): + class CapturingProvider: + def __init__(self): + self.calls = [] + + def get_generator(self, **_): + async def _generator(**kwargs): + self.calls.append(kwargs) + return {"replies": ["answer"], "metadata": []} + + return _generator + + def get_model(self): + return "test-model" + + provider = CapturingProvider() + pipeline = SQLAnswer(provider) + result = await pipeline.run( + query="show all the active users", + sql="SELECT username, name FROM users WHERE status = 'active'", + sql_data={ + "columns": [ + {"name": "username", "type": "varchar"}, + {"name": "name", "type": "varchar"}, + ], + "data": [["jdoe", "Jane Doe"]], + }, + language="English", + query_id="answer-table-task", + ) + + assert provider.calls[0]["query_id"] == "answer-table-task" + assert result["generate_answer"]["replies"] == ["answer"] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py new file mode 100644 index 0000000000..ed88d0d71f --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_generation_post_processor.py @@ -0,0 +1,212 @@ +from typing import Any + +import aiohttp +import pytest + +from src.core.engine import Engine +from src.providers.llm import ChatRole +from src.pipelines.generation.utils.sql import ( + SQL_GENERATION_MODEL_KWARGS, + SQLGenPostProcessor, + construct_ask_history_messages, + get_text_to_sql_rules, +) +from src.pipelines.retrieval.sql_knowledge import SqlKnowledge + + +class FakeEngine(Engine): + def __init__(self, dry_plan_success: bool = True, execute_success: bool = True): + self.dry_plan_success = dry_plan_success + self.execute_success = execute_success + self.dry_plan_calls: list[dict[str, Any]] = [] + self.execute_sql_calls: list[dict[str, Any]] = [] + + async def dry_plan( + self, + session: aiohttp.ClientSession, + sql: str, + data_source: str, + project_id: str | None = None, + mdl_hash: str | None = None, + allow_fallback: bool = True, + **kwargs, + ): + self.dry_plan_calls.append( + { + "sql": sql, + "data_source": data_source, + "project_id": project_id, + "mdl_hash": mdl_hash, + "allow_fallback": allow_fallback, + } + ) + return self.dry_plan_success, "" if self.dry_plan_success else "plan failed" + + async def execute_sql( + self, + sql: str, + session: aiohttp.ClientSession, + dry_run: bool = True, + **kwargs, + ): + self.execute_sql_calls.append( + { + "sql": sql, + "dry_run": dry_run, + **kwargs, + } + ) + return self.execute_success, {}, {"correlation_id": "correlation-id"} + + +def test_sql_generation_model_kwargs_preserve_strict_schema(): + assert SQL_GENERATION_MODEL_KWARGS["preserve_json_schema"] is True + assert SQL_GENERATION_MODEL_KWARGS["response_format"]["type"] == "json_schema" + assert SQL_GENERATION_MODEL_KWARGS["response_format"]["json_schema"]["strict"] is True + schema = SQL_GENERATION_MODEL_KWARGS["response_format"]["json_schema"]["schema"] + assert schema["additionalProperties"] is False + + +def test_construct_ask_history_messages_matches_legacy_context(): + histories = [{"question": "q", "sql": "SELECT 1"}] + + messages = construct_ask_history_messages(histories) + + assert [(message.role, message.content) for message in messages] == [ + (ChatRole.USER, "q"), + (ChatRole.ASSISTANT, "SELECT 1"), + ] + + +def test_connector_sql_knowledge_supplements_wren_sql_rules(): + sql_knowledge = SqlKnowledge( + { + "text_to_sql_rule": "Connector-only syntax guidance.", + "instructions": {}, + } + ) + + rules = get_text_to_sql_rules(sql_knowledge) + + assert "Generate Wren SQL only" in rules + assert "Connector-only syntax guidance." in rules + assert rules.index("### SQL RULES ###") < rules.index( + "Connector-only syntax guidance." + ) + + +@pytest.mark.asyncio +async def test_post_processor_extracts_tool_call_query_argument(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ['{"name":"query","arguments":{"query":"SELECT 1"}}'], + project_id="project-id", + mdl_hash="manifest-hash", + data_source="mssql", + ) + + assert result["valid_generation_result"] == { + "sql": "SELECT 1", + "correlation_id": "correlation-id", + } + assert engine.dry_plan_calls == [ + { + "sql": "SELECT 1", + "data_source": "mssql", + "project_id": "project-id", + "mdl_hash": "manifest-hash", + "allow_fallback": True, + } + ] + assert engine.execute_sql_calls[0]["dry_run"] is True + + +@pytest.mark.asyncio +async def test_post_processor_returns_no_relevant_sql_for_missing_sql_field(): + processor = SQLGenPostProcessor(FakeEngine()) + + result = await processor.run( + ['{"name":"query","arguments":{"value":"q"}}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "supported SQL field" in result["invalid_generation_result"]["error"] + + +@pytest.mark.asyncio +async def test_post_processor_treats_null_sql_as_no_relevant_sql(): + processor = SQLGenPostProcessor(FakeEngine()) + + result = await processor.run( + ['{"sql": null}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "No grounded SQL" in result["invalid_generation_result"]["error"] + + +@pytest.mark.asyncio +async def test_post_processor_rejects_code_tool_payload(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ['{"name":"execute_code","arguments":{"code":"SELECT 1"}}'], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + +@pytest.mark.asyncio +async def test_post_processor_rejects_plain_text_non_sql_response(): + engine = FakeEngine() + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ["q"], + project_id="project-id", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert "supported SQL JSON payload" in result["invalid_generation_result"]["error"] + assert engine.dry_plan_calls == [] + assert engine.execute_sql_calls == [] + + +@pytest.mark.asyncio +async def test_post_processor_dry_plans_before_preview_execution(): + engine = FakeEngine(dry_plan_success=False) + processor = SQLGenPostProcessor(engine) + + result = await processor.run( + ['{"sql":"SELECT 1"}'], + project_id="project-id", + mdl_hash="manifest-hash", + data_source="mssql", + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"] == { + "sql": "SELECT 1", + "original_sql": "SELECT 1", + "type": "DRY_PLAN", + "error": "plan failed", + "correlation_id": "", + } + assert len(engine.dry_plan_calls) == 1 + assert engine.execute_sql_calls == [] diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py new file mode 100644 index 0000000000..b72785cb3c --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_schema_grounding.py @@ -0,0 +1,1865 @@ +import asyncio + +from src.pipelines.generation.utils.sql import ( + SQLGenPostProcessor, + _SchemaCatalog, + _fallback_limit, + generate_simple_analytics_sql, + normalize_sql_with_schema_identifiers, + normalize_wren_sql_dialect, + sanitize_sql_generation_reasoning, + schema_grounding_failure_message, + unsupported_schema_generation_result, + unsupported_schema_message, + validate_sql_against_contexts, + validate_sql_semantic_coverage, +) + + +class _AcceptingEngine: + async def dry_plan(self, *args, **kwargs): + return True, "" + + async def execute_sql(self, *args, **kwargs): + return True, [], {"correlation_id": "test"} + + +SCHEMA_CONTEXTS = [ + """ + CREATE TABLE valid_invoice_comments ( + invoice_id VARCHAR, + comment_id VARCHAR + ); + """, + """ + CREATE TABLE "valid-order-lines" ( + order_id VARCHAR, + line_amount DECIMAL + ); + """, +] + + +def test_sql_reasoning_sanitizer_blocks_query_shaped_output(): + reasoning = """ + The SQL could look like this: + ```sql + SELECT * FROM dbo_dimOrderNumber WHERE id2 = 'CATERPILLAR S.A.R.L.'; + ``` + """ + + sanitized = sanitize_sql_generation_reasoning(reasoning) + + assert "SELECT" not in sanitized + assert "WHERE" not in sanitized + assert "assume" not in sanitized.lower() + assert "retrieved schema metadata" in sanitized + + +def test_schema_grounding_rejects_unretrieved_table_name(): + error = validate_sql_against_contexts( + "SELECT invoice_id, COUNT(comment_id) FROM comments GROUP BY invoice_id", + SCHEMA_CONTEXTS, + ) + + assert error is not None + assert "comments" in error + assert "valid_invoice_comments" in error + + +def test_schema_grounding_accepts_retrieved_table_name(): + error = validate_sql_against_contexts( + """ + SELECT invoice_id, COUNT(comment_id) + FROM valid_invoice_comments + GROUP BY invoice_id + """, + SCHEMA_CONTEXTS, + ) + + assert error is None + + +def test_supported_how_many_question_ignores_interrogative_fillers(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Invoice records by business unit."},"columns":[{"sql_column_name_use_exactly":"invoice_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Invoice identifier."},{"sql_column_name_use_exactly":"business_unit","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Business unit for invoice grouping."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE invoice_records ( + invoice_id VARCHAR, + business_unit VARCHAR + ); + """ + ] + + message = unsupported_schema_message( + "How many invoice records are there by business unit?", + contexts, + ) + + assert message is None + + +def test_unsupported_how_many_question_still_reports_missing_schema_terms(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Customer order records."},"columns":[{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."},{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + order_id VARCHAR, + customer_name VARCHAR + ); + """ + ] + + message = unsupported_schema_message( + "How many repair records are there by status?", + contexts, + ) + + assert message is not None + assert "repair" in message + assert "status" in message + assert "how" not in message + assert "there" not in message + + +def test_schema_grounding_rejects_invalid_qualified_column(): + error = validate_sql_against_contexts( + """ + SELECT c.invoice_number + FROM valid_invoice_comments c + """, + SCHEMA_CONTEXTS, + ) + + assert error is not None + assert "c.invoice_number" in error + + +def test_schema_identifier_normalization_quotes_special_identifiers(): + sql = normalize_sql_with_schema_identifiers( + "SELECT order_id FROM [valid-order-lines]", + SCHEMA_CONTEXTS, + ) + + assert 'FROM "valid-order-lines"' in sql + + +def test_schema_identifier_normalization_rewrites_verified_source_table_name(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"orders_model","sql_column_names_use_exactly":["customer_name"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"tblOrders"},"display_name":"New Orders"},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":{"source_column_name":"CustName","display_name":"CustName"}}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE orders_model ( + customer_name VARCHAR + ); + """ + ] + + sql = normalize_sql_with_schema_identifiers( + "SELECT customer_name FROM dbo.tblOrders", + contexts, + ) + + assert 'FROM "orders_model"' in sql + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + assert validate_sql_against_contexts(sql, contexts) is None + + +def test_schema_identifier_normalization_rewrites_verified_source_column_name(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"orders_model","sql_column_names_use_exactly":["customer_name"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"tblOrders"}},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":{"source_column_name":"CustName","display_name":"CustName"}}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE orders_model ( + customer_name VARCHAR + ); + """ + ] + + sql = normalize_sql_with_schema_identifiers( + """ + SELECT o.CustName + FROM dbo.tblOrders o + WHERE LOWER(o.CustName) = LOWER('lockheed martin') + """, + contexts, + ) + + assert 'FROM "orders_model" o' in sql + assert 'o."customer_name"' in sql + assert "CustName" not in sql + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + assert validate_sql_against_contexts(sql, contexts) is None + + +def test_schema_identifier_normalization_rewrites_verified_display_column_variant(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"orders_model","sql_column_names_use_exactly":["CustName"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"tblOrders"}},"columns":[{"sql_column_name_use_exactly":"CustName","data_type":"VARCHAR","display_name":"Customer name","semantic_context_not_sql_identifier":{"display_name":"Customer name"}}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE orders_model ( + CustName VARCHAR + ); + """ + ] + + sql = normalize_sql_with_schema_identifiers( + """ + SELECT o.CustomerName + FROM dbo.tblOrders o + WHERE LOWER(o.Customer_Name) = LOWER('lockheed martin') + """, + contexts, + ) + + assert 'o."CustName"' in sql + assert "CustomerName" not in sql + assert "Customer_Name" not in sql + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + + +def test_schema_identifier_normalization_keeps_ambiguous_source_table_invalid(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"order_archive","sql_column_names_use_exactly":["id"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"orders"}},"columns":[{"sql_column_name_use_exactly":"id","data_type":"VARCHAR"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_archive ( + id VARCHAR + ); + """, + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","sql_identifier_contract":{"sql_table_name_use_exactly":"order_current","sql_column_names_use_exactly":["id"]},"semantic_context_not_sql_identifiers":{"source_table_reference":{"schema":"dbo","table":"orders"}},"columns":[{"sql_column_name_use_exactly":"id","data_type":"VARCHAR"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_current ( + id VARCHAR + ); + """, + ] + + sql = normalize_sql_with_schema_identifiers("SELECT id FROM dbo.orders", contexts) + error = _SchemaCatalog.from_contexts(contexts).validate_sql(sql) + + assert "dbo.orders" in sql + assert error is not None + assert "dbo.orders" in error + + +def test_wren_sql_dialect_normalization_handles_top_and_joined_limit(): + assert ( + normalize_wren_sql_dialect("SELECT TOP 10 id1 FROM dbo_mbrTime") + == "SELECT id1 FROM dbo_mbrTime\nLIMIT 10" + ) + assert ( + normalize_wren_sql_dialect( + "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESCLIMIT 10" + ) + == "SELECT id1 FROM dbo_mbrTime ORDER BY metric DESC LIMIT 10" + ) + + +def test_semantic_coverage_rejects_unrepresented_query_terms(): + contexts = [ + """ + CREATE TABLE neutral_records ( + id1 INTEGER, + id2 INTEGER + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT id1, COUNT(*) AS record_count + FROM neutral_records + GROUP BY id1 + ORDER BY record_count DESC + LIMIT 10 + """, + "Show top 10 records by missing_dimension.", + contexts, + ) + + assert error is not None + assert "missing" in error or "dimension" in error + + +def test_unsupported_schema_message_reports_partial_coverage(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR + ); + """ + ] + + message = unsupported_schema_message( + "Show records by phase and unknown_segment.", + contexts, + ) + + assert message is not None + assert "unknown" in message or "segment" in message + + +def test_unsupported_schema_generation_result_has_no_invalid_sql(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR + ); + """ + ] + + result = unsupported_schema_generation_result( + "Show records by unknown_segment.", + contexts, + data_source="MSSQL", + ) + + assert result is not None + assert result["valid_generation_result"] == {} + invalid = result["invalid_generation_result"] + assert invalid["type"] == "NO_RELEVANT_SQL" + assert invalid["sql"] == "" + assert invalid["original_sql"] == "" + assert "unknown" in invalid["error"] or "segment" in invalid["error"] + + +def test_schema_grounding_failure_message_reports_terms_split_across_tables(): + contexts = [ + """ + CREATE TABLE user_activity ( + username VARCHAR + ); + """, + """ + CREATE TABLE record_counts ( + recordcnt INTEGER + ); + """, + ] + + message = schema_grounding_failure_message( + "Show top 5 username by recordcnt.", + contexts, + ) + + assert "active project" in message + assert "username" in message + assert "recordcnt" in message + assert "Generated SQL referenced" not in message + + +def test_schema_coverage_accepts_generic_word_form_variants(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = """ + SELECT + CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT) AS year, + CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) AS month, + COUNT(*) AS record_count + FROM work_update_log + GROUP BY + CAST(EXTRACT(YEAR FROM updated_at) AS BIGINT), + CAST(EXTRACT(MONTH FROM updated_at) AS BIGINT) + """ + + error = validate_sql_semantic_coverage( + sql, + "Show the number of work updates updated each month.", + contexts, + ) + + assert error is None + assert unsupported_schema_message( + "Show the number of work updates updated each month.", + contexts, + ) is None + + +def test_schema_fallback_uses_verified_monthly_update_timestamp(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of work updates updated each month.", + contexts, + ) + + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT)' in sql + assert "COUNT(*)" in sql + + +def test_schema_fallback_uses_unambiguous_implicit_text_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"product_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Product name. Use for product analysis."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + product_name VARCHAR, + order_id VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show orders from Lockheed Martine.", + contexts, + ) + + assert sql is not None + assert 'FROM "order_records"' in sql + assert 'LOWER("customer_name") LIKE \'%lockheed martine%\'' in sql + + +def test_schema_fallback_allows_punctuated_customer_value_after_for(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters.","sample_values":["LOCKHEED MARTIN CORPORATION"]},{"sql_column_name_use_exactly":"product_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Product name. Use for product analysis."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + product_name VARCHAR, + order_id VARCHAR + ); + """ + ] + query = "Show orders for CATERPILLAR S.A.R.L." + + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert validate_sql_semantic_coverage( + """ + SELECT customer_name, order_id + FROM order_records + WHERE LOWER(customer_name) LIKE '%caterpillar%' + """, + query, + contexts, + ) is None + assert sql is not None + assert 'FROM "order_records"' in sql + assert 'LOWER("customer_name") LIKE \'%caterpillar s.a.r.l%\'' in sql + + +def test_schema_fallback_keeps_customer_value_separate_from_grouping_phrase(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"product_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Product name. Use for product analysis."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + product_name VARCHAR, + order_id VARCHAR + ); + """ + ] + query = "Show orders from PACCAR PARTS DIVISION by product." + + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'LOWER("customer_name") LIKE \'%paccar parts division%\'' in sql + assert 'GROUP BY "product_name"' in sql + + +def test_schema_fallback_stops_grouping_phrase_before_from_table_name(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Business unit activity records."},"columns":[{"sql_column_name_use_exactly":"bunit","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Business unit grouping code."},{"sql_column_name_use_exactly":"record_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Record identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE business_unit_records ( + bunit VARCHAR, + record_id VARCHAR + ); + """ + ] + query = "Show row counts grouped by bunit from business_unit_records." + + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "business_unit_records"' in sql + assert 'GROUP BY "bunit"' in sql + assert validate_sql_semantic_coverage(sql, query, contexts) is None + + +def test_schema_fallback_groups_monthly_record_request_without_count_word(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"order_date","data_type":"DATE","semantic_context_not_sql_identifier":"Order placement date."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + order_date DATE, + order_id VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show monthly orders from JOHN DEERE COMMERCIAL PRODUCTS.", + contexts, + ) + + assert sql is not None + assert 'FROM "order_records"' in sql + assert 'CAST(EXTRACT(YEAR FROM "order_date") AS BIGINT)' in sql + assert 'CAST(EXTRACT(MONTH FROM "order_date") AS BIGINT)' in sql + assert 'COUNT(*) AS "record_count"' in sql + assert 'LOWER("customer_name") LIKE \'%john deere commercial products%\'' in sql + + +def test_schema_fallback_does_not_treat_customer_suffix_as_required_schema(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + order_id VARCHAR + ); + """, + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Refund timing records."},"columns":[{"sql_column_name_use_exactly":"day_s_from_refund","data_type":"INTEGER","semantic_context_not_sql_identifier":"Refund timing days."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE refund_records ( + day_s_from_refund INTEGER + ); + """, + ] + query = "Show orders for CATERPILLAR S.A.R.L." + + assert unsupported_schema_message(query, contexts) is None + + +def test_schema_fallback_prefers_customer_semantics_over_value_word_overlap(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Use this table for new order analysis."},"columns":[{"sql_column_name_use_exactly":"CustName","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use this column when the user asks for customer, customer name, account, or buyer.","display_name":"Customer name","source_column_name":"CustName"},{"sql_column_name_use_exactly":"Division","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Business division. Use for division-level reporting or grouping.","display_name":"Division","source_column_name":"Division"},{"sql_column_name_use_exactly":"OrdDate","data_type":"DATE","semantic_context_not_sql_identifier":"Order placement date."},{"sql_column_name_use_exactly":"OrdNo","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Sales order number."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE dbo_tblNewOrders ( + CustName VARCHAR, + Division VARCHAR, + OrdDate DATE, + OrdNo VARCHAR + ); + """ + ] + query = "Show recent orders from PACCAR PARTS DIVISION." + + sql = generate_simple_analytics_sql(query, contexts) + + assert sql is not None + assert 'FROM "dbo_tblNewOrders"' in sql + assert 'LOWER("CustName") LIKE \'%paccar parts division%\'' in sql + assert "Division) = 'paccar parts division'" not in sql + assert 'ORDER BY "OrdDate" DESC' in sql + + +def test_schema_fallback_prefers_direct_table_name_match_when_metadata_overlaps(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Use this table for order and customer questions."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE sales_history ( + customer_name VARCHAR, + order_id VARCHAR + ); + """, + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Use this table for order and customer questions."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE new_orders ( + customer_name VARCHAR, + order_id VARCHAR + ); + """, + ] + + sql = generate_simple_analytics_sql("Show orders from Acme.", contexts) + + assert sql is not None + assert 'FROM "new_orders"' in sql + + +def test_schema_coverage_uses_word_form_variants_per_table(): + contexts = [ + """ + CREATE TABLE report_failures ( + report_id VARCHAR, + failure_line VARCHAR + ); + """, + """ + CREATE TABLE failure_patterns ( + severity VARCHAR, + failure_type VARCHAR + ); + """, + ] + query = "Show failures by severity." + + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "failure_patterns"' in sql + assert 'GROUP BY "severity"' in sql + + +def test_schema_fallback_splits_compound_identifier_names_for_grouping_and_measure(): + contexts = [ + """ + CREATE TABLE payable_invoices ( + bunit VARCHAR, + suppliername VARCHAR, + grossamount DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show total gross amount by supplier name.", + contexts, + ) + + assert sql is not None + assert 'SELECT "suppliername", SUM("grossamount") AS "total_value"' in sql + assert 'GROUP BY "suppliername"' in sql + + +def test_semantic_validation_rejects_weaker_explicit_grouping_column(): + contexts = [ + """ + CREATE TABLE payable_invoices ( + bunit VARCHAR, + suppliername VARCHAR, + grossamount DECIMAL + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT bunit, SUM(grossamount) AS total_value + FROM payable_invoices + GROUP BY bunit + ORDER BY total_value DESC + """, + "Show total gross amount by supplier name.", + contexts, + ) + + assert error is not None + assert "weaker matching column" in error or "grouping dimension" in error + + +def test_semantic_validation_accepts_compound_status_grouping_column(): + contexts = [ + """ + CREATE TABLE task_rollups ( + taskstatus VARCHAR, + task_id VARCHAR + ); + """, + ] + + error = validate_sql_semantic_coverage( + """ + SELECT taskstatus, COUNT(*) AS record_count + FROM task_rollups + GROUP BY taskstatus + """, + "How many task records are there by task status?", + contexts, + ) + + assert error is None + + +def test_fallback_limit_accepts_spelled_out_top_number(): + assert _fallback_limit("Show the top five groups from that result.") == 5 + + +def test_followup_group_result_words_do_not_require_schema_columns(): + contexts = [ + """ + CREATE TABLE repair_logs ( + status VARCHAR, + repair_id VARCHAR + ); + """, + ] + + assert ( + unsupported_schema_message( + "Show the top five groups from the repair records by status.", + contexts, + ) + is None + ) + + +def test_unsupported_subject_still_blocks_cross_domain_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records for customer activity."},"columns":[{"sql_column_name_use_exactly":"customer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Customer name. Use for customer, account, or buyer filters."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + customer_name VARCHAR, + order_id VARCHAR + ); + """ + ] + query = "Show tickets from ACME CORP." + + message = unsupported_schema_message(query, contexts) + + assert message is not None + assert "ticket" in message.lower() + assert "acme" not in message.lower() + + +def test_schema_fallback_skips_ambiguous_implicit_text_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Order records."},"columns":[{"sql_column_name_use_exactly":"buyer_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Buyer name."},{"sql_column_name_use_exactly":"seller_name","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Seller name."},{"sql_column_name_use_exactly":"order_id","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Order identifier."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE order_records ( + buyer_name VARCHAR, + seller_name VARCHAR, + order_id VARCHAR + ); + """ + ] + + assert ( + generate_simple_analytics_sql("Show orders from Acme Industries.", contexts) + is None + ) + + +def test_schema_catalog_ignores_extract_from_column_clause(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + sql = """ + SELECT CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year", COUNT(*) AS "record_count" + FROM "work_update_log" + GROUP BY CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) + """ + + assert _SchemaCatalog.from_contexts(contexts).validate_sql(sql) is None + + +def test_post_processor_clears_sql_for_unsupported_schema(): + contexts = [ + """ + CREATE TABLE neutral_records ( + id1 INTEGER, + id2 INTEGER + ); + """ + ] + post_processor = SQLGenPostProcessor(engine=None) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT id1, COUNT(*) AS record_count + FROM neutral_records + GROUP BY id1 + ORDER BY record_count DESC + LIMIT 10 + """ + ], + contexts=contexts, + fallback_query="Show records by missing_dimension.", + data_source="MSSQL", + ) + ) + + assert result["valid_generation_result"] == {} + assert result["invalid_generation_result"]["type"] == "NO_RELEVANT_SQL" + assert result["invalid_generation_result"]["sql"] == "" + assert result["invalid_generation_result"]["original_sql"] == "" + + +def test_post_processor_prefers_fact_table_fallback_over_dimension_only_sql(): + contexts = [ + """ + CREATE TABLE invoice_records ( + invoice_id VARCHAR, + business_unit VARCHAR + ); + """, + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"invoice business unit grouping lookup"},"columns":[{"sql_column_name_use_exactly":"name","display_name":"Business Unit Group"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE business_unit_groups ( + name VARCHAR + ); + """, + ] + post_processor = SQLGenPostProcessor(engine=_AcceptingEngine()) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT name, COUNT(*) AS record_count + FROM business_unit_groups + GROUP BY name + ORDER BY record_count DESC + """ + ], + contexts=contexts, + fallback_query="How many invoice records are there by business unit?", + data_source="MSSQL", + use_dry_plan=True, + ) + ) + + sql = result["valid_generation_result"]["sql"] + assert result["invalid_generation_result"] == {} + assert "invoice_records" in sql + assert "business_unit" in sql + assert "business_unit_groups" not in sql + + +def test_post_processor_converts_invented_table_to_schema_message(): + contexts = [ + """ + CREATE TABLE user_activity ( + username VARCHAR + ); + """, + """ + CREATE TABLE record_counts ( + recordcnt INTEGER + ); + """, + ] + post_processor = SQLGenPostProcessor(engine=None) + + result = asyncio.run( + post_processor.run( + [ + """ + SELECT username + FROM records + ORDER BY recordcnt DESC + LIMIT 5 + """ + ], + contexts=contexts, + fallback_query="Show top 5 username by recordcnt.", + data_source="MSSQL", + ) + ) + + invalid = result["invalid_generation_result"] + assert result["valid_generation_result"] == {} + assert invalid["type"] == "NO_RELEVANT_SQL" + assert invalid["sql"] == "" + assert invalid["original_sql"] == "" + assert "active project" in invalid["error"] + assert "username" in invalid["error"] + assert "recordcnt" in invalid["error"] + assert "Generated SQL referenced" not in invalid["error"] + + +def test_schema_sample_value_filter_is_grounded_in_metadata(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done","In Progress"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE work_items ( + item_id VARCHAR, + State VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all work item records with In Progress.", + contexts, + ) + + assert sql is not None + assert 'FROM "work_items"' in sql + assert 'LOWER("State") = \'in progress\'' in sql + + +def test_user_values_are_allowed_for_single_verified_text_column(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + state_name VARCHAR(255), + updated_at TIMESTAMP + ); + """ + ] + + query = ( + "Show the distribution of work updates across completed and " + "in-progress state names." + ) + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'LOWER("state_name") IN (\'completed\', \'in-progress\')' in sql + assert 'GROUP BY "state_name"' in sql + + +def test_explicit_column_adjacent_value_after_column_is_filter_value(): + contexts = [ + """ + CREATE TABLE work_items ( + item_id VARCHAR, + status VARCHAR, + priority VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + query = "Show work items with status open." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_items"' in sql + assert 'LOWER("status") = \'open\'' in sql + + +def test_explicit_column_value_survives_multi_column_semantic_ambiguity(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Support ticket records."},"columns":[{"sql_column_name_use_exactly":"title","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket title."},{"sql_column_name_use_exactly":"description","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket description."},{"sql_column_name_use_exactly":"status","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket status."},{"sql_column_name_use_exactly":"priority","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket priority."},{"sql_column_name_use_exactly":"data","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket payload."},{"sql_column_name_use_exactly":"created_at","data_type":"TIMESTAMP","semantic_context_not_sql_identifier":"Ticket creation time."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE support_tickets ( + title VARCHAR, + description VARCHAR, + status VARCHAR, + priority VARCHAR, + data VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show tickets with status open." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "support_tickets"' in sql + assert 'LOWER("status") = \'open\'' in sql + + +def test_explicit_column_value_allows_value_token_seen_elsewhere_in_schema(): + contexts = [ + """ + CREATE TABLE production_batches ( + batch_id VARCHAR, + supplier VARCHAR, + inspection_status VARCHAR, + created_at TIMESTAMP + ); + """, + """ + CREATE TABLE debug_entries ( + entry_id VARCHAR, + failed_at TIMESTAMP + ); + """, + ] + + query = "Show batches with inspection status failed." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "production_batches"' in sql + assert 'LOWER("inspection_status") = \'failed\'' in sql + + +def test_explicit_column_adjacent_value_before_column_is_filter_value(): + contexts = [ + """ + CREATE TABLE support_tickets ( + ticket_id VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show high priority tickets." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "support_tickets"' in sql + assert 'LOWER("priority") = \'high\'' in sql + assert 'LOWER("title")' not in sql + assert 'LOWER("description")' not in sql + assert 'LOWER("status") = \'high priority\'' not in sql + assert 'LOWER("data")' not in sql + + +def test_explicit_preceding_value_survives_multi_column_semantic_ambiguity(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Support ticket records."},"columns":[{"sql_column_name_use_exactly":"title","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket title."},{"sql_column_name_use_exactly":"description","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket description."},{"sql_column_name_use_exactly":"status","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket status."},{"sql_column_name_use_exactly":"priority","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket priority."},{"sql_column_name_use_exactly":"data","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Ticket payload."},{"sql_column_name_use_exactly":"created_at","data_type":"TIMESTAMP","semantic_context_not_sql_identifier":"Ticket creation time."}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE support_tickets ( + title VARCHAR, + description VARCHAR, + status VARCHAR, + priority VARCHAR, + data VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show high priority tickets." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "support_tickets"' in sql + assert 'LOWER("priority") = \'high\'' in sql + + +def test_explicit_identifier_style_column_value_filter_is_grounded(): + contexts = [ + """ + CREATE TABLE repair_logs ( + board_model VARCHAR, + failure_code VARCHAR, + status VARCHAR, + priority VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show repairs for failure code BGA-001." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "repair_logs"' in sql + assert 'LOWER("failure_code") = \'bga-001\'' in sql + + +def test_open_filter_phrase_strips_explicit_column_name_from_value(): + contexts = [ + """ + CREATE TABLE production_batches ( + batch_id VARCHAR, + supplier VARCHAR, + board_model VARCHAR, + created_at TIMESTAMP + ); + """ + ] + + query = "Show batches from supplier Wurth Elektronik." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "production_batches"' in sql + assert 'LOWER("supplier") = \'wurth elektronik\'' in sql + assert "supplier wurth" not in sql.lower() + + +def test_column_value_label_is_not_treated_as_literal_filter_value(): + contexts = [ + """ + CREATE TABLE work_update_log ( + item_id VARCHAR, + state_name VARCHAR(255), + updated_at TIMESTAMP + ); + """ + ] + + query = "Show the distribution of work updates across state name values." + sql = generate_simple_analytics_sql(query, contexts) + + assert unsupported_schema_message(query, contexts) is None + assert sql is not None + assert 'FROM "work_update_log"' in sql + assert 'GROUP BY "state_name"' in sql + assert "WHERE" not in sql + + +def test_subject_noun_is_not_treated_as_literal_filter_value(): + contexts = [ + """ + CREATE TABLE purchase_order_records ( + purchase_order_id VARCHAR, + currency_code VARCHAR, + order_date DATE, + order_quantity DECIMAL, + record_type VARCHAR, + record_status VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "How many purchase orders are there by currency?", + contexts, + ) + + assert unsupported_schema_message( + "How many purchase orders are there by currency?", + contexts, + ) is None + assert sql is not None + assert 'FROM "purchase_order_records"' in sql + assert 'GROUP BY "currency_code"' in sql + assert "WHERE" not in sql + + +def test_grouping_dimension_is_not_treated_as_subject_column_literal(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"sales records with repair item status reporting metadata"},"columns":[{"sql_column_name_use_exactly":"RepairItem","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Repair item"}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE sales_records ( + RepairItem VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "How many repair records are there by repair status?", + contexts, + ) + + assert sql is None + assert validate_sql_semantic_coverage( + """ + SELECT "RepairItem", COUNT(*) AS "record_count" + FROM "sales_records" + WHERE LOWER("RepairItem") = 'status' + GROUP BY "RepairItem" + """, + "How many repair records are there by repair status?", + contexts, + ) + + +def test_subject_entity_is_not_grounded_by_sample_value_filter(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"Tariff liquidation records."},"columns":[{"sql_column_name_use_exactly":"LiquidationStatus","data_type":"VARCHAR","semantic_context_not_sql_identifier":"Liquidation status.","sample_values":["Repair","Complete"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE tariff_records ( + LiquidationStatus VARCHAR, + Entry_Date TIMESTAMP + ); + """ + ] + query = "How many repair records are there by repair status?" + + sql = generate_simple_analytics_sql(query, contexts) + message = unsupported_schema_message(query, contexts) + error = validate_sql_semantic_coverage( + """ + SELECT LiquidationStatus, COUNT(*) AS record_count + FROM tariff_records + WHERE LOWER(LiquidationStatus) = 'repair' + GROUP BY LiquidationStatus + """, + query, + contexts, + ) + + assert sql is None + assert message is not None + assert "repair" in message.lower() + assert error is not None + assert "repair" in error.lower() + + +def test_compact_accounting_identifiers_support_balance_and_currency_questions(): + balance_contexts = [ + """ + CREATE TABLE balance_records ( + bunit VARCHAR, + period VARCHAR, + endbalance FLOAT8 + ); + """ + ] + balance_sql = generate_simple_analytics_sql( + "Show total ending balance by period and business unit.", + balance_contexts, + ) + + assert balance_sql is not None + assert 'SUM("endbalance") AS "total_value"' in balance_sql + assert 'GROUP BY "bunit", "period"' in balance_sql or 'GROUP BY "period", "bunit"' in balance_sql + assert "WHERE" not in balance_sql + + exchange_contexts = [ + """ + CREATE TABLE exchange_rate_records ( + currencyfrom VARCHAR, + currencyto VARCHAR, + exchangerate FLOAT8 + ); + """ + ] + exchange_sql = generate_simple_analytics_sql( + "Show exchange rates by currency pair.", + exchange_contexts, + ) + + assert exchange_sql is not None + assert 'FROM "exchange_rate_records"' in exchange_sql + assert '"exchangerate"' in exchange_sql + assert '"currencyfrom"' in exchange_sql or '"currencyto"' in exchange_sql + + +def test_semantic_validation_rejects_subject_noun_literal_filter(): + contexts = [ + """ + CREATE TABLE purchase_order_records ( + purchase_order_id VARCHAR, + currency_code VARCHAR, + order_date DATE, + order_quantity DECIMAL, + record_type VARCHAR, + record_status VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT currency_code, COUNT(*) AS record_count + FROM purchase_order_records + WHERE LOWER(record_type) = 'purchase' + GROUP BY currency_code + """, + "How many purchase orders are there by currency?", + contexts, + ) + + assert error is not None + assert "not grounded as a filter value" in error + + +def test_unverified_filter_value_is_not_invented(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"work item records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Done"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE work_items ( + item_id VARCHAR, + State VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all work item records with Archived.", + contexts, + ) + message = unsupported_schema_message( + "Show all work item records with Archived.", + contexts, + ) + + assert sql is None + assert message is not None + assert "archived" in message.lower() + + +def test_grouped_count_uses_verified_dimension_only(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + phase VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql("Show records by phase.", contexts) + + assert sql is not None + assert 'SELECT "phase", COUNT(*) AS "record_count"' in sql + assert 'FROM "event_records"' in sql + assert 'GROUP BY "phase"' in sql + + +def test_average_uses_verified_numeric_measure_not_count(): + contexts = [ + """ + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR, + age_days DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql("Show average age by model.", contexts) + + assert sql is not None + assert 'SELECT "model_code", AVG("age_days") AS "average_value"' in sql + assert 'GROUP BY "model_code"' in sql + assert "COUNT(" not in sql + + +def test_average_without_verified_measure_is_unsupported(): + contexts = [ + """ + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show average age by model.", contexts) + message = unsupported_schema_message("Show average age by model.", contexts) + + assert sql is None + assert message is not None + assert "age" in message.lower() + + +def test_latest_uses_verified_temporal_column(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + event_time TIMESTAMP, + phase VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show latest event records.", contexts) + + assert sql is not None + assert 'FROM "event_records"' in sql + assert 'ORDER BY "event_time" DESC' in sql + + +def test_count_mentions_subject_without_implicit_grouping(): + contexts = [ + """ + CREATE TABLE invoice_records ( + invoice_number VARCHAR, + invoice_date TIMESTAMP, + invoice_type VARCHAR, + business_unit VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "How many invoice records are there?", + contexts, + ) + + assert sql is not None + assert sql.strip() == ( + 'SELECT COUNT(*) AS "record_count"\n' + 'FROM "invoice_records"' + ) + assert "GROUP BY" not in sql + + +def test_latest_by_temporal_column_does_not_become_grouped_count(): + contexts = [ + """ + CREATE TABLE account_reconciliation_records ( + account_number VARCHAR, + approval_date TIMESTAMP, + status VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the latest account reconciliation records by approval date.", + contexts, + ) + + assert sql is not None + assert 'FROM "account_reconciliation_records"' in sql + assert 'ORDER BY "approval_date" DESC' in sql + assert "COUNT(*)" not in sql + assert "GROUP BY" not in sql + + +def test_monthly_count_uses_requested_temporal_column_when_verified(): + contexts = [ + """ + CREATE TABLE event_records ( + event_id VARCHAR, + updated_at TIMESTAMP, + created_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show the number of event records updated each month.", + contexts, + ) + + assert sql is not None + assert 'CAST(EXTRACT(YEAR FROM "updated_at") AS BIGINT) AS "year"' in sql + assert 'CAST(EXTRACT(MONTH FROM "updated_at") AS BIGINT) AS "month"' in sql + assert 'COUNT(*) AS "record_count"' in sql + + +def test_order_by_uses_verified_column_and_sample_value(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open","Closed"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE case_records ( + case_id VARCHAR, + State VARCHAR, + updated_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show all case records with Open ordered by case ID.", + contexts, + ) + + assert sql is not None + assert 'LOWER("State") = \'open\'' in sql + assert 'ORDER BY "case_id" ASC' in sql + + +def test_top_grouped_count_is_schema_shape_based(): + contexts = [ + """ + CREATE TABLE occurrence_records ( + occurrence_id VARCHAR, + model_code VARCHAR, + reason_code VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show top 5 occurrence records by model.", + contexts, + ) + + assert sql is not None + assert 'SELECT "model_code", COUNT(*) AS "record_count"' in sql + assert 'ORDER BY "record_count" DESC' in sql + assert "LIMIT 5" in sql + + +def test_single_grouping_dimension_does_not_over_split_results(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show number of events by account.", contexts) + + assert sql is not None + assert 'GROUP BY "account_name"' in sql + assert "account_reference" not in sql + + +def test_missing_value_intent_uses_verified_plural_name_column(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + event_time TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show events with missing account names.", + contexts, + ) + + assert sql is not None + assert 'FROM "account_events"' in sql + assert '"account_name" IS NULL' in sql + + +def test_semantic_validation_rejects_weaker_null_check_column(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT account_name, account_reference + FROM account_events + WHERE account_reference IS NULL + """, + "Show events with missing account names.", + contexts, + ) + + assert error is not None + assert "weaker matching column" in error + + +def test_top_records_are_listed_without_implicit_grouped_aggregate(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show top 10 scored events from July.", contexts) + + assert sql is not None + assert 'FROM "scored_events"' in sql + assert "GROUP BY" not in sql + assert 'ORDER BY "score_value" DESC' in sql + assert "LIMIT 10" in sql + + +def test_semantic_validation_accepts_generated_month_date_range(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + query = "Show top 10 scored events from July." + sql = generate_simple_analytics_sql(query, contexts) + + assert sql is not None + assert validate_sql_semantic_coverage(sql, query, contexts) is None + + +def test_top_records_without_verified_rank_measure_uses_date_ordering(): + contexts = [ + """ + CREATE TABLE order_records ( + order_id VARCHAR, + order_date TIMESTAMP, + fx_currency DECIMAL, + customer_name VARCHAR + ); + """ + ] + + sql = generate_simple_analytics_sql("Show top 10 order records from July.", contexts) + + assert sql is not None + assert 'FROM "order_records"' in sql + assert 'ORDER BY "order_date" DESC' in sql + assert 'fx_currency' not in sql.split("ORDER BY", maxsplit=1)[1] + assert "LIMIT 10" in sql + + +def test_top_rows_by_numeric_column_orders_without_grouping(): + contexts = [ + """ + CREATE TABLE failure_patterns ( + pattern_id VARCHAR, + name VARCHAR, + severity VARCHAR, + occurrences INTEGER, + cost_impact DECIMAL + ); + """ + ] + + sql = generate_simple_analytics_sql( + "Show top 10 failure patterns by occurrences.", + contexts, + ) + + assert sql is not None + assert 'FROM "failure_patterns"' in sql + assert "GROUP BY" not in sql + assert 'ORDER BY "occurrences" DESC' in sql + assert "LIMIT 10" in sql + + +def test_sum_by_year_uses_verified_measure_and_temporal_column(): + contexts = [ + """ + CREATE TABLE transaction_records ( + transaction_id VARCHAR, + account_name VARCHAR, + amount_value DECIMAL, + posted_at TIMESTAMP + ); + """ + ] + + sql = generate_simple_analytics_sql("Show total amount by year.", contexts) + + assert sql is not None + assert 'CAST(EXTRACT(YEAR FROM "posted_at") AS BIGINT) AS "year"' in sql + assert 'SUM("amount_value") AS "total_value"' in sql + + +def test_semantic_coverage_rejects_count_for_average_intent(): + contexts = [ + """ + CREATE TABLE measurement_records ( + entity_id VARCHAR, + model_code VARCHAR, + age_days DECIMAL + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT model_code, COUNT(*) AS record_count + FROM measurement_records + GROUP BY model_code + """, + "Show average age by model.", + contexts, + ) + + assert error is not None + assert "average" in error.lower() + + +def test_literal_validation_rejects_values_outside_verified_samples(): + contexts = [ + """ + /* + WREN RETRIEVED SEMANTIC CONTEXT + {"object_type":"model","semantic_context_not_sql_identifiers":{"description":"case records"},"columns":[{"sql_column_name_use_exactly":"State","data_type":"VARCHAR","sample_values":["Open"]}]} + WREN SQL IDENTIFIER CONTRACT + */ + CREATE TABLE case_records ( + case_id VARCHAR, + State VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT case_id + FROM case_records + WHERE LOWER(State) = 'Closed' + """, + "Show case records with Open.", + contexts, + ) + + assert error is not None + assert "sample values" in error + + +def test_semantic_validation_rejects_multi_group_for_single_dimension(): + contexts = [ + """ + CREATE TABLE account_events ( + event_id VARCHAR, + account_name VARCHAR, + account_reference VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT account_name, account_reference, COUNT(*) AS record_count + FROM account_events + GROUP BY account_name, account_reference + """, + "Show number of events by account.", + contexts, + ) + + assert error is not None + assert "one grouping dimension" in error + + +def test_semantic_validation_rejects_top_record_grouped_aggregate(): + contexts = [ + """ + CREATE TABLE scored_events ( + event_id VARCHAR, + score_value DECIMAL, + event_date TIMESTAMP, + category_name VARCHAR + ); + """ + ] + + error = validate_sql_semantic_coverage( + """ + SELECT category_name, SUM(score_value) AS total_value + FROM scored_events + GROUP BY category_name + ORDER BY total_value DESC + LIMIT 10 + """, + "Show top 10 scored events from July.", + contexts, + ) + + assert error is not None + assert "grouped aggregate" in error diff --git a/wren-ai-service/tests/pytest/pipelines/generation/test_sql_tables_extraction.py b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_tables_extraction.py new file mode 100644 index 0000000000..7bf3be3b9b --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/generation/test_sql_tables_extraction.py @@ -0,0 +1,24 @@ +import pytest + +from src.pipelines.generation.sql_tables_extraction import post_process + + +@pytest.mark.asyncio +async def test_sql_tables_extraction_post_process_returns_tables(): + result = await post_process({"replies": ['{"tables": ["model_a", "model_b"]}']}) + + assert result == ["model_a", "model_b"] + + +@pytest.mark.asyncio +async def test_sql_tables_extraction_post_process_handles_missing_tables_key(): + result = await post_process({"replies": ['{"message": "no tables"}']}) + + assert result == [] + + +@pytest.mark.asyncio +async def test_sql_tables_extraction_post_process_handles_invalid_json(): + result = await post_process({"replies": ["not json"]}) + + assert result == [] diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py index 20bd8ac682..1cf5246ec6 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_db_schema.py @@ -1,3 +1,4 @@ +import ast from unittest.mock import AsyncMock import orjson @@ -5,7 +6,10 @@ from haystack import Document from pytest_mock import MockFixture -from src.pipelines.indexing.db_schema import DBSchema, DDLChunker +from src.pipelines.indexing.db_schema import ( + DBSchema, + DDLChunker, +) @pytest.mark.asyncio @@ -45,10 +49,51 @@ async def test_single_model(): "type": "TABLE", "comment": "\n/* {'alias': 'user', 'description': 'A table containing user information.'} */\n", "name": "user", + "properties": { + "description": "A table containing user information.", + "displayName": "user", + }, + "tableReference": None, + "refSql": None, } ) +@pytest.mark.asyncio +async def test_model_source_table_reference_is_indexed(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "orders_model", + "properties": { + "description": "Modeled order facts.", + "displayName": "New Orders", + }, + "tableReference": { + "catalog": "warehouse", + "schema": "dbo", + "table": "tblNewOrders", + }, + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=1) + + assert len(actual["documents"]) == 1 + content = ast.literal_eval(actual["documents"][0].content) + assert content["properties"]["displayName"] == "New Orders" + assert content["tableReference"] == { + "catalog": "warehouse", + "schema": "dbo", + "table": "tblNewOrders", + } + + @pytest.mark.asyncio async def test_multiple_models(): chunker = DDLChunker() @@ -84,6 +129,12 @@ async def test_multiple_models(): "type": "TABLE", "comment": "\n/* {'alias': 'user', 'description': 'A table containing user information.'} */\n", "name": "user", + "properties": { + "description": "A table containing user information.", + "displayName": "user", + }, + "tableReference": None, + "refSql": None, } ) @@ -94,6 +145,12 @@ async def test_multiple_models(): "type": "TABLE", "comment": "\n/* {'alias': 'order', 'description': 'A table containing order details.'} */\n", "name": "order", + "properties": { + "description": "A table containing order details.", + "displayName": "order", + }, + "tableReference": None, + "refSql": None, } ) @@ -134,6 +191,7 @@ async def test_column_is_primary_key(): "name": "id", "data_type": "INTEGER", "is_primary_key": True, + "properties": {}, } ], } @@ -175,10 +233,14 @@ async def test_column_with_properties(): "columns": [ { "type": "COLUMN", - "comment": '-- {"alias":"iid","description":"The unique identifier for a user."}\n ', + "comment": '-- {"alias":"iid","description":"The unique identifier for a user.","sourceColumnName":""}\n ', "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": { + "displayName": "iid", + "description": "The unique identifier for a user.", + }, } ], } @@ -191,6 +253,67 @@ async def test_column_with_properties(): "type": "TABLE", "comment": "\n/* {'alias': '', 'description': ''} */\n", "name": "user", + "properties": {}, + "tableReference": None, + "refSql": None, + } + ) + + +@pytest.mark.asyncio +async def test_null_metadata_properties_are_indexed_as_empty_text(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "user", + "properties": {"description": None, "displayName": None}, + "columns": [ + { + "name": "id", + "type": "INTEGER", + "properties": { + "displayName": None, + "description": None, + }, + } + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=1) + + assert len(actual["documents"]) == 2 + assert actual["documents"][0].content == str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "comment": '-- {"alias":null,"description":null,"sourceColumnName":""}\n ', + "name": "id", + "data_type": "INTEGER", + "is_primary_key": False, + "properties": { + "displayName": None, + "description": None, + }, + } + ], + } + ) + assert actual["documents"][1].content == str( + { + "type": "TABLE", + "comment": "\n/* {'alias': None, 'description': None} */\n", + "name": "user", + "properties": {"description": None, "displayName": None}, + "tableReference": None, + "refSql": None, } ) @@ -232,10 +355,16 @@ async def test_column_with_nested_columns(): "columns": [ { "type": "COLUMN", - "comment": '-- {"alias":"iid","description":"The unique identifier for a user.","nested_columns":{"nested.address":{"name":"address","type":"VARCHAR"},"nested.orders":{"name":"orders","type":"ARRAY"}}}\n ', + "comment": '-- {"alias":"iid","description":"The unique identifier for a user.","sourceColumnName":"","nested_columns":{"nested.address":{"name":"address","type":"VARCHAR"},"nested.orders":{"name":"orders","type":"ARRAY"}}}\n ', "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": { + "displayName": "iid", + "description": "The unique identifier for a user.", + "nested.address": {"name": "address", "type": "VARCHAR"}, + "nested.orders": {"name": "orders", "type": "ARRAY"}, + }, } ], } @@ -279,6 +408,7 @@ async def test_column_with_calculated_property(): "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": {}, } ], } @@ -343,22 +473,7 @@ async def test_column_with_relationship(): "name": "id", "data_type": "INTEGER", "is_primary_key": True, - } - ], - } - ) - - document_1: Document = actual["documents"][1] - assert document_1.meta == {"type": "TABLE_SCHEMA", "name": "user"} - assert document_1.content == str( - { - "type": "TABLE_COLUMNS", - "columns": [ - { - "type": "FOREIGN_KEY", - "comment": '-- {"condition": user.id = order.user_id, "joinType": ONE_TO_MANY}\n ', - "constraint": "FOREIGN KEY (id) REFERENCES order(user_id)", - "tables": ["user", "order"], + "properties": {}, } ], } @@ -414,6 +529,7 @@ async def test_column_batch_size(): "name": "id", "data_type": "INTEGER", "is_primary_key": False, + "properties": {}, }, { "type": "COLUMN", @@ -421,6 +537,7 @@ async def test_column_batch_size(): "name": "name", "data_type": "VARCHAR", "is_primary_key": False, + "properties": {}, }, ], } @@ -438,12 +555,73 @@ async def test_column_batch_size(): "name": "age", "data_type": "INTEGER", "is_primary_key": False, + "properties": {}, } ], } ) +@pytest.mark.asyncio +async def test_long_model_description_keeps_legacy_payload(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "user", + "properties": { + "displayName": "user", + "description": "x" * 5000, + }, + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=1) + + assert len(actual["documents"]) == 1 + document: Document = actual["documents"][0] + assert "x" * 5000 in document.content + assert "..." not in document.content + + +@pytest.mark.asyncio +async def test_table_columns_keep_legacy_comment_payload(): + chunker = DDLChunker() + mdl = { + "models": [ + { + "name": "user", + "columns": [ + { + "name": f"column_{index}", + "type": "VARCHAR", + "properties": { + "displayName": f"column_{index}", + "description": "x" * 6000, + }, + } + for index in range(2) + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = await chunker.run(mdl, column_batch_size=50) + table_column_documents = [ + document for document in actual["documents"] if "TABLE_COLUMNS" in document.content + ] + + assert len(table_column_documents) == 1 + assert "x" * 6000 in table_column_documents[0].content + + @pytest.mark.asyncio async def test_view(): chunker = DDLChunker() diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py index 3436a64c59..0ca5598ad9 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_sql_pairs.py @@ -3,9 +3,71 @@ from src.config import settings from src.core.provider import DocumentStoreProvider from src.pipelines.indexing import SqlPairs +from src.pipelines.indexing.sql_pairs import SqlPairsCleaner, embedding, sql_pairs, write from src.providers import generate_components +class _RecordingStore: + def __init__(self): + self.filters = [] + + async def delete_documents(self, filter): + self.filters.append(filter) + + +@pytest.mark.asyncio +async def test_sql_pairs_cleaner_delete_all_uses_project_scope(): + store = _RecordingStore() + cleaner = SqlPairsCleaner(store) + + await cleaner.run(sql_pair_ids=[], project_id="project-id", delete_all=True) + + assert store.filters == [ + { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-id"}, + ], + } + ] + + +def test_sql_pairs_can_skip_default_pairs(): + pairs = sql_pairs( + boilerplates={"default"}, + external_pairs={}, + include_default_pairs=False, + ) + + assert pairs == [] + + +@pytest.mark.asyncio +async def test_empty_sql_pairs_skip_embedding_and_write(): + class Embedder: + called = False + + async def run(self, documents): + self.called = True + return {"documents": documents} + + class Writer: + called = False + + async def run(self, documents): + self.called = True + + embedder = Embedder() + writer = Writer() + + result = await embedding({"documents": []}, embedder) + await write(result, writer) + + assert result == {"documents": []} + assert embedder.called is False + assert writer.called is False + + @pytest.mark.asyncio async def test_sql_pairs_indexing_saving_to_document_store(): pipe_components = generate_components(settings.components) @@ -94,3 +156,6 @@ async def test_sql_pairs_deletion(): await pipe.clean(sql_pairs=[], project_id="fake-id") assert await store.count_documents() == 2 + + await pipe.clean(sql_pairs=[], project_id="fake-id", delete_all=True) + assert await store.count_documents() == 0 diff --git a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py index 7def966fdb..fbdc1849e2 100644 --- a/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py +++ b/wren-ai-service/tests/pytest/pipelines/indexing/test_table_description.py @@ -24,8 +24,8 @@ def test_single_table_description(): mdl = { "models": [ { - "name": "user", - "properties": {"description": "A table containing user information."}, + "name": "entity", + "properties": {"description": "A generic entity resource."}, } ], "views": [], @@ -37,11 +37,11 @@ def test_single_table_description(): assert len(actual["documents"]) == 1 document: Document = actual["documents"][0] - assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "user"} + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "entity"} assert document.content == str( { - "name": "user", - "description": "A table containing user information.", + "name": "entity", + "description": "A generic entity resource.", "columns": "", } ) @@ -52,12 +52,12 @@ def test_multiple_table_descriptions(): mdl = { "models": [ { - "name": "user", - "properties": {"description": "A table containing user information."}, + "name": "entity", + "properties": {"description": "A generic entity resource."}, }, { - "name": "order", - "properties": {"description": "A table containing order details."}, + "name": "activity", + "properties": {"description": "A generic activity resource."}, }, ], "views": [], @@ -71,22 +71,22 @@ def test_multiple_table_descriptions(): document_1: Document = actual["documents"][0] assert document_1.meta == { "type": "TABLE_DESCRIPTION", - "name": "user", + "name": "entity", } assert document_1.content == str( { - "name": "user", - "description": "A table containing user information.", + "name": "entity", + "description": "A generic entity resource.", "columns": "", } ) document_2: Document = actual["documents"][1] - assert document_2.meta == {"type": "TABLE_DESCRIPTION", "name": "order"} + assert document_2.meta == {"type": "TABLE_DESCRIPTION", "name": "activity"} assert document_2.content == str( { - "name": "order", - "description": "A table containing order details.", + "name": "activity", + "description": "A generic activity resource.", "columns": "", } ) @@ -112,7 +112,7 @@ def test_table_description_missing_name(): def test_table_description_missing_description(): chunker = TableDescriptionChunker() mdl = { - "models": [{"name": "user"}], + "models": [{"name": "entity"}], "views": [], "relationships": [], "metrics": [], @@ -122,8 +122,187 @@ def test_table_description_missing_description(): assert len(actual["documents"]) == 1 document: Document = actual["documents"][0] - assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "user"} - assert document.content == str({"name": "user", "description": "", "columns": ""}) + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "entity"} + assert document.content == str( + {"name": "entity", "description": "", "columns": ""} + ) + + +def test_table_description_null_description(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + { + "name": "entity", + "properties": {"description": None, "displayName": None}, + "columns": [{"name": "id"}], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + assert len(actual["documents"]) == 1 + + document: Document = actual["documents"][0] + assert document.meta == {"type": "TABLE_DESCRIPTION", "name": "entity"} + assert document.content == str( + { + "name": "entity", + "description": None, + "columns": "id", + } + ) + + +def test_table_description_uses_legacy_column_list_without_semantic_context(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + { + "name": "resource", + "properties": { + "description": "A generic described resource.", + "displayName": "Resource", + }, + "columns": [ + { + "name": "AttributeOne", + "type": "varchar", + "properties": { + "description": "Generic generated attribute description." + }, + }, + { + "name": "MeasureOne", + "type": "float", + "properties": { + "description": "Generic generated measure description." + }, + }, + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + + document: Document = actual["documents"][0] + assert document.content == str( + { + "name": "resource", + "description": "A generic described resource.", + "columns": "AttributeOne, MeasureOne", + } + ) + + +def test_table_description_excludes_physical_source_metadata(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + { + "name": "orders_model", + "tableReference": { + "catalog": "warehouse", + "schema": "raw", + "table": "refunds", + }, + "properties": { + "description": "Business-facing order and refund activity." + }, + "columns": [ + { + "name": "refund_date", + "type": "date", + "expression": "raw_refund_timestamp", + "properties": { + "description": "Date when a refund was issued." + }, + } + ], + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + + document: Document = actual["documents"][0] + assert "orders_model" in document.content + assert "refund_date" in document.content + assert "Date when a refund was issued." not in document.content + assert "warehouse" not in document.content + assert "raw" not in document.content + assert "refunds" not in document.content + assert "raw_refund_timestamp" not in document.content + assert "'source':" not in document.content + assert "'expression':" not in document.content + + +def test_table_description_does_not_include_relationship_context(): + chunker = TableDescriptionChunker() + mdl = { + "models": [ + {"name": "source", "columns": [{"name": "source_id"}]}, + {"name": "target", "columns": [{"name": "source_id"}]}, + ], + "views": [], + "relationships": [ + { + "name": "source_to_target", + "models": ["source", "target"], + "joinType": "ONE_TO_MANY", + "condition": "source.source_id = target.source_id", + } + ], + "metrics": [], + } + + actual = chunker.run(mdl) + + assert len(actual["documents"]) == 2 + for document in actual["documents"]: + assert document.meta["type"] == "TABLE_DESCRIPTION" + assert ( + "source_to_target ONE_TO_MANY source.source_id = target.source_id" + not in document.content + ) + + +def test_table_description_keeps_complete_column_lists(): + chunker = TableDescriptionChunker() + columns = [{"name": f"column_{index}"} for index in range(205)] + mdl = { + "models": [ + { + "name": "entity", + "columns": columns, + } + ], + "views": [], + "relationships": [], + "metrics": [], + } + + actual = chunker.run(mdl) + + assert len(actual["documents"]) == 1 + document: Document = actual["documents"][0] + assert document.content == str( + { + "name": "entity", + "description": "", + "columns": ", ".join(column["name"] for column in columns), + } + ) @pytest.mark.asyncio @@ -131,12 +310,12 @@ async def test_pipeline_run(mocker: MockFixture): test_mdl = { "models": [ { - "name": "user", - "properties": {"description": "A table containing user information."}, + "name": "entity", + "properties": {"description": "A generic entity resource."}, }, { - "name": "order", - "properties": {"description": "A table containing order details."}, + "name": "activity", + "properties": {"description": "A generic activity resource."}, }, ], "views": [], diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/sql_function.py b/wren-ai-service/tests/pytest/pipelines/retrieval/sql_function.py index 433e04f680..9fc32ae3eb 100644 --- a/wren-ai-service/tests/pytest/pipelines/retrieval/sql_function.py +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/sql_function.py @@ -1,18 +1,28 @@ from unittest.mock import AsyncMock, MagicMock import pytest +from haystack import Document from src.pipelines.retrieval.sql_functions import SqlFunction, SqlFunctions MOCK_FUNCTION_DEFINITION = { "name": "test_func", - "param_types": "int,text", - "return_type": "boolean", + "function_type": "scalar", + "description": "Returns a test value.", } MOCK_FUNCTION_LIST = [ - {"name": "func1", "param_types": "int", "return_type": "text"}, - {"name": "func2", "param_types": "text,text", "return_type": "boolean"}, + { + "name": "func1", + "function_type": "scalar", + "description": "Returns the first value.", + }, + { + "name": "func2", + "function_type": "aggregate", + "description": "Aggregates values.", + }, + {"name": "func_without_description", "function_type": "scalar"}, ] @@ -25,19 +35,35 @@ def mock_engine(): @pytest.fixture def sql_functions_pipeline(mock_engine): - return SqlFunctions(engine=mock_engine) + retriever = MagicMock() + retriever.run = AsyncMock( + return_value={"documents": [Document(content="", meta={"data_source": "postgres"})]} + ) + document_store_provider = MagicMock() + document_store_provider.get_store.return_value = MagicMock() + document_store_provider.get_retriever.return_value = retriever + pipeline = SqlFunctions( + engine=mock_engine, + document_store_provider=document_store_provider, + ) + pipeline._test_retriever = retriever + return pipeline def test_sql_function_init(): func = SqlFunction(MOCK_FUNCTION_DEFINITION) - expected = "test_func($0: int, $1: text) -> boolean" + expected = "type: scalar, name: TEST_FUNC, description: Returns a test value." assert str(func) == expected assert repr(func) == expected -def test_sql_function_empty_params(): - func = SqlFunction({"name": "test_func", "return_type": "text"}) - assert str(func) == "test_func(any) -> text" +def test_sql_function_empty_requires_legacy_fields(): + assert SqlFunction.empty({"name": "test_func", "function_type": "scalar"}) + assert SqlFunction.empty({"name": "test_func", "description": "Returns a value."}) + assert SqlFunction.empty( + {"function_type": "scalar", "description": "Returns a value."} + ) + assert not SqlFunction.empty(MOCK_FUNCTION_DEFINITION) @pytest.mark.asyncio @@ -45,8 +71,12 @@ async def test_sql_functions_pipeline_run(sql_functions_pipeline): result = await sql_functions_pipeline.run("postgres") assert len(result) == 2 - assert str(result[0]) == "func1($0: int) -> text" - assert str(result[1]) == "func2($0: text, $1: text) -> boolean" + assert str(result[0]) == ( + "type: scalar, name: FUNC1, description: Returns the first value." + ) + assert str(result[1]) == ( + "type: aggregate, name: FUNC2, description: Aggregates values." + ) cached_result = await sql_functions_pipeline.run("postgres") assert result == cached_result @@ -56,6 +86,11 @@ async def test_sql_functions_pipeline_run(sql_functions_pipeline): @pytest.mark.asyncio async def test_sql_functions_pipeline_different_datasource(sql_functions_pipeline): + sql_functions_pipeline._test_retriever.run.side_effect = [ + {"documents": [Document(content="", meta={"data_source": "postgres"})]}, + {"documents": [Document(content="", meta={"data_source": "mysql"})]}, + ] + await sql_functions_pipeline.run("postgres") await sql_functions_pipeline.run("mysql") @@ -64,6 +99,11 @@ async def test_sql_functions_pipeline_different_datasource(sql_functions_pipelin @pytest.mark.asyncio async def test_sql_functions_pipeline_case_insensitive(sql_functions_pipeline): + sql_functions_pipeline._test_retriever.run.side_effect = [ + {"documents": [Document(content="", meta={"data_source": "POSTGRES"})]}, + {"documents": [Document(content="", meta={"data_source": "postgres"})]}, + ] + result1 = await sql_functions_pipeline.run("POSTGRES") result2 = await sql_functions_pipeline.run("postgres") @@ -71,26 +111,17 @@ async def test_sql_functions_pipeline_case_insensitive(sql_functions_pipeline): assert result1 == result2 -def test_sql_function_param_type_none(): - func = SqlFunction( - {"name": "test_func", "param_types": None, "return_type": "text"} - ) - assert str(func) == "test_func(any) -> text" - - -def test_sql_function_return_type_none(): - func = SqlFunction( - {"name": "test_func", "param_types": "int,text", "return_type": None} - ) - assert str(func) == "test_func($0: int, $1: text) -> any" - - -def test_sql_function_return_type_same_as_args(): +def test_sql_function_ignores_signature_metadata(): func = SqlFunction( { - "name": "test_func", - "param_types": "int,text", - "return_type": "same as arg types", + "name": "dateadd", + "function_type": "scalar", + "param_types": ["varchar", "int", "datetime"], + "return_type": "Datetime", + "description": "Adds a signed number of dateparts to a date.", } ) - assert str(func) == "test_func($0: int, $1: text) -> ['int', 'text']" + assert ( + str(func) + == "type: scalar, name: DATEADD, description: Adds a signed number of dateparts to a date." + ) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py new file mode 100644 index 0000000000..4b2a144a41 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_db_schema_retrieval.py @@ -0,0 +1,1557 @@ +import pytest +import tiktoken +from haystack import Document +from haystack.components.builders.prompt_builder import PromptBuilder + +from src.pipelines.common import build_table_ddl +from src.pipelines.retrieval.db_schema_retrieval import ( + _augment_retrieval_query, + _build_table_retrieval_context, + _build_view_ddl, + _lexical_columns_and_tables_needed, + _limit_retrieval_results_for_generation, + _parse_column_selection_response, + _rank_table_names_by_query, + check_using_db_schemas_without_pruning, + construct_db_schemas, + construct_retrieval_results, + dbschema_retrieval, + embedding, + table_columns_selection_system_prompt, + table_columns_selection_user_prompt_template, + table_retrieval, +) +from src.pipelines.retrieval.db_schema_retrieval import ( + prompt as build_column_selection_prompt, +) + + +@pytest.mark.asyncio +async def test_embedding_uses_current_query_without_history_text(): + class Embedder: + def __init__(self): + self.query = None + + async def run(self, query): + self.query = query + return {"embedding": [1.0]} + + embedder = Embedder() + + result = await embedding( + query="current request", + embedder=embedder, + histories=[{"question": "previous request"}], + ) + + assert result == {"embedding": [1.0]} + assert embedder.query == "current request" + + +def test_column_pruning_prompt_uses_current_query_without_history_text(): + result = build_column_selection_prompt( + query="current request", + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + prompt_builder=PromptBuilder( + template=table_columns_selection_user_prompt_template + ), + check_using_db_schemas_without_pruning={"db_schemas": []}, + histories=[{"question": "previous request"}], + ) + + assert "current request" in result["prompt"] + assert "previous request" not in result["prompt"] + + +def test_table_selection_prompt_prefers_best_schema_supported_dataset_set(): + assert "same business concept is represented by multiple modeled datasets" in ( + table_columns_selection_system_prompt + ) + assert "best support the current question" in ( + table_columns_selection_system_prompt + ) + + +def test_view_schema_context_uses_declared_view_columns_when_available(): + result = _build_view_ddl( + { + "type": "VIEW", + "comment": "Semantic description.", + "name": "retrieved_view", + "statement": "NON_EXECUTABLE_DEFINITION_TOKEN", + "columns": [ + { + "name": "visible_attribute", + "data_type": "VARCHAR", + "comment": "Semantic field.", + } + ], + } + ) + + assert "CREATE TABLE retrieved_view" in result + assert "visible_attribute VARCHAR" in result + assert "sql_column_names_use_exactly" in result + assert "NON_EXECUTABLE_DEFINITION_TOKEN" not in result + + +def test_view_schema_context_uses_deployed_view_statement_without_declared_columns(): + result = _build_view_ddl( + { + "type": "VIEW", + "comment": "Semantic description.", + "name": "retrieved_view", + "statement": "SELECT modeled_column FROM deployed_model", + } + ) + + assert "CREATE TABLE retrieved_view" in result + assert "modeled_column VARCHAR" in result + assert "sql_table_name_use_exactly: retrieved_view" in result + + +def test_table_schema_context_includes_source_identifier_metadata(): + result, _, _ = _build_table_retrieval_context( + { + "type": "TABLE", + "comment": "", + "name": "orders_model", + "properties": {"displayName": "New Orders"}, + "tableReference": { + "catalog": "warehouse", + "schema": "dbo", + "table": "tblOrders", + }, + "columns": [ + { + "type": "COLUMN", + "name": "customer_name", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + "properties": { + "displayName": "CustName", + "sourceColumnName": "CustName", + }, + } + ], + "primaryKey": "", + } + ) + + assert '"source_table_name":"dbo.tblOrders"' in result + assert '"source_table_reference":{"catalog":"warehouse","schema":"dbo","table":"tblOrders"}' in result + assert '"source_column_name":"CustName"' in result + + +def test_construct_db_schemas_keeps_deployed_views_for_column_pruning(): + result = construct_db_schemas( + [ + Document( + content=str( + { + "type": "VIEW", + "comment": "", + "name": "retrieved_view", + "statement": "SELECT modeled_column FROM deployed_model", + } + ), + meta={"type": "TABLE_SCHEMA", "name": "retrieved_view"}, + ) + ] + ) + + assert result == [ + { + "type": "VIEW", + "comment": "", + "name": "retrieved_view", + "statement": "SELECT modeled_column FROM deployed_model", + } + ] + + +@pytest.mark.asyncio +async def test_table_retrieval_fetches_explicit_table_descriptions(): + class Retriever: + def __init__(self): + self.filters = None + + async def run(self, query_embedding, filters): + self.filters = filters + return {"documents": []} + + retriever = Retriever() + + await table_retrieval( + embedding={}, + project_id="project-1", + tables=["orders"], + table_retriever=retriever, + ) + + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_DESCRIPTION"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + {"field": "name", "operator": "in", "value": ["orders"]}, + ], + } + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_loads_selected_active_project_schema(): + class Retriever: + def __init__(self): + self.filters = None + + async def run(self, query_embedding, filters): + self.filters = filters + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": "orders", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "orders"}, + ), + Document( + content=str( + { + "type": "TABLE", + "name": "customers", + "columns": [], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "customers"}, + ), + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document( + content=str({"name": "orders"}), + meta={"type": "TABLE_DESCRIPTION", "name": "orders"}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + embedding={}, + ) + + assert [document.meta["name"] for document in documents] == ["orders", "customers"] + assert retriever.filters == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": "orders"}, + ], + }, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_expands_declared_relationships(): + selected_model = "model_a" + related_model = "model_b" + downstream_model = "model_c" + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + names = [ + condition["value"] + for condition in filters["conditions"][1]["conditions"] + ] + self.calls.append(names) + + if names == [selected_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": selected_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": selected_model}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "FOREIGN_KEY", + "tables": [ + selected_model, + related_model, + ], + "column": "model_b_id", + "referenced_table": related_model, + "referenced_column": "id", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": selected_model}, + ), + ] + } + + if names == [related_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": related_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": related_model}, + ), + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "FOREIGN_KEY", + "tables": [ + related_model, + downstream_model, + ], + "column": "model_c_id", + "referenced_table": downstream_model, + "referenced_column": "id", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": related_model}, + ), + ] + } + + if names == [downstream_model]: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": downstream_model, + } + ), + meta={"type": "TABLE_SCHEMA", "name": downstream_model}, + ) + ] + } + + return {"documents": []} + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document( + content=str({"name": selected_model}), + meta={"type": "TABLE_DESCRIPTION", "name": selected_model}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + embedding={}, + ) + + assert retriever.calls == [[selected_model], [related_model]] + assert [document.meta["name"] for document in documents] == [ + selected_model, + selected_model, + related_model, + related_model, + ] + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_uses_semantic_schema_hits_when_table_retrieval_misses(): + semantic_model = "semantic_dataset" + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + if query_embedding: + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE_COLUMNS", + "columns": [ + { + "type": "COLUMN", + "name": "semantic_measure", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": semantic_model}, + ) + ] + } + + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": semantic_model, + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": semantic_model}, + ) + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.25]}, + ) + + assert retriever.calls[0] == { + "query_embedding": [0.25], + "filters": { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + }, + } + assert retriever.calls[1]["query_embedding"] == [] + assert retriever.calls[1]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": semantic_model}, + ], + }, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert [document.meta["name"] for document in documents] == [ + semantic_model, + semantic_model, + ] + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_lexically_recovers_active_project_schema_when_vector_misses(): + recovered_model = "customer_orders" + unrelated_model = "service_tickets" + + def table_document(name, comment="", properties=None, table_reference=None): + return Document( + content=str( + { + "type": "TABLE", + "name": name, + "comment": comment, + "columns": [], + "properties": properties or {}, + "tableReference": table_reference, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + + def columns_document(name, columns): + return Document( + content=str({"type": "TABLE_COLUMNS", "columns": columns}), + meta={"type": "TABLE_SCHEMA", "name": name}, + ) + + recovered_documents = [ + table_document( + recovered_model, + comment="Orders captured from customer purchase activity.", + properties={"displayName": "Customer Orders"}, + table_reference={"schema": "sales", "table": "tblCustomerOrders"}, + ), + columns_document( + recovered_model, + [ + { + "type": "COLUMN", + "name": "customer_name", + "data_type": "VARCHAR", + "comment": "Customer name on the order.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "order_total", + "data_type": "DOUBLE", + "comment": "Order revenue amount.", + "is_primary_key": False, + }, + ], + ), + ] + all_schema_documents = recovered_documents + [ + table_document(unrelated_model, comment="Support case tracking."), + columns_document( + unrelated_model, + [ + { + "type": "COLUMN", + "name": "ticket_status", + "data_type": "VARCHAR", + "comment": "Support ticket status.", + "is_primary_key": False, + } + ], + ), + ] + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + if query_embedding: + return {"documents": []} + + is_exact_fetch = any( + isinstance(condition, dict) + and condition.get("operator") == "OR" + and condition.get("conditions") + for condition in filters["conditions"] + ) + if not is_exact_fetch: + return {"documents": all_schema_documents} + + return {"documents": recovered_documents} + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={"documents": []}, + project_id="project-1", + dbschema_retriever=retriever, + query="show orders by customer", + embedding={"embedding": [0.25]}, + ) + + assert [call["query_embedding"] for call in retriever.calls] == [ + [0.25], + [], + [], + ] + assert retriever.calls[1]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert [document.meta["name"] for document in documents] == [ + recovered_model, + recovered_model, + ] + + +@pytest.mark.asyncio +async def test_dbschema_retrieval_prefers_table_description_hits_over_schema_chunk_hits(): + described_model = "described_dataset" + + class Retriever: + def __init__(self): + self.calls = [] + + async def run(self, query_embedding, filters): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + + return { + "documents": [ + Document( + content=str( + { + "type": "TABLE", + "name": described_model, + "comment": "", + "columns": [], + "properties": {}, + "primaryKey": "", + } + ), + meta={"type": "TABLE_SCHEMA", "name": described_model}, + ) + ] + } + + retriever = Retriever() + + documents = await dbschema_retrieval( + table_retrieval={ + "documents": [ + Document( + content=str({"name": described_model}), + meta={"type": "TABLE_DESCRIPTION", "name": described_model}, + ) + ] + }, + project_id="project-1", + dbschema_retriever=retriever, + embedding={"embedding": [0.25]}, + ) + + assert [call["query_embedding"] for call in retriever.calls] == [[0.25], []] + assert retriever.calls[1]["filters"] == { + "operator": "AND", + "conditions": [ + {"field": "type", "operator": "==", "value": "TABLE_SCHEMA"}, + { + "operator": "OR", + "conditions": [ + {"field": "name", "operator": "==", "value": described_model}, + ], + }, + {"field": "project_id", "operator": "==", "value": "project-1"}, + ], + } + assert [document.meta["name"] for document in documents] == [described_model] + + +def test_check_using_db_schemas_without_pruning_triggers_legacy_column_pruning(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "orders", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=True, + context_window_size=1000, + ) + + assert result["db_schemas"] == [] + assert result["tokens"] > 0 + + +def test_construct_retrieval_results_preserves_retrieved_metric_when_pruning(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + }, + { + "table_name": "semantic_metric", + "table_selection_reason": "Selected metric for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed metric."], + "columns": ["metric_value"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "name": "semantic_metric", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_metric"}, + ) + ], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset", + "semantic_metric", + ] + assert result["has_metric"] is True + + +def test_construct_retrieval_results_skips_retrieved_schema_without_name(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + }, + { + "table_name": "semantic_metric", + "table_selection_reason": "Selected metric for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed metric."], + "columns": ["metric_value"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA"}, + ), + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "name": "semantic_metric", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_metric"}, + ), + ], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset", + "semantic_metric", + ] + assert result["has_metric"] is True + + +def test_construct_retrieval_results_excludes_unselected_metric_when_pruning(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["stored_attribute"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[ + Document( + content=str( + { + "type": "METRIC", + "comment": "", + "name": "semantic_metric", + "columns": [ + { + "type": "COLUMN", + "name": "metric_value", + "data_type": "DOUBLE", + "comment": "", + } + ], + } + ), + meta={"type": "TABLE_SCHEMA", "name": "semantic_metric"}, + ) + ], + ) + + assert [item["table_name"] for item in result["retrieval_results"]] == [ + "modeled_dataset" + ] + assert result["has_metric"] is False + + +def test_construct_retrieval_results_keeps_schema_when_pruner_returns_unknown_columns(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed field."], + "columns": ["semantic_label"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_dimension", + "data_type": "VARCHAR", + "comment": "Semantic dimension label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "stored_measure", + "data_type": "DOUBLE", + "comment": "Semantic measure label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + table_ddl = result["retrieval_results"][0]["table_ddl"] + + assert "semantic_label" not in table_ddl + assert "stored_dimension VARCHAR" in table_ddl + assert "stored_measure DOUBLE" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_dimension\n- stored_measure" in ( + table_ddl + ) + + +def test_construct_retrieval_results_falls_back_when_pruner_omits_results(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={"replies": ['{"message":"not structured"}']}, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + assert result["retrieval_results"] == [] + assert result["has_metric"] is False + + +def test_construct_retrieval_results_keeps_schema_when_pruner_mixes_known_and_unknown_columns(): + result = construct_retrieval_results( + check_using_db_schemas_without_pruning={}, + filter_columns_in_tables={ + "replies": [ + """ + { + "results": [ + { + "table_name": "modeled_dataset", + "table_selection_reason": "Selected for the current request.", + "table_contents": { + "chain_of_thought_reasoning": ["Needed fields."], + "columns": ["stored_measure", "semantic_label"] + } + } + ] + } + """ + ] + }, + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "stored_dimension", + "data_type": "VARCHAR", + "comment": "Semantic dimension label.", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "stored_measure", + "data_type": "DOUBLE", + "comment": "Semantic measure label.", + "is_primary_key": False, + }, + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + ) + + table_ddl = result["retrieval_results"][0]["table_ddl"] + + assert "semantic_label" not in table_ddl + assert "stored_dimension VARCHAR" not in table_ddl + assert "stored_measure DOUBLE" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_measure" in table_ddl + + +def test_check_using_db_schemas_without_pruning_keeps_context_when_within_window(): + class Encoding: + def encode(self, value): + return value.split() + + def table_schema(name): + return { + "type": "TABLE", + "name": name, + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + table_schema("activity"), + table_schema("account"), + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + assert [schema["table_name"] for schema in result["db_schemas"]] == [ + "activity", + "account", + ] + assert all( + "WREN RETRIEVED SEMANTIC CONTEXT" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "WREN SQL IDENTIFIER CONTRACT" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "EXECUTABLE WREN IDENTIFIER CATALOG" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "sql_table_name_use_exactly" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "sql_column_name_use_exactly" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert all( + "semantic_context_not_sql_identifier" in schema["table_ddl"] + for schema in result["db_schemas"] + ) + assert result["tokens"] > 0 + + +def test_retrieved_schema_separates_exact_sql_names_from_semantic_context(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "modeled_dataset", + "comment": "Business-facing dataset description.", + "columns": [ + { + "type": "COLUMN", + "name": "stored_attribute", + "data_type": "VARCHAR", + "comment": "Business-facing attribute label.", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + table_ddl = result["db_schemas"][0]["table_ddl"] + executable_ddl = table_ddl.split("*/", maxsplit=1)[1] + + assert '"sql_table_name_use_exactly":"modeled_dataset"' in table_ddl + assert '"sql_column_name_use_exactly":"stored_attribute"' in table_ddl + assert "WREN SQL IDENTIFIER CONTRACT" in table_ddl + assert "sql_table_name_use_exactly: modeled_dataset" in table_ddl + assert "sql_column_names_use_exactly:\n- stored_attribute" in table_ddl + assert "END WREN SQL IDENTIFIER CONTRACT" in table_ddl + assert "EXECUTABLE WREN IDENTIFIER CATALOG" in table_ddl + assert "table: modeled_dataset" in table_ddl + assert "columns:\n- stored_attribute" in table_ddl + assert "Do not create identifiers from user wording" in table_ddl + assert ( + '"semantic_context_not_sql_identifier":"Business-facing attribute label."' + in table_ddl + ) + assert "Business-facing attribute label." not in executable_ddl + assert "Business-facing dataset description." not in executable_ddl + assert "CREATE TABLE modeled_dataset" in executable_ddl + assert "stored_attribute VARCHAR" in executable_ddl + + +def test_build_table_ddl_can_render_executable_schema_without_semantic_comments(): + ddl, has_calculated_field, has_json_field = build_table_ddl( + { + "comment": "/* semantic table context */\n", + "name": "modeled_dataset", + "columns": [ + { + "type": "COLUMN", + "comment": "-- semantic field context\n ", + "name": "stored_attribute", + "data_type": "VARCHAR", + "is_primary_key": False, + } + ], + }, + include_semantic_comments=False, + ) + + assert ddl == "CREATE TABLE modeled_dataset (\n stored_attribute VARCHAR\n);" + assert not has_calculated_field + assert not has_json_field + + +def test_check_using_db_schemas_without_pruning_keeps_explicit_table_fast_path(): + class Encoding: + def encode(self, value): + return value.split() + + result = check_using_db_schemas_without_pruning( + construct_db_schemas=[ + { + "type": "TABLE", + "name": "activity", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "primaryKey": "", + } + ], + dbschema_retrieval=[], + encoding=Encoding(), + enable_column_pruning=False, + context_window_size=1000, + ) + + assert [schema["table_name"] for schema in result["db_schemas"]] == ["activity"] + + +def test_build_table_ddl_preserves_join_columns_when_pruned(): + ddl, _, _ = build_table_ddl( + { + "type": "TABLE", + "name": "detail", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "detail_id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": True, + }, + { + "type": "COLUMN", + "name": "parent_id", + "data_type": "INTEGER", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "amount", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + }, + { + "type": "FOREIGN_KEY", + "comment": "", + "constraint": "FOREIGN KEY (parent_id) REFERENCES parent(parent_id)", + "tables": ["parent", "detail"], + "column": "parent_id", + "referenced_table": "parent", + "referenced_column": "parent_id", + }, + ], + }, + columns={"amount"}, + tables={"parent", "detail"}, + ) + + assert "detail_id INTEGER PRIMARY KEY" in ddl + assert "parent_id INTEGER" in ddl + assert "amount DOUBLE" in ddl + assert "FOREIGN KEY (parent_id) REFERENCES parent(parent_id)" in ddl + + +def _schema_document(name: str, columns: list[str]) -> Document: + return Document( + content=str( + { + "name": name, + "type": "TABLE", + "columns": [ + {"name": column, "type": "COLUMN", "data_type": "VARCHAR"} + for column in columns + ], + } + ), + meta={"name": name, "type": "TABLE"}, + ) + + +def test_column_selection_accepts_alternate_results_shape(): + parsed = _parse_column_selection_response( + { + "replies": [ + """ + { + "tables": [ + { + "table_name": "SalesOrderFact", + "columns": ["USDFXSalesValue", "OrderDate"] + } + ] + } + """ + ] + } + ) + + assert parsed == { + "SalesOrderFact": { + "table_name": "SalesOrderFact", + "columns": ["USDFXSalesValue", "OrderDate"], + } + } + + +def test_column_selection_returns_empty_dict_for_malformed_reply(): + parsed = _parse_column_selection_response({"replies": ["not-json"]}) + + assert parsed == {} + + +def test_retrieval_query_augmentation_is_schema_neutral(): + query = "show total amount by year" + + assert _augment_retrieval_query(query) == query + + +def test_table_ranking_prefers_direct_schema_metadata_overlap(): + documents = [ + _schema_document("neutral_records", ["id1", "id2"]), + _schema_document("amount_snapshots", ["amount_value", "snapshot_year"]), + _schema_document("event_records", ["event_code", "event_time"]), + ] + + ranked = _rank_table_names_by_query( + ["neutral_records", "event_records", "amount_snapshots"], + documents, + "show total amount by year", + ) + + assert ranked[0] == "amount_snapshots" + + +def test_table_ranking_prefers_meaningful_column_coverage_over_generic_overlap(): + documents = [ + _schema_document("year_total_staging", ["FY___Would_invoice_date", "total_cost"]), + _schema_document("sales_facts", ["Revenue", "Year", "CustomerName"]), + _schema_document("event_records", ["event_code", "event_time"]), + ] + + ranked = _rank_table_names_by_query( + ["year_total_staging", "event_records", "sales_facts"], + documents, + "Show total revenue by year.", + ) + + assert ranked[0] == "sales_facts" + + +def test_generation_context_limiter_skips_large_table_to_keep_later_candidates(): + encoding = tiktoken.get_encoding("cl100k_base") + retrieval_results = [ + {"table_name": "large_candidate", "table_ddl": "token " * 13_000}, + { + "table_name": "compact_candidate", + "table_ddl": "CREATE TABLE compact_candidate (Revenue DOUBLE, Year INTEGER);", + }, + ] + + limited, _, _, reason = _limit_retrieval_results_for_generation( + retrieval_results, + encoding, + ) + + assert [result["table_name"] for result in limited] == ["compact_candidate"] + assert reason == "ranked_top_k_skipped_token_budget" + + +def test_table_ranking_keeps_order_as_business_subject_token(): + documents = [ + _schema_document( + "tariff_missing_documents", + ["Missing_Document__1_", "Sold_to_Party_Name"], + ), + _schema_document("order_records", ["OrdNo", "OrdDate", "CustName"]), + _schema_document("customer_master", ["CustomerName"]), + ] + + ranked = _rank_table_names_by_query( + ["tariff_missing_documents", "customer_master", "order_records"], + documents, + "Show order records with missing customer names.", + ) + + assert ranked[0] == "order_records" + + +def test_table_ranking_splits_compound_schema_identifiers(): + documents = [ + _schema_document("account_groups", ["acctgroup"]), + _schema_document("balance_snapshots", ["endingbalance"]), + _schema_document("recon_status", ["glaccount", "acctgroup", "glbalance"]), + ] + + ranked = _rank_table_names_by_query( + ["account_groups", "balance_snapshots", "recon_status"], + documents, + "show total GL balance by account group", + ) + + assert ranked[0] == "recon_status" + + +def test_table_ranking_prefers_exact_schema_identifier_mention(): + documents = [ + _schema_document("dbo_AA", ["id", "taskdate"]), + _schema_document("dbo_View_Open_Invoices", ["invoice_number", "invoice_date"]), + _schema_document("dbo_Collections_Tickets_History", ["taskdate", "status"]), + ] + + ranked = _rank_table_names_by_query( + [ + "dbo_View_Open_Invoices", + "dbo_Collections_Tickets_History", + "dbo_AA", + ], + documents, + "How many dbo.AA records are there?", + ) + + assert ranked[0] == "dbo_AA" + + +def test_lexical_column_selection_splits_compound_schema_identifiers(): + result = _lexical_columns_and_tables_needed( + [ + { + "type": "TABLE", + "name": "account_groups", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "acctgroup", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + } + ], + "properties": {}, + "tableReference": None, + }, + { + "type": "TABLE", + "name": "recon_status", + "comment": "", + "columns": [ + { + "type": "COLUMN", + "name": "glaccount", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "acctgroup", + "data_type": "VARCHAR", + "comment": "", + "is_primary_key": False, + }, + { + "type": "COLUMN", + "name": "glbalance", + "data_type": "DOUBLE", + "comment": "", + "is_primary_key": False, + }, + ], + "properties": {}, + "tableReference": None, + }, + ], + "show total GL balance by account group", + ) + + assert list(result)[0] == "recon_status" + assert set(result["recon_status"]["columns"]) >= { + "acctgroup", + "glbalance", + "glaccount", + } diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py new file mode 100644 index 0000000000..b587e5e071 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_preprocess_sql_data.py @@ -0,0 +1,56 @@ +from src.pipelines.retrieval.preprocess_sql_data import preprocess + + +class _FakeEncoding: + def encode(self, value: str) -> list[str]: + return list(value) + + +def test_preprocess_maps_list_rows_to_column_named_records(): + sql_data = { + "columns": [ + {"name": "supplierid", "type": "integer"}, + {"name": "supplier_name", "type": "varchar"}, + {"name": "manufacturing_cost_per_unit", "type": "double"}, + ], + "data": [ + [2, "Supplier 1", 0.06], + [5, "Supplier 2", 0.06], + ], + } + + result = preprocess( + sql_data=sql_data, + encoding=_FakeEncoding(), + context_window_size=1000, + ) + + assert result["sql_data"]["row_records"] == [ + { + "supplierid": 2, + "supplier_name": "Supplier 1", + "manufacturing_cost_per_unit": 0.06, + }, + { + "supplierid": 5, + "supplier_name": "Supplier 2", + "manufacturing_cost_per_unit": 0.06, + }, + ] + assert "row_records" not in sql_data + + +def test_preprocess_keeps_row_records_in_sync_when_rows_are_reduced(): + sql_data = { + "columns": [{"name": "name"}, {"name": "value"}], + "data": [["a", 1], ["b", 2], ["c", 3]], + } + + result = preprocess( + sql_data=sql_data, + encoding=_FakeEncoding(), + context_window_size=1, + ) + + assert len(result["sql_data"]["row_records"]) == len(result["sql_data"]["data"]) + assert result["num_rows_used_in_llm"] == len(result["sql_data"]["data"]) diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py new file mode 100644 index 0000000000..35d016d7ef --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_project_scope_isolation.py @@ -0,0 +1,256 @@ +import pytest + +from src.pipelines.common import retrieve_metadata +from src.pipelines.retrieval import historical_question_retrieval, instructions +from src.pipelines.retrieval import sql_pairs_retrieval + + +PROJECT_FILTER = { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], +} + +PROJECT_DEPLOY_FILTER = { + "operator": "AND", + "conditions": [ + {"field": "project_id", "operator": "==", "value": "project-a"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-a"}, + ], +} + + +class StoreSpy: + def __init__(self, count=0): + self.count = count + self.filters = [] + + async def count_documents(self, filters=None): + self.filters.append(filters) + return self.count + + +class RetrieverSpy: + def __init__(self, documents=None): + self.documents = documents or [] + self.calls = [] + + async def run(self, query_embedding=None, filters=None): + self.calls.append( + { + "query_embedding": query_embedding, + "filters": filters, + } + ) + return {"documents": self.documents} + + +@pytest.mark.asyncio +async def test_metadata_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await retrieve_metadata("project-a", retriever) + + assert result == {} + assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_sql_pairs_count_stays_project_scoped_when_project_has_no_documents(): + store = StoreSpy(count=0) + + count = await sql_pairs_retrieval.count_documents(store, project_id="project-a") + + assert count == 0 + assert store.filters == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_sql_pairs_count_can_be_deploy_scoped(): + store = StoreSpy(count=0) + + count = await sql_pairs_retrieval.count_documents( + store, + project_id="project-a", + mdl_hash="deploy-a", + ) + + assert count == 0 + assert store.filters == [PROJECT_DEPLOY_FILTER] + + +@pytest.mark.asyncio +async def test_sql_pairs_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await sql_pairs_retrieval.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_sql_pairs_retrieval_can_be_deploy_scoped(): + retriever = RetrieverSpy() + + result = await sql_pairs_retrieval.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + mdl_hash="deploy-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [PROJECT_DEPLOY_FILTER] + + +@pytest.mark.asyncio +async def test_historical_question_count_stays_project_scoped_when_project_has_no_documents(): + store = StoreSpy(count=0) + + count = await historical_question_retrieval.count_documents( + store, + project_id="project-a", + ) + + assert count == 0 + assert store.filters == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_historical_question_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await historical_question_retrieval.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + view_questions_retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_instruction_count_stays_project_scoped_when_project_has_no_documents(): + store = StoreSpy(count=0) + + count = await instructions.count_documents(store, project_id="project-a") + + assert count == 0 + assert store.filters == [PROJECT_FILTER] + + +@pytest.mark.asyncio +async def test_instruction_count_can_be_deploy_scoped(): + store = StoreSpy(count=0) + + count = await instructions.count_documents( + store, + project_id="project-a", + mdl_hash="deploy-a", + ) + + assert count == 0 + assert store.filters == [PROJECT_DEPLOY_FILTER] + + +@pytest.mark.asyncio +async def test_instruction_retrieval_does_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await instructions.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": False}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], + } + ] + + +@pytest.mark.asyncio +async def test_instruction_retrieval_can_be_deploy_scoped(): + retriever = RetrieverSpy() + + result = await instructions.retrieval( + {"embedding": [0.1]}, + project_id="project-a", + mdl_hash="deploy-a", + retriever=retriever, + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": False}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-a"}, + ], + } + ] + + +@pytest.mark.asyncio +async def test_default_instructions_do_not_fall_back_to_global_documents(): + retriever = RetrieverSpy() + + result = await instructions.default_instructions( + count_documents=1, + retriever=retriever, + project_id="project-a", + scope_filter=instructions.ScopeFilter(), + scope="sql", + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": True}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + ], + } + ] + + +@pytest.mark.asyncio +async def test_default_instructions_can_be_deploy_scoped(): + retriever = RetrieverSpy() + + result = await instructions.default_instructions( + count_documents=1, + retriever=retriever, + project_id="project-a", + mdl_hash="deploy-a", + scope_filter=instructions.ScopeFilter(), + scope="sql", + ) + + assert result == {"documents": []} + assert [call["filters"] for call in retriever.calls] == [ + { + "operator": "AND", + "conditions": [ + {"field": "is_default", "operator": "==", "value": True}, + {"field": "project_id", "operator": "==", "value": "project-a"}, + {"field": "mdl_hash", "operator": "==", "value": "deploy-a"}, + ], + } + ] diff --git a/wren-ai-service/tests/pytest/pipelines/retrieval/test_sql_knowledge.py b/wren-ai-service/tests/pytest/pipelines/retrieval/test_sql_knowledge.py new file mode 100644 index 0000000000..a12c06cb05 --- /dev/null +++ b/wren-ai-service/tests/pytest/pipelines/retrieval/test_sql_knowledge.py @@ -0,0 +1,21 @@ +from src.pipelines.retrieval.sql_knowledge import SqlKnowledge + + +def test_sql_knowledge_exposes_legacy_structured_instructions(): + sql_knowledge = SqlKnowledge( + { + "text_to_sql_rule": "Use Wren SQL.", + "instructions": { + "date_and_time_functionality": "Use CURRENT_DATE.", + "bigquery": "Use BigQuery-specific guidance.", + "calculated_field_instructions": "Calculated field guidance.", + "metric_instructions": "Metric guidance.", + "json_field_instructions": "JSON guidance.", + }, + } + ) + + assert sql_knowledge.text_to_sql_rule == "Use Wren SQL." + assert sql_knowledge.calculated_field_instructions == "Calculated field guidance." + assert sql_knowledge.metric_instructions == "Metric guidance." + assert sql_knowledge.json_field_instructions == "JSON guidance." diff --git a/wren-ai-service/tests/pytest/providers/test_litellm_llm.py b/wren-ai-service/tests/pytest/providers/test_litellm_llm.py new file mode 100644 index 0000000000..357fe59b0d --- /dev/null +++ b/wren-ai-service/tests/pytest/providers/test_litellm_llm.py @@ -0,0 +1,371 @@ +from types import SimpleNamespace + +import pytest + +from src.providers.llm.litellm import LitellmLLMProvider + + +@pytest.mark.asyncio +async def test_model_kwargs_override_component_generation_defaults(mocker): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="test-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": null}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="test-model", + api_base="https://api.openai.com/v1", + kwargs={ + "temperature": 1, + "max_tokens": 2048, + "response_format": {"type": "text"}, + }, + ) + + generator = provider.get_generator( + generation_kwargs={ + "temperature": 0, + "response_format": { + "type": "json_schema", + "json_schema": {"name": "result", "schema": {}}, + }, + } + ) + + await generator(prompt="Return SQL") + + assert captured_kwargs["temperature"] == 1 + assert captured_kwargs["max_tokens"] == 2048 + assert "response_format" not in captured_kwargs + + +@pytest.mark.asyncio +async def test_runtime_generation_kwargs_override_model_defaults(mocker): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="test-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": null}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="test-model", + api_base="https://api.openai.com/v1", + kwargs={"temperature": 1}, + ) + + generator = provider.get_generator(generation_kwargs={"temperature": 0.5}) + + await generator(prompt="Return SQL", generation_kwargs={"temperature": 0}) + + assert captured_kwargs["temperature"] == 0 + + +@pytest.mark.asyncio +async def test_local_openai_compatible_endpoint_converts_component_json_schema_to_json_object( + mocker, +): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="openai/local-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": null}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="openai/local-model", + api_base="http://localhost/v1", + kwargs={"speed": 0}, + ) + + generator = provider.get_generator( + generation_kwargs={ + "response_format": { + "type": "json_schema", + "json_schema": {"name": "result", "schema": {}}, + }, + } + ) + + await generator(prompt="Return SQL", generation_kwargs={"max_tokens": 4096}) + + assert captured_kwargs["response_format"] == {"type": "json_object"} + assert captured_kwargs["max_tokens"] == 4096 + assert "speed" not in captured_kwargs + + +@pytest.mark.asyncio +async def test_local_openai_compatible_endpoint_preserves_sql_generation_json_schema( + mocker, +): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="openai/local-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": "SELECT 1"}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="openai/local-model", + api_base="http://localhost/v1", + kwargs={"speed": 0}, + ) + + generator = provider.get_generator( + generation_kwargs={ + "preserve_json_schema": True, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "sql_generation_result", + "strict": True, + "schema": { + "type": "object", + "properties": {"sql": {"type": "string"}}, + "required": ["sql"], + "additionalProperties": False, + }, + }, + }, + } + ) + + await generator(prompt="Return SQL") + + assert captured_kwargs["response_format"]["type"] == "json_schema" + assert captured_kwargs["response_format"]["json_schema"]["strict"] is True + assert "preserve_json_schema" not in captured_kwargs + assert "speed" not in captured_kwargs + + +@pytest.mark.asyncio +async def test_sql_generation_json_schema_overrides_loose_model_response_format( + mocker, +): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="openai/local-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": "SELECT 1"}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="openai/local-model", + api_base="http://localhost/v1", + kwargs={"response_format": {"type": "json_object"}}, + ) + + generator = provider.get_generator( + generation_kwargs={ + "preserve_json_schema": True, + "response_format": { + "type": "json_schema", + "json_schema": { + "name": "sql_generation_result", + "strict": True, + "schema": { + "type": "object", + "properties": {"sql": {"type": "string"}}, + "required": ["sql"], + "additionalProperties": False, + }, + }, + }, + } + ) + + await generator(prompt="Return SQL") + + assert captured_kwargs["response_format"]["type"] == "json_schema" + assert captured_kwargs["response_format"]["json_schema"]["strict"] is True + assert "preserve_json_schema" not in captured_kwargs + + +@pytest.mark.asyncio +async def test_litellm_provider_uses_configured_model_name_without_rewriting(mocker): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="configured-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": null}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="configured-model", + api_base="http://localhost/v1", + kwargs={}, + ) + + generator = provider.get_generator() + + await generator(prompt="Return SQL") + + assert captured_kwargs["model"] == "configured-model" + + +@pytest.mark.asyncio +async def test_local_openai_compatible_endpoint_preserves_configured_response_format( + mocker, +): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="openai/local-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": null}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="openai/local-model", + api_base="http://localhost/v1", + kwargs={"response_format": {"type": "json_object"}}, + ) + + generator = provider.get_generator( + generation_kwargs={ + "response_format": { + "type": "json_schema", + "json_schema": {"name": "result", "schema": {}}, + }, + } + ) + + await generator(prompt="Return SQL") + + assert captured_kwargs["response_format"]["type"] == "json_object" + + +@pytest.mark.asyncio +async def test_litellm_provider_does_not_infer_num_predict(mocker): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="configured-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": "SELECT 1"}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="configured-model", + api_base="http://localhost:11434", + ) + + generator = provider.get_generator() + + await generator(prompt="Return SQL", generation_kwargs={"max_tokens": 4096}) + + assert captured_kwargs["max_tokens"] == 4096 + assert "num_predict" not in captured_kwargs + + +@pytest.mark.asyncio +async def test_litellm_provider_preserves_configured_num_predict(mocker): + captured_kwargs = {} + + async def fake_acompletion(**kwargs): + captured_kwargs.update(kwargs) + return SimpleNamespace( + model="configured-model", + choices=[ + SimpleNamespace( + index=0, + finish_reason="stop", + message=SimpleNamespace(content='{"sql": "SELECT 1"}'), + ) + ], + ) + + mocker.patch("src.providers.llm.litellm.acompletion", side_effect=fake_acompletion) + + provider = LitellmLLMProvider( + model="configured-model", + api_base="http://localhost:11434", + kwargs={"num_predict": 2048}, + ) + + generator = provider.get_generator() + + await generator(prompt="Return SQL", generation_kwargs={"max_tokens": 4096}) + + assert captured_kwargs["max_tokens"] == 4096 + assert captured_kwargs["num_predict"] == 2048 diff --git a/wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py b/wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py new file mode 100644 index 0000000000..a21c65c588 --- /dev/null +++ b/wren-ai-service/tests/pytest/providers/test_wren_engine_provider.py @@ -0,0 +1,95 @@ +import aiohttp +import pytest +from aioresponses import CallbackResult, aioresponses + +from src.providers.engine.wren import WrenUI + + +@pytest.mark.asyncio +async def test_wren_ui_dry_plan_calls_graphql_planner(): + endpoint = "http://engine-host" + captured_request = {} + + def callback(_url, **kwargs): + captured_request.update(kwargs) + return CallbackResult(payload={"data": {"dryPlanSql": True}}) + + with aioresponses() as mocked: + mocked.post(f"{endpoint}/api/graphql", callback=callback) + + async with aiohttp.ClientSession() as session: + success, error_message = await WrenUI(endpoint=endpoint).dry_plan( + session, + sql="SELECT 1", + data_source="source", + project_id="project-id", + mdl_hash="deploy-hash", + allow_fallback=False, + ) + + assert success is True + assert error_message == "" + assert captured_request["json"] == { + "query": "mutation DryPlanSql($data: DryPlanSQLDataInput) { dryPlanSql(data: $data) }", + "variables": { + "data": { + "sql": "SELECT 1", + "projectId": "project-id", + "hash": "deploy-hash", + "allowFallback": False, + } + }, + } + + +@pytest.mark.asyncio +async def test_wren_ui_execute_sql_uses_project_preview_manifest(): + endpoint = "http://engine-host" + captured_request = {} + + def callback(_url, **kwargs): + captured_request.update(kwargs) + return CallbackResult(payload={"data": {"previewSql": {"data": [{"ok": 1}]}}}) + + with aioresponses() as mocked: + mocked.post(f"{endpoint}/api/graphql", callback=callback) + + async with aiohttp.ClientSession() as session: + success, _, addition = await WrenUI(endpoint=endpoint).execute_sql( + "SELECT 1", + session, + project_id="project-id", + mdl_hash="deploy-hash", + dry_run=True, + ) + + assert success is True + assert addition == {"correlation_id": ""} + assert captured_request["json"]["variables"]["data"] == { + "sql": "SELECT 1", + "projectId": "project-id", + "hash": "deploy-hash", + "dryRun": True, + "limit": 1, + } + + +@pytest.mark.asyncio +async def test_wren_ui_dry_plan_returns_graphql_error_message(): + endpoint = "http://engine-host" + + with aioresponses() as mocked: + mocked.post( + f"{endpoint}/api/graphql", + payload={"errors": [{"message": "planner failed"}]}, + ) + + async with aiohttp.ClientSession() as session: + success, error_message = await WrenUI(endpoint=endpoint).dry_plan( + session, + sql="SELECT 1", + data_source="source", + ) + + assert success is False + assert error_message == "planner failed" diff --git a/wren-ai-service/tests/pytest/services/mocks.py b/wren-ai-service/tests/pytest/services/mocks.py index a6ce06d323..7de9a4ffaa 100644 --- a/wren-ai-service/tests/pytest/services/mocks.py +++ b/wren-ai-service/tests/pytest/services/mocks.py @@ -9,7 +9,13 @@ class RetrievalMock(retrieval.DbSchemaRetrieval): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + **_, + ): return {"construct_retrieval_results": self._documents} @@ -17,7 +23,12 @@ class SqlPairsRetrievalMock(retrieval.SqlPairsRetrieval): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): return {"formatted_output": {"documents": self._documents}} @@ -25,7 +36,13 @@ class InstructionsRetrievalMock(retrieval.Instructions): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + scope: str = "sql", + mdl_hash: Optional[str] = None, + ): return {"formatted_output": {"documents": self._documents}} @@ -33,7 +50,12 @@ class HistoricalQuestionMock(retrieval.HistoricalQuestionRetrieval): def __init__(self, documents: list = []): self._documents = documents - async def run(self, query: str, project_id: Optional[str] = None): + async def run( + self, + query: str, + project_id: Optional[str] = None, + mdl_hash: Optional[str] = None, + ): return {"formatted_output": {"documents": self._documents}} diff --git a/wren-ai-service/tests/pytest/services/test_ask.py b/wren-ai-service/tests/pytest/services/test_ask.py index d479ac9550..02a346ae09 100644 --- a/wren-ai-service/tests/pytest/services/test_ask.py +++ b/wren-ai-service/tests/pytest/services/test_ask.py @@ -9,9 +9,13 @@ from src.providers import generate_components from src.utils import fetch_wren_ai_docs from src.web.v1.services.ask import ( + AskHistory, AskRequest, AskResultRequest, AskService, + _build_fast_path_grounding_query, + _history_sql_table_names, + _looks_like_simple_analytics_request, ) from src.web.v1.services.semantics_preparation import ( SemanticsPreparationRequest, @@ -19,6 +23,18 @@ ) +def test_ask_defaults_follow_legacy_grounded_sql_flow(): + request = AskRequest(query="q", mdl_hash="mdl-hash") + service = AskService({}) + + assert request.ignore_sql_generation_reasoning is True + assert request.use_dry_plan is True + assert request.allow_dry_plan_fallback is False + assert service._allow_sql_generation_reasoning is False + assert service._enable_column_pruning is True + assert service._max_sql_correction_retries == 0 + + @pytest.fixture def ask_service(): pipe_components = generate_components(settings.components) @@ -65,6 +81,20 @@ def ask_service(): @pytest.fixture def indexing_service(): pipe_components = generate_components(settings.components) + required_components = [ + "db_schema_indexing", + "historical_question_indexing", + "table_description_indexing", + ] + missing_components = [ + component + for component in required_components + if component not in pipe_components + ] + if missing_components: + pytest.skip( + f"Ask integration test requires configured components: {missing_components}" + ) return SemanticsPreparationService( { @@ -102,6 +132,38 @@ def mdl_str(): return orjson.dumps(json.load(f)).decode("utf-8") +def test_word_number_top_followup_uses_simple_fast_path(): + assert _looks_like_simple_analytics_request( + "Show the top five groups from that result." + ) + + +def test_followup_fast_path_grounding_query_includes_latest_history(): + grounding_query = _build_fast_path_grounding_query( + "Show the top five groups from that result.", + [ + AskHistory( + question="How many orders are there by customer?", + sql='SELECT "customer", COUNT(*) AS "record_count" FROM "orders" GROUP BY "customer"', + ) + ], + ) + + assert "How many orders are there by customer?" in grounding_query + assert "orders" in grounding_query + assert "customer" in grounding_query + assert "record_count" in grounding_query + assert "Show the top five groups from that result." in grounding_query + assert "Previous question" not in grounding_query + assert "SELECT" not in grounding_query + + +def test_history_sql_table_names_extracts_prior_verified_tables(): + assert _history_sql_table_names( + 'SELECT "customer", COUNT(*) FROM "orders" JOIN "regions" ON "orders"."region_id" = "regions"."id"' + ) == ["orders", "regions"] + + @pytest.mark.asyncio async def test_ask_with_successful_query( indexing_service: SemanticsPreparationService, diff --git a/wren-ai-service/tests/pytest/services/test_question_recommendation.py b/wren-ai-service/tests/pytest/services/test_question_recommendation.py new file mode 100644 index 0000000000..74a25e96e6 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_question_recommendation.py @@ -0,0 +1,115 @@ +import pytest + +from src.web.v1.services.question_recommendation import QuestionRecommendation + + +class DbSchemaRetrievalPipeline: + async def run(self, **_): + return { + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": "model_1", + "table_ddl": "CREATE TABLE model_1 (attribute_1 VARCHAR, measure_1 FLOAT)", + "manifest_column_names": ["attribute_1", "measure_1"], + } + ] + } + } + + +class EmptyFormattedPipeline: + async def run(self, **_): + return {"formatted_output": {}} + + +class CapturingSqlGenerationPipeline: + def __init__(self): + self.kwargs = {} + self.run_count = 0 + + async def run(self, **kwargs): + self.run_count += 1 + self.kwargs = kwargs + return { + "post_process": { + "valid_generation_result": {}, + "invalid_generation_result": { + "type": "DRY_RUN", + "sql": "", + "original_sql": "", + "error": "invalid", + }, + } + } + + +@pytest.mark.asyncio +async def test_question_recommendation_passes_ddl_context_to_sql_generation(): + sql_generation = CapturingSqlGenerationPipeline() + service = QuestionRecommendation( + pipelines={ + "db_schema_retrieval": DbSchemaRetrievalPipeline(), + "sql_pairs_retrieval": EmptyFormattedPipeline(), + "instructions_retrieval": EmptyFormattedPipeline(), + "sql_generation": sql_generation, + }, + allow_sql_functions_retrieval=False, + allow_sql_knowledge_retrieval=False, + ) + + await service._validate_question( + {"question": "show model records", "category": "General"}, + request_id="request-id", + max_questions=5, + max_categories=3, + project_id="11", + ) + + assert sql_generation.kwargs["contexts"] == [ + "CREATE TABLE model_1 (attribute_1 VARCHAR, measure_1 FLOAT)" + ] + assert "schema_contracts" not in sql_generation.kwargs + + +class MalformedRecommendationPipeline: + async def run(self, **_): + return { + "normalized": { + "questions": [ + {"category": "General"}, + {"question": "show orders"}, + "show users", + ] + } + } + + +@pytest.mark.asyncio +async def test_question_recommendation_skips_malformed_candidates(): + sql_generation = CapturingSqlGenerationPipeline() + service = QuestionRecommendation( + pipelines={ + "question_recommendation": MalformedRecommendationPipeline(), + "db_schema_retrieval": DbSchemaRetrievalPipeline(), + "sql_pairs_retrieval": EmptyFormattedPipeline(), + "instructions_retrieval": EmptyFormattedPipeline(), + "sql_generation": sql_generation, + }, + allow_sql_functions_retrieval=False, + allow_sql_knowledge_retrieval=False, + ) + + await service._recommend( + { + "event_id": "request-id", + "max_questions": 5, + "max_categories": 3, + "project_id": "11", + "allow_data_preview": False, + "use_dry_plan": True, + "allow_dry_plan_fallback": False, + } + ) + + assert sql_generation.run_count == 0 diff --git a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py index 65672afde6..2c655267a3 100644 --- a/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py +++ b/wren-ai-service/tests/pytest/services/test_relationship_recommendation.py @@ -1,7 +1,9 @@ +import asyncio from unittest.mock import AsyncMock import pytest +from src.pipelines.generation.relationship_recommendation import cleaned_models from src.web.v1.services.relationship_recommendation import RelationshipRecommendation @@ -16,6 +18,97 @@ def relationship_recommendation_service(mock_pipeline): return RelationshipRecommendation(pipelines) +@pytest.fixture +def mdl_with_project_relationship_candidate(): + return """ + { + "models": [ + { + "name": "project", + "primaryKey": "id", + "columns": [{"name": "id"}, {"name": "name"}] + }, + { + "name": "view", + "columns": [{"name": "id"}, {"name": "project_id"}] + } + ], + "relationships": [] + } + """ + + +@pytest.fixture +def mdl_with_prefixed_project_model(): + return """ + { + "models": [ + { + "name": "dbo_project", + "primaryKey": "id", + "tableReference": {"schema": "dbo", "table": "project"}, + "columns": [{"name": "id"}, {"name": "name"}] + }, + { + "name": "dbo_view", + "columns": [{"name": "id"}, {"name": "project_id"}] + } + ], + "relationships": [] + } + """ + + +@pytest.fixture +def mdl_with_one_to_one_profile_candidate(): + return """ + { + "models": [ + { + "name": "user", + "primaryKey": "id", + "columns": [{"name": "id"}, {"name": "email"}] + }, + { + "name": "profile", + "primaryKey": "user_id", + "columns": [{"name": "user_id"}, {"name": "display_name"}] + } + ], + "relationships": [] + } + """ + + +@pytest.fixture +def mdl_with_shared_key_candidates(): + return """ + { + "models": [ + { + "name": "employees", + "primaryKey": "emp_no", + "columns": [{"name": "emp_no"}, {"name": "first_name"}] + }, + { + "name": "titles", + "columns": [{"name": "emp_no"}, {"name": "title"}] + }, + { + "name": "departments", + "primaryKey": "dept_no", + "columns": [{"name": "dept_no"}, {"name": "dept_name"}] + }, + { + "name": "dept_emp", + "columns": [{"name": "emp_no"}, {"name": "dept_no"}] + } + ], + "relationships": [] + } + """ + + @pytest.mark.asyncio async def test_recommend_success(relationship_recommendation_service, mock_pipeline): request = RelationshipRecommendation.Input(id="test_id", mdl='{"key": "value"}') @@ -30,6 +123,40 @@ async def test_recommend_success(relationship_recommendation_service, mock_pipel mock_pipeline.run.assert_called_once_with(mdl={"key": "value"}, language="English") +@pytest.mark.asyncio +async def test_recommend_preserves_llm_relationship_reason( + relationship_recommendation_service, + mock_pipeline, + mdl_with_project_relationship_candidate, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_project_relationship_candidate + ) + mock_pipeline.run.return_value = { + "validated": { + "relationships": [ + { + "name": "view_project", + "fromModel": "view", + "fromColumn": "project_id", + "type": "MANY_TO_ONE", + "toModel": "project", + "toColumn": "id", + "reason": "view.project_id references project.id.", + } + ] + } + } + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response["relationships"][0]["reason"] == ( + "view.project_id references project.id." + ) + + @pytest.mark.asyncio async def test_recommend_invalid_mdl(relationship_recommendation_service): request = RelationshipRecommendation.Input(id="test_id", mdl="invalid_json") @@ -87,6 +214,130 @@ def test_getitem_not_found(relationship_recommendation_service): assert "not found" in response.error.message +def test_default_relationship_generation_timeout_is_long_running_safe( + relationship_recommendation_service, +): + assert relationship_recommendation_service._generation_timeout_seconds == 180.0 + + +@pytest.mark.asyncio +async def test_recommend_timeout_fails_without_relationship_suggestions( + mock_pipeline, mdl_with_project_relationship_candidate +): + service = RelationshipRecommendation( + {"relationship_recommendation": mock_pipeline}, + generation_timeout_seconds=0.01, + ) + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_project_relationship_candidate + ) + + async def never_finishes(**_kwargs): + await asyncio.sleep(1) + + mock_pipeline.run.side_effect = never_finishes + + await service.recommend(request) + response = service[request.id] + + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert ( + "An error occurred during relationship recommendation generation" + in response.error.message + ) + + +@pytest.mark.asyncio +async def test_recommend_empty_llm_result_stays_empty( + relationship_recommendation_service, + mock_pipeline, + mdl_with_project_relationship_candidate, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_project_relationship_candidate + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response == {"relationships": []} + + +@pytest.mark.asyncio +async def test_recommend_missing_validated_response_fails( + relationship_recommendation_service, + mock_pipeline, +): + request = RelationshipRecommendation.Input(id="test_id", mdl='{"relationships": []}') + mock_pipeline.run.return_value = {} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "returned no validated response" in response.error.message + + +@pytest.mark.asyncio +async def test_recommend_does_not_create_prefixed_model_relationships( + relationship_recommendation_service, + mock_pipeline, + mdl_with_prefixed_project_model, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_prefixed_project_model + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response == {"relationships": []} + + +@pytest.mark.asyncio +async def test_recommend_does_not_infer_one_to_one_relationships( + relationship_recommendation_service, + mock_pipeline, + mdl_with_one_to_one_profile_candidate, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_one_to_one_profile_candidate + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response == {"relationships": []} + + +@pytest.mark.asyncio +async def test_recommend_does_not_infer_shared_key_relationships( + relationship_recommendation_service, + mock_pipeline, + mdl_with_shared_key_candidates, +): + request = RelationshipRecommendation.Input( + id="test_id", mdl=mdl_with_shared_key_candidates + ) + mock_pipeline.run.return_value = {"validated": {"relationships": []}} + + await relationship_recommendation_service.recommend(request) + response = relationship_recommendation_service[request.id] + + assert response.status == "finished" + assert response.response == {"relationships": []} + + def test_setitem(relationship_recommendation_service): id = "test_id" value = RelationshipRecommendation.Resource(id="test_id", status="finished") @@ -94,3 +345,56 @@ def test_setitem(relationship_recommendation_service): relationship_recommendation_service[id] = value assert relationship_recommendation_service._cache["test_id"] == value + + +def test_cleaned_models_uses_compact_relationship_relevant_payload(): + result = cleaned_models( + { + "models": [ + { + "name": "model_alpha", + "primaryKey": "id", + "baseObject": "physical_source", + "properties": { + "displayName": "Alpha", + "description": "Business alpha records", + "unused": "omitted", + }, + "columns": [ + { + "name": "id", + "type": "INTEGER", + "expression": "complex expression", + "properties": { + "description": "Primary identifier", + "unused": "omitted", + }, + }, + { + "name": "model_beta_id", + "type": "INTEGER", + "relationship": "existing", + }, + ], + } + ] + } + ) + + assert result == [ + { + "name": "model_alpha", + "primaryKey": "id", + "columns": [ + { + "name": "id", + "type": "INTEGER", + "properties": {"description": "Primary identifier"}, + } + ], + "properties": { + "displayName": "Alpha", + "description": "Business alpha records", + }, + } + ] diff --git a/wren-ai-service/tests/pytest/services/test_semantics_description.py b/wren-ai-service/tests/pytest/services/test_semantics_description.py index 4b85ed914f..c6cf867ccd 100644 --- a/wren-ai-service/tests/pytest/services/test_semantics_description.py +++ b/wren-ai-service/tests/pytest/services/test_semantics_description.py @@ -13,8 +13,20 @@ def service(): mock_pipeline.run.return_value = { "output": { "model1": { - "columns": [], - "properties": {"description": "Test description"}, + "name": "model1", + "columns": [ + { + "name": "column1", + "properties": { + "description": "Customer segment for reporting.", + "displayName": "customer segment, segment", + }, + } + ], + "properties": { + "description": "Test description", + "displayName": "test model, model one", + }, } } } @@ -42,8 +54,21 @@ async def test_generate_semantics_description( assert response.status == "finished" assert response.response == { "model1": { - "columns": [], - "properties": {"description": "Test description"}, + "name": "model1", + "columns": [ + { + "name": "column1", + "type": "varchar", + "properties": { + "description": "Customer segment for reporting.", + "displayName": "customer segment, segment", + }, + } + ], + "properties": { + "description": "Test description", + "displayName": "test model, model one", + }, } } assert response.error is None @@ -100,6 +125,35 @@ async def test_generate_semantics_description_with_exception( ) +@pytest.mark.asyncio +async def test_generate_semantics_description_with_llm_timeout_fails_without_fallback(): + mock_pipeline = AsyncMock() + + async def never_returns(**_): + await asyncio.sleep(1) + + mock_pipeline.run.side_effect = never_returns + service = SemanticsDescription( + pipelines={"semantics_description": mock_pipeline}, + generation_timeout_seconds=0.01, + ) + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["model1"], + mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', + ) + + await service.generate(request) + response = service[request.id] + + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "timed out" in response.error.message + + def test_get_semantics_description_result( service: SemanticsDescription, ): @@ -127,6 +181,50 @@ def test_get_non_existent_semantics_description_result( assert "not found" in result.error.message +def test_semantics_description_uses_configured_timeout(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + generation_timeout_seconds=123, + ) + + assert service._generation_timeout_seconds == 123 + + +def test_semantics_description_uses_timeout_without_rewriting_ttl(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + ttl=120, + generation_timeout_seconds=600, + ) + + assert service._generation_timeout_seconds == 600 + assert service._cache.ttl == 120 + + +def test_semantics_description_uses_configured_batch_and_concurrency_limits(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + max_models_per_batch=2, + max_concurrent_tasks=3, + ) + + assert service._max_models_per_batch == 2 + assert service._max_columns_per_batch == 50 + assert service._max_concurrent_tasks == 3 + + +def test_semantics_description_scales_request_timeout_by_concurrency_waves(): + service = SemanticsDescription( + pipelines={"semantics_description": AsyncMock()}, + generation_timeout_seconds=120, + max_concurrent_tasks=4, + ) + + assert service._request_timeout_seconds(1) == 120 + assert service._request_timeout_seconds(4) == 120 + assert service._request_timeout_seconds(5) == 240 + + @pytest.mark.asyncio async def test_batch_processing_with_multiple_models( service: SemanticsDescription, @@ -139,32 +237,74 @@ async def test_batch_processing_with_multiple_models( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model3", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Mock pipeline responses for each chunk - service._pipelines["semantics_description"].run.side_effect = [ - {"output": {"model1": {"description": "Description 1"}}}, - {"output": {"model2": {"description": "Description 2"}}}, - {"output": {"model3": {"description": "Description 3"}}}, - ] + service._pipelines["semantics_description"].run.return_value = { + "output": { + "model1": { + "description": "Description 1", + "displayName": "model one, first model", + "columns": [ + { + "name": "column1", + "properties": { + "description": "Column description 1", + "displayName": "column one, first column", + }, + } + ], + }, + "model2": { + "description": "Description 2", + "displayName": "model two, second model", + "columns": [ + { + "name": "column1", + "properties": { + "description": "Column description 2", + "displayName": "column one, first column", + }, + } + ], + }, + "model3": { + "description": "Description 3", + "displayName": "model three, third model", + "columns": [ + { + "name": "column1", + "properties": { + "description": "Column description 3", + "displayName": "column one, first column", + }, + } + ], + }, + } + } await service.generate(request) response = service[request.id] assert response.id == "test_id" assert response.status == "finished" - assert response.response == { - "model1": {"description": "Description 1"}, - "model2": {"description": "Description 2"}, - "model3": {"description": "Description 3"}, - } + assert response.response["model1"]["properties"]["description"] == "Description 1" + assert response.response["model2"]["properties"]["description"] == "Description 2" + assert response.response["model3"]["properties"]["description"] == "Description 3" + assert len(response.response["model1"]["columns"]) == 1 + assert len(response.response["model2"]["columns"]) == 1 + assert len(response.response["model3"]["columns"]) == 1 chunks = service._chunking(orjson.loads(request.mdl), request) - assert len(chunks) == 3 # Default chunk_size=1 + assert len(chunks) == 3 assert all("user_prompt" in chunk for chunk in chunks) assert all("mdl" in chunk for chunk in chunks) - assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1] + assert [chunk["selected_models"] for chunk in chunks] == [ + ["model1"], + ["model2"], + ["model3"], + ] -def test_batch_processing_with_custom_chunk_size( +def test_batch_processing_groups_small_models_by_prompt( service: SemanticsDescription, ): service["test_id"] = SemanticsDescription.Resource(id="test_id") @@ -175,15 +315,733 @@ def test_batch_processing_with_custom_chunk_size( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model3", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model4", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Test chunking with custom chunk size chunks = service._chunking(orjson.loads(request.mdl), request, chunk_size=2) - assert len(chunks) == 4 - assert [len(chunk["selected_models"]) for chunk in chunks] == [1, 1, 1, 1] + assert len(chunks) == 2 + assert [chunk["selected_models"] for chunk in chunks] == [ + ["model1", "model2"], + ["model3", "model4"], + ] + + +def test_batch_processing_keeps_relationship_context( + service: SemanticsDescription, +): + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["orders", "customers"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [{"name": "customer_id", "type": "varchar"}], + }, + { + "name": "customers", + "columns": [{"name": "customer_id", "type": "varchar"}], + }, + { + "name": "products", + "columns": [{"name": "product_id", "type": "varchar"}], + }, + ], + "relationships": [ + { + "name": "OrdersCustomers", + "models": ["orders", "customers"], + "joinType": "MANY_TO_ONE", + "condition": "orders.customer_id = customers.customer_id", + }, + { + "name": "ProductsOnly", + "models": ["products"], + "joinType": "ONE_TO_ONE", + "condition": "", + }, + ], + } + ).decode(), + ) + + chunks = service._chunking(orjson.loads(request.mdl), request, chunk_size=2) + + assert len(chunks) == 1 + assert chunks[0]["mdl"]["relationships"] == [ + { + "name": "OrdersCustomers", + "models": ["orders", "customers"], + "joinType": "MANY_TO_ONE", + "condition": "orders.customer_id = customers.customer_id", + } + ] + + +def test_default_batch_splits_large_column_groups_by_model( + service: SemanticsDescription, +): + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["model1", "model2"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "model1", + "columns": [ + {"name": f"column_{index}", "type": "varchar"} + for index in range(500) + ], + }, + { + "name": "model2", + "columns": [ + {"name": f"field_{index}", "type": "varchar"} + for index in range(400) + ], + }, + ] + } + ).decode(), + ) + + chunks = service._chunking(orjson.loads(request.mdl), request) + + assert len(chunks) == 18 assert chunks[0]["selected_models"] == ["model1"] - assert chunks[1]["selected_models"] == ["model2"] - assert chunks[2]["selected_models"] == ["model3"] - assert chunks[3]["selected_models"] == ["model4"] + assert chunks[9]["selected_models"] == ["model1"] + assert chunks[10]["selected_models"] == ["model2"] + assert len(chunks[0]["mdl"]["models"][0]["columns"]) == 50 + assert len(chunks[9]["mdl"]["models"][0]["columns"]) == 50 + assert len(chunks[10]["mdl"]["models"][0]["columns"]) == 50 + + +@pytest.mark.asyncio +async def test_column_chunk_outputs_merge_into_single_model( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": f"column_{index}", "type": "varchar"} + for index in range(41) + ], + } + ] + } + ).decode(), + ) + + def response_for_chunk(**kwargs): + model = kwargs["mdl"]["models"][0] + return { + "output": { + "orders": { + "description": "Customer order transactions.", + "displayName": "orders, customer orders", + "columns": [ + { + "name": column["name"], + "properties": { + "description": f"Description for {column['name']}", + "displayName": f"Alias for {column['name']}", + }, + } + for column in model["columns"] + ], + } + } + } + + service._pipelines["semantics_description"].run.side_effect = response_for_chunk + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.response["orders"]["properties"]["description"] == ( + "Customer order transactions." + ) + assert len(response.response["orders"]["columns"]) == 41 + assert service._pipelines["semantics_description"].run.call_count == 1 + + +@pytest.mark.asyncio +async def test_generated_aliases_are_preserved(): + mock_pipeline = AsyncMock() + service = SemanticsDescription(pipelines={"semantics_description": mock_pipeline}) + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Sales reporting dataset", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "net_revenue", "type": "float"}, + ], + } + ] + } + ).decode(), + ) + mock_pipeline.run.return_value = { + "output": { + "orders": { + "name": "orders", + "properties": { + "description": "Customer order transaction records.", + "displayName": "orders, sales orders, transactions", + }, + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Unique identifier for an order.", + "displayName": "order id, order number", + }, + }, + { + "name": "net_revenue", + "properties": { + "description": "Net revenue measure for the order.", + "displayName": "net revenue, sales amount", + }, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.response["orders"]["properties"]["displayName"] == ( + "orders, sales orders, transactions" + ) + assert [ + column["properties"]["displayName"] + for column in response.response["orders"]["columns"] + ] == ["order id, order number", "net revenue, sales amount"] + + +@pytest.mark.asyncio +async def test_malformed_chunk_retries_with_smaller_column_groups( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + service._max_columns_per_batch = 3 + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": f"column_{index}", "type": "varchar"} + for index in range(3) + ], + } + ] + } + ).decode(), + ) + + async def response_for_chunk(**kwargs): + model = kwargs["mdl"]["models"][0] + if len(model["columns"]) > 1: + raise ValueError("Semantics description LLM returned malformed JSON.") + column = model["columns"][0] + return { + "output": { + "orders": { + "description": "Customer order transactions.", + "displayName": "orders, customer orders", + "columns": [ + { + "name": column["name"], + "properties": { + "description": f"Description for {column['name']}", + "displayName": f"Alias for {column['name']}", + }, + } + ], + } + } + } + + service._pipelines["semantics_description"].run.side_effect = response_for_chunk + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert len(response.response["orders"]["columns"]) == 3 + assert service._pipelines["semantics_description"].run.call_count == 5 + + +@pytest.mark.asyncio +async def test_timed_out_chunk_retries_with_smaller_column_groups_and_relationships(): + mock_pipeline = AsyncMock() + service = SemanticsDescription( + pipelines={"semantics_description": mock_pipeline}, + generation_timeout_seconds=0.01, + max_columns_per_batch=2, + ) + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "customer_id", "type": "varchar"}, + ], + } + ], + "relationships": [ + { + "name": "OrdersCustomers", + "models": ["orders", "customers"], + "joinType": "MANY_TO_ONE", + "condition": "orders.customer_id = customers.customer_id", + } + ], + } + ).decode(), + ) + observed_relationships = [] + + async def response_for_chunk(**kwargs): + model = kwargs["mdl"]["models"][0] + observed_relationships.append(kwargs["mdl"].get("relationships", [])) + if len(model["columns"]) > 1: + await asyncio.sleep(1) + column = model["columns"][0] + return { + "output": { + "orders": { + "description": "Customer order transactions.", + "displayName": "orders, customer orders", + "columns": [ + { + "name": column["name"], + "properties": { + "description": f"Description for {column['name']}", + "displayName": f"Alias for {column['name']}", + }, + } + ], + } + } + } + + mock_pipeline.run.side_effect = response_for_chunk + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert [ + column["properties"]["description"] + for column in response.response["orders"]["columns"] + ] == ["Description for order_id", "Description for customer_id"] + assert all( + relationships == [ + { + "name": "OrdersCustomers", + "models": ["orders", "customers"], + "joinType": "MANY_TO_ONE", + "condition": "orders.customer_id = customers.customer_id", + } + ] + for relationships in observed_relationships + ) + + +@pytest.mark.asyncio +async def test_truncated_smallest_chunk_fails_without_fallback( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "properties": {"description": "Existing order model."}, + "columns": [ + { + "name": "order_id", + "type": "varchar", + "properties": {"description": "Existing order id."}, + } + ], + } + ] + } + ).decode(), + ) + + async def truncated_response(**kwargs): + raise ValueError( + "The completion for index 0 has been truncated before reaching a " + "natural stopping point. Finish reason: length. unexpected end of data" + ) + + service._pipelines["semantics_description"].run.side_effect = truncated_response + + await service.generate(request) + response = service[request.id] + + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "truncated" in response.error.message + assert service._pipelines["semantics_description"].run.call_count == 1 + + +@pytest.mark.asyncio +async def test_incomplete_llm_output_fails_without_fallback( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["orders", "customers"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "order_date", "type": "date"}, + ], + }, + { + "name": "customers", + "columns": [ + {"name": "customer_id", "type": "varchar"}, + ], + }, + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "orders": { + "name": "orders", + "columns": [ + { + "name": "order_id", + "properties": {"description": "Unique order identifier."}, + } + ], + "properties": {"description": "Customer purchase transactions."}, + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "failed" + assert response.response is None + assert response.error.code == "OTHERS" + assert "omitted required metadata" in response.error.message + + +@pytest.mark.asyncio +async def test_llm_descriptions_are_not_rewritten_by_service( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Sales and operations reporting dataset", + selected_models=["dbo_xStageLoad2"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "dbo_xStageLoad2", + "columns": [ + {"name": "Division", "type": "varchar"}, + {"name": "SalesPerson", "type": "varchar"}, + {"name": "SalesAmount", "type": "float"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "dbo_xStageLoad2": { + "name": "dbo_xStageLoad2", + "columns": [ + { + "name": "Division", + "properties": { + "description": "Stores the Division value used to describe or analyze xStage records.", + "displayName": "division, business division", + }, + }, + { + "name": "SalesPerson", + "properties": { + "description": "SalesPerson", + "displayName": "sales person, account owner", + }, + }, + { + "name": "SalesAmount", + "properties": { + "description": "Stores the SalesAmount value.", + "displayName": "sales amount, revenue amount", + }, + }, + ], + "properties": { + "description": "Contains business records for xStageLoad2.", + "displayName": "stage load, staging records", + }, + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.response["dbo_xStageLoad2"]["properties"]["description"] == ( + "Contains business records for xStageLoad2." + ) + assert [ + column["properties"]["description"] + for column in response.response["dbo_xStageLoad2"]["columns"] + ] == [ + "Stores the Division value used to describe or analyze xStage records.", + "SalesPerson", + "Stores the SalesAmount value.", + ] + + +@pytest.mark.asyncio +async def test_single_model_output_is_bound_to_selected_schema_name( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the selected datasource", + selected_models=["schema_source_table"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "schema_source_table", + "columns": [ + {"name": "entity_code", "type": "varchar"}, + {"name": "event_date", "type": "timestamp"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "source table": { + "name": "source table", + "properties": { + "description": "Business records for source-table activity and reporting.", + "displayName": "source records, source activity", + }, + "columns": [ + { + "name": "entity_code", + "properties": { + "description": "Entity code used to group and filter records by business entity.", + "displayName": "entity code, business entity code", + }, + }, + { + "name": "event_date", + "properties": { + "description": "Date associated with the business event represented by the record.", + "displayName": "event date, record date", + }, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.error is None + assert list(response.response) == ["schema_source_table"] + assert response.response["schema_source_table"]["name"] == "schema_source_table" + assert [ + column["name"] + for column in response.response["schema_source_table"]["columns"] + ] == [ + "entity_code", + "event_date", + ] + assert response.response["schema_source_table"]["properties"]["displayName"] == ( + "source records, source activity" + ) + + +@pytest.mark.asyncio +async def test_single_column_output_is_bound_to_selected_schema_name( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the selected datasource", + selected_models=["users"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "users", + "columns": [ + {"name": "created_at", "type": "timestamp"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "users": { + "name": "users", + "properties": { + "description": "User account records for application access and profile management.", + "displayName": "users, user accounts", + }, + "columns": [ + { + "name": "created at", + "properties": { + "description": "Timestamp when the user account was created.", + "displayName": "created date, signup date", + }, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.error is None + column = response.response["users"]["columns"][0] + assert column["name"] == "created_at" + assert column["type"] == "timestamp" + assert column["properties"]["displayName"] == "created date, signup date" + + +@pytest.mark.asyncio +async def test_multi_model_name_mismatch_fails_without_ambiguous_binding( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the models", + selected_models=["orders", "customers"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [{"name": "order_id", "type": "varchar"}], + }, + { + "name": "customers", + "columns": [{"name": "customer_id", "type": "varchar"}], + }, + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "sales orders": { + "description": "Customer order transactions.", + "displayName": "orders, sales orders", + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Unique identifier for an order.", + "displayName": "order id, order number", + }, + } + ], + }, + "customer records": { + "description": "Customer master records.", + "displayName": "customers, customer records", + "columns": [ + { + "name": "customer_id", + "properties": { + "description": "Unique identifier for a customer.", + "displayName": "customer id, customer number", + }, + } + ], + }, + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "failed" + assert response.response is None + assert "omitted required metadata" in response.error.message @pytest.mark.asyncio @@ -198,11 +1056,9 @@ async def test_batch_processing_partial_failure( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Mock first chunk succeeds, second chunk fails - service._pipelines["semantics_description"].run.side_effect = [ - {"output": {"model1": {"description": "Description 1"}}}, - Exception("Failed processing model2"), - ] + service._pipelines["semantics_description"].run.side_effect = Exception( + "Failed processing selected models" + ) await service.generate(request) response = service[request.id] @@ -210,7 +1066,7 @@ async def test_batch_processing_partial_failure( assert response.id == "test_id" assert response.status == "failed" assert response.error.code == "OTHERS" - assert "Failed processing model2" in response.error.message + assert "Failed processing selected models" in response.error.message @pytest.mark.asyncio @@ -227,22 +1083,25 @@ async def test_concurrent_updates_no_race_condition( mdl='{"models": [{"name": "model1", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model2", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model3", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model4", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}, {"name": "model5", "columns": [{"name": "column1", "type": "varchar", "notNull": false}]}]}', ) - # Mock pipeline responses with delays to simulate concurrent execution - async def delayed_response(model_num, delay=0.1): - await asyncio.sleep(delay) # Add delay to increase chance of race condition - return { - "output": {f"model{model_num}": {"description": f"Description {model_num}"}} + service._pipelines["semantics_description"].run.return_value = { + "output": { + f"model{i}": { + "description": f"Description {i}", + "displayName": f"model {i}, model {i} alias", + "columns": [ + { + "name": "column1", + "properties": { + "description": f"Column description {i}", + "displayName": f"column one, column {i}", + }, + } + ], + } + for i in range(1, 6) } + } - service._pipelines["semantics_description"].run.side_effect = [ - await delayed_response(1), - await delayed_response(2), - await delayed_response(3), - await delayed_response(4), - await delayed_response(5), - ] - - # Generate response which will process chunks concurrently await service.generate(request) response = service[request.id] @@ -251,6 +1110,71 @@ async def delayed_response(model_num, delay=0.1): assert len(response.response) == 5 assert all(f"model{i}" in response.response for i in range(1, 6)) assert all( - response.response[f"model{i}"]["description"] == f"Description {i}" + response.response[f"model{i}"]["properties"]["description"] + == f"Description {i}" for i in range(1, 6) ) + + +@pytest.mark.asyncio +async def test_repeated_llm_column_descriptions_are_tolerated( + service: SemanticsDescription, +): + service["test_id"] = SemanticsDescription.Resource(id="test_id") + request = SemanticsDescription.GenerateRequest( + id="test_id", + user_prompt="Describe the model", + selected_models=["orders"], + mdl=orjson.dumps( + { + "models": [ + { + "name": "orders", + "columns": [ + {"name": "order_id", "type": "varchar"}, + {"name": "customer_id", "type": "varchar"}, + ], + } + ] + } + ).decode(), + ) + service._pipelines["semantics_description"].run.return_value = { + "output": { + "orders": { + "description": "Customer order transactions.", + "displayName": "orders, customer orders", + "columns": [ + { + "name": "order_id", + "properties": { + "description": "Identifier for reporting.", + "displayName": "order id, order identifier", + }, + }, + { + "name": "customer_id", + "properties": { + "description": "Identifier for reporting.", + "displayName": "customer id, customer identifier", + }, + }, + ], + } + } + } + + await service.generate(request) + response = service[request.id] + + assert response.status == "finished" + assert response.error is None + assert response.response["orders"]["properties"]["description"] == ( + "Customer order transactions." + ) + columns = response.response["orders"]["columns"] + assert [column["name"] for column in columns] == ["order_id", "customer_id"] + assert [ + column["properties"]["description"] for column in columns + ] == ["Identifier for reporting.", "Identifier for reporting."] + assert all(column["properties"]["displayName"] for column in columns) diff --git a/wren-ai-service/tests/pytest/services/test_semantics_preparation.py b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py new file mode 100644 index 0000000000..4a41fbabf8 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_semantics_preparation.py @@ -0,0 +1,126 @@ +import pytest + +from src.web.v1.services.semantics_preparation import ( + SemanticsPreparationService, + SemanticsPreparationStatusRequest, + SemanticsPreparationStatusResponse, +) + + +class CountPipeline: + def __init__(self, count: int): + self.count = count + self.calls = [] + + async def count_documents(self, project_id=None, mdl_hash=None): + self.calls.append({"project_id": project_id, "mdl_hash": mdl_hash}) + return self.count + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_fails_when_exact_schema_documents_are_missing(): + db_schema = CountPipeline(0) + table_description = CountPipeline(1) + service = SemanticsPreparationService( + { + "db_schema": db_schema, + "table_description": table_description, + } + ) + service._prepare_semantics_statuses["deploy-1"] = SemanticsPreparationStatusResponse( + status="finished" + ) + service._prepare_semantics_project_ids["deploy-1"] = "project-1" + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1") + ) + + assert status.status == "failed" + assert status.error.message == "Prepared schema documents are missing for this deployment" + assert db_schema.calls == [{"project_id": "project-1", "mdl_hash": "deploy-1"}] + assert table_description.calls == [ + {"project_id": "project-1", "mdl_hash": "deploy-1"} + ] + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_stays_finished_when_exact_schema_documents_exist(): + service = SemanticsPreparationService( + { + "db_schema": CountPipeline(1), + "table_description": CountPipeline(1), + } + ) + service._prepare_semantics_statuses["deploy-1"] = SemanticsPreparationStatusResponse( + status="finished" + ) + service._prepare_semantics_project_ids["deploy-1"] = "project-1" + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1") + ) + + assert status.status == "finished" + assert status.error is None + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_fails_when_table_descriptions_are_missing(): + service = SemanticsPreparationService( + { + "db_schema": CountPipeline(1), + "table_description": CountPipeline(0), + } + ) + service._prepare_semantics_statuses["deploy-1"] = SemanticsPreparationStatusResponse( + status="finished" + ) + service._prepare_semantics_project_ids["deploy-1"] = "project-1" + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1") + ) + + assert status.status == "failed" + assert status.error.message == "Prepared schema documents are missing for this deployment" + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_recovers_when_status_cache_is_missing(): + db_schema = CountPipeline(1) + table_description = CountPipeline(1) + service = SemanticsPreparationService( + { + "db_schema": db_schema, + "table_description": table_description, + } + ) + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1", project_id="project-1") + ) + + assert status.status == "finished" + assert status.error is None + assert db_schema.calls == [{"project_id": "project-1", "mdl_hash": "deploy-1"}] + assert table_description.calls == [ + {"project_id": "project-1", "mdl_hash": "deploy-1"} + ] + + +@pytest.mark.asyncio +async def test_prepare_semantics_status_does_not_recover_without_table_descriptions(): + service = SemanticsPreparationService( + { + "db_schema": CountPipeline(1), + "table_description": CountPipeline(0), + } + ) + + status = await service.get_prepare_semantics_status( + SemanticsPreparationStatusRequest(mdl_hash="deploy-1", project_id="project-1") + ) + + assert status.status == "failed" + assert status.error.message == "deploy-1 is not found" diff --git a/wren-ai-service/tests/pytest/services/test_sql_corrections.py b/wren-ai-service/tests/pytest/services/test_sql_corrections.py new file mode 100644 index 0000000000..db1cb463f9 --- /dev/null +++ b/wren-ai-service/tests/pytest/services/test_sql_corrections.py @@ -0,0 +1,71 @@ +import pytest + +from src.web.v1.services.sql_corrections import SqlCorrectionService + +DEPLOYED_ORDER_MODEL = "deployed_order_model" +DEPLOYED_SHIP_COUNTRY = "ShipCountry" +DEPLOYED_SELECTED_DDL = ( + f"CREATE TABLE {DEPLOYED_ORDER_MODEL} ({DEPLOYED_SHIP_COUNTRY} VARCHAR)" +) +DEPLOYED_COUNT_SQL = f"SELECT COUNT(*) FROM {DEPLOYED_ORDER_MODEL}" +DEPLOYED_PREVIEW_SQL = f"SELECT * FROM {DEPLOYED_ORDER_MODEL}" + + +class _CapturingRetrievalPipeline: + def __init__(self): + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + return { + "construct_retrieval_results": { + "retrieval_results": [ + { + "table_name": DEPLOYED_ORDER_MODEL, + "table_ddl": DEPLOYED_SELECTED_DDL, + } + ] + } + } + + +class _CapturingSqlCorrectionPipeline: + def __init__(self): + self.calls = [] + + async def run(self, **kwargs): + self.calls.append(kwargs) + return { + "post_process": { + "valid_generation_result": {"sql": DEPLOYED_COUNT_SQL}, + "invalid_generation_result": {}, + } + } + + +@pytest.mark.asyncio +async def test_sql_correction_retrieves_ddl_context_and_preserves_mdl_hash(): + retrieval = _CapturingRetrievalPipeline() + correction = _CapturingSqlCorrectionPipeline() + service = SqlCorrectionService( + { + "db_schema_retrieval": retrieval, + "sql_correction": correction, + }, + allow_sql_knowledge_retrieval=False, + ) + + request = SqlCorrectionService.CorrectionRequest( + event_id="event-1", + sql=DEPLOYED_PREVIEW_SQL, + error="dry run failed", + retrieved_tables=[DEPLOYED_ORDER_MODEL], + mdl_hash="deploy-hash", + ) + + await service.correct(request) + + assert retrieval.calls[0]["tables"] == [DEPLOYED_ORDER_MODEL] + assert retrieval.calls[0]["mdl_hash"] == "deploy-hash" + assert correction.calls[0]["contexts"] == [DEPLOYED_SELECTED_DDL] + assert correction.calls[0]["mdl_hash"] == "deploy-hash" diff --git a/wren-ai-service/tests/pytest/test_config.py b/wren-ai-service/tests/pytest/test_config.py index 70e3ddace3..957b7ae2a5 100644 --- a/wren-ai-service/tests/pytest/test_config.py +++ b/wren-ai-service/tests/pytest/test_config.py @@ -23,6 +23,9 @@ def test_settings_default_values(): assert settings.logging_level == "INFO" assert settings.development is False + assert settings.enable_column_pruning is False + assert settings.allow_sql_generation_reasoning is True + assert settings.max_sql_correction_retries == 3 assert settings.config_path == "config.yaml" diff --git a/wren-ai-service/tools/config/config.codestral-nomic.yaml b/wren-ai-service/tools/config/config.codestral-nomic.yaml new file mode 100644 index 0000000000..26e0375547 --- /dev/null +++ b/wren-ai-service/tools/config/config.codestral-nomic.yaml @@ -0,0 +1,179 @@ +type: llm +provider: litellm_llm +timeout: 600 +models: + - alias: default + model: openai/codestral:22b + api_base: http://10.104.74.10:18002/v1 + api_key_name: LLM_API_KEY + timeout: 600 + context_window_size: 256000 + kwargs: + max_tokens: 4096 + n: 1 + temperature: 0 + top_p: 1 + +--- +type: embedder +provider: litellm_embedder +models: + - alias: default + model: nomic-embed-text:v1.5 + api_base: http://10.104.74.10:8002/v1 + api_key_name: EMBEDDER_API_KEY + timeout: 600 + +--- +type: engine +provider: wren_ui +endpoint: http://localhost:3000 + +--- +type: engine +provider: wren_ibis +endpoint: http://localhost:8000 +source: mssql +manifest: "" +connection_info: "" + +--- +type: engine +provider: wren_engine +endpoint: http://localhost:8080 +manifest: "" + +--- +type: document_store +provider: qdrant +location: http://10.104.74.21:6333 +embedding_model_dim: 768 +timeout: 120 +recreate_index: true + +--- +type: pipeline +pipes: + - name: db_schema_indexing + embedder: litellm_embedder.default + document_store: qdrant + - name: historical_question_indexing + embedder: litellm_embedder.default + document_store: qdrant + - name: table_description_indexing + embedder: litellm_embedder.default + document_store: qdrant + - name: db_schema_retrieval + llm: litellm_llm.default + embedder: litellm_embedder.default + document_store: qdrant + - name: historical_question_retrieval + embedder: litellm_embedder.default + document_store: qdrant + - name: sql_generation + llm: litellm_llm.default + engine: wren_ui + document_store: qdrant + - name: sql_correction + llm: litellm_llm.default + engine: wren_ui + document_store: qdrant + - name: followup_sql_generation + llm: litellm_llm.default + engine: wren_ui + document_store: qdrant + - name: sql_answer + llm: litellm_llm.default + - name: semantics_description + llm: litellm_llm.default + - name: relationship_recommendation + llm: litellm_llm.default + - name: question_recommendation + llm: litellm_llm.default + - name: question_recommendation_sql_generation + llm: litellm_llm.default + engine: wren_ui + document_store: qdrant + - name: chart_generation + llm: litellm_llm.default + - name: chart_adjustment + llm: litellm_llm.default + - name: intent_classification + llm: litellm_llm.default + embedder: litellm_embedder.default + document_store: qdrant + - name: misleading_assistance + llm: litellm_llm.default + - name: data_assistance + llm: litellm_llm.default + - name: sql_pairs_indexing + document_store: qdrant + embedder: litellm_embedder.default + - name: sql_pairs_retrieval + document_store: qdrant + embedder: litellm_embedder.default + llm: litellm_llm.default + - name: preprocess_sql_data + llm: litellm_llm.default + - name: sql_executor + engine: wren_ui + - name: user_guide_assistance + llm: litellm_llm.default + - name: sql_question_generation + llm: litellm_llm.default + - name: sql_generation_reasoning + llm: litellm_llm.default + - name: followup_sql_generation_reasoning + llm: litellm_llm.default + - name: sql_regeneration + llm: litellm_llm.default + engine: wren_ui + - name: evaluation + llm: litellm_llm.default + - name: instructions_indexing + embedder: litellm_embedder.default + document_store: qdrant + - name: instructions_retrieval + embedder: litellm_embedder.default + document_store: qdrant + - name: sql_functions_retrieval + engine: wren_ibis + document_store: qdrant + - name: project_meta_indexing + document_store: qdrant + - name: sql_tables_extraction + llm: litellm_llm.default + - name: sql_diagnosis + llm: litellm_llm.default + - name: sql_knowledge_retrieval + engine: wren_ibis + document_store: qdrant + +--- +settings: + host: 0.0.0.0 + port: 5555 + doc_endpoint: https://docs.getwren.ai + is_oss: true + engine_timeout: 60 + column_indexing_batch_size: 50 + table_retrieval_size: 10 + table_column_retrieval_size: 100 + allow_intent_classification: false + allow_sql_generation_reasoning: true + allow_sql_functions_retrieval: true + enable_column_pruning: false + max_sql_correction_retries: 3 + query_cache_maxsize: 1000 + query_cache_ttl: 3600 + langfuse_host: https://cloud.langfuse.com + langfuse_enable: false + logging_level: DEBUG + development: true + historical_question_retrieval_similarity_threshold: 0.9 + sql_pairs_similarity_threshold: 0.7 + sql_pairs_retrieval_max_size: 10 + instructions_similarity_threshold: 0.7 + instructions_top_k: 25 + pipeline_timeout_seconds: 600 + allow_sql_knowledge_retrieval: false diff --git a/wren-ai-service/tools/config/config.example.yaml b/wren-ai-service/tools/config/config.example.yaml index dd3904529b..eb61e59d42 100644 --- a/wren-ai-service/tools/config/config.example.yaml +++ b/wren-ai-service/tools/config/config.example.yaml @@ -188,13 +188,18 @@ settings: column_indexing_batch_size: 50 table_retrieval_size: 10 table_column_retrieval_size: 100 - allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_intent_classification: false + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 + sql_generation_timeout_seconds: 30 query_cache_maxsize: 1000 query_cache_ttl: 3600 + semantics_description_max_models_per_batch: 4 + semantics_description_max_columns_per_batch: 50 + semantics_description_max_concurrent_tasks: 4 + semantics_description_generation_timeout_seconds: 120 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: DEBUG diff --git a/wren-ai-service/tools/config/config.full.yaml b/wren-ai-service/tools/config/config.full.yaml index 8fc543403e..07959dfb23 100644 --- a/wren-ai-service/tools/config/config.full.yaml +++ b/wren-ai-service/tools/config/config.full.yaml @@ -186,12 +186,17 @@ settings: table_retrieval_size: 10 table_column_retrieval_size: 100 query_cache_maxsize: 1000 - allow_intent_classification: true - allow_sql_generation_reasoning: true + allow_intent_classification: false + allow_sql_generation_reasoning: false allow_sql_functions_retrieval: true enable_column_pruning: false - max_sql_correction_retries: 3 + max_sql_correction_retries: 1 + sql_generation_timeout_seconds: 30 query_cache_ttl: 3600 + semantics_description_max_models_per_batch: 4 + semantics_description_max_columns_per_batch: 50 + semantics_description_max_concurrent_tasks: 4 + semantics_description_generation_timeout_seconds: 120 langfuse_host: https://cloud.langfuse.com langfuse_enable: true logging_level: INFO diff --git a/wren-engine b/wren-engine index 47ca29ebba..44d0811961 160000 --- a/wren-engine +++ b/wren-engine @@ -1 +1 @@ -Subproject commit 47ca29ebba291100ba5d70ce1790f9887eaed7a0 +Subproject commit 44d08119612dca9c8ff007fa12b305fd1aad7593 diff --git a/wren-mdl/mdl.schema.json b/wren-mdl/mdl.schema.json index b2f6d95e73..33a42fade1 100644 --- a/wren-mdl/mdl.schema.json +++ b/wren-mdl/mdl.schema.json @@ -4,6 +4,81 @@ "title": "WrenMDL Manifest Schema", "description": "A schema for WrenMDL manifest file", "$defs": { + "semanticProperties": { + "description": "Optional semantic metadata used by WrenAI retrieval and SQL generation. These fields add business context while preserving exact MDL model, column, and relationship identifiers.", + "type": "object", + "properties": { + "displayName": { + "description": "Business-friendly display name or alias.", + "type": "string" + }, + "description": { + "description": "Business meaning and purpose.", + "type": "string" + }, + "aliases": { + "description": "Alternative names users may use in natural-language questions.", + "type": "array", + "items": { + "type": "string" + } + }, + "synonyms": { + "description": "Synonyms that refer to this object.", + "type": "array", + "items": { + "type": "string" + } + }, + "businessContext": { + "description": "Business context and common use cases.", + "type": "string" + }, + "dataMeaning": { + "description": "Meaning of the data based on schema context.", + "type": "string" + }, + "semanticType": { + "description": "Semantic role such as identifier, date, measure, dimension, status, currency, quantity, cost, or revenue.", + "oneOf": [ + { + "type": "string" + }, + { + "type": "array", + "items": { + "type": "string" + } + } + ] + }, + "useCases": { + "description": "Common analytical uses.", + "type": "array", + "items": { + "type": "string" + } + }, + "aggregationDefault": { + "description": "Default aggregation intent when relevant.", + "type": "string" + }, + "format": { + "description": "Expected value format when relevant.", + "type": "string" + }, + "examples": { + "description": "Representative values or phrasing examples.", + "type": "array", + "items": { + "type": ["string", "number", "boolean", "null"] + } + } + }, + "additionalProperties": { + "type": ["string", "number", "boolean", "object", "array", "null"] + } + }, "column": { "type": "object", "properties": { @@ -91,11 +166,8 @@ "additionalProperties": false }, "properties": { - "description": "the customize properties of the column", - "type": "object", - "additionalProperties": { - "type": "string" - } + "description": "the customized semantic properties of the column", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "type"], @@ -236,11 +308,8 @@ } }, "properties": { - "description": "the customize properties of the model", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the model", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name"], @@ -285,11 +354,8 @@ "minLength": 1 }, "properties": { - "description": "the customize properties of the relationship", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the relationship", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "models", "joinType", "condition"] @@ -376,10 +442,7 @@ "pattern": "^\\s*(\\d+(?:\\.\\d+)?)\\s*([a-zA-Z]+)\\s*$" }, "properties": { - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "baseObject", "dimension", "measure"] @@ -403,11 +466,8 @@ "minLength": 1 }, "properties": { - "description": "the customize properties of the view", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the view", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "statement"] @@ -444,22 +504,16 @@ "minLength": 1 }, "properties": { - "description": "the customize properties of the member", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the member", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name"] } }, "properties": { - "description": "the customize properties of the enum", - "type": "object", - "additionalProperties": { - "type": ["string", "number", "boolean", "object", "array", "null"] - } + "description": "the customized semantic properties of the enum", + "$ref": "#/$defs/semanticProperties" } }, "required": ["name", "values"] diff --git a/wren-ui/docs/rbac-architecture.md b/wren-ui/docs/rbac-architecture.md new file mode 100644 index 0000000000..9f824c252b --- /dev/null +++ b/wren-ui/docs/rbac-architecture.md @@ -0,0 +1,100 @@ +# RBAC Foundation + +This document describes the application-level RBAC foundation in Wren UI. It does not enforce permissions yet; it establishes durable role, user, and user-role assignment primitives for future governance work. + +## Scope + +Implemented: + +- Roles: `Admin`, `Manager`, `Analyst`, `Viewer`, plus custom roles. +- Users with local identity metadata. +- User-role assignments. +- GraphQL APIs for role, user, and assignment management. +- Administration UI for user management, role management, and role assignment. + +Deferred: + +- Authorization middleware for the existing AI/query APIs. +- Governance policies. +- Data scoping. +- SQL validation. +- Schema/table-level permissions. +- Teams, LDAP, and Azure AD synchronization. + +## Database Model + +Tables: + +- `roles` + - `id` + - `name` + - `description` + - timestamps +- `users` + - `id` + - `name` + - `email` + - `external_id` + - `identity_provider` + - `is_active` + - timestamps +- `user_roles` + - `id` + - `user_id` + - `role_id` + - timestamps +`external_id` and `identity_provider` are intentionally present now so Teams, LDAP, and Azure AD integrations can later attach external identities without replacing the RBAC tables. + +## Backend Layers + +- Migration: `migrations/20250602000000_create_rbac_tables.js` +- Models: `src/apollo/server/models/rbac.ts` +- Repositories: `src/apollo/server/repositories/rbacRepository.ts` +- Service: `src/apollo/server/services/rbacService.ts` +- Resolver: `src/apollo/server/resolvers/rbacResolver.ts` +- GraphQL schema: `src/apollo/server/schema.ts` + +The service owns validation and duplicate checks. The repository owns persistence and joined user-role mapping queries. + +## GraphQL API + +Queries: + +- `roles` +- `users` +- `userRoleMappings` + +Mutations: + +- `createRole` +- `updateRole` +- `createUser` +- `updateUser` +- `assignRoleToUser` +- `updateUserRoles` +- `removeRoleFromUser` + +## UI + +Navigation: + +- Header tab: `Admin` +- Sidebar section: `Administration` + +Screens: + +- `/administration/users` +- `/administration/roles` +- `/administration/assignments` + +The UI uses the existing Next.js, Apollo Client, Ant Design, and `SiderLayout`/`PageLayout` patterns. + +## Future Permission Model + +Future schema-level or table-level permissions should be added as separate tables referencing `roles.id`, for example: + +- `role_schema_permissions` +- `role_table_permissions` +- `role_policy_bindings` + +This keeps identity and assignment management stable while allowing governance policies to evolve independently. diff --git a/wren-ui/jest.config.js b/wren-ui/jest.config.js index 9cac8ffa54..2bdc6395ce 100644 --- a/wren-ui/jest.config.js +++ b/wren-ui/jest.config.js @@ -3,6 +3,7 @@ module.exports = { preset: 'ts-jest', testEnvironment: 'node', moduleNameMapper: { + '^@/(.*)$': '/src/$1', '^@server/(.*)$': '/src/apollo/server/$1', }, modulePathIgnorePatterns: ['/e2e/'], diff --git a/wren-ui/knexfile.js b/wren-ui/knexfile.js index 72c26263a4..9d537ea031 100644 --- a/wren-ui/knexfile.js +++ b/wren-ui/knexfile.js @@ -1,19 +1,90 @@ // Update with your config settings. +const normalizeDbType = (dbType) => + (dbType || 'mssql').trim().toLowerCase().replace(/[-_ ]/g, ''); + +const ensureMssqlConfig = () => { + if (process.env.MSSQL_URL) { + return; + } + + const missingFields = [ + ['MSSQL_HOST', process.env.MSSQL_HOST], + ['MSSQL_DATABASE', process.env.MSSQL_DATABASE], + ['MSSQL_USER', process.env.MSSQL_USER], + ['MSSQL_PASSWORD', process.env.MSSQL_PASSWORD], + ].filter(([, value]) => !value); + + if (missingFields.length > 0) { + throw new Error( + `MSSQL is the required Wren UI application database. Missing configuration: ${missingFields + .map(([key]) => key) + .join(', ')}`, + ); + } +}; + +const parseBooleanUrlParam = (searchParams, key, fallback) => { + const value = searchParams.get(key); + if (value === null) return fallback; + return value.toLowerCase() === 'true'; +}; + +const getMssqlConnection = () => { + if (process.env.MSSQL_URL) { + const url = new URL(process.env.MSSQL_URL); + return { + server: url.hostname, + port: url.port ? parseInt(url.port) : 1433, + database: decodeURIComponent(url.pathname.replace(/^\//, '')), + user: decodeURIComponent(url.username), + password: decodeURIComponent(url.password), + options: { + encrypt: parseBooleanUrlParam(url.searchParams, 'encrypt', false), + trustServerCertificate: parseBooleanUrlParam( + url.searchParams, + 'trustServerCertificate', + true, + ), + }, + }; + } + + return { + server: process.env.MSSQL_HOST || 'localhost', + port: process.env.MSSQL_PORT ? parseInt(process.env.MSSQL_PORT) : 1433, + database: process.env.MSSQL_DATABASE || 'wren_ui', + user: process.env.MSSQL_USER, + password: process.env.MSSQL_PASSWORD, + options: { + encrypt: process.env.MSSQL_ENCRYPT === 'true', + trustServerCertificate: + process.env.MSSQL_TRUST_SERVER_CERTIFICATE !== 'false', + }, + }; +}; + +const dbType = normalizeDbType(process.env.DB_TYPE); + /** * @type { Object. } */ -if (process.env.DB_TYPE === 'pg') { +if (dbType === 'pg' || dbType === 'postgres' || dbType === 'postgresql') { console.log('Using Postgres'); module.exports = { client: 'pg', connection: process.env.PG_URL, }; -} else { - console.log('Using SQLite'); +} else if (dbType === 'mssql' || dbType === 'sqlserver') { + console.log('Using MSSQL'); + ensureMssqlConfig(); module.exports = { - client: 'better-sqlite3', - connection: process.env.SQLITE_FILE || './db.sqlite3', - useNullAsDefault: true, + client: 'mssql', + connection: getMssqlConnection(), + pool: { min: 2, max: 10 }, }; +} else { + throw new Error( + `Unsupported DB_TYPE "${process.env.DB_TYPE || ''}". Wren UI application storage now requires MSSQL.`, + ); } diff --git a/wren-ui/migrations/20250602000000_create_rbac_tables.js b/wren-ui/migrations/20250602000000_create_rbac_tables.js new file mode 100644 index 0000000000..92bd99a2bc --- /dev/null +++ b/wren-ui/migrations/20250602000000_create_rbac_tables.js @@ -0,0 +1,76 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + await knex.schema.createTable('roles', (table) => { + table.increments('id').primary(); + table.string('name', 80).notNullable().unique(); + table + .text('description') + .nullable() + .comment('Human-readable role purpose and future governance notes'); + table.timestamps(true, true); + }); + + await knex.schema.createTable('users', (table) => { + table.increments('id').primary(); + table.string('name', 160).notNullable(); + table.string('email', 320).notNullable().unique(); + table.string('external_id', 255).nullable().unique(); + table.string('identity_provider', 80).nullable(); + table.boolean('is_active').notNullable().defaultTo(true); + table.timestamps(true, true); + }); + + await knex.schema.createTable('user_roles', (table) => { + table.increments('id').primary(); + table.integer('user_id').notNullable(); + table.integer('role_id').notNullable(); + table.timestamps(true, true); + + table.foreign('user_id').references('users.id').onDelete('CASCADE'); + table.foreign('role_id').references('roles.id').onDelete('CASCADE'); + table.unique(['user_id', 'role_id']); + }); + + const now = new Date().toISOString(); + await knex('roles').insert([ + { + name: 'Admin', + description: 'Full administration access foundation role.', + created_at: now, + updated_at: now, + }, + { + name: 'Manager', + description: + 'Manages users, assignments, and future governance workflows.', + created_at: now, + updated_at: now, + }, + { + name: 'Analyst', + description: 'Creates and analyzes project content.', + created_at: now, + updated_at: now, + }, + { + name: 'Viewer', + description: + 'Read-only foundation role for future permission enforcement.', + created_at: now, + updated_at: now, + }, + ]); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.dropTableIfExists('user_roles'); + await knex.schema.dropTableIfExists('users'); + await knex.schema.dropTableIfExists('roles'); +}; diff --git a/wren-ui/migrations/20250603000000_create_organization_table.js b/wren-ui/migrations/20250603000000_create_organization_table.js new file mode 100644 index 0000000000..9f87776f3b --- /dev/null +++ b/wren-ui/migrations/20250603000000_create_organization_table.js @@ -0,0 +1,33 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = function (knex) { + return knex.schema.createTable('organization', (table) => { + table.increments('id').comment('ID'); + table.string('name').notNullable().comment('Organization display name'); + table + .string('identifier') + .notNullable() + .unique() + .comment('Organization identifier used in selectors and APIs'); + table + .text('description') + .nullable() + .comment('Optional organization description'); + table + .boolean('is_current') + .notNullable() + .defaultTo(false) + .comment('Whether the organization is currently selected'); + table.timestamps(true, true); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = function (knex) { + return knex.schema.dropTable('organization'); +}; diff --git a/wren-ui/migrations/20250604000000_create_organization_member_tables.js b/wren-ui/migrations/20250604000000_create_organization_member_tables.js new file mode 100644 index 0000000000..8361f107f0 --- /dev/null +++ b/wren-ui/migrations/20250604000000_create_organization_member_tables.js @@ -0,0 +1,44 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + await knex.schema.createTable('organization_members', (table) => { + table.increments('id').primary(); + table.integer('organization_id').notNullable(); + table.integer('user_id').notNullable(); + table.string('organization_role', 80).notNullable(); + table.timestamps(true, true); + + table + .foreign('organization_id') + .references('organization.id') + .onDelete('CASCADE'); + table.foreign('user_id').references('users.id').onDelete('CASCADE'); + table.unique(['organization_id', 'user_id']); + }); + + await knex.schema.createTable('organization_member_projects', (table) => { + table.increments('id').primary(); + table.integer('organization_member_id').notNullable(); + table.integer('project_id').notNullable(); + table.string('permission', 80).notNullable(); + table.timestamps(true, true); + + table + .foreign('organization_member_id') + .references('organization_members.id') + .onDelete('CASCADE'); + table.foreign('project_id').references('project.id').onDelete('CASCADE'); + table.unique(['organization_member_id', 'project_id']); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.dropTableIfExists('organization_member_projects'); + await knex.schema.dropTableIfExists('organization_members'); +}; diff --git a/wren-ui/migrations/20250604010000_create_organization_invitation_tables.js b/wren-ui/migrations/20250604010000_create_organization_invitation_tables.js new file mode 100644 index 0000000000..b5db435fa6 --- /dev/null +++ b/wren-ui/migrations/20250604010000_create_organization_invitation_tables.js @@ -0,0 +1,51 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + await knex.schema.createTable('organization_invitations', (table) => { + table.increments('id').primary(); + table.integer('organization_id').notNullable(); + table.integer('invited_by_user_id').nullable(); + table.string('email', 255).notNullable(); + table.string('organization_role', 80).notNullable(); + table.string('token', 255).notNullable().unique(); + table.string('status', 40).notNullable().defaultTo('Pending'); + table.timestamp('expires_at').notNullable(); + table.timestamp('accepted_at').nullable(); + table.timestamps(true, true); + + table + .foreign('organization_id') + .references('organization.id') + .onDelete('CASCADE'); + table + .foreign('invited_by_user_id') + .references('users.id') + .onDelete('SET NULL'); + }); + + await knex.schema.createTable('organization_invitation_projects', (table) => { + table.increments('id').primary(); + table.integer('organization_invitation_id').notNullable(); + table.integer('project_id').notNullable(); + table.string('permission', 80).notNullable(); + table.timestamps(true, true); + + table + .foreign('organization_invitation_id') + .references('organization_invitations.id') + .onDelete('CASCADE'); + table.foreign('project_id').references('project.id').onDelete('CASCADE'); + table.unique(['organization_invitation_id', 'project_id']); + }); +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.dropTableIfExists('organization_invitation_projects'); + await knex.schema.dropTableIfExists('organization_invitations'); +}; diff --git a/wren-ui/migrations/20260604000000_add_workspace_project_metadata.js b/wren-ui/migrations/20260604000000_add_workspace_project_metadata.js new file mode 100644 index 0000000000..9de02f63b8 --- /dev/null +++ b/wren-ui/migrations/20260604000000_add_workspace_project_metadata.js @@ -0,0 +1,35 @@ +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.up = async function (knex) { + await knex.schema.alterTable('project', (table) => { + table + .string('project_type') + .notNullable() + .defaultTo('CLASSIC') + .comment('Workspace project type, either AGENTIC or CLASSIC'); + table + .boolean('is_current') + .notNullable() + .defaultTo(false) + .comment('Whether this project is the active project in the workspace'); + }); + + const projects = await knex('project').select('id').orderBy('id', 'asc'); + if (projects.length > 0) { + await knex('project').update({ is_current: false }); + await knex('project').where({ id: projects[0].id }).update({ is_current: true }); + } +}; + +/** + * @param { import("knex").Knex } knex + * @returns { Promise } + */ +exports.down = async function (knex) { + await knex.schema.alterTable('project', (table) => { + table.dropColumn('project_type'); + table.dropColumn('is_current'); + }); +}; diff --git a/wren-ui/next.config.js b/wren-ui/next.config.js index c2791bec9b..a751e072e8 100644 --- a/wren-ui/next.config.js +++ b/wren-ui/next.config.js @@ -7,12 +7,14 @@ const withBundleAnalyzer = require('@next/bundle-analyzer')({ const resolveAlias = { antd$: path.resolve(__dirname, 'src/import/antd'), + 'rc-util/es': path.resolve(__dirname, 'node_modules/rc-util/lib'), }; /** @type {import('next').NextConfig} */ const nextConfig = withLess({ output: 'standalone', staticPageGenerationTimeout: 1000, + transpilePackages: ['rc-util'], compiler: { // Enables the styled-components SWC transform styledComponents: { diff --git a/wren-ui/openapi.yaml b/wren-ui/openapi.yaml index 8424a08a53..6b0e5f3b11 100644 --- a/wren-ui/openapi.yaml +++ b/wren-ui/openapi.yaml @@ -288,6 +288,106 @@ paths: application/json: schema: $ref: '#/components/schemas/ErrorResponse' + /organizations: + get: + summary: List organizations + description: Returns all organizations stored in the local WrenAI workspace + responses: + '200': + description: Organization list + content: + application/json: + schema: + type: object + properties: + organizations: + type: array + items: + $ref: '#/components/schemas/Organization' + currentProjectName: + type: string + post: + summary: Create organization + description: Creates a new organization in the local WrenAI workspace + requestBody: + required: true + content: + application/json: + schema: + type: object + required: + - name + properties: + name: + type: string + identifier: + type: string + description: + type: string + responses: + '201': + description: Organization created + content: + application/json: + schema: + $ref: '#/components/schemas/Organization' + '400': + description: Invalid payload + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' + '409': + description: Duplicate organization identifier + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' + /organizations/current: + get: + summary: Get current organization + description: Returns the current organization, the organization list, and the current project label + responses: + '200': + description: Current organization payload + content: + application/json: + schema: + type: object + properties: + currentOrganization: + allOf: + - $ref: '#/components/schemas/Organization' + nullable: true + organizations: + type: array + items: + $ref: '#/components/schemas/Organization' + currentProjectName: + type: string + /organizations/{id}/select: + post: + summary: Select current organization + description: Marks the specified organization as the current active organization + parameters: + - name: id + in: path + required: true + schema: + type: integer + responses: + '200': + description: Selected organization + content: + application/json: + schema: + $ref: '#/components/schemas/Organization' + '404': + description: Organization not found + content: + application/json: + schema: + $ref: '#/components/schemas/ErrorResponse' /stream_explanation: get: summary: Stream an explanation @@ -356,3 +456,21 @@ components: properties: type: object description: Additional column properties + Organization: + type: object + properties: + id: + type: integer + name: + type: string + identifier: + type: string + description: + type: string + nullable: true + isCurrent: + type: boolean + createdAt: + type: string + updatedAt: + type: string diff --git a/wren-ui/package.json b/wren-ui/package.json index 3f76b0e43b..7e89f39b2e 100644 --- a/wren-ui/package.json +++ b/wren-ui/package.json @@ -10,9 +10,10 @@ "test": "jest", "test:e2e": "npx playwright install chromium && npx playwright test", "check-types": "tsc --noEmit", - "migrate": "yarn knex migrate:latest", - "rollback": "yarn knex migrate:rollback", - "generate-gql": "yarn graphql-codegen --config codegen.yaml" + "migrate": "knex migrate:latest", + "migrate:sqlite-to-mssql": "node -e \"process.env.MIGRATE_SQLITE_TO_MSSQL='true'; require('./tools/knex.js')\"", + "rollback": "knex migrate:rollback", + "generate-gql": "graphql-codegen --config codegen.yaml" }, "dependencies": { "@google-cloud/bigquery": "^6.0.3", @@ -24,6 +25,7 @@ "cron-parser": "^5.1.1", "graphql": "^16.6.0", "graphql-type-json": "^0.3.2", + "is-mobile": "^5.0.0", "knex": "^3.1.0", "lodash": "^4.17.23", "log4js": "^6.9.1", @@ -34,6 +36,7 @@ "pg-cursor": "^2.7.4", "posthog-node": "^4.3.2", "sql-formatter": "^15.3.0", + "tedious": "^18.6.1", "uuid": "^11.1.0" }, "devDependencies": { diff --git a/wren-ui/src/apollo/client/graphql/__types__.ts b/wren-ui/src/apollo/client/graphql/__types__.ts index a7f054a896..8d99acaf67 100644 --- a/wren-ui/src/apollo/client/graphql/__types__.ts +++ b/wren-ui/src/apollo/client/graphql/__types__.ts @@ -11,6 +11,7 @@ export type Scalars = { Boolean: boolean; Int: number; Float: number; + BigIntString: string; DialectSQL: any; JSON: any; }; @@ -77,22 +78,45 @@ export type ApiHistoryResponse = { }; export enum ApiType { + ACCEPT_ORGANIZATION_INVITATION = 'ACCEPT_ORGANIZATION_INVITATION', + ADD_PROJECT_MEMBER = 'ADD_PROJECT_MEMBER', ASK = 'ASK', + CREATE_ORGANIZATION = 'CREATE_ORGANIZATION', CREATE_INSTRUCTION = 'CREATE_INSTRUCTION', CREATE_SQL_PAIR = 'CREATE_SQL_PAIR', + DELETE_CURRENT_ORGANIZATION = 'DELETE_CURRENT_ORGANIZATION', + DELETE_CURRENT_USER = 'DELETE_CURRENT_USER', DELETE_INSTRUCTION = 'DELETE_INSTRUCTION', DELETE_SQL_PAIR = 'DELETE_SQL_PAIR', GENERATE_SQL = 'GENERATE_SQL', GENERATE_SUMMARY = 'GENERATE_SUMMARY', GENERATE_VEGA_CHART = 'GENERATE_VEGA_CHART', + GET_CURRENT_ORGANIZATION = 'GET_CURRENT_ORGANIZATION', + GET_PROJECT_ACCESS = 'GET_PROJECT_ACCESS', + GET_CURRENT_USER = 'GET_CURRENT_USER', GET_INSTRUCTIONS = 'GET_INSTRUCTIONS', GET_MODELS = 'GET_MODELS', + GET_ORGANIZATION_MEMBERS = 'GET_ORGANIZATION_MEMBERS', + GET_ORGANIZATIONS = 'GET_ORGANIZATIONS', + GET_CURRENT_PROJECT = 'GET_CURRENT_PROJECT', + GET_PROJECTS = 'GET_PROJECTS', GET_SQL_PAIRS = 'GET_SQL_PAIRS', + INVITE_ORGANIZATION_MEMBER = 'INVITE_ORGANIZATION_MEMBER', + LEAVE_ORGANIZATION = 'LEAVE_ORGANIZATION', + REMOVE_ORGANIZATION_INVITATION = 'REMOVE_ORGANIZATION_INVITATION', + REMOVE_ORGANIZATION_MEMBER = 'REMOVE_ORGANIZATION_MEMBER', RUN_SQL = 'RUN_SQL', + SELECT_ORGANIZATION = 'SELECT_ORGANIZATION', + SELECT_PROJECT = 'SELECT_PROJECT', STREAM_ASK = 'STREAM_ASK', STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', + UPDATE_CURRENT_ORGANIZATION = 'UPDATE_CURRENT_ORGANIZATION', + UPDATE_CURRENT_USER = 'UPDATE_CURRENT_USER', UPDATE_INSTRUCTION = 'UPDATE_INSTRUCTION', - UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR' + UPDATE_PROJECT_MEMBER = 'UPDATE_PROJECT_MEMBER', + UPDATE_ORGANIZATION_MEMBER = 'UPDATE_ORGANIZATION_MEMBER', + UPDATE_SQL_PAIR = 'UPDATE_SQL_PAIR', + REMOVE_PROJECT_MEMBER = 'REMOVE_PROJECT_MEMBER' } export type AskingTask = { @@ -230,12 +254,22 @@ export type CreateSqlPairInput = { }; export type CreateThreadInput = { + answerContent?: InputMaybe; + answerErrorCode?: InputMaybe; + answerErrorMessage?: InputMaybe; + answerErrorShortMessage?: InputMaybe; + answerStatus?: InputMaybe; question?: InputMaybe; sql?: InputMaybe; taskId?: InputMaybe; }; export type CreateThreadResponseInput = { + answerContent?: InputMaybe; + answerErrorCode?: InputMaybe; + answerErrorMessage?: InputMaybe; + answerErrorShortMessage?: InputMaybe; + answerStatus?: InputMaybe; question?: InputMaybe; sql?: InputMaybe; taskId?: InputMaybe; @@ -255,10 +289,10 @@ export type CustomFieldInput = { export type Dashboard = { __typename?: 'Dashboard'; cacheEnabled: Scalars['Boolean']; - id: Scalars['Int']; + id: Scalars['BigIntString']; name: Scalars['String']; nextScheduledAt?: Maybe; - projectId: Scalars['Int']; + projectId: Scalars['BigIntString']; scheduleCron?: Maybe; scheduleFrequency?: Maybe; scheduleTimezone?: Maybe; @@ -266,10 +300,10 @@ export type Dashboard = { export type DashboardItem = { __typename?: 'DashboardItem'; - dashboardId: Scalars['Int']; + dashboardId: Scalars['BigIntString']; detail: DashboardItemDetail; displayName?: Maybe; - id: Scalars['Int']; + id: Scalars['BigIntString']; layout: DashboardItemLayout; type: DashboardItemType; }; @@ -301,7 +335,7 @@ export enum DashboardItemType { } export type DashboardItemWhereInput = { - id: Scalars['Int']; + id: Scalars['BigIntString']; }; export type DashboardSchedule = { @@ -323,6 +357,7 @@ export type DataSource = { export type DataSourceInput = { properties: Scalars['JSON']; + projectType?: InputMaybe; type: DataSourceName; }; @@ -347,7 +382,7 @@ export enum DatabricksConnectionType { } export type DeleteDashboardItemInput = { - itemId: Scalars['Int']; + itemId: Scalars['BigIntString']; }; export type DetailStep = { @@ -402,7 +437,7 @@ export type DetailedDashboard = { __typename?: 'DetailedDashboard'; cacheEnabled: Scalars['Boolean']; description?: Maybe; - id: Scalars['Int']; + id: Scalars['BigIntString']; items: Array; name: Scalars['String']; nextScheduledAt?: Maybe; @@ -569,9 +604,7 @@ export enum ExpressionName { AVG = 'AVG', CBRT = 'CBRT', CEIL = 'CEIL', - CEILING = 'CEILING', COUNT = 'COUNT', - COUNT_IF = 'COUNT_IF', EXP = 'EXP', FLOOR = 'FLOOR', LENGTH = 'LENGTH', @@ -630,7 +663,7 @@ export type InstructionWhereInput = { export type ItemLayoutInput = { h: Scalars['Int']; - itemId: Scalars['Int']; + itemId: Scalars['BigIntString']; w: Scalars['Int']; x: Scalars['Int']; y: Scalars['Int']; @@ -1103,7 +1136,7 @@ export type PreviewItemResponse = { }; export type PreviewItemSqlInput = { - itemId: Scalars['Int']; + itemId: Scalars['BigIntString']; limit?: InputMaybe; refresh?: InputMaybe; }; @@ -1270,9 +1303,13 @@ export type Relation = { export type RelationInput = { fromColumnId: Scalars['Int']; + fromColumnReferenceName?: InputMaybe; fromModelId: Scalars['Int']; + fromModelReferenceName?: InputMaybe; toColumnId: Scalars['Int']; + toColumnReferenceName?: InputMaybe; toModelId: Scalars['Int']; + toModelReferenceName?: InputMaybe; type: RelationType; }; @@ -1318,6 +1355,11 @@ export enum SampleDatasetName { NBA = 'NBA' } +export enum WorkspaceProjectType { + AGENTIC = 'AGENTIC', + CLASSIC = 'CLASSIC' +} + export type SaveLearningRecordInput = { path: Scalars['String']; }; @@ -1523,6 +1565,7 @@ export type UpdateColumnMetadataInput = { }; export type UpdateCurrentProjectInput = { + displayName?: InputMaybe; language: ProjectLanguage; }; diff --git a/wren-ui/src/apollo/client/graphql/dashboard.generated.ts b/wren-ui/src/apollo/client/graphql/dashboard.generated.ts index 1014a11ee8..67bde922c3 100644 --- a/wren-ui/src/apollo/client/graphql/dashboard.generated.ts +++ b/wren-ui/src/apollo/client/graphql/dashboard.generated.ts @@ -3,19 +3,19 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type CommonDashboardItemFragment = { __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }; +export type CommonDashboardItemFragment = { __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }; export type DashboardItemsQueryVariables = Types.Exact<{ [key: string]: never; }>; -export type DashboardItemsQuery = { __typename?: 'Query', dashboardItems: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; +export type DashboardItemsQuery = { __typename?: 'Query', dashboardItems: Array<{ __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; export type CreateDashboardItemMutationVariables = Types.Exact<{ data: Types.CreateDashboardItemInput; }>; -export type CreateDashboardItemMutation = { __typename?: 'Mutation', createDashboardItem: { __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; +export type CreateDashboardItemMutation = { __typename?: 'Mutation', createDashboardItem: { __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; export type UpdateDashboardItemMutationVariables = Types.Exact<{ where: Types.DashboardItemWhereInput; @@ -23,14 +23,14 @@ export type UpdateDashboardItemMutationVariables = Types.Exact<{ }>; -export type UpdateDashboardItemMutation = { __typename?: 'Mutation', updateDashboardItem: { __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; +export type UpdateDashboardItemMutation = { __typename?: 'Mutation', updateDashboardItem: { __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } } }; export type UpdateDashboardItemLayoutsMutationVariables = Types.Exact<{ data: Types.UpdateDashboardItemLayoutsInput; }>; -export type UpdateDashboardItemLayoutsMutation = { __typename?: 'Mutation', updateDashboardItemLayouts: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; +export type UpdateDashboardItemLayoutsMutation = { __typename?: 'Mutation', updateDashboardItemLayouts: Array<{ __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> }; export type DeleteDashboardItemMutationVariables = Types.Exact<{ where: Types.DashboardItemWhereInput; @@ -51,12 +51,12 @@ export type SetDashboardScheduleMutationVariables = Types.Exact<{ }>; -export type SetDashboardScheduleMutation = { __typename?: 'Mutation', setDashboardSchedule: { __typename?: 'Dashboard', id: number, projectId: number, name: string, cacheEnabled: boolean, scheduleFrequency?: Types.ScheduleFrequencyEnum | null, scheduleTimezone?: string | null, scheduleCron?: string | null, nextScheduledAt?: string | null } }; +export type SetDashboardScheduleMutation = { __typename?: 'Mutation', setDashboardSchedule: { __typename?: 'Dashboard', id: string, projectId: string, name: string, cacheEnabled: boolean, scheduleFrequency?: Types.ScheduleFrequencyEnum | null, scheduleTimezone?: string | null, scheduleCron?: string | null, nextScheduledAt?: string | null } }; export type DashboardQueryVariables = Types.Exact<{ [key: string]: never; }>; -export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: number, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule?: { __typename?: 'DashboardSchedule', frequency?: Types.ScheduleFrequencyEnum | null, hour?: number | null, minute?: number | null, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null } | null, items: Array<{ __typename?: 'DashboardItem', id: number, dashboardId: number, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } }; +export type DashboardQuery = { __typename?: 'Query', dashboard: { __typename?: 'DetailedDashboard', id: string, name: string, description?: string | null, cacheEnabled: boolean, nextScheduledAt?: string | null, schedule?: { __typename?: 'DashboardSchedule', frequency?: Types.ScheduleFrequencyEnum | null, hour?: number | null, minute?: number | null, day?: Types.CacheScheduleDayEnum | null, timezone?: string | null, cron?: string | null } | null, items: Array<{ __typename?: 'DashboardItem', id: string, dashboardId: string, type: Types.DashboardItemType, displayName?: string | null, layout: { __typename?: 'DashboardItemLayout', x: number, y: number, w: number, h: number }, detail: { __typename?: 'DashboardItemDetail', sql: string, chartSchema?: any | null } }> } }; export const CommonDashboardItemFragmentDoc = gql` fragment CommonDashboardItem on DashboardItem { @@ -366,4 +366,4 @@ export function useDashboardLazyQuery(baseOptions?: Apollo.LazyQueryHookOptions< } export type DashboardQueryHookResult = ReturnType; export type DashboardLazyQueryHookResult = ReturnType; -export type DashboardQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DashboardQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/client/graphql/deploy.generated.ts b/wren-ui/src/apollo/client/graphql/deploy.generated.ts index 83d2456a17..13cee2c28d 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.generated.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.generated.ts @@ -3,7 +3,9 @@ import * as Types from './__types__'; import { gql } from '@apollo/client'; import * as Apollo from '@apollo/client'; const defaultOptions = {} as const; -export type DeployMutationVariables = Types.Exact<{ [key: string]: never; }>; +export type DeployMutationVariables = Types.Exact<{ + force?: Types.InputMaybe; +}>; export type DeployMutation = { __typename?: 'Mutation', deploy: any }; @@ -15,8 +17,8 @@ export type DeployStatusQuery = { __typename?: 'Query', modelSync: { __typename? export const DeployDocument = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; export type DeployMutationFn = Apollo.MutationFunction; @@ -77,4 +79,4 @@ export function useDeployStatusLazyQuery(baseOptions?: Apollo.LazyQueryHookOptio } export type DeployStatusQueryHookResult = ReturnType; export type DeployStatusLazyQueryHookResult = ReturnType; -export type DeployStatusQueryResult = Apollo.QueryResult; \ No newline at end of file +export type DeployStatusQueryResult = Apollo.QueryResult; diff --git a/wren-ui/src/apollo/client/graphql/deploy.ts b/wren-ui/src/apollo/client/graphql/deploy.ts index 75fe308d2f..0967269ebd 100644 --- a/wren-ui/src/apollo/client/graphql/deploy.ts +++ b/wren-ui/src/apollo/client/graphql/deploy.ts @@ -1,8 +1,8 @@ import { gql } from '@apollo/client'; export const DEPLOY = gql` - mutation Deploy { - deploy + mutation Deploy($force: Boolean) { + deploy(force: $force) } `; diff --git a/wren-ui/src/apollo/client/graphql/home.generated.ts b/wren-ui/src/apollo/client/graphql/home.generated.ts index 5a477d401a..a8ecfbd961 100644 --- a/wren-ui/src/apollo/client/graphql/home.generated.ts +++ b/wren-ui/src/apollo/client/graphql/home.generated.ts @@ -23,7 +23,7 @@ export type SuggestedQuestionsQueryVariables = Types.Exact<{ [key: string]: neve export type SuggestedQuestionsQuery = { __typename?: 'Query', suggestedQuestions: { __typename?: 'SuggestedQuestionResponse', questions: Array<{ __typename?: 'SuggestedQuestion', label: string, question: string } | null> } }; export type AskingTaskQueryVariables = Types.Exact<{ - taskId: Types.Scalars['String']; + taskId?: Types.InputMaybe; }>; @@ -42,11 +42,11 @@ export type ThreadQueryVariables = Types.Exact<{ export type ThreadQuery = { __typename?: 'Query', thread: { __typename?: 'DetailedThread', id: number, responses: Array<{ __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null }> } }; export type ThreadResponseQueryVariables = Types.Exact<{ - responseId: Types.Scalars['Int']; + responseId?: Types.InputMaybe; }>; -export type ThreadResponseQuery = { __typename?: 'Query', threadResponse: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null } }; +export type ThreadResponseQuery = { __typename?: 'Query', threadResponse?: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null } | null }; export type CreateAskingTaskMutationVariables = Types.Exact<{ data: Types.AskingTaskInput; @@ -144,7 +144,7 @@ export type CreateInstantRecommendedQuestionsMutationVariables = Types.Exact<{ export type CreateInstantRecommendedQuestionsMutation = { __typename?: 'Mutation', createInstantRecommendedQuestions: { __typename?: 'Task', id: string } }; export type InstantRecommendedQuestionsQueryVariables = Types.Exact<{ - taskId: Types.Scalars['String']; + taskId?: Types.InputMaybe; }>; @@ -197,7 +197,7 @@ export type AdjustThreadResponseChartMutationVariables = Types.Exact<{ export type AdjustThreadResponseChartMutation = { __typename?: 'Mutation', adjustThreadResponseChart: { __typename?: 'ThreadResponse', id: number, threadId: number, question: string, sql?: string | null, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, breakdownDetail?: { __typename?: 'ThreadResponseBreakdownDetail', queryId?: string | null, status: Types.AskingTaskStatus, description?: string | null, steps?: Array<{ __typename?: 'DetailStep', summary: string, sql: string, cteName?: string | null }> | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, answerDetail?: { __typename?: 'ThreadResponseAnswerDetail', queryId?: string | null, status?: Types.ThreadResponseAnswerStatus | null, content?: string | null, numRowsUsedInLLM?: number | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, chartDetail?: { __typename?: 'ThreadResponseChartDetail', queryId?: string | null, status: Types.ChartTaskStatus, description?: string | null, chartType?: Types.ChartType | null, chartSchema?: any | null, adjustment?: boolean | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, askingTask?: { __typename?: 'AskingTask', status: Types.AskingTaskStatus, type?: Types.AskingTaskType | null, rephrasedQuestion?: string | null, intentReasoning?: string | null, sqlGenerationReasoning?: string | null, retrievedTables?: Array | null, invalidSql?: string | null, traceId?: string | null, queryId?: string | null, candidates: Array<{ __typename?: 'ResultCandidate', sql: string, type: Types.ResultCandidateType, view?: { __typename?: 'ViewInfo', id: number, name: string, statement: string, displayName: string } | null, sqlPair?: { __typename?: 'SqlPair', id: number, question: string, sql: string, projectId: number } | null }>, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null, adjustment?: { __typename?: 'ThreadResponseAdjustment', type: Types.ThreadResponseAdjustmentType, payload?: any | null } | null, adjustmentTask?: { __typename?: 'AdjustmentTask', queryId?: string | null, status?: Types.AskingTaskStatus | null, sql?: string | null, traceId?: string | null, invalidSql?: string | null, error?: { __typename?: 'Error', code?: string | null, shortMessage?: string | null, message?: string | null, stacktrace?: Array | null } | null } | null } }; export type AdjustmentTaskQueryVariables = Types.Exact<{ - taskId: Types.Scalars['String']; + taskId?: Types.InputMaybe; }>; @@ -391,7 +391,7 @@ export type SuggestedQuestionsQueryHookResult = ReturnType; export type SuggestedQuestionsQueryResult = Apollo.QueryResult; export const AskingTaskDocument = gql` - query AskingTask($taskId: String!) { + query AskingTask($taskId: String) { askingTask(taskId: $taskId) { ...CommonAskingTask } @@ -499,7 +499,7 @@ export type ThreadQueryHookResult = ReturnType; export type ThreadLazyQueryHookResult = ReturnType; export type ThreadQueryResult = Apollo.QueryResult; export const ThreadResponseDocument = gql` - query ThreadResponse($responseId: Int!) { + query ThreadResponse($responseId: Int) { threadResponse(responseId: $responseId) { ...CommonResponse } @@ -960,7 +960,7 @@ export type CreateInstantRecommendedQuestionsMutationHookResult = ReturnType; export type CreateInstantRecommendedQuestionsMutationOptions = Apollo.BaseMutationOptions; export const InstantRecommendedQuestionsDocument = gql` - query InstantRecommendedQuestions($taskId: String!) { + query InstantRecommendedQuestions($taskId: String) { instantRecommendedQuestions(taskId: $taskId) { ...CommonRecommendedQuestionsTask } @@ -1225,9 +1225,9 @@ export type AdjustThreadResponseChartMutationHookResult = ReturnType; export type AdjustThreadResponseChartMutationOptions = Apollo.BaseMutationOptions; export const AdjustmentTaskDocument = gql` - query AdjustmentTask($taskId: String!) { + query AdjustmentTask($taskId: String) { adjustmentTask(taskId: $taskId) { - queryId + queryId status error { code @@ -1330,4 +1330,4 @@ export function useRerunAdjustmentTaskMutation(baseOptions?: Apollo.MutationHook } export type RerunAdjustmentTaskMutationHookResult = ReturnType; export type RerunAdjustmentTaskMutationResult = Apollo.MutationResult; -export type RerunAdjustmentTaskMutationOptions = Apollo.BaseMutationOptions; \ No newline at end of file +export type RerunAdjustmentTaskMutationOptions = Apollo.BaseMutationOptions; diff --git a/wren-ui/src/apollo/client/graphql/home.ts b/wren-ui/src/apollo/client/graphql/home.ts index 1fe6086912..762e41d00a 100644 --- a/wren-ui/src/apollo/client/graphql/home.ts +++ b/wren-ui/src/apollo/client/graphql/home.ts @@ -164,7 +164,7 @@ export const SUGGESTED_QUESTIONS = gql` `; export const ASKING_TASK = gql` - query AskingTask($taskId: String!) { + query AskingTask($taskId: String) { askingTask(taskId: $taskId) { ...CommonAskingTask } @@ -194,7 +194,7 @@ export const THREAD = gql` `; export const THREAD_RESPONSE = gql` - query ThreadResponse($responseId: Int!) { + query ThreadResponse($responseId: Int) { threadResponse(responseId: $responseId) { ...CommonResponse } @@ -317,7 +317,7 @@ export const CREATE_INSTANT_RECOMMENDED_QUESTIONS = gql` `; export const INSTANT_RECOMMENDED_QUESTIONS = gql` - query InstantRecommendedQuestions($taskId: String!) { + query InstantRecommendedQuestions($taskId: String) { instantRecommendedQuestions(taskId: $taskId) { ...CommonRecommendedQuestionsTask } @@ -389,7 +389,7 @@ export const ADJUST_THREAD_RESPONSE_CHART = gql` `; export const ADJUSTMENT_TASK = gql` - query AdjustmentTask($taskId: String!) { + query AdjustmentTask($taskId: String) { adjustmentTask(taskId: $taskId) { queryId status diff --git a/wren-ui/src/apollo/client/graphql/rbac.ts b/wren-ui/src/apollo/client/graphql/rbac.ts new file mode 100644 index 0000000000..6da6060285 --- /dev/null +++ b/wren-ui/src/apollo/client/graphql/rbac.ts @@ -0,0 +1,169 @@ +import { gql } from '@apollo/client'; + +export const ROLE_FIELDS = gql` + fragment RoleFields on Role { + id + name + description + createdAt + updatedAt + } +`; + +export const USER_FIELDS = gql` + fragment UserFields on User { + id + name + email + externalId + identityProvider + isActive + createdAt + updatedAt + } +`; + +export const LIST_RBAC_USERS = gql` + query RbacUsers { + users { + ...UserFields + roles { + ...RoleFields + } + } + roles { + ...RoleFields + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const LIST_RBAC_ROLES = gql` + query RbacRoles { + roles { + ...RoleFields + users { + ...UserFields + } + } + users { + ...UserFields + } + } + + ${ROLE_FIELDS} + ${USER_FIELDS} +`; + +export const LIST_USER_ROLE_MAPPINGS = gql` + query UserRoleMappings { + userRoleMappings { + id + userId + roleId + createdAt + updatedAt + user { + ...UserFields + roles { + ...RoleFields + } + } + role { + ...RoleFields + } + } + users { + ...UserFields + roles { + ...RoleFields + } + } + roles { + ...RoleFields + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const CREATE_ROLE = gql` + mutation CreateRole($data: CreateRoleInput!) { + createRole(data: $data) { + ...RoleFields + } + } + + ${ROLE_FIELDS} +`; + +export const UPDATE_ROLE = gql` + mutation UpdateRole($where: RoleWhereInput!, $data: UpdateRoleInput!) { + updateRole(where: $where, data: $data) { + ...RoleFields + } + } + + ${ROLE_FIELDS} +`; + +export const CREATE_USER = gql` + mutation CreateUser($data: CreateUserInput!) { + createUser(data: $data) { + ...UserFields + roles { + ...RoleFields + } + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const UPDATE_USER = gql` + mutation UpdateUser($where: UserWhereInput!, $data: UpdateUserInput!) { + updateUser(where: $where, data: $data) { + ...UserFields + roles { + ...RoleFields + } + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const ASSIGN_ROLE_TO_USER = gql` + mutation AssignRoleToUser($data: UserRoleInput!) { + assignRoleToUser(data: $data) { + id + userId + roleId + } + } +`; + +export const UPDATE_USER_ROLES = gql` + mutation UpdateUserRoles($data: UpdateUserRolesInput!) { + updateUserRoles(data: $data) { + ...UserFields + roles { + ...RoleFields + } + } + } + + ${USER_FIELDS} + ${ROLE_FIELDS} +`; + +export const REMOVE_ROLE_FROM_USER = gql` + mutation RemoveRoleFromUser($data: UserRoleInput!) { + removeRoleFromUser(data: $data) + } +`; diff --git a/wren-ui/src/apollo/client/index.ts b/wren-ui/src/apollo/client/index.ts index df6d75ffde..c36af89098 100644 --- a/wren-ui/src/apollo/client/index.ts +++ b/wren-ui/src/apollo/client/index.ts @@ -1,7 +1,8 @@ import { ApolloClient, HttpLink, InMemoryCache, from } from '@apollo/client'; -import { onError } from '@apollo/client/link/error'; import errorHandler from '@/utils/errorHandler'; +const { onError } = require('@apollo/client/link/error/error.cjs'); + const apolloErrorLink = onError((error) => errorHandler(error)); const httpLink = new HttpLink({ diff --git a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts index 0d2d4a9c5e..fa8f1bc18c 100644 --- a/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/ibisAdaptor.ts @@ -5,7 +5,7 @@ import { DataSourceName } from '@server/types'; import { Manifest } from '@server/mdl/type'; import * as Errors from '@server/utils/error'; import { getConfig } from '@server/config'; -import { toDockerHost } from '@server/utils'; +import { encodeManifest, toDockerHost } from '@server/utils'; import { CompactColumn, CompactTable, @@ -192,6 +192,7 @@ export interface IbisDryPlanOptions { mdl: Manifest; // TODO: replace sql type with WrenSQL sql: string; + allowFallback?: boolean; } export interface IIbisAdaptor { @@ -268,19 +269,29 @@ export class IbisAdaptor implements IIbisAdaptor { this.ibisServerEndpoint = ibisServerEndpoint; } public async getNativeSql(options: IbisDryPlanOptions): Promise { - const { dataSource, mdl, sql } = options; + const { dataSource, mdl, sql, allowFallback } = options; const body = { sql, - manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + manifestStr: encodeManifest(mdl), }; try { const res = await axios.post( `${this.ibisServerEndpoint}/${this.getIbisApiVersion(IBIS_API_TYPE.DRY_PLAN)}/connector/${dataSourceUrlMap[dataSource]}/dry-plan`, body, + { + headers: { + 'x-wren-fallback_disable': + allowFallback === true ? 'false' : 'true', + }, + }, ); return res.data; } catch (e) { - logger.debug(`Dry plan error: ${e.response?.data || e.message}`); + logger.debug( + `Dry plan error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error during dry plan execution'); } } @@ -296,7 +307,7 @@ export class IbisAdaptor implements IIbisAdaptor { const body = { sql: query, connectionInfo: ibisConnectionInfo, - manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + manifestStr: encodeManifest(mdl), }; try { const res = await axios.post( @@ -322,7 +333,9 @@ export class IbisAdaptor implements IIbisAdaptor { override: res.headers['x-cache-override'] === 'true', }; } catch (e) { - logger.debug(`Query error: ${e.response?.data || e.message}`); + logger.debug( + `Query error: ${this.stringifyDebugValue(e.response?.data || e.message)}`, + ); this.throwError(e, 'Error querying ibis server'); } } @@ -337,7 +350,7 @@ export class IbisAdaptor implements IIbisAdaptor { const body = { sql: query, connectionInfo: ibisConnectionInfo, - manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + manifestStr: encodeManifest(mdl), }; logger.debug(`Dry run sql from ibis with body:`); try { @@ -351,7 +364,11 @@ export class IbisAdaptor implements IIbisAdaptor { processTime: response.headers['x-process-time'], }; } catch (err) { - logger.debug(`Dry run error: ${err.response?.data || err.message}`); + logger.debug( + `Dry run error: ${this.stringifyDebugValue( + err.response?.data || err.message, + )}`, + ); this.throwError(err, 'Error during dry run execution'); } } @@ -395,7 +412,11 @@ export class IbisAdaptor implements IIbisAdaptor { ); return await getTablesByConnectionInfo(ibisConnectionInfo); } catch (e) { - logger.debug(`Get tables error: ${e.response?.data || e.message}`); + logger.debug( + `Get tables error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error getting table from ibis server'); } } @@ -417,7 +438,11 @@ export class IbisAdaptor implements IIbisAdaptor { ); return res.data; } catch (e) { - logger.debug(`Get constraints error: ${e.response?.data || e.message}`); + logger.debug( + `Get constraints error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error getting constraint from ibis server'); } } @@ -433,7 +458,7 @@ export class IbisAdaptor implements IIbisAdaptor { const ibisConnectionInfo = toIbisConnectionInfo(dataSource, connectionInfo); const body = { connectionInfo: ibisConnectionInfo, - manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + manifestStr: encodeManifest(mdl), parameters, }; try { @@ -470,7 +495,7 @@ export class IbisAdaptor implements IIbisAdaptor { const body = { sql, connectionInfo: ibisConnectionInfo, - manifestStr: Buffer.from(JSON.stringify(mdl)).toString('base64'), + manifestStr: encodeManifest(mdl), }; try { logger.debug(`Running model substitution with ibis`); @@ -511,11 +536,24 @@ export class IbisAdaptor implements IIbisAdaptor { ); return res.data; } catch (e) { - logger.debug(`Get version error: ${e.response?.data || e.message}`); + logger.debug( + `Get version error: ${this.stringifyDebugValue( + e.response?.data || e.message, + )}`, + ); this.throwError(e, 'Error getting version from ibis server'); } } + private stringifyDebugValue(value: any): string { + if (typeof value === 'string') return value; + try { + return JSON.stringify(value); + } catch { + return String(value); + } + } + private updateConnectionInfo(connectionInfo: any) { if ( config.otherServiceUsingDocker && @@ -571,7 +609,10 @@ export class IbisAdaptor implements IIbisAdaptor { IBIS_API_TYPE.VALIDATION, IBIS_API_TYPE.MODEL_SUBSTITUTE, ].includes(apiType); - if (useV3) logger.debug('Using ibis v3 api'); + + if (useV3) { + logger.debug('Using ibis v3 api'); + } return useV3 ? 'v3' : 'v2'; } diff --git a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts index 44bf4dc5ee..3755b9f8d4 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/ibisAdaptor.test.ts @@ -400,7 +400,7 @@ describe('IbisAdaptor', () => { expect(result).toEqual({ valid: true, message: null }); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v3/connector/postgres/validate/column_is_valid`, + `${ibisServerEndpoint}/v2/connector/postgres/validate/column_is_valid`, { connectionInfo: { connectionUrl: postgresConnectionUrl }, manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( @@ -432,7 +432,7 @@ describe('IbisAdaptor', () => { expect(result).toEqual({ valid: false, message: 'Error' }); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v3/connector/postgres/validate/column_is_valid`, + `${ibisServerEndpoint}/v2/connector/postgres/validate/column_is_valid`, { connectionInfo: { connectionUrl: postgresConnectionUrl }, manifestStr: Buffer.from(JSON.stringify(mockManifest)).toString( @@ -512,7 +512,7 @@ describe('IbisAdaptor', () => { expect(res.override).toEqual(false); expect(new Date(res.cacheOverrodeAt).getTime()).toBeGreaterThan(0); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v3/connector/postgres/query?cacheEnable=true`, + `${ibisServerEndpoint}/v2/connector/postgres/query?cacheEnable=true`, expect.any(Object), expect.any(Object), ); @@ -547,7 +547,7 @@ describe('IbisAdaptor', () => { expect(res.data).toEqual([[1]]); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v3/connector/postgres/query?cacheEnable=true&overrideCache=true`, + `${ibisServerEndpoint}/v2/connector/postgres/query?cacheEnable=true&overrideCache=true`, expect.any(Object), expect.any(Object), ); @@ -728,7 +728,7 @@ describe('IbisAdaptor', () => { expect(result).toEqual('SELECT * FROM substituted_table'); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v3/connector/postgres/model-substitute`, + `${ibisServerEndpoint}/v2/connector/postgres/model-substitute`, { sql: 'SELECT * FROM test_table', connectionInfo: { connectionUrl: postgresConnectionUrl }, @@ -974,7 +974,7 @@ describe('IbisAdaptor', () => { expect(result).toEqual('SELECT * FROM substituted_table'); expect(mockedAxios.post).toHaveBeenCalledWith( - `${ibisServerEndpoint}/v3/connector/postgres/model-substitute`, + `${ibisServerEndpoint}/v2/connector/postgres/model-substitute`, { sql: 'SELECT * FROM test_table', connectionInfo: { connectionUrl: postgresConnectionUrl }, diff --git a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts index 69252623ff..3f9300a858 100644 --- a/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts +++ b/wren-ui/src/apollo/server/adaptors/tests/wrenAIAdaptor.test.ts @@ -1,6 +1,7 @@ import axios from 'axios'; import { WrenAIAdaptor } from '../wrenAIAdaptor'; import { + AskInput, RecommendationQuestionsInput, RecommendationQuestionStatus, } from '@server/models/adaptor'; @@ -33,6 +34,71 @@ describe('WrenAIAdaptor', () => { jest.clearAllMocks(); }); + describe('deploy', () => { + it('should send the project id to scope indexed semantics', async () => { + const mockInput = { + manifest: sampleManifest, + hash: 'deploy-hash', + projectId: 123, + }; + mockedAxios.post.mockResolvedValueOnce({ + data: { id: mockInput.hash }, + }); + mockedAxios.get.mockResolvedValueOnce({ + data: { status: 'finished' }, + }); + + await adaptor.deploy(mockInput); + + expect(mockedAxios.post).toHaveBeenCalledWith( + `${baseEndpoint}/v1/semantics-preparations`, + { + mdl: JSON.stringify(mockInput.manifest), + id: mockInput.hash, + project_id: mockInput.projectId.toString(), + }, + ); + expect(mockedAxios.get).toHaveBeenCalledWith( + `${baseEndpoint}/v1/semantics-preparations/${mockInput.hash}/status`, + { + params: { + project_id: mockInput.projectId.toString(), + }, + }, + ); + }); + }); + + describe('ask', () => { + const mockInput: AskInput = { + query: 'Show active records', + deployId: 'deploy-hash', + projectId: 'project-123', + histories: [], + configurations: { + language: 'English', + }, + }; + + it('should send the project id to scope AI retrieval', async () => { + const mockQueryId = 'query-123'; + mockedAxios.post.mockResolvedValueOnce({ + data: { query_id: mockQueryId }, + }); + + const result = await adaptor.ask(mockInput); + + expect(result).toEqual({ queryId: mockQueryId }); + expect(mockedAxios.post).toHaveBeenCalledWith(`${baseEndpoint}/v1/asks`, { + query: mockInput.query, + id: mockInput.deployId, + project_id: mockInput.projectId, + histories: [], + configurations: mockInput.configurations, + }); + }); + }); + describe('generateRecommendationQuestions', () => { const mockInput: RecommendationQuestionsInput = { manifest: sampleManifest, diff --git a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts index 2d0675633e..6b095f33c9 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenAIAdaptor.ts @@ -34,7 +34,7 @@ import { AskFeedbackResult, AskFeedbackStatus, } from '@server/models/adaptor'; -import { getLogger } from '@server/utils'; +import { getLogger, normalizeManifest } from '@server/utils'; import * as Errors from '@server/utils/error'; import { SqlPair } from '../repositories'; import { ThreadResponse } from '@server/repositories'; @@ -51,6 +51,10 @@ const getAIServiceError = (error: any) => { export interface IWrenAIAdaptor { deploy(deployData: DeployData): Promise; + getDeployStatus( + deployId: string, + projectId?: string | number, + ): Promise; delete(projectId: number): Promise; /** @@ -83,6 +87,19 @@ export interface IWrenAIAdaptor { queryId: string, ): Promise; + generateSemanticsDescription(input: { + manifest: any; + selectedModels: string[]; + userPrompt: string; + projectId: number; + }): Promise; + getSemanticsDescriptionResult(queryId: string): Promise; + generateRelationshipRecommendations(input: { + manifest: any; + projectId: number; + }): Promise; + getRelationshipRecommendationResult(queryId: string): Promise; + /** * Get text-based answer from SQL */ @@ -108,6 +125,7 @@ export interface IWrenAIAdaptor { deploySqlPair( projectId: number, sqlPair: { question: string; sql: string }, + mdlHash?: string, ): Promise; getSqlPairResult(queryId: string): Promise; deleteSqlPairs(projectId: number, sqlPairIds: number[]): Promise; @@ -119,6 +137,7 @@ export interface IWrenAIAdaptor { */ generateInstruction( input: GenerateInstructionInput[], + mdlHash?: string, ): Promise; getInstructionResult(queryId: string): Promise; deleteInstructions(ids: number[], projectId: number): Promise; @@ -165,9 +184,10 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async deploySqlPair( projectId: number, sqlPair: Partial, + mdlHash?: string, ): Promise { try { - const body = { + const body: any = { sql_pairs: [ { id: `${sqlPair.id}`, @@ -177,6 +197,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { ], project_id: projectId.toString(), }; + if (mdlHash) { + body['mdl_hash'] = mdlHash; + } return axios .post(`${this.wrenAIBaseEndpoint}/v1/sql-pairs`, body) @@ -233,15 +256,42 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { */ public async ask(input: AskInput): Promise { + const startedAt = Date.now(); try { - const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, { + const body: Record = { query: input.query, id: input.deployId, + project_id: input.projectId, histories: this.transformHistoryInput(input.histories), configurations: input.configurations, - }); + }; + if (input.ignoreSqlGenerationReasoning !== undefined) { + body['ignore_sql_generation_reasoning'] = + input.ignoreSqlGenerationReasoning; + } + if (input.enableColumnPruning !== undefined) { + body['enable_column_pruning'] = input.enableColumnPruning; + } + if (input.useDryPlan !== undefined) { + body['use_dry_plan'] = input.useDryPlan; + } + if (input.allowDryPlanFallback !== undefined) { + body['allow_dry_plan_fallback'] = input.allowDryPlanFallback; + } + + const res = await axios.post(`${this.wrenAIBaseEndpoint}/v1/asks`, body); + logger.info( + `Ask timing stage=ai_ask_request project_id=${ + input.projectId ?? '' + } query_id=${res.data.query_id} elapsed_ms=${Date.now() - startedAt}`, + ); return { queryId: res.data.query_id }; } catch (err: any) { + logger.info( + `Ask timing stage=ai_ask_request project_id=${ + input.projectId ?? '' + } elapsed_ms=${Date.now() - startedAt} status=failed`, + ); logger.debug(`Got error when asking wren AI: ${getAIServiceError(err)}`); throw err; } @@ -249,11 +299,22 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async cancelAsk(queryId: string): Promise { // make PATCH request /v1/asks/:query_id to cancel the query + const startedAt = Date.now(); try { await axios.patch(`${this.wrenAIBaseEndpoint}/v1/asks/${queryId}`, { status: 'stopped', }); + logger.info( + `Ask timing stage=cancel_request query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + }`, + ); } catch (err: any) { + logger.info( + `Ask timing stage=cancel_request query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + } status=failed`, + ); logger.debug(`Got error when canceling ask: ${getAIServiceError(err)}`); throw err; } @@ -333,20 +394,21 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } public async deploy(deployData: DeployData): Promise { - const { manifest, hash } = deployData; + const { manifest, hash, projectId } = deployData; try { const res = await axios.post( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations`, { - mdl: JSON.stringify(manifest), + mdl: JSON.stringify(normalizeManifest(manifest)), id: hash, + project_id: projectId.toString(), }, ); const deployId = res.data.id; logger.debug( `Wren AI: Deploying wren AI, hash: ${hash}, deployId: ${deployId}`, ); - const deploySuccess = await this.waitDeployFinished(deployId); + const deploySuccess = await this.waitDeployFinished(deployId, projectId); if (deploySuccess) { logger.debug(`Wren AI: Deploy wren AI success, hash: ${hash}`); return { status: WrenAIDeployStatusEnum.SUCCESS }; @@ -371,8 +433,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { input: RecommendationQuestionsInput, ): Promise { const body = { - mdl: JSON.stringify(input.manifest), + mdl: JSON.stringify(normalizeManifest(input.manifest)), previous_questions: input.previousQuestions, + project_id: input.projectId, max_questions: input.maxQuestions, max_categories: input.maxCategories, configuration: input.configuration, @@ -411,9 +474,84 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } } + public async generateSemanticsDescription(input: { + manifest: any; + selectedModels: string[]; + userPrompt: string; + projectId: number; + }): Promise { + try { + const res = await axios.post( + `${this.wrenAIBaseEndpoint}/v1/semantics-descriptions`, + { + mdl: JSON.stringify(normalizeManifest(input.manifest)), + selected_models: input.selectedModels, + user_prompt: input.userPrompt, + project_id: String(input.projectId), + }, + ); + return { queryId: res.data.id }; + } catch (err: any) { + logger.debug( + `Got error when generating semantics descriptions: ${getAIServiceError(err)}`, + ); + throw err; + } + } + + public async getSemanticsDescriptionResult(queryId: string): Promise { + try { + const res = await axios.get( + `${this.wrenAIBaseEndpoint}/v1/semantics-descriptions/${queryId}`, + ); + return res.data; + } catch (err: any) { + logger.debug( + `Got error when getting semantics descriptions: ${getAIServiceError(err)}`, + ); + throw err; + } + } + + public async generateRelationshipRecommendations(input: { + manifest: any; + projectId: number; + }): Promise { + try { + const res = await axios.post( + `${this.wrenAIBaseEndpoint}/v1/relationship-recommendations`, + { + mdl: JSON.stringify(normalizeManifest(input.manifest)), + project_id: String(input.projectId), + }, + ); + return { queryId: res.data.id }; + } catch (err: any) { + logger.debug( + `Got error when generating relationship recommendations: ${getAIServiceError(err)}`, + ); + throw err; + } + } + + public async getRelationshipRecommendationResult(queryId: string): Promise { + try { + const res = await axios.get( + `${this.wrenAIBaseEndpoint}/v1/relationship-recommendations/${queryId}`, + ); + return res.data; + } catch (err: any) { + logger.debug( + `Got error when getting relationship recommendations: ${getAIServiceError(err)}`, + ); + throw err; + } + } + public async createTextBasedAnswer( input: TextBasedAnswerInput, ): Promise { + const startedAt = Date.now(); const body = { query: input.query, sql: input.sql, @@ -428,8 +566,18 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { `${this.wrenAIBaseEndpoint}/v1/sql-answers`, body, ); + logger.info( + `Ask timing stage=answer_formatting_request thread_id=${ + input.threadId ?? '' + } query_id=${res.data.query_id} elapsed_ms=${Date.now() - startedAt}`, + ); return { queryId: res.data.query_id }; } catch (err: any) { + logger.info( + `Ask timing stage=answer_formatting_request thread_id=${ + input.threadId ?? '' + } elapsed_ms=${Date.now() - startedAt} status=failed`, + ); logger.debug( `Got error when creating text-based answer: ${getAIServiceError(err)}`, ); @@ -441,12 +589,23 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { queryId: string, ): Promise { // make GET request /v1/sql-answers/:query_id to get the result + const startedAt = Date.now(); try { const res = await axios.get( `${this.wrenAIBaseEndpoint}/v1/sql-answers/${queryId}`, ); + logger.info( + `Ask timing stage=answer_formatting_poll query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + } status=${res.data.status}`, + ); return this.transformTextBasedAnswerResult(res.data); } catch (err: any) { + logger.info( + `Ask timing stage=answer_formatting_poll query_id=${queryId} elapsed_ms=${ + Date.now() - startedAt + } status=failed`, + ); logger.debug( `Got error when getting text-based answer result: ${getAIServiceError(err)}`, ); @@ -577,8 +736,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { public async generateInstruction( input: GenerateInstructionInput[], + mdlHash?: string, ): Promise { - const body = { + const body: any = { instructions: input.map((item) => ({ id: item.id.toString(), instruction: item.instruction, @@ -587,6 +747,9 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { })), project_id: input[0]?.projectId.toString(), }; + if (mdlHash) { + body['mdl_hash'] = mdlHash; + } try { const res = await axios.post( `${this.wrenAIBaseEndpoint}/v1/instructions`, @@ -730,10 +893,12 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } private transformChartAdjustmentInput(input: ChartAdjustmentInput) { - const { query, sql, adjustmentOption, chartSchema, configurations } = input; + const { query, sql, data, adjustmentOption, chartSchema, configurations } = + input; return { query, sql, + data, adjustment_option: { chart_type: adjustmentOption.chartType.toLowerCase(), x_axis: adjustmentOption.xAxis, @@ -769,13 +934,20 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { }; } - private async waitDeployFinished(deployId: string): Promise { + private async waitDeployFinished( + deployId: string, + projectId?: string | number, + ): Promise { let deploySuccess = false; - // timeout after 30 seconds - for (let waitTime = 1; waitTime <= 7; waitTime++) { + const maxAttempts = 90; + const pollingIntervalMs = 2000; + + for (let attempt = 1; attempt <= maxAttempts; attempt++) { try { - const status = await this.getDeployStatus(deployId); - logger.debug(`Wren AI: Deploy status: ${status}`); + const status = await this.getDeployStatus(deployId, projectId); + logger.debug( + `Wren AI: Deploy status: ${status}, attempt: ${attempt}/${maxAttempts}`, + ); if (status === WrenAISystemStatus.FINISHED) { deploySuccess = true; break; @@ -790,19 +962,28 @@ export class WrenAIAdaptor implements IWrenAIAdaptor { } catch (err: any) { throw err; } - await new Promise((resolve) => setTimeout(resolve, waitTime * 1000)); + await new Promise((resolve) => setTimeout(resolve, pollingIntervalMs)); } return deploySuccess; } - private async getDeployStatus(deployId: string): Promise { + public async getDeployStatus( + deployId: string, + projectId?: string | number, + ): Promise { try { const res = await axios.get( `${this.wrenAIBaseEndpoint}/v1/semantics-preparations/${deployId}/status`, + { + params: projectId ? { project_id: projectId.toString() } : undefined, + }, ); if (res.data.error) { - // passing AI response error string to catch block - throw new Error(res.data.error); + const error = + typeof res.data.error === 'string' + ? res.data.error + : res.data.error.message || JSON.stringify(res.data.error); + throw new Error(error); } return res.data?.status.toUpperCase() as WrenAISystemStatus; } catch (err: any) { diff --git a/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts b/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts index f849cffb88..731d0a93a0 100644 --- a/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts +++ b/wren-ui/src/apollo/server/adaptors/wrenEngineAdaptor.ts @@ -1,6 +1,6 @@ import axios, { AxiosResponse } from 'axios'; import { Manifest } from '../mdl/type'; -import { getLogger } from '@server/utils'; +import { getLogger, normalizeManifest } from '@server/utils'; import * as Errors from '@server/utils/error'; import { CompactTable, DEFAULT_PREVIEW_LIMIT } from '../services'; @@ -268,7 +268,7 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { try { const props = { modelingOnly: options?.modelingOnly ? true : false, - manifest: options?.manifest, + manifest: normalizeManifest(options?.manifest), }; const url = new URL(this.dryPlanUrlPath, this.wrenEngineBaseEndpoint); @@ -302,7 +302,7 @@ export class WrenEngineAdaptor implements IWrenEngineAdaptor { const { manifest } = options; const body = { sql, - manifest, + manifest: normalizeManifest(manifest), }; logger.debug( `Dry run wren engine with body: ${JSON.stringify(sql, null, 2)}`, diff --git a/wren-ui/src/apollo/server/backgrounds/chart.ts b/wren-ui/src/apollo/server/backgrounds/chart.ts index 8034d67c8c..d4a2600b58 100644 --- a/wren-ui/src/apollo/server/backgrounds/chart.ts +++ b/wren-ui/src/apollo/server/backgrounds/chart.ts @@ -22,13 +22,19 @@ const isFinalized = (status: ChartStatus) => { ); }; +const MIN_POLL_DELAY = 1000; +const MAX_POLL_DELAY = 10000; + export class ChartBackgroundTracker { private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private threadResponseRepository: IThreadResponseRepository; private runningJobs = new Set(); private telemetry: PostHogTelemetry; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -47,8 +53,11 @@ export class ChartBackgroundTracker { } private start() { + if (this.intervalId) { + return; + } logger.info('Chart background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map( (threadResponse) => async () => { // check if same job is running @@ -56,70 +65,83 @@ export class ChartBackgroundTracker { return; } + if (Date.now() < (this.nextPollAt[threadResponse.id] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(threadResponse.id); - // get the chart detail - const chartDetail = threadResponse.chartDetail; + try { + // get the chart detail + const chartDetail = threadResponse.chartDetail; - // get the latest result from AI service - const result = await this.wrenAIAdaptor.getChartResult( - chartDetail.queryId, - ); + // get the latest result from AI service + const result = await this.wrenAIAdaptor.getChartResult( + chartDetail.queryId, + ); - // check if status change - if (chartDetail.status === result.status) { - // mark the job as finished - logger.debug( - `Job ${threadResponse.id} chart status not changed, finished`, + const statusChanged = chartDetail.status !== result.status; + this.scheduleNextPoll( + threadResponse.id, + result.status, + statusChanged, ); - this.runningJobs.delete(threadResponse.id); - return; - } - // update database - const updatedChartDetail = { - queryId: chartDetail.queryId, - status: result?.status, - error: result?.error, - description: result?.response?.reasoning, - chartType: result?.response?.chartType?.toUpperCase() || null, - chartSchema: result?.response?.chartSchema, - }; - logger.debug( - `Job ${threadResponse.id} chart status changed, updating`, - ); - await this.threadResponseRepository.updateOne(threadResponse.id, { - chartDetail: updatedChartDetail, - }); - - // remove the task from tracker if it is finalized - if (isFinalized(result.status)) { - const eventProperties = { - question: threadResponse.question, - error: result.error, - }; - if (result.status === ChartStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_CHART, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_CHART, - eventProperties, - WrenService.AI, - false, + if (isFinalized(result.status) && !statusChanged) { + this.finalizeTask(threadResponse, result); + return; + } + + // check if status change + if (!statusChanged) { + // mark the job as finished + logger.debug( + `Job ${threadResponse.id} chart status not changed, finished`, ); + return; } + + // update database + const updatedChartDetail = { + queryId: chartDetail.queryId, + status: result?.status, + error: result?.error, + description: result?.response?.reasoning, + chartType: result?.response?.chartType?.toUpperCase() || null, + chartSchema: result?.response?.chartSchema, + }; logger.debug( - `Job ${threadResponse.id} chart is finalized, removing`, + `Job ${threadResponse.id} chart status changed, updating`, ); - delete this.tasks[threadResponse.id]; - } + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: updatedChartDetail, + }); + threadResponse.chartDetail = updatedChartDetail; - // mark the job as finished - this.runningJobs.delete(threadResponse.id); + // remove the task from tracker if it is finalized + if (isFinalized(result.status)) { + this.finalizeTask(threadResponse, result); + } + } catch (error) { + logger.error(`Chart job ${threadResponse.id} failed: ${error}`); + const failedChartDetail = { + ...threadResponse.chartDetail, + status: ChartStatus.FAILED, + error: error?.extensions || error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: failedChartDetail, + }); + threadResponse.chartDetail = failedChartDetail; + this.finalizeTask(threadResponse, { + status: ChartStatus.FAILED, + error, + }); + throw error; + } finally { + this.runningJobs.delete(threadResponse.id); + } }, ); @@ -135,22 +157,79 @@ export class ChartBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(threadResponse: ThreadResponse) { this.tasks[threadResponse.id] = threadResponse; + this.nextPollAt[threadResponse.id] = Date.now(); + this.pollDelay[threadResponse.id] = MIN_POLL_DELAY; } public getTasks() { return this.tasks; } + + private scheduleNextPoll( + taskId: number, + status: ChartStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskId] = Number.MAX_SAFE_INTEGER; + return; + } + + const baseDelay = status === ChartStatus.FETCHING ? MIN_POLL_DELAY : 3000; + this.pollDelay[taskId] = resultChanged + ? baseDelay + : Math.min( + Math.max((this.pollDelay[taskId] || baseDelay) * 1.5, baseDelay), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskId] = Date.now() + this.pollDelay[taskId]; + } + + private finalizeTask(threadResponse: ThreadResponse, result) { + const eventProperties = { + question: threadResponse.question, + error: result.error, + }; + if (result.status === ChartStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_CHART, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_CHART, + eventProperties, + WrenService.AI, + false, + ); + } + logger.debug(`Job ${threadResponse.id} chart is finalized, removing`); + delete this.tasks[threadResponse.id]; + delete this.nextPollAt[threadResponse.id]; + delete this.pollDelay[threadResponse.id]; + } } export class ChartAdjustmentBackgroundTracker { private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private threadResponseRepository: IThreadResponseRepository; private runningJobs = new Set(); private telemetry: PostHogTelemetry; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -169,8 +248,11 @@ export class ChartAdjustmentBackgroundTracker { } private start() { + if (this.intervalId) { + return; + } logger.info('Chart adjustment background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map( (threadResponse) => async () => { // check if same job is running @@ -178,71 +260,87 @@ export class ChartAdjustmentBackgroundTracker { return; } + if (Date.now() < (this.nextPollAt[threadResponse.id] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(threadResponse.id); - // get the chart detail - const chartDetail = threadResponse.chartDetail; + try { + // get the chart detail + const chartDetail = threadResponse.chartDetail; - // get the latest result from AI service - const result = await this.wrenAIAdaptor.getChartAdjustmentResult( - chartDetail.queryId, - ); + // get the latest result from AI service + const result = await this.wrenAIAdaptor.getChartAdjustmentResult( + chartDetail.queryId, + ); - // check if status change - if (chartDetail.status === result.status) { - // mark the job as finished - logger.debug( - `Job ${threadResponse.id} chart status not changed, finished`, + const statusChanged = chartDetail.status !== result.status; + this.scheduleNextPoll( + threadResponse.id, + result.status, + statusChanged, ); - this.runningJobs.delete(threadResponse.id); - return; - } - // update database - const updatedChartDetail = { - queryId: chartDetail.queryId, - status: result?.status, - error: result?.error, - description: result?.response?.reasoning, - chartType: result?.response?.chartType?.toUpperCase() || null, - chartSchema: result?.response?.chartSchema, - adjustment: true, - }; - logger.debug( - `Job ${threadResponse.id} chart status changed, updating`, - ); - await this.threadResponseRepository.updateOne(threadResponse.id, { - chartDetail: updatedChartDetail, - }); - - // remove the task from tracker if it is finalized - if (isFinalized(result.status)) { - const eventProperties = { - question: threadResponse.question, - error: result.error, - }; - if (result.status === ChartStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_ADJUST_CHART, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_ANSWER_ADJUST_CHART, - eventProperties, - WrenService.AI, - false, + if (isFinalized(result.status) && !statusChanged) { + this.finalizeTask(threadResponse, result); + return; + } + + // check if status change + if (!statusChanged) { + // mark the job as finished + logger.debug( + `Job ${threadResponse.id} chart status not changed, finished`, ); + return; } + + // update database + const updatedChartDetail = { + queryId: chartDetail.queryId, + status: result?.status, + error: result?.error, + description: result?.response?.reasoning, + chartType: result?.response?.chartType?.toUpperCase() || null, + chartSchema: result?.response?.chartSchema, + adjustment: true, + }; logger.debug( - `Job ${threadResponse.id} chart is finalized, removing`, + `Job ${threadResponse.id} chart status changed, updating`, ); - delete this.tasks[threadResponse.id]; - } + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: updatedChartDetail, + }); + threadResponse.chartDetail = updatedChartDetail; - // mark the job as finished - this.runningJobs.delete(threadResponse.id); + // remove the task from tracker if it is finalized + if (isFinalized(result.status)) { + this.finalizeTask(threadResponse, result); + } + } catch (error) { + logger.error( + `Chart adjustment job ${threadResponse.id} failed: ${error}`, + ); + const failedChartDetail = { + ...threadResponse.chartDetail, + status: ChartStatus.FAILED, + error: error?.extensions || error, + adjustment: true, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + chartDetail: failedChartDetail, + }); + threadResponse.chartDetail = failedChartDetail; + this.finalizeTask(threadResponse, { + status: ChartStatus.FAILED, + error, + }); + throw error; + } finally { + this.runningJobs.delete(threadResponse.id); + } }, ); @@ -258,11 +356,65 @@ export class ChartAdjustmentBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(threadResponse: ThreadResponse) { this.tasks[threadResponse.id] = threadResponse; + this.nextPollAt[threadResponse.id] = Date.now(); + this.pollDelay[threadResponse.id] = MIN_POLL_DELAY; } public getTasks() { return this.tasks; } + + private scheduleNextPoll( + taskId: number, + status: ChartStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskId] = Number.MAX_SAFE_INTEGER; + return; + } + + const baseDelay = status === ChartStatus.FETCHING ? MIN_POLL_DELAY : 3000; + this.pollDelay[taskId] = resultChanged + ? baseDelay + : Math.min( + Math.max((this.pollDelay[taskId] || baseDelay) * 1.5, baseDelay), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskId] = Date.now() + this.pollDelay[taskId]; + } + + private finalizeTask(threadResponse: ThreadResponse, result) { + const eventProperties = { + question: threadResponse.question, + error: result.error, + }; + if (result.status === ChartStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_ADJUST_CHART, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_ANSWER_ADJUST_CHART, + eventProperties, + WrenService.AI, + false, + ); + } + logger.debug(`Job ${threadResponse.id} chart is finalized, removing`); + delete this.tasks[threadResponse.id]; + delete this.nextPollAt[threadResponse.id]; + delete this.pollDelay[threadResponse.id]; + } } diff --git a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts index baa496b21f..472897ec07 100644 --- a/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/dashboardCacheBackgroundTracker.ts @@ -24,7 +24,8 @@ export class DashboardCacheBackgroundTracker { private projectService: IProjectService; private deployService: IDeployService; private queryService: IQueryService; - private runningJobs = new Set(); + private runningJobs = new Set(); + private intervalId?: NodeJS.Timeout; constructor({ dashboardRepository, @@ -52,12 +53,23 @@ export class DashboardCacheBackgroundTracker { } private start(): void { + if (this.intervalId) { + return; + } logger.info('Dashboard cache background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { this.checkAndRefreshCaches(); }, this.intervalTime); } + public stop(): void { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + private async checkAndRefreshCaches(): Promise { try { // Get all dashboards with cache enabled @@ -102,7 +114,9 @@ export class DashboardCacheBackgroundTracker { }); // Get project and deployment info - const project = await this.projectService.getCurrentProject(); + const project = await this.projectService.getProjectById( + dashboard.projectId, + ); const deployment = await this.deployService.getLastDeployment(project.id); const mdl = deployment.manifest; const hash = uuidv4(); diff --git a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts index da33377deb..6575fa83c8 100644 --- a/wren-ui/src/apollo/server/backgrounds/recommend-question.ts +++ b/wren-ui/src/apollo/server/backgrounds/recommend-question.ts @@ -20,15 +20,22 @@ const isFinalized = (status: RecommendationQuestionStatus) => { ].includes(status); }; +const MIN_POLL_DELAY = 1000; +const MAX_POLL_DELAY = 10000; + export class ProjectRecommendQuestionBackgroundTracker { // tasks is a kv pair of task id and thread response private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private projectRepository: IProjectRepository; private runningJobs = new Set(); private telemetry: ITelemetry; private logger: Logger; + private initialized = false; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -49,14 +56,21 @@ export class ProjectRecommendQuestionBackgroundTracker { } public start() { + if (this.intervalId) { + return; + } this.logger.info('Recommend question background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map((project) => async () => { // check if same job is running if (this.runningJobs.has(this.taskKey(project))) { return; } + if (Date.now() < (this.nextPollAt[this.taskKey(project)] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(this.taskKey(project)); @@ -67,11 +81,20 @@ export class ProjectRecommendQuestionBackgroundTracker { project.queryId, ); + const changed = + project.questionsStatus !== result.status || + result.response?.questions.length !== + (project.questions || []).length; + this.scheduleNextPoll(this.taskKey(project), result.status, changed); + + if (isFinalized(result.status) && !changed) { + this.finalizeTask(project, result); + this.runningJobs.delete(this.taskKey(project)); + return; + } + // check if status change - if ( - project.questionsStatus === result.status && - result.response?.questions.length === (project.questions || []).length - ) { + if (!changed) { // mark the job as finished this.logger.debug( `${loggerPrefix}job ${this.taskKey(project)} status not changed, returning question count: ${result.response?.questions.length || 0}`, @@ -81,10 +104,7 @@ export class ProjectRecommendQuestionBackgroundTracker { } // update database - if ( - result.status !== project.questionsStatus || - result.response?.questions.length !== (project.questions || []).length - ) { + if (changed) { this.logger.debug( `${loggerPrefix}job ${this.taskKey(project)} have changes, returning question count: ${result.response?.questions.length || 0}, updating`, ); @@ -99,30 +119,7 @@ export class ProjectRecommendQuestionBackgroundTracker { // remove the task from tracker if it is finalized if (isFinalized(result.status)) { - const eventProperties = { - projectId: project.id, - projectType: project.type, - status: result.status, - questions: project.questions, - error: result.error, - }; - if (result.status === RecommendationQuestionStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, - eventProperties, - WrenService.AI, - false, - ); - } - this.logger.debug( - `${loggerPrefix}job ${this.taskKey(project)} is finalized, removing`, - ); - delete this.tasks[this.taskKey(project)]; + this.finalizeTask(project, result); } // mark the job as finished @@ -141,8 +138,18 @@ export class ProjectRecommendQuestionBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(project: Project) { this.tasks[this.taskKey(project)] = project; + this.nextPollAt[this.taskKey(project)] = Date.now(); + this.pollDelay[this.taskKey(project)] = MIN_POLL_DELAY; } public getTasks() { @@ -150,6 +157,10 @@ export class ProjectRecommendQuestionBackgroundTracker { } public async initialize() { + if (this.initialized) { + return; + } + const projects = await this.projectRepository.findAll(); for (const project of projects) { if ( @@ -159,6 +170,8 @@ export class ProjectRecommendQuestionBackgroundTracker { this.addTask(project); } } + + this.initialized = true; } public taskKey(project: Project) { @@ -168,17 +181,71 @@ export class ProjectRecommendQuestionBackgroundTracker { public isExist(project: Project) { return this.tasks[this.taskKey(project)]; } + + private scheduleNextPoll( + taskKey: number, + status: RecommendationQuestionStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskKey] = Number.MAX_SAFE_INTEGER; + return; + } + + this.pollDelay[taskKey] = resultChanged + ? MIN_POLL_DELAY + : Math.min( + Math.max( + (this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, + MIN_POLL_DELAY, + ), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskKey] = Date.now() + this.pollDelay[taskKey]; + } + + private finalizeTask(project: Project, result) { + const eventProperties = { + projectId: project.id, + projectType: project.type, + status: result.status, + questions: project.questions, + error: result.error, + }; + if (result.status === RecommendationQuestionStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_PROJECT_RECOMMENDATION_QUESTIONS, + eventProperties, + WrenService.AI, + false, + ); + } + const taskKey = this.taskKey(project); + this.logger.debug(`${loggerPrefix}job ${taskKey} is finalized, removing`); + delete this.tasks[taskKey]; + delete this.nextPollAt[taskKey]; + delete this.pollDelay[taskKey]; + } } export class ThreadRecommendQuestionBackgroundTracker { // tasks is a kv pair of task id and thread response private tasks: Record = {}; + private nextPollAt: Record = {}; + private pollDelay: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; private threadRepository: IThreadRepository; private runningJobs = new Set(); private telemetry: ITelemetry; private logger: Logger; + private initialized = false; + private intervalId?: NodeJS.Timeout; constructor({ telemetry, @@ -199,14 +266,21 @@ export class ThreadRecommendQuestionBackgroundTracker { } public start() { + if (this.intervalId) { + return; + } this.logger.info('Recommend question background tracker started'); - setInterval(() => { + this.intervalId = setInterval(() => { const jobs = Object.values(this.tasks).map((thread) => async () => { // check if same job is running if (this.runningJobs.has(this.taskKey(thread))) { return; } + if (Date.now() < (this.nextPollAt[this.taskKey(thread)] || 0)) { + return; + } + // mark the job as running this.runningJobs.add(this.taskKey(thread)); @@ -217,11 +291,19 @@ export class ThreadRecommendQuestionBackgroundTracker { thread.queryId, ); + const changed = + thread.questionsStatus !== result.status || + result.response?.questions.length !== (thread.questions || []).length; + this.scheduleNextPoll(this.taskKey(thread), result.status, changed); + + if (isFinalized(result.status) && !changed) { + this.finalizeTask(thread, result); + this.runningJobs.delete(this.taskKey(thread)); + return; + } + // check if status change - if ( - thread.questionsStatus === result.status && - result.response?.questions.length === (thread.questions || []).length - ) { + if (!changed) { // mark the job as finished this.logger.debug( `${loggerPrefix}job ${this.taskKey(thread)} status not changed, returning question count: ${result.response?.questions.length || 0}`, @@ -231,10 +313,7 @@ export class ThreadRecommendQuestionBackgroundTracker { } // update database - if ( - result.status !== thread.questionsStatus || - result.response?.questions.length !== (thread.questions || []).length - ) { + if (changed) { this.logger.debug( `${loggerPrefix}job ${this.taskKey(thread)} have changes, returning question count: ${result.response?.questions.length || 0}, updating`, ); @@ -249,29 +328,7 @@ export class ThreadRecommendQuestionBackgroundTracker { // remove the task from tracker if it is finalized if (isFinalized(result.status)) { - const eventProperties = { - thread_id: thread.id, - status: result.status, - questions: thread.questions, - error: result.error, - }; - if (result.status === RecommendationQuestionStatus.FINISHED) { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, - eventProperties, - ); - } else { - this.telemetry.sendEvent( - TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, - eventProperties, - WrenService.AI, - false, - ); - } - this.logger.debug( - `${loggerPrefix}job ${this.taskKey(thread)} is finalized, removing`, - ); - delete this.tasks[this.taskKey(thread)]; + this.finalizeTask(thread, result); } // mark the job as finished @@ -290,8 +347,18 @@ export class ThreadRecommendQuestionBackgroundTracker { }, this.intervalTime); } + public stop() { + if (!this.intervalId) { + return; + } + clearInterval(this.intervalId); + this.intervalId = undefined; + } + public addTask(thread: Thread) { this.tasks[this.taskKey(thread)] = thread; + this.nextPollAt[this.taskKey(thread)] = Date.now(); + this.pollDelay[this.taskKey(thread)] = MIN_POLL_DELAY; } public getTasks() { @@ -299,6 +366,10 @@ export class ThreadRecommendQuestionBackgroundTracker { } public async initialize() { + if (this.initialized) { + return; + } + const threads = await this.threadRepository.findAll(); for (const thread of threads) { if ( @@ -309,6 +380,8 @@ export class ThreadRecommendQuestionBackgroundTracker { this.addTask(thread); } } + + this.initialized = true; } public taskKey(thread: Thread) { @@ -318,4 +391,53 @@ export class ThreadRecommendQuestionBackgroundTracker { public isExist(thread: Thread) { return this.tasks[this.taskKey(thread)]; } + + private scheduleNextPoll( + taskKey: number, + status: RecommendationQuestionStatus, + resultChanged: boolean, + ) { + if (isFinalized(status)) { + this.nextPollAt[taskKey] = Number.MAX_SAFE_INTEGER; + return; + } + + this.pollDelay[taskKey] = resultChanged + ? MIN_POLL_DELAY + : Math.min( + Math.max( + (this.pollDelay[taskKey] || MIN_POLL_DELAY) * 1.5, + MIN_POLL_DELAY, + ), + MAX_POLL_DELAY, + ); + this.nextPollAt[taskKey] = Date.now() + this.pollDelay[taskKey]; + } + + private finalizeTask(thread: Thread, result) { + const eventProperties = { + thread_id: thread.id, + status: result.status, + questions: thread.questions, + error: result.error, + }; + if (result.status === RecommendationQuestionStatus.FINISHED) { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, + eventProperties, + ); + } else { + this.telemetry.sendEvent( + TelemetryEvent.HOME_GENERATE_THREAD_RECOMMENDATION_QUESTIONS, + eventProperties, + WrenService.AI, + false, + ); + } + const taskKey = this.taskKey(thread); + this.logger.debug(`${loggerPrefix}job ${taskKey} is finalized, removing`); + delete this.tasks[taskKey]; + delete this.nextPollAt[taskKey]; + delete this.pollDelay[taskKey]; + } } diff --git a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts index 91f14bcab1..80f32ce7f2 100644 --- a/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts +++ b/wren-ui/src/apollo/server/backgrounds/textBasedAnswerBackgroundTracker.ts @@ -4,7 +4,11 @@ import { TextBasedAnswerResult, TextBasedAnswerStatus, } from '../models/adaptor'; -import { ThreadResponse, IThreadResponseRepository } from '../repositories'; +import { + ThreadResponse, + IThreadRepository, + IThreadResponseRepository, +} from '../repositories'; import { IProjectService, IDeployService, @@ -22,6 +26,7 @@ export class TextBasedAnswerBackgroundTracker { private tasks: Record = {}; private intervalTime: number; private wrenAIAdaptor: IWrenAIAdaptor; + private threadRepository: IThreadRepository; private threadResponseRepository: IThreadResponseRepository; private projectService: IProjectService; private deployService: IDeployService; @@ -30,18 +35,21 @@ export class TextBasedAnswerBackgroundTracker { constructor({ wrenAIAdaptor, + threadRepository, threadResponseRepository, projectService, deployService, queryService, }: { wrenAIAdaptor: IWrenAIAdaptor; + threadRepository: IThreadRepository; threadResponseRepository: IThreadResponseRepository; projectService: IProjectService; deployService: IDeployService; queryService: IQueryService; }) { this.wrenAIAdaptor = wrenAIAdaptor; + this.threadRepository = threadRepository; this.threadResponseRepository = threadResponseRepository; this.projectService = projectService; this.deployService = deployService; @@ -62,88 +70,149 @@ export class TextBasedAnswerBackgroundTracker { } this.runningJobs.add(threadResponse.id); - // update the status to fetching data - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: { - ...threadResponse.answerDetail, - status: ThreadResponseAnswerStatus.FETCHING_DATA, - }, - }); - - // get sql data - const project = await this.projectService.getCurrentProject(); - const deployment = await this.deployService.getLastDeployment( - project.id, - ); - const mdl = deployment.manifest; - let data: PreviewDataResponse; try { - data = (await this.queryService.preview(threadResponse.sql, { - project, - manifest: mdl, - modelingOnly: false, - limit: 500, - })) as PreviewDataResponse; - } catch (error) { - logger.error(`Error when query sql data: ${error}`); - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: { - ...threadResponse.answerDetail, - status: ThreadResponseAnswerStatus.FAILED, - error: error?.extensions || error, - }, - }); - throw error; - } + const answerDetail = threadResponse.answerDetail; - // request AI service - const response = await this.wrenAIAdaptor.createTextBasedAnswer({ - query: threadResponse.question, - sql: threadResponse.sql, - sqlData: data, - threadId: threadResponse.threadId.toString(), - configurations: { - language: WrenAILanguage[project.language] || WrenAILanguage.EN, - }, - }); + if ( + !answerDetail.queryId && + answerDetail.status !== ThreadResponseAnswerStatus.FETCHING_DATA + ) { + const fetchingDetail = { + ...answerDetail, + status: ThreadResponseAnswerStatus.FETCHING_DATA, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: fetchingDetail, + }); + threadResponse.answerDetail = fetchingDetail; - // update the status to preprocessing - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: { - ...threadResponse.answerDetail, - status: ThreadResponseAnswerStatus.PREPROCESSING, - }, - }); + const thread = await this.threadRepository.findOneBy({ + id: threadResponse.threadId, + }); + if (!thread) { + throw new Error(`Thread ${threadResponse.threadId} not found`); + } + const project = await this.projectService.getProjectById( + thread.projectId, + ); + const deployment = await this.deployService.getLastDeployment( + project.id, + ); + const mdl = deployment.manifest; + let data: PreviewDataResponse; + try { + const executionStartedAt = Date.now(); + data = (await this.queryService.preview(threadResponse.sql, { + project, + manifest: mdl, + modelingOnly: false, + limit: 500, + cacheEnabled: false, + })) as PreviewDataResponse; + logger.info( + `Ask timing stage=sql_execution response_id=${ + threadResponse.id + } project_id=${project.id} elapsed_ms=${ + Date.now() - executionStartedAt + } row_count=${data?.data?.length ?? ''}`, + ); + } catch (error) { + logger.error(`Error when query sql data: ${error}`); + const failedDetail = { + ...threadResponse.answerDetail, + status: ThreadResponseAnswerStatus.FAILED, + error: error?.extensions || error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: failedDetail, + }); + threadResponse.answerDetail = failedDetail; + delete this.tasks[threadResponse.id]; + throw error; + } - // polling query id to check the status - let result: TextBasedAnswerResult; - do { - result = await this.wrenAIAdaptor.getTextBasedAnswerResult( - response.queryId, - ); - if (result.status === TextBasedAnswerStatus.PREPROCESSING) { - await new Promise((resolve) => setTimeout(resolve, 500)); + const answerRequestStartedAt = Date.now(); + const response = await this.wrenAIAdaptor.createTextBasedAnswer({ + query: threadResponse.question, + sql: threadResponse.sql, + sqlData: data, + threadId: threadResponse.threadId.toString(), + configurations: { + language: WrenAILanguage[project.language] || WrenAILanguage.EN, + }, + }); + logger.info( + `Ask timing stage=answer_formatting_request response_id=${ + threadResponse.id + } project_id=${project.id} query_id=${ + response.queryId + } elapsed_ms=${Date.now() - answerRequestStartedAt}`, + ); + + const preprocessingDetail = { + ...threadResponse.answerDetail, + queryId: response.queryId, + status: ThreadResponseAnswerStatus.PREPROCESSING, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: preprocessingDetail, + }); + threadResponse.answerDetail = preprocessingDetail; + return; } - } while (result.status === TextBasedAnswerStatus.PREPROCESSING); - // update the status to final - const updatedAnswerDetail = { - queryId: response.queryId, - status: - result.status === TextBasedAnswerStatus.SUCCEEDED - ? ThreadResponseAnswerStatus.STREAMING - : ThreadResponseAnswerStatus.FAILED, - numRowsUsedInLLM: result.numRowsUsedInLLM, - error: result.error, - }; - await this.threadResponseRepository.updateOne(threadResponse.id, { - answerDetail: updatedAnswerDetail, - }); + if ( + answerDetail.queryId && + answerDetail.status === ThreadResponseAnswerStatus.PREPROCESSING + ) { + const answerPollStartedAt = Date.now(); + const result: TextBasedAnswerResult = + await this.wrenAIAdaptor.getTextBasedAnswerResult( + answerDetail.queryId, + ); + logger.info( + `Ask timing stage=answer_formatting_poll response_id=${ + threadResponse.id + } query_id=${answerDetail.queryId} elapsed_ms=${ + Date.now() - answerPollStartedAt + } status=${result.status}`, + ); - delete this.tasks[threadResponse.id]; + if (result.status === TextBasedAnswerStatus.PREPROCESSING) { + return; + } - // Mark the job as finished - this.runningJobs.delete(threadResponse.id); + const updatedAnswerDetail = { + queryId: answerDetail.queryId, + status: + result.status === TextBasedAnswerStatus.SUCCEEDED + ? ThreadResponseAnswerStatus.STREAMING + : ThreadResponseAnswerStatus.FAILED, + numRowsUsedInLLM: result.numRowsUsedInLLM, + error: result.error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: updatedAnswerDetail, + }); + threadResponse.answerDetail = updatedAnswerDetail; + delete this.tasks[threadResponse.id]; + } + } catch (error) { + logger.error(`Answer job ${threadResponse.id} failed: ${error}`); + const failedDetail = { + ...threadResponse.answerDetail, + status: ThreadResponseAnswerStatus.FAILED, + error: error?.extensions || error, + }; + await this.threadResponseRepository.updateOne(threadResponse.id, { + answerDetail: failedDetail, + }); + threadResponse.answerDetail = failedDetail; + delete this.tasks[threadResponse.id]; + throw error; + } finally { + this.runningJobs.delete(threadResponse.id); + } }, ); diff --git a/wren-ui/src/apollo/server/config.ts b/wren-ui/src/apollo/server/config.ts index ad43e2da84..558bea67c0 100644 --- a/wren-ui/src/apollo/server/config.ts +++ b/wren-ui/src/apollo/server/config.ts @@ -9,6 +9,15 @@ export interface IConfig { // pg pgUrl?: string; debug?: boolean; + // mssql + mssqlUrl?: string; + mssqlHost?: string; + mssqlPort?: number; + mssqlDatabase?: string; + mssqlUser?: string; + mssqlPassword?: string; + mssqlEncrypt?: boolean; + mssqlTrustServerCertificate?: boolean; // sqlite sqliteFile?: string; @@ -53,12 +62,19 @@ const defaultConfig = { otherServiceUsingDocker: false, // database - dbType: 'sqlite', + dbType: 'mssql', // pg pgUrl: 'postgres://postgres:postgres@localhost:5432/admin_ui', debug: false, + // mssql + mssqlHost: 'localhost', + mssqlPort: 1433, + mssqlDatabase: 'wren_ui', + mssqlEncrypt: false, + mssqlTrustServerCertificate: true, + // sqlite sqliteFile: './db.sqlite3', @@ -88,6 +104,21 @@ const config = { // pg pgUrl: process.env.PG_URL, debug: process.env.DEBUG === 'true', + // mssql + mssqlUrl: process.env.MSSQL_URL, + mssqlHost: process.env.MSSQL_HOST, + mssqlPort: process.env.MSSQL_PORT + ? parseInt(process.env.MSSQL_PORT) + : undefined, + mssqlDatabase: process.env.MSSQL_DATABASE, + mssqlUser: process.env.MSSQL_USER, + mssqlPassword: process.env.MSSQL_PASSWORD, + mssqlEncrypt: process.env.MSSQL_ENCRYPT + ? process.env.MSSQL_ENCRYPT === 'true' + : undefined, + mssqlTrustServerCertificate: process.env.MSSQL_TRUST_SERVER_CERTIFICATE + ? process.env.MSSQL_TRUST_SERVER_CERTIFICATE === 'true' + : undefined, // sqlite sqliteFile: process.env.SQLITE_FILE, @@ -147,9 +178,12 @@ const config = { threadRecommendationQuestionsMaxQuestions: process.env .THREAD_RECOMMENDATION_QUESTIONS_MAX_QUESTIONS ? parseInt(process.env.THREAD_RECOMMENDATION_QUESTIONS_MAX_QUESTIONS) - : 1, + : 3, }; export function getConfig(): IConfig { - return { ...defaultConfig, ...pickBy(config) }; + return { + ...defaultConfig, + ...pickBy(config, (value) => value !== undefined && value !== null), + }; } diff --git a/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts new file mode 100644 index 0000000000..869f7b3299 --- /dev/null +++ b/wren-ui/src/apollo/server/managers/tests/dataSourceSchemaDetector.test.ts @@ -0,0 +1,169 @@ +import DataSourceSchemaDetector from '../dataSourceSchemaDetector'; + +describe('DataSourceSchemaDetector', () => { + const projectId = 1; + + const createContext = ({ + models = [], + columns = [], + latestTables = [], + lastSchemaChange = null, + }: { + models?: any[]; + columns?: any[]; + latestTables?: any[]; + lastSchemaChange?: any; + }) => + ({ + projectRepository: { + findOneBy: jest.fn().mockResolvedValue({ id: projectId }), + }, + projectService: { + getProjectDataSourceTables: jest.fn().mockResolvedValue(latestTables), + }, + schemaChangeRepository: { + findLastSchemaChange: jest.fn().mockResolvedValue(lastSchemaChange), + createOne: jest.fn(), + updateOne: jest.fn(), + }, + modelRepository: { + findAllBy: jest.fn().mockResolvedValue(models), + }, + modelColumnRepository: { + findColumnsByModelIds: jest.fn().mockResolvedValue(columns), + createOne: jest.fn().mockImplementation((data) => + Promise.resolve({ + id: 99, + ...data, + }), + ), + updateOne: jest.fn().mockImplementation((id, data) => + Promise.resolve({ + id, + modelId: 1, + sourceColumnName: 'amount', + ...data, + }), + ), + }, + modelNestedColumnRepository: { + createMany: jest.fn(), + deleteAllBy: jest.fn(), + }, + }) as any; + + it('syncs newly added datasource columns into existing models', async () => { + const model = { + id: 10, + projectId, + sourceTableName: 'orders', + }; + const existingColumn = { + id: 20, + modelId: 10, + isCalculated: false, + displayName: 'id', + referenceName: 'id', + sourceColumnName: 'id', + type: 'int', + notNull: true, + isPk: true, + properties: null, + }; + const ctx = createContext({ + models: [model], + columns: [existingColumn], + latestTables: [ + { + name: 'orders', + columns: [ + { name: 'id', type: 'int', notNull: true }, + { name: 'status', type: 'varchar', notNull: false }, + ], + }, + ], + }); + + const detector = new DataSourceSchemaDetector({ ctx, projectId }); + + await expect(detector.detectSchemaChange()).resolves.toBe(true); + expect(ctx.modelColumnRepository.createOne).toHaveBeenCalledWith({ + modelId: 10, + isCalculated: false, + displayName: 'status', + referenceName: 'status', + sourceColumnName: 'status', + type: 'varchar', + notNull: false, + isPk: false, + properties: null, + }); + expect(ctx.schemaChangeRepository.createOne).not.toHaveBeenCalled(); + }); + + it('updates existing column schema attributes without overwriting user properties', async () => { + const model = { + id: 10, + projectId, + sourceTableName: 'orders', + }; + const existingColumn = { + id: 20, + modelId: 10, + isCalculated: false, + displayName: 'Amount', + referenceName: 'amount', + sourceColumnName: 'amount', + type: 'int', + notNull: false, + isPk: false, + properties: JSON.stringify({ description: 'User description' }), + }; + const ctx = createContext({ + models: [model], + columns: [existingColumn], + latestTables: [ + { + name: 'orders', + columns: [ + { + name: 'amount', + type: 'decimal', + notNull: true, + properties: { description: 'Datasource description' }, + }, + ], + }, + ], + }); + + const detector = new DataSourceSchemaDetector({ ctx, projectId }); + + await expect(detector.detectSchemaChange()).resolves.toBe(true); + expect(ctx.modelColumnRepository.updateOne).toHaveBeenCalledWith(20, { + type: 'decimal', + notNull: true, + properties: JSON.stringify({ description: 'User description' }), + }); + expect(ctx.schemaChangeRepository.createOne).toHaveBeenCalledWith( + expect.objectContaining({ + projectId, + change: { + modifiedColumns: [ + { + name: 'orders', + columns: [ + { + name: 'amount', + type: 'decimal', + notNull: true, + properties: { description: 'Datasource description' }, + }, + ], + }, + ], + }, + }), + ); + }); +}); diff --git a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts index 4f7a87671a..77d5c0e41f 100644 --- a/wren-ui/src/apollo/server/mdl/mdlBuilder.ts +++ b/wren-ui/src/apollo/server/mdl/mdlBuilder.ts @@ -8,6 +8,7 @@ import { View, } from '../repositories'; import { + ColumnMDL, Manifest, ModelMDL, TableReference, @@ -16,9 +17,12 @@ import { import { getLogger } from '@server/utils'; import { getConfig } from '@server/config'; import { DataSourceName } from '../types'; +import { + getUniqueReferenceName, + replaceInvalidReferenceName, +} from '../utils/model'; const logger = getLogger('MDLBuilder'); -logger.level = 'debug'; const config = getConfig(); @@ -41,6 +45,13 @@ export interface IMDLBuilder { // responsible to generate a valid manifest json export class MDLBuilder implements IMDLBuilder { private manifest: Manifest; + private skippedDuplicateColumns = new Map>(); + private readonly columnNameAliases = new Map(); + private readonly manifestColumnNamesByModel = new Map>(); + private readonly manifestColumnNameBySourceByModel = new Map< + string, + Map + >(); private project: Project; private readonly models: Model[]; @@ -83,7 +94,40 @@ export class MDLBuilder implements IMDLBuilder { this.manifest = {}; } + private recordSkippedDuplicateColumn(modelName: string, columnName: string) { + const columns = + this.skippedDuplicateColumns.get(modelName) || new Set(); + columns.add(columnName); + this.skippedDuplicateColumns.set(modelName, columns); + } + + private logSkippedDuplicateColumns() { + if (this.skippedDuplicateColumns.size === 0) { + return; + } + + const duplicateCount = [...this.skippedDuplicateColumns.values()].reduce( + (count, columns) => count + columns.size, + 0, + ); + const examples = [...this.skippedDuplicateColumns.entries()] + .slice(0, 5) + .map( + ([modelName, columns]) => + `${modelName}: ${[...columns].slice(0, 5).join(', ')}`, + ) + .join('; '); + + logger.debug( + `Skipped ${duplicateCount} duplicated MDL columns across ${this.skippedDuplicateColumns.size} models. Examples: ${examples}`, + ); + } + public build(): Manifest { + this.skippedDuplicateColumns.clear(); + this.columnNameAliases.clear(); + this.manifestColumnNamesByModel.clear(); + this.manifestColumnNameBySourceByModel.clear(); this.addProject(); this.addModel(); this.addNormalField(); @@ -91,6 +135,7 @@ export class MDLBuilder implements IMDLBuilder { this.addCalculatedField(); this.addView(); this.postProcessManifest(); + this.logSkippedDuplicateColumns(); return this.getManifest(); } @@ -110,16 +155,12 @@ export class MDLBuilder implements IMDLBuilder { } const tableReference = this.buildTableReference(model); - return { + const modelMdl = { name: model.referenceName, columns: [], tableReference, // can only have one of refSql or tableReference - refSql: this.useRustWrenEngine() - ? null - : tableReference - ? null - : model.refSql, + refSql: tableReference ? null : model.refSql, cached: model.cached ? true : false, refreshTime: model.refreshTime, properties: { @@ -128,6 +169,23 @@ export class MDLBuilder implements IMDLBuilder { }, primaryKey: '', // will be modified in addColumn } as ModelMDL; + + if (tableReference && this.hasDuplicateSourceColumns(model.id)) { + const refSql = this.buildDedupedTableReferenceSql( + model.id, + modelMdl, + tableReference, + ); + if (refSql) { + logger.debug( + `Using deduped explicit projection for model "${model.referenceName}" because its source table contains duplicate column names.`, + ); + modelMdl.tableReference = null; + modelMdl.refSql = refSql; + } + } + + return modelMdl; }); } @@ -184,11 +242,6 @@ export class MDLBuilder implements IMDLBuilder { (model: any) => model.name === modelRefName, ); - // modify model primary key - if (column.isPk) { - model.primaryKey = column.referenceName; - } - // add column into model if (!model.columns) { model.columns = []; @@ -200,6 +253,9 @@ export class MDLBuilder implements IMDLBuilder { if (column.displayName) { properties.displayName = column.displayName; } + if (column.sourceColumnName) { + properties.sourceColumnName = column.sourceColumnName; + } // put nested columns in properties if (column.type.includes('STRUCT')) { const nestedColumns = this.nestedColumns.filter( @@ -216,9 +272,31 @@ export class MDLBuilder implements IMDLBuilder { } }, {}); } - const expression = this.getColumnExpression(column, model); + const sourceColumnName = column.sourceColumnName || column.referenceName; + const sourceColumnNames = this.getManifestSourceColumnNameMap(model); + const existingColumnName = sourceColumnNames.get( + sourceColumnName.toLowerCase(), + ); + if (existingColumnName) { + this.columnNameAliases.set(column.id, existingColumnName); + if (column.isPk) { + model.primaryKey = existingColumnName; + } + this.recordSkippedDuplicateColumn(modelRefName, sourceColumnName); + return; + } + + const columnName = this.getManifestColumnName(column, model); + sourceColumnNames.set(sourceColumnName.toLowerCase(), columnName); + + // modify model primary key + if (column.isPk) { + model.primaryKey = columnName; + } + + const expression = this.getColumnExpression(column, model, columnName); model.columns.push({ - name: column.referenceName, + name: columnName, type: column.type, isCalculated: column.isCalculated ? true : false, notNull: column.notNull ? true : false, @@ -250,9 +328,10 @@ export class MDLBuilder implements IMDLBuilder { ); return; } - const expression = this.getColumnExpression(column, model); + const columnName = this.getManifestColumnName(column, model); + const expression = this.getColumnExpression(column, model, columnName); const columnValue = { - name: column.referenceName, + name: columnName, type: column.type, isCalculated: true, expression, @@ -274,17 +353,14 @@ export class MDLBuilder implements IMDLBuilder { logger.debug(`Can not find model "${modelName}" to add calculated field`); return; } - // if calculated field is already in the model, skip - if ( - model.columns.find( - (column: any) => column.name === calculatedField.referenceName, - ) - ) { - return; - } - const expression = this.getColumnExpression(calculatedField, model); + const columnName = this.getManifestColumnName(calculatedField, model); + const expression = this.getColumnExpression( + calculatedField, + model, + columnName, + ); const columnValue = { - name: calculatedField.referenceName, + name: columnName, type: calculatedField.type, isCalculated: true, expression, @@ -302,18 +378,22 @@ export class MDLBuilder implements IMDLBuilder { joinType, fromModelName, fromColumnName, + fromColumnId, toModelName, toColumnName, + toColumnId, } = relation; const condition = this.getRelationCondition(relation); this.addRelationColumn(fromModelName, { modelReferenceName: toModelName, - columnReferenceName: toColumnName, + columnReferenceName: + this.columnNameAliases.get(toColumnId) || toColumnName, relation: name, }); this.addRelationColumn(toModelName, { modelReferenceName: fromModelName, - columnReferenceName: fromColumnName, + columnReferenceName: + this.columnNameAliases.get(fromColumnId) || fromColumnName, relation: name, }); @@ -359,14 +439,18 @@ export class MDLBuilder implements IMDLBuilder { if (!model.columns) { model.columns = []; } - // check if the modelReferenceName is already in the model column - const modelNameDuplicated = model.columns.find( - (column: any) => column.name === columnData.modelReferenceName, + const modelColumnNames = this.getManifestColumnNames(model); + const modelNameDuplicated = modelColumnNames.has( + columnData.modelReferenceName.toLowerCase(), ); - const column = { - name: modelNameDuplicated + const columnName = getUniqueReferenceName( + modelNameDuplicated ? `${columnData.modelReferenceName}_${columnData.columnReferenceName}` : columnData.modelReferenceName, + modelColumnNames, + ); + const column = { + name: columnName, type: columnData.modelReferenceName, properties: null, relationship: columnData.relation, @@ -379,12 +463,14 @@ export class MDLBuilder implements IMDLBuilder { protected getColumnExpression( column: ModelColumn, currentModel?: Partial, + columnReferenceName = column.referenceName, ): string { if (!column.isCalculated) { // columns existed in the data source. // Provide original column name in expression to MDL if referenceName has converted. - if (column.sourceColumnName !== column.referenceName) { - return `"${column.sourceColumnName}"`; + const sourceColumnName = column.sourceColumnName || column.referenceName; + if (sourceColumnName !== columnReferenceName) { + return `"${sourceColumnName}"`; } return ''; } @@ -396,9 +482,13 @@ export class MDLBuilder implements IMDLBuilder { const isLast = parseInt(index) == lineage.length - 1; if (isLast) { // id is columnId - const columnReferenceName = this.relatedColumns.find( + const relatedColumn = this.relatedColumns.find( (relatedColumn) => relatedColumn.id === id, - )?.referenceName; + ); + const columnReferenceName = relatedColumn + ? this.columnNameAliases.get(relatedColumn.id) || + relatedColumn.referenceName + : null; acc.push(`\"${columnReferenceName}\"`); return acc; } @@ -428,9 +518,19 @@ export class MDLBuilder implements IMDLBuilder { protected getRelationCondition(relation: RelationInfo): string { //TODO phase2: implement the expression for relation condition - const { fromColumnName, toColumnName, fromModelName, toModelName } = - relation; - return `"${fromModelName}".${fromColumnName} = "${toModelName}".${toColumnName}`; + const { + fromColumnId, + fromColumnName, + toColumnId, + toColumnName, + fromModelName, + toModelName, + } = relation; + const fromColumnReferenceName = + this.columnNameAliases.get(fromColumnId) || fromColumnName; + const toColumnReferenceName = + this.columnNameAliases.get(toColumnId) || toColumnName; + return `"${fromModelName}".${fromColumnReferenceName} = "${toModelName}".${toColumnReferenceName}`; } private buildTableReference(model: Model): TableReference | null { @@ -438,15 +538,170 @@ export class MDLBuilder implements IMDLBuilder { model.properties && typeof model.properties === 'string' ? JSON.parse(model.properties) : {}; - if (!modelProps.table) { + + if (modelProps.table) { + return this.parseTableReference( + modelProps.table, + modelProps.schema || null, + modelProps.catalog || null, + ); + } + + if (model.sourceTableName) { + const sourceTableReference = this.parseTableReference( + model.sourceTableName, + ); + if ( + sourceTableReference && + (sourceTableReference.schema || + sourceTableReference.catalog || + model.sourceTableName !== model.referenceName) + ) { + return sourceTableReference; + } + } + + if ( + model.displayName && + model.displayName.includes('.') && + replaceInvalidReferenceName(model.displayName) === model.referenceName + ) { + return this.parseTableReference(model.displayName); + } + + return model.sourceTableName + ? this.parseTableReference(model.sourceTableName) + : null; + } + + private parseTableReference( + qualifiedTableName: string, + schema: string | null = null, + catalog: string | null = null, + ): TableReference | null { + const tableParts = qualifiedTableName.split('.').filter(Boolean); + const table = tableParts.pop(); + if (!table) { return null; } + return { - catalog: modelProps.catalog || null, - schema: modelProps.schema || null, - table: modelProps.table, + catalog: + catalog || + (tableParts.length > 1 ? tableParts.slice(0, -1).join('.') : null), + schema: schema || tableParts.pop() || null, + table, }; } + + private hasDuplicateSourceColumns(modelId: number): boolean { + const sourceColumnNames = new Set(); + for (const column of this.columns.filter( + ({ isCalculated, modelId: columnModelId }) => + !isCalculated && columnModelId === modelId, + )) { + const sourceColumnName = ( + column.sourceColumnName || column.referenceName + ).toLowerCase(); + if (sourceColumnNames.has(sourceColumnName)) { + return true; + } + sourceColumnNames.add(sourceColumnName); + } + return false; + } + + private buildDedupedTableReferenceSql( + modelId: number, + model: Partial, + tableReference: TableReference, + ): string | null { + const sourceColumnNames = new Map(); + const projections: string[] = []; + + this.columns + .filter( + ({ isCalculated, modelId: columnModelId }) => + !isCalculated && columnModelId === modelId, + ) + .forEach((column) => { + const sourceColumnName = column.sourceColumnName || column.referenceName; + const normalizedSourceColumnName = sourceColumnName.toLowerCase(); + const existingColumnName = sourceColumnNames.get( + normalizedSourceColumnName, + ); + + if (existingColumnName) { + this.columnNameAliases.set(column.id, existingColumnName); + return; + } + + const columnName = this.getManifestColumnName(column, model); + sourceColumnNames.set(normalizedSourceColumnName, columnName); + const sourceExpression = this.quoteSqlIdentifier(sourceColumnName); + projections.push( + sourceColumnName === columnName + ? sourceExpression + : `${sourceExpression} AS ${this.quoteSqlIdentifier(columnName)}`, + ); + }); + + if (!projections.length) { + return null; + } + + const tableParts = [ + tableReference.catalog, + tableReference.schema, + tableReference.table, + ].filter((part): part is string => Boolean(part)); + return `SELECT ${projections.join(', ')} FROM ${tableParts + .map((part) => this.quoteSqlIdentifier(part)) + .join('.')}`; + } + + private quoteSqlIdentifier(identifier: string): string { + return `"${identifier.replace(/"/g, '""')}"`; + } + + private getManifestColumnName( + column: ModelColumn, + model: Partial, + ): string { + if (this.columnNameAliases.has(column.id)) { + return this.columnNameAliases.get(column.id)!; + } + + const columnName = getUniqueReferenceName( + column.referenceName, + this.getManifestColumnNames(model), + ); + this.columnNameAliases.set(column.id, columnName); + return columnName; + } + + private getManifestColumnNames(model: Partial): Set { + const modelName = model.name || ''; + if (!this.manifestColumnNamesByModel.has(modelName)) { + const existingColumns = (model.columns || []) as ColumnMDL[]; + this.manifestColumnNamesByModel.set( + modelName, + new Set(existingColumns.map((column) => column.name.toLowerCase())), + ); + } + return this.manifestColumnNamesByModel.get(modelName)!; + } + + private getManifestSourceColumnNameMap( + model: Partial, + ): Map { + const modelName = model.name || ''; + if (!this.manifestColumnNameBySourceByModel.has(modelName)) { + this.manifestColumnNameBySourceByModel.set(modelName, new Map()); + } + return this.manifestColumnNameBySourceByModel.get(modelName)!; + } + private postProcessManifest() { if (this.useRustWrenEngine()) { // 1. remove all the key that the value is null diff --git a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts index 4c7976d665..76f4b1dff5 100644 --- a/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts +++ b/wren-ui/src/apollo/server/mdl/test/mdlBuilder.test.ts @@ -32,6 +32,165 @@ describe('MDLBuilder', () => { expect(manifest).toBeDefined(); }); + it('should skip duplicated model column reference names', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'my project', + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'orders', + sourceTableName: 'dbo.orders', + referenceName: 'orders', + refSql: null, + cached: false, + refreshTime: null, + properties: JSON.stringify({ table: 'dbo.orders' }), + }, + ] as Model[]; + const columns = [ + { + id: 1, + modelId: 1, + isCalculated: false, + displayName: 'status', + referenceName: 'status', + sourceColumnName: 'status', + type: 'STRING', + notNull: false, + isPk: false, + properties: null, + }, + { + id: 2, + modelId: 1, + isCalculated: false, + displayName: 'Status', + referenceName: 'Status', + sourceColumnName: 'Status', + type: 'STRING', + notNull: false, + isPk: false, + properties: null, + }, + ] as ModelColumn[]; + const builderOptions = { + project, + models, + columns, + nestedColumns: [], + relations: [], + views: [], + relatedModels: models, + relatedColumns: columns, + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].columns).toHaveLength(1); + expect((manifest.models[0] as ModelMDL).columns[0].name).toBe('status'); + }); + + it('should build an explicit projection when source columns are duplicated', () => { + const catalogName = 'catalog_fixture'; + const schemaName = 'schema_fixture'; + const physicalTableName = 'physical_entity_fixture'; + const modelName = 'model_entity_fixture'; + const duplicateSourceName = 'field_alpha'; + const duplicateSourceNameWithDifferentCase = 'Field_Alpha'; + const retainedSourceName = 'field_beta'; + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'my project', + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Model Entity Fixture', + sourceTableName: physicalTableName, + referenceName: modelName, + refSql: null, + cached: false, + refreshTime: null, + properties: JSON.stringify({ + catalog: catalogName, + schema: schemaName, + table: physicalTableName, + }), + }, + ] as Model[]; + const columns = [ + { + id: 1, + modelId: 1, + isCalculated: false, + displayName: duplicateSourceName, + referenceName: duplicateSourceName, + sourceColumnName: duplicateSourceName, + type: 'STRING', + notNull: false, + isPk: false, + properties: null, + }, + { + id: 2, + modelId: 1, + isCalculated: false, + displayName: duplicateSourceNameWithDifferentCase, + referenceName: duplicateSourceNameWithDifferentCase, + sourceColumnName: duplicateSourceNameWithDifferentCase, + type: 'STRING', + notNull: false, + isPk: false, + properties: null, + }, + { + id: 3, + modelId: 1, + isCalculated: false, + displayName: retainedSourceName, + referenceName: retainedSourceName, + sourceColumnName: retainedSourceName, + type: 'DATE', + notNull: false, + isPk: false, + properties: null, + }, + ] as ModelColumn[]; + const builderOptions = { + project, + models, + columns, + nestedColumns: [], + relations: [], + views: [], + relatedModels: models, + relatedColumns: columns, + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toBeNull(); + expect(manifest.models[0].refSql).toEqual( + `SELECT "${duplicateSourceName}", "${retainedSourceName}" FROM "${catalogName}"."${schemaName}"."${physicalTableName}"`, + ); + expect(manifest.models[0].columns.map((column) => column.name)).toEqual([ + duplicateSourceName, + retainedSourceName, + ]); + }); + it('should return a manifest with models & columns & relations.', () => { // Arrange const project = { @@ -539,6 +698,355 @@ describe('MDLBuilder', () => { expect(manifest.views).toEqual(expectedViews); }); + it('should use tableReference when a model has no refSql.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Search Queries', + sourceTableName: 'search_queries', + referenceName: 'search_queries', + refSql: null, + cached: false, + refreshTime: null, + properties: JSON.stringify({ + schema: 'public', + table: 'search_queries', + }), + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'public', + table: 'search_queries', + }); + expect(manifest.models[0].refSql).toBeFalsy(); + }); + + it('should build tableReference from sourceTableName for imported models without stored table metadata.', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'wren ai project', + connectionInfo: {}, + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'dbo.PBI_View_Apex_Invoice', + sourceTableName: 'dbo.PBI_View_Apex_Invoice', + referenceName: 'dbo_PBI_View_Apex_Invoice', + refSql: null, + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'dbo', + table: 'PBI_View_Apex_Invoice', + }); + expect(manifest.models[0].refSql).toBeFalsy(); + }); + + it('should split qualified table names stored in model properties.', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'wren ai project', + connectionInfo: {}, + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'dbo.orders', + sourceTableName: 'dbo.orders', + referenceName: 'dbo_orders', + refSql: null, + cached: false, + refreshTime: null, + properties: JSON.stringify({ table: 'dbo.orders' }), + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'dbo', + table: 'orders', + }); + }); + + it('should recover tableReference from displayName for legacy sanitized models.', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'wren ai project', + connectionInfo: {}, + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'dbo.vw_HL_Invoice_Details', + sourceTableName: 'dbo_vw_HL_Invoice_Details', + referenceName: 'dbo_vw_HL_Invoice_Details', + refSql: null, + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'dbo', + table: 'vw_HL_Invoice_Details', + }); + }); + + it('should prefer physical tableReference over stale sanitized refSql.', () => { + const project = { + id: 1, + type: DataSourceName.MSSQL, + displayName: 'wren ai project', + connectionInfo: {}, + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'source.Actual_Table', + sourceTableName: 'source.Actual_Table', + referenceName: 'source_Actual_Table', + refSql: 'SELECT * FROM source_Actual_Table', + cached: false, + refreshTime: null, + properties: JSON.stringify({ + table: 'source.Actual_Table', + }), + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toEqual({ + catalog: null, + schema: 'source', + table: 'Actual_Table', + }); + expect(manifest.models[0].refSql).toBeFalsy(); + }); + + it('should build refSql from tableReference metadata when a model has columns but no refSql.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Imported Model', + sourceTableName: 'imported_model', + referenceName: 'imported_model', + refSql: null, + cached: false, + refreshTime: null, + properties: JSON.stringify({ + catalog: 'physical_catalog', + schema: 'physical_schema', + table: 'physical_table', + }), + }, + ] as Model[]; + const columns = [ + { + id: 1, + modelId: 1, + isCalculated: false, + displayName: 'Source Name', + referenceName: 'SourceName', + sourceColumnName: 'Source Name', + type: 'VARCHAR', + notNull: false, + isPk: false, + properties: null, + }, + { + id: 2, + modelId: 1, + isCalculated: false, + displayName: 'Source Type', + referenceName: 'SourceType', + sourceColumnName: 'Source Type', + type: 'VARCHAR', + notNull: false, + isPk: false, + properties: null, + }, + ] as ModelColumn[]; + const builderOptions = { + project, + models, + columns, + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toBeFalsy(); + expect(manifest.models[0].refSql).toEqual( + 'SELECT "Source Name" AS "SourceName", "Source Type" AS "SourceType" FROM "physical_catalog"."physical_schema"."physical_table"', + ); + expect(manifest.models[0].columns[0].properties.sourceColumnName).toEqual( + 'Source Name', + ); + expect(manifest.models[0].columns[1].properties.sourceColumnName).toEqual( + 'Source Type', + ); + }); + + it('should preserve refSql when a model has no tableReference.', () => { + const project = { + id: 1, + type: DataSourceName.POSTGRES, + displayName: 'wren ai project', + connectionInfo: {}, + catalog: 'wrenai', + schema: 'public', + sampleDataset: null, + } as Project; + const models = [ + { + id: 1, + projectId: 1, + displayName: 'Semantic Model', + sourceTableName: 'physical_table', + referenceName: 'semantic_model', + refSql: 'SELECT * FROM physical_schema.physical_table', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const builderOptions = { + project, + models, + columns: [], + nestedColumns: [], + relations: [], + views: [], + relatedModels: [], + relatedColumns: [], + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + expect(manifest.models[0].tableReference).toBeFalsy(); + expect(manifest.models[0].refSql).toEqual( + 'SELECT * FROM physical_schema.physical_table', + ); + }); + it('should return correct expression in calculated field.', () => { const models = [ // customer model @@ -617,7 +1125,7 @@ describe('MDLBuilder', () => { displayName: 'total_payment', referenceName: 'total_payment', sourceColumnName: 'total_payment', - aggregation: 'sum', + aggregation: 'SUM', lineage: JSON.stringify([1, 2, 8]), customExpression: null, type: 'FLOAT', @@ -758,6 +1266,74 @@ describe('MDLBuilder', () => { ); }); + it('should skip calculated fields with expressions outside the documented function list.', () => { + const models = [ + { + id: 1, + projectId: 1, + displayName: 'orders', + sourceTableName: 'orders', + referenceName: 'orders', + refSql: 'SELECT * FROM orders', + cached: false, + refreshTime: null, + properties: null, + }, + ] as Model[]; + const columns = [ + { + id: 1, + modelId: 1, + isCalculated: false, + displayName: 'id', + referenceName: 'id', + sourceColumnName: 'id', + aggregation: null, + lineage: null, + customExpression: null, + type: 'INTEGER', + notNull: true, + isPk: true, + properties: null, + }, + { + id: 2, + modelId: 1, + isCalculated: true, + displayName: 'unsupported_count_if', + referenceName: 'unsupported_count_if', + sourceColumnName: 'unsupported_count_if', + aggregation: 'COUNT_IF', + lineage: JSON.stringify([1]), + customExpression: null, + type: 'BIGINT', + notNull: false, + isPk: false, + properties: null, + }, + ] as ModelColumn[]; + const builderOptions = { + project: { + schema: 'public', + catalog: 'wrenai', + }, + models, + columns, + relations: [], + relatedModels: models, + relatedColumns: columns, + relatedRelations: [], + } as MDLBuilderBuildFromOptions; + mdlBuilder = new MDLBuilder(builderOptions); + + const manifest = mdlBuilder.build(); + + const ordersModel = manifest.models.find((m) => m.name === 'orders'); + expect( + ordersModel.columns.find((c) => c.name === 'unsupported_count_if'), + ).toBeUndefined(); + }); + it.each(Object.values(DataSourceName))( `should return correct data source type`, (type) => { diff --git a/wren-ui/src/apollo/server/mdl/type.ts b/wren-ui/src/apollo/server/mdl/type.ts index bb7c8be10f..fc2bea59fe 100644 --- a/wren-ui/src/apollo/server/mdl/type.ts +++ b/wren-ui/src/apollo/server/mdl/type.ts @@ -7,6 +7,7 @@ export interface ColumnMDL { properties?: { description?: string; // eg: "the key of each order" displayName?: string; // eg: "Order Key" + sourceColumnName?: string; }; expression?: string; // eg: "SUM(orders.totalprice)" } diff --git a/wren-ui/src/apollo/server/middlewares/organizationApi.ts b/wren-ui/src/apollo/server/middlewares/organizationApi.ts new file mode 100644 index 0000000000..aa7f6e34ef --- /dev/null +++ b/wren-ui/src/apollo/server/middlewares/organizationApi.ts @@ -0,0 +1,50 @@ +import { NextApiRequest } from 'next'; +import { ApiError } from '../utils/apiUtils'; + +const getProjectService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.projectService; +}; + +export const getCurrentProjectContext = async () => { + try { + const projectService = getProjectService(); + const project = await projectService.getCurrentProject(); + return { + id: project?.id ?? null, + displayName: project?.displayName || 'Default Project', + }; + } catch { + return { + id: null, + displayName: 'Default Project', + }; + } +}; + +export const assertAllowedMethods = ( + req: NextApiRequest, + methods: string[], +) => { + if (!req.method || !methods.includes(req.method)) { + throw new ApiError('Method not allowed', 405); + } +}; + +export const parseOrganizationId = (value: string | string[] | undefined) => { + const rawValue = Array.isArray(value) ? value[0] : value; + const parsed = Number(rawValue); + if (!rawValue || !Number.isInteger(parsed) || parsed <= 0) { + throw new ApiError('Invalid organization id', 400); + } + return parsed; +}; + +export const getCurrentProjectName = async () => { + const project = await getCurrentProjectContext(); + return project.displayName; +}; diff --git a/wren-ui/src/apollo/server/models/adaptor.ts b/wren-ui/src/apollo/server/models/adaptor.ts index 8135c07160..218ad79523 100644 --- a/wren-ui/src/apollo/server/models/adaptor.ts +++ b/wren-ui/src/apollo/server/models/adaptor.ts @@ -51,6 +51,7 @@ export enum WrenAILanguage { export interface DeployData { manifest: Manifest; hash: string; + projectId: number; } // ask @@ -73,8 +74,13 @@ export interface ProjectConfigurations { export interface AskInput { query: string; deployId: string; + projectId?: string; histories?: ThreadResponse[]; configurations?: ProjectConfigurations; + ignoreSqlGenerationReasoning?: boolean; + enableColumnPruning?: boolean; + useDryPlan?: boolean; + allowDryPlanFallback?: boolean; } export interface AsyncQueryResponse { @@ -223,6 +229,7 @@ export enum ChartType { export interface ChartInput { query: string; sql: string; + data?: Record; projectId?: string; configurations?: ProjectConfigurations; } @@ -239,6 +246,7 @@ export interface ChartAdjustmentOption { export interface ChartAdjustmentInput { query: string; sql: string; + data?: Record; adjustmentOption: ChartAdjustmentOption; chartSchema: Record; projectId?: string; diff --git a/wren-ui/src/apollo/server/models/dashboard.ts b/wren-ui/src/apollo/server/models/dashboard.ts index d4c7e9d57b..f7da453868 100644 --- a/wren-ui/src/apollo/server/models/dashboard.ts +++ b/wren-ui/src/apollo/server/models/dashboard.ts @@ -40,8 +40,8 @@ export interface SetDashboardCacheData { } export interface DetailedDashboard { - id: number; - projectId: number; + id: string | number; + projectId: string | number; name: string; cacheEnabled: boolean; scheduleFrequency: ScheduleFrequencyEnum | null; diff --git a/wren-ui/src/apollo/server/models/index.ts b/wren-ui/src/apollo/server/models/index.ts index 349625fd76..2a3049bb90 100644 --- a/wren-ui/src/apollo/server/models/index.ts +++ b/wren-ui/src/apollo/server/models/index.ts @@ -2,3 +2,4 @@ export * from './model'; export * from './instruction'; export * from './adaptor'; export * from './dashboard'; +export * from './rbac'; diff --git a/wren-ui/src/apollo/server/models/model.ts b/wren-ui/src/apollo/server/models/model.ts index fd763d3de6..1cec4870bb 100644 --- a/wren-ui/src/apollo/server/models/model.ts +++ b/wren-ui/src/apollo/server/models/model.ts @@ -11,44 +11,53 @@ export interface UpdateModelData { export interface NestedColumnMetadataInput { id: number; - displayName: string; - description: string; + displayName?: string; + description?: string; } export interface ColumnMetadataInput { id: number; - displayName: string; - description: string; + referenceName?: string; + displayName?: string; + description?: string; } export interface CalculatedFieldMetadataInput { id: number; - description: string; + description?: string; } export interface RelationshipMetadataInput { id: number; - description: string; + description?: string; } export interface ViewColumnMetadataInput { referenceName: string; - description: string; + description?: string; } export interface UpdateModelMetadataInput { - displayName: string; - description: string; + displayName?: string; + description?: string; + columns?: Array; + nestedColumns?: Array; + calculatedFields?: Array; + relationships?: Array; +} + +export interface SaveModelingSemanticInput { + modelId: number; + referenceName?: string; + displayName?: string; + description?: string; columns: Array; - nestedColumns: Array; - calculatedFields: Array; - relationships: Array; } export interface UpdateViewMetadataInput { - displayName: string; - description: string; - columns: Array; + displayName?: string; + description?: string; + columns?: Array; } export enum ExpressionName { @@ -92,6 +101,14 @@ export interface CheckCalculatedFieldCanQueryData { export interface PreviewSQLData { sql: string; projectId?: string; + hash?: string; limit?: number; dryRun?: boolean; } + +export interface DryPlanSQLData { + sql: string; + projectId?: string; + hash?: string; + allowFallback?: boolean; +} diff --git a/wren-ui/src/apollo/server/models/rbac.ts b/wren-ui/src/apollo/server/models/rbac.ts new file mode 100644 index 0000000000..52b1a98ce9 --- /dev/null +++ b/wren-ui/src/apollo/server/models/rbac.ts @@ -0,0 +1,38 @@ +export interface CreateRoleInput { + name: string; + description?: string | null; +} + +export interface UpdateRoleInput { + id: number; + name?: string | null; + description?: string | null; +} + +export interface CreateUserInput { + name: string; + email: string; + externalId?: string | null; + identityProvider?: string | null; + isActive?: boolean; + roleIds?: number[]; +} + +export interface UpdateUserInput { + id: number; + name?: string | null; + email?: string | null; + externalId?: string | null; + identityProvider?: string | null; + isActive?: boolean | null; +} + +export interface UserRoleInput { + userId: number; + roleId: number; +} + +export interface UpdateUserRolesInput { + userId: number; + roleIds: number[]; +} diff --git a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts index 7de5e5c097..46364bbeff 100644 --- a/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts +++ b/wren-ui/src/apollo/server/repositories/apiHistoryRepository.ts @@ -5,7 +5,11 @@ import { mapValues, snakeCase, } from 'lodash'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { Knex } from 'knex'; export enum ApiType { @@ -25,8 +29,60 @@ export enum ApiType { GET_MODELS = 'GET_MODELS', STREAM_ASK = 'STREAM_ASK', STREAM_GENERATE_SQL = 'STREAM_GENERATE_SQL', + GET_ORGANIZATIONS = 'GET_ORGANIZATIONS', + GET_CURRENT_ORGANIZATION = 'GET_CURRENT_ORGANIZATION', + CREATE_ORGANIZATION = 'CREATE_ORGANIZATION', + UPDATE_CURRENT_ORGANIZATION = 'UPDATE_CURRENT_ORGANIZATION', + SELECT_ORGANIZATION = 'SELECT_ORGANIZATION', + GET_ORGANIZATION_MEMBERS = 'GET_ORGANIZATION_MEMBERS', + INVITE_ORGANIZATION_MEMBER = 'INVITE_ORGANIZATION_MEMBER', + UPDATE_ORGANIZATION_MEMBER = 'UPDATE_ORGANIZATION_MEMBER', + REMOVE_ORGANIZATION_MEMBER = 'REMOVE_ORGANIZATION_MEMBER', + REMOVE_ORGANIZATION_INVITATION = 'REMOVE_ORGANIZATION_INVITATION', + ACCEPT_ORGANIZATION_INVITATION = 'ACCEPT_ORGANIZATION_INVITATION', + LEAVE_ORGANIZATION = 'LEAVE_ORGANIZATION', + DELETE_CURRENT_ORGANIZATION = 'DELETE_CURRENT_ORGANIZATION', + GET_PROJECTS = 'GET_PROJECTS', + GET_CURRENT_PROJECT = 'GET_CURRENT_PROJECT', + SELECT_PROJECT = 'SELECT_PROJECT', + GET_CURRENT_USER = 'GET_CURRENT_USER', + UPDATE_CURRENT_USER = 'UPDATE_CURRENT_USER', + DELETE_CURRENT_USER = 'DELETE_CURRENT_USER', + GET_PROJECT_ACCESS = 'GET_PROJECT_ACCESS', + ADD_PROJECT_MEMBER = 'ADD_PROJECT_MEMBER', + UPDATE_PROJECT_MEMBER = 'UPDATE_PROJECT_MEMBER', + REMOVE_PROJECT_MEMBER = 'REMOVE_PROJECT_MEMBER', } +export const INTERNAL_API_HISTORY_TYPES = [ + ApiType.GET_ORGANIZATIONS, + ApiType.GET_CURRENT_ORGANIZATION, + ApiType.CREATE_ORGANIZATION, + ApiType.UPDATE_CURRENT_ORGANIZATION, + ApiType.SELECT_ORGANIZATION, + ApiType.GET_ORGANIZATION_MEMBERS, + ApiType.INVITE_ORGANIZATION_MEMBER, + ApiType.UPDATE_ORGANIZATION_MEMBER, + ApiType.REMOVE_ORGANIZATION_MEMBER, + ApiType.REMOVE_ORGANIZATION_INVITATION, + ApiType.ACCEPT_ORGANIZATION_INVITATION, + ApiType.LEAVE_ORGANIZATION, + ApiType.DELETE_CURRENT_ORGANIZATION, + ApiType.GET_PROJECTS, + ApiType.GET_CURRENT_PROJECT, + ApiType.SELECT_PROJECT, + ApiType.GET_CURRENT_USER, + ApiType.UPDATE_CURRENT_USER, + ApiType.DELETE_CURRENT_USER, + ApiType.GET_PROJECT_ACCESS, + ApiType.ADD_PROJECT_MEMBER, + ApiType.UPDATE_PROJECT_MEMBER, + ApiType.REMOVE_PROJECT_MEMBER, +]; + +export const shouldPersistApiHistory = (apiType?: ApiType) => + Boolean(apiType && !INTERNAL_API_HISTORY_TYPES.includes(apiType)); + export interface ApiHistory { id?: string; projectId: number; @@ -73,6 +129,20 @@ export class ApiHistoryRepository super({ knexPg, tableName: 'api_history' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany(data.map(this.withTimestamps), queryOptions); + } + /** * Count API history records with filtering */ @@ -81,6 +151,7 @@ export class ApiHistoryRepository dateFilter?: { startDate?: Date; endDate?: Date }, ): Promise { let query = this.knex(this.tableName).count('id as count'); + query = this.excludeInternalApiTypes(query); if (filter) { query = query.where(this.transformToDBData(filter)); @@ -109,6 +180,7 @@ export class ApiHistoryRepository pagination?: PaginationOptions, ): Promise { let query = this.knex(this.tableName).select('*'); + query = this.excludeInternalApiTypes(query); if (filter) { query = query.where(this.transformToDBData(filter)); @@ -166,6 +238,15 @@ export class ApiHistoryRepository return formattedData; }; + private withTimestamps = (data: Partial): Partial => { + const now = new Date().toISOString(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + protected override transformToDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); @@ -186,4 +267,8 @@ export class ApiHistoryRepository private camelToSnakeCase(str: string): string { return str.replace(/[A-Z]/g, (letter) => `_${letter.toLowerCase()}`); } + + private excludeInternalApiTypes(query: Knex.QueryBuilder) { + return query.whereNotIn('api_type', INTERNAL_API_HISTORY_TYPES); + } } diff --git a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts index 942f445ee1..ff98b7502a 100644 --- a/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts +++ b/wren-ui/src/apollo/server/repositories/askingTaskRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -44,6 +48,35 @@ export class AskingTaskRepository return this.findOneBy({ queryId }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany(data.map(this.withTimestamps), queryOptions); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + protected override transformFromDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); @@ -74,4 +107,15 @@ export class AskingTaskRepository }); return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; } diff --git a/wren-ui/src/apollo/server/repositories/baseRepository.ts b/wren-ui/src/apollo/server/repositories/baseRepository.ts index f76fe27654..da7c54b1d2 100644 --- a/wren-ui/src/apollo/server/repositories/baseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/baseRepository.ts @@ -41,9 +41,25 @@ export interface IBasicRepository { ) => Promise; } +export const coerceBoolean = (value: unknown): boolean => { + if (typeof value === 'boolean') { + return value; + } + if (typeof value === 'number') { + return value === 1; + } + if (typeof value === 'string') { + return ['1', 'true'].includes(value.toLowerCase()); + } + return Boolean(value); +}; + export class BaseRepository implements IBasicRepository { + private static manualIdInsertLocks = new Map>(); protected knex: Knex; protected tableName: string; + private hasIdColumnCache: boolean | null = null; + private hasIdentityIdPromise?: Promise; constructor({ knexPg, tableName }: { knexPg: Knex; tableName: string }) { this.knex = knexPg; @@ -65,7 +81,7 @@ export class BaseRepository implements IBasicRepository { public async findOneBy(filter: Partial, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const query = executer(this.tableName).where( - this.transformToDBData(filter), + this.normalizeMssqlBindings(this.transformToDBData(filter)), ); if (queryOptions?.limit) { query.limit(queryOptions.limit); @@ -81,7 +97,7 @@ export class BaseRepository implements IBasicRepository { // format filter keys to snake_case const query = executer(this.tableName).where( - this.transformToDBData(filter), + this.normalizeMssqlBindings(this.transformToDBData(filter)), ); if (queryOptions?.order) { query.orderBy(queryOptions.order); @@ -105,28 +121,77 @@ export class BaseRepository implements IBasicRepository { public async createOne(data: Partial, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const [result] = await executer(this.tableName) - .insert(this.transformToDBData(data)) - .returning('*'); - return this.transformFromDBData(result); + try { + const insertValue = await this.prepareInsertData(data, executer); + const [result] = await this.insertOne(executer, insertValue); + return this.transformFromDBData(result); + } catch (error) { + if (!this.shouldRetryManualId(error, data, executer)) { + throw error; + } + + return await this.withManualIdInsertLock(executer, async () => { + let lastError: unknown = error; + for (let attempt = 0; attempt < 3; attempt++) { + try { + const insertValue = await this.prepareInsertData( + data, + executer, + true, + ); + const [result] = await this.insertOne(executer, insertValue); + return this.transformFromDBData(result); + } catch (retryError) { + if (!this.shouldRetryManualId(retryError, data, executer)) { + throw retryError; + } + lastError = retryError; + } + } + throw lastError; + }); + } } public async createMany(data: Partial[], queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const batchSize = 100; - const batchCount = Math.ceil(data.length / batchSize); - const result = []; - for (let i = 0; i < batchCount; i++) { - const start = i * batchSize; - const end = Math.min((i + 1) * batchSize, data.length); - const batchValues = data.slice(start, end); - const chunk = await executer(this.tableName) - .insert(batchValues.map(this.transformToDBData)) - .returning('*'); - result.push(...chunk); + let preparedData: any[]; + try { + preparedData = await this.prepareInsertManyData(data, executer); + } catch (error) { + throw error; + } + if (preparedData.length === 0) { + return []; } - return result.map((data) => this.transformFromDBData(data)); + try { + return await this.insertMany(executer, preparedData); + } catch (error) { + if (!this.shouldRetryManualId(error, data, executer)) { + throw error; + } + + return await this.withManualIdInsertLock(executer, async () => { + let lastError: unknown = error; + for (let attempt = 0; attempt < 3; attempt++) { + try { + preparedData = await this.prepareInsertManyData( + data, + executer, + true, + ); + return await this.insertMany(executer, preparedData); + } catch (retryError) { + if (!this.shouldRetryManualId(retryError, data, executer)) { + throw retryError; + } + lastError = retryError; + } + } + throw lastError; + }); + } } public async updateOne( @@ -135,16 +200,21 @@ export class BaseRepository implements IBasicRepository { queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const normalizedId = this.normalizeMssqlBindings({ id }).id; const [result] = await executer(this.tableName) - .where({ id }) - .update(this.transformToDBData(data)) + .where({ id: normalizedId }) + .update(this.normalizeMssqlBindings(this.transformToDBData(data))) .returning('*'); return this.transformFromDBData(result); } public async deleteOne(id: string, queryOptions?: IQueryOptions) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer.from(this.tableName).where({ id }).delete(); + const normalizedId = this.normalizeMssqlBindings({ id }).id; + const builder = executer + .from(this.tableName) + .where({ id: normalizedId }) + .delete(); return await builder; } @@ -153,8 +223,15 @@ export class BaseRepository implements IBasicRepository { queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer.from(this.tableName).whereIn('id', ids).delete(); - return await builder; + const normalizedIds = this.normalizeMssqlBindings(ids); + let deleted = 0; + for (const batch of this.toWhereInBatches(normalizedIds)) { + deleted += await executer + .from(this.tableName) + .whereIn('id', batch) + .delete(); + } + return deleted; } public deleteAllBy = async ( @@ -163,7 +240,7 @@ export class BaseRepository implements IBasicRepository { ) => { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const builder = executer(this.tableName) - .where(this.transformToDBData(where)) + .where(this.normalizeMssqlBindings(this.transformToDBData(where))) .delete(); return await builder; }; @@ -175,11 +252,303 @@ export class BaseRepository implements IBasicRepository { return mapKeys(data, (_value, key) => snakeCase(key)); }; - protected transformFromDBData = (data: any): T => { + protected defaultTransformFromDBData(data: any): T { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); } const camelCaseData = mapKeys(data, (_value, key) => camelCase(key)); return camelCaseData as T; - }; + } + + protected transformFromDBData = (data: any): T => + this.defaultTransformFromDBData(data); + + protected getMssqlWhereInBatchSize() { + return 2000; + } + + protected toWhereInBatches(values: TValue[]) { + const client = String(this.knex.client.config.client || '').toLowerCase(); + const batchSize = + client === 'mssql' + ? this.getMssqlWhereInBatchSize() + : Math.max(values.length, 1); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; + } + + protected getCreateManyBatchSize(insertValues: any[]) { + const defaultBatchSize = 100; + if (insertValues.length === 0) { + return defaultBatchSize; + } + + const client = String(this.knex.client.config.client || '').toLowerCase(); + if (client !== 'mssql') { + return defaultBatchSize; + } + + const parameterLimit = 2100; + const safetyMargin = 100; + const columnCount = Math.max( + ...insertValues.map((value) => Object.keys(value).length), + 1, + ); + + return Math.max( + 1, + Math.min( + defaultBatchSize, + Math.floor((parameterLimit - safetyMargin) / columnCount), + ), + ); + } + + private isMssql(executer: Knex | Knex.Transaction) { + return executer.client.config.client === 'mssql'; + } + + private async hasIdColumn(executer: Knex | Knex.Transaction) { + if (this.hasIdColumnCache !== null) { + return this.hasIdColumnCache; + } + + const hasIdColumn = await executer.schema.hasColumn(this.tableName, 'id'); + this.hasIdColumnCache = hasIdColumn; + return hasIdColumn; + } + + private serializeIdValue(value: number | string | bigint) { + const bigintValue = BigInt(value); + return bigintValue <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(bigintValue) + : bigintValue.toString(); + } + + private toNextIdValue(maxId: number | string | bigint | null | undefined) { + const nextId = BigInt(maxId ?? 0) + 1n; + return this.serializeIdValue(nextId); + } + + private async getNextId(executer: Knex | Knex.Transaction) { + const [row] = await executer(this.tableName).max<{ + maxId: number | string | bigint | null; + }>('id as maxId'); + return this.toNextIdValue(row?.maxId); + } + + private async withManualIdInsertLock( + executer: Knex | Knex.Transaction, + task: () => Promise, + ): Promise { + const lockKey = `${executer.client.config.client}:${this.tableName}`; + const previousLock = + BaseRepository.manualIdInsertLocks.get(lockKey) ?? Promise.resolve(); + let releaseLock: () => void = () => undefined; + const currentLock = new Promise((resolve) => { + releaseLock = resolve; + }); + const chainedLock = previousLock.catch(() => undefined).then( + () => currentLock, + ); + BaseRepository.manualIdInsertLocks.set(lockKey, chainedLock); + + await previousLock.catch(() => undefined); + try { + return await task(); + } finally { + releaseLock(); + if (BaseRepository.manualIdInsertLocks.get(lockKey) === chainedLock) { + BaseRepository.manualIdInsertLocks.delete(lockKey); + } + } + } + + private async hasIdentityId(executer: Knex | Knex.Transaction) { + if (!this.isMssql(executer)) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = executer('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + } + + private async prepareInsertData( + data: Partial, + executer: Knex | Knex.Transaction, + forceManualId = false, + ) { + const dbData = this.transformToDBData(data); + if (!this.isMssql(executer)) { + return dbData; + } + + if ( + !(await this.hasIdColumn(executer)) || + dbData.id !== undefined || + !forceManualId + ) { + return dbData; + } + + return { + ...dbData, + id: await this.getNextId(executer), + }; + } + + private async prepareInsertManyData( + data: Partial[], + executer: Knex | Knex.Transaction, + forceManualId = false, + ) { + const dbData = data.map((item) => this.transformToDBData(item)); + if ( + !this.isMssql(executer) || + !(await this.hasIdColumn(executer)) || + !forceManualId + ) { + return dbData; + } + + const missingIdIndexes = dbData.reduce((acc, item, index) => { + if (item.id === undefined) { + acc.push(index); + } + return acc; + }, []); + + if (!missingIdIndexes.length) { + return dbData; + } + + let nextId = BigInt((await this.getNextId(executer)) ?? 0); + for (const index of missingIdIndexes) { + dbData[index] = { + ...dbData[index], + id: this.serializeIdValue(nextId++), + }; + } + + return dbData; + } + + private async insertMany( + executer: Knex | Knex.Transaction, + preparedData: any[], + ) { + const batchSize = this.getCreateManyBatchSize(preparedData); + const batchCount = Math.ceil(preparedData.length / batchSize); + const result = []; + for (let i = 0; i < batchCount; i++) { + const start = i * batchSize; + const end = Math.min((i + 1) * batchSize, preparedData.length); + const batchValues = preparedData.slice(start, end); + const chunk = await executer(this.tableName) + .insert(this.normalizeMssqlBindings(batchValues)) + .returning('*'); + result.push(...chunk); + } + + return result.map((data) => this.transformFromDBData(data)); + } + + private async insertOne(executer: Knex | Knex.Transaction, preparedData: any) { + return await executer(this.tableName) + .insert(this.normalizeMssqlBindings(preparedData)) + .returning('*'); + } + + private normalizeMssqlBindings(value: any): any { + if (!this.isMssql(this.knex)) { + return value; + } + + if (Array.isArray(value)) { + return value.map((item) => this.normalizeMssqlBindings(item)); + } + + if (typeof value === 'bigint') { + return value.toString(); + } + + if (!isPlainObject(value)) { + if ( + typeof value === 'number' && + Number.isInteger(value) && + !Number.isSafeInteger(value) + ) { + return String(value); + } + return value; + } + + return Object.fromEntries( + Object.entries(value).map(([key, entryValue]) => { + if ( + typeof entryValue === 'number' && + Number.isInteger(entryValue) && + (!Number.isSafeInteger(entryValue) || + key === 'id' || + key.endsWith('_id')) + ) { + return [key, String(entryValue)]; + } + + if (typeof entryValue === 'bigint') { + return [key, entryValue.toString()]; + } + + return [key, this.normalizeMssqlBindings(entryValue)]; + }), + ); + } + + private shouldRetryManualId( + error: unknown, + data: Partial | Partial[], + executer: Knex | Knex.Transaction, + ) { + if (!this.isMssql(executer)) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item: any) => item?.id === undefined || item?.id === null) + : (data as any)?.id === undefined || (data as any)?.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return ( + message.includes("Cannot insert the value NULL into column 'id'") || + message.includes('Violation of PRIMARY KEY constraint') || + message.includes('Cannot insert duplicate key') + ); + } } diff --git a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts index 12959f0167..fb9a5b9a73 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardItemRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -31,15 +35,19 @@ export interface DashboardItemLayout { export interface DashboardItemDetail { sql: string; chartSchema?: Record; + previewDataSnapshot?: Record[]; } export interface DashboardItem { - id: number; - dashboardId: number; + id: string | number; + dashboardId: string | number; type: DashboardItemType; layout: DashboardItemLayout; detail: DashboardItemDetail; displayName?: string; + title?: string; + createdAt?: Date; + updatedAt?: Date; } export interface IDashboardItemRepository @@ -50,11 +58,37 @@ export class DashboardItemRepository implements IDashboardItemRepository { private readonly jsonbColumns = ['layout', 'detail']; + private hasIdColumnCache: boolean | null = null; + private hasTitleColumnCache: boolean | null = null; + private hasDisplayNameColumnCache: boolean | null = null; + private columnCache = new Map(); constructor(knexPg: Knex) { super({ knexPg, tableName: 'dashboard_item' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return await super.createOne( + await this.normalizeWriteData(data, queryOptions, true), + queryOptions, + ); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return await super.updateOne( + id, + await this.normalizeWriteData(data, queryOptions), + queryOptions, + ); + } + protected override transformFromDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); @@ -70,7 +104,10 @@ export class DashboardItemRepository } return value; }); - return transformData as DashboardItem; + return { + ...transformData, + displayName: transformData.displayName || transformData.title, + } as DashboardItem; }; protected override transformToDBData = (data: any) => { @@ -86,4 +123,112 @@ export class DashboardItemRepository }); return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; + + private async normalizeWriteData( + data: Partial, + queryOptions?: IQueryOptions, + includeGeneratedId = false, + forceManualId = false, + ): Promise> { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [ + hasIdColumn, + hasTitleColumn, + hasDisplayNameColumn, + hasCreatedAtColumn, + hasUpdatedAtColumn, + ] = + await Promise.all([ + this.hasColumn('id', executer), + this.hasColumn('title', executer), + this.hasColumn('display_name', executer), + this.hasColumn('created_at', executer), + this.hasColumn('updated_at', executer), + ]); + const normalizedData: Partial = { ...data }; + const displayName = + typeof data.displayName === 'string' ? data.displayName.trim() : ''; + const chartTitle = + typeof data.detail?.chartSchema?.title === 'string' + ? data.detail.chartSchema.title.trim() + : ''; + const title = displayName || chartTitle || 'Untitled dashboard item'; + + if (hasTitleColumn && !normalizedData.title) { + normalizedData.title = title; + } + if (!hasDisplayNameColumn) { + delete normalizedData.displayName; + } + if ( + includeGeneratedId && + hasIdColumn && + normalizedData.id === undefined && + this.isMssqlLike(executer) && + forceManualId + ) { + normalizedData.id = await this.getNextId(executer); + } + if (includeGeneratedId && hasCreatedAtColumn && !normalizedData.createdAt) { + normalizedData.createdAt = new Date(); + } + if (includeGeneratedId && hasUpdatedAtColumn && !normalizedData.updatedAt) { + normalizedData.updatedAt = normalizedData.createdAt || new Date(); + } + + return normalizedData; + } + + private async hasColumn(column: string, executer: Knex | Knex.Transaction) { + if (this.columnCache.has(column)) { + return this.columnCache.get(column); + } + if (column === 'id' && this.hasIdColumnCache !== null) { + return this.hasIdColumnCache; + } + if (column === 'title' && this.hasTitleColumnCache !== null) { + return this.hasTitleColumnCache; + } + if ( + column === 'display_name' && + this.hasDisplayNameColumnCache !== null + ) { + return this.hasDisplayNameColumnCache; + } + + const result = await executer.schema.hasColumn(this.tableName, column); + this.columnCache.set(column, result); + if (column === 'id') { + this.hasIdColumnCache = result; + } + if (column === 'title') { + this.hasTitleColumnCache = result; + } + if (column === 'display_name') { + this.hasDisplayNameColumnCache = result; + } + return result; + } + + private isMssqlLike(executer: Knex | Knex.Transaction) { + const clientName = String(executer.client.config.client || '').toLowerCase(); + const dialect = String((executer.client as any).dialect || '').toLowerCase(); + const driverName = String( + (executer.client as any).driverName || '', + ).toLowerCase(); + + return [clientName, dialect, driverName].some((value) => + value.includes('mssql'), + ); + } + + private async getNextId(executer: Knex | Knex.Transaction) { + const [row] = await executer(this.tableName).max<{ + maxId: number | string | bigint | null; + }>('id as maxId'); + const nextId = BigInt(row?.maxId ?? 0) + 1n; + return nextId <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(nextId) + : nextId.toString(); + } } diff --git a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts index d88fe35c68..86c0abaf45 100644 --- a/wren-ui/src/apollo/server/repositories/dashboardRepository.ts +++ b/wren-ui/src/apollo/server/repositories/dashboardRepository.ts @@ -1,16 +1,23 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, + IQueryOptions, +} from './baseRepository'; import { ScheduleFrequencyEnum } from '@server/models/dashboard'; export interface Dashboard { - id: number; - projectId: number; + id: string | number; + projectId: string | number; name: string; cacheEnabled: boolean; scheduleFrequency: ScheduleFrequencyEnum | null; scheduleTimezone: string | null; // e.g. 'America/New_York', 'Asia/Taipei' scheduleCron: string | null; // cron expression string nextScheduledAt: Date | null; // Next scheduled run timestamp + createdAt?: Date; + updatedAt?: Date; } export interface IDashboardRepository extends IBasicRepository {} @@ -19,7 +26,205 @@ export class DashboardRepository extends BaseRepository implements IDashboardRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'dashboard' }); } + + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + + protected override transformFromDBData = (data: any): Dashboard => { + const dashboard = this.defaultTransformFromDBData(data) as Dashboard; + return { + ...dashboard, + cacheEnabled: coerceBoolean(dashboard.cacheEnabled), + }; + }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + !forceManualId + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ + maxId: 'id', + }); + return { + ...data, + id: this.toNextIdValue(row?.maxId), + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + !forceManualId + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ + maxId: 'id', + }); + let nextId = BigInt(row?.maxId ?? 0) + 1n; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: this.serializeIdValue(nextId++), + }; + }); + }; + + private toNextIdValue = (maxId: number | string | bigint | null | undefined) => { + const nextId = BigInt(maxId ?? 0) + 1n; + return this.serializeIdValue(nextId); + }; + + private serializeIdValue = (value: number | string | bigint) => { + const bigintValue = BigInt(value); + return bigintValue <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(bigintValue) + : bigintValue.toString(); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } diff --git a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts index 6e2e0471fc..40a8b61095 100644 --- a/wren-ui/src/apollo/server/repositories/deployLogRepository.ts +++ b/wren-ui/src/apollo/server/repositories/deployLogRepository.ts @@ -1,6 +1,12 @@ import { Knex } from 'knex'; import { BaseRepository, IBasicRepository } from './baseRepository'; -import { camelCase, isPlainObject, mapKeys, mapValues } from 'lodash'; +import { + camelCase, + isPlainObject, + mapKeys, + mapValues, + snakeCase, +} from 'lodash'; export interface Deploy { id: number; // ID @@ -9,6 +15,8 @@ export interface Deploy { hash: string; status: string; // Deploy status error: string; // Error message + createdAt?: Date; + updatedAt?: Date; } export enum DeployStatusEnum { @@ -19,6 +27,7 @@ export enum DeployStatusEnum { export interface IDeployLogRepository extends IBasicRepository { findLastProjectDeployLog(projectId: number): Promise; + findLatestProjectDeployLog(projectId: number): Promise; findInProgressProjectDeployLog(projectId: number): Promise; } @@ -26,10 +35,28 @@ export class DeployLogRepository extends BaseRepository implements IDeployLogRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'deploy_log' }); } + public async createOne(data: Partial): Promise { + if (!this.isMssql() || (await this.hasIdentityId())) { + return super.createOne(this.withTimestamps(data)); + } + + const id = await this.nextDeployLogId(); + return super.createOne(this.withTimestamps({ id, ...data })); + } + + public async updateOne(id: number, data: Partial) { + return super.updateOne(id, { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }); + } + public async findLastProjectDeployLog(projectId: number) { const res = await this.knex .select('*') @@ -38,6 +65,18 @@ export class DeployLogRepository this.transformToDBData({ projectId, status: DeployStatusEnum.SUCCESS }), ) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') + .first(); + return (res && this.transformFromDBData(res)) || null; + } + + public async findLatestProjectDeployLog(projectId: number) { + const res = await this.knex + .select('*') + .from(this.tableName) + .where(this.transformToDBData({ projectId })) + .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -53,6 +92,7 @@ export class DeployLogRepository }), ) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -71,4 +111,55 @@ export class DeployLogRepository }); return formattedData as Deploy; }; + + protected override transformToDBData = (data: Partial) => { + if (!isPlainObject(data)) { + throw new Error('Unexpected dbdata'); + } + const snakeCaseData = mapKeys(data, (_value, key) => snakeCase(key)); + const formattedData = mapValues(snakeCaseData, (value, key) => { + if (key === 'manifest' && typeof value !== 'string') { + return JSON.stringify(value); + } + return value; + }); + return formattedData; + }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private async hasIdentityId() { + this.hasIdentityIdPromise ??= this.knex('sys.columns as c') + .join('sys.tables as t', 'c.object_id', 't.object_id') + .where('t.name', this.tableName) + .where('c.name', 'id') + .select( + this.knex.raw( + 'COLUMNPROPERTY(c.object_id, c.name, ?) as isIdentity', + ['IsIdentity'], + ), + ) + .first() + .then((row) => Number(row?.isIdentity ?? 0) === 1); + + return this.hasIdentityIdPromise; + } + + private async nextDeployLogId() { + const row = await this.knex(this.tableName) + .max<{ maxId?: number | string }>('id as maxId') + .first(); + + return Number(row?.maxId ?? 0) + 1; + } } diff --git a/wren-ui/src/apollo/server/repositories/index.ts b/wren-ui/src/apollo/server/repositories/index.ts index 0dd1cc5905..0ecfef7f79 100644 --- a/wren-ui/src/apollo/server/repositories/index.ts +++ b/wren-ui/src/apollo/server/repositories/index.ts @@ -19,3 +19,7 @@ export * from './askingTaskRepository'; export * from './instructionRepository'; export * from './apiHistoryRepository'; export * from './dashboardItemRefreshJobRepository'; +export * from './rbacRepository'; +export * from './organizationRepository'; +export * from './organizationMemberRepository'; +export * from './organizationInvitationRepository'; diff --git a/wren-ui/src/apollo/server/repositories/instructionRepository.ts b/wren-ui/src/apollo/server/repositories/instructionRepository.ts index d3cef2cbc2..05d12d4eeb 100644 --- a/wren-ui/src/apollo/server/repositories/instructionRepository.ts +++ b/wren-ui/src/apollo/server/repositories/instructionRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -43,6 +47,9 @@ export class InstructionRepository return value; } } + if (key === 'isDefault') { + return coerceBoolean(value); + } return value; }); return transformData as Instruction; diff --git a/wren-ui/src/apollo/server/repositories/metricsRepository.ts b/wren-ui/src/apollo/server/repositories/metricsRepository.ts index 57bec8f920..411ae9f45c 100644 --- a/wren-ui/src/apollo/server/repositories/metricsRepository.ts +++ b/wren-ui/src/apollo/server/repositories/metricsRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + coerceBoolean, +} from './baseRepository'; export interface Metric { id: number; // ID @@ -24,4 +28,12 @@ export class MetricRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'metric' }); } + + protected override transformFromDBData = (data: any): Metric => { + const metric = this.defaultTransformFromDBData(data) as Metric; + return { + ...metric, + cached: coerceBoolean(metric.cached), + }; + }; } diff --git a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts index df4198e214..43f6a89f97 100644 --- a/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelColumnRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, IQueryOptions, + coerceBoolean, } from './baseRepository'; export interface ModelColumn { @@ -19,6 +20,8 @@ export interface ModelColumn { notNull: boolean; // Is not null isPk: boolean; // Is primary key of the table properties?: string; // Column properties, a json string, the description and displayName should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IModelColumnRepository extends IBasicRepository { @@ -51,33 +54,92 @@ export class ModelColumnRepository extends BaseRepository implements IModelColumnRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'model_column' }); } - public async findColumnsByModelIds(modelIds, queryOptions?: IQueryOptions) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - const result = await tx(this.tableName) - .whereIn('model_id', modelIds) - .select('*'); - return result.map((r) => this.transformFromDBData(r)); + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); } - const result = await this.knex('model_column') - .whereIn('model_id', modelIds) - .select('*'); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + + protected override transformFromDBData = (data: any): ModelColumn => { + const column = this.defaultTransformFromDBData(data) as ModelColumn; + return { + ...column, + isCalculated: coerceBoolean(column.isCalculated), + notNull: coerceBoolean(column.notNull), + isPk: coerceBoolean(column.isPk), + }; + }; + + public async findColumnsByModelIds(modelIds, queryOptions?: IQueryOptions) { + const result = await this.findByColumnIn( + 'model_id', + modelIds, + queryOptions, + ); return result.map((r) => this.transformFromDBData(r)); } public async findColumnsByIds(ids: number[], queryOptions?: IQueryOptions) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - const result = await tx(this.tableName).whereIn('id', ids).select('*'); - return result.map((r) => this.transformFromDBData(r)); - } - const result = await this.knex('model_column') - .whereIn('id', ids) - .select('*'); + const result = await this.findByColumnIn('id', ids, queryOptions); return result.map((r) => this.transformFromDBData(r)); } @@ -85,14 +147,7 @@ export class ModelColumnRepository modelIds: number[], queryOptions?: IQueryOptions, ) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - await tx(this.tableName).whereIn('model_id', modelIds).delete(); - return; - } - await this.knex('model_column') - .whereIn('model_id', modelIds) - .delete(); + await this.deleteByColumnIn('model_id', modelIds, queryOptions); } public async resetModelPrimaryKey(modelId: number) { @@ -111,21 +166,177 @@ export class ModelColumnRepository sourceColumnNames: string[], queryOptions?: IQueryOptions, ): Promise { - const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer(this.tableName) - .where(this.transformToDBData({ modelId })) - .whereIn('source_column_name', sourceColumnNames) - .delete(); - return await builder; + return await this.deleteByColumnIn( + 'source_column_name', + sourceColumnNames, + queryOptions, + this.transformToDBData({ modelId }), + ); } public async deleteAllByColumnIds( columnIds: number[], queryOptions?: IQueryOptions, ): Promise { + await this.deleteByColumnIn('id', columnIds, queryOptions); + } + + private async findByColumnIn( + columnName: string, + values: Array, + queryOptions?: IQueryOptions, + ) { + if (values.length === 0) { + return []; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = []; + for (const batch of this.toWhereInBatches(values)) { + const result = await executer(this.tableName) + .whereIn(columnName, batch) + .select('*'); + rows.push(...result); + } + + return rows; + } + + private async deleteByColumnIn( + columnName: string, + values: Array, + queryOptions?: IQueryOptions, + extraWhere?: Record, + ) { + if (values.length === 0) { + return 0; + } + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - await executer(this.tableName) - .whereIn('id', columnIds) - .delete(); + let deleted = 0; + for (const batch of this.toWhereInBatches(values)) { + const query = executer(this.tableName).whereIn(columnName, batch); + if (extraWhere) { + query.where(extraWhere); + } + deleted += await query.delete(); + } + + return deleted; } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } diff --git a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts index 5438ff47ec..b46952b5ab 100644 --- a/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelNestedColumnRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -18,6 +22,8 @@ export interface ModelNestedColumn { sourceColumnName: string; // The nested column name in the datasource type: string; // Data type, refer to the nested column type in the datasource properties?: Record; // Nested column properties, a json string, the description should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IModelNestedColumnRepository @@ -30,24 +36,99 @@ export class ModelNestedColumnRepository extends BaseRepository implements IModelNestedColumnRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'model_nested_column' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public findNestedColumnsByModelIds = async (modelIds: number[]) => { - const result = await this.knex(this.tableName) - .select('*') - .whereIn('model_id', modelIds); + const result = await this.findByColumnIn('model_id', modelIds); return result.map((r) => this.transformFromDBData(r)); }; public findNestedColumnsByIds = async (ids: number[]) => { - const result = await this.knex(this.tableName) - .select('*') - .whereIn('id', ids); + const result = await this.findByColumnIn('id', ids); return result.map((r) => this.transformFromDBData(r)); }; + private async findByColumnIn( + columnName: string, + values: Array, + ) { + if (values.length === 0) { + return []; + } + + const rows = []; + for (const batch of this.toWhereInBatches(values)) { + const result = await this.knex(this.tableName) + .select('*') + .whereIn(columnName, batch); + rows.push(...result); + } + return rows; + } + protected override transformToDBData = (data: any) => { if (!isPlainObject(data)) { throw new Error('Unexpected dbdata'); @@ -79,4 +160,119 @@ export class ModelNestedColumnRepository }) as ModelNestedColumn; return formattedData; }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } diff --git a/wren-ui/src/apollo/server/repositories/modelRepository.ts b/wren-ui/src/apollo/server/repositories/modelRepository.ts index 57b92c9983..bbd1bd93d1 100644 --- a/wren-ui/src/apollo/server/repositories/modelRepository.ts +++ b/wren-ui/src/apollo/server/repositories/modelRepository.ts @@ -3,6 +3,7 @@ import { BaseRepository, IBasicRepository, IQueryOptions, + coerceBoolean, } from './baseRepository'; export interface Model { @@ -15,6 +16,8 @@ export interface Model { cached: boolean; // Model is cached or not refreshTime: string | null; // Contain a number followed by a time unit (ns, us, ms, s, m, h, d). For example, "2h" properties: string | null; // Model properties, a json string, the description and displayName should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IModelRepository extends IBasicRepository { @@ -29,12 +32,91 @@ export class ModelRepository extends BaseRepository implements IModelRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'model' }); } + + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + + protected override transformFromDBData = (data: any): Model => { + const model = this.defaultTransformFromDBData(data) as Model; + return { + ...model, + cached: coerceBoolean(model.cached), + }; + }; + public async findAllByIds(ids: number[]) { - const res = await this.knex(this.tableName).whereIn('id', ids); - return res.map((r) => this.transformFromDBData(r)); + if (ids.length === 0) { + return []; + } + + const rows = []; + for (const batch of this.toWhereInBatches(ids)) { + const res = await this.knex(this.tableName).whereIn('id', batch); + rows.push(...res); + } + + return rows.map((r) => this.transformFromDBData(r)); } public async deleteAllBySourceTableNames( @@ -42,9 +124,125 @@ export class ModelRepository queryOptions?: IQueryOptions, ) { const executer = queryOptions?.tx ? queryOptions.tx : this.knex; - const builder = executer(this.tableName) - .whereIn('source_table_name', sourceTableNames) - .delete(); - return await builder; + let deleted = 0; + for (const batch of this.toWhereInBatches(sourceTableNames)) { + deleted += await executer(this.tableName) + .whereIn('source_table_name', batch) + .delete(); + } + return deleted; } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } diff --git a/wren-ui/src/apollo/server/repositories/organizationInvitationRepository.ts b/wren-ui/src/apollo/server/repositories/organizationInvitationRepository.ts new file mode 100644 index 0000000000..80791fba1e --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/organizationInvitationRepository.ts @@ -0,0 +1,152 @@ +import { Knex } from 'knex'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; +import { Project } from './projectRepository'; +import { DataSourceName } from '@server/types'; + +export interface OrganizationInvitation { + id: number; + organizationId: number; + invitedByUserId?: number | null; + email: string; + organizationRole: string; + token: string; + status: string; + expiresAt: string; + acceptedAt?: string | null; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationInvitationProject { + id: number; + organizationInvitationId: number; + projectId: number; + permission: string; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationInvitationProjectAssignment { + id: number; + projectId: number; + permission: string; + project: Project; +} + +export interface OrganizationInvitationMapping extends OrganizationInvitation { + projects: OrganizationInvitationProjectAssignment[]; +} + +export interface IOrganizationInvitationRepository + extends IBasicRepository { + findByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export interface IOrganizationInvitationProjectRepository + extends IBasicRepository { + findByOrganizationInvitationId( + organizationInvitationId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export class OrganizationInvitationRepository + extends BaseRepository + implements IOrganizationInvitationRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_invitations' }); + } + + public async findByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ) { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = await executer(this.tableName) + .where({ organization_id: organizationId }) + .orderBy('created_at', 'desc'); + + return rows.map(this.transformFromDBData); + } +} + +export class OrganizationInvitationProjectRepository + extends BaseRepository + implements IOrganizationInvitationProjectRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_invitation_projects' }); + } + + public async findByOrganizationInvitationId( + organizationInvitationId: number, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = await executer('organization_invitation_projects') + .select( + 'organization_invitation_projects.*', + 'project.id as project__id', + 'project.display_name as project__display_name', + 'project.type as project__type', + 'project.version as project__version', + 'project.catalog as project__catalog', + 'project.schema as project__schema', + 'project.sample_dataset as project__sample_dataset', + 'project.connection_info as project__connection_info', + 'project.language as project__language', + 'project.query_id as project__query_id', + 'project.questions as project__questions', + 'project.questions_status as project__questions_status', + 'project.questions_error as project__questions_error', + ) + .join( + 'project', + 'organization_invitation_projects.project_id', + 'project.id', + ) + .where( + 'organization_invitation_projects.organization_invitation_id', + organizationInvitationId, + ) + .orderBy('project.id'); + + return rows.map((row) => ({ + id: row.id, + projectId: row.project_id, + permission: row.permission, + project: { + id: row.project__id, + displayName: row.project__display_name, + type: DataSourceName[row.project__type], + version: row.project__version, + catalog: row.project__catalog, + schema: row.project__schema, + sampleDataset: row.project__sample_dataset, + connectionInfo: + typeof row.project__connection_info === 'string' + ? JSON.parse(row.project__connection_info || '{}') + : row.project__connection_info, + language: row.project__language, + queryId: row.project__query_id, + questions: + typeof row.project__questions === 'string' + ? JSON.parse(row.project__questions || '[]') + : row.project__questions, + questionsStatus: row.project__questions_status, + questionsError: + typeof row.project__questions_error === 'string' + ? JSON.parse(row.project__questions_error || '{}') + : row.project__questions_error, + }, + })); + } +} diff --git a/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts b/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts new file mode 100644 index 0000000000..4506c3df63 --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/organizationMemberRepository.ts @@ -0,0 +1,176 @@ +import { Knex } from 'knex'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, + coerceBoolean, +} from './baseRepository'; +import { Project } from './projectRepository'; +import { DataSourceName } from '@server/types'; +import { RbacUser } from './rbacRepository'; + +export interface OrganizationMember { + id: number; + organizationId: number; + userId: number; + organizationRole: string; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMemberProject { + id: number; + organizationMemberId: number; + projectId: number; + permission: string; + createdAt: string; + updatedAt: string; +} + +export interface OrganizationMemberProjectAssignment { + id: number; + projectId: number; + permission: string; + project: Project; +} + +export interface OrganizationMemberMapping extends OrganizationMember { + user: RbacUser; + projects: OrganizationMemberProjectAssignment[]; +} + +export interface IOrganizationMemberRepository + extends IBasicRepository { + findMappingsByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export interface IOrganizationMemberProjectRepository + extends IBasicRepository { + findByOrganizationMemberId( + organizationMemberId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export class OrganizationMemberRepository + extends BaseRepository + implements IOrganizationMemberRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_members' }); + } + + public async findMappingsByOrganizationId( + organizationId: number, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = await executer('organization_members') + .select( + 'organization_members.*', + 'users.id as user__id', + 'users.name as user__name', + 'users.email as user__email', + 'users.external_id as user__external_id', + 'users.identity_provider as user__identity_provider', + 'users.is_active as user__is_active', + 'users.created_at as user__created_at', + 'users.updated_at as user__updated_at', + ) + .join('users', 'organization_members.user_id', 'users.id') + .where('organization_members.organization_id', organizationId) + .orderBy('users.email'); + + return rows.map((row) => ({ + id: row.id, + organizationId: row.organization_id, + userId: row.user_id, + organizationRole: row.organization_role, + createdAt: row.created_at, + updatedAt: row.updated_at, + user: { + id: row.user__id, + name: row.user__name, + email: row.user__email, + externalId: row.user__external_id, + identityProvider: row.user__identity_provider, + isActive: coerceBoolean(row.user__is_active), + createdAt: row.user__created_at, + updatedAt: row.user__updated_at, + }, + projects: [], + })); + } +} + +export class OrganizationMemberProjectRepository + extends BaseRepository + implements IOrganizationMemberProjectRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization_member_projects' }); + } + + public async findByOrganizationMemberId( + organizationMemberId: number, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const rows = await executer('organization_member_projects') + .select( + 'organization_member_projects.*', + 'project.id as project__id', + 'project.display_name as project__display_name', + 'project.type as project__type', + 'project.version as project__version', + 'project.catalog as project__catalog', + 'project.schema as project__schema', + 'project.sample_dataset as project__sample_dataset', + 'project.connection_info as project__connection_info', + 'project.language as project__language', + 'project.query_id as project__query_id', + 'project.questions as project__questions', + 'project.questions_status as project__questions_status', + 'project.questions_error as project__questions_error', + ) + .join('project', 'organization_member_projects.project_id', 'project.id') + .where( + 'organization_member_projects.organization_member_id', + organizationMemberId, + ) + .orderBy('project.id'); + + return rows.map((row) => ({ + id: row.id, + projectId: row.project_id, + permission: row.permission, + project: { + id: row.project__id, + displayName: row.project__display_name, + type: DataSourceName[row.project__type], + version: row.project__version, + catalog: row.project__catalog, + schema: row.project__schema, + sampleDataset: row.project__sample_dataset, + connectionInfo: + typeof row.project__connection_info === 'string' + ? JSON.parse(row.project__connection_info || '{}') + : row.project__connection_info, + language: row.project__language, + queryId: row.project__query_id, + questions: + typeof row.project__questions === 'string' + ? JSON.parse(row.project__questions || '[]') + : row.project__questions, + questionsStatus: row.project__questions_status, + questionsError: + typeof row.project__questions_error === 'string' + ? JSON.parse(row.project__questions_error || '{}') + : row.project__questions_error, + }, + })); + } +} diff --git a/wren-ui/src/apollo/server/repositories/organizationRepository.ts b/wren-ui/src/apollo/server/repositories/organizationRepository.ts new file mode 100644 index 0000000000..7f35ac49e4 --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/organizationRepository.ts @@ -0,0 +1,58 @@ +import { Knex } from 'knex'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, + coerceBoolean, +} from './baseRepository'; + +export interface Organization { + id: number; + name: string; + identifier: string; + description?: string | null; + isCurrent: boolean; + createdAt?: string; + updatedAt?: string; +} + +export interface IOrganizationRepository extends IBasicRepository { + getCurrentOrganization: ( + queryOptions?: IQueryOptions, + ) => Promise; + setCurrentOrganization: ( + id: number, + queryOptions?: IQueryOptions, + ) => Promise; +} + +export class OrganizationRepository + extends BaseRepository + implements IOrganizationRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'organization' }); + } + + protected override transformFromDBData = (data: any): Organization => { + const organization = this.defaultTransformFromDBData(data) as Organization; + return { + ...organization, + isCurrent: coerceBoolean(organization.isCurrent), + }; + }; + + public async getCurrentOrganization(queryOptions?: IQueryOptions) { + return await this.findOneBy({ isCurrent: true }, queryOptions); + } + + public async setCurrentOrganization(id: number, queryOptions?: IQueryOptions) { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + await executer(this.tableName).update({ is_current: false }); + const [result] = await executer(this.tableName) + .where({ id }) + .update({ is_current: true }) + .returning('*'); + return this.transformFromDBData(result); + } +} diff --git a/wren-ui/src/apollo/server/repositories/projectRepository.ts b/wren-ui/src/apollo/server/repositories/projectRepository.ts index 5dd054e273..5e48177ba1 100644 --- a/wren-ui/src/apollo/server/repositories/projectRepository.ts +++ b/wren-ui/src/apollo/server/repositories/projectRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -13,6 +17,10 @@ import { IbisRedshiftConnectionType, IbisDatabricksConnectionType, } from '@server/adaptors/ibisAdaptor'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('ProjectRepository'); +logger.level = 'debug'; export interface BIG_QUERY_CONNECTION_INFO { projectId: string; @@ -164,6 +172,11 @@ export interface RecommendationQuestionResult { sql: string; // validated sql for this question, can be used in generateAskDetail } +export enum WorkspaceProjectType { + AGENTIC = 'AGENTIC', + CLASSIC = 'CLASSIC', +} + export interface Project { id: number; // ID type: DataSourceName; // Project datasource type. ex: bigquery, mysql, postgresql, mongodb, etc @@ -180,10 +193,17 @@ export interface Project { questions?: RecommendationQuestionResult[]; questionsStatus?: string; questionsError?: object; + projectType?: WorkspaceProjectType; + isCurrent?: boolean; + createdAt?: Date; + updatedAt?: Date; } export interface IProjectRepository extends IBasicRepository { getCurrentProject: () => Promise; + listProjects: () => Promise; + findCurrentProject: () => Promise; + setCurrentProject: (projectId: string | number) => Promise; } export class ProjectRepository @@ -191,20 +211,137 @@ export class ProjectRepository implements IProjectRepository { private jsonTypeColumns = ['questions', 'questions_error', 'connection_info']; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'project' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = this.withTimestamps(data); + try { + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createOne( + await this.withMssqlId(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + const timestamped = data.map(this.withTimestamps); + try { + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions), + queryOptions, + ); + } catch (error) { + if (!this.shouldRetryManualId(error, timestamped)) { + throw error; + } + + return await super.createMany( + await this.withMssqlIds(timestamped, queryOptions, true), + queryOptions, + ); + } + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public async getCurrentProject() { + const currentProject = await this.findCurrentProject(); + if (currentProject) { + logger.debug( + `Resolved current project ${String(currentProject.id)} (${currentProject.type || 'unknown'})`, + ); + return currentProject; + } + const projects = await this.findAll({ order: 'id', - limit: 1, }); + logger.warn( + `Current project marker missing. Available projects: ${projects + .map((project) => String(project.id)) + .join(', ') || 'none'}`, + ); if (!projects.length) { throw new Error('No project found'); } - return projects[0]; + if (projects.length === 1) { + logger.warn( + `Repairing missing current project marker by selecting the only project ${String( + projects[0].id, + )}`, + ); + return await this.setCurrentProject(projects[0].id); + } + throw new Error('No current project selected'); + } + + public async listProjects() { + return await this.findAll({ + order: 'id', + }); + } + + public async findCurrentProject() { + return await this.findOneBy({ + isCurrent: true, + } as Partial); + } + + public async setCurrentProject(projectId: string | number) { + const tx = await this.transaction(); + try { + logger.debug(`Selecting current project ${String(projectId)}`); + await tx(this.tableName).update({ is_current: false }); + await tx(this.tableName) + .where({ id: this.normalizeProjectId(projectId) }) + .update({ is_current: true }); + const project = await this.findOneBy({ id: projectId } as Partial, { + tx, + }); + if (!project) { + throw new Error(`Project ${projectId} not found`); + } + await this.commit(tx); + logger.debug( + `Selected current project ${String(project.id)} (${project.type || 'unknown'})`, + ); + return project; + } catch (error) { + await this.rollback(tx); + throw error; + } } public override transformFromDBData: (data: any) => Project = (data: any) => { @@ -219,16 +356,22 @@ export class ProjectRepository if (key === 'type') { return DataSourceName[value]; } + if (key === 'project_type' && value) { + return WorkspaceProjectType[value]; + } return value; }); const camelCaseData = mapKeys(formattedData, (_value, key) => camelCase(key), ); + if (Object.prototype.hasOwnProperty.call(camelCaseData, 'isCurrent')) { + camelCaseData.isCurrent = this.normalizeProjectBoolean(camelCaseData.isCurrent); + } return camelCaseData as Project; }; - public override transformToDBData: (data: Project) => any = ( - data: Project, + public override transformToDBData: (data: Partial) => any = ( + data: Partial, ) => { if (!isPlainObject(data)) { throw new Error('Unexpected db data'); @@ -242,4 +385,151 @@ export class ProjectRepository }); return formattedData; }; + + private normalizeProjectBoolean(value: unknown): boolean { + if (typeof value === 'boolean') { + return value; + } + if (typeof value === 'number') { + return value === 1; + } + if (typeof value === 'string') { + return ['1', 'true'].includes(value.toLowerCase()); + } + return Boolean(value); + } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ + maxId: 'id', + }); + return { + ...data, + id: this.toNextIdValue(row?.maxId), + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + forceManualId = false, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (!forceManualId && (await this.hasIdentityId())) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ + maxId?: number | string | bigint | null; + }>({ + maxId: 'id', + }); + let nextId = BigInt(row?.maxId ?? 0) + 1n; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: this.serializeIdValue(nextId++), + }; + }); + }; + + private serializeIdValue(value: number | string | bigint) { + const bigintValue = BigInt(value); + return bigintValue <= BigInt(Number.MAX_SAFE_INTEGER) + ? Number(bigintValue) + : bigintValue.toString(); + } + + private toNextIdValue(maxId: number | string | bigint | null | undefined) { + const nextId = BigInt(maxId ?? 0) + 1n; + return this.serializeIdValue(nextId); + } + + private normalizeProjectId(projectId: string | number) { + return typeof projectId === 'string' ? projectId : projectId; + } + + private shouldRetryManualId = ( + error: unknown, + data: Partial | Partial[], + ): boolean => { + if (!this.isMssql()) { + return false; + } + + const hasMissingId = Array.isArray(data) + ? data.some((item) => item.id === undefined || item.id === null) + : data.id === undefined || data.id === null; + + if (!hasMissingId) { + return false; + } + + const message = + error instanceof Error + ? error.message + : typeof error === 'string' + ? error + : ''; + + return message.includes("Cannot insert the value NULL into column 'id'"); + }; } + diff --git a/wren-ui/src/apollo/server/repositories/rbacRepository.ts b/wren-ui/src/apollo/server/repositories/rbacRepository.ts new file mode 100644 index 0000000000..d6c2dbf55a --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/rbacRepository.ts @@ -0,0 +1,169 @@ +import { Knex } from 'knex'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, + coerceBoolean, +} from './baseRepository'; + +export interface Role { + id: number; + name: string; + description?: string | null; + createdAt: string; + updatedAt: string; +} + +export interface RbacUser { + id: number; + name: string; + email: string; + externalId?: string | null; + identityProvider?: string | null; + isActive: boolean; + createdAt: string; + updatedAt: string; +} + +export interface UserRole { + id: number; + userId: number; + roleId: number; + createdAt: string; + updatedAt: string; +} + +export interface UserRoleMapping extends UserRole { + user: RbacUser; + role: Role; +} + +export interface IRoleRepository extends IBasicRepository {} + +export interface IUserRepository extends IBasicRepository {} + +export interface IUserRoleRepository extends IBasicRepository { + findMappings(queryOptions?: IQueryOptions): Promise; + findMappingsByUserId( + userId: number, + queryOptions?: IQueryOptions, + ): Promise; + findMappingsByRoleId( + roleId: number, + queryOptions?: IQueryOptions, + ): Promise; +} + +export class RoleRepository + extends BaseRepository + implements IRoleRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'roles' }); + } +} + +export class UserRepository + extends BaseRepository + implements IUserRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'users' }); + } + + protected override transformFromDBData = (data: any): RbacUser => { + const user = this.defaultTransformFromDBData(data) as RbacUser; + return { + ...user, + isActive: coerceBoolean(user.isActive), + }; + }; +} + +export class UserRoleRepository + extends BaseRepository + implements IUserRoleRepository +{ + constructor(knexPg: Knex) { + super({ knexPg, tableName: 'user_roles' }); + } + + public async findMappings(queryOptions?: IQueryOptions) { + return this.queryMappings({}, queryOptions); + } + + public async findMappingsByUserId( + userId: number, + queryOptions?: IQueryOptions, + ) { + return this.queryMappings({ userId }, queryOptions); + } + + public async findMappingsByRoleId( + roleId: number, + queryOptions?: IQueryOptions, + ) { + return this.queryMappings({ roleId }, queryOptions); + } + + private async queryMappings( + filter: Partial, + queryOptions?: IQueryOptions, + ): Promise { + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const query = executer('user_roles') + .select( + 'user_roles.*', + 'users.id as user__id', + 'users.name as user__name', + 'users.email as user__email', + 'users.external_id as user__external_id', + 'users.identity_provider as user__identity_provider', + 'users.is_active as user__is_active', + 'users.created_at as user__created_at', + 'users.updated_at as user__updated_at', + 'roles.id as role__id', + 'roles.name as role__name', + 'roles.description as role__description', + 'roles.created_at as role__created_at', + 'roles.updated_at as role__updated_at', + ) + .join('users', 'user_roles.user_id', 'users.id') + .join('roles', 'user_roles.role_id', 'roles.id') + .orderBy('users.email') + .orderBy('roles.name'); + + if (filter.userId) { + query.where('user_roles.user_id', filter.userId); + } + if (filter.roleId) { + query.where('user_roles.role_id', filter.roleId); + } + + const rows = await query; + return rows.map((row) => ({ + id: row.id, + userId: row.user_id, + roleId: row.role_id, + createdAt: row.created_at, + updatedAt: row.updated_at, + user: { + id: row.user__id, + name: row.user__name, + email: row.user__email, + externalId: row.user__external_id, + identityProvider: row.user__identity_provider, + isActive: coerceBoolean(row.user__is_active), + createdAt: row.user__created_at, + updatedAt: row.user__updated_at, + }, + role: { + id: row.role__id, + name: row.role__name, + description: row.role__description, + createdAt: row.role__created_at, + updatedAt: row.role__updated_at, + }, + })); + } +} diff --git a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts index 9aba1b6bc7..d02aa4f6fc 100644 --- a/wren-ui/src/apollo/server/repositories/relationshipRepository.ts +++ b/wren-ui/src/apollo/server/repositories/relationshipRepository.ts @@ -15,6 +15,8 @@ export interface Relation { fromColumnId: number; // from column id, "{fromColumn} {joinType} {toColumn}" toColumnId: number; // to column id, "{fromColumn} {joinType} {toColumn}" properties: string | null; // Model properties, a json string, the description should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface ExtraRelationInfo { @@ -62,49 +64,89 @@ export class RelationRepository extends BaseRepository implements IRelationRepository { + private hasIdentityIdPromise?: Promise; + constructor(knexPg: Knex) { super({ knexPg, tableName: 'relation' }); } + public async createOne(data: Partial): Promise { + if (!this.isMssql() || (await this.hasIdentityId())) { + return super.createOne(this.withTimestamps(data)); + } + + const id = await this.nextRelationId(); + return super.createOne(this.withTimestamps({ id, ...data })); + } + + public async createMany(data: Partial[]): Promise { + if (data.length === 0) { + return []; + } + + const dataWithTimestamps = data.map((relation) => + this.withTimestamps(relation), + ); + + if (!this.isMssql() || (await this.hasIdentityId())) { + return super.createMany(dataWithTimestamps); + } + + const startId = await this.nextRelationId(); + return super.createMany( + dataWithTimestamps.map((relation, index) => ({ + id: startId + index, + ...relation, + })), + ); + } + + public async updateOne(id: number, data: Partial) { + return super.updateOne(id, { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }); + } + public async findRelationsBy( { columnIds, modelIds }, queryOptions?: IQueryOptions, ) { - let executer = this.knex; - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - executer = tx; - } - // select the leftModel name and rightModel name along with relation - const builder = executer(this.tableName) - .join( - 'model_column AS fmc', - `${this.tableName}.from_column_id`, - '=', - 'fmc.id', - ) - .join( - 'model_column AS tmc', - `${this.tableName}.to_column_id`, - '=', - 'tmc.id', + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const selectRows = (builder) => + builder.select( + `${this.tableName}.*`, + 'fmc.model_id AS fromModelId', + 'tmc.model_id AS toModelId', ); + const rows = []; if (columnIds && columnIds.length > 0) { - builder - .whereIn(`${this.tableName}.from_column_id`, columnIds) - .orWhereIn(`${this.tableName}.to_column_id`, columnIds); + for (const batch of this.toDualWhereInBatches(columnIds)) { + const result = await selectRows(this.relationJoinBuilder(executer)).where( + (builder) => + builder + .whereIn(`${this.tableName}.from_column_id`, batch) + .orWhereIn(`${this.tableName}.to_column_id`, batch), + ); + rows.push(...result); + } + return rows.map((r) => this.transformFromDBData(r)); } + if (modelIds && modelIds.length > 0) { - builder - .whereIn('fmc.model_id', modelIds) - .orWhereIn('tmc.model_id', modelIds); + for (const batch of this.toDualWhereInBatches(modelIds)) { + const result = await selectRows(this.relationJoinBuilder(executer)).where( + (builder) => + builder + .whereIn('fmc.model_id', batch) + .orWhereIn('tmc.model_id', batch), + ); + rows.push(...result); + } + return rows.map((r) => this.transformFromDBData(r)); } - const result = await builder.select( - `${this.tableName}.*`, - 'fmc.model_id AS fromModelId', - 'tmc.model_id AS toModelId', - ); + const result = await selectRows(this.relationJoinBuilder(executer)); return result.map((r) => this.transformFromDBData(r)); } @@ -115,9 +157,11 @@ export class RelationRepository executer = tx; } - const result = await executer(this.tableName) - .whereIn('id', ids) - .select('*'); + const result = []; + for (const batch of this.toWhereInBatches(ids)) { + const rows = await executer(this.tableName).whereIn('id', batch).select('*'); + result.push(...rows); + } return result.map((r) => this.transformFromDBData(r)); } @@ -125,70 +169,69 @@ export class RelationRepository columnIds: number[], queryOptions?: IQueryOptions, ) { - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - await tx(this.tableName) - .whereIn('from_column_id', columnIds) - .orWhereIn('to_column_id', columnIds) + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + for (const batch of this.toDualWhereInBatches(columnIds)) { + await executer(this.tableName) + .where((builder) => + builder + .whereIn('from_column_id', batch) + .orWhereIn('to_column_id', batch), + ) .delete(); - return; } - await this.knex(this.tableName) - .whereIn('from_column_id', columnIds) - .orWhereIn('to_column_id', columnIds) - .delete(); } public async findRelationInfoBy(filter, queryOptions) { const { projectId, columnIds, modelIds } = filter; - let executer = this.knex; - if (queryOptions && queryOptions.tx) { - const { tx } = queryOptions; - executer = tx; - } - // select the leftModel name and rightModel name along with relation - const builder = executer(this.tableName) - .join( - 'model_column AS fmc', - `${this.tableName}.from_column_id`, - '=', - 'fmc.id', - ) - .join( - 'model_column AS tmc', - `${this.tableName}.to_column_id`, - '=', - 'tmc.id', - ) - .join('model AS fm', 'fmc.model_id', '=', 'fm.id') - .join('model AS tm', 'tmc.model_id', '=', 'tm.id'); + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const selectRows = (builder) => + builder.select( + `${this.tableName}.*`, + 'fm.id AS fromModelId', + 'fm.reference_name AS fromModelName', + 'fm.display_name AS fromModelDisplayName', + 'tm.id AS toModelId', + 'tm.reference_name AS toModelName', + 'tm.display_name AS toModelDisplayName', + 'fmc.reference_name AS fromColumnName', + 'fmc.display_name AS fromColumnDisplayName', + 'tmc.reference_name AS toColumnName', + 'tmc.display_name AS toColumnDisplayName', + ); if (projectId) { + const builder = this.relationInfoJoinBuilder(executer); builder.where(`${this.tableName}.project_id`, projectId); - } else if (columnIds && columnIds.length > 0) { - builder - .whereIn(`${this.tableName}.from_column_id`, columnIds) - .orWhereIn(`${this.tableName}.to_column_id`, columnIds); + const result = await selectRows(builder); + return result.map((r) => this.transformFromDBData(r)) as RelationInfo[]; + } + + const rows = []; + if (columnIds && columnIds.length > 0) { + for (const batch of this.toDualWhereInBatches(columnIds)) { + const result = await selectRows(this.relationInfoJoinBuilder(executer)).where( + (builder) => + builder + .whereIn(`${this.tableName}.from_column_id`, batch) + .orWhereIn(`${this.tableName}.to_column_id`, batch), + ); + rows.push(...result); + } } else if (modelIds && modelIds.length > 0) { - builder - .whereIn('fmc.model_id', modelIds) - .orWhereIn('tmc.model_id', modelIds); + for (const batch of this.toDualWhereInBatches(modelIds)) { + const result = await selectRows(this.relationInfoJoinBuilder(executer)).where( + (builder) => + builder + .whereIn('fmc.model_id', batch) + .orWhereIn('tmc.model_id', batch), + ); + rows.push(...result); + } + } else { + rows.push(...(await selectRows(this.relationInfoJoinBuilder(executer)))); } - const result = await builder.select( - `${this.tableName}.*`, - 'fm.id AS fromModelId', - 'fm.reference_name AS fromModelName', - 'fm.display_name AS fromModelDisplayName', - 'tm.id AS toModelId', - 'tm.reference_name AS toModelName', - 'tm.display_name AS toModelDisplayName', - 'fmc.reference_name AS fromColumnName', - 'fmc.display_name AS fromColumnDisplayName', - 'tmc.reference_name AS toColumnName', - 'tmc.display_name AS toColumnDisplayName', - ); - return result.map((r) => this.transformFromDBData(r)) as RelationInfo[]; + return rows.map((r) => this.transformFromDBData(r)) as RelationInfo[]; } public async findExistedRelationBetweenModels(relation: RelationData) { @@ -220,4 +263,79 @@ export class RelationRepository const result = await query; return result.map((r) => this.transformFromDBData(r)) as RelationInfo[]; } + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private toDualWhereInBatches(values: TValue[]) { + if (!this.isMssql()) { + return this.toWhereInBatches(values); + } + + const batchSize = Math.max( + Math.floor(this.getMssqlWhereInBatchSize() / 2), + 1, + ); + const batches: TValue[][] = []; + for (let index = 0; index < values.length; index += batchSize) { + batches.push(values.slice(index, index + batchSize)); + } + return batches; + } + + private relationJoinBuilder(executer: Knex | Knex.Transaction) { + return executer(this.tableName) + .join( + 'model_column AS fmc', + `${this.tableName}.from_column_id`, + '=', + 'fmc.id', + ) + .join( + 'model_column AS tmc', + `${this.tableName}.to_column_id`, + '=', + 'tmc.id', + ); + } + + private relationInfoJoinBuilder(executer: Knex | Knex.Transaction) { + return this.relationJoinBuilder(executer) + .join('model AS fm', 'fmc.model_id', '=', 'fm.id') + .join('model AS tm', 'tmc.model_id', '=', 'tm.id'); + } + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private async hasIdentityId() { + this.hasIdentityIdPromise ??= this.knex('sys.columns as c') + .join('sys.tables as t', 'c.object_id', 't.object_id') + .where('t.name', this.tableName) + .where('c.name', 'id') + .select( + this.knex.raw( + 'COLUMNPROPERTY(c.object_id, c.name, ?) as isIdentity', + ['IsIdentity'], + ), + ) + .first() + .then((row) => Number(row?.isIdentity ?? 0) === 1); + + return this.hasIdentityIdPromise; + } + + private async nextRelationId() { + const row = await this.knex(this.tableName) + .max<{ maxId?: number | string }>('id as maxId') + .first(); + + return Number(row?.maxId ?? 0) + 1; + } } diff --git a/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts b/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts index 0a7a7d379d..846e578c50 100644 --- a/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts +++ b/wren-ui/src/apollo/server/repositories/schemaChangeRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -17,8 +21,8 @@ export interface SchemaChange { projectId: number; // Reference to project.id change: DataSourceSchemaChange; // Schema change resolve: DataSourceSchemaResolve; // Save resolve - createdAt: string; // Created at - updateAt: string; // Updated at + createdAt?: Date | string; // Created at + updatedAt?: Date | string; // Updated at } export interface ISchemaChangeRepository @@ -34,12 +38,35 @@ export class SchemaChangeRepository super({ knexPg, tableName: 'schema_change' }); } + public async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ) { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ) { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public async findLastSchemaChange(projectId: number) { const res = await this.knex .select('*') .from(this.tableName) .where(this.transformToDBData({ projectId })) .orderBy('created_at', 'desc') + .orderBy('id', 'desc') .first(); return (res && this.transformFromDBData(res)) || null; } @@ -75,4 +102,13 @@ export class SchemaChangeRepository }) as SchemaChange; return formattedData; }; + + private withTimestamps(data: Partial): Partial { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + } } diff --git a/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts b/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts index 9f02f48232..e4b7f3c945 100644 --- a/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts +++ b/wren-ui/src/apollo/server/repositories/sqlPairRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; export interface SqlPair { id: number; // ID @@ -19,4 +23,35 @@ export class SqlPairRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'sql_pair' }); } + + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date().toISOString(), + }, + queryOptions, + ); + } + + private withTimestamps = (data: Partial): Partial => { + const now = new Date().toISOString(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; } diff --git a/wren-ui/src/apollo/server/repositories/tests/deployLogRepository.test.ts b/wren-ui/src/apollo/server/repositories/tests/deployLogRepository.test.ts new file mode 100644 index 0000000000..d6c641ab57 --- /dev/null +++ b/wren-ui/src/apollo/server/repositories/tests/deployLogRepository.test.ts @@ -0,0 +1,39 @@ +import { Deploy, DeployLogRepository } from '../deployLogRepository'; + +class TestDeployLogRepository extends DeployLogRepository { + public toDBData(data: Partial) { + return this.transformToDBData(data); + } + + public fromDBData(data: any) { + return this.transformFromDBData(data); + } +} + +describe('DeployLogRepository', () => { + const repository = new TestDeployLogRepository({ + client: { config: { client: 'mssql' } }, + } as any); + + it('serializes manifest before writing deploy logs', () => { + const manifest = { + models: [{ name: 'dbo_xStageNewOrders', columns: [{ name: 'CustName' }] }], + }; + + expect(repository.toDBData({ manifest, projectId: 11 })).toEqual({ + manifest: JSON.stringify(manifest), + project_id: 11, + }); + }); + + it('keeps string manifests unchanged and parses them after reading', () => { + const manifest = JSON.stringify({ + models: [{ name: 'dbo_ytblRefund', columns: [{ name: 'RefundDate' }] }], + }); + + expect(repository.toDBData({ manifest } as any)).toEqual({ manifest }); + expect(repository.fromDBData({ id: 1, manifest }).manifest).toEqual( + JSON.parse(manifest), + ); + }); +}); diff --git a/wren-ui/src/apollo/server/repositories/threadRepository.ts b/wren-ui/src/apollo/server/repositories/threadRepository.ts index 65720ba5fb..d3fc858ae4 100644 --- a/wren-ui/src/apollo/server/repositories/threadRepository.ts +++ b/wren-ui/src/apollo/server/repositories/threadRepository.ts @@ -1,5 +1,9 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, +} from './baseRepository'; import { camelCase, isPlainObject, @@ -24,6 +28,8 @@ export interface Thread { questions?: ThreadRecommendationQuestionResult[]; // Recommended questions questionsStatus?: string; // Status of the recommended questions questionsError?: object; // Error of the recommended questions + createdAt?: Date; + updatedAt?: Date; } export interface IThreadRepository extends IBasicRepository { @@ -35,11 +41,47 @@ export class ThreadRepository implements IThreadRepository { private readonly jsonbColumns = ['questions', 'questionsError']; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'thread' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne( + await this.withMssqlId(this.withTimestamps(data), queryOptions), + queryOptions, + ); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany( + await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), + queryOptions, + ); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + public async listAllTimeDescOrder(projectId: number): Promise { const threads = await this.knex(this.tableName) .where(this.transformToDBData({ projectId })) @@ -78,4 +120,89 @@ export class ThreadRepository }); return mapKeys(transformedData, (_value, key) => snakeCase(key)); }; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; } diff --git a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts index f51d1eda91..a425068bc9 100644 --- a/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts +++ b/wren-ui/src/apollo/server/repositories/threadResponseRepository.ts @@ -73,6 +73,8 @@ export interface ThreadResponse { breakdownDetail?: ThreadResponseBreakdownDetail; // Thread response breakdown detail chartDetail?: ThreadResponseChartDetail; // Thread response chart detail adjustment?: ThreadResponseAdjustment; // Thread response adjustment + createdAt?: Date; + updatedAt?: Date; } export interface IThreadResponseRepository @@ -93,11 +95,32 @@ export class ThreadResponseRepository 'chartDetail', 'adjustment', ]; + private hasIdentityIdPromise?: Promise; constructor(knexPg: Knex) { super({ knexPg, tableName: 'thread_response' }); } + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne( + await this.withMssqlId(this.withTimestamps(data), queryOptions), + queryOptions, + ); + } + + public override async createMany( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise { + return super.createMany( + await this.withMssqlIds(data.map(this.withTimestamps), queryOptions), + queryOptions, + ); + } + public async getResponsesWithThread(threadId: number, limit?: number) { const query = this.knex(this.tableName) .select('thread_response.*') @@ -105,7 +128,9 @@ export class ThreadResponseRepository .leftJoin('thread', 'thread.id', 'thread_response.thread_id'); if (limit) { - query.orderBy('created_at', 'desc').limit(limit); + query.orderBy('thread_response.id', 'desc').limit(limit); + } else { + query.orderBy('thread_response.id', 'asc'); } return (await query) @@ -146,10 +171,10 @@ export class ThreadResponseRepository data: Partial<{ status: AskResultStatus; sql: string; - viewId: number; - answerDetail: ThreadResponseAnswerDetail; - breakdownDetail: ThreadResponseBreakdownDetail; - chartDetail: ThreadResponseChartDetail; + viewId: number | null; + answerDetail: ThreadResponseAnswerDetail | null; + breakdownDetail: ThreadResponseBreakdownDetail | null; + chartDetail: ThreadResponseChartDetail | null; adjustment: ThreadResponseAdjustment; }>, queryOptions?: IQueryOptions, @@ -157,17 +182,29 @@ export class ThreadResponseRepository const transformedData = { status: data.status ? data.status : undefined, sql: data.sql ? data.sql : undefined, - viewId: data.viewId ? data.viewId : undefined, - answerDetail: data.answerDetail - ? JSON.stringify(data.answerDetail) + viewId: Object.prototype.hasOwnProperty.call(data, 'viewId') + ? data.viewId : undefined, - breakdownDetail: data.breakdownDetail - ? JSON.stringify(data.breakdownDetail) + answerDetail: Object.prototype.hasOwnProperty.call(data, 'answerDetail') + ? data.answerDetail + ? JSON.stringify(data.answerDetail) + : null : undefined, - chartDetail: data.chartDetail - ? JSON.stringify(data.chartDetail) + breakdownDetail: Object.prototype.hasOwnProperty.call( + data, + 'breakdownDetail', + ) + ? data.breakdownDetail + ? JSON.stringify(data.breakdownDetail) + : null + : undefined, + chartDetail: Object.prototype.hasOwnProperty.call(data, 'chartDetail') + ? data.chartDetail + ? JSON.stringify(data.chartDetail) + : null : undefined, adjustment: data.adjustment ? JSON.stringify(data.adjustment) : undefined, + updatedAt: new Date(), }; const executer = queryOptions?.tx ? queryOptions.tx : this.knex; const [result] = await executer(this.tableName) @@ -195,4 +232,91 @@ export class ThreadResponseRepository }) as ThreadResponse; return formattedData; }; + + private withTimestamps = ( + data: Partial, + ): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; + + private isMssql = () => + String(this.knex.client.config.client || '').toLowerCase() === 'mssql'; + + private hasIdentityId = async (): Promise => { + if (!this.isMssql()) { + return true; + } + + if (!this.hasIdentityIdPromise) { + this.hasIdentityIdPromise = this.knex('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ + TABLE_SCHEMA: 'dbo', + TABLE_NAME: this.tableName, + COLUMN_NAME: 'id', + }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ) + .first() + .then(Boolean); + } + + return this.hasIdentityIdPromise; + }; + + private withMssqlId = async ( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise> => { + if ( + (data.id !== undefined && data.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + return { + ...data, + id: Number(row?.maxId || 0) + 1, + }; + }; + + private withMssqlIds = async ( + data: Partial[], + queryOptions?: IQueryOptions, + ): Promise[]> => { + if ( + data.every((item) => item.id !== undefined && item.id !== null) || + !this.isMssql() || + (await this.hasIdentityId()) + ) { + return data; + } + + const executer = queryOptions?.tx ? queryOptions.tx : this.knex; + const [row] = await executer(this.tableName).max<{ maxId?: number }>({ + maxId: 'id', + }); + let nextId = Number(row?.maxId || 0) + 1; + return data.map((item) => { + if (item.id !== undefined && item.id !== null) { + return item; + } + return { + ...item, + id: nextId++, + }; + }); + }; } diff --git a/wren-ui/src/apollo/server/repositories/viewRepository.ts b/wren-ui/src/apollo/server/repositories/viewRepository.ts index 39ed245c35..c056ebd0e0 100644 --- a/wren-ui/src/apollo/server/repositories/viewRepository.ts +++ b/wren-ui/src/apollo/server/repositories/viewRepository.ts @@ -1,5 +1,10 @@ import { Knex } from 'knex'; -import { BaseRepository, IBasicRepository } from './baseRepository'; +import { + BaseRepository, + IBasicRepository, + IQueryOptions, + coerceBoolean, +} from './baseRepository'; export interface View { id: number; // ID @@ -9,6 +14,8 @@ export interface View { cached: boolean; // View is cached or not refreshTime?: string; // Contain a number followed by a time unit (ns, us, ms, s, m, h, d). For example, "2h" properties?: string; // View properties, a json string, the description and displayName should be stored here + createdAt?: Date; + updatedAt?: Date; } export interface IViewRepository extends IBasicRepository {} @@ -20,4 +27,43 @@ export class ViewRepository constructor(knexPg: Knex) { super({ knexPg, tableName: 'view' }); } + + public override async createOne( + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.createOne(this.withTimestamps(data), queryOptions); + } + + public override async updateOne( + id: string | number, + data: Partial, + queryOptions?: IQueryOptions, + ): Promise { + return super.updateOne( + id, + { + ...data, + updatedAt: data.updatedAt ?? new Date(), + }, + queryOptions, + ); + } + + protected override transformFromDBData = (data: any): View => { + const view = this.defaultTransformFromDBData(data) as View; + return { + ...view, + cached: coerceBoolean(view.cached), + }; + }; + + private withTimestamps = (data: Partial): Partial => { + const now = new Date(); + return { + ...data, + createdAt: data.createdAt ?? now, + updatedAt: data.updatedAt ?? now, + }; + }; } diff --git a/wren-ui/src/apollo/server/resolvers.ts b/wren-ui/src/apollo/server/resolvers.ts index 4c4c9af187..11e3487602 100644 --- a/wren-ui/src/apollo/server/resolvers.ts +++ b/wren-ui/src/apollo/server/resolvers.ts @@ -8,8 +8,9 @@ import { DashboardResolver } from './resolvers/dashboardResolver'; import { SqlPairResolver } from './resolvers/sqlPairResolver'; import { InstructionResolver } from './resolvers/instructionResolver'; import { ApiHistoryResolver } from './resolvers/apiHistoryResolver'; +import { RbacResolver } from './resolvers/rbacResolver'; import { convertColumnType } from '@server/utils'; -import { DialectSQLScalar } from './scalars'; +import { BigIntStringScalar, DialectSQLScalar } from './scalars'; const projectResolver = new ProjectResolver(); const modelResolver = new ModelResolver(); @@ -20,9 +21,11 @@ const dashboardResolver = new DashboardResolver(); const sqlPairResolver = new SqlPairResolver(); const instructionResolver = new InstructionResolver(); const apiHistoryResolver = new ApiHistoryResolver(); +const rbacResolver = new RbacResolver(); const resolvers = { JSON: GraphQLJSON, DialectSQL: DialectSQLScalar, + BigIntString: BigIntStringScalar, Query: { listDataSourceTables: projectResolver.listDataSourceTables, autoGenerateRelation: projectResolver.autoGenerateRelation, @@ -54,6 +57,8 @@ const resolvers = { // Settings settings: projectResolver.getSettings, getMDL: modelResolver.getMDL, + modelingSemanticsResult: modelResolver.getModelingSemanticsResult, + modelingRelationshipsResult: modelResolver.getModelingRelationshipsResult, // Learning learningRecord: learningResolver.getLearningRecord, @@ -75,6 +80,11 @@ const resolvers = { // API History apiHistory: apiHistoryResolver.getApiHistory, + + // Administration / RBAC + roles: rbacResolver.listRoles, + users: rbacResolver.listUsers, + userRoleMappings: rbacResolver.listUserRoleMappings, }, Mutation: { deploy: modelResolver.deploy, @@ -141,6 +151,10 @@ const resolvers = { previewViewData: modelResolver.previewViewData, validateView: modelResolver.validateView, updateViewMetadata: modelResolver.updateViewMetadata, + generateModelingSemantics: modelResolver.generateModelingSemantics, + saveModelingSemantics: modelResolver.saveModelingSemantics, + generateModelingRelationships: modelResolver.generateModelingRelationships, + saveModelingRelationships: modelResolver.saveModelingRelationships, // Settings resetCurrentProject: projectResolver.resetCurrentProject, @@ -149,6 +163,7 @@ const resolvers = { // preview previewSql: modelResolver.previewSql, + dryPlanSql: modelResolver.dryPlanSql, // Learning saveLearningRecord: learningResolver.saveLearningRecord, @@ -177,6 +192,15 @@ const resolvers = { createInstruction: instructionResolver.createInstruction, updateInstruction: instructionResolver.updateInstruction, deleteInstruction: instructionResolver.deleteInstruction, + + // Administration / RBAC + createRole: rbacResolver.createRole, + updateRole: rbacResolver.updateRole, + createUser: rbacResolver.createUser, + updateUser: rbacResolver.updateUser, + assignRoleToUser: rbacResolver.assignRoleToUser, + updateUserRoles: rbacResolver.updateUserRoles, + removeRoleFromUser: rbacResolver.removeRoleFromUser, }, ThreadResponse: askingResolver.getThreadResponseNestedResolver(), DetailStep: askingResolver.getDetailStepNestedResolver(), @@ -196,6 +220,9 @@ const resolvers = { // Add ApiHistoryResponse nested resolvers ApiHistoryResponse: apiHistoryResolver.getApiHistoryNestedResolver(), + + Role: rbacResolver.getRoleNestedResolver(), + User: rbacResolver.getUserNestedResolver(), }; export default resolvers; diff --git a/wren-ui/src/apollo/server/resolvers/askingResolver.ts b/wren-ui/src/apollo/server/resolvers/askingResolver.ts index e674c0e278..a11b5b539f 100644 --- a/wren-ui/src/apollo/server/resolvers/askingResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/askingResolver.ts @@ -19,6 +19,7 @@ import { safeFormatSQL } from '@server/utils/sqlFormat'; import { AskingDetailTaskInput, constructCteSql, + RecommendQuestionResultStatus, ThreadRecommendQuestionResult, } from '../services/askingService'; import { @@ -65,6 +66,17 @@ export interface AskingTask { queryId?: string; } +type ThreadCreationArgs = { + question?: string; + taskId?: string; + sql?: string; + answerContent?: string; + answerStatus?: 'FINISHED' | 'FAILED' | 'INTERRUPTED' | null; + answerErrorCode?: string; + answerErrorShortMessage?: string; + answerErrorMessage?: string; +}; + // DetailedThread is a type that represents a detailed thread, which is a thread with responses. export interface DetailedThread { id: number; // ID @@ -78,7 +90,7 @@ export interface RecommendedQuestionsTask { category: string; sql: string; }[]; - status: RecommendationQuestionStatus; + status: RecommendationQuestionStatus | RecommendQuestionResultStatus; error: WrenAIError | null; } @@ -176,6 +188,7 @@ export class AskingResolver { args: { data: { question: string; threadId?: number } }, ctx: IContext, ): Promise { + const startedAt = Date.now(); const { question, threadId } = args.data; const project = await ctx.projectService.getCurrentProject(); @@ -185,6 +198,11 @@ export class AskingResolver { threadId, language: WrenAILanguage[project.language] || WrenAILanguage.EN, }); + logger.info( + `Ask timing stage=frontend_request project_id=${project.id} thread_id=${ + threadId ?? '' + } elapsed_ms=${Date.now() - startedAt}`, + ); ctx.telemetry.sendEvent(TelemetryEvent.HOME_ASK_CANDIDATE, { question, taskId: task.id, @@ -205,10 +223,13 @@ export class AskingResolver { public async getAskingTask( _root: any, - args: { taskId: string }, + args: { taskId?: string | null }, ctx: IContext, ): Promise { const { taskId } = args; + if (!taskId) { + return null; + } const askingService = ctx.askingService; const askResult = await askingService.getAskingTask(taskId); @@ -243,14 +264,7 @@ export class AskingResolver { public async createThread( _root: any, - args: { - data: { - question?: string; - taskId?: string; - // if we use recommendation questions, sql will be provided - sql?: string; - }; - }, + args: { data: ThreadCreationArgs }, ctx: IContext, ): Promise { const { data } = args; @@ -275,6 +289,20 @@ export class AskingResolver { threadInput = data; } + if (data.answerStatus) { + threadInput.answerDetail = { + status: data.answerStatus, + content: data.answerContent, + error: data.answerErrorCode + ? { + code: data.answerErrorCode as any, + shortMessage: data.answerErrorShortMessage, + message: data.answerErrorMessage, + } + : undefined, + }; + } + const eventName = TelemetryEvent.HOME_CREATE_THREAD; try { const thread = await askingService.createThread(threadInput); @@ -385,15 +413,7 @@ export class AskingResolver { public async createThreadResponse( _root: any, - args: { - threadId: number; - data: { - question?: string; - taskId?: string; - // if we use recommendation questions, sql will be provided - sql?: string; - }; - }, + args: { threadId: number; data: ThreadCreationArgs }, ctx: IContext, ): Promise { const { threadId, data } = args; @@ -419,6 +439,20 @@ export class AskingResolver { threadResponseInput = data; } + if (data.answerStatus) { + threadResponseInput.answerDetail = { + status: data.answerStatus, + content: data.answerContent, + error: data.answerErrorCode + ? { + code: data.answerErrorCode as any, + shortMessage: data.answerErrorShortMessage, + message: data.answerErrorMessage, + } + : undefined, + }; + } + try { const response = await askingService.createThreadResponse( threadResponseInput, @@ -543,10 +577,13 @@ export class AskingResolver { public async getAdjustmentTask( _root: any, - args: { taskId: string }, + args: { taskId?: string | null }, ctx: IContext, ): Promise { const { taskId } = args; + if (!taskId) { + return null; + } const askingService = ctx.askingService; const adjustmentTask = await askingService.getAdjustmentTask(taskId); return { @@ -617,10 +654,13 @@ export class AskingResolver { public async getResponse( _root: any, - args: { responseId: number }, + args: { responseId?: number | null }, ctx: IContext, ): Promise { const { responseId } = args; + if (!responseId) { + return null; + } const askingService = ctx.askingService; const response = await askingService.getResponse(responseId); @@ -665,10 +705,17 @@ export class AskingResolver { public async getInstantRecommendedQuestions( _root: any, - args: { taskId: string }, + args: { taskId?: string | null }, ctx: IContext, ): Promise { const { taskId } = args; + if (!taskId) { + return { + questions: [], + status: RecommendQuestionResultStatus.NOT_STARTED, + error: null, + }; + } const askingService = ctx.askingService; const result = await askingService.getInstantRecommendedQuestions(taskId); return { diff --git a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts index 8411b5409c..2f4f88db80 100644 --- a/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/dashboardResolver.ts @@ -16,9 +16,12 @@ import { DashboardSchedule, PreviewItemResponse, } from '@server/models/dashboard'; +import { Manifest } from '@server/mdl/type'; const logger = getLogger('DashboardResolver'); logger.level = 'debug'; +const DASHBOARD_SNAPSHOT_TIMEOUT_MS = 5000; +const DASHBOARD_PREVIEW_TIMEOUT_MS = 15000; export class DashboardResolver { constructor() { @@ -48,6 +51,10 @@ export class DashboardResolver { if (!dashboard) { throw new Error('Dashboard not found.'); } + const project = await ctx.projectService.getCurrentProject(); + logger.debug( + `Resolving dashboard ${String(dashboard.id)} for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const schedule = ctx.dashboardService.parseCronExpression(dashboard); const items = await ctx.dashboardService.getDashboardItems(dashboard.id); return { @@ -69,6 +76,10 @@ export class DashboardResolver { if (!dashboard) { throw new Error('Dashboard not found.'); } + const project = await ctx.projectService.getCurrentProject(); + logger.debug( + `Resolving dashboard items for dashboard ${String(dashboard.id)} in project ${String(project.id)} (${project.type || 'unknown'})`, + ); return await ctx.dashboardService.getDashboardItems(dashboard.id); } @@ -77,14 +88,20 @@ export class DashboardResolver { args: { data: { itemType: DashboardItemType; responseId: number } }, ctx: IContext, ): Promise { - const { responseId, itemType } = args.data; + const { responseId } = args.data; + const itemType = this.normalizeDashboardItemType(args.data.itemType); const dashboard = await ctx.dashboardService.getCurrentDashboard(); const response = await ctx.askingService.getResponse(responseId); + const project = await ctx.projectService.getCurrentProject(); + + logger.debug( + `Pinning response ${responseId} into dashboard ${String(dashboard.id)} for project ${String(project.id)} (${project.type || 'unknown'})`, + ); if (!response) { throw new Error(`Thread response not found. responseId: ${responseId}`); } - if (!Object.keys(ChartType).includes(itemType)) { + if (!itemType) { throw new Error(`Chart type not supported. responseId: ${responseId}`); } if (!response.chartDetail?.chartSchema) { @@ -92,30 +109,56 @@ export class DashboardResolver { `Chart schema not found in thread response. responseId: ${responseId}`, ); } + if (!response.sql) { + throw new Error(`Chart SQL not found in thread response. responseId: ${responseId}`); + } - // query with cache enabled - const project = await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; - await ctx.queryService.preview(response.sql, { - project, - manifest: mdl, - limit: DEFAULT_PREVIEW_LIMIT, - cacheEnabled: true, - refresh: true, - }); + const previewDataSnapshot = await this.capturePreviewSnapshot( + ctx, + response.sql, + ); - return await ctx.dashboardService.createDashboardItem({ + const dashboardItem = await ctx.dashboardService.createDashboardItem({ dashboardId: dashboard.id, type: itemType, sql: response.sql, chartSchema: response.chartDetail?.chartSchema, + previewDataSnapshot, }); + + // Warm dashboard cache asynchronously so datasource latency does not block pinning. + void this.warmDashboardCache(ctx, response.sql, dashboardItem.id); + + return dashboardItem; + } + + private normalizeDashboardItemType( + itemType: DashboardItemType | ChartType | string, + ): DashboardItemType | null { + const rawValue = String(itemType || ''); + const normalizedKey = rawValue.toUpperCase() as keyof typeof DashboardItemType; + const normalizedValue = DashboardItemType[normalizedKey]; + if (normalizedValue) { + return normalizedValue; + } + + const chartTypeValue = Object.values(ChartType).find( + (value) => value === rawValue, + ); + if (!chartTypeValue) { + return null; + } + + return ( + DashboardItemType[ + chartTypeValue.toUpperCase() as keyof typeof DashboardItemType + ] || null + ); } public async updateDashboardItem( _root: any, - args: { where: { id: number }; data: { displayName: string } }, + args: { where: { id: string | number }; data: { displayName: string } }, ctx: IContext, ): Promise { const { id } = args.where; @@ -129,7 +172,7 @@ export class DashboardResolver { public async deleteDashboardItem( _root: any, - args: { where: { id: number } }, + args: { where: { id: string | number } }, ctx: IContext, ): Promise { const { id } = args.where; @@ -154,7 +197,9 @@ export class DashboardResolver { public async previewItemSQL( _root: any, - args: { data: { itemId: number; limit?: number; refresh?: boolean } }, + args: { + data: { itemId: string | number; limit?: number; refresh?: boolean }; + }, ctx: IContext, ): Promise { const { itemId, limit, refresh } = args.data; @@ -162,30 +207,43 @@ export class DashboardResolver { const item = await ctx.dashboardService.getDashboardItem(itemId); const { cacheEnabled } = await ctx.dashboardService.getCurrentDashboard(); const project = await ctx.projectService.getCurrentProject(); - const deployment = await ctx.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; - const data = (await ctx.queryService.preview(item.detail.sql, { - project, - manifest: mdl, - limit: limit || DEFAULT_PREVIEW_LIMIT, - cacheEnabled, - refresh: refresh || false, - })) as PreviewDataResponse; - - // handle data to [{ column1: value1, column2: value2, ... }] - const values = data.data.map((val) => { - return data.columns.reduce((acc, col, index) => { - acc[col.name] = val[index]; - return acc; - }, {}); - }); - return { - cacheHit: data.cacheHit || false, - cacheCreatedAt: data.cacheCreatedAt || null, - cacheOverrodeAt: data.cacheOverrodeAt || null, - override: data.override || false, - data: values, - } as PreviewItemResponse; + logger.debug( + `Previewing dashboard item ${String(itemId)} for project ${String(project.id)} (${project.type || 'unknown'})`, + ); + const manifest = await this.getPreviewManifest(ctx, project.id); + + try { + const data = (await this.withTimeout( + ctx.queryService.preview(item.detail.sql, { + project, + manifest, + limit: limit || DEFAULT_PREVIEW_LIMIT, + cacheEnabled, + refresh: refresh || false, + }), + DASHBOARD_PREVIEW_TIMEOUT_MS, + `Dashboard preview timed out for item ${itemId}`, + )) as PreviewDataResponse; + + return this.formatPreviewItemResponse(data); + } catch (error) { + const snapshot = item.detail.previewDataSnapshot; + if (snapshot?.length) { + logger.warn( + `Using stored dashboard preview snapshot for item ${itemId}: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return { + cacheHit: false, + cacheCreatedAt: null, + cacheOverrodeAt: null, + override: false, + data: snapshot, + } as PreviewItemResponse; + } + throw error; + } } catch (error) { logger.error(`Error previewing SQL item ${itemId}: ${error}`); throw error; @@ -212,4 +270,117 @@ export class DashboardResolver { throw error; } } + + private async capturePreviewSnapshot( + ctx: IContext, + sql: string, + ): Promise[] | undefined> { + try { + const project = await ctx.projectService.getCurrentProject(); + const manifest = await this.getPreviewManifest(ctx, project.id); + const data = (await this.withTimeout( + ctx.queryService.preview(sql, { + project, + manifest, + limit: DEFAULT_PREVIEW_LIMIT, + cacheEnabled: false, + refresh: false, + }), + DASHBOARD_SNAPSHOT_TIMEOUT_MS, + 'Dashboard snapshot preview timed out', + )) as PreviewDataResponse; + + return this.formatPreviewItemResponse(data).data; + } catch (error) { + logger.warn( + `Failed to capture dashboard preview snapshot: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + return undefined; + } + } + + private async warmDashboardCache( + ctx: IContext, + sql: string, + dashboardItemId: string | number, + ): Promise { + try { + const project = await ctx.projectService.getCurrentProject(); + const manifest = await this.getPreviewManifest(ctx, project.id); + await this.withTimeout( + ctx.queryService.preview(sql, { + project, + manifest, + limit: DEFAULT_PREVIEW_LIMIT, + cacheEnabled: true, + refresh: true, + }), + DASHBOARD_PREVIEW_TIMEOUT_MS, + `Dashboard cache warm-up timed out for item ${dashboardItemId}`, + ); + } catch (error) { + logger.warn( + `Dashboard item ${dashboardItemId} was pinned but cache warm-up failed: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + } + } + + private async getPreviewManifest( + ctx: IContext, + projectId: string | number, + ): Promise { + const deployment = await ctx.deployService.getLastDeployment(projectId); + if (deployment?.manifest) { + logger.debug(`Using deployed manifest for project ${projectId}`); + return deployment.manifest as Manifest; + } + + logger.debug(`Using current model manifest fallback for project ${projectId}`); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + return manifest; + } + + private formatPreviewItemResponse(data: PreviewDataResponse): PreviewItemResponse { + const values = data.data.map((val) => { + return data.columns.reduce((acc, col, index) => { + acc[col.name] = val[index]; + return acc; + }, {}); + }); + + return { + cacheHit: data.cacheHit || false, + cacheCreatedAt: data.cacheCreatedAt || null, + cacheOverrodeAt: data.cacheOverrodeAt || null, + override: data.override || false, + data: values, + } as PreviewItemResponse; + } + + private async withTimeout( + promise: Promise, + timeoutMs: number, + timeoutMessage: string, + ): Promise { + let timeoutHandle: ReturnType | null = null; + try { + return await Promise.race([ + promise, + new Promise((_resolve, reject) => { + timeoutHandle = setTimeout( + () => reject(new Error(timeoutMessage)), + timeoutMs, + ); + }), + ]); + } finally { + if (timeoutHandle) { + clearTimeout(timeoutHandle); + } + } + } } diff --git a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts index f57b6980cb..e8d5a0be9f 100644 --- a/wren-ui/src/apollo/server/resolvers/diagramResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/diagramResolver.ts @@ -89,7 +89,7 @@ export class DiagramResolver { const allColumns = modelColumns.filter( (column) => column.modelId === model.id, ); - const modelMDL = manifest.models.find( + const modelMDL = manifest.models?.find( (modelMDL) => modelMDL.name === model.referenceName, ); allColumns.forEach((column) => { @@ -113,9 +113,13 @@ export class DiagramResolver { } if (column.isCalculated) { - transformedModel.calculatedFields.push( - this.transformCalculatedField(column, modelMDL.columns), + const transformedCalculatedField = this.transformCalculatedField( + column, + modelMDL?.columns || [], ); + if (transformedCalculatedField) { + transformedModel.calculatedFields.push(transformedCalculatedField); + } } else { const nestedColumns = modelNestedColumns.filter( (nestedColumn) => nestedColumn.columnId === column.id, @@ -128,17 +132,18 @@ export class DiagramResolver { return transformedModel; }); - const diagramViews = views.map(this.transformView); + const diagramViews = views.map((view) => this.transformView(view)); return { models: diagramModels, views: diagramViews }; } private transformModel(model: Model): DiagramModel { - const properties = JSON.parse(model.properties); + const properties = this.parseProperties(model.properties); + const displayName = model.displayName || model.referenceName; return { id: uuidv4(), modelId: model.id, nodeType: NodeType.MODEL, - displayName: model.displayName, + displayName, referenceName: model.referenceName, sourceTableName: model.sourceTableName, refSql: model.refSql, @@ -155,7 +160,8 @@ export class DiagramResolver { column: ModelColumn, nestedColumns: ModelNestedColumn[], ): DiagramModelField { - const properties = JSON.parse(column.properties); + const properties = this.parseProperties(column.properties); + const displayName = column.displayName || column.referenceName; return { id: uuidv4(), columnId: column.id, @@ -163,7 +169,7 @@ export class DiagramResolver { ? NodeType.CALCULATED_FIELD : NodeType.FIELD, type: column.type, - displayName: column.displayName, + displayName, referenceName: column.referenceName, description: properties?.description, isPrimaryKey: column.isPk, @@ -174,7 +180,7 @@ export class DiagramResolver { nestedColumnId: nestedColumn.id, columnPath: nestedColumn.columnPath, type: nestedColumn.type, - displayName: nestedColumn.displayName, + displayName: nestedColumn.displayName || nestedColumn.referenceName, referenceName: nestedColumn.referenceName, description: nestedColumn.properties?.description, })) @@ -185,12 +191,25 @@ export class DiagramResolver { private transformCalculatedField( column: ModelColumn, columnsMDL: ColumnMDL[], - ): DiagramModelField { - const properties = JSON.parse(column.properties); - const lineage = JSON.parse(column.lineage); + ): DiagramModelField | null { + const properties = this.parseProperties(column.properties); + const lineage = this.parseLineage(column.lineage); + const displayName = column.displayName || column.referenceName; const columnMDL = columnsMDL.find( ({ name }) => name === column.referenceName, ); + if (!columnMDL) { + logger.debug( + `Skip diagram calculated field "${column.referenceName}" because it is missing from built MDL`, + ); + return null; + } + if (columnMDL.expression == null) { + logger.debug( + `Skip diagram calculated field "${column.referenceName}" because its MDL expression is missing`, + ); + return null; + } return { id: uuidv4(), columnId: column.id, @@ -198,7 +217,7 @@ export class DiagramResolver { aggregation: column.aggregation, lineage, type: column.type, - displayName: column.displayName, + displayName, referenceName: column.referenceName, description: properties?.description, isPrimaryKey: column.isPk, @@ -222,34 +241,35 @@ export class DiagramResolver { const displayName = models.find( (model) => model.referenceName === referenceName, )?.displayName; - const properties = relation.properties - ? JSON.parse(relation.properties) - : null; + const properties = this.parseProperties(relation.properties); return { id: uuidv4(), relationId: relation.id, nodeType: NodeType.RELATION, - displayName, + displayName: displayName || referenceName, referenceName, type: relation.joinType as RelationType, fromModelId: relation.fromModelId, fromModelName: relation.fromModelName, - fromModelDisplayName: relation.fromModelDisplayName, + fromModelDisplayName: + relation.fromModelDisplayName || relation.fromModelName, fromColumnId: relation.fromColumnId, fromColumnName: relation.fromColumnName, - fromColumnDisplayName: relation.fromColumnDisplayName, + fromColumnDisplayName: + relation.fromColumnDisplayName || relation.fromColumnName, toModelId: relation.toModelId, toModelName: relation.toModelName, - toModelDisplayName: relation.toModelDisplayName, + toModelDisplayName: relation.toModelDisplayName || relation.toModelName, toColumnId: relation.toColumnId, toColumnName: relation.toColumnName, - toColumnDisplayName: relation.toColumnDisplayName, + toColumnDisplayName: + relation.toColumnDisplayName || relation.toColumnName, description: properties?.description, }; } private transformView(view: View): DiagramView { - const properties = JSON.parse(view.properties); + const properties = this.parseProperties(view.properties); const fields = (properties?.columns || []).map((column: any) => ({ id: uuidv4(), nodeType: NodeType.FIELD, @@ -270,4 +290,35 @@ export class DiagramResolver { description: properties?.description, }; } + + private parseProperties( + properties?: string | Record | null, + ): Record { + if (!properties) { + return {}; + } + if (typeof properties === 'object') { + return properties; + } + try { + const parsed = JSON.parse(properties); + return parsed && typeof parsed === 'object' ? parsed : {}; + } catch (error) { + logger.debug(`Can not parse properties "${properties}"`); + return {}; + } + } + + private parseLineage(lineage?: string | null): number[] { + if (!lineage) { + return []; + } + try { + const parsed = JSON.parse(lineage); + return Array.isArray(parsed) ? parsed : []; + } catch (error) { + logger.debug(`Can not parse lineage "${lineage}"`); + return []; + } + } } diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts index 6facf755ff..dcaf404ddd 100644 --- a/wren-ui/src/apollo/server/resolvers/modelResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/modelResolver.ts @@ -2,23 +2,26 @@ import { CreateModelData, UpdateModelData, UpdateModelMetadataInput, + SaveModelingSemanticInput, CreateCalculatedFieldData, UpdateCalculatedFieldData, UpdateViewMetadataInput, + DryPlanSQLData, PreviewSQLData, } from '../models'; import { DataSourceName, IContext, + ModelingRelationshipData, RelationData, UpdateRelationData, } from '../types'; -import { getLogger, transformInvalidColumnName } from '@server/utils'; +import { getLogger, transformUniqueInvalidColumnName } from '@server/utils'; import { DeployResponse } from '../services/deployService'; import { safeFormatSQL } from '@server/utils/sqlFormat'; import { isEmpty, isNil } from 'lodash'; import { replaceAllowableSyntax, validateDisplayName } from '../utils/regex'; -import { Model, ModelColumn } from '../repositories'; +import { Model, ModelColumn, Project } from '../repositories'; import { findColumnsToUpdate, getPreviewColumnsStr, @@ -28,10 +31,38 @@ import { } from '../utils/model'; import { CompactTable, PreviewDataResponse } from '@server/services'; import { TelemetryEvent } from '../telemetry/telemetry'; +import DataSourceSchemaDetector, { + SchemaChangeType, +} from '@server/managers/dataSourceSchemaDetector'; const logger = getLogger('ModelResolver'); logger.level = 'debug'; +const dirtyProjectIds = new Set(); + +const isSameId = (left: string | number, right: string | number) => + String(left) === String(right); + +const firstNonEmptyString = (...values: unknown[]) => + values.find( + (value): value is string => + typeof value === 'string' && value.trim().length > 0, + ) || ''; + +const normalizeModelDisplayName = (model: Model) => + firstNonEmptyString( + model.displayName, + model.referenceName, + model.sourceTableName, + ); + +const normalizeColumnDisplayName = (column: ModelColumn) => + firstNonEmptyString( + column.displayName, + column.referenceName, + column.sourceColumnName, + ); + export enum SyncStatusEnum { IN_PROGRESS = 'IN_PROGRESS', SYNCRONIZED = 'SYNCRONIZED', @@ -49,6 +80,15 @@ export class ModelResolver { this.updateModelMetadata = this.updateModelMetadata.bind(this); this.deploy = this.deploy.bind(this); this.getMDL = this.getMDL.bind(this); + this.generateModelingSemantics = this.generateModelingSemantics.bind(this); + this.getModelingSemanticsResult = + this.getModelingSemanticsResult.bind(this); + this.saveModelingSemantics = this.saveModelingSemantics.bind(this); + this.generateModelingRelationships = + this.generateModelingRelationships.bind(this); + this.getModelingRelationshipsResult = + this.getModelingRelationshipsResult.bind(this); + this.saveModelingRelationships = this.saveModelingRelationships.bind(this); this.checkModelSync = this.checkModelSync.bind(this); // view @@ -63,6 +103,7 @@ export class ModelResolver { this.previewModelData = this.previewModelData.bind(this); this.previewViewData = this.previewViewData.bind(this); this.previewSql = this.previewSql.bind(this); + this.dryPlanSql = this.dryPlanSql.bind(this); this.getNativeSql = this.getNativeSql.bind(this); // calculated field @@ -75,6 +116,11 @@ export class ModelResolver { this.createRelation = this.createRelation.bind(this); this.updateRelation = this.updateRelation.bind(this); this.deleteRelation = this.deleteRelation.bind(this); + this.prepareProjectForDeploy = this.prepareProjectForDeploy.bind(this); + this.resolveModifiedSchemaChanges = + this.resolveModifiedSchemaChanges.bind(this); + this.refreshProjectDataSourceVersion = + this.refreshProjectDataSourceVersion.bind(this); } public async createRelation( @@ -87,6 +133,7 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_RELATION; try { const relation = await ctx.modelService.createRelation(data); + this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -109,6 +156,7 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_UPDATE_RELATION; try { const relation = await ctx.modelService.updateRelation(data, where.id); + this.markProjectDirty(relation.projectId); ctx.telemetry.sendEvent(eventName, { data }); return relation; } catch (err: any) { @@ -127,8 +175,10 @@ export class ModelResolver { args: { where: { id: number } }, ctx: IContext, ) { + const project = await ctx.projectService.getCurrentProject(); const relationId = args.where.id; await ctx.modelService.deleteRelation(relationId); + this.markProjectDirty(project.id); return true; } @@ -140,6 +190,8 @@ export class ModelResolver { const eventName = TelemetryEvent.MODELING_CREATE_CF; try { const column = await ctx.modelService.createCalculatedField(_args.data); + const project = await ctx.projectService.getCurrentProject(); + this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data: _args.data }); return column; } catch (err: any) { @@ -175,6 +227,8 @@ export class ModelResolver { data, where.id, ); + const project = await ctx.projectService.getCurrentProject(); + this.markProjectDirty(project.id); ctx.telemetry.sendEvent(eventName, { data }); return column; } catch (err: any) { @@ -195,24 +249,34 @@ export class ModelResolver { if (!column || !column.isCalculated) { throw new Error('Calculated field not found'); } + const project = await ctx.projectService.getCurrentProject(); await ctx.modelColumnRepository.deleteOne(columnId); + this.markProjectDirty(project.id); return true; } public async checkModelSync(_root: any, _args: any, ctx: IContext) { - const { id } = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); - const currentHash = ctx.deployService.createMDLHash(manifest, id); - const lastDeploy = await ctx.deployService.getLastDeployment(id); - const lastDeployHash = lastDeploy?.hash; - const inProgressDeployment = - await ctx.deployService.getInProgressDeployment(id); - if (inProgressDeployment) { - return { status: SyncStatusEnum.IN_PROGRESS }; - } - return currentHash == lastDeployHash - ? { status: SyncStatusEnum.SYNCRONIZED } - : { status: SyncStatusEnum.UNSYNCRONIZED }; + try { + const { id } = await ctx.projectService.getCurrentProject(); + const inProgressDeployment = + await ctx.deployService.getInProgressDeployment(id); + if (inProgressDeployment) { + return { status: SyncStatusEnum.IN_PROGRESS }; + } + + const project = await ctx.projectService.getCurrentProject(); + if (dirtyProjectIds.has(project.id)) { + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } + + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + return lastDeploy + ? { status: SyncStatusEnum.SYNCRONIZED } + : { status: SyncStatusEnum.UNSYNCRONIZED }; + } catch (err: any) { + logger.error(`checkModelSync failed: ${err.message}`, err); + return { status: SyncStatusEnum.UNSYNCRONIZED }; + } } public async deploy( @@ -220,28 +284,194 @@ export class ModelResolver { args: { force: boolean }, ctx: IContext, ): Promise { - const project = await ctx.projectService.getCurrentProject(); - if (!project.version && project.type !== DataSourceName.DUCKDB) { - const version = - await ctx.projectService.getProjectDataSourceVersion(project); - await ctx.projectService.updateProject(project.id, { - version, - }); - } + const project = await this.prepareProjectForDeploy(ctx); const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + const lastDeploy = await ctx.deployService.getLastDeployment(project.id); + const hasModelingChangesAfterDeploy = + !(await this.isLastDeployNewerThanModelingChanges( + ctx, + project.id, + lastDeploy, + )); + const shouldForceDeploy = + args.force || + hasModelingChangesAfterDeploy || + !ctx.deployService.isSameDeployment(manifest, project.id, lastDeploy); const deployRes = await ctx.deployService.deploy( manifest, project.id, - args.force, + shouldForceDeploy, ); - - // only generating for user's data source - if (project.sampleDataset === null) { - await ctx.projectService.generateProjectRecommendationQuestions(); + if (deployRes.status === 'SUCCESS') { + dirtyProjectIds.delete(project.id); + } + if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { + ctx.projectService.generateProjectRecommendationQuestions().catch((err) => + logger.warn( + `Failed to generate project recommendation questions after deploy: ${err.message}`, + ), + ); } return deployRes; } + private async prepareProjectForDeploy(ctx: IContext): Promise { + let project = await ctx.projectService.getCurrentProject(); + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + + await schemaDetector.detectSchemaChange(); + await this.resolveModifiedSchemaChanges(ctx, project.id); + project = await this.refreshProjectDataSourceVersion(ctx, project); + return project; + } + + private async isLastDeployNewerThanModelingChanges( + ctx: IContext, + projectId: number, + lastDeploy?: { createdAt?: Date; updatedAt?: Date } | null, + ): Promise { + if (!lastDeploy) { + return false; + } + + const deployedAt = this.toTime(lastDeploy.updatedAt || lastDeploy.createdAt); + if (!deployedAt) { + return false; + } + + const models = await ctx.modelRepository.findAllBy({ projectId }); + const modelIds = models.map((model) => model.id); + const [columns, nestedColumns, relations, views] = await Promise.all([ + modelIds.length + ? ctx.modelColumnRepository.findColumnsByModelIds(modelIds) + : Promise.resolve([]), + modelIds.length + ? ctx.modelNestedColumnRepository.findNestedColumnsByModelIds(modelIds) + : Promise.resolve([]), + ctx.relationRepository.findRelationInfoBy({ projectId }), + ctx.viewRepository.findAllBy({ projectId }), + ]); + + const latestModelingChangeAt = [ + ...models, + ...columns, + ...nestedColumns, + ...relations, + ...views, + ].reduce((latest, item: any) => { + return Math.max(latest, this.toTime(item.updatedAt || item.createdAt)); + }, 0); + + return deployedAt >= latestModelingChangeAt; + } + + private toTime(value?: Date | string | null): number { + if (!value) { + return 0; + } + const time = new Date(value).getTime(); + return Number.isFinite(time) ? time : 0; + } + + private isSameDeploymentIgnoringColumnNullability( + manifest: any, + lastDeploy?: { manifest?: any } | null, + ): boolean { + if (!lastDeploy?.manifest) { + return false; + } + + return ( + this.stableStringify(this.omitColumnNullability(lastDeploy.manifest)) === + this.stableStringify(this.omitColumnNullability(manifest)) + ); + } + + private omitColumnNullability(value: any): any { + if (Array.isArray(value)) { + return value.map((item) => this.omitColumnNullability(item)); + } + if (!value || typeof value !== 'object') { + return value; + } + + const result: Record = {}; + for (const key of Object.keys(value)) { + if (key === 'notNull') { + continue; + } + result[key] = this.omitColumnNullability(value[key]); + } + return result; + } + + private stableStringify(value: any): string { + if (Array.isArray(value)) { + const serializedItems = value.map((item) => this.stableStringify(item)); + if (value.every((item) => item && typeof item === 'object')) { + serializedItems.sort(); + } + return `[${serializedItems.join(',')}]`; + } + + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${this.stableStringify(value[key])}`) + .join(',')}}`; + } + + return JSON.stringify(value); + } + + private markProjectDirty(projectId: number) { + dirtyProjectIds.add(projectId); + } + + private async resolveModifiedSchemaChanges(ctx: IContext, projectId: number) { + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(projectId); + const hasUnresolvedModifiedColumns = + lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && + !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; + + if (!hasUnresolvedModifiedColumns) { + return; + } + + const schemaDetector = new DataSourceSchemaDetector({ ctx, projectId }); + await schemaDetector.resolveSchemaChange(SchemaChangeType.MODIFIED_COLUMNS); + } + + private async refreshProjectDataSourceVersion( + ctx: IContext, + project: Project, + ): Promise { + if (project.type === DataSourceName.DUCKDB) { + return project; + } + + try { + const version = await ctx.projectService.getProjectDataSourceVersion( + project, + ); + if (version && version !== project.version) { + return await ctx.projectService.updateProject(project.id, { + version, + updatedAt: project.updatedAt, + }); + } + } catch (err: any) { + logger.warn( + `Failed to refresh project datasource version before deploy: ${err.message}`, + ); + } + return project; + } + public async getMDL(_root: any, args: { hash: string }, ctx: IContext) { const mdl = await ctx.deployService.getMDLByHash(args.hash); return { @@ -250,6 +480,258 @@ export class ModelResolver { }; } + public async generateModelingSemantics( + _root: any, + args: { data: { selectedModels: string[]; userPrompt: string } }, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + return await ctx.wrenAIAdaptor.generateSemanticsDescription({ + manifest, + selectedModels: args.data.selectedModels, + userPrompt: args.data.userPrompt, + projectId: project.id, + }); + } + + public async getModelingSemanticsResult( + _root: any, + args: { queryId: string }, + ctx: IContext, + ) { + return await ctx.wrenAIAdaptor.getSemanticsDescriptionResult(args.queryId); + } + + public async saveModelingSemantics( + _root: any, + args: { data: SaveModelingSemanticInput[] }, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const models = await ctx.modelRepository.findAllBy({ + projectId: project.id, + }); + const modelById = new Map(models.map((model) => [model.id, model])); + const modelByReferenceName = new Map( + models.map((model) => [model.referenceName, model]), + ); + const requestedItems = args.data || []; + + const resolveModel = (item: SaveModelingSemanticInput) => { + const model = modelById.get(item.modelId); + if (model) return model; + return item.referenceName + ? modelByReferenceName.get(item.referenceName) + : undefined; + }; + const hasMetadataText = (value?: string) => + typeof value === 'string' && value.trim().length > 0; + + for (const item of requestedItems) { + if (!resolveModel(item)) { + throw new Error( + `Model not found: ${item.referenceName || item.modelId}`, + ); + } + } + + await Promise.all( + requestedItems.map(async (item) => { + if ( + !hasMetadataText(item.description) && + !hasMetadataText(item.displayName) + ) { + return; + } + + const model = resolveModel(item); + const modelMetadata: Partial = {}; + + if (hasMetadataText(item.displayName)) { + modelMetadata.displayName = this.determineMetadataValue( + item.displayName, + ); + } + + if (hasMetadataText(item.description)) { + const properties = model?.properties + ? JSON.parse(model.properties) + : {}; + properties.description = this.determineMetadataValue( + item.description, + ); + modelMetadata.properties = JSON.stringify(properties); + } + + if (!isEmpty(modelMetadata)) { + await ctx.modelRepository.updateOne(model.id, modelMetadata); + } + }), + ); + + const requestedColumns = requestedItems.flatMap((item) => { + const model = resolveModel(item); + return (item.columns || []).map((column) => ({ + ...column, + modelId: model?.id, + })); + }); + if (!isEmpty(requestedColumns)) { + const columnIds = requestedColumns.map((column) => column.id); + const modelIds = models.map((model) => model.id); + const [columnsByIdSource, columnsByModelSource] = await Promise.all([ + ctx.modelColumnRepository.findColumnsByIds(columnIds), + ctx.modelColumnRepository.findColumnsByModelIds(modelIds), + ]); + const columns = [...columnsByIdSource, ...columnsByModelSource]; + const columnById = new Map( + columns.map((column) => [String(column.id), column]), + ); + const columnByModelAndReferenceName = new Map( + columns.map((column) => [ + `${column.modelId}:${column.referenceName}`, + column, + ]), + ); + + await Promise.all( + requestedColumns.map(async (requestedColumn) => { + const columnByRequestedId = columnById.get( + String(requestedColumn.id), + ); + const column = + columnByRequestedId?.modelId === requestedColumn.modelId + ? columnByRequestedId + : columnByModelAndReferenceName.get( + `${requestedColumn.modelId}:${requestedColumn.referenceName}`, + ); + if (!column) return; + + const columnMetadata: Partial = {}; + if (hasMetadataText(requestedColumn.displayName)) { + columnMetadata.displayName = this.determineMetadataValue( + requestedColumn.displayName, + ); + } + + if (hasMetadataText(requestedColumn.description)) { + const properties = column.properties + ? JSON.parse(column.properties) + : {}; + properties.description = this.determineMetadataValue( + requestedColumn.description, + ); + columnMetadata.properties = JSON.stringify(properties); + } + + if (!isEmpty(columnMetadata)) { + await ctx.modelColumnRepository.updateOne( + column.id, + columnMetadata, + ); + } + }), + ); + } + + if (!isEmpty(requestedItems)) { + this.markProjectDirty(project.id); + } + + return { + savedCount: requestedItems.length, + columnSavedCount: requestedColumns.length, + }; + } + + public async generateModelingRelationships( + _root: any, + _args: any, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); + return await ctx.wrenAIAdaptor.generateRelationshipRecommendations({ + manifest, + projectId: project.id, + }); + } + + public async getModelingRelationshipsResult( + _root: any, + args: { queryId: string }, + ctx: IContext, + ) { + return await ctx.wrenAIAdaptor.getRelationshipRecommendationResult( + args.queryId, + ); + } + + public async saveModelingRelationships( + _root: any, + args: { data: ModelingRelationshipData[] }, + ctx: IContext, + ) { + const project = await ctx.projectService.getCurrentProject(); + const models = await ctx.modelRepository.findAllBy({ + projectId: project.id, + }); + const modelIds = models.map((model) => model.id); + const columns = + await ctx.modelColumnRepository.findColumnsByModelIds(modelIds); + let createdCount = 0; + let skippedCount = 0; + + for (const relationship of args.data || []) { + const fromModel = models.find( + (model) => model.referenceName === relationship.fromModel, + ); + const toModel = models.find( + (model) => model.referenceName === relationship.toModel, + ); + const fromColumn = fromModel + ? columns.find( + (column) => + column.modelId === fromModel.id && + column.referenceName === relationship.fromColumn, + ) + : null; + const toColumn = toModel + ? columns.find( + (column) => + column.modelId === toModel.id && + column.referenceName === relationship.toColumn, + ) + : null; + + if (!fromModel || !toModel || !fromColumn || !toColumn) { + skippedCount += 1; + continue; + } + + try { + const savedRelation = await ctx.modelService.createRelation({ + fromModelId: fromModel.id, + fromColumnId: fromColumn.id, + toModelId: toModel.id, + toColumnId: toColumn.id, + type: relationship.type, + description: relationship.description, + }); + this.markProjectDirty(savedRelation.projectId); + createdCount += 1; + } catch (err: any) { + logger.warn( + `Skip Modeling AI Assistant relationship ${relationship.fromModel}.${relationship.fromColumn} -> ${relationship.toModel}.${relationship.toColumn}: ${err.message}`, + ); + skippedCount += 1; + } + } + + return { createdCount, skippedCount }; + } + public async listModels(_root: any, _args: any, ctx: IContext) { const { id: projectId } = await ctx.projectService.getCurrentProject(); const models = await ctx.modelRepository.findAllBy({ projectId }); @@ -266,15 +748,27 @@ export class ModelResolver { .filter((c) => c.modelId === model.id) .map((c) => ({ ...c, + displayName: normalizeColumnDisplayName(c), properties: JSON.parse(c.properties), nestedColumns: c.type.includes('STRUCT') - ? modelNestedColumnList.filter((nc) => nc.columnId === c.id) + ? modelNestedColumnList + .filter((nc) => nc.columnId === c.id) + .map((nc) => ({ + ...nc, + displayName: firstNonEmptyString( + nc.displayName, + nc.referenceName, + nc.sourceColumnName, + nc.columnPath?.join('.'), + ), + })) : undefined, })); const fields = modelFields.filter((c) => !c.isCalculated); const calculatedFields = modelFields.filter((c) => c.isCalculated); result.push({ ...model, + displayName: normalizeModelDisplayName(model), fields, calculatedFields, properties: { @@ -301,9 +795,20 @@ export class ModelResolver { const columns = modelColumns.map((c) => ({ ...c, + displayName: normalizeColumnDisplayName(c), properties: JSON.parse(c.properties), nestedColumns: c.type.includes('STRUCT') - ? modelNestedColumns.filter((nc) => nc.columnId === c.id) + ? modelNestedColumns + .filter((nc) => nc.columnId === c.id) + .map((nc) => ({ + ...nc, + displayName: firstNonEmptyString( + nc.displayName, + nc.referenceName, + nc.sourceColumnName, + nc.columnPath?.join('.'), + ), + })) : undefined, })); const relations = ( @@ -318,6 +823,7 @@ export class ModelResolver { return { ...model, + displayName: normalizeModelDisplayName(model), fields: columns.filter((c) => !c.isCalculated), calculatedFields: columns.filter((c) => c.isCalculated), relations, @@ -343,6 +849,7 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_CREATE_MODEL, { data: args.data, }); + this.markProjectDirty(model.projectId); return model; } catch (error: any) { ctx.telemetry.sendEvent( @@ -390,13 +897,17 @@ export class ModelResolver { const compactColumns = dataSourceTable.columns.filter((c) => fields.includes(c.name), ); + const usedReferenceNames = new Set(); const columnValues = compactColumns.map( (column) => ({ modelId: model.id, isCalculated: false, displayName: column.name, - referenceName: transformInvalidColumnName(column.name), + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), sourceColumnName: column.name, type: column.type || 'string', notNull: column.notNull || false, @@ -437,6 +948,7 @@ export class ModelResolver { ctx.telemetry.sendEvent(TelemetryEvent.MODELING_UPDATE_MODEL, { data: args.data, }); + this.markProjectDirty(model.projectId); return model; } catch (err: any) { ctx.telemetry.sendEvent( @@ -485,6 +997,11 @@ export class ModelResolver { // create columns if (toCreateColumns.length) { + const usedReferenceNames = new Set( + existingColumns + .filter(({ id }) => !toDeleteColumnIds.includes(id)) + .map(({ referenceName }) => referenceName.toLowerCase()), + ); const compactColumns = sourceTableColumns.filter((sourceColumn) => toCreateColumns.includes(sourceColumn.name), ); @@ -494,7 +1011,10 @@ export class ModelResolver { isCalculated: false, displayName: column.name, sourceColumnName: column.name, - referenceName: transformInvalidColumnName(column.name), + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), type: column.type || 'string', notNull: column.notNull, isPk: primaryKey === column.name, @@ -558,6 +1078,7 @@ export class ModelResolver { // related columns and relationships will be deleted in cascade await ctx.modelRepository.deleteOne(modelId); + this.markProjectDirty(model.projectId); return true; } @@ -603,6 +1124,7 @@ export class ModelResolver { } ctx.telemetry.sendEvent(eventName, { data }); + this.markProjectDirty(model.projectId); return true; } catch (err: any) { ctx.telemetry.sendEvent( @@ -651,7 +1173,10 @@ export class ModelResolver { const relationships = await ctx.relationRepository.findRelationsByIds(relationshipIds); for (const rel of relationships) { - const requestedMetadata = data.relationships.find((r) => r.id === rel.id); + const requestedMetadata = data.relationships.find((r) => + isSameId(r.id, rel.id), + ); + if (!requestedMetadata) continue; const relationMetadata: any = {}; @@ -678,8 +1203,9 @@ export class ModelResolver { await ctx.modelColumnRepository.findColumnsByIds(calculatedFieldIds); for (const col of modelColumns) { const requestedMetadata = data.calculatedFields.find( - (c) => c.id === col.id, + (c) => isSameId(c.id, col.id), ); + if (!requestedMetadata) continue; const columnMetadata: any = {}; // check if description is empty @@ -707,7 +1233,10 @@ export class ModelResolver { const modelColumns = await ctx.modelColumnRepository.findColumnsByIds(columnIds); for (const col of modelColumns) { - const requestedMetadata = data.columns.find((c) => c.id === col.id); + const requestedMetadata = data.columns.find((c) => + isSameId(c.id, col.id), + ); + if (!requestedMetadata) continue; // update metadata const columnMetadata: any = {}; @@ -742,7 +1271,10 @@ export class ModelResolver { nestedColumnIds, ); for (const col of modelNestedColumns) { - const requestedMetadata = data.nestedColumns.find((c) => c.id === col.id); + const requestedMetadata = data.nestedColumns.find((c) => + isSameId(c.id, col.id), + ); + if (!requestedMetadata) continue; const nestedColumnMetadata: any = {}; @@ -812,7 +1344,7 @@ export class ModelResolver { // create view const project = await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.deployService.getLastDeployment(project.id); + const manifest = await this.getLastDeployedManifest(ctx, project.id); // get sql statement of a response const response = await ctx.askingService.getResponse(responseId); @@ -862,6 +1394,7 @@ export class ModelResolver { // telemetry ctx.telemetry.sendEvent(eventName, eventProperties); + this.markProjectDirty(project.id); return { ...view, displayName }; } catch (err: any) { @@ -887,6 +1420,7 @@ export class ModelResolver { throw new Error('View not found'); } await ctx.viewRepository.deleteOne(viewId); + this.markProjectDirty(view.projectId); return true; } @@ -937,18 +1471,71 @@ export class ModelResolver { args: { data: PreviewSQLData }, ctx: IContext, ) { - const { sql, projectId, limit, dryRun } = args.data; + const startedAt = Date.now(); + const { sql, projectId, hash, limit, dryRun } = args.data; const project = projectId ? await ctx.projectService.getProjectById(parseInt(projectId)) : await ctx.projectService.getCurrentProject(); - const { manifest } = await ctx.deployService.getLastDeployment(project.id); - return await ctx.queryService.preview(sql, { + const manifest = hash + ? await ctx.deployService.getManifestByHash(hash) + : await this.getLastDeployedManifest(ctx, project.id); + if (!manifest) { + throw new Error( + 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', + ); + } + const result = await ctx.queryService.preview(sql, { project, limit: limit, modelingOnly: false, manifest, dryRun, }); + logger.info( + `Ask timing stage=preview_sql_request project_id=${project.id} dry_run=${ + dryRun ?? false + } elapsed_ms=${Date.now() - startedAt}`, + ); + return result; + } + + public async dryPlanSql( + _root: any, + args: { data: DryPlanSQLData }, + ctx: IContext, + ): Promise { + const { sql, projectId, hash, allowFallback } = args.data; + const project = projectId + ? await ctx.projectService.getProjectById(parseInt(projectId)) + : await ctx.projectService.getCurrentProject(); + const manifest = hash + ? await ctx.deployService.getManifestByHash(hash) + : await this.getLastDeployedManifest(ctx, project.id); + if (!manifest) { + throw new Error( + 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', + ); + } + + try { + if (project.type === DataSourceName.DUCKDB) { + await ctx.wrenEngineAdaptor.getNativeSQL(sql, { + manifest, + modelingOnly: false, + }); + } else { + await ctx.ibisServerAdaptor.getNativeSql({ + dataSource: project.type, + sql, + mdl: manifest, + allowFallback, + }); + } + } catch (error) { + throw error; + } + + return true; } public async getNativeSql( @@ -1043,6 +1630,7 @@ export class ModelResolver { name: newName, properties: JSON.stringify(properties), }); + this.markProjectDirty(view.projectId); return true; } @@ -1089,6 +1677,17 @@ export class ModelResolver { }; } + private async getLastDeployedManifest(ctx: IContext, projectId: number) { + const deployment = await ctx.deployService.getLastDeployment(projectId); + if (!deployment?.manifest) { + throw new Error( + 'Project has not been deployed successfully yet. Deploy the model before previewing or validating SQL.', + ); + } + + return deployment.manifest; + } + private validateTableExist( tableName: string, dataSourceTables: CompactTable[], diff --git a/wren-ui/src/apollo/server/resolvers/projectResolver.ts b/wren-ui/src/apollo/server/resolvers/projectResolver.ts index 0d6d6ba0e0..e82741d5e5 100644 --- a/wren-ui/src/apollo/server/resolvers/projectResolver.ts +++ b/wren-ui/src/apollo/server/resolvers/projectResolver.ts @@ -12,7 +12,7 @@ import { trim, getLogger, replaceInvalidReferenceName, - transformInvalidColumnName, + transformUniqueInvalidColumnName, handleNestedColumns, } from '@server/utils'; import { @@ -20,6 +20,7 @@ import { Model, ModelColumn, Project, + WorkspaceProjectType, } from '../repositories'; import { SampleDatasetName, @@ -39,6 +40,7 @@ import { TelemetryEvent } from '../telemetry/telemetry'; const logger = getLogger('DataSourceResolver'); logger.level = 'debug'; +const DEFAULT_PROJECT_LANGUAGE = 'EN'; export enum OnboardingStatusEnum { NOT_STARTED = 'NOT_STARTED', @@ -83,7 +85,7 @@ export class ProjectResolver { } as DataSourceProperties, sampleDataset: project.sampleDataset, }, - language: project.language, + language: project.language || DEFAULT_PROJECT_LANGUAGE, }; } @@ -97,14 +99,20 @@ export class ProjectResolver { public async updateCurrentProject( _root: any, - arg: { data: { language: string } }, + arg: { data: { language: string; displayName?: string } }, ctx: IContext, ) { - const { language } = arg.data; + const { language, displayName } = arg.data; const project = await ctx.projectService.getCurrentProject(); - await ctx.projectRepository.updateOne(project.id, { + const changes: Record = { language, - }); + }; + + if (typeof displayName === 'string' && trim(displayName)) { + changes.displayName = trim(displayName); + } + + await ctx.projectRepository.updateOne(project.id, changes); // only generating for user's data source if (project.sampleDataset === null) { @@ -258,19 +266,17 @@ export class ProjectResolver { public async saveDataSource( _root: any, args: { - data: DataSource; + data: DataSource & { projectType?: WorkspaceProjectType }; }, ctx: IContext, ) { const { type, properties } = args.data; - // Currently only can create one project - await this.resetCurrentProject(_root, args, ctx); - const { displayName, ...connectionInfo } = properties; const project = await ctx.projectService.createProject({ displayName, type, connectionInfo, + projectType: args.data.projectType || WorkspaceProjectType.CLASSIC, } as ProjectData); logger.debug(`Project created.`); @@ -372,6 +378,14 @@ export class ProjectResolver { displayName, connectionInfo: { ...project.connectionInfo, ...toUpdateConnectionInfo }, }); + + await ctx.schemaChangeRepository.deleteAllBy({ projectId: project.id }); + await ctx.deployService.deleteAllByProjectId(project.id); + await ctx.askingService.deleteAllByProjectId(project.id); + await ctx.modelService.deleteAllViewsByProjectId(project.id); + await ctx.modelService.deleteAllModelsByProjectId(project.id); + await ctx.wrenAIAdaptor.delete(project.id); + return { type: updatedProject.type, properties: { @@ -633,13 +647,50 @@ export class ProjectResolver { } private async deploy(ctx: IContext) { - const project = await ctx.projectService.getCurrentProject(); + let project = await ctx.projectService.getCurrentProject(); + const schemaDetector = new DataSourceSchemaDetector({ + ctx, + projectId: project.id, + }); + await schemaDetector.detectSchemaChange(); + + const lastSchemaChange = + await ctx.schemaChangeRepository.findLastSchemaChange(project.id); + const hasUnresolvedModifiedColumns = + lastSchemaChange?.resolve?.[SchemaChangeType.MODIFIED_COLUMNS] === false && + !!lastSchemaChange?.change?.[SchemaChangeType.MODIFIED_COLUMNS]?.length; + if (hasUnresolvedModifiedColumns) { + await schemaDetector.resolveSchemaChange( + SchemaChangeType.MODIFIED_COLUMNS, + ); + } + + if (project.type !== DataSourceName.DUCKDB) { + try { + const version = + await ctx.projectService.getProjectDataSourceVersion(project); + if (version && version !== project.version) { + project = await ctx.projectService.updateProject(project.id, { + version, + updatedAt: project.updatedAt, + }); + } + } catch (err: any) { + logger.warn( + `Failed to refresh project datasource version before deploy: ${err.message}`, + ); + } + } + const { manifest } = await ctx.mdlService.makeCurrentModelMDL(); const deployRes = await ctx.deployService.deploy(manifest, project.id); - // only generating for user's data source - if (project.sampleDataset === null) { - await ctx.projectService.generateProjectRecommendationQuestions(); + if (deployRes.status === 'SUCCESS' && project.sampleDataset === null) { + ctx.projectService.generateProjectRecommendationQuestions().catch((err) => + logger.warn( + `Failed to generate project recommendation questions after deploy: ${err.message}`, + ), + ); } return deployRes; } @@ -727,13 +778,17 @@ export class ProjectResolver { const compactColumns = table.columns; const primaryKey = table.primaryKey; const model = models.find((m) => m.sourceTableName === table.name); + const usedReferenceNames = new Set(); return compactColumns.map( (column) => ({ modelId: model.id, isCalculated: false, displayName: column.name, - referenceName: transformInvalidColumnName(column.name), + referenceName: transformUniqueInvalidColumnName( + column.name, + usedReferenceNames, + ), sourceColumnName: column.name, type: column.type || 'string', notNull: column.notNull || false, @@ -747,15 +802,19 @@ export class ProjectResolver { const columns = await ctx.modelColumnRepository.createMany(columnValues); // create nested columns - const compactColumns = selectedTables.flatMap((table) => table.columns); - const nestedColumnValues = compactColumns.flatMap((compactColumn) => { - const column = columns.find( - (c) => c.sourceColumnName === compactColumn.name, - ); - return handleNestedColumns(compactColumn, { - modelId: column.modelId, - columnId: column.id, - sourceColumnName: column.sourceColumnName, + const nestedColumnValues = selectedTables.flatMap((table) => { + const model = models.find((m) => m.sourceTableName === table.name); + const tableColumns = columns.filter((c) => c.modelId === model.id); + return table.columns.flatMap((compactColumn) => { + const column = tableColumns.find( + (c) => c.sourceColumnName === compactColumn.name, + ); + if (!column) return []; + return handleNestedColumns(compactColumn, { + modelId: column.modelId, + columnId: column.id, + sourceColumnName: column.sourceColumnName, + }); }); }); await ctx.modelNestedColumnRepository.createMany(nestedColumnValues); diff --git a/wren-ui/src/apollo/server/resolvers/rbacResolver.ts b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts new file mode 100644 index 0000000000..26f4f761ab --- /dev/null +++ b/wren-ui/src/apollo/server/resolvers/rbacResolver.ts @@ -0,0 +1,149 @@ +import { IContext } from '@server/types'; +import { + CreateRoleInput, + CreateUserInput, + UpdateRoleInput, + UpdateUserInput, + UpdateUserRolesInput, + UserRoleInput, +} from '@server/models'; +import { + RbacUser, + Role, + UserRole, + UserRoleMapping, +} from '@server/repositories'; +import { RbacUserWithRoles, RoleWithUsers } from '@server/services'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('RbacResolver'); +logger.level = 'debug'; + +export class RbacResolver { + constructor() { + this.listRoles = this.listRoles.bind(this); + this.listUsers = this.listUsers.bind(this); + this.listUserRoleMappings = this.listUserRoleMappings.bind(this); + this.createRole = this.createRole.bind(this); + this.updateRole = this.updateRole.bind(this); + this.createUser = this.createUser.bind(this); + this.updateUser = this.updateUser.bind(this); + this.assignRoleToUser = this.assignRoleToUser.bind(this); + this.updateUserRoles = this.updateUserRoles.bind(this); + this.removeRoleFromUser = this.removeRoleFromUser.bind(this); + } + + public getRoleNestedResolver() { + return { + users: async (role: RoleWithUsers, _args: any, ctx: IContext) => { + if (role.users) return role.users; + const mappings = await ctx.rbacService.getUserRoleMappings(); + return mappings + .filter((mapping) => mapping.roleId === role.id) + .map((mapping) => mapping.user); + }, + }; + } + + public getUserNestedResolver() { + return { + roles: async (user: RbacUserWithRoles, _args: any, ctx: IContext) => { + if (user.roles) return user.roles; + const mappings = await ctx.rbacService.getUserRoleMappings(); + return mappings + .filter((mapping) => mapping.userId === user.id) + .map((mapping) => mapping.role); + }, + }; + } + + public async listRoles( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + try { + return await ctx.rbacService.listRoles(); + } catch (error) { + logger.error(`Error listing roles: ${error}`); + throw error; + } + } + + public async listUsers( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + try { + return await ctx.rbacService.listUsers(); + } catch (error) { + logger.error(`Error listing users: ${error}`); + throw error; + } + } + + public async listUserRoleMappings( + _root: any, + _args: any, + ctx: IContext, + ): Promise { + return ctx.rbacService.getUserRoleMappings(); + } + + public async createRole( + _root: any, + args: { data: CreateRoleInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.createRole(args.data); + } + + public async updateRole( + _root: any, + args: { where: { id: number }; data: Omit }, + ctx: IContext, + ): Promise { + return ctx.rbacService.updateRole({ id: args.where.id, ...args.data }); + } + + public async createUser( + _root: any, + args: { data: CreateUserInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.createUser(args.data); + } + + public async updateUser( + _root: any, + args: { where: { id: number }; data: Omit }, + ctx: IContext, + ): Promise { + return ctx.rbacService.updateUser({ id: args.where.id, ...args.data }); + } + + public async assignRoleToUser( + _root: any, + args: { data: UserRoleInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.assignRoleToUser(args.data); + } + + public async updateUserRoles( + _root: any, + args: { data: UpdateUserRolesInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.updateUserRoles(args.data); + } + + public async removeRoleFromUser( + _root: any, + args: { data: UserRoleInput }, + ctx: IContext, + ): Promise { + return ctx.rbacService.removeRoleFromUser(args.data); + } +} diff --git a/wren-ui/src/apollo/server/scalars.ts b/wren-ui/src/apollo/server/scalars.ts index d137671730..3c8be92bdc 100644 --- a/wren-ui/src/apollo/server/scalars.ts +++ b/wren-ui/src/apollo/server/scalars.ts @@ -23,3 +23,35 @@ export const DialectSQLScalar = new GraphQLScalarType({ return ast.value as DialectSQL; }, }); + +export const BigIntStringScalar = new GraphQLScalarType({ + name: 'BigIntString', + description: + 'A bigint-compatible scalar serialized as a string to preserve precision across GraphQL and MSSQL.', + serialize(value: unknown): string { + if ( + typeof value !== 'string' && + typeof value !== 'number' && + typeof value !== 'bigint' + ) { + throw new Error('BigIntString must be a string, number, or bigint'); + } + return String(value); + }, + parseValue(value: unknown): string { + if ( + typeof value !== 'string' && + typeof value !== 'number' && + typeof value !== 'bigint' + ) { + throw new Error('BigIntString must be a string, number, or bigint'); + } + return String(value); + }, + parseLiteral(ast: any): string { + if (ast.kind !== 'StringValue' && ast.kind !== 'IntValue') { + throw new Error('BigIntString must be a string or int literal'); + } + return ast.value; + }, +}); diff --git a/wren-ui/src/apollo/server/schema.ts b/wren-ui/src/apollo/server/schema.ts index 64698ad56d..4d79bb4237 100644 --- a/wren-ui/src/apollo/server/schema.ts +++ b/wren-ui/src/apollo/server/schema.ts @@ -3,6 +3,7 @@ import { gql } from 'apollo-server-micro'; export const typeDefs = gql` scalar JSON scalar DialectSQL + scalar BigIntString enum ApiType { GENERATE_SQL @@ -21,6 +22,29 @@ export const typeDefs = gql` GET_MODELS STREAM_ASK STREAM_GENERATE_SQL + GET_ORGANIZATIONS + GET_CURRENT_ORGANIZATION + CREATE_ORGANIZATION + UPDATE_CURRENT_ORGANIZATION + SELECT_ORGANIZATION + GET_ORGANIZATION_MEMBERS + INVITE_ORGANIZATION_MEMBER + UPDATE_ORGANIZATION_MEMBER + REMOVE_ORGANIZATION_MEMBER + REMOVE_ORGANIZATION_INVITATION + ACCEPT_ORGANIZATION_INVITATION + LEAVE_ORGANIZATION + DELETE_CURRENT_ORGANIZATION + GET_PROJECTS + GET_CURRENT_PROJECT + SELECT_PROJECT + GET_CURRENT_USER + UPDATE_CURRENT_USER + DELETE_CURRENT_USER + GET_PROJECT_ACCESS + ADD_PROJECT_MEMBER + UPDATE_PROJECT_MEMBER + REMOVE_PROJECT_MEMBER } input ApiHistoryFilterInput { @@ -86,13 +110,11 @@ export const typeDefs = gql` ABS AVG COUNT - COUNT_IF MAX MIN SUM CBRT CEIL - CEILING EXP FLOOR LN @@ -110,6 +132,11 @@ export const typeDefs = gql` MUSIC } + enum WorkspaceProjectType { + AGENTIC + CLASSIC + } + enum SyncStatus { IN_PROGRESS SYNCRONIZED @@ -155,6 +182,7 @@ export const typeDefs = gql` input DataSourceInput { type: DataSourceName! properties: JSON! + projectType: WorkspaceProjectType } input SampleDatasetInput { @@ -216,9 +244,13 @@ export const typeDefs = gql` input RelationInput { fromModelId: Int! + fromModelReferenceName: String fromColumnId: Int! + fromColumnReferenceName: String toModelId: Int! + toModelReferenceName: String toColumnId: Int! + toColumnReferenceName: String type: RelationType! } @@ -230,6 +262,15 @@ export const typeDefs = gql` relations: [RelationInput]! } + input ModelingRelationshipInput { + fromModel: String! + fromColumn: String! + toModel: String! + toColumn: String! + type: RelationType! + description: String + } + input SaveTablesInput { tables: [String!]! } @@ -304,6 +345,7 @@ export const typeDefs = gql` input UpdateColumnMetadataInput { id: Int! + referenceName: String displayName: String description: String } @@ -338,6 +380,19 @@ export const typeDefs = gql` columns: [UpdateViewColumnMetadataInput!] } + input GenerateModelingSemanticsInput { + selectedModels: [String!]! + userPrompt: String! + } + + input SaveModelingSemanticInput { + modelId: Int! + referenceName: String + displayName: String + description: String + columns: [UpdateColumnMetadataInput!]! + } + type NestedFieldInfo { id: Int! displayName: String! @@ -698,12 +753,22 @@ export const typeDefs = gql` question: String sql: String taskId: String + answerContent: String + answerStatus: ThreadResponseAnswerStatus + answerErrorCode: String + answerErrorShortMessage: String + answerErrorMessage: String } input CreateThreadResponseInput { question: String sql: String taskId: String + answerContent: String + answerStatus: ThreadResponseAnswerStatus + answerErrorCode: String + answerErrorShortMessage: String + answerErrorMessage: String } input ThreadUniqueWhereInput { @@ -849,6 +914,7 @@ export const typeDefs = gql` input UpdateCurrentProjectInput { language: ProjectLanguage! + displayName: String } type Settings { @@ -865,10 +931,18 @@ export const typeDefs = gql` input PreviewSQLDataInput { sql: String! projectId: String + hash: String limit: Int dryRun: Boolean } + input DryPlanSQLDataInput { + sql: String! + projectId: String + hash: String + allowFallback: Boolean + } + # Schema Change type SchemaChange { deletedTables: [DetailedChangeTable!] @@ -929,7 +1003,7 @@ export const typeDefs = gql` } input DashboardItemWhereInput { - id: Int! + id: BigIntString! } input CreateDashboardItemInput { @@ -942,7 +1016,7 @@ export const typeDefs = gql` } input ItemLayoutInput { - itemId: Int! + itemId: BigIntString! x: Int! y: Int! w: Int! @@ -954,11 +1028,11 @@ export const typeDefs = gql` } input DeleteDashboardItemInput { - itemId: Int! + itemId: BigIntString! } input PreviewItemSQLInput { - itemId: Int! + itemId: BigIntString! limit: Int refresh: Boolean = false } @@ -1024,8 +1098,8 @@ export const typeDefs = gql` } type DashboardItem { - id: Int! - dashboardId: Int! + id: BigIntString! + dashboardId: BigIntString! type: DashboardItemType! layout: DashboardItemLayout! detail: DashboardItemDetail! @@ -1033,8 +1107,8 @@ export const typeDefs = gql` } type Dashboard { - id: Int! - projectId: Int! + id: BigIntString! + projectId: BigIntString! name: String! cacheEnabled: Boolean! scheduleFrequency: ScheduleFrequencyEnum @@ -1044,7 +1118,7 @@ export const typeDefs = gql` } type DetailedDashboard { - id: Int! + id: BigIntString! name: String! description: String cacheEnabled: Boolean! @@ -1110,6 +1184,90 @@ export const typeDefs = gql` id: Int! } + type Role { + id: Int! + name: String! + description: String + users: [User!]! + createdAt: String! + updatedAt: String! + } + + type User { + id: Int! + name: String! + email: String! + externalId: String + identityProvider: String + isActive: Boolean! + roles: [Role!]! + createdAt: String! + updatedAt: String! + } + + type UserRole { + id: Int! + userId: Int! + roleId: Int! + createdAt: String! + updatedAt: String! + } + + type UserRoleMapping { + id: Int! + userId: Int! + roleId: Int! + user: User! + role: Role! + createdAt: String! + updatedAt: String! + } + + input RoleWhereInput { + id: Int! + } + + input UserWhereInput { + id: Int! + } + + input CreateRoleInput { + name: String! + description: String + } + + input UpdateRoleInput { + name: String + description: String + } + + input CreateUserInput { + name: String! + email: String! + externalId: String + identityProvider: String + isActive: Boolean + roleIds: [Int!] + } + + input UpdateUserInput { + name: String + email: String + externalId: String + identityProvider: String + isActive: Boolean + } + + input UserRoleInput { + userId: Int! + roleId: Int! + } + + input UpdateUserRolesInput { + userId: Int! + roleIds: [Int!]! + } + # Query and Mutation type Query { # On Boarding Steps @@ -1129,21 +1287,23 @@ export const typeDefs = gql` view(where: ViewWhereUniqueInput!): ViewInfo! # Ask - askingTask(taskId: String!): AskingTask + askingTask(taskId: String): AskingTask suggestedQuestions: SuggestedQuestionResponse! threads: [Thread!]! thread(threadId: Int!): DetailedThread! - threadResponse(responseId: Int!): ThreadResponse! + threadResponse(responseId: Int): ThreadResponse nativeSql(responseId: Int!): String! # Adjustment - adjustmentTask(taskId: String!): AdjustmentTask + adjustmentTask(taskId: String): AdjustmentTask # Settings settings: Settings! # System getMDL(hash: String!): GetMDLResult! + modelingSemanticsResult(queryId: String!): JSON! + modelingRelationshipsResult(queryId: String!): JSON! # Learning learningRecord: LearningRecord! @@ -1151,7 +1311,7 @@ export const typeDefs = gql` # Recommendation questions getThreadRecommendationQuestions(threadId: Int!): RecommendedQuestionsTask! getProjectRecommendationQuestions: RecommendedQuestionsTask! - instantRecommendedQuestions(taskId: String!): RecommendedQuestionsTask! + instantRecommendedQuestions(taskId: String): RecommendedQuestionsTask! # Dashboard dashboardItems: [DashboardItem!]! @@ -1167,6 +1327,11 @@ export const typeDefs = gql` filter: ApiHistoryFilterInput pagination: ApiHistoryPaginationInput! ): ApiHistoryPaginatedResponse! + + # Administration / RBAC + roles: [Role!]! + users: [User!]! + userRoleMappings: [UserRoleMapping!]! } type Mutation { @@ -1194,6 +1359,10 @@ export const typeDefs = gql` where: ViewWhereUniqueInput! data: UpdateViewMetadataInput! ): Boolean! + generateModelingSemantics(data: GenerateModelingSemanticsInput!): JSON! + saveModelingSemantics(data: [SaveModelingSemanticInput!]!): JSON! + generateModelingRelationships: JSON! + saveModelingRelationships(data: [ModelingRelationshipInput!]!): JSON! # Relation createRelation(data: RelationInput!): JSON! @@ -1272,6 +1441,7 @@ export const typeDefs = gql` # preview previewSql(data: PreviewSQLDataInput): JSON! + dryPlanSql(data: DryPlanSQLDataInput): Boolean! # Learning saveLearningRecord(data: SaveLearningRecordInput!): LearningRecord! @@ -1312,5 +1482,14 @@ export const typeDefs = gql` data: UpdateInstructionInput! ): Instruction! deleteInstruction(where: InstructionWhereInput!): Boolean! + + # Administration / RBAC + createRole(data: CreateRoleInput!): Role! + updateRole(where: RoleWhereInput!, data: UpdateRoleInput!): Role! + createUser(data: CreateUserInput!): User! + updateUser(where: UserWhereInput!, data: UpdateUserInput!): User! + assignRoleToUser(data: UserRoleInput!): UserRole! + updateUserRoles(data: UpdateUserRolesInput!): User! + removeRoleFromUser(data: UserRoleInput!): Boolean! } `; diff --git a/wren-ui/src/apollo/server/services/askingService.ts b/wren-ui/src/apollo/server/services/askingService.ts index 82c9375856..3dc51efc45 100644 --- a/wren-ui/src/apollo/server/services/askingService.ts +++ b/wren-ui/src/apollo/server/services/askingService.ts @@ -1,9 +1,11 @@ import { IWrenAIAdaptor } from '@server/adaptors/wrenAIAdaptor'; import { AskResultStatus, + AskResultType, + AskCandidateType, + type RecommendationQuestion, RecommendationQuestionsResult, RecommendationQuestionsInput, - RecommendationQuestion, WrenAIError, RecommendationQuestionStatus, ChartStatus, @@ -12,11 +14,12 @@ import { } from '@server/models/adaptor'; import { IDeployService } from './deployService'; import { IProjectService } from './projectService'; +import { IMDLService } from './mdlService'; import { IThreadRepository, Thread } from '../repositories/threadRepository'; import { IThreadResponseRepository, ThreadResponse, - ThreadResponseAdjustmentType, + ThreadResponseAnswerDetail, } from '../repositories/threadResponseRepository'; import { getLogger } from '@server/utils'; import { isEmpty, isNil } from 'lodash'; @@ -32,7 +35,6 @@ import { Project, } from '../repositories'; import { IQueryService, PreviewDataResponse } from './queryService'; -import { IMDLService } from './mdlService'; import { ThreadRecommendQuestionBackgroundTracker, ChartBackgroundTracker, @@ -58,6 +60,7 @@ export interface Task { export interface AskingPayload { threadId?: number; language: string; + projectId?: number; } export interface AskingTaskInput { @@ -68,6 +71,7 @@ export interface AskingDetailTaskInput { question?: string; sql?: string; trackedAskingResult?: TrackedAskingResult; + answerDetail?: ThreadResponseAnswerDetail; } export interface AskingDetailTaskUpdateInput { @@ -101,6 +105,30 @@ export enum ThreadResponseAnswerStatus { INTERRUPTED = 'INTERRUPTED', } +const isAnswerGenerationInProgress = ( + status?: ThreadResponseAnswerStatus | string | null, +) => + ( + [ + ThreadResponseAnswerStatus.NOT_STARTED, + ThreadResponseAnswerStatus.FETCHING_DATA, + ThreadResponseAnswerStatus.PREPROCESSING, + ThreadResponseAnswerStatus.STREAMING, + ] as string[] + ).includes(status || ''); + +const isChartGenerationInProgress = (status?: ChartStatus | string | null) => + ([ChartStatus.FETCHING, ChartStatus.GENERATING] as string[]).includes( + status || '', + ); + +const isContextualFollowUpQuestion = (question?: string) => { + if (!question) return false; + return /\b(previous|last|above|earlier|same|that|those|them|it|this)\b/i.test( + question, + ); +}; + // adjustment input export interface AdjustmentReasoningInput { tables: string[]; @@ -301,7 +329,7 @@ class BreakdownBackgroundTracker { this.telemetry = telemetry; this.wrenAIAdaptor = wrenAIAdaptor; this.threadResponseRepository = threadResponseRepository; - this.intervalTime = 1000; + this.intervalTime = 2000; this.start(); } @@ -348,6 +376,7 @@ class BreakdownBackgroundTracker { await this.threadResponseRepository.updateOne(threadResponse.id, { breakdownDetail: updatedBreakdownDetail, }); + threadResponse.breakdownDetail = updatedBreakdownDetail; // remove the task from tracker if it is finalized if (isFinalized(result.status)) { @@ -401,6 +430,7 @@ class BreakdownBackgroundTracker { export class AskingService implements IAskingService { private wrenAIAdaptor: IWrenAIAdaptor; private deployService: IDeployService; + private mdlService: IMDLService; private projectService: IProjectService; private viewRepository: IViewRepository; private threadRepository: IThreadRepository; @@ -412,38 +442,47 @@ export class AskingService implements IAskingService { private threadRecommendQuestionBackgroundTracker: ThreadRecommendQuestionBackgroundTracker; private queryService: IQueryService; private telemetry: PostHogTelemetry; - private mdlService: IMDLService; private askingTaskTracker: IAskingTaskTracker; private askingTaskRepository: IAskingTaskRepository; private adjustmentBackgroundTracker: AdjustmentBackgroundTaskTracker; + private instantRecommendationJobs = new Map>(); + private threadRecommendationJobs = new Map>(); + private initialized = false; constructor({ telemetry, wrenAIAdaptor, deployService, + mdlService, projectService, viewRepository, threadRepository, threadResponseRepository, askingTaskRepository, queryService, - mdlService, askingTaskTracker, + chartBackgroundTracker, + chartAdjustmentBackgroundTracker, + threadRecommendQuestionBackgroundTracker, }: { telemetry: PostHogTelemetry; wrenAIAdaptor: IWrenAIAdaptor; deployService: IDeployService; + mdlService: IMDLService; projectService: IProjectService; viewRepository: IViewRepository; threadRepository: IThreadRepository; threadResponseRepository: IThreadResponseRepository; askingTaskRepository: IAskingTaskRepository; queryService: IQueryService; - mdlService: IMDLService; askingTaskTracker: IAskingTaskTracker; + chartBackgroundTracker?: ChartBackgroundTracker; + chartAdjustmentBackgroundTracker?: ChartAdjustmentBackgroundTracker; + threadRecommendQuestionBackgroundTracker?: ThreadRecommendQuestionBackgroundTracker; }) { this.wrenAIAdaptor = wrenAIAdaptor; this.deployService = deployService; + this.mdlService = mdlService; this.projectService = projectService; this.viewRepository = viewRepository; this.threadRepository = threadRepository; @@ -458,23 +497,28 @@ export class AskingService implements IAskingService { this.textBasedAnswerBackgroundTracker = new TextBasedAnswerBackgroundTracker({ wrenAIAdaptor, + threadRepository, threadResponseRepository, projectService, deployService, queryService, }); - this.chartBackgroundTracker = new ChartBackgroundTracker({ - telemetry, - wrenAIAdaptor, - threadResponseRepository, - }); + this.chartBackgroundTracker = + chartBackgroundTracker ?? + new ChartBackgroundTracker({ + telemetry, + wrenAIAdaptor, + threadResponseRepository, + }); this.chartAdjustmentBackgroundTracker = + chartAdjustmentBackgroundTracker ?? new ChartAdjustmentBackgroundTracker({ telemetry, wrenAIAdaptor, threadResponseRepository, }); this.threadRecommendQuestionBackgroundTracker = + threadRecommendQuestionBackgroundTracker ?? new ThreadRecommendQuestionBackgroundTracker({ telemetry, wrenAIAdaptor, @@ -488,17 +532,20 @@ export class AskingService implements IAskingService { }); this.askingTaskRepository = askingTaskRepository; - this.mdlService = mdlService; this.askingTaskTracker = askingTaskTracker; } + public dispose(): void { + this.chartBackgroundTracker.stop(); + this.chartAdjustmentBackgroundTracker.stop(); + this.threadRecommendQuestionBackgroundTracker.stop(); + this.askingTaskTracker.stopPolling(); + } + public async getThreadRecommendationQuestions( threadId: number, ): Promise { - const thread = await this.threadRepository.findOneBy({ id: threadId }); - if (!thread) { - throw new Error(`Thread ${threadId} not found`); - } + const thread = await this.ensureThreadInCurrentProject(threadId); // handle not started const res: ThreadRecommendQuestionResult = { @@ -519,10 +566,27 @@ export class AskingService implements IAskingService { public async generateThreadRecommendationQuestions( threadId: number, ): Promise { - const thread = await this.threadRepository.findOneBy({ id: threadId }); - if (!thread) { - throw new Error(`Thread ${threadId} not found`); + const existingJob = this.threadRecommendationJobs.get(threadId); + if (existingJob) { + logger.debug( + `thread "${threadId}" recommended questions are already being requested, reusing in-flight job`, + ); + return existingJob; + } + + const job = this.doGenerateThreadRecommendationQuestions(threadId); + this.threadRecommendationJobs.set(threadId, job); + try { + return await job; + } finally { + this.threadRecommendationJobs.delete(threadId); } + } + + private async doGenerateThreadRecommendationQuestions( + threadId: number, + ): Promise { + const thread = await this.ensureThreadInCurrentProject(threadId); if (this.threadRecommendQuestionBackgroundTracker.isExist(thread)) { logger.debug( @@ -531,8 +595,12 @@ export class AskingService implements IAskingService { return; } - const project = await this.projectService.getCurrentProject(); - const { manifest } = await this.mdlService.makeCurrentModelMDL(); + const project = await this.projectService.getProjectById(thread.projectId); + const deployment = await this.deployService.getLastDeployment(project.id); + if (!deployment) { + throw new Error(`No deployment found for project ${project.id}`); + } + const manifest = deployment.manifest; const threadResponses = await this.threadResponseRepository.findAllBy({ threadId, @@ -544,6 +612,7 @@ export class AskingService implements IAskingService { const questions = slicedThreadResponses.map(({ question }) => question); const recommendQuestionData: RecommendationQuestionsInput = { manifest, + projectId: project.id.toString(), previousQuestions: questions, ...this.getThreadRecommendationQuestionsConfig(project), }; @@ -563,6 +632,18 @@ export class AskingService implements IAskingService { } public async initialize() { + if (this.initialized) { + return; + } + + if (typeof this.askingTaskTracker?.initialize === 'function') { + await this.askingTaskTracker.initialize(); + } else { + logger.warn( + 'Asking task tracker does not expose initialize(); skipping tracker restoration', + ); + } + // list thread responses from database // filter status not finalized and put them into background tracker const threadResponses = await this.threadResponseRepository.findAll(); @@ -579,6 +660,8 @@ export class AskingService implements IAskingService { for (const threadResponse of unfininshedBreakdownThreadResponses) { this.breakdownBackgroundTracker.addTask(threadResponse); } + + this.initialized = true; } /** @@ -591,27 +674,94 @@ export class AskingService implements IAskingService { previousTaskId?: number, threadResponseId?: number, ): Promise { + const startedAt = Date.now(); const { threadId, language } = payload; - const deployId = await this.getDeployId(); + const currentProject = await this.projectService.getCurrentProject(); + let projectId = payload.projectId ?? currentProject.id; + let previousTaskState = null; + if (threadId) { + const thread = await this.ensureThreadInCurrentProject(threadId); + if (payload.projectId && payload.projectId !== thread.projectId) { + throw new Error( + `Thread ${threadId} does not belong to project ${payload.projectId}`, + ); + } + projectId = thread.projectId; + previousTaskState = await this.getLatestThreadTaskState(threadId); + } else if (projectId !== currentProject.id) { + throw new Error(`Project ${projectId} is not the active project`); + } + const deployId = await this.getDeployId(projectId); // if it's a follow-up question, then the input will have a threadId // then use the threadId to get the sql and get the steps of last thread response // construct it into AskHistory and pass to ask - const histories = threadId + const histories = threadId && isContextualFollowUpQuestion(input.question) ? await this.getAskingHistory(threadId, threadResponseId) : null; + logger.info( + `Ask timing stage=task_creation_context project_id=${projectId} thread_id=${ + threadId ?? '' + } history_count=${histories?.length ?? 0} elapsed_ms=${ + Date.now() - startedAt + }`, + ); + const trackerStartedAt = Date.now(); const response = await this.askingTaskTracker.createAskingTask({ query: input.question, histories, deployId, + projectId: projectId.toString(), configurations: { language }, rerunFromCancelled, previousTaskId, threadResponseId, }); + logger.info( + `Ask timing stage=task_creation project_id=${projectId} thread_id=${ + threadId ?? '' + } elapsed_ms=${Date.now() - trackerStartedAt} total_ms=${ + Date.now() - startedAt + }`, + ); return { id: response.queryId, }; + logger.info( + `Creating asking task: ${JSON.stringify({ + ...logContext, + question: input.question, + })}`, + ); + + try { + const response = await this.askingTaskTracker.createAskingTask({ + query: input.question, + histories, + deployId, + projectId: projectId.toString(), + configurations: { language }, + rerunFromCancelled, + previousTaskId, + threadResponseId, + }); + logger.info( + `Created asking task: ${JSON.stringify({ + ...logContext, + queryId: response.queryId, + })}`, + ); + return { + id: response.queryId, + }; + } catch (err: any) { + logger.error( + `Failed to create asking task: ${JSON.stringify(logContext)} reason=${ + err?.stack || err?.message || err + }`, + ); + throw err; + } } public async rerunAskingTask( @@ -679,7 +829,14 @@ export class AskingService implements IAskingService { */ public async createThread(input: AskingDetailTaskInput): Promise { // 1. create a thread and the first thread response - const { id } = await this.projectService.getCurrentProject(); + const project = await this.projectService.getCurrentProject(); + const { id } = project; + const trackedAskingResult = + await this.createTrackedAskingResultFromShortcutSql(input, { + projectId: id, + language: WrenAILanguage[project.language] || WrenAILanguage.EN, + }); + const thread = await this.threadRepository.createOne({ projectId: id, summary: input.question, @@ -688,15 +845,16 @@ export class AskingService implements IAskingService { const threadResponse = await this.threadResponseRepository.createOne({ threadId: thread.id, question: input.question, - sql: input.sql, - askingTaskId: input.trackedAskingResult?.taskId, + answerDetail: input.answerDetail, + sql: trackedAskingResult ? undefined : input.sql, + askingTaskId: trackedAskingResult?.taskId, }); // if queryId is provided, update asking task - if (input.trackedAskingResult?.taskId) { + if (trackedAskingResult?.taskId) { await this.askingTaskTracker.bindThreadResponse( - input.trackedAskingResult.taskId, - input.trackedAskingResult.queryId, + trackedAskingResult.taskId, + trackedAskingResult.queryId, thread.id, threadResponse.id, ); @@ -720,12 +878,14 @@ export class AskingService implements IAskingService { throw new Error('Update thread input is empty'); } + await this.ensureThreadInCurrentProject(threadId); return this.threadRepository.updateOne(threadId, { summary: input.summary, }); } public async deleteThread(threadId: number): Promise { + await this.ensureThreadInCurrentProject(threadId); await this.threadRepository.deleteOne(threadId); } @@ -733,26 +893,28 @@ export class AskingService implements IAskingService { input: AskingDetailTaskInput, threadId: number, ): Promise { - const thread = await this.threadRepository.findOneBy({ - id: threadId, - }); + const thread = await this.ensureThreadInCurrentProject(threadId); - if (!thread) { - throw new Error(`Thread ${threadId} not found`); - } + const project = await this.projectService.getProjectById(thread.projectId); + const trackedAskingResult = + await this.createTrackedAskingResultFromShortcutSql(input, { + projectId: thread.projectId, + language: WrenAILanguage[project.language] || WrenAILanguage.EN, + }); const threadResponse = await this.threadResponseRepository.createOne({ threadId: thread.id, question: input.question, - sql: input.sql, - askingTaskId: input.trackedAskingResult?.taskId, + answerDetail: input.answerDetail, + sql: trackedAskingResult ? undefined : input.sql, + askingTaskId: trackedAskingResult?.taskId, }); // if queryId is provided, update asking task - if (input.trackedAskingResult?.taskId) { + if (trackedAskingResult?.taskId) { await this.askingTaskTracker.bindThreadResponse( - input.trackedAskingResult.taskId, - input.trackedAskingResult.queryId, + trackedAskingResult.taskId, + trackedAskingResult.queryId, thread.id, threadResponse.id, ); @@ -771,9 +933,14 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${responseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); return await this.threadResponseRepository.updateOne(responseId, { sql: data.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, }); } @@ -826,6 +993,13 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${threadResponseId} not found`); } + if (isAnswerGenerationInProgress(threadResponse.answerDetail?.status)) { + logger.debug( + `Thread response ${threadResponseId} answer generation already in progress, skipping duplicate request`, + ); + return threadResponse; + } + // update with initial status const updatedThreadResponse = await this.threadResponseRepository.updateOne( threadResponse.id, @@ -854,10 +1028,42 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${threadResponseId} not found`); } + if (isChartGenerationInProgress(threadResponse.chartDetail?.status)) { + logger.debug( + `Thread response ${threadResponseId} chart generation already in progress, skipping duplicate request`, + ); + return threadResponse; + } + + const project = await this.getProjectForThreadResponse(threadResponse); + const deployment = await this.deployService.getLastDeployment(project.id); + if (!deployment) { + throw new Error(`No deployment found for project ${project.id}`); + } + const manifest = deployment.manifest; + let chartData: PreviewDataResponse; + try { + chartData = (await this.queryService.preview(threadResponse.sql, { + project, + manifest, + modelingOnly: false, + limit: 100, + cacheEnabled: false, + })) as PreviewDataResponse; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn( + `Preview failed before chart generation for response ${threadResponse.id}. ${message}`, + ); + throw error; + } + // 1. create a task on AI service to generate the chart const response = await this.wrenAIAdaptor.generateChart({ query: threadResponse.question, sql: threadResponse.sql, + data: chartData, + projectId: project.id.toString(), configurations, }); @@ -891,10 +1097,45 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${threadResponseId} not found`); } + if ( + isChartGenerationInProgress(threadResponse.chartDetail?.status) && + threadResponse.chartDetail?.adjustment + ) { + logger.debug( + `Thread response ${threadResponseId} chart adjustment already in progress, skipping duplicate request`, + ); + return threadResponse; + } + + const project = await this.getProjectForThreadResponse(threadResponse); + const deployment = await this.deployService.getLastDeployment(project.id); + if (!deployment) { + throw new Error(`No deployment found for project ${project.id}`); + } + const manifest = deployment.manifest; + let chartData: PreviewDataResponse; + try { + chartData = (await this.queryService.preview(threadResponse.sql, { + project, + manifest, + modelingOnly: false, + limit: 500, + cacheEnabled: false, + })) as PreviewDataResponse; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.warn( + `Preview failed before chart adjustment for response ${threadResponse.id}. ${message}`, + ); + throw error; + } + // 1. create a task on AI service to adjust the chart const response = await this.wrenAIAdaptor.adjustChart({ query: threadResponse.question, sql: threadResponse.sql, + data: chartData, + projectId: project.id.toString(), adjustmentOption: input, chartSchema: threadResponse.chartDetail?.chartSchema, configurations, @@ -919,11 +1160,19 @@ export class AskingService implements IAskingService { } public async getResponsesWithThread(threadId: number) { + await this.ensureThreadInCurrentProject(threadId); return this.threadResponseRepository.getResponsesWithThread(threadId); } public async getResponse(responseId: number) { - return this.threadResponseRepository.findOneBy({ id: responseId }); + const response = await this.threadResponseRepository.findOneBy({ + id: responseId, + }); + if (!response) { + return null; + } + await this.ensureThreadInCurrentProject(response.threadId); + return response; } public async previewData(responseId: number, limit?: number) { @@ -931,15 +1180,19 @@ export class AskingService implements IAskingService { if (!response) { throw new Error(`Thread response ${responseId} not found`); } - const project = await this.projectService.getCurrentProject(); + const project = await this.getProjectForThreadResponse(response); const deployment = await this.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; + if (!deployment) { + throw new Error(`No deployment found for project ${project.id}`); + } + const manifest = deployment.manifest; const eventName = TelemetryEvent.HOME_PREVIEW_ANSWER; try { const data = (await this.queryService.preview(response.sql, { project, - manifest: mdl, + manifest, limit, + cacheEnabled: false, })) as PreviewDataResponse; this.telemetry.sendEvent(eventName, { sql: response.sql }); return data; @@ -971,17 +1224,21 @@ export class AskingService implements IAskingService { if (!response) { throw new Error(`Thread response ${responseId} not found`); } - const project = await this.projectService.getCurrentProject(); + const project = await this.getProjectForThreadResponse(response); const deployment = await this.deployService.getLastDeployment(project.id); - const mdl = deployment.manifest; + if (!deployment) { + throw new Error(`No deployment found for project ${project.id}`); + } + const manifest = deployment.manifest; const steps = response?.breakdownDetail?.steps; const sql = safeFormatSQL(constructCteSql(steps, stepIndex)); const eventName = TelemetryEvent.HOME_PREVIEW_ANSWER; try { const data = (await this.queryService.preview(sql, { project, - manifest: mdl, + manifest, limit, + cacheEnabled: false, })) as PreviewDataResponse; this.telemetry.sendEvent(eventName, { sql }); return data; @@ -1000,12 +1257,44 @@ export class AskingService implements IAskingService { input: InstantRecommendedQuestionsInput, ): Promise { const project = await this.projectService.getCurrentProject(); - const { manifest } = await this.deployService.getLastDeployment(project.id); + const key = JSON.stringify({ + projectId: project.id, + previousQuestions: input.previousQuestions || [], + }); + const existingJob = this.instantRecommendationJobs.get(key); + if (existingJob) { + logger.debug('instant recommended questions are already being requested'); + return existingJob; + } + + const job = this.doCreateInstantRecommendedQuestions(input, project); + this.instantRecommendationJobs.set(key, job); + try { + return await job; + } finally { + this.instantRecommendationJobs.delete(key); + } + } + + private async doCreateInstantRecommendedQuestions( + input: InstantRecommendedQuestionsInput, + project?: Project, + ): Promise { + const currentProject = + project ?? (await this.projectService.getCurrentProject()); + const deployment = await this.deployService.getLastDeployment( + currentProject.id, + ); + if (!deployment) { + throw new Error(`No deployment found for project ${currentProject.id}`); + } + const manifest = deployment.manifest; const response = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, + projectId: currentProject.id.toString(), previousQuestions: input.previousQuestions, - ...this.getThreadRecommendationQuestionsConfig(project), + ...this.getThreadRecommendationQuestionsConfig(currentProject), }); return { id: response.queryId }; } @@ -1035,8 +1324,11 @@ export class AskingService implements IAskingService { throw new Error(`Thread response ${responseId} not found`); } - if (response.answerDetail?.status === status) { - return; + if ( + response.answerDetail?.status === status && + (!content || response.answerDetail?.content === content) + ) { + return response; } const updatedResponse = await this.threadResponseRepository.updateOne( @@ -1045,7 +1337,7 @@ export class AskingService implements IAskingService { answerDetail: { ...response.answerDetail, status, - content, + content: content ?? response.answerDetail?.content, }, }, ); @@ -1053,10 +1345,45 @@ export class AskingService implements IAskingService { return updatedResponse; } - private async getDeployId() { - const { id } = await this.projectService.getCurrentProject(); - const lastDeploy = await this.deployService.getLastDeployment(id); - return lastDeploy.hash; + private async getDeployId(projectId?: number) { + const project = projectId + ? await this.projectService.getProjectById(projectId) + : await this.projectService.getCurrentProject(); + const { manifest } = await this.mdlService.makeModelMDL(project); + return this.deployService.ensureDeploymentPrepared(project.id, manifest); + } + + private async getProjectForThreadResponse(threadResponse: ThreadResponse) { + const thread = await this.threadRepository.findOneBy({ + id: threadResponse.threadId, + }); + if (!thread) { + logger.warn( + `Thread ${threadResponse.threadId} for response ${threadResponse.id} not found; falling back to current project`, + ); + return this.projectService.getCurrentProject(); + } + + return this.projectService.getProjectById(thread.projectId); + } + + private async createTrackedAskingResultFromShortcutSql( + input: AskingDetailTaskInput, + payload: AskingPayload, + ): Promise { + if (input.trackedAskingResult || !input.sql || !input.question) { + return input.trackedAskingResult; + } + + const task = await this.createAskingTask( + { question: input.question }, + payload, + ); + const trackedAskingResult = await this.getAskingTask(task.id); + if (!trackedAskingResult?.taskId) { + throw new Error(`Asking task ${task.id} not found`); + } + return trackedAskingResult; } public async adjustThreadResponseWithSQL( @@ -1069,19 +1396,51 @@ export class AskingService implements IAskingService { if (!response) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(response.threadId); - return await this.threadResponseRepository.createOne({ - sql: input.sql, - threadId: response.threadId, - question: response.question, - adjustment: { - type: ThreadResponseAdjustmentType.APPLY_SQL, - payload: { - originalThreadResponseId: response.id, - sql: input.sql, - }, - }, + const project = await this.getProjectForThreadResponse(response); + const deployment = await this.deployService.getLastDeployment(project.id); + if (!deployment) { + throw new Error(`No deployment found for project ${project.id}`); + } + const manifest = deployment.manifest; + await this.queryService.preview(input.sql, { + project, + manifest, + modelingOnly: false, + limit: 1, + cacheEnabled: false, }); + + const updatedResponse = await this.threadResponseRepository.updateOne( + response.id, + { + sql: input.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, + }, + ); + + if (response.askingTaskId) { + await this.askingTaskRepository.updateOne(response.askingTaskId, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FINISHED, + response: [ + { + type: AskCandidateType.LLM, + sql: input.sql, + }, + ], + error: null, + invalidSql: null, + }, + }); + } + + return updatedResponse; } public async adjustThreadResponseAnswer( @@ -1096,6 +1455,7 @@ export class AskingService implements IAskingService { if (!originalThreadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(originalThreadResponse.threadId); const { createdThreadResponse } = await this.adjustmentBackgroundTracker.createAdjustmentTask({ @@ -1127,6 +1487,7 @@ export class AskingService implements IAskingService { if (!threadResponse) { throw new Error(`Thread response ${threadResponseId} not found`); } + await this.ensureThreadInCurrentProject(threadResponse.threadId); const { queryId } = await this.adjustmentBackgroundTracker.rerunAdjustmentTask({ @@ -1183,9 +1544,56 @@ export class AskingService implements IAskingService { return { maxCategories: config.threadRecommendationQuestionMaxCategories, maxQuestions: config.threadRecommendationQuestionsMaxQuestions, + regenerate: true, configuration: { language: WrenAILanguage[project.language] || WrenAILanguage.EN, }, }; } + + private async getLatestThreadTaskState(threadId: number) { + try { + const [latestResponse] = + await this.threadResponseRepository.getResponsesWithThread(threadId, 1); + if (!latestResponse) { + return null; + } + + const task = latestResponse.askingTaskId + ? await this.askingTaskRepository.findOneBy({ + id: latestResponse.askingTaskId, + }) + : null; + const detail = task?.detail as any; + + return { + threadResponseId: latestResponse.id, + askingTaskId: latestResponse.askingTaskId ?? null, + queryId: task?.queryId ?? null, + status: detail?.status ?? null, + type: detail?.type ?? null, + hasSql: !!latestResponse.sql, + }; + } catch (err: any) { + logger.warn( + `Failed to inspect latest thread task state for thread ${threadId}: ${ + err?.message || err + }`, + ); + return null; + } + } + + private async ensureThreadInCurrentProject(threadId: number): Promise { + const [thread, project] = await Promise.all([ + this.threadRepository.findOneBy({ id: threadId }), + this.projectService.getCurrentProject(), + ]); + + if (!thread || thread.projectId !== project.id) { + throw new Error(`Thread ${threadId} not found in current project`); + } + + return thread; + } } diff --git a/wren-ui/src/apollo/server/services/askingTaskTracker.ts b/wren-ui/src/apollo/server/services/askingTaskTracker.ts index a2f18a61f0..5eceb1188e 100644 --- a/wren-ui/src/apollo/server/services/askingTaskTracker.ts +++ b/wren-ui/src/apollo/server/services/askingTaskTracker.ts @@ -1,4 +1,5 @@ import { getLogger } from '@server/utils'; +import { isEqual } from 'lodash'; import { AskResult, AskResultType, @@ -21,6 +22,8 @@ interface TrackedTask { queryId: string; taskId?: number; lastPolled: number; + nextPollAt: number; + pollDelay: number; question?: string; result?: AskResult; isFinalized: boolean; @@ -45,6 +48,8 @@ export interface IAskingTaskTracker { getAskingResult(queryId: string): Promise; getAskingResultById(id: number): Promise; cancelAskingTask(queryId: string): Promise; + initialize(): Promise; + stopPolling(): void; bindThreadResponse( id: number, queryId: string, @@ -54,16 +59,19 @@ export interface IAskingTaskTracker { } export class AskingTaskTracker implements IAskingTaskTracker { + private readonly minPollDelay = 1000; + private readonly maxPollDelay = 10000; private wrenAIAdaptor: IWrenAIAdaptor; private askingTaskRepository: IAskingTaskRepository; private trackedTasks: Map = new Map(); private trackedTasksById: Map = new Map(); private pollingInterval: number; private memoryRetentionTime: number; - private pollingIntervalId: NodeJS.Timeout; + private pollingIntervalId?: NodeJS.Timeout; private runningJobs = new Set(); private threadResponseRepository: IThreadResponseRepository; private viewRepository: IViewRepository; + private initialized = false; constructor({ wrenAIAdaptor, @@ -93,8 +101,27 @@ export class AskingTaskTracker implements IAskingTaskTracker { input: CreateAskingTaskInput, ): Promise<{ queryId: string }> { try { + logger.info( + `Creating asking task request: ${JSON.stringify({ + projectId: input.projectId ?? null, + deployId: input.deployId ?? null, + hasHistories: !!input.histories?.length, + historyCount: input.histories?.length ?? 0, + rerunFromCancelled: !!input.rerunFromCancelled, + previousTaskId: input.previousTaskId ?? null, + threadResponseId: input.threadResponseId ?? null, + question: input.query, + })}`, + ); // Call the AI service to create a task + const startedAt = Date.now(); + const aiRequestStartedAt = Date.now(); const response = await this.wrenAIAdaptor.ask(input); + logger.info( + `Ask timing stage=task_creation_ai_request project_id=${ + input.projectId ?? '' + } elapsed_ms=${Date.now() - aiRequestStartedAt}`, + ); const queryId = response.queryId; // validate the input @@ -111,7 +138,15 @@ export class AskingTaskTracker implements IAskingTaskTracker { const task = { queryId, lastPolled: Date.now(), + nextPollAt: Date.now(), + pollDelay: this.minPollDelay, question: input.query, + result: { + type: null, + status: AskResultStatus.UNDERSTANDING, + response: null, + error: null, + }, isFinalized: false, rerunFromCancelled: input.rerunFromCancelled, } as TrackedTask; @@ -140,13 +175,41 @@ export class AskingTaskTracker implements IAskingTaskTracker { // update the query id in database await this.askingTaskRepository.updateOne(input.previousTaskId, { queryId, + detail: task.result, + }); + } else { + const dbStartedAt = Date.now(); + const createdTask = await this.askingTaskRepository.createOne({ + queryId, + question: input.query, + detail: task.result, }); + logger.info( + `Ask timing stage=task_creation_db project_id=${ + input.projectId ?? '' + } elapsed_ms=${Date.now() - dbStartedAt}`, + ); + task.taskId = createdTask.id; + this.trackedTasksById.set(createdTask.id, task); } - logger.info(`Created asking task with queryId: ${queryId}`); + logger.info( + `Created asking task with queryId: ${queryId}`, + ); + logger.info( + `Ask timing stage=task_creation_tracker project_id=${ + input.projectId ?? '' + } query_id=${queryId} elapsed_ms=${Date.now() - startedAt}`, + ); return { queryId }; - } catch (err) { - logger.error(`Failed to create asking task: ${err}`); + } catch (err: any) { + logger.error( + `Failed to create asking task for projectId=${ + input.projectId ?? 'unknown' + }, deployId=${input.deployId ?? 'unknown'}: ${ + err?.stack || err?.message || err + }`, + ); throw err; } } @@ -167,7 +230,11 @@ export class AskingTaskTracker implements IAskingTaskTracker { } // If not in memory or no result yet, check the database - return this.getAskingResultFromDB({ queryId }); + const result = await this.getAskingResultFromDB({ queryId }); + if (result && !this.isTaskFinalized(result.status)) { + this.restoreTrackedTask(result); + } + return result; } public async getAskingResultById( @@ -178,16 +245,53 @@ export class AskingTaskTracker implements IAskingTaskTracker { return this.getAskingResult(task.queryId); } - return this.getAskingResultFromDB({ taskId: id }); + const result = await this.getAskingResultFromDB({ taskId: id }); + if (result && !this.isTaskFinalized(result.status)) { + this.restoreTrackedTask(result); + } + return result; } public async cancelAskingTask(queryId: string): Promise { await this.wrenAIAdaptor.cancelAsk(queryId); } + public async initialize(): Promise { + if (this.initialized) return; + + const taskRecords = await this.askingTaskRepository.findAll(); + await Promise.all( + taskRecords.map(async (taskRecord) => { + const detail = taskRecord.detail as AskResult | undefined; + if ( + !taskRecord.queryId || + !detail || + this.isTaskFinalized(detail.status) + ) { + return; + } + + if (this.isStaleUnfinishedTask(taskRecord)) { + await this.finalizeStaleTask(taskRecord); + return; + } + + this.restoreTrackedTask({ + ...detail, + queryId: taskRecord.queryId, + question: taskRecord.question, + taskId: taskRecord.id, + }); + }), + ); + + this.initialized = true; + } + public stopPolling(): void { if (this.pollingIntervalId) { clearInterval(this.pollingIntervalId); + this.pollingIntervalId = undefined; } } @@ -216,6 +320,10 @@ export class AskingTaskTracker implements IAskingTaskTracker { } private startPolling(): void { + if (this.pollingIntervalId) { + return; + } + this.pollingIntervalId = setInterval(() => { this.pollTasks(); }, this.pollingInterval); @@ -249,16 +357,32 @@ export class AskingTaskTracker implements IAskingTaskTracker { return; } + if (now < task.nextPollAt) { + return; + } + // Mark the job as running this.runningJobs.add(queryId); // Poll for updates - logger.info(`Polling for updates for task ${queryId}`); - const result = await this.wrenAIAdaptor.getAskResult(queryId); + logger.debug(`Polling for updates for task ${queryId}`); + const pollStartedAt = Date.now(); + const resultFromAIService = + await this.wrenAIAdaptor.getAskResult(queryId); + logger.info( + `Ask timing stage=task_poll query_id=${queryId} elapsed_ms=${ + Date.now() - pollStartedAt + }`, + ); + const result = this.isMissingInAIService(resultFromAIService) + ? this.createExpiredTaskResult() + : resultFromAIService; task.lastPolled = now; + const resultChanged = this.isResultChanged(task.result, result); + this.scheduleNextPoll(task, result.status, resultChanged); // if result is not changed, we don't need to update the database - if (!this.isResultChanged(task.result, result)) { + if (!resultChanged) { this.runningJobs.delete(queryId); return; } @@ -303,6 +427,8 @@ export class AskingTaskTracker implements IAskingTaskTracker { }, }, ); + } else { + await this.updateTaskInDatabase({ queryId }, task); } this.runningJobs.delete(queryId); return; @@ -339,26 +465,26 @@ export class AskingTaskTracker implements IAskingTaskTracker { ); // Run all jobs in parallel - Promise.allSettled(jobs.map((job) => job())).then((results) => { - // Log any rejected promises - results.forEach((result, index) => { - if (result.status === 'rejected') { - logger.error(`Job ${index} failed: ${result.reason}`); - } - }); + const results = await Promise.allSettled(jobs.map((job) => job())); - // Clean up tasks that have been in memory too long - if (tasksToRemove.length > 0) { - logger.info( - `Cleaning up tasks that have been in memory too long. Tasks: ${tasksToRemove.join( - ', ', - )}`, - ); - } - for (const queryId of tasksToRemove) { - this.trackedTasks.delete(queryId); + // Log any rejected promises + results.forEach((result, index) => { + if (result.status === 'rejected') { + logger.error(`Job ${index} failed: ${result.reason}`); } }); + + // Clean up tasks that have been in memory too long + if (tasksToRemove.length > 0) { + logger.info( + `Cleaning up tasks that have been in memory too long. Tasks: ${tasksToRemove.join( + ', ', + )}`, + ); + } + for (const queryId of tasksToRemove) { + this.trackedTasks.delete(queryId); + } } private async updateThreadResponseWhenTaskFinalized( @@ -377,10 +503,17 @@ export class AskingTaskTracker implements IAskingTaskTracker { await this.threadResponseRepository.updateOne(task.threadResponseId, { sql: view.statement, viewId: response.viewId, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, }); } else { await this.threadResponseRepository.updateOne(task.threadResponseId, { sql: response?.sql, + viewId: null, + answerDetail: null, + breakdownDetail: null, + chartDetail: null, }); } } @@ -403,6 +536,20 @@ export class AskingTaskTracker implements IAskingTaskTracker { return null; } + if ( + taskRecord.detail && + !this.isTaskFinalized((taskRecord.detail as AskResult).status) && + this.isStaleUnfinishedTask(taskRecord) + ) { + await this.finalizeStaleTask(taskRecord); + taskRecord = await this.askingTaskRepository.findOneBy({ + id: taskRecord.id, + }); + if (!taskRecord) { + return null; + } + } + return { ...(taskRecord?.detail as AskResult), queryId: queryId || taskRecord?.queryId, @@ -457,15 +604,153 @@ export class AskingTaskTracker implements IAskingTaskTracker { ].includes(status); } + private isStaleUnfinishedTask(taskRecord: AskingTask): boolean { + const updatedAt = this.toTimestamp(taskRecord.updatedAt); + const createdAt = this.toTimestamp(taskRecord.createdAt); + const lastKnownUpdate = updatedAt ?? createdAt; + + if (!lastKnownUpdate) { + logger.warn( + `Treating unfinished task ${taskRecord.queryId} as stale because it has no valid timestamp`, + ); + return true; + } + + return Date.now() - lastKnownUpdate > this.memoryRetentionTime; + } + + private toTimestamp(value?: Date | string | number): number | null { + if (!value) { + return null; + } + + if (value instanceof Date) { + const timestamp = value.getTime(); + return Number.isFinite(timestamp) ? timestamp : null; + } + + if (typeof value === 'number') { + const timestamp = value < 10_000_000_000 ? value * 1000 : value; + return Number.isFinite(timestamp) ? timestamp : null; + } + + const trimmed = value.trim(); + if (!trimmed) { + return null; + } + + const numericValue = Number(trimmed); + if (Number.isFinite(numericValue)) { + return this.toTimestamp(numericValue); + } + + const timestamp = Date.parse(trimmed); + if (Number.isFinite(timestamp)) { + return timestamp; + } + + const normalizedTimestamp = Date.parse(trimmed.replace(' ', 'T')); + return Number.isFinite(normalizedTimestamp) ? normalizedTimestamp : null; + } + + private async finalizeStaleTask(taskRecord: AskingTask): Promise { + const result = this.createExpiredTaskResult(); + + await this.askingTaskRepository.updateOne(taskRecord.id, { + detail: result, + }); + } + + private createExpiredTaskResult(): AskResult { + return { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }; + } + + private isMissingInAIService(result: AskResult): boolean { + return ( + result.status === AskResultStatus.FAILED && + result.error?.code === Errors.GeneralErrorCodes.ASK_RESULT_NOT_FOUND + ); + } + private isResultChanged( previousResult: AskResult, newResult: AskResult, ): boolean { - // check status change - if (previousResult?.status !== newResult.status) { - return true; + return !isEqual( + this.getComparableAskResult(previousResult), + this.getComparableAskResult(newResult), + ); + } + + private restoreTrackedTask(result: TrackedAskingResult) { + const restoredTask: TrackedTask = { + queryId: result.queryId, + taskId: result.taskId, + lastPolled: Date.now(), + nextPollAt: Date.now(), + pollDelay: this.minPollDelay, + question: result.question, + result: this.getComparableAskResult(result), + isFinalized: false, + }; + this.trackedTasks.set(result.queryId, restoredTask); + if (result.taskId) { + this.trackedTasksById.set(result.taskId, restoredTask); + } + } + + private getComparableAskResult(result: AskResult | TrackedAskingResult) { + return { + status: result?.status ?? null, + type: result?.type ?? null, + response: result?.response ?? null, + error: result?.error ?? null, + rephrasedQuestion: result?.rephrasedQuestion ?? null, + intentReasoning: result?.intentReasoning ?? null, + sqlGenerationReasoning: result?.sqlGenerationReasoning ?? null, + retrievedTables: result?.retrievedTables ?? null, + invalidSql: result?.invalidSql ?? null, + traceId: result?.traceId ?? null, + } as AskResult; + } + + private scheduleNextPoll( + task: TrackedTask, + status: AskResultStatus, + resultChanged: boolean, + ) { + if (this.isTaskFinalized(status)) { + task.nextPollAt = Number.MAX_SAFE_INTEGER; + return; + } + + const baseDelay = this.getBasePollDelay(status); + task.pollDelay = resultChanged + ? baseDelay + : Math.min(Math.max(task.pollDelay * 1.5, baseDelay), this.maxPollDelay); + task.nextPollAt = Date.now() + task.pollDelay; + } + + private getBasePollDelay(status: AskResultStatus) { + if ( + [ + AskResultStatus.UNDERSTANDING, + AskResultStatus.SEARCHING, + AskResultStatus.PLANNING, + ].includes(status) + ) { + return this.minPollDelay; } - return false; + return 1500; } } diff --git a/wren-ui/src/apollo/server/services/dashboardService.ts b/wren-ui/src/apollo/server/services/dashboardService.ts index 59818eef9f..190e168549 100644 --- a/wren-ui/src/apollo/server/services/dashboardService.ts +++ b/wren-ui/src/apollo/server/services/dashboardService.ts @@ -22,10 +22,11 @@ const logger = getLogger('DashboardService'); logger.level = 'debug'; export interface CreateDashboardItemInput { - dashboardId: number; + dashboardId: string | number; type: DashboardItemType; sql: string; chartSchema: DashboardItemDetail['chartSchema']; + previewDataSnapshot?: DashboardItemDetail['previewDataSnapshot']; } export interface UpdateDashboardItemInput { @@ -33,25 +34,25 @@ export interface UpdateDashboardItemInput { } export type UpdateDashboardItemLayouts = (DashboardItemLayout & { - itemId: number; + itemId: string | number; })[]; export interface IDashboardService { initDashboard(): Promise; getCurrentDashboard(): Promise; - getDashboardItem(dashboardItemId: number): Promise; - getDashboardItems(dashboardId: number): Promise; + getDashboardItem(dashboardItemId: string | number): Promise; + getDashboardItems(dashboardId: string | number): Promise; createDashboardItem(input: CreateDashboardItemInput): Promise; updateDashboardItem( - dashboardItemId: number, + dashboardItemId: string | number, input: UpdateDashboardItemInput, ): Promise; - deleteDashboardItem(dashboardItemId: number): Promise; + deleteDashboardItem(dashboardItemId: string | number): Promise; updateDashboardItemLayouts( layouts: UpdateDashboardItemLayouts, ): Promise; setDashboardSchedule( - dashboardId: number, + dashboardId: string | number, data: SetDashboardCacheData, ): Promise; parseCronExpression(dashboard: Dashboard): DashboardSchedule; @@ -77,7 +78,7 @@ export class DashboardService implements IDashboardService { } public async setDashboardSchedule( - dashboardId: number, + dashboardId: string | number, data: SetDashboardCacheData, ): Promise { try { @@ -127,6 +128,9 @@ export class DashboardService implements IDashboardService { public async initDashboard(): Promise { const project = await this.projectService.getCurrentProject(); + logger.debug( + `Initializing dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const existingDashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); @@ -140,49 +144,75 @@ export class DashboardService implements IDashboardService { public async getCurrentDashboard(): Promise { const project = await this.projectService.getCurrentProject(); + logger.debug( + `Loading current dashboard for project ${String(project.id)} (${project.type || 'unknown'})`, + ); const dashboard = await this.dashboardRepository.findOneBy({ projectId: project.id, }); - return { ...dashboard }; + if (dashboard) { + logger.debug( + `Resolved dashboard ${String(dashboard.id)} for project ${String(project.id)}`, + ); + return dashboard; + } + + logger.debug( + `Dashboard not found for project ${project.id}; initializing a default dashboard.`, + ); + return await this.initDashboard(); } public async getDashboardItem( - dashboardItemId: number, + dashboardItemId: string | number, ): Promise { + const dashboard = await this.getCurrentDashboard(); const item = await this.dashboardItemRepository.findOneBy({ id: dashboardItemId, }); - if (!item) { + if (!item || String(item.dashboardId) !== String(dashboard.id)) { throw new Error('Dashboard item not found.'); } return item; } public async getDashboardItems( - dashboardId: number, + dashboardId: string | number, ): Promise { - return await this.dashboardItemRepository.findAllBy({ + const items = await this.dashboardItemRepository.findAllBy({ dashboardId, }); + logger.debug( + `Loaded ${items.length} dashboard item(s) for dashboard ${String(dashboardId)}`, + ); + return items; } public async createDashboardItem( input: CreateDashboardItemInput, ): Promise { const layout = await this.calculateNewLayout(input.dashboardId); - return await this.dashboardItemRepository.createOne({ + logger.debug( + `Creating dashboard item for dashboard ${String(input.dashboardId)} with type ${input.type}`, + ); + const dashboardItem = await this.dashboardItemRepository.createOne({ dashboardId: input.dashboardId, type: input.type, detail: { sql: input.sql, chartSchema: input.chartSchema, + previewDataSnapshot: input.previewDataSnapshot, }, layout, }); + logger.debug( + `Created dashboard item ${String(dashboardItem.id)} for dashboard ${String(input.dashboardId)}`, + ); + return dashboardItem; } public async updateDashboardItem( - dashboardItemId: number, + dashboardItemId: string | number, input: UpdateDashboardItemInput, ): Promise { return await this.dashboardItemRepository.updateOne(dashboardItemId, { @@ -193,10 +223,18 @@ export class DashboardService implements IDashboardService { public async updateDashboardItemLayouts( layouts: UpdateDashboardItemLayouts, ): Promise { + const dashboard = await this.getCurrentDashboard(); + const dashboardItems = await this.dashboardItemRepository.findAllBy({ + dashboardId: dashboard.id, + }); + const dashboardItemIds = new Set( + dashboardItems.map((item) => String(item.id)), + ); const updatedItems: DashboardItem[] = []; const isValidLayouts = layouts.every( (layout) => layout.itemId && + dashboardItemIds.has(String(layout.itemId)) && layout.x >= 0 && layout.y >= 0 && layout.w > 0 && @@ -224,13 +262,15 @@ export class DashboardService implements IDashboardService { return updatedItems; } - public async deleteDashboardItem(dashboardItemId: number): Promise { + public async deleteDashboardItem( + dashboardItemId: string | number, + ): Promise { await this.dashboardItemRepository.deleteOne(dashboardItemId); return true; } private async calculateNewLayout( - dashboardId: number, + dashboardId: string | number, ): Promise { const dashboardItems = await this.dashboardItemRepository.findAllBy({ dashboardId, diff --git a/wren-ui/src/apollo/server/services/deployService.ts b/wren-ui/src/apollo/server/services/deployService.ts index a6885146cb..6aad6810c4 100644 --- a/wren-ui/src/apollo/server/services/deployService.ts +++ b/wren-ui/src/apollo/server/services/deployService.ts @@ -1,4 +1,7 @@ -import { WrenAIDeployStatusEnum } from '@server/models/adaptor'; +import { + WrenAIDeployStatusEnum, + WrenAISystemStatus, +} from '@server/models/adaptor'; import { IWrenAIAdaptor } from '../adaptors/wrenAIAdaptor'; import { Deploy, @@ -7,7 +10,7 @@ import { } from '../repositories/deployLogRepository'; import { Manifest } from '../mdl/type'; import { createHash } from 'node:crypto'; -import { getLogger } from '@server/utils'; +import { getLogger, normalizeManifest } from '@server/utils'; import { PostHogTelemetry, TelemetryEvent, @@ -17,9 +20,12 @@ import { const logger = getLogger('DeployService'); logger.level = 'debug'; +const STALE_DEPLOYMENT_MS = 10 * 60 * 1000; + export interface DeployResponse { status: DeployStatusEnum; error?: string; + hash?: string; } export interface MDLSyncResponse { @@ -33,9 +39,16 @@ export interface IDeployService { force?: boolean, ): Promise; getLastDeployment(projectId: number): Promise; + ensureDeploymentPrepared(projectId: number, manifest?: Manifest): Promise; getInProgressDeployment(projectId: number): Promise; createMDLHash(manifest: Manifest, projectId: number): string; - getMDLByHash(hash: string): Promise; + isSameDeployment( + manifest: Manifest, + projectId: number, + deployment?: Deploy | null, + ): boolean; + getManifestByHash(hash: string): Promise; + getMDLByHash(hash: string): Promise; deleteAllByProjectId(projectId: number): Promise; } @@ -67,15 +80,100 @@ export class DeployService implements IDeployService { return lastDeploy; } + public async ensureDeploymentPrepared( + projectId: number, + manifest?: Manifest, + ): Promise { + const lastDeploy = + await this.deployLogRepository.findLastProjectDeployLog(projectId); + if (!lastDeploy) { + throw new Error(`No deployment found for project ${projectId}`); + } + const activeManifest = manifest || lastDeploy.manifest; + const activeHash = this.createMDLHash(activeManifest, projectId); + + if (lastDeploy.hash === activeHash) { + try { + const status = await this.wrenAIAdaptor.getDeployStatus( + activeHash, + projectId, + ); + if (status === WrenAISystemStatus.FINISHED) { + return activeHash; + } + logger.warn( + `Deployment ${activeHash} is not ready in AI service: ${status}`, + ); + } catch (err: any) { + logger.warn( + `Deployment ${activeHash} is not available in AI service: ${err.message}`, + ); + } + } else { + logger.warn( + `Deployment ${lastDeploy.hash} does not match current manifest hash ${activeHash}; preparing current hash.`, + ); + try { + const status = await this.wrenAIAdaptor.getDeployStatus( + activeHash, + projectId, + ); + if (status === WrenAISystemStatus.FINISHED) { + await this.deployLogRepository.updateOne(lastDeploy.id, { + manifest: normalizeManifest(activeManifest), + hash: activeHash, + status: DeployStatusEnum.SUCCESS, + error: undefined, + }); + return activeHash; + } + logger.warn( + `Deployment ${activeHash} is not ready in AI service: ${status}`, + ); + } catch (err: any) { + logger.warn( + `Deployment ${activeHash} is not available in AI service: ${err.message}`, + ); + } + } + + const result = await this.deploy(activeManifest, projectId); + if (result.status !== DeployStatusEnum.SUCCESS) { + throw new Error( + result.error || + `Failed to prepare deployment ${lastDeploy.hash} for project ${projectId}`, + ); + } + + return result.hash || activeHash; + } + public async getInProgressDeployment(projectId) { - return await this.deployLogRepository.findInProgressProjectDeployLog( + const inProgressDeploy = await this.deployLogRepository.findInProgressProjectDeployLog( projectId, ); + if (!inProgressDeploy) { + return null; + } + + const updatedAt = inProgressDeploy.updatedAt || inProgressDeploy.createdAt; + const updatedAtTime = updatedAt ? new Date(updatedAt).getTime() : 0; + if (updatedAtTime > 0 && Date.now() - updatedAtTime > STALE_DEPLOYMENT_MS) { + await this.markDeploymentFailed( + inProgressDeploy, + 'Deployment timed out before completion.', + ); + return null; + } + + return inProgressDeploy; } public async deploy(manifest, projectId, force = false) { const eventName = TelemetryEvent.MODELING_DEPLOY_MDL; + let deploy: Deploy | null = null; try { + manifest = normalizeManifest(manifest); // generate hash of manifest const hash = this.createMDLHash(manifest, projectId); logger.debug(`Deploying model, hash: ${hash}`); @@ -85,23 +183,48 @@ export class DeployService implements IDeployService { const lastDeploy = await this.deployLogRepository.findLastProjectDeployLog(projectId); if (lastDeploy && lastDeploy.hash === hash) { - logger.log(`Model has been deployed, hash: ${hash}`); - return { status: DeployStatusEnum.SUCCESS }; + logger.log(`Model has been deployed, refreshing AI index, hash: ${hash}`); + deploy = lastDeploy; + const { status: aiStatus, error: aiError } = + await this.wrenAIAdaptor.deploy({ + manifest, + hash, + projectId, + }); + const status = + aiStatus === WrenAIDeployStatusEnum.SUCCESS + ? DeployStatusEnum.SUCCESS + : DeployStatusEnum.FAILED; + await this.deployLogRepository.updateOne(lastDeploy.id, { + status, + error: aiError, + }); + return { status, error: aiError, hash }; } } + const previousInProgressDeploy = + await this.deployLogRepository.findInProgressProjectDeployLog(projectId); + if (previousInProgressDeploy) { + await this.markDeploymentFailed( + previousInProgressDeploy, + 'Deployment was superseded by a new deployment.', + ); + } + const deployData = { manifest, hash, projectId, status: DeployStatusEnum.IN_PROGRESS, } as Deploy; - const deploy = await this.deployLogRepository.createOne(deployData); + deploy = await this.deployLogRepository.createOne(deployData); // deploy to AI-service const { status: aiStatus, error: aiError } = await this.wrenAIAdaptor.deploy({ manifest, hash, + projectId, }); // update deploy status @@ -125,9 +248,16 @@ export class DeployService implements IDeployService { false, ); } - return { status, error: aiError }; + return { status, error: aiError, hash }; } catch (err: any) { logger.error(`Error deploying model: ${err.message}`); + if (deploy?.id) { + try { + await this.markDeploymentFailed(deploy, err.message); + } catch (updateErr: any) { + logger.error(`Error marking deployment failed: ${updateErr.message}`); + } + } this.telemetry.sendEvent( eventName, { mdl: manifest, error: err.message }, @@ -138,20 +268,76 @@ export class DeployService implements IDeployService { } } + private async markDeploymentFailed(deploy: Deploy, error: string) { + await this.deployLogRepository.updateOne(deploy.id, { + status: DeployStatusEnum.FAILED, + error, + }); + } + public createMDLHash(manifest: Manifest, projectId: number) { - const manifestStr = JSON.stringify(manifest); + const manifestStr = this.canonicalStringify(normalizeManifest(manifest)); const content = `${projectId} ${manifestStr}`; const hash = createHash('sha1').update(content).digest('hex'); return hash; } + public isSameDeployment( + manifest: Manifest, + projectId: number, + deployment?: Deploy | null, + ) { + if (!deployment) { + return false; + } + + if (deployment.hash === this.createMDLHash(manifest, projectId)) { + return true; + } + + return ( + this.canonicalStringify(normalizeManifest(deployment.manifest)) === + this.canonicalStringify(normalizeManifest(manifest)) + ); + } + + private canonicalStringify(value: any): string { + if (Array.isArray(value)) { + const serializedItems = value.map((item) => this.canonicalStringify(item)); + if (value.every((item) => item && typeof item === 'object')) { + serializedItems.sort(); + } + return `[${serializedItems.join(',')}]`; + } + + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${this.canonicalStringify(value[key])}`) + .join(',')}}`; + } + + return JSON.stringify(value); + } + public async getMDLByHash(hash: string) { const deploy = await this.deployLogRepository.findOneBy({ hash }); if (!deploy) { return null; } // return base64 encoded manifest - return Buffer.from(JSON.stringify(deploy.manifest)).toString('base64'); + return Buffer.from(JSON.stringify(normalizeManifest(deploy.manifest))).toString( + 'base64', + ); + } + + public async getManifestByHash(hash: string) { + const deploy = await this.deployLogRepository.findOneBy({ hash }); + if (!deploy?.manifest) { + return null; + } + + return normalizeManifest(deploy.manifest); } public async deleteAllByProjectId(projectId: number): Promise { diff --git a/wren-ui/src/apollo/server/services/index.ts b/wren-ui/src/apollo/server/services/index.ts index ce1bb78175..8d341f923e 100644 --- a/wren-ui/src/apollo/server/services/index.ts +++ b/wren-ui/src/apollo/server/services/index.ts @@ -8,3 +8,6 @@ export * from './metadataService'; export * from './dashboardService'; export * from './askingTaskTracker'; export * from './instructionService'; +export * from './rbacService'; +export * from './organizationService'; +export * from './organizationMemberService'; diff --git a/wren-ui/src/apollo/server/services/instructionService.ts b/wren-ui/src/apollo/server/services/instructionService.ts index 37112ad5f4..e8ff322ca5 100644 --- a/wren-ui/src/apollo/server/services/instructionService.ts +++ b/wren-ui/src/apollo/server/services/instructionService.ts @@ -9,6 +9,7 @@ import { import { IInstructionRepository, Instruction } from '@server/repositories'; import * as Errors from '@server/utils/error'; import { GeneralErrorCodes } from '@server/utils/error'; +import { IDeployService } from './deployService'; export interface IInstructionService { getInstructions(projectId: number): Promise; getInstruction(id: number): Promise; @@ -21,15 +22,19 @@ export interface IInstructionService { export class InstructionService implements IInstructionService { private readonly instructionRepository: IInstructionRepository; private readonly wrenAIAdaptor: IWrenAIAdaptor; + private readonly deployService: IDeployService; constructor({ instructionRepository, wrenAIAdaptor, + deployService, }: { instructionRepository: IInstructionRepository; wrenAIAdaptor: IWrenAIAdaptor; + deployService: IDeployService; }) { this.instructionRepository = instructionRepository; this.wrenAIAdaptor = wrenAIAdaptor; + this.deployService = deployService; } public async getInstructions(projectId: number): Promise { @@ -56,9 +61,11 @@ export class InstructionService implements IInstructionService { tx, }, ); - const { queryId } = await this.wrenAIAdaptor.generateInstruction([ - this.pickGenerateInstructionInput(newInstruction), - ]); + const mdlHash = await this.getDeployHash(input.projectId); + const { queryId } = await this.wrenAIAdaptor.generateInstruction( + [this.pickGenerateInstructionInput(newInstruction)], + mdlHash, + ); const res = await this.waitDeployInstruction(queryId); if (res.error) { await tx.rollback(); @@ -90,8 +97,10 @@ export class InstructionService implements IInstructionService { tx, }, ); + const mdlHash = await this.getDeployHash(inputs[0]?.projectId); const { queryId } = await this.wrenAIAdaptor.generateInstruction( newInstructions.map(this.pickGenerateInstructionInput), + mdlHash, ); const res = await this.waitDeployInstruction(queryId); if (res.error) { @@ -133,9 +142,11 @@ export class InstructionService implements IInstructionService { instructionData, { tx }, ); - const { queryId } = await this.wrenAIAdaptor.generateInstruction([ - this.pickGenerateInstructionInput(updatedInstruction), - ]); + const mdlHash = await this.getDeployHash(input.projectId); + const { queryId } = await this.wrenAIAdaptor.generateInstruction( + [this.pickGenerateInstructionInput(updatedInstruction)], + mdlHash, + ); const res = await this.waitDeployInstruction(queryId); if (res.error) { await tx.rollback(); @@ -215,4 +226,12 @@ export class InstructionService implements IInstructionService { throw new Error('Instruction is too long'); } } + + private async getDeployHash(projectId?: number): Promise { + if (!projectId) { + return undefined; + } + const deployment = await this.deployService.getLastDeployment(projectId); + return deployment?.hash; + } } diff --git a/wren-ui/src/apollo/server/services/mdlService.ts b/wren-ui/src/apollo/server/services/mdlService.ts index c60199aad0..5d65604a5b 100644 --- a/wren-ui/src/apollo/server/services/mdlService.ts +++ b/wren-ui/src/apollo/server/services/mdlService.ts @@ -3,6 +3,7 @@ import { IModelNestedColumnRepository, IModelColumnRepository, IModelRepository, + Project, IProjectRepository, IRelationRepository, IViewRepository, @@ -15,6 +16,7 @@ export interface MakeCurrentModelMDLResult { } export interface IMDLService { makeCurrentModelMDL(): Promise; + makeModelMDL(project: Project): Promise; } export class MDLService implements IMDLService { @@ -50,6 +52,10 @@ export class MDLService implements IMDLService { public async makeCurrentModelMDL() { const project = await this.projectRepository.getCurrentProject(); + return this.makeModelMDL(project); + } + + public async makeModelMDL(project: Project) { const projectId = project.id; const models = await this.modelRepository.findAllBy({ projectId }); const modelIds = models.map((m) => m.id); diff --git a/wren-ui/src/apollo/server/services/modelService.ts b/wren-ui/src/apollo/server/services/modelService.ts index 5299ade75c..e98581c3b7 100644 --- a/wren-ui/src/apollo/server/services/modelService.ts +++ b/wren-ui/src/apollo/server/services/modelService.ts @@ -260,7 +260,7 @@ export class ModelService implements IModelService { projectId: id, }); - await Promise.all([ + await Promise.all( tables.map(async (table) => { const model = models.find((m) => m.sourceTableName === table.tableName); if (!model) { @@ -275,7 +275,7 @@ export class ModelService implements IModelService { properties: JSON.stringify(properties), }); }), - ]); + ); } public async batchUpdateColumnProperties(tables: SampleDatasetTable[]) { @@ -298,7 +298,7 @@ export class ModelService implements IModelService { return acc; }, []); - await Promise.all([ + await Promise.all( transformedColumns.map(async (column) => { if (!column.properties) { return; @@ -327,7 +327,7 @@ export class ModelService implements IModelService { properties: JSON.stringify(properties), }); }), - ]); + ); } public generateReferenceName(data: GenerateReferenceNameData): string { @@ -353,32 +353,44 @@ export class ModelService implements IModelService { .flat(); const columns = await this.modelColumnRepository.findColumnsByIds(columnIds); - const relationValues = relations.map((relation) => { + const relationValues = relations.flatMap((relation) => { const fromColumn = columns.find( (column) => column.id === relation.fromColumnId, ); if (!fromColumn) { - throw new Error(`Column not found, column Id ${relation.fromColumnId}`); + logger.warn( + `Skip relation because column ${relation.fromColumnId} was not found`, + ); + return []; } const toColumn = columns.find( (column) => column.id === relation.toColumnId, ); if (!toColumn) { - throw new Error(`Column not found, column Id ${relation.toColumnId}`); + logger.warn( + `Skip relation because column ${relation.toColumnId} was not found`, + ); + return []; } const relationName = this.generateRelationName(relation, models, columns); - return { - projectId: id, - name: relationName, - fromColumnId: relation.fromColumnId, - toColumnId: relation.toColumnId, - joinType: relation.type, - properties: relation.description - ? JSON.stringify({ description: relation.description }) - : null, - } as Partial; + return [ + { + projectId: id, + name: relationName, + fromColumnId: relation.fromColumnId, + toColumnId: relation.toColumnId, + joinType: relation.type, + properties: relation.description + ? JSON.stringify({ description: relation.description }) + : null, + } as Partial, + ]; }); + if (isEmpty(relationValues)) { + return []; + } + const savedRelations = await this.relationRepository.createMany(relationValues); @@ -387,31 +399,113 @@ export class ModelService implements IModelService { public async createRelation(relation: RelationData): Promise { const { id } = await this.projectService.getCurrentProject(); - const modelIds = [relation.fromModelId, relation.toModelId]; - const models = await this.modelRepository.findAllByIds(modelIds); - const columnIds = [relation.fromColumnId, relation.toColumnId]; + const models = await this.modelRepository.findAllBy({ projectId: id }); + const modelIds = models.map((model) => model.id); + const projectColumns = + await this.modelColumnRepository.findColumnsByModelIds(modelIds); + const resolvedRelation = this.resolveRelationEndpoints( + relation, + models, + projectColumns, + ); + const columnIds = [ + resolvedRelation.fromColumnId, + resolvedRelation.toColumnId, + ]; const columns = await this.modelColumnRepository.findColumnsByIds(columnIds); const { valid, message } = await this.validateCreateRelation( models, columns, - relation, + resolvedRelation, ); if (!valid) { throw new Error(message); } - const relationName = this.generateRelationName(relation, models, columns); + const relationName = this.generateRelationName( + resolvedRelation, + models, + columns, + ); const savedRelation = await this.relationRepository.createOne({ projectId: id, name: relationName, - fromColumnId: relation.fromColumnId, - toColumnId: relation.toColumnId, - joinType: relation.type, + fromColumnId: resolvedRelation.fromColumnId, + toColumnId: resolvedRelation.toColumnId, + joinType: resolvedRelation.type, + properties: resolvedRelation.description + ? JSON.stringify({ description: resolvedRelation.description }) + : null, }); return savedRelation; } + private resolveRelationEndpoints( + relation: RelationData, + models: Model[], + columns: ModelColumn[], + ): RelationData { + const resolveModel = ( + id: number, + referenceName?: string, + ): Model | undefined => + models.find((model) => model.id === id) || + (referenceName + ? models.find((model) => model.referenceName === referenceName) + : undefined); + + let fromModel = resolveModel( + relation.fromModelId, + relation.fromModelReferenceName, + ); + let toModel = resolveModel( + relation.toModelId, + relation.toModelReferenceName, + ); + + const resolveColumn = ( + id: number, + modelId: number | undefined, + referenceName?: string, + ): ModelColumn | undefined => + columns.find( + (column) => column.id === id && (!modelId || column.modelId === modelId), + ) || + (referenceName && modelId + ? columns.find( + (column) => + column.modelId === modelId && column.referenceName === referenceName, + ) + : undefined); + + const fromColumn = resolveColumn( + relation.fromColumnId, + fromModel?.id, + relation.fromColumnReferenceName, + ); + const toColumn = resolveColumn( + relation.toColumnId, + toModel?.id, + relation.toColumnReferenceName, + ); + + if (!fromModel && fromColumn) { + fromModel = models.find((model) => model.id === fromColumn.modelId); + } + if (!toModel && toColumn) { + toModel = models.find((model) => model.id === toColumn.modelId); + } + + return { + ...relation, + fromModelId: fromModel?.id ?? relation.fromModelId, + fromColumnId: fromColumn?.id ?? relation.fromColumnId, + toModelId: toModel?.id ?? relation.toModelId, + toColumnId: toColumn?.id ?? relation.toColumnId, + }; + } + public async updateRelation( relation: UpdateRelationData, id: number, diff --git a/wren-ui/src/apollo/server/services/organizationMemberService.ts b/wren-ui/src/apollo/server/services/organizationMemberService.ts new file mode 100644 index 0000000000..cc1524844d --- /dev/null +++ b/wren-ui/src/apollo/server/services/organizationMemberService.ts @@ -0,0 +1,1100 @@ +import { ApiError } from '../utils/apiUtils'; +import { v4 as uuidv4 } from 'uuid'; +import { + IOrganizationMemberProjectRepository, + IOrganizationMemberRepository, + OrganizationMemberMapping, +} from '../repositories/organizationMemberRepository'; +import { + IOrganizationInvitationProjectRepository, + IOrganizationInvitationRepository, + OrganizationInvitation, +} from '../repositories/organizationInvitationRepository'; +import { IOrganizationRepository } from '../repositories/organizationRepository'; +import { IProjectRepository } from '../repositories/projectRepository'; +import { IUserRepository, RbacUser } from '../repositories/rbacRepository'; + +export const ORGANIZATION_MEMBER_ROLES = ['Admin', 'Member'] as const; +export const PROJECT_PERMISSION_ROLES = ['Owner', 'Editor', 'Viewer'] as const; +export const ORGANIZATION_INVITATION_STATUSES = [ + 'Pending', + 'Accepted', + 'Expired', +] as const; + +export type OrganizationMemberRole = + (typeof ORGANIZATION_MEMBER_ROLES)[number]; +export type ProjectPermissionRole = (typeof PROJECT_PERMISSION_ROLES)[number]; +export type OrganizationInvitationStatus = + (typeof ORGANIZATION_INVITATION_STATUSES)[number]; + +export interface MemberProjectInput { + projectId: number; + permission: ProjectPermissionRole; +} + +export interface InviteOrganizationMemberInput { + email: string; + organizationRole: OrganizationMemberRole; + projects: MemberProjectInput[]; +} + +export interface UpdateOrganizationMemberInput { + organizationRole: OrganizationMemberRole; +} + +export interface UpdateCurrentUserProfileInput { + name: string; +} + +export interface CurrentUserProfile { + id: number; + name: string; + email: string; + isActive: boolean; +} + +export interface OrganizationMemberSummary { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationMemberRole; + projects: Array<{ + projectId: number; + displayName: string; + permission: ProjectPermissionRole; + }>; +} + +export interface ProjectAccessMemberSummary { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationMemberRole; + permission: 'Owner' | 'Contributor' | 'Viewer'; + isCurrentUser: boolean; + canEditPermission: boolean; + canRemove: boolean; +} + +export interface ProjectAccessAvailableMember { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationMemberRole; +} + +export interface AddProjectMemberInput { + organizationMemberId: number; + permission: 'Owner' | 'Contributor' | 'Viewer'; +} + +export interface UpdateProjectMemberPermissionInput { + permission: 'Owner' | 'Contributor' | 'Viewer'; +} + +export interface OrganizationInvitationSummary { + id: number; + email: string; + organizationRole: OrganizationMemberRole; + status: OrganizationInvitationStatus; + token: string; + expiresAt: string; + acceptedAt?: string | null; + createdAt: string; + projects: Array<{ + projectId: number; + displayName: string; + permission: ProjectPermissionRole; + }>; +} + +export interface IOrganizationMemberService { + listCurrentOrganizationMembers(): Promise<{ + members: OrganizationMemberSummary[]; + invitations: OrganizationInvitationSummary[]; + projects: Array<{ id: number; displayName: string }>; + currentUserId: number | null; + }>; + inviteMember( + input: InviteOrganizationMemberInput, + ): Promise; + updateMember( + id: number, + input: UpdateOrganizationMemberInput, + ): Promise; + removeMember(id: number): Promise; + removeInvitation(id: number): Promise; + acceptInvitation(token: string): Promise; + leaveCurrentOrganization(): Promise; + deleteCurrentOrganization(): Promise; + getCurrentUserProfile(): Promise; + updateCurrentUserProfile( + input: UpdateCurrentUserProfileInput, + ): Promise; + deleteCurrentUserAccount(): Promise; + listCurrentProjectAccess(): Promise<{ + members: ProjectAccessMemberSummary[]; + availableMembers: ProjectAccessAvailableMember[]; + currentUserId: number | null; + canManageAccess: boolean; + }>; + addProjectMember( + input: AddProjectMemberInput, + ): Promise; + updateProjectMemberPermission( + id: number, + input: UpdateProjectMemberPermissionInput, + ): Promise; + removeProjectMember(id: number): Promise; +} + +const EMAIL_PATTERN = /^[^\s@]+@[^\s@]+\.[^\s@]+$/; +const INVITATION_TTL_DAYS = 3; + +export class OrganizationMemberService implements IOrganizationMemberService { + constructor( + private readonly organizationRepository: IOrganizationRepository, + private readonly organizationMemberRepository: IOrganizationMemberRepository, + private readonly organizationMemberProjectRepository: IOrganizationMemberProjectRepository, + private readonly organizationInvitationRepository: IOrganizationInvitationRepository, + private readonly organizationInvitationProjectRepository: IOrganizationInvitationProjectRepository, + private readonly userRepository: IUserRepository, + private readonly projectRepository: IProjectRepository, + ) {} + + public async listCurrentOrganizationMembers() { + const organization = await this.getCurrentOrganizationOrThrow(); + const [members, listedProjects] = await Promise.all([ + this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ), + this.projectRepository.findAll({ order: 'id' }), + ]); + const projects = + listedProjects.length > 0 + ? listedProjects + : await this.getFallbackProjects(); + const invitations = await this.organizationInvitationRepository.findByOrganizationId( + organization.id, + ); + const normalizedInvitations = await Promise.all( + invitations.map((invite) => this.normalizeInvitation(invite)), + ); + + const hydratedMembers = await Promise.all( + members.map((member) => this.serializeMember(member)), + ); + const currentUserId = + hydratedMembers.find((member) => member.organizationRole === 'Admin') + ?.userId ?? + hydratedMembers[0]?.userId ?? + null; + + return { + members: hydratedMembers, + invitations: await Promise.all( + normalizedInvitations + .filter((invite) => invite.status !== 'Accepted') + .map((invite) => this.serializeInvitation(invite)), + ), + projects: projects.map((project) => ({ + id: project.id, + displayName: project.displayName, + })), + currentUserId, + }; + } + + public async inviteMember(input: InviteOrganizationMemberInput) { + const organization = await this.getCurrentOrganizationOrThrow(); + const payload = await this.validateInvitePayload(input); + const existingMembership = + await this.findExistingMemberByEmail(organization.id, payload.email); + if (existingMembership) { + throw new ApiError('Member already exists in this organization', 409); + } + const existingInvitation = + await this.findActiveInvitationByEmail(organization.id, payload.email); + if (existingInvitation) { + throw new ApiError( + 'A pending invitation already exists for this email', + 409, + ); + } + + const currentUserId = await this.getCurrentUserId(organization.id); + + const tx = await this.organizationMemberRepository.transaction(); + try { + const now = new Date().toISOString(); + const invitation = await this.organizationInvitationRepository.createOne( + { + organizationId: organization.id, + invitedByUserId: currentUserId, + email: payload.email, + organizationRole: payload.organizationRole, + token: uuidv4(), + status: 'Pending', + expiresAt: this.buildInvitationExpiry(now), + acceptedAt: null, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + + if (payload.projects.length) { + await this.organizationInvitationProjectRepository.createMany( + payload.projects.map((project) => ({ + organizationInvitationId: invitation.id, + projectId: project.projectId, + permission: project.permission, + createdAt: now, + updatedAt: now, + })), + { tx }, + ); + } + + await tx.commit(); + const createdInvitation = + await this.organizationInvitationRepository.findOneBy({ + id: invitation.id, + }); + return this.serializeInvitation(createdInvitation); + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async updateMember( + id: number, + input: UpdateOrganizationMemberInput, + ): Promise { + const member = await this.organizationMemberRepository.findOneBy({ id }); + if (!member) { + throw new ApiError('Member not found', 404); + } + + const organization = await this.getCurrentOrganizationOrThrow(); + if (member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + await this.assertNotCurrentUser(member.userId, organization.id); + + const role = this.validateOrganizationRole(input.organizationRole); + await this.organizationMemberRepository.updateOne(member.id, { + organizationRole: role, + updatedAt: new Date().toISOString(), + }); + + const mapping = await this.getMemberMappingOrThrow(member.id); + return this.serializeMember(mapping); + } + + public async removeMember(id: number): Promise { + const member = await this.organizationMemberRepository.findOneBy({ id }); + if (!member) return true; + + const organization = await this.getCurrentOrganizationOrThrow(); + if (member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + await this.assertNotCurrentUser(member.userId, organization.id); + + await this.organizationMemberRepository.deleteOne(member.id); + return true; + } + + public async removeInvitation(id: number): Promise { + const invitation = await this.organizationInvitationRepository.findOneBy({ + id, + }); + if (!invitation) return true; + + const organization = await this.getCurrentOrganizationOrThrow(); + if (invitation.organizationId !== organization.id) { + throw new ApiError('Invitation not found', 404); + } + + await this.organizationInvitationRepository.deleteOne(invitation.id); + return true; + } + + public async leaveCurrentOrganization(): Promise { + const organization = await this.getCurrentOrganizationOrThrow(); + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + const currentUserId = await this.getCurrentUserId(organization.id); + const currentMember = members.find( + (member) => member.userId === currentUserId, + ); + + if (!currentMember) { + throw new ApiError('Current organization member not found', 404); + } + + if (currentMember.organizationRole === 'Admin') { + const adminCount = members.filter( + (member) => member.organizationRole === 'Admin', + ).length; + if (adminCount <= 1) { + throw new ApiError( + 'If you are the last Organization admin, you cannot leave the organization. You will need to delete the organization to remove yourself from it.', + 400, + ); + } + } + + await this.organizationMemberRepository.deleteOne(currentMember.id); + return true; + } + + public async deleteCurrentOrganization(): Promise { + const organization = await this.getCurrentOrganizationOrThrow(); + await this.assertCurrentUserAdmin(organization.id); + + const tx = await this.organizationRepository.transaction(); + try { + await this.organizationRepository.deleteOne(organization.id, { tx }); + const remainingOrganizations = await this.organizationRepository.findAll({ + order: 'id', + tx, + }); + if (remainingOrganizations.length > 0) { + await this.organizationRepository.setCurrentOrganization( + remainingOrganizations[0].id, + { tx }, + ); + } + await tx.commit(); + return true; + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async getCurrentUserProfile(): Promise { + const organization = await this.getCurrentOrganizationOrThrow(); + const currentUserId = await this.getCurrentUserId(organization.id); + if (!currentUserId) { + throw new ApiError('Current user not found', 404); + } + const user = await this.userRepository.findOneBy({ id: currentUserId }); + if (!user) { + throw new ApiError('Current user not found', 404); + } + return this.serializeCurrentUserProfile(user); + } + + public async updateCurrentUserProfile( + input: UpdateCurrentUserProfileInput, + ): Promise { + const organization = await this.getCurrentOrganizationOrThrow(); + const currentUserId = await this.getCurrentUserId(organization.id); + if (!currentUserId) { + throw new ApiError('Current user not found', 404); + } + + const name = `${input.name || ''}`.trim(); + if (!name) { + throw new ApiError('Name is required', 400); + } + if (name.length > 160) { + throw new ApiError('Name must be 160 characters or fewer', 400); + } + + const user = await this.userRepository.updateOne(currentUserId, { + name, + updatedAt: new Date().toISOString(), + }); + return this.serializeCurrentUserProfile(user); + } + + public async deleteCurrentUserAccount(): Promise { + const organization = await this.getCurrentOrganizationOrThrow(); + const currentUserId = await this.getCurrentUserId(organization.id); + if (!currentUserId) { + throw new ApiError('Current user not found', 404); + } + + await this.assertCurrentUserCanDeleteAccount(currentUserId); + await this.userRepository.deleteOne(currentUserId); + return true; + } + + public async listCurrentProjectAccess() { + const organization = await this.getCurrentOrganizationOrThrow(); + const project = await this.projectRepository.getCurrentProject(); + const currentUserId = await this.getCurrentUserId(organization.id); + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + + const serializedMembers = await Promise.all( + members.map(async (member) => { + const projects = + await this.organizationMemberProjectRepository.findByOrganizationMemberId( + member.id, + ); + const assignment = projects.find( + (mappedProject) => mappedProject.projectId === project.id, + ); + const isOrganizationAdmin = member.organizationRole === 'Admin'; + const isCurrentUser = member.userId === currentUserId; + const permission = isOrganizationAdmin + ? 'Owner' + : assignment + ? this.presentProjectPermission(assignment.permission) + : null; + + return permission + ? { + id: member.id, + userId: member.userId, + name: member.user.name, + email: member.user.email, + organizationRole: + member.organizationRole as OrganizationMemberRole, + permission, + isCurrentUser, + canEditPermission: !isOrganizationAdmin && !isCurrentUser, + canRemove: !isOrganizationAdmin && !isCurrentUser, + } + : null; + }), + ); + + const availableMembers = members + .filter( + (member) => + member.organizationRole !== 'Admin' && + !serializedMembers.some( + (serializedMember) => serializedMember?.id === member.id, + ), + ) + .map((member) => ({ + id: member.id, + userId: member.userId, + name: member.user.name, + email: member.user.email, + organizationRole: member.organizationRole as OrganizationMemberRole, + })); + + const currentUserMember = serializedMembers.find( + (member) => member?.userId === currentUserId, + ); + return { + members: serializedMembers.filter(Boolean), + availableMembers, + currentUserId, + canManageAccess: + currentUserMember?.permission === 'Owner' || + currentUserMember?.organizationRole === 'Admin', + }; + } + + public async addProjectMember(input: AddProjectMemberInput) { + const organization = await this.getCurrentOrganizationOrThrow(); + const project = await this.projectRepository.getCurrentProject(); + await this.assertCanManageCurrentProject(organization.id, project.id); + + const member = await this.organizationMemberRepository.findOneBy({ + id: input.organizationMemberId, + }); + if (!member || member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + if (member.organizationRole === 'Admin') { + throw new ApiError( + 'Organization admins are owners of all projects by default', + 400, + ); + } + + const existingAssignment = await this.findProjectAssignment(member.id, project.id); + if (existingAssignment) { + throw new ApiError('Member already has access to this project', 409); + } + + const now = new Date().toISOString(); + await this.organizationMemberProjectRepository.createOne({ + organizationMemberId: member.id, + projectId: project.id, + permission: this.storeProjectPermission(input.permission), + createdAt: now, + updatedAt: now, + }); + + return this.serializeProjectAccessMember(member.id, project.id, organization.id); + } + + public async updateProjectMemberPermission( + id: number, + input: UpdateProjectMemberPermissionInput, + ) { + const organization = await this.getCurrentOrganizationOrThrow(); + const project = await this.projectRepository.getCurrentProject(); + const currentUserId = await this.getCurrentUserId(organization.id); + await this.assertCanManageCurrentProject(organization.id, project.id); + + const member = await this.organizationMemberRepository.findOneBy({ id }); + if (!member || member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + if (member.userId === currentUserId) { + throw new ApiError('You cannot modify your own role', 400); + } + if (member.organizationRole === 'Admin') { + throw new ApiError( + 'You cannot change the role of the organization admin', + 400, + ); + } + + const assignment = await this.findProjectAssignment(member.id, project.id); + if (!assignment) { + throw new ApiError('Project member not found', 404); + } + + await this.organizationMemberProjectRepository.updateOne(assignment.id, { + permission: this.storeProjectPermission(input.permission), + updatedAt: new Date().toISOString(), + }); + + return this.serializeProjectAccessMember(member.id, project.id, organization.id); + } + + public async removeProjectMember(id: number) { + const organization = await this.getCurrentOrganizationOrThrow(); + const project = await this.projectRepository.getCurrentProject(); + const currentUserId = await this.getCurrentUserId(organization.id); + await this.assertCanManageCurrentProject(organization.id, project.id); + + const member = await this.organizationMemberRepository.findOneBy({ id }); + if (!member || member.organizationId !== organization.id) { + throw new ApiError('Member not found', 404); + } + if (member.userId === currentUserId) { + throw new ApiError('You cannot remove yourself from the project', 400); + } + if (member.organizationRole === 'Admin') { + throw new ApiError( + 'You cannot remove the organization admin from the project', + 400, + ); + } + + const assignment = await this.findProjectAssignment(member.id, project.id); + if (!assignment) { + throw new ApiError('Project member not found', 404); + } + await this.organizationMemberProjectRepository.deleteOne(assignment.id); + return true; + } + + public async acceptInvitation( + token: string, + ): Promise { + const invitation = await this.organizationInvitationRepository.findOneBy({ + token, + }); + if (!invitation) { + throw new ApiError('Invitation not found', 404); + } + + const normalizedInvitation = await this.normalizeInvitation(invitation); + if (normalizedInvitation.status === 'Expired') { + throw new ApiError('Invitation has expired', 410); + } + + const existingMembership = + await this.findExistingMemberByEmail( + normalizedInvitation.organizationId, + normalizedInvitation.email, + ); + if (existingMembership) { + await this.organizationInvitationRepository.updateOne( + normalizedInvitation.id, + { + status: 'Accepted', + acceptedAt: new Date().toISOString(), + updatedAt: new Date().toISOString(), + }, + ); + return existingMembership; + } + + const tx = await this.organizationInvitationRepository.transaction(); + try { + const now = new Date().toISOString(); + const existingUser = await this.userRepository.findOneBy( + { email: normalizedInvitation.email }, + { tx }, + ); + const user = existingUser + ? existingUser + : await this.createUserFromInvite(normalizedInvitation.email, tx); + const member = await this.organizationMemberRepository.createOne( + { + organizationId: normalizedInvitation.organizationId, + userId: user.id, + organizationRole: normalizedInvitation.organizationRole, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + const invitationProjects = + await this.organizationInvitationProjectRepository.findByOrganizationInvitationId( + normalizedInvitation.id, + { tx }, + ); + if (invitationProjects.length) { + await this.organizationMemberProjectRepository.createMany( + invitationProjects.map((project) => ({ + organizationMemberId: member.id, + projectId: project.projectId, + permission: project.permission, + createdAt: now, + updatedAt: now, + })), + { tx }, + ); + } + await this.organizationInvitationRepository.updateOne( + normalizedInvitation.id, + { + status: 'Accepted', + acceptedAt: now, + updatedAt: now, + }, + { tx }, + ); + await tx.commit(); + const mapping = await this.getMemberMappingOrThrow(member.id); + return this.serializeMember(mapping); + } catch (error) { + await tx.rollback(); + throw error; + } + } + + private async getCurrentOrganizationOrThrow() { + const organization = + await this.organizationRepository.getCurrentOrganization(); + if (!organization) { + throw new ApiError('Current organization not found', 404); + } + return organization; + } + + private async createUserFromInvite( + email: string, + tx?: any, + ): Promise { + const now = new Date().toISOString(); + const localPart = email.split('@')[0] || 'user'; + return this.userRepository.createOne( + { + name: localPart, + email, + isActive: true, + createdAt: now, + updatedAt: now, + }, + tx ? { tx } : undefined, + ); + } + + private async validateInvitePayload(input: InviteOrganizationMemberInput) { + const email = `${input.email || ''}`.trim().toLowerCase(); + if (!EMAIL_PATTERN.test(email)) { + throw new ApiError('A valid email address is required', 400); + } + + const organizationRole = this.validateOrganizationRole( + input.organizationRole, + ); + const allProjects = await this.projectRepository.findAll({ order: 'id' }); + const projectMap = new Map(allProjects.map((project) => [project.id, project])); + + let projects = (input.projects || []).map((project) => ({ + projectId: Number(project.projectId), + permission: this.validateProjectPermission(project.permission), + })); + + if (organizationRole === 'Admin') { + projects = allProjects.map((project) => ({ + projectId: project.id, + permission: 'Owner' as ProjectPermissionRole, + })); + } else if (!projects.length) { + throw new ApiError( + 'Select at least one project for organization members', + 400, + ); + } + + for (const project of projects) { + if (!projectMap.has(project.projectId)) { + throw new ApiError(`Project ${project.projectId} not found`, 400); + } + } + + const dedupedProjects = Array.from( + new Map(projects.map((project) => [project.projectId, project])).values(), + ); + + return { + email, + organizationRole, + projects: dedupedProjects, + }; + } + + private validateOrganizationRole(role: string): OrganizationMemberRole { + const normalized = `${role || ''}`.trim(); + if ( + !ORGANIZATION_MEMBER_ROLES.includes( + normalized as OrganizationMemberRole, + ) + ) { + throw new ApiError('Invalid organization role', 400); + } + return normalized as OrganizationMemberRole; + } + + private validateProjectPermission(role: string): ProjectPermissionRole { + const normalized = `${role || ''}`.trim(); + if ( + !PROJECT_PERMISSION_ROLES.includes(normalized as ProjectPermissionRole) + ) { + throw new ApiError('Invalid project permission', 400); + } + return normalized as ProjectPermissionRole; + } + + private async getMemberMappingOrThrow(id: number) { + const organization = await this.getCurrentOrganizationOrThrow(); + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + const member = members.find((item) => item.id === id); + if (!member) { + throw new ApiError('Member not found', 404); + } + return member; + } + + private async serializeMember( + member: OrganizationMemberMapping, + ): Promise { + const projects = + await this.organizationMemberProjectRepository.findByOrganizationMemberId( + member.id, + ); + return { + id: member.id, + userId: member.userId, + name: member.user.name, + email: member.user.email, + organizationRole: member.organizationRole as OrganizationMemberRole, + projects: projects.map((project) => ({ + projectId: project.projectId, + displayName: project.project.displayName, + permission: project.permission as ProjectPermissionRole, + })), + }; + } + + private async serializeInvitation( + invitation: OrganizationInvitation | null, + ): Promise { + if (!invitation) { + throw new ApiError('Invitation not found', 404); + } + const projects = + await this.organizationInvitationProjectRepository.findByOrganizationInvitationId( + invitation.id, + ); + return { + id: invitation.id, + email: invitation.email, + organizationRole: invitation.organizationRole as OrganizationMemberRole, + status: invitation.status as OrganizationInvitationStatus, + token: invitation.token, + expiresAt: invitation.expiresAt, + acceptedAt: invitation.acceptedAt, + createdAt: invitation.createdAt, + projects: projects.map((project) => ({ + projectId: project.projectId, + displayName: project.project.displayName, + permission: project.permission as ProjectPermissionRole, + })), + }; + } + + private buildInvitationExpiry(createdAt: string) { + const expiresAt = new Date(createdAt); + expiresAt.setDate(expiresAt.getDate() + INVITATION_TTL_DAYS); + return expiresAt.toISOString(); + } + + private async normalizeInvitation(invitation: OrganizationInvitation) { + if ( + invitation.status === 'Pending' && + new Date(invitation.expiresAt).getTime() < Date.now() + ) { + return await this.organizationInvitationRepository.updateOne( + invitation.id, + { + status: 'Expired', + updatedAt: new Date().toISOString(), + }, + ); + } + return invitation; + } + + private async findExistingMemberByEmail( + organizationId: number, + email: string, + ): Promise { + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organizationId, + ); + const member = members.find( + (item) => item.user.email.toLowerCase() === email.toLowerCase(), + ); + return member ? this.serializeMember(member) : null; + } + + private async findActiveInvitationByEmail( + organizationId: number, + email: string, + ) { + const invitations = + await this.organizationInvitationRepository.findByOrganizationId( + organizationId, + ); + for (const invitation of invitations) { + const normalized = await this.normalizeInvitation(invitation); + if ( + normalized.email.toLowerCase() === email.toLowerCase() && + normalized.status === 'Pending' + ) { + return normalized; + } + } + return null; + } + + private async getCurrentUserId( + organizationId: number, + ): Promise { + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organizationId, + ); + return ( + members.find((member) => member.organizationRole === 'Admin')?.userId ?? + members[0]?.userId ?? + null + ); + } + + private async assertNotCurrentUser(userId: number, organizationId: number) { + const currentUserId = await this.getCurrentUserId(organizationId); + if (currentUserId && currentUserId === userId) { + throw new ApiError('You cannot modify your own organization role', 400); + } + } + + private serializeCurrentUserProfile(user: RbacUser): CurrentUserProfile { + return { + id: user.id, + name: user.name, + email: user.email, + isActive: Boolean(user.isActive), + }; + } + + private async assertCurrentUserAdmin(organizationId: number) { + const currentUserId = await this.getCurrentUserId(organizationId); + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organizationId, + ); + const currentMember = members.find( + (member) => member.userId === currentUserId, + ); + if (!currentMember || currentMember.organizationRole !== 'Admin') { + throw new ApiError( + 'Only Organization admins can delete the organization', + 403, + ); + } + } + + private async assertCurrentUserCanDeleteAccount(userId: number) { + const organizations = await this.organizationRepository.findAll({ + order: 'id', + }); + + for (const organization of organizations) { + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organization.id, + ); + const currentMember = members.find((member) => member.userId === userId); + if (!currentMember || currentMember.organizationRole !== 'Admin') { + continue; + } + + const adminCount = members.filter( + (member) => member.organizationRole === 'Admin', + ).length; + if (adminCount <= 1) { + throw new ApiError( + 'If you are the last Organization admin, you cannot delete your account. Assign another Organization admin or delete the organization first.', + 400, + ); + } + } + } + + private async getFallbackProjects() { + try { + const currentProject = await this.projectRepository.getCurrentProject(); + return currentProject ? [currentProject] : []; + } catch { + return []; + } + } + + private async serializeProjectAccessMember( + organizationMemberId: number, + projectId: number, + organizationId: number, + ): Promise { + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organizationId, + ); + const member = members.find((item) => item.id === organizationMemberId); + if (!member) { + throw new ApiError('Member not found', 404); + } + const currentUserId = await this.getCurrentUserId(organizationId); + const assignment = + await this.organizationMemberProjectRepository.findByOrganizationMemberId( + organizationMemberId, + ); + const projectAssignment = assignment.find( + (item) => item.projectId === projectId, + ); + const isOrganizationAdmin = member.organizationRole === 'Admin'; + const permission = isOrganizationAdmin + ? 'Owner' + : projectAssignment + ? this.presentProjectPermission(projectAssignment.permission) + : null; + + if (!permission) { + throw new ApiError('Project member not found', 404); + } + + return { + id: member.id, + userId: member.userId, + name: member.user.name, + email: member.user.email, + organizationRole: member.organizationRole as OrganizationMemberRole, + permission, + isCurrentUser: member.userId === currentUserId, + canEditPermission: !isOrganizationAdmin && member.userId !== currentUserId, + canRemove: !isOrganizationAdmin && member.userId !== currentUserId, + }; + } + + private async assertCanManageCurrentProject( + organizationId: number, + projectId: number, + ) { + const currentUserId = await this.getCurrentUserId(organizationId); + const members = + await this.organizationMemberRepository.findMappingsByOrganizationId( + organizationId, + ); + const currentMember = members.find( + (member) => member.userId === currentUserId, + ); + if (!currentMember) { + throw new ApiError('Current user not found', 404); + } + if (currentMember.organizationRole === 'Admin') { + return; + } + + const assignments = + await this.organizationMemberProjectRepository.findByOrganizationMemberId( + currentMember.id, + ); + const currentProjectAssignment = assignments.find( + (assignment) => assignment.projectId === projectId, + ); + const currentPermission = currentProjectAssignment + ? this.presentProjectPermission(currentProjectAssignment.permission) + : null; + + if (currentPermission !== 'Owner') { + throw new ApiError( + 'Only project owners can manage project access', + 403, + ); + } + } + + private async findProjectAssignment( + organizationMemberId: number, + projectId: number, + ) { + const assignments = + await this.organizationMemberProjectRepository.findByOrganizationMemberId( + organizationMemberId, + ); + return assignments.find((assignment) => assignment.projectId === projectId); + } + + private presentProjectPermission(permission: string): + | 'Owner' + | 'Contributor' + | 'Viewer' { + if (permission === 'Editor') { + return 'Contributor'; + } + if (permission === 'Owner' || permission === 'Viewer') { + return permission; + } + return 'Contributor'; + } + + private storeProjectPermission(permission: string): ProjectPermissionRole { + if (permission === 'Contributor') { + return 'Editor'; + } + return this.validateProjectPermission(permission); + } +} diff --git a/wren-ui/src/apollo/server/services/organizationService.ts b/wren-ui/src/apollo/server/services/organizationService.ts new file mode 100644 index 0000000000..5c463b4e3e --- /dev/null +++ b/wren-ui/src/apollo/server/services/organizationService.ts @@ -0,0 +1,178 @@ +import { ApiError } from '../utils/apiUtils'; +import { + IOrganizationRepository, + Organization, +} from '../repositories/organizationRepository'; + +export interface CreateOrganizationData { + name: string; + identifier?: string; + description?: string; +} + +export interface UpdateOrganizationData { + name: string; +} + +export interface IOrganizationService { + listOrganizations: () => Promise; + getCurrentOrganization: () => Promise; + createOrganization: (data: CreateOrganizationData) => Promise; + selectCurrentOrganization: (id: number) => Promise; + updateCurrentOrganization: (data: UpdateOrganizationData) => Promise; +} + +const NAME_MAX_LENGTH = 64; +const IDENTIFIER_MAX_LENGTH = 64; +const DESCRIPTION_MAX_LENGTH = 255; +const IDENTIFIER_PATTERN = /^[a-z0-9]+(?:[_-][a-z0-9]+)*$/; + +export const normalizeOrganizationName = (value: string) => value.trim(); + +export const normalizeOrganizationIdentifier = (value: string) => + value + .trim() + .toLowerCase() + .replace(/[^a-z0-9]+/g, '_') + .replace(/_+/g, '_') + .replace(/^_+|_+$/g, ''); + +export const validateOrganizationPayload = ( + data: CreateOrganizationData, +): Required => { + const name = normalizeOrganizationName(data.name || ''); + const identifier = normalizeOrganizationIdentifier( + data.identifier || data.name || '', + ); + const description = (data.description || '').trim(); + + if (!name) { + throw new ApiError('Organization name is required', 400); + } + if (name.length > NAME_MAX_LENGTH) { + throw new ApiError( + `Organization name is too long (max ${NAME_MAX_LENGTH} characters)`, + 400, + ); + } + if (!identifier) { + throw new ApiError('Organization identifier is required', 400); + } + if (identifier.length > IDENTIFIER_MAX_LENGTH) { + throw new ApiError( + `Organization identifier is too long (max ${IDENTIFIER_MAX_LENGTH} characters)`, + 400, + ); + } + if (!IDENTIFIER_PATTERN.test(identifier)) { + throw new ApiError( + 'Organization identifier must contain only lowercase letters, numbers, underscores, or hyphens', + 400, + ); + } + if (description.length > DESCRIPTION_MAX_LENGTH) { + throw new ApiError( + `Organization description is too long (max ${DESCRIPTION_MAX_LENGTH} characters)`, + 400, + ); + } + + return { + name, + identifier, + description, + }; +}; + +export const validateOrganizationUpdatePayload = ( + data: UpdateOrganizationData, +): UpdateOrganizationData => { + const name = normalizeOrganizationName(data.name || ''); + + if (!name) { + throw new ApiError('Organization name is required', 400); + } + if (name.length > NAME_MAX_LENGTH) { + throw new ApiError( + `Organization name is too long (max ${NAME_MAX_LENGTH} characters)`, + 400, + ); + } + + return { name }; +}; + +export class OrganizationService implements IOrganizationService { + private organizationRepository: IOrganizationRepository; + + constructor({ + organizationRepository, + }: { + organizationRepository: IOrganizationRepository; + }) { + this.organizationRepository = organizationRepository; + } + + public async listOrganizations() { + return await this.organizationRepository.findAll({ order: 'id' }); + } + + public async getCurrentOrganization() { + return await this.organizationRepository.getCurrentOrganization(); + } + + public async createOrganization(data: CreateOrganizationData) { + const payload = validateOrganizationPayload(data); + const existingByIdentifier = await this.organizationRepository.findOneBy({ + identifier: payload.identifier, + }); + if (existingByIdentifier) { + throw new ApiError('Organization identifier already exists', 409); + } + + const tx = await this.organizationRepository.transaction(); + try { + const currentOrganization = + await this.organizationRepository.getCurrentOrganization({ tx }); + const created = await this.organizationRepository.createOne( + { + ...payload, + isCurrent: !currentOrganization, + }, + { tx }, + ); + const result = currentOrganization + ? created + : await this.organizationRepository.setCurrentOrganization(created.id, { + tx, + }); + await this.organizationRepository.commit(tx); + return result; + } catch (error) { + await this.organizationRepository.rollback(tx); + throw error; + } + } + + public async selectCurrentOrganization(id: number) { + const organization = await this.organizationRepository.findOneBy({ id }); + if (!organization) { + throw new ApiError('Organization not found', 404); + } + + return await this.organizationRepository.setCurrentOrganization(id); + } + + public async updateCurrentOrganization(data: UpdateOrganizationData) { + const currentOrganization = + await this.organizationRepository.getCurrentOrganization(); + if (!currentOrganization) { + throw new ApiError('Current organization not found', 404); + } + + const payload = validateOrganizationUpdatePayload(data); + return await this.organizationRepository.updateOne(currentOrganization.id, { + name: payload.name, + }); + } +} diff --git a/wren-ui/src/apollo/server/services/projectService.ts b/wren-ui/src/apollo/server/services/projectService.ts index e6e70e014e..7d63f0f528 100644 --- a/wren-ui/src/apollo/server/services/projectService.ts +++ b/wren-ui/src/apollo/server/services/projectService.ts @@ -2,7 +2,11 @@ import crypto from 'crypto'; import * as fs from 'fs'; import path from 'path'; import { getLogger } from '@server/utils'; -import { IProjectRepository, WREN_AI_CONNECTION_INFO } from '../repositories'; +import { + IProjectRepository, + WREN_AI_CONNECTION_INFO, + WorkspaceProjectType, +} from '../repositories'; import { Project } from '../repositories'; import { CompactTable, @@ -42,17 +46,18 @@ export interface ProjectData { displayName: string; type: DataSourceName; connectionInfo: WREN_AI_CONNECTION_INFO; + projectType?: WorkspaceProjectType; } -export interface ProjectRecommendationQuestionsResult { +export type ProjectRecommendationQuestionsResult = { status: RecommendQuestionResultStatus; questions: RecommendationQuestion[]; - error: WrenAIError; -} + error: WrenAIError | null; +}; export interface IProjectService { createProject: (projectData: ProjectData) => Promise; updateProject: ( - projectId: number, + projectId: string | number, projectData: Partial, ) => Promise; getGeneralConnectionInfo: (project: Project) => Record; @@ -70,7 +75,9 @@ export interface IProjectService { ) => Promise; getCurrentProject: () => Promise; - getProjectById: (projectId: number) => Promise; + listProjects: () => Promise; + selectProject: (projectId: string | number) => Promise; + getProjectById: (projectId: string | number) => Promise; writeCredentialFile: ( credentials: JSON, persistCredentialDir: string, @@ -88,32 +95,40 @@ export class ProjectService implements IProjectService { private mdlService: IMDLService; private wrenAIAdaptor: IWrenAIAdaptor; private projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; + private projectRecommendationJobs = new Map>(); constructor({ projectRepository, metadataService, mdlService, wrenAIAdaptor, telemetry, + projectRecommendQuestionBackgroundTracker, }: { projectRepository: IProjectRepository; metadataService: IDataSourceMetadataService; mdlService: IMDLService; wrenAIAdaptor: IWrenAIAdaptor; telemetry: ITelemetry; + projectRecommendQuestionBackgroundTracker?: ProjectRecommendQuestionBackgroundTracker; }) { this.projectRepository = projectRepository; this.metadataService = metadataService; this.mdlService = mdlService; this.wrenAIAdaptor = wrenAIAdaptor; this.projectRecommendQuestionBackgroundTracker = + projectRecommendQuestionBackgroundTracker ?? new ProjectRecommendQuestionBackgroundTracker({ projectRepository, telemetry, wrenAIAdaptor, }); } + + public dispose(): void { + this.projectRecommendQuestionBackgroundTracker.stop(); + } public async updateProject( - projectId: number, + projectId: string | number, projectData: Partial, ): Promise { return await this.projectRepository.updateOne(projectId, projectData); @@ -136,10 +151,33 @@ export class ProjectService implements IProjectService { if (!project) { throw new Error(`Project not found`); } - const { manifest } = await this.mdlService.makeCurrentModelMDL(); + + const projectJobKey = String(project.id); + const existingJob = this.projectRecommendationJobs.get(projectJobKey); + if (existingJob) { + logger.debug( + `project "${project.id}" recommended questions are already being requested, reusing in-flight job`, + ); + return existingJob; + } + + const job = this.doGenerateProjectRecommendationQuestions(project); + this.projectRecommendationJobs.set(projectJobKey, job); + try { + return await job; + } finally { + this.projectRecommendationJobs.delete(projectJobKey); + } + } + + private async doGenerateProjectRecommendationQuestions( + project: Project, + ): Promise { + const { manifest } = await this.mdlService.makeModelMDL(project); const recommendQuestionResult = await this.wrenAIAdaptor.generateRecommendationQuestions({ manifest, + projectId: project.id.toString(), ...this.getProjectRecommendationQuestionsConfig(project), }); @@ -185,7 +223,16 @@ export class ProjectService implements IProjectService { return await this.projectRepository.getCurrentProject(); } - public async getProjectById(projectId: number) { + public async listProjects() { + return await this.projectRepository.listProjects(); + } + + public async selectProject(projectId: string | number) { + logger.debug(`Selecting active project ${String(projectId)}`); + return await this.projectRepository.setCurrentProject(projectId); + } + + public async getProjectById(projectId: string | number) { return await this.projectRepository.findOneBy({ id: projectId }); } @@ -219,6 +266,8 @@ export class ProjectService implements IProjectService { type: projectData.type, catalog: 'wrenai', schema: 'public', + projectType: projectData.projectType || WorkspaceProjectType.CLASSIC, + isCurrent: false, connectionInfo: encryptConnectionInfo( projectData.type, projectData.connectionInfo, @@ -226,7 +275,7 @@ export class ProjectService implements IProjectService { }; logger.debug('Creating project...'); const project = await this.projectRepository.createOne(projectValue); - return project; + return await this.projectRepository.setCurrentProject(project.id); } public writeCredentialFile(credentials: JSON, persistCredentialDir: string) { @@ -254,7 +303,33 @@ export class ProjectService implements IProjectService { } public async deleteProject(projectId: number): Promise { + const [projectToDelete, currentProject, remainingProjects] = + await Promise.all([ + this.projectRepository.findOneBy({ id: projectId }), + this.projectRepository.findCurrentProject(), + this.projectRepository.listProjects(), + ]); + await this.projectRepository.deleteOne(projectId); + + if (currentProject?.id !== projectId) { + return; + } + + const nextProject = remainingProjects + .filter((project) => String(project.id) !== String(projectId)) + .sort((a, b) => this.compareProjectIdsDescending(a.id, b.id))[0]; + + if (nextProject) { + await this.projectRepository.setCurrentProject(nextProject.id); + return; + } + + if (projectToDelete) { + logger.debug( + `Deleted the last current project ${projectToDelete.id}; no active project remains.`, + ); + } } public getGeneralConnectionInfo(project) { @@ -279,4 +354,16 @@ export class ProjectService implements IProjectService { }, }; } + + private compareProjectIdsDescending( + left: string | number, + right: string | number, + ) { + const leftId = BigInt(left); + const rightId = BigInt(right); + if (leftId === rightId) { + return 0; + } + return leftId > rightId ? -1 : 1; + } } diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts index 08c78d6132..b22532040b 100644 --- a/wren-ui/src/apollo/server/services/queryService.ts +++ b/wren-ui/src/apollo/server/services/queryService.ts @@ -16,6 +16,38 @@ const logger = getLogger('QueryService'); logger.level = 'debug'; export const DEFAULT_PREVIEW_LIMIT = 500; +const MSSQL_DEADLOCK_RETRY_LIMIT = 2; +const MSSQL_DEADLOCK_RETRY_BASE_DELAY_MS = 150; + +const delay = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); + +const errorText = (err: any) => + [ + err?.message, + err?.response?.data?.message, + err?.response?.data?.detail, + err?.extensions?.message, + err?.extensions?.originalError?.message, + ] + .filter(Boolean) + .join(' ') + .toLowerCase(); + +const isSqlServerDeadlockError = (err: any) => { + const text = errorText(err); + if (!text) { + return false; + } + + return ( + text.includes('deadlock') && + (text.includes('1205') || + text.includes('deadlock victim') || + text.includes('sqlexecdirectw') || + text.includes('sql server')) + ); +}; export interface ColumnMetadata { name: string; @@ -111,21 +143,42 @@ export class QueryService implements IQueryService { if (this.useEngine(dataSource)) { if (dryRun) { logger.debug('Using wren engine to dry run'); + const startedAt = Date.now(); await this.wrenEngineAdaptor.dryRun(sql, { manifest: mdl, limit, }); + logger.info( + `Ask timing stage=sql_validation project_id=${project.id} data_source=${dataSource} engine=wren elapsed_ms=${ + Date.now() - startedAt + }`, + ); return true; } else { logger.debug('Using wren engine to preview'); - const data = await this.wrenEngineAdaptor.previewData(sql, mdl, limit); + const startedAt = Date.now(); + const data = await this.wrenEngineAdaptor.previewData( + sql, + mdl, + limit, + ); + logger.info( + `Ask timing stage=sql_execution project_id=${project.id} data_source=${dataSource} engine=wren elapsed_ms=${ + Date.now() - startedAt + } row_count=${(data as PreviewDataResponse)?.data?.length ?? ''}`, + ); return data as PreviewDataResponse; } } else { this.checkDataSourceIsSupported(dataSource); logger.debug('Use ibis adaptor to preview'); if (dryRun) { - return await this.ibisDryRun(sql, dataSource, connectionInfo, mdl); + return await this.ibisDryRun( + sql, + dataSource, + connectionInfo, + mdl, + ); } else { return await this.ibisQuery( sql, @@ -195,6 +248,7 @@ export class QueryService implements IQueryService { mdl: Manifest, ): Promise { const event = TelemetryEvent.IBIS_DRY_RUN; + const startedAt = Date.now(); try { const res = await this.ibisAdaptor.dryRun(sql, { dataSource, @@ -202,11 +256,24 @@ export class QueryService implements IQueryService { mdl, }); this.sendIbisEvent(event, res, { dataSource, sql }); + logger.info( + `Ask timing stage=sql_validation data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + }`, + ); return { correlationId: res.correlationId, }; } catch (err: any) { - this.sendIbisFailedEvent(event, err, { dataSource, sql }); + logger.info( + `Ask timing stage=sql_validation data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + } status=failed`, + ); + this.sendIbisFailedEvent(event, err, { + dataSource, + sql, + }); throw err; } } @@ -221,27 +288,73 @@ export class QueryService implements IQueryService { cacheEnabled?: boolean, ): Promise { const event = TelemetryEvent.IBIS_QUERY; + let attempt = 0; + const startedAt = Date.now(); try { - const res = await this.ibisAdaptor.query(sql, { + let res: IbisQueryResponse | undefined; + while (true) { + try { + res = await this.ibisAdaptor.query(sql, { + dataSource, + connectionInfo, + mdl, + limit, + refresh, + cacheEnabled, + }); + break; + } catch (err: any) { + const canRetry = + dataSource === DataSourceName.MSSQL && + isSqlServerDeadlockError(err) && + attempt < MSSQL_DEADLOCK_RETRY_LIMIT; + + if (!canRetry) { + throw err; + } + + attempt += 1; + logger.warn( + `MSSQL deadlock while querying ibis; retrying attempt ${attempt}/${MSSQL_DEADLOCK_RETRY_LIMIT}`, + ); + await delay(MSSQL_DEADLOCK_RETRY_BASE_DELAY_MS * attempt); + } + } + + if (!res) { + throw new Error('Ibis query did not return a response'); + } + + this.sendIbisEvent(event, res, { dataSource, - connectionInfo, - mdl, - limit, - refresh, - cacheEnabled, + sql, }); - this.sendIbisEvent(event, res, { dataSource, sql }); const data = this.transformDataType(res); + logger.info( + `Ask timing stage=sql_execution data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + } row_count=${data.data?.length ?? ''} cache_hit=${ + res.cacheHit ?? false + } attempts=${attempt + 1}`, + ); return { correlationId: res.correlationId, - cacheHit: res.cacheHit, + cacheHit: res.cacheHit ?? false, cacheCreatedAt: res.cacheCreatedAt, cacheOverrodeAt: res.cacheOverrodeAt, - override: res.override, + override: res.override ?? false, ...data, }; } catch (err: any) { - this.sendIbisFailedEvent(event, err, { dataSource, sql }); + logger.info( + `Ask timing stage=sql_execution data_source=${dataSource} engine=ibis elapsed_ms=${ + Date.now() - startedAt + } status=failed attempts=${attempt + 1}`, + ); + this.sendIbisFailedEvent(event, err, { + dataSource, + sql, + }); throw err; } } diff --git a/wren-ui/src/apollo/server/services/rbacService.ts b/wren-ui/src/apollo/server/services/rbacService.ts new file mode 100644 index 0000000000..76064df4b5 --- /dev/null +++ b/wren-ui/src/apollo/server/services/rbacService.ts @@ -0,0 +1,297 @@ +import { Knex } from 'knex'; +import { + CreateRoleInput, + CreateUserInput, + UpdateRoleInput, + UpdateUserInput, + UpdateUserRolesInput, + UserRoleInput, +} from '@server/models'; +import { + IRoleRepository, + IUserRepository, + IUserRoleRepository, + RbacUser, + Role, + UserRole, + UserRoleMapping, +} from '@server/repositories'; + +export interface RbacUserWithRoles extends RbacUser { + roles: Role[]; +} + +export interface RoleWithUsers extends Role { + users: RbacUser[]; +} + +export interface IRbacService { + listRoles(): Promise; + createRole(input: CreateRoleInput): Promise; + updateRole(input: UpdateRoleInput): Promise; + listUsers(): Promise; + createUser(input: CreateUserInput): Promise; + updateUser(input: UpdateUserInput): Promise; + assignRoleToUser(input: UserRoleInput): Promise; + updateUserRoles(input: UpdateUserRolesInput): Promise; + removeRoleFromUser(input: UserRoleInput): Promise; + getUserRoleMappings(): Promise; +} + +const EMAIL_PATTERN = /^[^\s@]+@[^\s@]+\.[^\s@]+$/; + +export class RbacService implements IRbacService { + private readonly roleRepository: IRoleRepository; + private readonly userRepository: IUserRepository; + private readonly userRoleRepository: IUserRoleRepository; + + constructor({ + roleRepository, + userRepository, + userRoleRepository, + }: { + roleRepository: IRoleRepository; + userRepository: IUserRepository; + userRoleRepository: IUserRoleRepository; + }) { + this.roleRepository = roleRepository; + this.userRepository = userRepository; + this.userRoleRepository = userRoleRepository; + } + + public async listRoles(): Promise { + const roles = await this.roleRepository.findAll({ order: 'name' }); + const mappings = await this.userRoleRepository.findMappings(); + return roles.map((role) => ({ + ...role, + users: mappings + .filter((mapping) => mapping.roleId === role.id) + .map((mapping) => mapping.user), + })); + } + + public async createRole(input: CreateRoleInput): Promise { + const name = this.validateRoleName(input.name); + await this.assertUniqueRoleName(name); + const now = new Date().toISOString(); + return this.roleRepository.createOne({ + name, + description: this.normalizeNullable(input.description), + createdAt: now, + updatedAt: now, + }); + } + + public async updateRole(input: UpdateRoleInput): Promise { + const role = await this.getRoleOrThrow(input.id); + const data: Partial = { updatedAt: new Date().toISOString() }; + + if (input.name !== undefined && input.name !== null) { + const name = this.validateRoleName(input.name); + await this.assertUniqueRoleName(name, role.id); + data.name = name; + } + if (input.description !== undefined) { + data.description = this.normalizeNullable(input.description); + } + + return this.roleRepository.updateOne(role.id, data); + } + + public async listUsers(): Promise { + const users = await this.userRepository.findAll({ order: 'email' }); + const mappings = await this.userRoleRepository.findMappings(); + return users.map((user) => ({ + ...user, + roles: mappings + .filter((mapping) => mapping.userId === user.id) + .map((mapping) => mapping.role), + })); + } + + public async createUser(input: CreateUserInput): Promise { + const name = this.validateRequiredText(input.name, 'User name'); + const email = this.validateEmail(input.email); + await this.assertUniqueUserEmail(email); + const now = new Date().toISOString(); + + const tx = await this.userRepository.transaction(); + try { + const user = await this.userRepository.createOne( + { + name, + email, + externalId: this.normalizeNullable(input.externalId), + identityProvider: this.normalizeNullable(input.identityProvider), + isActive: input.isActive ?? true, + createdAt: now, + updatedAt: now, + }, + { tx }, + ); + + if (input.roleIds?.length) { + await this.createUserRoleAssignments(user.id, input.roleIds, tx); + } + + await tx.commit(); + return user; + } catch (error) { + await tx.rollback(); + throw error; + } + } + + public async updateUser(input: UpdateUserInput): Promise { + const user = await this.getUserOrThrow(input.id); + const data: Partial = { updatedAt: new Date().toISOString() }; + + if (input.name !== undefined && input.name !== null) { + data.name = this.validateRequiredText(input.name, 'User name'); + } + if (input.email !== undefined && input.email !== null) { + const email = this.validateEmail(input.email); + await this.assertUniqueUserEmail(email, user.id); + data.email = email; + } + if (input.externalId !== undefined) { + data.externalId = this.normalizeNullable(input.externalId); + } + if (input.identityProvider !== undefined) { + data.identityProvider = this.normalizeNullable(input.identityProvider); + } + if (input.isActive !== undefined && input.isActive !== null) { + data.isActive = input.isActive; + } + + return this.userRepository.updateOne(user.id, data); + } + + public async assignRoleToUser(input: UserRoleInput): Promise { + await this.getUserOrThrow(input.userId); + await this.getRoleOrThrow(input.roleId); + const existing = await this.userRoleRepository.findOneBy(input); + if (existing) return existing; + + const now = new Date().toISOString(); + return this.userRoleRepository.createOne({ + ...input, + createdAt: now, + updatedAt: now, + }); + } + + public async updateUserRoles( + input: UpdateUserRolesInput, + ): Promise { + const user = await this.getUserOrThrow(input.userId); + const roleIds = this.uniqueIds(input.roleIds); + const tx = await this.userRoleRepository.transaction(); + + try { + await this.userRoleRepository.deleteAllBy({ userId: user.id }, { tx }); + if (roleIds.length) { + await this.createUserRoleAssignments(user.id, roleIds, tx); + } + await tx.commit(); + } catch (error) { + await tx.rollback(); + throw error; + } + + const mappings = await this.userRoleRepository.findMappingsByUserId( + user.id, + ); + return { ...user, roles: mappings.map((mapping) => mapping.role) }; + } + + public async removeRoleFromUser(input: UserRoleInput): Promise { + const existing = await this.userRoleRepository.findOneBy(input); + if (!existing) return true; + await this.userRoleRepository.deleteOne(existing.id); + return true; + } + + public async getUserRoleMappings(): Promise { + return this.userRoleRepository.findMappings(); + } + + private async createUserRoleAssignments( + userId: number, + roleIds: number[], + tx: Knex.Transaction, + ): Promise { + const uniqueRoleIds = this.uniqueIds(roleIds); + for (const roleId of uniqueRoleIds) { + await this.getRoleOrThrow(roleId); + } + const now = new Date().toISOString(); + await this.userRoleRepository.createMany( + uniqueRoleIds.map((roleId) => ({ + userId, + roleId, + createdAt: now, + updatedAt: now, + })), + { tx }, + ); + } + + private async getRoleOrThrow(id: number): Promise { + const role = await this.roleRepository.findOneBy({ id }); + if (!role) throw new Error(`Role ${id} was not found.`); + return role; + } + + private async getUserOrThrow(id: number): Promise { + const user = await this.userRepository.findOneBy({ id }); + if (!user) throw new Error(`User ${id} was not found.`); + return user; + } + + private async assertUniqueRoleName(name: string, exceptId?: number) { + const roles = await this.roleRepository.findAll(); + const duplicate = roles.find( + (role) => + role.name.toLowerCase() === name.toLowerCase() && role.id !== exceptId, + ); + if (duplicate) throw new Error(`Role "${name}" already exists.`); + } + + private async assertUniqueUserEmail(email: string, exceptId?: number) { + const users = await this.userRepository.findAll(); + const duplicate = users.find( + (user) => + user.email.toLowerCase() === email.toLowerCase() && + user.id !== exceptId, + ); + if (duplicate) throw new Error(`User "${email}" already exists.`); + } + + private validateRoleName(name: string): string { + return this.validateRequiredText(name, 'Role name'); + } + + private validateEmail(email: string): string { + const normalized = this.validateRequiredText(email, 'Email').toLowerCase(); + if (!EMAIL_PATTERN.test(normalized)) { + throw new Error('A valid email address is required.'); + } + return normalized; + } + + private validateRequiredText(value: string, label: string): string { + const normalized = `${value || ''}`.trim(); + if (!normalized) throw new Error(`${label} is required.`); + return normalized; + } + + private normalizeNullable(value?: string | null): string | null { + const normalized = `${value || ''}`.trim(); + return normalized || null; + } + + private uniqueIds(ids: number[]): number[] { + return Array.from(new Set((ids || []).filter(Boolean))); + } +} diff --git a/wren-ui/src/apollo/server/services/sqlPairService.ts b/wren-ui/src/apollo/server/services/sqlPairService.ts index fd9c8b186d..51eb915486 100644 --- a/wren-ui/src/apollo/server/services/sqlPairService.ts +++ b/wren-ui/src/apollo/server/services/sqlPairService.ts @@ -17,6 +17,7 @@ import { } from '../models/adaptor'; import { Manifest } from '@server/mdl/type'; import { DataSourceName } from '@server/types'; +import { IDeployService } from './deployService'; const logger = getLogger('SqlPairService'); @@ -60,19 +61,23 @@ export class SqlPairService implements ISqlPairService { private sqlPairRepository: ISqlPairRepository; private wrenAIAdaptor: IWrenAIAdaptor; private ibisAdaptor: IIbisAdaptor; + private deployService: IDeployService; constructor({ sqlPairRepository, wrenAIAdaptor, ibisAdaptor, + deployService, }: { sqlPairRepository: ISqlPairRepository; wrenAIAdaptor: IWrenAIAdaptor; ibisAdaptor: IIbisAdaptor; + deployService: IDeployService; }) { this.sqlPairRepository = sqlPairRepository; this.wrenAIAdaptor = wrenAIAdaptor; this.ibisAdaptor = ibisAdaptor; + this.deployService = deployService; } public async modelSubstitute( @@ -145,9 +150,11 @@ export class SqlPairService implements ISqlPairService { }, { tx }, ); + const mdlHash = await this.getDeployHash(projectId); const { queryId } = await this.wrenAIAdaptor.deploySqlPair( projectId, newPair, + mdlHash, ); const deployResult = await this.waitUntilSqlPairResult(queryId); if (deployResult.error) { @@ -179,12 +186,14 @@ export class SqlPairService implements ISqlPairService { const successPairs = []; const errorPairs = []; const chunks = chunk(newPairs, 10); + const mdlHash = await this.getDeployHash(projectId); for (const pairs of chunks) { await Promise.allSettled( pairs.map(async (pair) => { const { queryId } = await this.wrenAIAdaptor.deploySqlPair( projectId, pair, + mdlHash, ); const deployResult = await this.waitUntilSqlPairResult(queryId); if (deployResult.error) { @@ -249,9 +258,11 @@ export class SqlPairService implements ISqlPairService { updatedData, { tx }, ); + const mdlHash = await this.getDeployHash(projectId); const { queryId } = await this.wrenAIAdaptor.deploySqlPair( projectId, updatedSqlPair, + mdlHash, ); const deployResult = await this.waitUntilSqlPairResult(queryId); if (deployResult.error) { @@ -326,4 +337,9 @@ export class SqlPairService implements ISqlPairService { private isFinishedState(status: SqlPairStatus) { return [SqlPairStatus.FINISHED, SqlPairStatus.FAILED].includes(status); } + + private async getDeployHash(projectId: number): Promise { + const deployment = await this.deployService.getLastDeployment(projectId); + return deployment?.hash; + } } diff --git a/wren-ui/src/apollo/server/services/tests/askingService.test.ts b/wren-ui/src/apollo/server/services/tests/askingService.test.ts index e6fa374c14..1d88c85131 100644 --- a/wren-ui/src/apollo/server/services/tests/askingService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/askingService.test.ts @@ -1,4 +1,4 @@ -import { constructCteSql } from '../askingService'; +import { AskingService, constructCteSql } from '../askingService'; describe('AskingService', () => { describe('utility: constructCteSql', () => { @@ -94,4 +94,161 @@ describe('AskingService', () => { ); }); }); + + describe('project lookup for thread response', () => { + test('falls back to current project when parent thread is missing', async () => { + const currentProject = { id: 2, type: 'mssql' }; + const service = Object.create(AskingService.prototype) as any; + service.threadRepository = { + findOneBy: jest.fn().mockResolvedValue(null), + }; + service.projectService = { + getCurrentProject: jest.fn().mockResolvedValue(currentProject), + getProjectById: jest.fn(), + }; + + const project = await service.getProjectForThreadResponse({ + id: 10, + threadId: 530, + }); + + expect(project).toBe(currentProject); + expect(service.projectService.getCurrentProject).toHaveBeenCalledTimes(1); + expect(service.projectService.getProjectById).not.toHaveBeenCalled(); + }); + }); + + describe('recommendation question shortcut SQL', () => { + const trackedAskingResult = { + taskId: 42, + queryId: 'ask-query-id', + question: 'Show monthly record count', + status: 'UNDERSTANDING', + response: null, + error: null, + }; + + const createService = () => { + const manifest = { models: [{ name: 'current_model' }] }; + const service = Object.create(AskingService.prototype) as any; + service.projectService = { + getCurrentProject: jest.fn().mockResolvedValue({ + id: 1, + language: 'EN', + }), + getProjectById: jest.fn().mockResolvedValue({ + id: 1, + language: 'EN', + }), + }; + service.deployService = { + ensureDeploymentPrepared: jest + .fn() + .mockResolvedValue('latest-deploy-hash'), + }; + service.mdlService = { + makeModelMDL: jest.fn().mockResolvedValue({ manifest }), + }; + service.threadRepository = { + createOne: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), + findOneBy: jest.fn().mockResolvedValue({ id: 7, projectId: 1 }), + }; + service.threadResponseRepository = { + createOne: jest.fn().mockResolvedValue({ id: 11, threadId: 7 }), + getResponsesWithThread: jest.fn().mockResolvedValue([]), + }; + service.askingTaskTracker = { + createAskingTask: jest.fn().mockResolvedValue({ + queryId: trackedAskingResult.queryId, + }), + getAskingResult: jest.fn().mockResolvedValue(trackedAskingResult), + bindThreadResponse: jest.fn().mockResolvedValue(undefined), + }; + service.currentManifest = manifest; + return service; + }; + + test('creates a normal asking task for a new thread instead of storing shortcut SQL', async () => { + const service = createService(); + + await service.createThread({ + question: trackedAskingResult.question, + sql: 'SELECT stale_recommendation_sql', + }); + + expect(service.mdlService.makeModelMDL).toHaveBeenCalledWith({ + id: 1, + language: 'EN', + }); + expect(service.deployService.ensureDeploymentPrepared).toHaveBeenCalledWith( + 1, + service.currentManifest, + ); + expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith({ + query: trackedAskingResult.question, + histories: null, + deployId: 'latest-deploy-hash', + projectId: '1', + configurations: { language: 'EN' }, + ignoreSqlGenerationReasoning: true, + enableColumnPruning: true, + useDryPlan: true, + allowDryPlanFallback: true, + rerunFromCancelled: undefined, + previousTaskId: undefined, + threadResponseId: undefined, + }); + expect(service.threadResponseRepository.createOne).toHaveBeenCalledWith({ + threadId: 7, + question: trackedAskingResult.question, + sql: undefined, + askingTaskId: trackedAskingResult.taskId, + }); + expect(service.askingTaskTracker.bindThreadResponse).toHaveBeenCalledWith( + trackedAskingResult.taskId, + trackedAskingResult.queryId, + 7, + 11, + ); + }); + + test('creates a standalone asking task for current-thread recommendations', async () => { + const service = createService(); + service.threadResponseRepository.getResponsesWithThread.mockResolvedValue([ + { id: 1, question: 'Previous question', sql: 'SELECT 1' }, + ]); + + await service.createThreadResponse( + { + question: trackedAskingResult.question, + sql: 'SELECT stale_recommendation_sql', + }, + 7, + ); + + expect(service.projectService.getProjectById).toHaveBeenCalledWith(1); + expect(service.askingTaskTracker.createAskingTask).toHaveBeenCalledWith( + expect.objectContaining({ + query: trackedAskingResult.question, + histories: null, + deployId: 'latest-deploy-hash', + projectId: '1', + configurations: { language: 'EN' }, + ignoreSqlGenerationReasoning: true, + enableColumnPruning: true, + useDryPlan: true, + allowDryPlanFallback: true, + }), + ); + expect( + service.threadResponseRepository.getResponsesWithThread, + ).not.toHaveBeenCalled(); + expect(service.threadResponseRepository.createOne).toHaveBeenCalledWith({ + threadId: 7, + question: trackedAskingResult.question, + sql: undefined, + askingTaskId: trackedAskingResult.taskId, + }); + }); + }); }); diff --git a/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts b/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts new file mode 100644 index 0000000000..99bdb59811 --- /dev/null +++ b/wren-ui/src/apollo/server/services/tests/askingTaskTracker.test.ts @@ -0,0 +1,222 @@ +import { + AskingTaskTracker, + TrackedAskingResult, +} from '../askingTaskTracker'; +import { + AskResultStatus, + AskResultType, +} from '@server/models/adaptor'; +import * as Errors from '@server/utils/error'; + +describe('AskingTaskTracker', () => { + const createTracker = ({ + taskRecords = [], + memoryRetentionTime = 1000, + }: { + taskRecords?: any[]; + memoryRetentionTime?: number; + }) => { + const askingTaskRepository = { + findAll: jest.fn().mockResolvedValue(taskRecords), + findByQueryId: jest.fn(), + findOneBy: jest.fn(), + createOne: jest.fn(), + updateOne: jest.fn(), + }; + const wrenAIAdaptor = { + getAskResult: jest.fn(), + ask: jest.fn(), + cancelAsk: jest.fn(), + }; + const tracker = new AskingTaskTracker({ + wrenAIAdaptor: wrenAIAdaptor as any, + askingTaskRepository: askingTaskRepository as any, + threadResponseRepository: { updateOne: jest.fn() } as any, + viewRepository: { findOneBy: jest.fn() } as any, + pollingInterval: 1000, + memoryRetentionTime, + }); + tracker.stopPolling(); + + return { + tracker, + askingTaskRepository, + wrenAIAdaptor, + }; + }; + + test('finalizes stale unfinished tasks during initialization without polling AI service', async () => { + const oldDate = new Date(Date.now() - 60_000); + const staleTask = { + id: 7, + queryId: 'expired-query-id', + question: 'expired question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: oldDate, + updatedAt: oldDate, + }; + const { tracker, askingTaskRepository, wrenAIAdaptor } = createTracker({ + taskRecords: [staleTask], + memoryRetentionTime: 1000, + }); + + await tracker.initialize(); + + expect(wrenAIAdaptor.getAskResult).not.toHaveBeenCalled(); + expect(askingTaskRepository.updateOne).toHaveBeenCalledWith(7, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }, + }); + }); + + test('restores recent unfinished tasks so active AI service requests can continue', async () => { + const recentDate = new Date(); + const recentTask = { + id: 8, + queryId: 'recent-query-id', + question: 'recent question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: recentDate, + updatedAt: recentDate, + }; + const { tracker, askingTaskRepository } = createTracker({ + taskRecords: [recentTask], + memoryRetentionTime: 60_000, + }); + + await tracker.initialize(); + const result = (await tracker.getAskingResult( + recentTask.queryId, + )) as TrackedAskingResult; + + expect(askingTaskRepository.updateOne).not.toHaveBeenCalled(); + expect(result.queryId).toBe(recentTask.queryId); + expect(result.status).toBe(AskResultStatus.GENERATING); + }); + + test('finalizes unfinished tasks with invalid timestamps during initialization', async () => { + const taskWithInvalidTimestamp = { + id: 10, + queryId: 'invalid-timestamp-query-id', + question: 'expired question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: 'not-a-date', + updatedAt: 'not-a-date', + }; + const { tracker, askingTaskRepository, wrenAIAdaptor } = createTracker({ + taskRecords: [taskWithInvalidTimestamp], + memoryRetentionTime: 1000, + }); + + await tracker.initialize(); + + expect(wrenAIAdaptor.getAskResult).not.toHaveBeenCalled(); + expect(askingTaskRepository.updateOne).toHaveBeenCalledWith(10, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }, + }); + }); + + test('treats numeric second timestamps as valid freshness timestamps', async () => { + const recentSeconds = Math.floor(Date.now() / 1000); + const recentTask = { + id: 11, + queryId: 'recent-seconds-query-id', + question: 'recent question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: recentSeconds, + updatedAt: recentSeconds, + }; + const { tracker, askingTaskRepository } = createTracker({ + taskRecords: [recentTask], + memoryRetentionTime: 60_000, + }); + + await tracker.initialize(); + + expect(askingTaskRepository.updateOne).not.toHaveBeenCalled(); + }); + + test('finalizes restored tasks when AI service no longer has the query id', async () => { + const recentDate = new Date(); + const recentTask = { + id: 9, + queryId: 'missing-query-id', + question: 'recent question', + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.GENERATING, + response: null, + error: null, + }, + createdAt: recentDate, + updatedAt: recentDate, + }; + const { tracker, askingTaskRepository, wrenAIAdaptor } = createTracker({ + taskRecords: [recentTask], + memoryRetentionTime: 60_000, + }); + wrenAIAdaptor.getAskResult.mockResolvedValue({ + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.ASK_RESULT_NOT_FOUND, + message: 'The asking task result is no longer available', + }, + }); + + await tracker.initialize(); + await (tracker as any).pollTasks(); + + expect(wrenAIAdaptor.getAskResult).toHaveBeenCalledWith(recentTask.queryId); + expect(askingTaskRepository.updateOne).toHaveBeenCalledWith(9, { + detail: { + type: AskResultType.TEXT_TO_SQL, + status: AskResultStatus.FAILED, + response: null, + error: { + code: Errors.GeneralErrorCodes.POLLING_TIMEOUT, + message: + 'The previous asking task expired after the service restarted. Please ask again.', + }, + }, + }); + }); +}); diff --git a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts index 55814d5610..11e09af2c6 100644 --- a/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/dashboardService.test.ts @@ -85,6 +85,98 @@ describe('DashboardService', () => { }); }); + describe('dashboard initialization', () => { + it('should return the existing dashboard for the current project', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + + await expect(dashboardService.getCurrentDashboard()).resolves.toEqual( + dashboard, + ); + expect(mockDashboardRepository.createOne).not.toHaveBeenCalled(); + }); + + it('should initialize a dashboard when the current project has none', async () => { + const project = { id: 42 }; + const createdDashboard = { + id: '1002', + projectId: 42, + name: 'Dashboard', + }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValueOnce(null); + mockDashboardRepository.findOneBy.mockResolvedValueOnce(null); + mockDashboardRepository.createOne.mockResolvedValue(createdDashboard); + + await expect(dashboardService.getCurrentDashboard()).resolves.toEqual( + createdDashboard, + ); + expect(mockDashboardRepository.createOne).toHaveBeenCalledWith({ + name: 'Dashboard', + projectId: project.id, + }); + }); + }); + + describe('dashboard scoping', () => { + it('should only return dashboard items from the current dashboard', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + const dashboardItem = { + id: '2001', + dashboardId: '1001', + layout: { x: 0, y: 0, w: 3, h: 2 }, + }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + mockDashboardItemRepository.findOneBy.mockResolvedValue(dashboardItem); + + await expect(dashboardService.getDashboardItem('2001')).resolves.toEqual( + dashboardItem, + ); + }); + + it('should reject dashboard items from another dashboard context', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + const dashboardItem = { + id: '2001', + dashboardId: '9999', + layout: { x: 0, y: 0, w: 3, h: 2 }, + }; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + mockDashboardItemRepository.findOneBy.mockResolvedValue(dashboardItem); + + await expect(dashboardService.getDashboardItem('2001')).rejects.toThrow( + 'Dashboard item not found.', + ); + }); + + it('should reject layout updates for items outside the current dashboard', async () => { + const project = { id: 42 }; + const dashboard = { id: '1001', projectId: 42, name: 'Dashboard' }; + const dashboardItems = [ + { + id: '2001', + dashboardId: '1001', + layout: { x: 0, y: 0, w: 3, h: 2 }, + }, + ]; + mockProjectService.getCurrentProject.mockResolvedValue(project); + mockDashboardRepository.findOneBy.mockResolvedValue(dashboard); + mockDashboardItemRepository.findAllBy.mockResolvedValue(dashboardItems); + + await expect( + dashboardService.updateDashboardItemLayouts([ + { itemId: '9999', x: 0, y: 0, w: 3, h: 2 }, + ]), + ).rejects.toThrow('Invalid layouts boundaries.'); + }); + }); + describe('generateCronExpression', () => { it('should generate correct cron expression for daily schedule', () => { const schedule = { diff --git a/wren-ui/src/apollo/server/services/tests/deployService.test.ts b/wren-ui/src/apollo/server/services/tests/deployService.test.ts index 14f922de85..f749b9b6f7 100644 --- a/wren-ui/src/apollo/server/services/tests/deployService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/deployService.test.ts @@ -11,8 +11,11 @@ describe('DeployService', () => { beforeEach(() => { mockTelemetry = { sendEvent: jest.fn() }; mockWrenAIAdaptor = { deploy: jest.fn() }; + mockWrenAIAdaptor.getDeployStatus = jest.fn(); mockDeployLogRepository = { findLastProjectDeployLog: jest.fn(), + findInProgressProjectDeployLog: jest.fn(), + findOneBy: jest.fn(), createOne: jest.fn(), updateOne: jest.fn(), }; @@ -35,6 +38,11 @@ describe('DeployService', () => { const response = await deployService.deploy(manifest, projectId); expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash: deployService.createMDLHash(manifest, projectId), + projectId, + }); expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { status: DeployStatusEnum.SUCCESS, error: undefined, @@ -58,19 +66,361 @@ describe('DeployService', () => { expect(response.error).toEqual('AI error'); }); - it('should skip deployment if an existing deployment with the same hash exists', async () => { + it('should mark deployment failed if ai-service deployment throws', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockRejectedValue(new Error('AI unavailable')); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.FAILED); + expect(response.error).toEqual('AI unavailable'); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.FAILED, + error: 'AI unavailable', + }); + }); + + it('should refresh ai-service deployment if an existing deployment with the same hash exists', async () => { const manifest = { key: 'value' }; const projectId = 1; + const hash = deployService.createMDLHash(manifest, 1); mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ - hash: deployService.createMDLHash(manifest, 1), + id: 123, + hash, }); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); const response = await deployService.deploy(manifest, projectId); expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash, + projectId, + }); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.SUCCESS, + error: undefined, + }); + }); + + it('should return the deployment hash when ai-service already has the exact deployment prepared', async () => { + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ + hash: 'deploy-hash', + manifest: { key: 'value' }, + }); + mockWrenAIAdaptor.getDeployStatus.mockResolvedValue('FINISHED'); + + const hash = await deployService.ensureDeploymentPrepared(1); + + expect(hash).toEqual('deploy-hash'); + expect(mockWrenAIAdaptor.getDeployStatus).toHaveBeenCalledWith( + 'deploy-hash', + 1, + ); expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); }); - // Add more tests here to cover other scenarios and error handling + it('should redeploy the saved manifest when ai-service no longer has the exact deployment prepared', async () => { + const manifest = { key: 'value' }; + const hash = deployService.createMDLHash(manifest, 1); + mockDeployLogRepository.findLastProjectDeployLog + .mockResolvedValueOnce({ + id: 123, + hash, + manifest, + }) + .mockResolvedValueOnce({ + id: 123, + hash, + manifest, + }); + mockWrenAIAdaptor.getDeployStatus.mockRejectedValue(new Error('not found')); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + + const preparedHash = await deployService.ensureDeploymentPrepared(1); + + expect(preparedHash).toEqual(hash); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash, + projectId: 1, + }); + }); + + it('should prepare the supplied current manifest instead of reusing a stale saved manifest', async () => { + const staleManifest = { models: [{ name: 'stale_model' }] }; + const currentManifest = { + models: [ + { + name: 'current_model', + tableReference: { schema: 'dbo', table: 'current_model' }, + }, + ], + }; + const activeHash = deployService.createMDLHash(currentManifest, 1); + + mockDeployLogRepository.findLastProjectDeployLog + .mockResolvedValueOnce({ + id: 123, + hash: deployService.createMDLHash(staleManifest, 1), + manifest: staleManifest, + }) + .mockResolvedValueOnce(null); + mockWrenAIAdaptor.getDeployStatus.mockRejectedValue(new Error('not found')); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 456 }); + + const preparedHash = await deployService.ensureDeploymentPrepared( + 1, + currentManifest, + ); + + expect(preparedHash).toEqual(activeHash); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest: currentManifest, + hash: activeHash, + projectId: 1, + }); + }); + + it('should return the current manifest hash after redeploying a stale saved hash', async () => { + const manifest = { key: 'value' }; + const activeHash = deployService.createMDLHash(manifest, 1); + + mockDeployLogRepository.findLastProjectDeployLog + .mockResolvedValueOnce({ + id: 123, + hash: 'legacy-saved-hash', + manifest, + }) + .mockResolvedValueOnce({ + id: 456, + hash: activeHash, + manifest, + }); + mockWrenAIAdaptor.getDeployStatus.mockRejectedValue(new Error('not found')); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 456 }); + + const preparedHash = await deployService.ensureDeploymentPrepared(1); + + expect(preparedHash).toEqual(activeHash); + expect(mockWrenAIAdaptor.getDeployStatus).toHaveBeenCalledWith(activeHash, 1); + expect(mockWrenAIAdaptor.deploy).toHaveBeenCalledWith({ + manifest, + hash: activeHash, + projectId: 1, + }); + }); + + it('should persist the current hash when ai-service already prepared a stale saved deployment', async () => { + const manifest = { key: 'value' }; + const activeHash = deployService.createMDLHash(manifest, 1); + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue({ + id: 123, + hash: 'legacy-saved-hash', + manifest, + }); + mockWrenAIAdaptor.getDeployStatus.mockResolvedValue('FINISHED'); + + const preparedHash = await deployService.ensureDeploymentPrepared(1); + + expect(preparedHash).toEqual(activeHash); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + manifest, + hash: activeHash, + status: DeployStatusEnum.SUCCESS, + error: undefined, + }); + expect(mockWrenAIAdaptor.deploy).not.toHaveBeenCalled(); + }); + + it('should create the same deployment hash for equivalent manifests', () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + { + name: 'customers', + columns: [{ name: 'id' }, { name: 'name' }], + }, + ], + }; + const reorderedManifest = { + models: [ + { + columns: [{ name: 'name' }, { name: 'id' }], + name: 'customers', + }, + { + columns: [{ name: 'amount' }, { name: 'id' }], + name: 'orders', + }, + ], + }; + + expect(deployService.createMDLHash(manifest, 1)).toEqual( + deployService.createMDLHash(reorderedManifest, 1), + ); + }); + + it('should treat equivalent deployed manifests as the same deployment', () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + { + name: 'customers', + columns: [{ name: 'id' }, { name: 'name' }], + }, + ], + }; + const reorderedManifest = { + models: [ + { + columns: [{ name: 'name' }, { name: 'id' }], + name: 'customers', + }, + { + columns: [{ name: 'amount' }, { name: 'id' }], + name: 'orders', + }, + ], + }; + + expect( + deployService.isSameDeployment(manifest, 1, { + hash: 'different-hash-version', + manifest: reorderedManifest, + }), + ).toBe(true); + }); + + it('should not treat changed deployed manifests as the same deployment', () => { + const manifest = { + models: [{ name: 'orders', columns: [{ name: 'id' }] }], + }; + const changedManifest = { + models: [{ name: 'orders', columns: [{ name: 'id' }, { name: 'amount' }] }], + }; + + expect( + deployService.isSameDeployment(manifest, 1, { + hash: 'different-hash-version', + manifest: changedManifest, + }), + ).toBe(false); + }); + + it('should return manifest object by hash for runtime SQL validation', async () => { + const manifest = { + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'ID' }, { name: 'amount' }], + }, + ], + }; + mockDeployLogRepository.findOneBy.mockResolvedValue({ manifest }); + + const result = await deployService.getManifestByHash('deploy-hash'); + + expect(mockDeployLogRepository.findOneBy).toHaveBeenCalledWith({ + hash: 'deploy-hash', + }); + expect(result).toEqual({ + models: [ + { + name: 'orders', + columns: [{ name: 'id' }, { name: 'amount' }], + }, + ], + }); + }); + + it('should keep getMDLByHash as base64 for the public getMDL API', async () => { + const manifest = { + models: [{ name: 'orders', columns: [{ name: 'id' }] }], + }; + mockDeployLogRepository.findOneBy.mockResolvedValue({ manifest }); + + const result = await deployService.getMDLByHash('deploy-hash'); + + expect(result).not.toBeNull(); + expect(Buffer.from(result as string, 'base64').toString()).toEqual( + JSON.stringify(manifest), + ); + }); + + it('should clear stale in-progress deployments', async () => { + const oldDate = new Date(Date.now() - 11 * 60 * 1000); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ + id: 122, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: oldDate, + }); + + const deployment = await deployService.getInProgressDeployment(1); + + expect(deployment).toBeNull(); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(122, { + status: DeployStatusEnum.FAILED, + error: 'Deployment timed out before completion.', + }); + }); + + it('should clear previous in-progress deployment before creating a new one', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue({ + id: 122, + status: DeployStatusEnum.IN_PROGRESS, + updatedAt: new Date(), + }); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockResolvedValue({ status: 'SUCCESS' }); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.SUCCESS); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(122, { + status: DeployStatusEnum.FAILED, + error: 'Deployment was superseded by a new deployment.', + }); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.SUCCESS, + error: undefined, + }); + }); + + it('should mark created deployment failed when deployment throws', async () => { + const manifest = { key: 'value' }; + const projectId = 1; + + mockDeployLogRepository.findLastProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.findInProgressProjectDeployLog.mockResolvedValue(null); + mockDeployLogRepository.createOne.mockResolvedValue({ id: 123 }); + mockWrenAIAdaptor.deploy.mockRejectedValue(new Error('network error')); + + const response = await deployService.deploy(manifest, projectId); + + expect(response.status).toEqual(DeployStatusEnum.FAILED); + expect(mockDeployLogRepository.updateOne).toHaveBeenCalledWith(123, { + status: DeployStatusEnum.FAILED, + error: 'network error', + }); + }); }); diff --git a/wren-ui/src/apollo/server/services/tests/organizationService.test.ts b/wren-ui/src/apollo/server/services/tests/organizationService.test.ts new file mode 100644 index 0000000000..41830424d9 --- /dev/null +++ b/wren-ui/src/apollo/server/services/tests/organizationService.test.ts @@ -0,0 +1,106 @@ +import { + OrganizationService, + normalizeOrganizationIdentifier, + validateOrganizationPayload, +} from '../organizationService'; + +describe('OrganizationService', () => { + const mockOrganizationRepository = () => ({ + findAll: jest.fn(), + findOneBy: jest.fn(), + createOne: jest.fn(), + updateOne: jest.fn(), + transaction: jest.fn(), + commit: jest.fn(), + rollback: jest.fn(), + getCurrentOrganization: jest.fn(), + setCurrentOrganization: jest.fn(), + }); + + describe('validateOrganizationPayload', () => { + it('should normalize valid payloads', () => { + expect( + validateOrganizationPayload({ + name: ' My Org ', + identifier: 'My Org', + description: ' Team workspace ', + }), + ).toEqual({ + name: 'My Org', + identifier: 'my_org', + description: 'Team workspace', + }); + }); + + it('should reject invalid identifiers', () => { + expect(() => + validateOrganizationPayload({ + name: 'Org', + identifier: '***', + }), + ).toThrow('Organization identifier is required'); + }); + }); + + describe('createOrganization', () => { + it('should make the first organization current', async () => { + const repository = mockOrganizationRepository(); + const tx = {}; + repository.transaction.mockResolvedValue(tx); + repository.getCurrentOrganization.mockResolvedValue(null); + repository.findOneBy.mockResolvedValue(null); + repository.createOne.mockResolvedValue({ + id: 1, + name: 'My Org', + identifier: 'my_org', + description: '', + isCurrent: true, + }); + repository.setCurrentOrganization.mockResolvedValue({ + id: 1, + name: 'My Org', + identifier: 'my_org', + description: '', + isCurrent: true, + }); + + const service = new OrganizationService({ + organizationRepository: repository as any, + }); + + const result = await service.createOrganization({ name: 'My Org' }); + + expect(repository.createOne).toHaveBeenCalledWith( + { + name: 'My Org', + identifier: 'my_org', + description: '', + isCurrent: true, + }, + { tx }, + ); + expect(repository.setCurrentOrganization).toHaveBeenCalledWith(1, { tx }); + expect(result.identifier).toBe('my_org'); + }); + + it('should reject duplicate identifiers', async () => { + const repository = mockOrganizationRepository(); + repository.findOneBy.mockResolvedValue({ id: 1, identifier: 'my_org' }); + const service = new OrganizationService({ + organizationRepository: repository as any, + }); + + await expect( + service.createOrganization({ name: 'My Org', identifier: 'my_org' }), + ).rejects.toThrow('Organization identifier already exists'); + }); + }); + + describe('normalizeOrganizationIdentifier', () => { + it('should collapse separators into underscores', () => { + expect(normalizeOrganizationIdentifier('My cool-org!!')).toBe( + 'my_cool_org', + ); + }); + }); +}); diff --git a/wren-ui/src/apollo/server/services/tests/queryService.test.ts b/wren-ui/src/apollo/server/services/tests/queryService.test.ts index 011b56176b..736a2b0f01 100644 --- a/wren-ui/src/apollo/server/services/tests/queryService.test.ts +++ b/wren-ui/src/apollo/server/services/tests/queryService.test.ts @@ -3,6 +3,11 @@ import { DataSourceName } from '../../types'; import { QueryService } from '../queryService'; describe('QueryService', () => { + const sql = 'selectable statement'; + const dataSource = DataSourceName.POSTGRES; + const project = { type: dataSource, connectionInfo: {} }; + const manifest = {}; + let mockIbisAdaptor; let mockWrenEngineAdaptor; let mockTelemetry; @@ -28,138 +33,241 @@ describe('QueryService', () => { jest.clearAllMocks(); }); - it('should return true and send event when previewing via ibis dry run succeeds', async () => { + it('passes dry-run requests to ibis and records success telemetry', async () => { mockIbisAdaptor.dryRun.mockResolvedValue({ - correlationId: '123', - processTime: '1s', + correlationId: 'correlation-id', + processTime: 'process-time', }); - const res = await queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, + const res: any = await queryService.preview(sql, { + project, + manifest, dryRun: true, }); - expect(res).toEqual({ correlationId: '123' }); - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_DRY_RUN, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, + expect(res).toEqual({ correlationId: 'correlation-id' }); + expect(mockIbisAdaptor.dryRun).toHaveBeenCalledWith( + sql, + expect.objectContaining({ + dataSource, + connectionInfo: project.connectionInfo, + mdl: manifest, + }), + ); + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_DRY_RUN, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, + }, + actionSuccess: true, }, - actionSuccess: true, - }); + ]); }); - it('should send event when previewing via ibis dry run fails', async () => { - mockIbisAdaptor.dryRun.mockRejectedValue({ - message: 'Error message', + it('records dry-run failure telemetry and rethrows the adaptor error', async () => { + const error = { + message: 'adaptor failure', extensions: { other: { - correlationId: '123', - processTime: '1s', + correlationId: 'correlation-id', + processTime: 'process-time', }, }, - }); + }; + mockIbisAdaptor.dryRun.mockRejectedValue(error); - try { - await queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, + await expect( + queryService.preview(sql, { + project, + manifest, dryRun: true, - }); - } catch (e) { - expect(e.message).toEqual('Error message'); - expect(e.extensions.other.correlationId).toEqual('123'); - expect(e.extensions.other.processTime).toEqual('1s'); - } - - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_DRY_RUN, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, - error: 'Error message', + }), + ).rejects.toMatchObject(error); + + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_DRY_RUN, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, + error: 'adaptor failure', + }, + actionSuccess: false, + service: undefined, }, - actionSuccess: false, - service: undefined, - }); + ]); }); - it('should return data and send event when previewing via ibis query succeeds', async () => { + it('passes query requests to ibis, transforms column metadata, and records success telemetry', async () => { mockIbisAdaptor.query.mockResolvedValue({ - data: [], - columns: [], - dtypes: [], - correlationId: '123', - processTime: '1s', + data: [['value']], + columns: ['field'], + dtypes: { field: 'object' }, + correlationId: 'correlation-id', + processTime: 'process-time', }); - const res = await queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, - limit: 10, + const res = await queryService.preview(sql, { + project, + manifest, + limit: 1, }); - expect(res.data).toEqual([]); - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_QUERY, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, + expect(res).toEqual({ + columns: [{ name: 'field', type: 'string' }], + data: [['value']], + correlationId: 'correlation-id', + cacheHit: false, + cacheCreatedAt: undefined, + cacheOverrodeAt: undefined, + override: false, + }); + expect(mockIbisAdaptor.query).toHaveBeenCalledWith( + sql, + expect.objectContaining({ + dataSource, + connectionInfo: project.connectionInfo, + mdl: manifest, + limit: 1, + }), + ); + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, + }, + actionSuccess: true, }, - actionSuccess: true, + ]); + }); + + it('retries transient MSSQL deadlock query failures', async () => { + const mssqlProject = { + type: DataSourceName.MSSQL, + connectionInfo: {}, + }; + mockIbisAdaptor.query + .mockRejectedValueOnce({ + message: + "[SQL Server]Transaction was deadlocked on lock resources and has been chosen as the deadlock victim. Rerun the transaction. (1205) (SQLExecDirectW)", + }) + .mockResolvedValueOnce({ + data: [['value']], + columns: ['field'], + dtypes: { field: 'object' }, + correlationId: 'correlation-id', + processTime: 'process-time', + }); + + const res = await queryService.preview(sql, { + project: mssqlProject, + manifest, + limit: 1, }); + + expect(mockIbisAdaptor.query).toHaveBeenCalledTimes(2); + expect(res).toMatchObject({ + columns: [{ name: 'field', type: 'string' }], + data: [['value']], + correlationId: 'correlation-id', + }); + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource: DataSourceName.MSSQL, + }, + actionSuccess: true, + }, + ]); }); - it('should send event when previewing via ibis query fails', async () => { - mockIbisAdaptor.query.mockRejectedValue({ - message: 'Error message', + it('does not retry non-deadlock query failures', async () => { + const mssqlProject = { + type: DataSourceName.MSSQL, + connectionInfo: {}, + }; + const error = { + message: 'syntax error near from', extensions: { other: { - correlationId: '123', - processTime: '1s', + correlationId: 'correlation-id', + processTime: 'process-time', }, }, - }); + }; + mockIbisAdaptor.query.mockRejectedValue(error); await expect( - queryService.preview('SELECT * FROM test', { - project: { type: DataSourceName.POSTGRES, connectionInfo: {} }, - manifest: {}, + queryService.preview(sql, { + project: mssqlProject, + manifest, }), - ).rejects.toMatchObject({ - message: 'Error message', + ).rejects.toMatchObject(error); + + expect(mockIbisAdaptor.query).toHaveBeenCalledTimes(1); + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource: DataSourceName.MSSQL, + error: 'syntax error near from', + }, + actionSuccess: false, + service: undefined, + }, + ]); + }); + + it('records query failure telemetry and rethrows the adaptor error', async () => { + const error = { + message: 'adaptor failure', extensions: { other: { - correlationId: '123', - processTime: '1s', + correlationId: 'correlation-id', + processTime: 'process-time', }, }, - }); + }; + mockIbisAdaptor.query.mockRejectedValue(error); - expect(mockTelemetry.records).toHaveLength(1); - expect(mockTelemetry.records[0]).toEqual({ - event: TelemetryEvent.IBIS_QUERY, - properties: { - correlationId: '123', - processTime: '1s', - sql: 'SELECT * FROM test', - dataSource: DataSourceName.POSTGRES, - error: 'Error message', + await expect( + queryService.preview(sql, { + project, + manifest, + }), + ).rejects.toMatchObject(error); + + expect(mockTelemetry.records).toEqual([ + { + event: TelemetryEvent.IBIS_QUERY, + properties: { + correlationId: 'correlation-id', + processTime: 'process-time', + sql, + dataSource, + error: 'adaptor failure', + }, + actionSuccess: false, + service: undefined, }, - actionSuccess: false, - service: undefined, - }); + ]); }); }); @@ -168,7 +276,7 @@ class MockTelemetry { sendEvent( event: TelemetryEvent, properties: Record = {}, - service: any, + service: any = undefined, actionSuccess: boolean = true, ) { this.records.push({ event, properties, service, actionSuccess }); diff --git a/wren-ui/src/apollo/server/types/context.ts b/wren-ui/src/apollo/server/types/context.ts index a037ec75ee..4aed639176 100644 --- a/wren-ui/src/apollo/server/types/context.ts +++ b/wren-ui/src/apollo/server/types/context.ts @@ -20,6 +20,10 @@ import { IInstructionRepository, IApiHistoryRepository, IDashboardItemRefreshJobRepository, + IRoleRepository, + IUserRepository, + IUserRoleRepository, + IOrganizationRepository, } from '@server/repositories'; import { IQueryService, @@ -30,6 +34,8 @@ import { IProjectService, IDashboardService, IInstructionService, + IRbacService, + IOrganizationService, } from '@server/services'; import { ITelemetry } from '@server/telemetry/telemetry'; import { @@ -59,6 +65,8 @@ export interface IContext { dashboardService: IDashboardService; sqlPairService: ISqlPairService; instructionService: IInstructionService; + rbacService: IRbacService; + organizationService: IOrganizationService; // repository projectRepository: IProjectRepository; @@ -76,6 +84,10 @@ export interface IContext { instructionRepository: IInstructionRepository; apiHistoryRepository: IApiHistoryRepository; dashboardItemRefreshJobRepository: IDashboardItemRefreshJobRepository; + roleRepository: IRoleRepository; + userRepository: IUserRepository; + userRoleRepository: IUserRoleRepository; + organizationRepository: IOrganizationRepository; // background trackers projectRecommendQuestionBackgroundTracker: ProjectRecommendQuestionBackgroundTracker; diff --git a/wren-ui/src/apollo/server/types/diagram.ts b/wren-ui/src/apollo/server/types/diagram.ts index 22f52a82a1..ca91c9aa2f 100644 --- a/wren-ui/src/apollo/server/types/diagram.ts +++ b/wren-ui/src/apollo/server/types/diagram.ts @@ -68,7 +68,7 @@ export interface DiagramModelField { description: string; isPrimaryKey?: boolean; expression?: string; - lineage?: string; + lineage?: number[]; aggregation?: string; nestedFields?: DiagramModelNestedField[]; } diff --git a/wren-ui/src/apollo/server/types/manifest.ts b/wren-ui/src/apollo/server/types/manifest.ts index e96a55d600..50320b644b 100644 --- a/wren-ui/src/apollo/server/types/manifest.ts +++ b/wren-ui/src/apollo/server/types/manifest.ts @@ -31,6 +31,7 @@ export interface Measure { export interface CumulativeMetricProperties { description?: string; + sourceColumnName?: string; } export interface Window { diff --git a/wren-ui/src/apollo/server/types/relationship.ts b/wren-ui/src/apollo/server/types/relationship.ts index c57103e443..5ab216168d 100644 --- a/wren-ui/src/apollo/server/types/relationship.ts +++ b/wren-ui/src/apollo/server/types/relationship.ts @@ -1,8 +1,12 @@ export interface RelationData { fromModelId: number; + fromModelReferenceName?: string; fromColumnId: number; + fromColumnReferenceName?: string; toModelId: number; + toModelReferenceName?: string; toColumnId: number; + toColumnReferenceName?: string; type: RelationType; description?: string; } @@ -11,6 +15,15 @@ export interface UpdateRelationData { type: RelationType; } +export interface ModelingRelationshipData { + fromModel: string; + fromColumn: string; + toModel: string; + toColumn: string; + type: RelationType; + description?: string; +} + export interface AnalysisRelationInfo { name: string; fromModelId: number; diff --git a/wren-ui/src/apollo/server/utils/apiUtils.ts b/wren-ui/src/apollo/server/utils/apiUtils.ts index a30b1ab372..b915b68c25 100644 --- a/wren-ui/src/apollo/server/utils/apiUtils.ts +++ b/wren-ui/src/apollo/server/utils/apiUtils.ts @@ -1,8 +1,11 @@ import { NextApiResponse } from 'next'; import { v4 as uuidv4 } from 'uuid'; -import { ApiType, ApiHistory } from '@server/repositories/apiHistoryRepository'; +import { + ApiType, + ApiHistory, + shouldPersistApiHistory, +} from '@server/repositories/apiHistoryRepository'; import * as Errors from '@server/utils/error'; -import { components } from '@/common'; import { AskResult, AskResultStatus, @@ -12,7 +15,34 @@ import { TextBasedAnswerStatus, } from '@/apollo/server/models/adaptor'; -const { apiHistoryRepository } = components; +const getComponentGraph = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph; +}; + +const getApiHistoryRepository = () => { + const repository = getComponentGraph().apiHistoryRepository; + if (!repository) { + throw new Error('API history repository is not initialized'); + } + return repository; +}; + +const persistApiHistory = async (payload: Record) => { + if (!shouldPersistApiHistory(payload.apiType)) { + return; + } + + try { + await getApiHistoryRepository().createOne(payload); + } catch (error) { + console.error('Failed to persist API history:', error); + } +}; export const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes @@ -195,7 +225,7 @@ export const respondWith = async ({ }) => { const durationMs = startTime ? Date.now() - startTime : undefined; const responseId = uuidv4(); - await apiHistoryRepository.createOne({ + await persistApiHistory({ id: responseId, projectId, apiType, @@ -238,7 +268,7 @@ export const respondWithSimple = async ({ }) => { const durationMs = startTime ? Date.now() - startTime : undefined; const responseId = uuidv4(); - await apiHistoryRepository.createOne({ + await persistApiHistory({ id: responseId, projectId, apiType, diff --git a/wren-ui/src/apollo/server/utils/error.ts b/wren-ui/src/apollo/server/utils/error.ts index 4bf793c49e..9108bd21f9 100644 --- a/wren-ui/src/apollo/server/utils/error.ts +++ b/wren-ui/src/apollo/server/utils/error.ts @@ -45,6 +45,7 @@ export enum GeneralErrorCodes { WREN_ENGINE_ERROR = 'WREN_ENGINE_ERROR', // asking task error + ASK_RESULT_NOT_FOUND = 'ASK_RESULT_NOT_FOUND', // when rerun from cancelled, the task is identified as general or misleading query IDENTIED_AS_GENERAL = 'IDENTIED_AS_GENERAL', IDENTIED_AS_MISLEADING_QUERY = 'IDENTIED_AS_MISLEADING_QUERY', @@ -109,6 +110,8 @@ export const errorMessages = { 'Invalid SQL, please check your SQL syntax', // asking task error + [GeneralErrorCodes.ASK_RESULT_NOT_FOUND]: + 'The asking task result is no longer available', [GeneralErrorCodes.IDENTIED_AS_GENERAL]: 'The question is identified as a general question, please follow-up ask with more specific questions.', [GeneralErrorCodes.IDENTIED_AS_MISLEADING_QUERY]: @@ -150,6 +153,7 @@ export const shortMessages = { [GeneralErrorCodes.GENERATE_QUESTIONS_ERROR]: 'Generate questions error', [GeneralErrorCodes.INVALID_SQL_ERROR]: 'Invalid SQL, please check your SQL syntax', + [GeneralErrorCodes.ASK_RESULT_NOT_FOUND]: 'Ask result not found', [GeneralErrorCodes.IDENTIED_AS_GENERAL]: 'Identified as general question', [GeneralErrorCodes.IDENTIED_AS_MISLEADING_QUERY]: 'Identified as misleading query', diff --git a/wren-ui/src/apollo/server/utils/index.ts b/wren-ui/src/apollo/server/utils/index.ts index f9628a9dc3..afce703da7 100644 --- a/wren-ui/src/apollo/server/utils/index.ts +++ b/wren-ui/src/apollo/server/utils/index.ts @@ -4,6 +4,7 @@ export * from './encode'; export * from './string'; export * from './docker'; export * from './model'; +export * from './manifest'; export * from './helper'; export * from './regex'; export * from './sseTypes'; diff --git a/wren-ui/src/apollo/server/utils/knex.ts b/wren-ui/src/apollo/server/utils/knex.ts index b7c74bba53..f5b2c8d74c 100644 --- a/wren-ui/src/apollo/server/utils/knex.ts +++ b/wren-ui/src/apollo/server/utils/knex.ts @@ -2,11 +2,88 @@ interface KnexOptions { dbType: string; pgUrl?: string; debug?: boolean; + mssqlUrl?: string; + mssqlHost?: string; + mssqlPort?: number; + mssqlDatabase?: string; + mssqlUser?: string; + mssqlPassword?: string; + mssqlEncrypt?: boolean; + mssqlTrustServerCertificate?: boolean; sqliteFile?: string; } +const normalizeDbType = (dbType?: string) => + (dbType || 'mssql').trim().toLowerCase().replace(/[-_ ]/g, ''); + +const ensureMssqlConfig = (options: KnexOptions) => { + if (options.mssqlUrl) { + return; + } + + const missingFields = [ + ['MSSQL_HOST', options.mssqlHost], + ['MSSQL_DATABASE', options.mssqlDatabase], + ['MSSQL_USER', options.mssqlUser], + ['MSSQL_PASSWORD', options.mssqlPassword], + ].filter(([, value]) => !value); + + if (missingFields.length > 0) { + throw new Error( + `MSSQL is the required Wren UI application database. Missing configuration: ${missingFields + .map(([key]) => key) + .join(', ')}`, + ); + } +}; + +const parseBooleanUrlParam = ( + searchParams: URLSearchParams, + key: string, + fallback: boolean, +) => { + const value = searchParams.get(key); + if (value === null) return fallback; + return value.toLowerCase() === 'true'; +}; + +const getMssqlConnection = (options: KnexOptions) => { + if (options.mssqlUrl) { + const url = new URL(options.mssqlUrl); + return { + server: url.hostname, + port: url.port ? parseInt(url.port) : 1433, + database: decodeURIComponent(url.pathname.replace(/^\//, '')), + user: decodeURIComponent(url.username), + password: decodeURIComponent(url.password), + options: { + encrypt: parseBooleanUrlParam(url.searchParams, 'encrypt', false), + trustServerCertificate: parseBooleanUrlParam( + url.searchParams, + 'trustServerCertificate', + true, + ), + }, + }; + } + + return { + server: options.mssqlHost, + port: options.mssqlPort || 1433, + database: options.mssqlDatabase, + user: options.mssqlUser, + password: options.mssqlPassword, + options: { + encrypt: options.mssqlEncrypt ?? false, + trustServerCertificate: options.mssqlTrustServerCertificate ?? true, + }, + }; +}; + export const bootstrapKnex = (options: KnexOptions) => { - if (options.dbType === 'pg') { + const dbType = normalizeDbType(options.dbType); + + if (dbType === 'pg' || dbType === 'postgres' || dbType === 'postgresql') { const { pgUrl, debug } = options; console.log('using pg'); /* eslint-disable @typescript-eslint/no-var-requires */ @@ -16,15 +93,21 @@ export const bootstrapKnex = (options: KnexOptions) => { debug, pool: { min: 2, max: 10 }, }); - } else { - console.log('using sqlite'); + } + + if (dbType === 'mssql' || dbType === 'sqlserver') { + console.log('using mssql'); + ensureMssqlConfig(options); /* eslint-disable @typescript-eslint/no-var-requires */ return require('knex')({ - client: 'better-sqlite3', - connection: { - filename: options.sqliteFile, - }, - useNullAsDefault: true, + client: 'mssql', + connection: getMssqlConnection(options), + debug: options.debug, + pool: { min: 2, max: 10 }, }); } + + throw new Error( + `Unsupported DB_TYPE "${options.dbType}". Wren UI application storage now requires MSSQL.`, + ); }; diff --git a/wren-ui/src/apollo/server/utils/logger.ts b/wren-ui/src/apollo/server/utils/logger.ts index f19e429efb..7218acf7ee 100644 --- a/wren-ui/src/apollo/server/utils/logger.ts +++ b/wren-ui/src/apollo/server/utils/logger.ts @@ -1 +1,47 @@ -export { getLogger } from 'log4js'; +import { + getLogger as getLog4jsLogger, + type Logger, + type LoggingEvent, +} from 'log4js'; + +const DEFAULT_LOG_LEVEL = process.env.LOG_LEVEL || 'info'; +const DEBUG_ENABLED = DEFAULT_LOG_LEVEL.toLowerCase() === 'debug'; +const wrappedLoggers = new WeakMap(); + +const normalizeLevel = (level: unknown) => { + const requestedLevel = + typeof level === 'string' ? level.toLowerCase() : String(level); + + if (requestedLevel === 'debug' && !DEBUG_ENABLED) { + return DEFAULT_LOG_LEVEL; + } + + return level; +}; + +export const getLogger = (category?: string): Logger => { + const logger = getLog4jsLogger(category); + const cachedLogger = wrappedLoggers.get(logger); + + if (cachedLogger) { + return cachedLogger; + } + + logger.level = DEFAULT_LOG_LEVEL; + + const wrappedLogger = new Proxy(logger, { + set(target, property, value, receiver) { + if (property === 'level') { + target.level = normalizeLevel(value) as string; + return true; + } + + return Reflect.set(target, property, value, receiver); + }, + }); + + wrappedLoggers.set(logger, wrappedLogger); + return wrappedLogger; +}; + +export type { Logger, LoggingEvent }; diff --git a/wren-ui/src/apollo/server/utils/manifest.ts b/wren-ui/src/apollo/server/utils/manifest.ts new file mode 100644 index 0000000000..7d41415b1a --- /dev/null +++ b/wren-ui/src/apollo/server/utils/manifest.ts @@ -0,0 +1,42 @@ +import type { Manifest } from '@server/mdl/type'; + +const normalizeColumns = (columns?: T[]) => { + if (!Array.isArray(columns)) { + return columns; + } + + const seen = new Set(); + return columns.filter((column) => { + const name = column?.name; + if (!name) { + return true; + } + + const normalizedName = name.toLowerCase(); + if (seen.has(normalizedName)) { + return false; + } + + seen.add(normalizedName); + return true; + }); +}; + +export function normalizeManifest(manifest: Manifest): Manifest; +export function normalizeManifest(manifest: undefined): undefined; +export function normalizeManifest(manifest?: Manifest): Manifest | undefined { + if (!manifest?.models) { + return manifest; + } + + return { + ...manifest, + models: manifest.models.map((model) => ({ + ...model, + columns: normalizeColumns(model.columns), + })), + }; +} + +export const encodeManifest = (manifest: Manifest): string => + Buffer.from(JSON.stringify(normalizeManifest(manifest))).toString('base64'); diff --git a/wren-ui/src/apollo/server/utils/model.ts b/wren-ui/src/apollo/server/utils/model.ts index 05e79cb2a3..0a5fb52806 100644 --- a/wren-ui/src/apollo/server/utils/model.ts +++ b/wren-ui/src/apollo/server/utils/model.ts @@ -8,7 +8,15 @@ import { CompactColumn } from '@server/services/metadataService'; export function getPreviewColumnsStr(modelColumns: ModelColumn[]) { if (modelColumns.length === 0) return '*'; - const columns = modelColumns.map((column) => `"${column.referenceName}"`); + const usedReferenceNames = new Set(); + const columns = modelColumns + .filter((column) => { + const referenceName = column.referenceName.toLowerCase(); + if (usedReferenceNames.has(referenceName)) return false; + usedReferenceNames.add(referenceName); + return true; + }) + .map((column) => `"${column.referenceName}"`); return columns.join(','); } @@ -22,6 +30,33 @@ export function transformInvalidColumnName(columnName: string) { return referenceName; } +export function getUniqueReferenceName( + referenceName: string, + usedReferenceNames: Set, +) { + const baseName = referenceName || 'column'; + let uniqueName = baseName; + let suffix = 2; + + while (usedReferenceNames.has(uniqueName.toLowerCase())) { + uniqueName = `${baseName}_${suffix}`; + suffix += 1; + } + + usedReferenceNames.add(uniqueName.toLowerCase()); + return uniqueName; +} + +export function transformUniqueInvalidColumnName( + columnName: string, + usedReferenceNames: Set, +) { + return getUniqueReferenceName( + transformInvalidColumnName(columnName), + usedReferenceNames, + ); +} + export function replaceInvalidReferenceName(referenceName: string) { // replace dot with underscore return referenceName.replace(/\./g, '_'); diff --git a/wren-ui/src/apollo/server/utils/tests/manifest.test.ts b/wren-ui/src/apollo/server/utils/tests/manifest.test.ts new file mode 100644 index 0000000000..cea03fc57f --- /dev/null +++ b/wren-ui/src/apollo/server/utils/tests/manifest.test.ts @@ -0,0 +1,39 @@ +import { normalizeManifest } from '../manifest'; + +describe('normalizeManifest', () => { + it('removes duplicate model columns case-insensitively', () => { + const duplicateName = 'field_alpha'; + const duplicateNameWithDifferentCase = 'Field_Alpha'; + const retainedName = 'field_beta'; + const manifest = normalizeManifest({ + models: [ + { + name: 'model_fixture', + cached: false, + columns: [ + { + name: duplicateName, + type: 'STRING', + isCalculated: false, + }, + { + name: duplicateNameWithDifferentCase, + type: 'STRING', + isCalculated: false, + }, + { + name: retainedName, + type: 'DATE', + isCalculated: false, + }, + ], + }, + ], + })!; + + expect(manifest.models[0].columns.map((column) => column.name)).toEqual([ + duplicateName, + retainedName, + ]); + }); +}); diff --git a/wren-ui/src/apollo/server/utils/tests/model.test.ts b/wren-ui/src/apollo/server/utils/tests/model.test.ts new file mode 100644 index 0000000000..b55f8f144f --- /dev/null +++ b/wren-ui/src/apollo/server/utils/tests/model.test.ts @@ -0,0 +1,30 @@ +import { ModelColumn } from '@server/repositories'; +import { getPreviewColumnsStr } from '../model'; + +describe('model utils', () => { + describe('getPreviewColumnsStr', () => { + it('returns wildcard when no model columns exist', () => { + expect(getPreviewColumnsStr([])).toBe('*'); + }); + + it('quotes unique model reference names for preview SQL', () => { + const columns = [ + { referenceName: 'id' }, + { referenceName: 'status' }, + ] as ModelColumn[]; + + expect(getPreviewColumnsStr(columns)).toBe('"id","status"'); + }); + + it('deduplicates model reference names case-insensitively', () => { + const columns = [ + { referenceName: 'status' }, + { referenceName: 'Status' }, + { referenceName: 'bunit' }, + { referenceName: 'BUNIT' }, + ] as ModelColumn[]; + + expect(getPreviewColumnsStr(columns)).toBe('"status","bunit"'); + }); + }); +}); diff --git a/wren-ui/src/common.ts b/wren-ui/src/common.ts index 6ad7384bf8..62fd2db038 100644 --- a/wren-ui/src/common.ts +++ b/wren-ui/src/common.ts @@ -19,6 +19,14 @@ import { InstructionRepository, ApiHistoryRepository, DashboardItemRefreshJobRepository, + RoleRepository, + UserRepository, + UserRoleRepository, + OrganizationRepository, + OrganizationMemberRepository, + OrganizationMemberProjectRepository, + OrganizationInvitationRepository, + OrganizationInvitationProjectRepository, } from '@server/repositories'; import { WrenEngineAdaptor, @@ -35,23 +43,101 @@ import { DashboardService, AskingTaskTracker, InstructionService, + RbacService, + OrganizationService, + OrganizationMemberService, } from '@server/services'; import { PostHogTelemetry } from './apollo/server/telemetry/telemetry'; import { ProjectRecommendQuestionBackgroundTracker, ThreadRecommendQuestionBackgroundTracker, DashboardCacheBackgroundTracker, + ChartBackgroundTracker, + ChartAdjustmentBackgroundTracker, } from './apollo/server/backgrounds'; import { SqlPairService } from './apollo/server/services/sqlPairService'; export const serverConfig = getConfig(); +type Initializable = { + initialize?: unknown; +}; + +type Disposable = { + dispose?: unknown; + stop?: unknown; +}; + +type ReusableComponentGraph = { + askingTaskTracker?: Initializable & Disposable; + askingService?: Initializable & Disposable; + projectService?: Disposable; + projectRecommendQuestionBackgroundTracker?: Initializable & Disposable; + threadRecommendQuestionBackgroundTracker?: Initializable & Disposable; + dashboardCacheBackgroundTracker?: Disposable; + knex?: { + destroy?: () => unknown; + }; +}; + +const hasInitialize = ( + value: Initializable | null | undefined, +): value is { initialize: () => Promise | void } => { + return typeof value?.initialize === 'function'; +}; + +const isReusableComponentGraph = (graph?: ReusableComponentGraph): boolean => { + const nestedAskingTaskTracker = ( + graph?.askingService as { askingTaskTracker?: Initializable } | undefined + )?.askingTaskTracker; + + return Boolean( + graph && + hasInitialize(graph.askingTaskTracker) && + hasInitialize(graph.askingService) && + hasInitialize(nestedAskingTaskTracker) && + hasInitialize(graph.projectRecommendQuestionBackgroundTracker) && + hasInitialize(graph.threadRecommendQuestionBackgroundTracker), + ); +}; + +const disposeComponentGraph = (graph?: ReusableComponentGraph): void => { + const disposables = [ + graph?.askingService, + graph?.askingTaskTracker, + graph?.projectService, + graph?.projectRecommendQuestionBackgroundTracker, + graph?.threadRecommendQuestionBackgroundTracker, + graph?.dashboardCacheBackgroundTracker, + ]; + + disposables.forEach((disposable) => { + if (typeof disposable?.dispose === 'function') { + disposable.dispose(); + } else if (typeof disposable?.stop === 'function') { + disposable.stop(); + } + }); + + if (typeof graph?.knex?.destroy === 'function') { + void Promise.resolve(graph.knex.destroy()).catch(() => undefined); + } +}; + export const initComponents = () => { const telemetry = new PostHogTelemetry(); const knex = bootstrapKnex({ dbType: serverConfig.dbType, pgUrl: serverConfig.pgUrl, debug: serverConfig.debug, + mssqlUrl: serverConfig.mssqlUrl, + mssqlHost: serverConfig.mssqlHost, + mssqlPort: serverConfig.mssqlPort, + mssqlDatabase: serverConfig.mssqlDatabase, + mssqlUser: serverConfig.mssqlUser, + mssqlPassword: serverConfig.mssqlPassword, + mssqlEncrypt: serverConfig.mssqlEncrypt, + mssqlTrustServerCertificate: serverConfig.mssqlTrustServerCertificate, sqliteFile: serverConfig.sqliteFile, }); @@ -75,6 +161,17 @@ export const initComponents = () => { const apiHistoryRepository = new ApiHistoryRepository(knex); const dashboardItemRefreshJobRepository = new DashboardItemRefreshJobRepository(knex); + const roleRepository = new RoleRepository(knex); + const userRepository = new UserRepository(knex); + const userRoleRepository = new UserRoleRepository(knex); + const organizationRepository = new OrganizationRepository(knex); + const organizationMemberRepository = new OrganizationMemberRepository(knex); + const organizationMemberProjectRepository = + new OrganizationMemberProjectRepository(knex); + const organizationInvitationRepository = + new OrganizationInvitationRepository(knex); + const organizationInvitationProjectRepository = + new OrganizationInvitationProjectRepository(knex); // adaptors const wrenEngineAdaptor = new WrenEngineAdaptor({ @@ -87,6 +184,32 @@ export const initComponents = () => { ibisServerEndpoint: serverConfig.ibisServerEndpoint, }); + // background trackers + const projectRecommendQuestionBackgroundTracker = + new ProjectRecommendQuestionBackgroundTracker({ + telemetry, + wrenAIAdaptor, + projectRepository, + }); + const threadRecommendQuestionBackgroundTracker = + new ThreadRecommendQuestionBackgroundTracker({ + telemetry, + wrenAIAdaptor, + threadRepository, + }); + const chartBackgroundTracker = new ChartBackgroundTracker({ + telemetry, + wrenAIAdaptor, + threadResponseRepository, + }); + const chartAdjustmentBackgroundTracker = new ChartAdjustmentBackgroundTracker( + { + telemetry, + wrenAIAdaptor, + threadResponseRepository, + }, + ); + // services const metadataService = new DataSourceMetadataService({ ibisAdaptor, @@ -114,8 +237,10 @@ export const initComponents = () => { projectRepository, metadataService, mdlService, + queryService, wrenAIAdaptor, telemetry, + projectRecommendQuestionBackgroundTracker, }); const askingTaskTracker = new AskingTaskTracker({ wrenAIAdaptor, @@ -127,14 +252,17 @@ export const initComponents = () => { telemetry, wrenAIAdaptor, deployService, + mdlService, projectService, viewRepository, threadRepository, threadResponseRepository, queryService, - mdlService, askingTaskTracker, askingTaskRepository, + chartBackgroundTracker, + chartAdjustmentBackgroundTracker, + threadRecommendQuestionBackgroundTracker, }); const dashboardService = new DashboardService({ projectService, @@ -145,25 +273,31 @@ export const initComponents = () => { sqlPairRepository, wrenAIAdaptor, ibisAdaptor, + deployService, }); const instructionService = new InstructionService({ instructionRepository, wrenAIAdaptor, + deployService, }); + const rbacService = new RbacService({ + roleRepository, + userRepository, + userRoleRepository, + }); + const organizationService = new OrganizationService({ + organizationRepository, + }); + const organizationMemberService = new OrganizationMemberService( + organizationRepository, + organizationMemberRepository, + organizationMemberProjectRepository, + organizationInvitationRepository, + organizationInvitationProjectRepository, + userRepository, + projectRepository, + ); - // background trackers - const projectRecommendQuestionBackgroundTracker = - new ProjectRecommendQuestionBackgroundTracker({ - telemetry, - wrenAIAdaptor, - projectRepository, - }); - const threadRecommendQuestionBackgroundTracker = - new ThreadRecommendQuestionBackgroundTracker({ - telemetry, - wrenAIAdaptor, - threadRepository, - }); const dashboardCacheBackgroundTracker = new DashboardCacheBackgroundTracker({ dashboardRepository, dashboardItemRepository, @@ -196,6 +330,14 @@ export const initComponents = () => { apiHistoryRepository, instructionRepository, dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, + organizationRepository, + organizationMemberRepository, + organizationMemberProjectRepository, + organizationInvitationRepository, + organizationInvitationProjectRepository, // adaptors wrenEngineAdaptor, @@ -212,6 +354,9 @@ export const initComponents = () => { dashboardService, sqlPairService, instructionService, + rbacService, + organizationService, + organizationMemberService, askingTaskTracker, // background trackers @@ -221,5 +366,17 @@ export const initComponents = () => { }; }; -// singleton components -export const components = initComponents(); +declare global { + // eslint-disable-next-line no-var + var __wrenComponents: ReturnType | undefined; +} + +// Keep a single server-side component graph across Next.js dev reloads. +const existingComponents = globalThis.__wrenComponents; + +if (!isReusableComponentGraph(existingComponents)) { + disposeComponentGraph(existingComponents); + globalThis.__wrenComponents = initComponents(); +} + +export const components = globalThis.__wrenComponents; diff --git a/wren-ui/src/components/HeaderBar.tsx b/wren-ui/src/components/HeaderBar.tsx index a3a0292c49..da6c85b0a8 100644 --- a/wren-ui/src/components/HeaderBar.tsx +++ b/wren-ui/src/components/HeaderBar.tsx @@ -1,9 +1,11 @@ +import { useEffect, useMemo, useState } from 'react'; import { useRouter } from 'next/router'; -import { Button, Layout, Space } from 'antd'; +import { Button, Dropdown, Layout, Menu, Space } from 'antd'; import styled from 'styled-components'; import LogoBar from '@/components/LogoBar'; import { Path } from '@/utils/enum'; import Deploy from '@/components/deploy/Deploy'; +import OrganizationSwitcher from '@/components/OrganizationSwitcher'; const { Header } = Layout; @@ -31,11 +33,94 @@ const StyledHeader = styled(Header)` padding: 10px 16px; `; +const HeaderLeft = styled.div` + display: flex; + align-items: center; + gap: 48px; +`; + +const HeaderCenter = styled.div` + display: flex; + align-items: center; + justify-content: center; + flex: 1; + min-width: 0; +`; + +const HeaderRight = styled.div` + display: flex; + align-items: center; + justify-content: flex-end; + gap: 12px; + min-width: 120px; +`; + +const UserAvatar = styled.div` + width: 32px; + height: 32px; + border-radius: 50%; + display: flex; + align-items: center; + justify-content: center; + cursor: pointer; + background: var(--geekblue-6); + color: var(--gray-1); + font-weight: 600; + line-height: 1; +`; + +interface CurrentUserProfile { + name?: string; + email?: string; +} + export default function HeaderBar() { const router = useRouter(); const { pathname } = router; + const currentPath = router.asPath.split(/[?#]/)[0]; const showNav = !pathname.startsWith(Path.Onboarding); const isModeling = pathname.startsWith(Path.Modeling); + const navigateTo = (path: Path) => { + if (currentPath !== path) { + router.push(path); + } + }; + const [currentUser, setCurrentUser] = useState( + null, + ); + + useEffect(() => { + if (!showNav) return; + + const loadCurrentUser = async () => { + try { + const response = await fetch('/api/v1/users/current'); + if (!response.ok) return; + const payload = (await response.json()) as CurrentUserProfile; + setCurrentUser(payload); + } catch { + setCurrentUser(null); + } + }; + + loadCurrentUser(); + }, [showNav]); + + const userInitials = useMemo(() => { + const displayName = currentUser?.name || currentUser?.email || 'User'; + const parts = displayName.trim().split(/\s+/).filter(Boolean); + if (!parts.length) return 'U'; + if (parts.length === 1) return parts[0].slice(0, 2).toUpperCase(); + return `${parts[0][0]}${parts[parts.length - 1][0]}`.toUpperCase(); + }, [currentUser]); + + const userMenu = ( + + router.push(Path.UserProfile)}> + {currentUser?.email || 'Profile'} + + + ); return ( @@ -43,7 +128,7 @@ export default function HeaderBar() { className="d-flex justify-space-between align-center" style={{ marginTop: -2 }} > - + {showNav && ( @@ -51,7 +136,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.Home)} - onClick={() => router.push(Path.Home)} + onClick={() => navigateTo(Path.Home)} > Home @@ -59,7 +144,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.Modeling)} - onClick={() => router.push(Path.Modeling)} + onClick={() => navigateTo(Path.Modeling)} > Modeling @@ -67,7 +152,7 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.Knowledge)} - onClick={() => router.push(Path.KnowledgeQuestionSQLPairs)} + onClick={() => navigateTo(Path.KnowledgeQuestionSQLPairs)} > Knowledge @@ -75,18 +160,26 @@ export default function HeaderBar() { shape="round" size="small" $isHighlight={pathname.startsWith(Path.APIManagement)} - onClick={() => router.push(Path.APIManagementHistory)} + onClick={() => navigateTo(Path.APIManagementHistory)} > API )} - - {isModeling && ( - - - - )} + + {showNav && } + + {isModeling && } + {showNav && ( + + {userInitials} + + )} + ); diff --git a/wren-ui/src/components/OrganizationSwitcher.tsx b/wren-ui/src/components/OrganizationSwitcher.tsx new file mode 100644 index 0000000000..03c9340c4a --- /dev/null +++ b/wren-ui/src/components/OrganizationSwitcher.tsx @@ -0,0 +1,609 @@ +import { useEffect, useMemo, useState } from 'react'; +import { + Button, + Dropdown, + Form, + Input, + message, + Modal, + Radio, + Space, + Spin, + Tooltip, + Typography, +} from 'antd'; +import styled from 'styled-components'; +import PlusOutlined from '@ant-design/icons/PlusOutlined'; +import DownOutlined from '@ant-design/icons/DownOutlined'; +import SettingOutlined from '@ant-design/icons/SettingOutlined'; +import SearchOutlined from '@ant-design/icons/SearchOutlined'; +import { useRouter } from 'next/router'; +import { Path } from '@/utils/enum'; +import { WorkspaceProjectType } from '@/apollo/client/graphql/__types__'; +import apolloClient from '@/apollo/client'; + +interface OrganizationRecord { + id: number; + name: string; + identifier: string; + description?: string | null; + isCurrent: boolean; +} + +interface OrganizationResponse { + organizations: OrganizationRecord[]; + currentOrganization: OrganizationRecord | null; + currentProjectName: string; +} + +interface ProjectRecord { + id: string; + displayName: string; + projectType: WorkspaceProjectType; + isCurrent: boolean; + hasDataSource: boolean; +} + +interface ProjectResponse { + projects: ProjectRecord[]; + currentProject: ProjectRecord | null; +} + +const TriggerButton = styled(Button)` + display: flex; + align-items: center; + color: var(--gray-2); + border: none; + background: transparent; + padding: 0 8px; + + &:hover, + &:focus { + color: var(--gray-1); + background: rgba(255, 255, 255, 0.05); + } +`; + +const OrganizationBadge = styled.div` + width: 28px; + height: 28px; + border-radius: 6px; + background: #b7eb8f; + color: #274916; + display: flex; + align-items: center; + justify-content: center; + font-weight: 600; + text-transform: uppercase; +`; + +const Overlay = styled.div` + width: 280px; + background: var(--gray-1); + border-radius: 10px; + box-shadow: 0 10px 30px rgba(0, 0, 0, 0.16); + overflow: hidden; +`; + +const OverlayHeader = styled.div` + display: flex; + align-items: center; + justify-content: space-between; + padding: 12px 16px 8px; +`; + +const OverlayLabel = styled.div` + font-size: 12px; + letter-spacing: 0.02em; + color: var(--gray-6); + text-transform: uppercase; +`; + +const OverlayBody = styled.div` + padding: 0 8px 8px; +`; + +const OverlayFooter = styled.div` + border-top: 1px solid var(--gray-4); + padding: 8px; +`; + +const MenuRow = styled.button<{ $active?: boolean }>` + width: 100%; + display: flex; + align-items: center; + gap: 12px; + padding: 10px 16px; + border: none; + background: ${(props) => (props.$active ? 'var(--gray-3)' : 'transparent')}; + text-align: left; + cursor: pointer; + + &:hover { + background: var(--gray-3); + } +`; + +const SectionLabel = styled.div` + font-size: 12px; + letter-spacing: 0.02em; + color: var(--gray-6); + text-transform: uppercase; + margin: 12px 8px 8px; +`; + +const SearchBox = styled(Input)` + margin-bottom: 10px; +`; + +const ProjectTypeTag = styled.span` + font-size: 12px; + color: var(--gray-6); +`; + +const ProjectCard = styled.button<{ $selected?: boolean }>` + width: 100%; + text-align: left; + border: 1px solid ${(props) => (props.$selected ? '#3b59e6' : 'var(--gray-4)')}; + border-radius: 12px; + padding: 16px; + background: ${(props) => (props.$selected ? 'rgba(59, 89, 230, 0.04)' : '#fff')}; + cursor: pointer; +`; + +const ProjectOptionRow = styled.div` + display: flex; + align-items: flex-start; + justify-content: space-between; + gap: 12px; +`; + +const getBadgeText = (name?: string) => + (name || 'O') + .trim() + .split(/\s+/) + .slice(0, 1) + .map((segment) => segment[0]) + .join('') + .toUpperCase(); + +export default function OrganizationSwitcher() { + const router = useRouter(); + const [form] = Form.useForm(); + const [classicProjectForm] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [visible, setVisible] = useState(false); + const [projectModalVisible, setProjectModalVisible] = useState(false); + const [classicProjectModalVisible, setClassicProjectModalVisible] = + useState(false); + const [projectSearch, setProjectSearch] = useState(''); + const [selectedProjectType, setSelectedProjectType] = + useState(WorkspaceProjectType.AGENTIC); + const [currentProjectName, setCurrentProjectName] = useState('Default Project'); + const [organizations, setOrganizations] = useState([]); + const [currentOrganization, setCurrentOrganization] = + useState(null); + const [projects, setProjects] = useState([]); + + const loadOrganizations = async () => { + setLoading(true); + try { + const [organizationResponse, projectResponse] = await Promise.all([ + fetch('/api/v1/organizations/current'), + fetch('/api/v1/projects/current'), + ]); + const organizationPayload = (await organizationResponse.json()) as OrganizationResponse & { + error?: string; + }; + const projectPayload = (await projectResponse.json()) as ProjectResponse & { + error?: string; + }; + if (!organizationResponse.ok) { + throw new Error( + organizationPayload.error || 'Failed to load organizations', + ); + } + if (!projectResponse.ok && projectResponse.status !== 500) { + throw new Error(projectPayload.error || 'Failed to load projects'); + } + setOrganizations(organizationPayload.organizations || []); + setCurrentOrganization(organizationPayload.currentOrganization || null); + setCurrentProjectName( + projectPayload.currentProject?.displayName || + organizationPayload.currentProjectName || + 'Default Project', + ); + setProjects(projectPayload.projects || []); + } catch (error: any) { + message.error(error.message || 'Failed to load organizations'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadOrganizations(); + }, []); + + const hasOrganizations = organizations.length > 0; + const currentLabel = currentOrganization?.identifier || 'Create organization'; + const projectLabel = useMemo( + () => currentProjectName || 'Default Project', + [currentProjectName], + ); + const filteredProjects = useMemo(() => { + const keyword = projectSearch.trim().toLowerCase(); + if (!keyword) { + return projects; + } + return projects.filter((project) => + project.displayName.toLowerCase().includes(keyword), + ); + }, [projectSearch, projects]); + + const createOrganization = async () => { + try { + const values = await form.validateFields(); + setSaving(true); + const response = await fetch('/api/v1/organizations', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify(values), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to create organization'); + } + message.success('Organization created successfully.'); + setVisible(false); + form.resetFields(); + await loadOrganizations(); + await router.push(Path.OrganizationGeneral); + } catch (error: any) { + if (error?.errorFields) { + return; + } + message.error(error.message || 'Failed to create organization'); + } finally { + setSaving(false); + } + }; + + const selectOrganization = async (organizationId: number) => { + try { + const response = await fetch( + `/api/v1/organizations/${organizationId}/select`, + { + method: 'POST', + }, + ); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to switch organization'); + } + message.success('Organization switched successfully.'); + await loadOrganizations(); + } catch (error: any) { + message.error(error.message || 'Failed to switch organization'); + } + }; + + const selectProject = async (projectId: string) => { + try { + const response = await fetch(`/api/v1/projects/${projectId}/select`, { + method: 'POST', + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to switch project'); + } + message.success('Project switched successfully.'); + await apolloClient.clearStore(); + await loadOrganizations(); + await router.replace(Path.Home); + } catch (error: any) { + message.error(error.message || 'Failed to switch project'); + } + }; + + const startNewProjectFlow = async () => { + setProjectModalVisible(false); + if (selectedProjectType === WorkspaceProjectType.CLASSIC) { + classicProjectForm.resetFields(); + setClassicProjectModalVisible(true); + return; + } + await router.push({ + pathname: Path.OnboardingConnection, + query: { + newProject: '1', + projectType: selectedProjectType, + }, + }); + }; + + const startClassicProjectFlow = async () => { + try { + const values = await classicProjectForm.validateFields(); + setClassicProjectModalVisible(false); + await router.push({ + pathname: Path.OnboardingConnection, + query: { + newProject: '1', + projectType: WorkspaceProjectType.CLASSIC, + projectName: values.projectName, + }, + }); + } catch (error: any) { + if (error?.errorFields) { + return; + } + message.error(error.message || 'Failed to start project setup'); + } + }; + + const overlay = ( + + + Projects + + + {filteredProjects.map((project) => ( + void selectProject(project.id)} + > +
+
{project.displayName}
+ + {project.projectType === WorkspaceProjectType.AGENTIC + ? 'Agentic' + : 'Classic'} + +
+ {project.isCurrent && Current} + + ))} + Organizations + {organizations.map((organization) => ( +
+ void selectOrganization(organization.id)} + > + + {getBadgeText(organization.name)} + +
+
{organization.identifier}
+ + {organization.name} + +
+ +
+ ))} + + {currentOrganization && ( + + + + )} + + ); + + if (loading) { + return ; + } + + return ( + <> + {hasOrganizations ? ( + + + + {currentLabel} + / + {projectLabel} + + + + + ) : ( + + )} + + void createOrganization()} + onCancel={() => { + setVisible(false); + form.resetFields(); + }} + destroyOnClose + > +
+ + + + + + + + + + + + + setProjectModalVisible(false)}> + Cancel + , + , + ]} + onCancel={() => setProjectModalVisible(false)} + destroyOnClose + > + + setSelectedProjectType(WorkspaceProjectType.AGENTIC)} + > + +
+
+ Agentic project{' '} + + RECOMMENDED + +
+ + Agentic mode, Knowledge, Skills, Memory, and artifacts for + charts, dashboards, and reports. + +
+ + + + setSelectedProjectType(WorkspaceProjectType.CLASSIC)} + > + +
+
Classic project
+ + Dashboards, spreadsheets, and the classic BI workflow. + +
+ + + + + + + { + setClassicProjectModalVisible(false); + classicProjectForm.resetFields(); + }} + > + Cancel + , + , + ]} + onCancel={() => { + setClassicProjectModalVisible(false); + classicProjectForm.resetFields(); + }} + destroyOnClose + > +
+ + + + + + + ); +} diff --git a/wren-ui/src/components/chart/handler.ts b/wren-ui/src/components/chart/handler.ts index bb150b9878..48300930dc 100644 --- a/wren-ui/src/components/chart/handler.ts +++ b/wren-ui/src/components/chart/handler.ts @@ -185,6 +185,7 @@ export default class ChartSpecHandler { // avoid mutating the original spec const clonedSpec = cloneDeep(spec); + this.normalizeSpecFields(clonedSpec); this.parseSpec(clonedSpec); } @@ -258,6 +259,42 @@ export default class ChartSpecHandler { } } + private normalizeSpecFields(spec: TopLevelSpec) { + const values = ((this.data as any)?.values || []) as Record[]; + if (!values.length) return; + + const columns = Object.keys(values[0]); + const normalizeField = (field?: string) => { + if (!field) return field; + if (columns.includes(field)) return field; + + const lowered = columns.find( + (column) => column.toLowerCase() === field.toLowerCase(), + ); + if (lowered) return lowered; + + const compactField = field.replace(/[\s_]+/g, '').toLowerCase(); + return ( + columns.find( + (column) => + column.replace(/[\s_]+/g, '').toLowerCase() === compactField, + ) || field + ); + }; + + const encoding = (spec as any).encoding as EncodingSpec; + ['x', 'y', 'theta', 'color', 'xOffset'].forEach((key) => { + const axis = encoding?.[key] as { field?: string } | undefined; + if (axis?.field) axis.field = normalizeField(axis.field); + }); + + (spec.transform || []).forEach((transform: any) => { + if (Array.isArray(transform?.fold)) { + transform.fold = transform.fold.map(normalizeField); + } + }); + } + private addMark(mark: MarkSpec) { let additionalProps = {}; @@ -415,30 +452,31 @@ export default class ChartSpecHandler { return encoding[axis]?.title || undefined; } - private transformDataValues( - data: DataSpec, - encoding: { - x?: { type?: string; field?: string }; - y?: { type?: string; field?: string }; - }, - ) { - // If axis x is temporal - if (encoding?.x?.type === 'temporal') { - const transformedValues = data.values.map((val) => ({ - ...val, - [encoding.x.field]: this.transformTemporalValue(val[encoding.x.field]), - })); - return { ...data, values: transformedValues }; - } - // If axis y is temporal - if (encoding?.y?.type === 'temporal') { - const transformedValues = data.values.map((val) => ({ - ...val, - [encoding.y.field]: this.transformTemporalValue(val[encoding.y.field]), - })); - return { ...data, values: transformedValues }; - } - return data; + private transformDataValues(data: DataSpec, encoding?: EncodingSpec) { + const encodingKeys = ['x', 'y', 'theta', 'color', 'xOffset']; + const transformedValues = data.values.map((val) => { + const next = { ...val }; + + encodingKeys.forEach((key) => { + const axis = encoding?.[key] as + | { type?: string; field?: string } + | undefined; + if (!axis || typeof axis.field !== 'string') return; + + if (axis.type === 'temporal') { + next[axis.field] = this.transformTemporalValue(val[axis.field]); + return; + } + + if (axis.type === 'quantitative') { + next[axis.field] = this.transformQuantitativeValue(val[axis.field]); + } + }); + + return next; + }); + + return { ...data, values: transformedValues }; } private transformTemporalValue(value: string | any) { @@ -453,6 +491,18 @@ export default class ChartSpecHandler { } return strValue; } + + private transformQuantitativeValue(value: any) { + if (value === null || value === undefined || value === '') { + return value; + } + if (typeof value === 'number') { + return value; + } + + const numericValue = Number(String(value).replace(/,/g, '')); + return Number.isFinite(numericValue) ? numericValue : value; + } } export const convertToChartType = ( diff --git a/wren-ui/src/components/chart/index.tsx b/wren-ui/src/components/chart/index.tsx index dfc81a7c7e..955ca891cd 100644 --- a/wren-ui/src/components/chart/index.tsx +++ b/wren-ui/src/components/chart/index.tsx @@ -125,6 +125,7 @@ export default function Chart(props: VegaLiteProps) { }; const getChartContent = () => { + if (!values) return null; if (values.length === 0) return
No available data
; if (parsedError) { diff --git a/wren-ui/src/components/deploy/Deploy.tsx b/wren-ui/src/components/deploy/Deploy.tsx index 14560b42db..a266968711 100644 --- a/wren-ui/src/components/deploy/Deploy.tsx +++ b/wren-ui/src/components/deploy/Deploy.tsx @@ -1,4 +1,4 @@ -import { useEffect } from 'react'; +import { useEffect, useState } from 'react'; import { Button, Space, Typography, message } from 'antd'; import CheckCircleOutlined from '@ant-design/icons/CheckCircleOutlined'; import LoadingOutlined from '@ant-design/icons/LoadingOutlined'; @@ -39,20 +39,33 @@ const getDeployStatus = (deploying: boolean, status: SyncStatus) => { export default function Deploy() { const deployContext = useDeployStatusContext(); const { data, loading, startPolling, stopPolling } = deployContext; + const [deployedSuccessfully, setDeployedSuccessfully] = useState(false); const [deployMutation, { data: deployResult, loading: deploying }] = useDeployMutation({ onError: (error) => console.error(error), onCompleted: (data) => { if (data.deploy?.status === 'FAILED') { + setDeployedSuccessfully(false); console.error('Failed to deploy - ', data.deploy?.error); message.error( 'Failed to deploy. Please check the log for more details.', ); + } else if (data.deploy?.status === 'SUCCESS') { + setDeployedSuccessfully(true); + stopPolling(); } }, }); + useEffect(() => { + const resetDeploySuccess = () => setDeployedSuccessfully(false); + window.addEventListener('wren:modeling-changed', resetDeploySuccess); + return () => { + window.removeEventListener('wren:modeling-changed', resetDeploySuccess); + }; + }, []); + useEffect(() => { // Stop polling deploy status if deploy failed if ( @@ -63,10 +76,16 @@ export default function Deploy() { } }, [deployResult, data]); - const syncStatus = data?.modelSync.status; + const serverSyncStatus = data?.modelSync.status; + const syncStatus = deployedSuccessfully + ? SyncStatus.SYNCRONIZED + : serverSyncStatus; const onDeploy = () => { - deployMutation(); + setDeployedSuccessfully(false); + deployMutation({ + variables: { force: serverSyncStatus === SyncStatus.UNSYNCRONIZED }, + }); startPolling(1000); }; diff --git a/wren-ui/src/components/organization/SettingsLayout.tsx b/wren-ui/src/components/organization/SettingsLayout.tsx new file mode 100644 index 0000000000..323b6f7b8d --- /dev/null +++ b/wren-ui/src/components/organization/SettingsLayout.tsx @@ -0,0 +1,136 @@ +import { ReactNode } from 'react'; +import Link from 'next/link'; +import { Layout, Typography } from 'antd'; +import styled from 'styled-components'; +import SimpleLayout from '@/components/layouts/SimpleLayout'; +import { Path } from '@/utils/enum'; + +const { Sider, Content } = Layout; + +const linkStyle = { color: 'inherit', transition: 'none' }; + +const StyledSider = styled(Sider)` + height: calc(100vh - 48px); + background: var(--gray-2); + border-right: 1px solid var(--gray-4); + overflow-y: auto; +`; + +const StyledContent = styled(Content)` + height: calc(100vh - 48px); + overflow-y: auto; + background: white; +`; + +const SidebarSection = styled.div` + padding: 14px 16px 8px; + font-size: 12px; + font-weight: 700; + color: var(--gray-7); +`; + +const SidebarItem = styled.div<{ $active?: boolean; $disabled?: boolean }>` + padding: 6px 20px; + color: ${(props) => + props.$active + ? 'var(--gray-10)' + : props.$disabled + ? 'var(--gray-7)' + : 'var(--gray-8)'}; + background: ${(props) => (props.$active ? 'var(--gray-4)' : 'transparent')}; + font-weight: ${(props) => (props.$active ? 600 : 400)}; +`; + +const PageBody = styled.div` + padding: 24px 48px; +`; + +const PlaceholderItem = ({ children }: { children: ReactNode }) => ( + {children} +); + +export default function OrganizationSettingsLayout({ + section, + title, + titleExtra, + children, +}: { + section: + | 'project-general' + | 'project-access-control' + | 'project-danger-zone' + | 'general' + | 'members' + | 'danger-zone' + | 'user-profile' + | 'user-danger-zone'; + title: ReactNode; + titleExtra?: ReactNode; + children: ReactNode; +}) { + return ( + + + + Project + + + General + + + + + Access control + + + Data source + + + Danger zone + + + + Organization + + + General + + + + + Members + + + + + Danger zone + + + + User + + + Profile + + + + + Danger zone + + + + + +
+ + {title} + + {titleExtra} +
+ {children} + + + + + ); +} diff --git a/wren-ui/src/components/pages/administration/types.tsx b/wren-ui/src/components/pages/administration/types.tsx new file mode 100644 index 0000000000..896bf8cf40 --- /dev/null +++ b/wren-ui/src/components/pages/administration/types.tsx @@ -0,0 +1,51 @@ +import { Tag } from 'antd'; + +export interface Role { + id: number; + name: string; + description?: string | null; + createdAt: string; + updatedAt: string; + users?: User[]; +} + +export interface User { + id: number; + name: string; + email: string; + externalId?: string | null; + identityProvider?: string | null; + isActive: boolean; + roles?: Role[]; + createdAt: string; + updatedAt: string; +} + +export interface UserRoleMapping { + id: number; + userId: number; + roleId: number; + user: User; + role: Role; + createdAt: string; + updatedAt: string; +} + +export const RoleTags = ({ roles = [] }: { roles?: Role[] }) => { + if (!roles.length) return No roles; + return ( + <> + {roles.map((role) => ( + + {role.name} + + ))} + + ); +}; + +export const StatusTag = ({ active }: { active: boolean }) => ( + + {active ? 'Active' : 'Inactive'} + +); diff --git a/wren-ui/src/components/pages/home/dashboardGrid/index.tsx b/wren-ui/src/components/pages/home/dashboardGrid/index.tsx index e27b4c6b00..b320be8342 100644 --- a/wren-ui/src/components/pages/home/dashboardGrid/index.tsx +++ b/wren-ui/src/components/pages/home/dashboardGrid/index.tsx @@ -129,7 +129,7 @@ const getLayoutToGrid = (item: DashboardItem) => { const getLayoutToUpdateItem = (layout: Layout) => { return { - itemId: Number(layout.i), + itemId: layout.i, x: layout.x, y: layout.y, w: layout.w, @@ -141,7 +141,7 @@ interface Props { items: DashboardItem[]; isSupportCached: boolean; onUpdateChange: (layouts: ItemLayoutInput[]) => void; - onDelete: (id: number) => Promise; + onDelete: (id: string) => Promise; } const DashboardGrid = forwardRef( @@ -240,7 +240,7 @@ const DashboardGrid = forwardRef( export default DashboardGrid; -const PinnedItemTitle = (props: { id: number; title: string }) => { +const PinnedItemTitle = (props: { id: string; title: string }) => { const { title } = props; const [form] = Form.useForm(); @@ -248,7 +248,7 @@ const PinnedItemTitle = (props: { id: number; title: string }) => { onError: (error) => console.error(error), }); - const handleSave = (dashboardItemId: number, values: { title: string }) => { + const handleSave = (dashboardItemId: string, values: { title: string }) => { if (values.title === title) return; updateDashboardItem({ variables: { @@ -280,7 +280,7 @@ const PinnedItem = forwardRef( props: { item: DashboardItem; isSupportCached: boolean; - onDelete: (id: number) => Promise; + onDelete: (id: string) => Promise; }, ref: React.RefObject<{ onRefresh: () => void }>, ) => { diff --git a/wren-ui/src/components/pages/home/prompt/index.tsx b/wren-ui/src/components/pages/home/prompt/index.tsx index c65a0ed2f6..d8e3260277 100644 --- a/wren-ui/src/components/pages/home/prompt/index.tsx +++ b/wren-ui/src/components/pages/home/prompt/index.tsx @@ -4,6 +4,7 @@ import { useState, forwardRef, useImperativeHandle, + useRef, } from 'react'; import styled from 'styled-components'; import { PROCESS_STATE } from '@/utils/enum'; @@ -16,6 +17,7 @@ import { AskPromptData } from '@/hooks/useAskPrompt'; import { CreateThreadInput, CreateThreadResponseInput, + ThreadResponseAnswerStatus, } from '@/apollo/client/graphql/__types__'; interface Props { @@ -84,6 +86,7 @@ export default forwardRef(function Prompt(props, ref) { const error = useMemo(() => askingTask?.error || null, [askingTask?.error]); const [showResult, setShowResult] = useState(false); const [question, setQuestion] = useState(''); + const persistedTaskIdRef = useRef(null); const currentProcessState = useMemo( () => askProcessState.currentState, [askProcessState.currentState], @@ -118,14 +121,71 @@ export default forwardRef(function Prompt(props, ref) { // create thread response for text to sql const intentSQLAnswer = async () => { - onCreateResponse && - (await onCreateResponse({ question, taskId: askingTask?.queryId })); + const taskId = askingTask?.queryId; + if (!taskId || !question) return; + + persistedTaskIdRef.current = taskId; + onCreateResponse && (await onCreateResponse({ question, taskId })); setShowResult(false); }; + useEffect(() => { + const taskId = askingTask?.queryId; + if (!taskId || !question || persistedTaskIdRef.current === taskId) return; + + const type = askingTask?.type; + const status = askingTask?.status; + + const persistResult = async ( + payload: CreateThreadInput | CreateThreadResponseInput, + ) => { + persistedTaskIdRef.current = taskId; + await onCreateResponse?.(payload); + setShowResult(false); + }; + + if ( + type === 'GENERAL' && + askingStreamTask && + status === 'FINISHED' && + !loading + ) { + void persistResult({ + question, + taskId, + answerStatus: ThreadResponseAnswerStatus.FINISHED, + answerContent: askingStreamTask, + }); + return; + } + + if (status === 'FAILED' && askingTask?.error) { + void persistResult({ + question, + taskId, + answerStatus: ThreadResponseAnswerStatus.FAILED, + answerErrorCode: askingTask.error.code, + answerErrorShortMessage: askingTask.error.shortMessage, + answerErrorMessage: askingTask.error.message, + }); + } + }, [ + askingTask?.queryId, + askingTask?.type, + askingTask?.status, + askingTask?.error?.code, + askingTask?.error?.shortMessage, + askingTask?.error?.message, + askingStreamTask, + question, + loading, + onCreateResponse, + ]); + const closeResult = () => { askProcessState.resetState(); setQuestion(''); + persistedTaskIdRef.current = null; onStopStreaming && onStopStreaming(); onStopRecommend && onStopRecommend(); }; @@ -142,7 +202,13 @@ export default forwardRef(function Prompt(props, ref) { // start the state as understanding when user submit question askProcessState.transitionTo(PROCESS_STATE.UNDERSTANDING); setShowResult(true); - onSubmit && (await onSubmit(value)); + try { + onSubmit && (await onSubmit(value)); + } catch (error) { + console.error(error); + setShowResult(false); + askProcessState.resetState(); + } }; useImperativeHandle( diff --git a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx index 6a0e628206..6dd1b6af03 100644 --- a/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx +++ b/wren-ui/src/components/pages/home/promptThread/AnswerResult.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo } from 'react'; +import { useEffect, useMemo, useRef } from 'react'; import { isEmpty, debounce } from 'lodash'; import clsx from 'clsx'; import { Button, Typography, Tabs, Tag, Tooltip } from 'antd'; @@ -21,10 +21,13 @@ import ViewSQLTabContent from '@/components/pages/home/promptThread/ViewSQLTabCo import TextBasedAnswer, { getAnswerIsFinished, } from '@/components/pages/home/promptThread/TextBasedAnswer'; -import ChartAnswer from '@/components/pages/home/promptThread/ChartAnswer'; +import ChartAnswer, { + getIsChartFinished, +} from '@/components/pages/home/promptThread/ChartAnswer'; import Preparation from '@/components/pages/home/preparation'; import { AskingTaskStatus, + ChartTaskStatus, ThreadResponse, ThreadResponseAnswerDetail, ThreadResponseAnswerStatus, @@ -176,6 +179,13 @@ const AdjustmentInformation = (props: { }; const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => { + if (!answerDetail) return true; + + const isFinishedWithoutContent = + answerDetail?.status === ThreadResponseAnswerStatus.FINISHED && + !answerDetail?.content?.trim(); + if (isFinishedWithoutContent) return true; + const isFinished = getAnswerIsFinished(answerDetail?.status); // it means the background task has not started yet, but answer is pending for generating const isProcessing = [ @@ -183,9 +193,22 @@ const isNeedGenerateAnswer = (answerDetail: ThreadResponseAnswerDetail) => { ThreadResponseAnswerStatus.PREPROCESSING, ThreadResponseAnswerStatus.FETCHING_DATA, ].includes(answerDetail?.status); - return answerDetail?.queryId === null && !isFinished && !isProcessing; + return !answerDetail?.queryId && !isFinished && !isProcessing; }; +const isAnswerGenerationInProgress = ( + status?: ThreadResponseAnswerStatus | null, +) => + [ + ThreadResponseAnswerStatus.NOT_STARTED, + ThreadResponseAnswerStatus.FETCHING_DATA, + ThreadResponseAnswerStatus.PREPROCESSING, + ThreadResponseAnswerStatus.STREAMING, + ].includes(status); + +const isChartGenerationActive = (status?: ChartTaskStatus | null) => + [ChartTaskStatus.FETCHING, ChartTaskStatus.GENERATING].includes(status); + export default function AnswerResult(props: Props) { const { threadResponse, isLastThreadResponse, isOpeningQuestion } = props; @@ -213,6 +236,8 @@ export default function AnswerResult(props: Props) { view, adjustment, } = threadResponse; + const autoGenerateAnswerRef = useRef(null); + const autoGenerateChartRef = useRef(null); const resultStyle = isLastThreadResponse ? { minHeight: 'calc(100vh - (194px))' } @@ -225,20 +250,33 @@ export default function AnswerResult(props: Props) { showRecommendedQuestions, ); - const isAnswerPrepared = !!answerDetail?.queryId || !!answerDetail?.status; const isBreakdownOnly = useMemo(() => { // we support rendering different types of answers now, so we need to check if it's old data. // existing thread response's answerDetail is null. return answerDetail === null && !isEmpty(breakdownDetail); }, [answerDetail, breakdownDetail]); + const isAnswerPrepared = !!answerDetail?.queryId || !!answerDetail?.status; + const showTextOnlyAnswer = + isAnswerPrepared && + !sql && + !view && + !isBreakdownOnly && + !threadResponse.chartDetail; // initialize generate answer useEffect(() => { if (isBreakdownOnly) return; + if ( + autoGenerateAnswerRef.current === id && + isAnswerGenerationInProgress(answerDetail?.status) + ) { + return; + } if ( canGenerateAnswer(askingTask, adjustmentTask) && isNeedGenerateAnswer(answerDetail) ) { + autoGenerateAnswerRef.current = id; const debouncedGenerateAnswer = debounce( () => { onGenerateTextBasedAnswer(id); @@ -254,22 +292,71 @@ export default function AnswerResult(props: Props) { }; } }, [ + id, isBreakdownOnly, askingTask?.status, adjustmentTask?.status, answerDetail?.status, + answerDetail?.content, + ]); + + useEffect(() => { + if (getAnswerIsFinished(answerDetail?.status)) { + autoGenerateAnswerRef.current = null; + } + }, [answerDetail?.status]); + + useEffect(() => { + if (!sql) return; + if (autoGenerateChartRef.current === id) return; + if (threadResponse.chartDetail) return; + if (!canGenerateAnswer(askingTask, adjustmentTask)) return; + + autoGenerateChartRef.current = id; + const debouncedGenerateChart = debounce( + () => { + onGenerateChartAnswer(id); + }, + 250, + { leading: false, trailing: true }, + ); + debouncedGenerateChart(); + + return () => { + debouncedGenerateChart.cancel(); + }; + }, [ + id, + sql, + askingTask?.status, + adjustmentTask?.status, + threadResponse.chartDetail?.status, ]); + useEffect(() => { + if ( + threadResponse.chartDetail?.status && + getIsChartFinished(threadResponse.chartDetail.status) + ) { + autoGenerateChartRef.current = null; + } + }, [threadResponse.chartDetail?.status]); + const onTabClick = (activeKey: string) => { - if (activeKey === ANSWER_TAB_KEYS.CHART && !threadResponse.chartDetail) { + if ( + activeKey === ANSWER_TAB_KEYS.CHART && + !threadResponse.chartDetail && + !isChartGenerationActive(threadResponse.chartDetail?.status) + ) { onGenerateChartAnswer(id); } }; const showAnswerTabs = - askingTask?.status === AskingTaskStatus.FINISHED || - isAnswerPrepared || - isBreakdownOnly; + !showTextOnlyAnswer && + (askingTask?.status === AskingTaskStatus.FINISHED || + isAnswerPrepared || + isBreakdownOnly); const rephrasedQuestion = threadResponse?.askingTask?.rephrasedQuestion || question; @@ -292,6 +379,18 @@ export default function AnswerResult(props: Props) { data={threadResponse} minimized={isAnswerPrepared} /> + {showTextOnlyAnswer && ( + <> +
+ +
+ {renderRecommendedQuestions( + isLastThreadResponse, + recommendedQuestionProps, + onSelectRecommendedQuestion, + )} + + )} {showAnswerTabs && ( <> diff --git a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx index 851d7f2ae4..2db457ba35 100644 --- a/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ChartAnswer.tsx @@ -11,7 +11,11 @@ import LineProperties from '@/components/chart/properties/LineProperties'; import StackedBarProperties from '@/components/chart/properties/StackedBarProperties'; import GroupedBarProperties from '@/components/chart/properties/GroupedBarProperties'; import { Props as AnswerResultProps } from '@/components/pages/home/promptThread/AnswerResult'; -import { ChartTaskStatus, ChartType } from '@/apollo/client/graphql/__types__'; +import { + ChartTaskStatus, + ChartType, + DashboardItemType, +} from '@/apollo/client/graphql/__types__'; import { usePreviewDataMutation } from '@/apollo/client/graphql/home.generated'; import { isEmpty, isEqual } from 'lodash'; import { @@ -19,7 +23,6 @@ import { getChartSpecOptionValues, } from '@/components/chart/handler'; import { useCreateDashboardItemMutation } from '@/apollo/client/graphql/dashboard.generated'; -import { DashboardItemType } from '@/apollo/server/repositories'; import usePromptThreadStore from './store'; const Chart = dynamic(() => import('@/components/chart'), { @@ -77,6 +80,15 @@ const getDynamicProperties = (chartType: ChartType) => { return propertiesMap[chartType] || BasicProperties; }; +const chartTypeToDashboardItemType = ( + chartType: ChartType | string, +): DashboardItemType | null => { + const normalized = String(chartType || '') + .replace(/-/g, '_') + .toUpperCase() as keyof typeof DashboardItemType; + return DashboardItemType[normalized] || null; +}; + export default function ChartAnswer(props: AnswerResultProps) { const { onGenerateChartAnswer, onAdjustChartAnswer } = usePromptThreadStore(); const { threadResponse } = props; @@ -94,18 +106,25 @@ export default function ChartAnswer(props: AnswerResultProps) { }); const [createDashboardItem] = useCreateDashboardItemMutation({ - onError: (error) => console.error(error), + onError: (error) => { + console.error(error); + message.error( + error.message || 'Failed to pin chart to dashboard. Please try again.', + ); + }, onCompleted: () => { message.success('Successfully pinned chart to dashboard.'); }, }); - // initial trigger when render + // Fetch preview data after the chart task has a terminal schema/result. useEffect(() => { + if (!getIsChartFinished(status)) return; + previewData({ variables: { where: { responseId: threadResponse.id } }, }); - }, []); + }, [previewData, status, threadResponse.id, threadResponse.sql]); const chartSpec = useMemo(() => { if ( @@ -179,6 +198,14 @@ export default function ChartAnswer(props: AnswerResultProps) { }; const onPin = () => { + const dashboardItemType = chartTypeToDashboardItemType( + chartType || chartDetail?.chartType, + ); + if (!dashboardItemType) { + message.error('Chart type is not supported for dashboard pinning.'); + return; + } + Modal.confirm({ title: 'Are you sure you want to pin this chart to the dashboard?', okText: 'Save', @@ -186,8 +213,7 @@ export default function ChartAnswer(props: AnswerResultProps) { await createDashboardItem({ variables: { data: { - // DashboardItemType is compatible with ChartType - itemType: chartType as unknown as DashboardItemType, + itemType: dashboardItemType, responseId: threadResponse.id, }, }, diff --git a/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx b/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx index 37126e24c2..78d740d022 100644 --- a/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx +++ b/wren-ui/src/components/pages/home/promptThread/TextBasedAnswer.tsx @@ -135,7 +135,7 @@ export default function TextBasedAnswer(props: AnswerResultProps) { onInitPreviewDone(); } - }, [isLastThreadResponse, allowPreviewData]); + }, [isLastThreadResponse, allowPreviewData, threadResponse.sql]); const loading = !getIsLoadingFinished(status); diff --git a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx index d52c103f4d..1cee69660a 100644 --- a/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx +++ b/wren-ui/src/components/pages/home/promptThread/ViewSQLTabContent.tsx @@ -71,12 +71,15 @@ export default function ViewSQLTabContent(props: AnswerResultProps) { if (isLastThreadResponse) { autoTriggerPreviewDataButton(); } - }, [isLastThreadResponse]); + }, [isLastThreadResponse, threadResponse.sql]); const { id, sql } = threadResponse; const { hasNativeSQL, dataSourceType } = nativeSQLResult; const showNativeSQL = hasNativeSQL; + const dataSourceOption = dataSourceType + ? DATA_SOURCE_OPTIONS[dataSourceType] + : undefined; const sqls = nativeSQLResult.nativeSQLMode && nativeSQLResult.loading === false @@ -132,15 +135,17 @@ export default function ViewSQLTabContent(props: AnswerResultProps) {
{nativeSQLResult.nativeSQLMode ? ( <> - + {dataSourceOption?.logo && ( + + )} - {DATA_SOURCE_OPTIONS[dataSourceType].label} + {dataSourceOption?.label || 'Original SQL'} ) : ( diff --git a/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx b/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx index c40410c30b..5a14c57071 100644 --- a/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx +++ b/wren-ui/src/components/pages/modeling/metadata/EditModelMetadata.tsx @@ -36,6 +36,12 @@ const CalculatedFieldEditableTable = makeEditableBaseTable(CalculatedFieldTable); const RelationshipEditableTable = makeEditableBaseTable(RelationTable); +const getMetadataId = (item: any) => + item?.relationId || item?.columnId || item?.nestedColumnId; + +const getFieldDisplayName = (item: any) => + item?.displayName ?? item?.referenceName ?? item?.sourceColumnName ?? ''; + export default function EditModelMetadata(props: Props) { const { formNamespace, @@ -61,17 +67,26 @@ export default function EditModelMetadata(props: Props) { }); }; - const handleMetadataChange = (fieldsName: string) => (value: any[]) => { + const handleMetadataChange = (fieldsName: string) => (value: any[] = []) => { // bind changeable metadata values onChange({ - [fieldsName]: value.map((item) => ({ - id: item.relationId || item.columnId || item.nestedColumnId, - description: item.description, - // Only models & fields, nested fields have alias - ...([FIELDS_NAME.FIELDS, FIELDS_NAME.NESTED_FIELDS].includes(fieldsName) - ? { displayName: item.displayName } - : {}), - })), + [fieldsName]: value + .map((item) => { + const id = getMetadataId(item); + if (!id) return null; + + return { + id, + description: item?.description, + // Only models & fields, nested fields have alias + ...([FIELDS_NAME.FIELDS, FIELDS_NAME.NESTED_FIELDS].includes( + fieldsName, + ) + ? { displayName: getFieldDisplayName(item) } + : {}), + }; + }) + .filter(Boolean), }); }; diff --git a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx index a95d73c48c..7908318a55 100644 --- a/wren-ui/src/components/pages/setup/ConnectDataSource.tsx +++ b/wren-ui/src/components/pages/setup/ConnectDataSource.tsx @@ -1,9 +1,11 @@ import Image from 'next/image'; import Link from 'next/link'; +import { useRouter } from 'next/router'; import { Alert, Typography, Form, Row, Col, Button } from 'antd'; import styled from 'styled-components'; import { DATA_SOURCES } from '@/utils/enum/dataSources'; import { getDataSource, getPostgresErrorMessage } from './utils'; +import { useEffect } from 'react'; const StyledForm = styled(Form)` border: 1px var(--gray-4) solid; @@ -26,8 +28,15 @@ interface Props { export default function ConnectDataSource(props: Props) { const { connectError, dataSource, submitting, onNext, onBack } = props; const [form] = Form.useForm(); + const router = useRouter(); const current = getDataSource(dataSource); + useEffect(() => { + if (typeof router.query.projectName === 'string') { + form.setFieldsValue({ displayName: router.query.projectName }); + } + }, [form, router.query.projectName]); + const submit = () => { form .validateFields() diff --git a/wren-ui/src/components/selectors/CombineFieldSelector.tsx b/wren-ui/src/components/selectors/CombineFieldSelector.tsx index 5b3da42537..ed41fe4117 100644 --- a/wren-ui/src/components/selectors/CombineFieldSelector.tsx +++ b/wren-ui/src/components/selectors/CombineFieldSelector.tsx @@ -39,6 +39,13 @@ export default function CombineFieldSelector(props: Props) { ...value, }); + useEffect(() => { + setInternalValue((currentValue) => ({ + model: value?.model || modelValue || currentValue.model, + field: value?.field || fieldValue || currentValue.field, + })); + }, [value?.model, value?.field, modelValue, fieldValue]); + const syncOnChange = () => { if (internalValue?.model && internalValue?.field) { onChange && onChange(internalValue); @@ -56,7 +63,11 @@ export default function CombineFieldSelector(props: Props) { const changeField = (field: string) => { onFieldChange && onFieldChange(field); - setInternalValue({ ...internalValue, field }); + setInternalValue({ + ...internalValue, + model: internalValue.model || value?.model || modelValue, + field, + }); }; return ( diff --git a/wren-ui/src/components/sidebar/Administration.tsx b/wren-ui/src/components/sidebar/Administration.tsx new file mode 100644 index 0000000000..ecec7dd7ea --- /dev/null +++ b/wren-ui/src/components/sidebar/Administration.tsx @@ -0,0 +1,73 @@ +import Link from 'next/link'; +import { useRouter } from 'next/router'; +import styled from 'styled-components'; +import TeamOutlined from '@ant-design/icons/TeamOutlined'; +import IdcardOutlined from '@ant-design/icons/IdcardOutlined'; +import SafetyCertificateOutlined from '@ant-design/icons/SafetyCertificateOutlined'; +import { Path, MENU_KEY } from '@/utils/enum'; +import SidebarMenu from '@/components/sidebar/SidebarMenu'; + +const Layout = styled.div` + padding: 16px 0; + position: absolute; + z-index: 1; + left: 0; + top: 0; + width: 100%; + background-color: var(--gray-2); + overflow: hidden; +`; + +const MENU_KEY_MAP = { + [Path.AdministrationUsers]: MENU_KEY.ADMIN_USERS, + [Path.AdministrationRoles]: MENU_KEY.ADMIN_ROLES, + [Path.AdministrationAssignments]: MENU_KEY.ADMIN_ASSIGNMENTS, +}; + +const linkStyle = { color: 'inherit', transition: 'none' }; + +export default function Administration() { + const router = useRouter(); + + const menuItems = [ + { + label: ( + + User Management + + ), + icon: , + key: MENU_KEY.ADMIN_USERS, + className: 'pl-4', + }, + { + label: ( + + Role Management + + ), + icon: , + key: MENU_KEY.ADMIN_ROLES, + className: 'pl-4', + }, + { + label: ( + + User Role Assignment + + ), + icon: , + key: MENU_KEY.ADMIN_ASSIGNMENTS, + className: 'pl-4', + }, + ]; + + return ( + + + + ); +} diff --git a/wren-ui/src/components/sidebar/index.tsx b/wren-ui/src/components/sidebar/index.tsx index 442c7255ed..b076a18e94 100644 --- a/wren-ui/src/components/sidebar/index.tsx +++ b/wren-ui/src/components/sidebar/index.tsx @@ -75,11 +75,10 @@ const DynamicSidebar = ( }; export default function Sidebar(props: Props) { - const { onOpenSettings } = props; const router = useRouter(); const onSettingsClick = (event) => { - onOpenSettings && onOpenSettings(); + router.push(Path.OrganizationGeneral); event.target.blur(); }; diff --git a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx index 56217dd7a3..db0a30c8d1 100644 --- a/wren-ui/src/components/sidebar/modeling/ModelTree.tsx +++ b/wren-ui/src/components/sidebar/modeling/ModelTree.tsx @@ -49,6 +49,7 @@ export default function ModelTree(props: Props) { const [triggerDataSourceDetection, { loading: isDetecting }] = useTriggerDataSourceDetectionMutation({ onError: (error) => console.error(error), + refetchQueries: [{ query: DIAGRAM }, { query: LIST_MODELS }], onCompleted: async (data) => { if (data.triggerDataSourceDetection) { message.warning('Schema change detected.'); diff --git a/wren-ui/src/hooks/useAdjustAnswer.tsx b/wren-ui/src/hooks/useAdjustAnswer.tsx index 832af9d7fa..7a960785f0 100644 --- a/wren-ui/src/hooks/useAdjustAnswer.tsx +++ b/wren-ui/src/hooks/useAdjustAnswer.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo } from 'react'; +import { useCallback, useEffect, useMemo, useRef } from 'react'; import { cloneDeep } from 'lodash'; import { ApolloClient, NormalizedCacheObject } from '@apollo/client'; import { THREAD } from '@/apollo/client/graphql/home'; @@ -15,6 +15,9 @@ import { ThreadResponse, } from '@/apollo/client/graphql/__types__'; +const ADJUSTMENT_POLL_INTERVAL_MS = 2000; +const ADJUSTMENT_POLL_MAX_INTERVAL_MS = 10000; + export const getIsFinished = (status: AskingTaskStatus) => [ AskingTaskStatus.FINISHED, @@ -71,9 +74,15 @@ export default function useAdjustAnswer(threadId?: number) { onError: (error) => console.error(error), }); const [fetchThreadResponse, threadResponseResult] = - useThreadResponseLazyQuery({ - pollInterval: 1000, - }); + useThreadResponseLazyQuery(); + const threadResponsePollingRef = useRef | null>( + null, + ); + const threadResponsePollingSessionRef = useRef(0); + const threadResponsePollingTargetRef = useRef(null); + const threadResponsePollingRequestRef = useRef | null>(null); + const threadResponsePollingDelayRef = useRef(ADJUSTMENT_POLL_INTERVAL_MS); + const lastAdjustmentTaskFingerprintRef = useRef(null); const loading = adjustThreadResponseResult.loading; @@ -87,15 +96,109 @@ export default function useAdjustAnswer(threadId?: number) { }; }, [adjustmentTask]); + const stopThreadResponsePolling = useCallback(() => { + threadResponsePollingSessionRef.current += 1; + threadResponsePollingTargetRef.current = null; + if (threadResponsePollingRef.current) { + clearTimeout(threadResponsePollingRef.current); + threadResponsePollingRef.current = null; + } + threadResponsePollingDelayRef.current = ADJUSTMENT_POLL_INTERVAL_MS; + }, []); + + const startThreadResponsePolling = useCallback( + async (responseId?: number) => { + if (!responseId) return; + if ( + threadResponsePollingTargetRef.current === responseId && + (threadResponsePollingRequestRef.current || threadResponsePollingRef.current) + ) { + return; + } + + stopThreadResponsePolling(); + threadResponsePollingTargetRef.current = responseId; + const pollingSessionId = threadResponsePollingSessionRef.current; + + const run = async () => { + if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + if (threadResponsePollingRequestRef.current) { + await threadResponsePollingRequestRef.current; + if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + } + + let shouldContinuePolling = true; + try { + const request = fetchThreadResponse({ + variables: { responseId }, + }); + threadResponsePollingRequestRef.current = request.then( + () => undefined, + ); + const result = await request; + const task = result.data?.threadResponse?.adjustmentTask; + if (!task || getIsFinished(task.status)) { + shouldContinuePolling = false; + stopThreadResponsePolling(); + } + } catch (error) { + console.error(error); + } finally { + threadResponsePollingRequestRef.current = null; + if ( + shouldContinuePolling && + threadResponsePollingSessionRef.current === pollingSessionId + ) { + threadResponsePollingRef.current = setTimeout( + run, + threadResponsePollingDelayRef.current, + ); + } + } + }; + + await run(); + }, + [fetchThreadResponse, stopThreadResponsePolling], + ); + useEffect(() => { const isFinished = getIsFinished(adjustmentTask?.status); - if (isFinished) threadResponseResult.stopPolling(); + if (isFinished) stopThreadResponsePolling(); }, [adjustmentTask?.status]); + useEffect(() => { + const fingerprint = JSON.stringify({ + queryId: adjustmentTask?.queryId || null, + status: adjustmentTask?.status || null, + sql: adjustmentTask?.sql || null, + errorCode: adjustmentTask?.error?.code || null, + invalidSql: adjustmentTask?.invalidSql || null, + }); + + if (lastAdjustmentTaskFingerprintRef.current === fingerprint) { + threadResponsePollingDelayRef.current = Math.min( + threadResponsePollingDelayRef.current * 2, + ADJUSTMENT_POLL_MAX_INTERVAL_MS, + ); + } else { + threadResponsePollingDelayRef.current = ADJUSTMENT_POLL_INTERVAL_MS; + lastAdjustmentTaskFingerprintRef.current = fingerprint; + } + }, [ + adjustmentTask?.queryId, + adjustmentTask?.status, + adjustmentTask?.sql, + adjustmentTask?.error?.code, + adjustmentTask?.invalidSql, + ]); + const onAdjustReasoningSteps = async ( responseId: number, input: { tables: string[]; sqlGenerationReasoning: string }, ) => { + if (!responseId) return; + const response = await adjustThreadResponse({ variables: { responseId, @@ -108,9 +211,9 @@ export default function useAdjustAnswer(threadId?: number) { // start polling new thread response const nextThreadResponse = response.data?.adjustThreadResponse; - await fetchThreadResponse({ - variables: { responseId: nextThreadResponse.id }, - }); + if (!nextThreadResponse?.id) return; + + await startThreadResponsePolling(nextThreadResponse.id); // update new thread response to cache handleUpdateThreadCache( @@ -121,12 +224,16 @@ export default function useAdjustAnswer(threadId?: number) { }; const onAdjustSQL = async (responseId: number, sql: string) => { + if (!responseId) return; + const response = await adjustThreadResponse({ variables: { responseId, data: { sql } }, }); // update thread cache const nextThreadResponse = response.data?.adjustThreadResponse; + if (!nextThreadResponse) return; + handleUpdateThreadCache( threadId, nextThreadResponse, @@ -150,10 +257,18 @@ export default function useAdjustAnswer(threadId?: number) { const onReRun = async (threadResponse: ThreadResponse) => { const responseId = threadResponse.id; + if (!responseId) return; + await rerunAdjustmentTask({ variables: { responseId } }); - await fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; + useEffect(() => { + return () => { + stopThreadResponsePolling(); + }; + }, [stopThreadResponsePolling]); + return { data, loading, diff --git a/wren-ui/src/hooks/useAskPrompt.tsx b/wren-ui/src/hooks/useAskPrompt.tsx index ac45713287..1a1e45a6b8 100644 --- a/wren-ui/src/hooks/useAskPrompt.tsx +++ b/wren-ui/src/hooks/useAskPrompt.tsx @@ -1,4 +1,4 @@ -import { useCallback, useEffect, useMemo, useState } from 'react'; +import { useCallback, useEffect, useMemo, useRef, useState } from 'react'; import { cloneDeep, uniq } from 'lodash'; import { AdjustmentTask, @@ -30,6 +30,11 @@ export interface AskPromptData { recommendedQuestions?: RecommendedQuestionsTask; } +const ASKING_TASK_POLL_INTERVAL_MS = 1000; +const RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS = 1000; +const ASKING_TASK_POLL_MAX_INTERVAL_MS = 10000; +const RECOMMENDED_QUESTIONS_POLL_MAX_INTERVAL_MS = 10000; + export const getIsFinished = (status: AskingTaskStatus) => [ AskingTaskStatus.FINISHED, @@ -130,11 +135,6 @@ const handleUpdateRerunAskingTaskCache = ( if (result?.thread) { const task = cloneDeep(askingTask); - // bypass understanding status to thread response - if (task.status === AskingTaskStatus.UNDERSTANDING) { - task.status = AskingTaskStatus.SEARCHING; - task.type = AskingTaskType.TEXT_TO_SQL; - } client.cache.updateQuery( { query: THREAD, @@ -170,7 +170,8 @@ export default function useAskPrompt(threadId?: number) { onError: (error) => console.error(error), }); const [fetchAskingTask, askingTaskResult] = useAskingTaskLazyQuery({ - pollInterval: 1000, + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const [fetchAskingStreamTask, askingStreamTaskResult] = useAskingStreamTask(); const [createInstantRecommendedQuestions] = @@ -179,12 +180,46 @@ export default function useAskPrompt(threadId?: number) { }); const [fetchInstantRecommendedQuestions, instantRecommendedQuestionsResult] = useInstantRecommendedQuestionsLazyQuery({ - pollInterval: 1000, + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); + const askingTaskPollingRef = useRef | null>( + null, + ); + const askingTaskPollingSessionRef = useRef(0); + const askingTaskPollingTargetRef = useRef(null); + const askingTaskPollingRequestRef = useRef | null>(null); + const askingTaskPollingDelayRef = useRef(ASKING_TASK_POLL_INTERVAL_MS); + const lastAskingTaskFingerprintRef = useRef(null); + const recommendedPollingRef = useRef | null>( + null, + ); + const recommendedPollingSessionRef = useRef(0); + const recommendedPollingTargetRef = useRef(null); + const recommendedPollingRequestRef = useRef | null>(null); + const recommendedPollingDelayRef = useRef( + RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS, + ); + const lastRecommendedFingerprintRef = useRef(null); + const recommendedCreationKeyRef = useRef(null); + const recommendedCreationRequestRef = useRef | null>(null); + const [activeAskingTaskId, setActiveAskingTaskId] = useState( + null, + ); const askingTask = useMemo( - () => askingTaskResult.data?.askingTask || null, - [askingTaskResult.data], + () => { + const task = askingTaskResult.data?.askingTask || null; + if ( + activeAskingTaskId && + task?.queryId && + task.queryId !== activeAskingTaskId + ) { + return null; + } + return task; + }, + [activeAskingTaskId, askingTaskResult.data], ); const askingTaskType = useMemo(() => askingTask?.type, [askingTask?.type]); const askingStreamTask = askingStreamTaskResult.data; @@ -197,6 +232,137 @@ export default function useAskPrompt(threadId?: number) { const loading = askingStreamTaskResult.loading; + const stopAskingTaskPolling = useCallback(() => { + askingTaskPollingSessionRef.current += 1; + askingTaskPollingTargetRef.current = null; + if (askingTaskPollingRef.current) { + clearTimeout(askingTaskPollingRef.current); + askingTaskPollingRef.current = null; + } + askingTaskPollingDelayRef.current = ASKING_TASK_POLL_INTERVAL_MS; + }, []); + + const stopRecommendedPolling = useCallback(() => { + recommendedPollingSessionRef.current += 1; + recommendedPollingTargetRef.current = null; + if (recommendedPollingRef.current) { + clearTimeout(recommendedPollingRef.current); + recommendedPollingRef.current = null; + } + recommendedPollingDelayRef.current = RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS; + recommendedCreationKeyRef.current = null; + recommendedCreationRequestRef.current = null; + }, []); + + const startAskingTaskPolling = useCallback( + async (taskId?: string) => { + if (!taskId) return; + if ( + askingTaskPollingTargetRef.current === taskId && + (askingTaskPollingRequestRef.current || askingTaskPollingRef.current) + ) { + return; + } + + stopAskingTaskPolling(); + setActiveAskingTaskId(taskId); + askingTaskPollingTargetRef.current = taskId; + const pollingSessionId = askingTaskPollingSessionRef.current; + + const run = async () => { + if (askingTaskPollingSessionRef.current !== pollingSessionId) return; + if (askingTaskPollingRequestRef.current) { + await askingTaskPollingRequestRef.current; + if (askingTaskPollingSessionRef.current !== pollingSessionId) return; + } + + let shouldContinuePolling = true; + try { + const request = fetchAskingTask({ + variables: { taskId }, + }); + askingTaskPollingRequestRef.current = request.then(() => undefined); + const result = await request; + const task = result.data?.askingTask; + if (!task || getIsFinished(task.status)) { + shouldContinuePolling = false; + stopAskingTaskPolling(); + } + } catch (error) { + console.error(error); + } finally { + askingTaskPollingRequestRef.current = null; + if ( + shouldContinuePolling && + askingTaskPollingSessionRef.current === pollingSessionId + ) { + askingTaskPollingRef.current = setTimeout( + run, + askingTaskPollingDelayRef.current, + ); + } + } + }; + + await run(); + }, + [fetchAskingTask, stopAskingTaskPolling], + ); + + const startRecommendedPolling = useCallback( + async (taskId?: string) => { + if (!taskId) return; + if ( + recommendedPollingTargetRef.current === taskId && + (recommendedPollingRequestRef.current || recommendedPollingRef.current) + ) { + return; + } + + stopRecommendedPolling(); + recommendedPollingTargetRef.current = taskId; + const pollingSessionId = recommendedPollingSessionRef.current; + + const run = async () => { + if (recommendedPollingSessionRef.current !== pollingSessionId) return; + if (recommendedPollingRequestRef.current) { + await recommendedPollingRequestRef.current; + if (recommendedPollingSessionRef.current !== pollingSessionId) return; + } + + let shouldContinuePolling = true; + try { + const request = fetchInstantRecommendedQuestions({ + variables: { taskId }, + }); + recommendedPollingRequestRef.current = request.then(() => undefined); + const result = await request; + const task = result.data?.instantRecommendedQuestions; + if (!task || isRecommendedFinished(task.status)) { + shouldContinuePolling = false; + stopRecommendedPolling(); + } + } catch (error) { + console.error(error); + } finally { + recommendedPollingRequestRef.current = null; + if ( + shouldContinuePolling && + recommendedPollingSessionRef.current === pollingSessionId + ) { + recommendedPollingRef.current = setTimeout( + run, + recommendedPollingDelayRef.current, + ); + } + } + }; + + await run(); + }, + [fetchInstantRecommendedQuestions, stopRecommendedPolling], + ); + const data = useMemo( () => ({ originalQuestion, @@ -208,22 +374,57 @@ export default function useAskPrompt(threadId?: number) { ); const startRecommendedQuestions = useCallback(async () => { + if (!originalQuestion?.trim()) return; + const previousQuestions = [ // slice the last 5 questions in threadQuestions ...uniq(threadQuestions).slice(-5), originalQuestion, ]; - const response = await createInstantRecommendedQuestions({ - variables: { data: { previousQuestions } }, + const creationKey = JSON.stringify({ + askingQueryId: askingTask?.queryId || null, + previousQuestions, }); - fetchInstantRecommendedQuestions({ - variables: { taskId: response.data.createInstantRecommendedQuestions.id }, - }); - }, [originalQuestion]); + + if (recommendedCreationKeyRef.current === creationKey) { + if (recommendedCreationRequestRef.current) { + await recommendedCreationRequestRef.current; + } + return; + } + + recommendedCreationKeyRef.current = creationKey; + const request = (async () => { + const response = await createInstantRecommendedQuestions({ + variables: { data: { previousQuestions } }, + }); + const taskId = response.data?.createInstantRecommendedQuestions?.id; + if (!taskId) return; + + await startRecommendedPolling(taskId); + })(); + + recommendedCreationRequestRef.current = request; + try { + await request; + } finally { + recommendedCreationRequestRef.current = null; + } + }, [ + originalQuestion, + threadQuestions, + askingTask?.queryId, + createInstantRecommendedQuestions, + startRecommendedPolling, + ]); const checkFetchAskingStreamTask = useCallback( (task: AskingTask) => { - if (!askingStreamTask && task.status === AskingTaskStatus.PLANNING) { + if ( + !askingStreamTask && + task?.queryId && + task.status === AskingTaskStatus.PLANNING + ) { fetchAskingStreamTask(task.queryId); } }, @@ -232,7 +433,7 @@ export default function useAskPrompt(threadId?: number) { useEffect(() => { const isFinished = getIsFinished(askingTask?.status); - if (isFinished) askingTaskResult.stopPolling(); + if (isFinished) stopAskingTaskPolling(); // handle update cache for preparing component if (isNeedPreparing(askingTask)) { @@ -243,6 +444,34 @@ export default function useAskPrompt(threadId?: number) { } }, [askingTask?.status, threadId, checkFetchAskingStreamTask]); + useEffect(() => { + const fingerprint = JSON.stringify({ + queryId: askingTask?.queryId || null, + status: askingTask?.status || null, + type: askingTask?.type || null, + candidateCount: askingTask?.candidates?.length || 0, + errorCode: askingTask?.error?.code || null, + traceId: askingTask?.traceId || null, + }); + + if (lastAskingTaskFingerprintRef.current === fingerprint) { + askingTaskPollingDelayRef.current = Math.min( + askingTaskPollingDelayRef.current * 2, + ASKING_TASK_POLL_MAX_INTERVAL_MS, + ); + } else { + askingTaskPollingDelayRef.current = ASKING_TASK_POLL_INTERVAL_MS; + lastAskingTaskFingerprintRef.current = fingerprint; + } + }, [ + askingTask?.queryId, + askingTask?.status, + askingTask?.type, + askingTask?.candidates?.length, + askingTask?.error?.code, + askingTask?.traceId, + ]); + useEffect(() => { // handle instant recommended questions if (isNeedRecommendedQuestions(askingTask)) { @@ -251,8 +480,31 @@ export default function useAskPrompt(threadId?: number) { }, [askingTask?.type]); useEffect(() => { - if (isRecommendedFinished(recommendedQuestions?.status)) - instantRecommendedQuestionsResult.stopPolling(); + const fingerprint = JSON.stringify({ + status: recommendedQuestions?.status || null, + count: recommendedQuestions?.questions?.length || 0, + errorCode: recommendedQuestions?.error?.code || null, + }); + + if (lastRecommendedFingerprintRef.current === fingerprint) { + recommendedPollingDelayRef.current = Math.min( + recommendedPollingDelayRef.current * 2, + RECOMMENDED_QUESTIONS_POLL_MAX_INTERVAL_MS, + ); + } else { + recommendedPollingDelayRef.current = RECOMMENDED_QUESTIONS_POLL_INTERVAL_MS; + lastRecommendedFingerprintRef.current = fingerprint; + } + }, [ + recommendedQuestions?.status, + recommendedQuestions?.questions?.length, + recommendedQuestions?.error?.code, + ]); + + useEffect(() => { + if (isRecommendedFinished(recommendedQuestions?.status)) { + stopRecommendedPolling(); + } }, [recommendedQuestions]); useEffect(() => { @@ -268,21 +520,29 @@ export default function useAskPrompt(threadId?: number) { await cancelAskingTask({ variables: { taskId } }).catch((error) => console.error(error), ); + stopAskingTaskPolling(); // waiting for polling fetching stop await nextTick(1000); } }; const onReRun = async (threadResponse: ThreadResponse) => { + if (!threadResponse?.id) return; + askingStreamTaskResult.reset(); setOriginalQuestion(threadResponse.question); try { const response = await rerunAskingTask({ variables: { responseId: threadResponse.id }, }); - const { data } = await fetchAskingTask({ - variables: { taskId: response.data.rerunAskingTask.id }, - }); + const taskId = response.data?.rerunAskingTask?.id; + if (!taskId) return; + + const { data } = await fetchAskingTask({ variables: { taskId } }); + if (!data?.askingTask) return; + + await startAskingTaskPolling(taskId); + // update the asking task in cache manually handleUpdateRerunAskingTaskCache( threadId, @@ -302,25 +562,37 @@ export default function useAskPrompt(threadId?: number) { const response = await createAskingTask({ variables: { data: { question: value, threadId } }, }); - await fetchAskingTask({ - variables: { taskId: response.data.createAskingTask.id }, - }); + const taskId = response.data?.createAskingTask?.id; + if (!taskId) { + throw new Error('CreateAskingTask returned no task id'); + } + + await startAskingTaskPolling(taskId); } catch (error) { console.error(error); + stopAskingTaskPolling(); + throw error; } }; const onFetching = async (queryId: string) => { - await fetchAskingTask({ - variables: { taskId: queryId }, - }); + if (!queryId) return; + + await startAskingTaskPolling(queryId); }; - const onStopPolling = () => askingTaskResult.stopPolling(); + const onStopPolling = () => stopAskingTaskPolling(); const onStopStreaming = () => askingStreamTaskResult.reset(); - const onStopRecommend = () => instantRecommendedQuestionsResult.stopPolling(); + const onStopRecommend = () => stopRecommendedPolling(); + + useEffect(() => { + return () => { + stopAskingTaskPolling(); + stopRecommendedPolling(); + }; + }, [stopAskingTaskPolling, stopRecommendedPolling]); const onStoreThreadQuestions = (questions: string[]) => setThreadQuestions(questions); diff --git a/wren-ui/src/hooks/useCheckOnboarding.tsx b/wren-ui/src/hooks/useCheckOnboarding.tsx index fb62e2480a..5f76ab711e 100644 --- a/wren-ui/src/hooks/useCheckOnboarding.tsx +++ b/wren-ui/src/hooks/useCheckOnboarding.tsx @@ -6,7 +6,7 @@ import { Path } from '@/utils/enum'; const redirectRoute = { [OnboardingStatus.DATASOURCE_SAVED]: Path.OnboardingModels, - [OnboardingStatus.NOT_STARTED]: Path.OnboardingConnection, + [OnboardingStatus.NOT_STARTED]: Path.ProjectCreate, [OnboardingStatus.ONBOARDING_FINISHED]: Path.Modeling, [OnboardingStatus.WITH_SAMPLE_DATASET]: Path.Modeling, }; @@ -21,9 +21,14 @@ export const useWithOnboarding = () => { if (onboardingStatus) { const newPath = redirectRoute[onboardingStatus]; const pathname = router.pathname; + const isNewProjectFlow = router.query.newProject === '1'; // redirect to new path if onboarding is not completed if (newPath && newPath !== Path.Modeling) { + if (pathname.startsWith(Path.Onboarding) && isNewProjectFlow) { + return; + } + // do not redirect if the new path and router pathname are the same if (newPath === pathname) { return; @@ -64,11 +69,14 @@ export const useWithOnboarding = () => { pathname as Path, ) ) { + if (isNewProjectFlow) { + return; + } router.push(newPath); return; } } - }, [onboardingStatus, router.pathname]); + }, [onboardingStatus, router.pathname, router.query.newProject]); return { loading, diff --git a/wren-ui/src/hooks/useHomeSidebar.tsx b/wren-ui/src/hooks/useHomeSidebar.tsx index 8611af4795..c914fa6837 100644 --- a/wren-ui/src/hooks/useHomeSidebar.tsx +++ b/wren-ui/src/hooks/useHomeSidebar.tsx @@ -10,7 +10,8 @@ import { export default function useHomeSidebar() { const router = useRouter(); const { data, refetch } = useThreadsQuery({ - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const [updateThread] = useUpdateThreadMutation({ onError: (error) => console.error(error), diff --git a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx index 9a04a605a3..243843006d 100644 --- a/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx +++ b/wren-ui/src/hooks/useRecommendedQuestionsInstruction.tsx @@ -1,4 +1,4 @@ -import { useMemo, useState, useEffect } from 'react'; +import { useCallback, useMemo, useRef, useState, useEffect } from 'react'; import { groupBy, orderBy, flatMap } from 'lodash'; import { message } from 'antd'; import Icon from '@/import/icon'; @@ -20,6 +20,9 @@ export interface GroupedQuestion { sql: string; } +const RECOMMENDATION_POLL_INTERVAL_MS = 2000; +const RECOMMENDATION_POLL_MAX_INTERVAL_MS = 10000; + const getGroupedQuestions = ( questions: ResultQuestion[], ): GroupedQuestion[] => { @@ -42,10 +45,16 @@ export default function useRecommendedQuestionsInstruction() { const [recommendedQuestions, setRecommendedQuestions] = useState< GroupedQuestion[] >([]); + const pollingRef = useRef | null>(null); + const pollingSessionRef = useRef(0); + const pollingRequestRef = useRef | null>(null); + const pollingDelayRef = useRef(RECOMMENDATION_POLL_INTERVAL_MS); + const lastFingerprintRef = useRef(null); const [fetchRecommendationQuestions, recommendationQuestionsResult] = useGetProjectRecommendationQuestionsLazyQuery({ - pollInterval: 2000, + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); // Handle errors via try/catch blocks rather than onError callback @@ -59,12 +68,64 @@ export default function useRecommendedQuestionsInstruction() { [recommendationQuestionsResult.data], ); + const stopPolling = useCallback(() => { + pollingSessionRef.current += 1; + if (pollingRef.current) { + clearTimeout(pollingRef.current); + pollingRef.current = null; + } + pollingDelayRef.current = RECOMMENDATION_POLL_INTERVAL_MS; + }, []); + + const startPolling = useCallback(async () => { + if (pollingRequestRef.current || pollingRef.current) { + return; + } + + stopPolling(); + const pollingSessionId = pollingSessionRef.current; + + const run = async () => { + if (pollingSessionRef.current !== pollingSessionId) return; + if (pollingRequestRef.current) { + await pollingRequestRef.current; + if (pollingSessionRef.current !== pollingSessionId) return; + } + + let shouldContinuePolling = true; + try { + const request = fetchRecommendationQuestions(); + pollingRequestRef.current = request.then(() => undefined); + const result = await request; + const task = result.data?.getProjectRecommendationQuestions; + if (!task || isRecommendedFinished(task.status)) { + shouldContinuePolling = false; + stopPolling(); + } + } catch (error) { + console.error(error); + } finally { + pollingRequestRef.current = null; + if ( + shouldContinuePolling && + pollingSessionRef.current === pollingSessionId + ) { + pollingRef.current = setTimeout(run, pollingDelayRef.current); + } + } + }; + + await run(); + }, [fetchRecommendationQuestions, stopPolling]); + useEffect(() => { const fetchRecommendationQuestionsData = async () => { const result = await fetchRecommendationQuestions(); const data = result.data?.getProjectRecommendationQuestions; + if (!data) { + return; + } - // for existing projects that do not have to generate recommended questions yet if (isRecommendedFinished(data.status)) { if (data.questions.length > 0) { // for regenerate then leave and go back to the home page @@ -72,15 +133,23 @@ export default function useRecommendedQuestionsInstruction() { setShowRecommendedQuestionsPromptMode(true); } + } else { + setGenerating(true); + await startPolling(); } }; fetchRecommendationQuestionsData(); - }, []); + return () => stopPolling(); + }, [fetchRecommendationQuestions, startPolling, stopPolling]); useEffect(() => { + if (!recommendedQuestionsTask) { + return; + } + if (isRecommendedFinished(recommendedQuestionsTask?.status)) { - recommendationQuestionsResult.stopPolling(); + stopPolling(); if (recommendedQuestionsTask.questions.length === 0) { isRegenerate && setShowRetry(true); @@ -106,14 +175,41 @@ export default function useRecommendedQuestionsInstruction() { setGenerating(false); } - }, [recommendedQuestionsTask]); + }, [ + isRegenerate, + recommendedQuestionsTask, + showRecommendedQuestionsPromptMode, + stopPolling, + ]); + + useEffect(() => { + const fingerprint = JSON.stringify({ + status: recommendedQuestionsTask?.status || null, + count: recommendedQuestionsTask?.questions?.length || 0, + errorCode: recommendedQuestionsTask?.error?.code || null, + }); + + if (lastFingerprintRef.current === fingerprint) { + pollingDelayRef.current = Math.min( + pollingDelayRef.current * 2, + RECOMMENDATION_POLL_MAX_INTERVAL_MS, + ); + } else { + pollingDelayRef.current = RECOMMENDATION_POLL_INTERVAL_MS; + lastFingerprintRef.current = fingerprint; + } + }, [ + recommendedQuestionsTask?.status, + recommendedQuestionsTask?.questions?.length, + recommendedQuestionsTask?.error?.code, + ]); const onGetRecommendationQuestions = async () => { setGenerating(true); setIsRegenerate(true); try { await generateProjectRecommendationQuestions(); - fetchRecommendationQuestions(); + await startPolling(); } catch (error) { console.error(error); } diff --git a/wren-ui/src/hooks/useSetupConnectionDataSource.tsx b/wren-ui/src/hooks/useSetupConnectionDataSource.tsx index 5e8900fde5..5d516dcc69 100644 --- a/wren-ui/src/hooks/useSetupConnectionDataSource.tsx +++ b/wren-ui/src/hooks/useSetupConnectionDataSource.tsx @@ -6,7 +6,10 @@ import { DATABRICKS_AUTH_METHOD, } from '@/utils/enum'; import { useSaveDataSourceMutation } from '@/apollo/client/graphql/dataSource.generated'; -import { DataSourceName } from '@/apollo/client/graphql/__types__'; +import { + DataSourceName, + WorkspaceProjectType, +} from '@/apollo/client/graphql/__types__'; const PASSWORD_PLACEHOLDER = '************'; @@ -30,16 +33,25 @@ export default function useSetupConnectionDataSource() { const saveDataSource = useCallback( async (properties?: Record) => { + const requestedProjectType = + typeof router.query.projectType === 'string' + ? router.query.projectType + : WorkspaceProjectType.CLASSIC; + await saveDataSourceMutation({ variables: { data: { type: selected, properties: transformFormToProperties(properties, selected), + projectType: + requestedProjectType === WorkspaceProjectType.AGENTIC + ? WorkspaceProjectType.AGENTIC + : WorkspaceProjectType.CLASSIC, }, }, }); }, - [selected, saveDataSourceMutation], + [router.query.projectType, saveDataSourceMutation, selected], ); const completedDataSourceSave = useCallback(async () => { diff --git a/wren-ui/src/pages/administration/assignments.tsx b/wren-ui/src/pages/administration/assignments.tsx new file mode 100644 index 0000000000..252691c282 --- /dev/null +++ b/wren-ui/src/pages/administration/assignments.tsx @@ -0,0 +1,281 @@ +import { useEffect, useState } from 'react'; +import { useMutation, useQuery } from '@apollo/client'; +import { + Button, + Form, + Modal, + Popconfirm, + Select, + Table, + TableColumnsType, + Typography, + message, +} from 'antd'; +import SafetyCertificateOutlined from '@ant-design/icons/SafetyCertificateOutlined'; +import EditOutlined from '@ant-design/icons/EditOutlined'; +import DeleteOutlined from '@ant-design/icons/DeleteOutlined'; +import SiderLayout from '@/components/layouts/SiderLayout'; +import PageLayout from '@/components/layouts/PageLayout'; +import { + ASSIGN_ROLE_TO_USER, + LIST_USER_ROLE_MAPPINGS, + REMOVE_ROLE_FROM_USER, + UPDATE_USER_ROLES, +} from '@/apollo/client/graphql/rbac'; +import { + Role, + RoleTags, + User, + UserRoleMapping, +} from '@/components/pages/administration/types'; +import { getAbsoluteTime } from '@/utils/time'; + +const { Text } = Typography; + +type AssignmentModalState = { + visible: boolean; + user?: User; +}; + +const AssignmentModal = ({ + users, + roles, + state, + loading, + onClose, + onSubmit, +}: { + users: User[]; + roles: Role[]; + state: AssignmentModalState; + loading: boolean; + onClose: () => void; + onSubmit: (values: any, user?: User) => Promise; +}) => { + const [form] = Form.useForm(); + const isUpdate = !!state.user; + + useEffect(() => { + if (!state.visible) return; + form.setFieldsValue({ + userId: state.user?.id, + roleId: undefined, + roleIds: state.user?.roles?.map((role) => role.id) || [], + }); + }, [form, state.visible, state.user]); + + const submit = async () => { + const values = await form.validateFields(); + await onSubmit(values, state.user); + form.resetFields(); + onClose(); + }; + + return ( + form.resetFields()} + > +
+ + ({ + label: role.name, + value: role.id, + }))} + /> + + ) : ( + +
+ + + + ); +} diff --git a/wren-ui/src/pages/administration/index.tsx b/wren-ui/src/pages/administration/index.tsx new file mode 100644 index 0000000000..653c921b87 --- /dev/null +++ b/wren-ui/src/pages/administration/index.tsx @@ -0,0 +1,14 @@ +import { useEffect } from 'react'; +import { useRouter } from 'next/router'; +import PageLoading from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; + +export default function AdministrationIndex() { + const router = useRouter(); + + useEffect(() => { + router.replace(Path.AdministrationUsers); + }, [router]); + + return ; +} diff --git a/wren-ui/src/pages/administration/roles.tsx b/wren-ui/src/pages/administration/roles.tsx new file mode 100644 index 0000000000..a17f1457e1 --- /dev/null +++ b/wren-ui/src/pages/administration/roles.tsx @@ -0,0 +1,245 @@ +import { useEffect, useState } from 'react'; +import { useMutation, useQuery } from '@apollo/client'; +import { + Button, + Form, + Input, + Modal, + Table, + TableColumnsType, + Tag, + Typography, + message, +} from 'antd'; +import IdcardOutlined from '@ant-design/icons/IdcardOutlined'; +import EditOutlined from '@ant-design/icons/EditOutlined'; +import SiderLayout from '@/components/layouts/SiderLayout'; +import PageLayout from '@/components/layouts/PageLayout'; +import { + CREATE_ROLE, + LIST_RBAC_ROLES, + UPDATE_ROLE, +} from '@/apollo/client/graphql/rbac'; +import { Role, User } from '@/components/pages/administration/types'; +import { getAbsoluteTime } from '@/utils/time'; + +const { Paragraph, Text } = Typography; + +type RoleModalState = { + visible: boolean; + role?: Role; +}; + +const RoleModal = ({ + state, + loading, + onClose, + onSubmit, +}: { + state: RoleModalState; + loading: boolean; + onClose: () => void; + onSubmit: (values: any, role?: Role) => Promise; +}) => { + const [form] = Form.useForm(); + const isEdit = !!state.role; + + useEffect(() => { + if (!state.visible) return; + form.setFieldsValue({ + name: state.role?.name, + description: state.role?.description, + }); + }, [form, state.visible, state.role]); + + const submit = async () => { + const values = await form.validateFields(); + await onSubmit(values, state.role); + form.resetFields(); + onClose(); + }; + + return ( + form.resetFields()} + > + + + + + + + + + + ); +}; + +const AssignedUsers = ({ users = [] }: { users?: User[] }) => { + if (!users.length) return No users assigned; + return ( + <> + {users.map((user) => ( + + {user.name} - {user.email} + + ))} + + ); +}; + +export default function RoleManagement() { + const [modalState, setModalState] = useState({ + visible: false, + }); + const { data, loading, refetch } = useQuery(LIST_RBAC_ROLES, { + fetchPolicy: 'cache-and-network', + }); + const roles: Role[] = data?.roles || []; + + const mutationOptions = { + onError: (error) => message.error(error.message), + }; + const [createRole, createRoleState] = useMutation( + CREATE_ROLE, + mutationOptions, + ); + const [updateRole, updateRoleState] = useMutation( + UPDATE_ROLE, + mutationOptions, + ); + + const closeModal = () => setModalState({ visible: false }); + + const submitRole = async (values: any, role?: Role) => { + if (role) { + await updateRole({ + variables: { + where: { id: role.id }, + data: { + name: values.name, + description: values.description || null, + }, + }, + }); + message.success('Successfully updated role.'); + } else { + await createRole({ + variables: { + data: { + name: values.name, + description: values.description || null, + }, + }, + }); + message.success('Successfully created role.'); + } + await refetch(); + }; + + const columns: TableColumnsType = [ + { + title: 'Role', + dataIndex: 'name', + width: 220, + render: (name) => {name}, + }, + { + title: 'Description', + dataIndex: 'description', + render: (description) => ( + + {description || No description} + + ), + }, + { + title: 'Assigned users', + dataIndex: 'users', + render: (users) => , + }, + { + title: 'Created', + dataIndex: 'createdAt', + width: 180, + render: (value) => ( + {getAbsoluteTime(value)} + ), + }, + { + title: 'Actions', + width: 110, + align: 'center', + render: (_, record) => ( + + ), + }, + ]; + + return ( + + + + Role Management + + } + description="Maintain reusable role definitions. Permission enforcement can be added later without changing these assignments." + titleExtra={ + + } + > +
+ + + + ); +} diff --git a/wren-ui/src/pages/administration/users.tsx b/wren-ui/src/pages/administration/users.tsx new file mode 100644 index 0000000000..ad4d605bde --- /dev/null +++ b/wren-ui/src/pages/administration/users.tsx @@ -0,0 +1,296 @@ +import { useEffect, useState } from 'react'; +import { useMutation, useQuery } from '@apollo/client'; +import { + Button, + Form, + Input, + Modal, + Select, + Switch, + Table, + TableColumnsType, + Typography, + message, +} from 'antd'; +import TeamOutlined from '@ant-design/icons/TeamOutlined'; +import EditOutlined from '@ant-design/icons/EditOutlined'; +import SiderLayout from '@/components/layouts/SiderLayout'; +import PageLayout from '@/components/layouts/PageLayout'; +import { + CREATE_USER, + LIST_RBAC_USERS, + UPDATE_USER, + UPDATE_USER_ROLES, +} from '@/apollo/client/graphql/rbac'; +import { + Role, + RoleTags, + StatusTag, + User, +} from '@/components/pages/administration/types'; +import { getAbsoluteTime } from '@/utils/time'; + +const { Text } = Typography; + +type UserModalState = { + visible: boolean; + user?: User; +}; + +const UserModal = ({ + roles, + state, + loading, + onClose, + onSubmit, +}: { + roles: Role[]; + state: UserModalState; + loading: boolean; + onClose: () => void; + onSubmit: (values: any, user?: User) => Promise; +}) => { + const [form] = Form.useForm(); + const isEdit = !!state.user; + + useEffect(() => { + if (!state.visible) return; + form.setFieldsValue({ + name: state.user?.name, + email: state.user?.email, + externalId: state.user?.externalId, + identityProvider: state.user?.identityProvider, + isActive: state.user?.isActive ?? true, + roleIds: state.user?.roles?.map((role) => role.id) || [], + }); + }, [form, state.visible, state.user]); + + const submit = async () => { + const values = await form.validateFields(); + await onSubmit(values, state.user); + form.resetFields(); + onClose(); + }; + + return ( + form.resetFields()} + > +
+ + + + + + + + + + + + + + + + + + ); +}; + +export default function UserManagement() { + const [modalState, setModalState] = useState({ + visible: false, + }); + const { data, loading, refetch } = useQuery(LIST_RBAC_USERS, { + fetchPolicy: 'cache-and-network', + }); + const users: User[] = data?.users || []; + const roles: Role[] = data?.roles || []; + + const mutationOptions = { + onError: (error) => message.error(error.message), + }; + const [createUser, createUserState] = useMutation( + CREATE_USER, + mutationOptions, + ); + const [updateUser, updateUserState] = useMutation( + UPDATE_USER, + mutationOptions, + ); + const [updateUserRoles, updateRolesState] = useMutation( + UPDATE_USER_ROLES, + mutationOptions, + ); + + const closeModal = () => setModalState({ visible: false }); + + const submitUser = async (values: any, user?: User) => { + if (user) { + await updateUser({ + variables: { + where: { id: user.id }, + data: { + name: values.name, + email: values.email, + externalId: values.externalId || null, + identityProvider: values.identityProvider || null, + isActive: values.isActive, + }, + }, + }); + await updateUserRoles({ + variables: { + data: { userId: user.id, roleIds: values.roleIds || [] }, + }, + }); + message.success('Successfully updated user.'); + } else { + await createUser({ + variables: { + data: { + name: values.name, + email: values.email, + externalId: values.externalId || null, + identityProvider: values.identityProvider || null, + isActive: values.isActive, + roleIds: values.roleIds || [], + }, + }, + }); + message.success('Successfully created user.'); + } + await refetch(); + }; + + const columns: TableColumnsType = [ + { + title: 'User', + dataIndex: 'name', + render: (_, record) => ( +
+
{record.name}
+ {record.email} +
+ ), + }, + { + title: 'Status', + dataIndex: 'isActive', + width: 120, + render: (active) => , + }, + { + title: 'Assigned roles', + dataIndex: 'roles', + render: (assignedRoles) => , + }, + { + title: 'Identity provider', + dataIndex: 'identityProvider', + width: 180, + render: (value) => value || Local, + }, + { + title: 'Created', + dataIndex: 'createdAt', + width: 180, + render: (value) => ( + {getAbsoluteTime(value)} + ), + }, + { + title: 'Actions', + width: 110, + align: 'center', + render: (_, record) => ( + + ), + }, + ]; + + return ( + + + + User Management + + } + description="Create users, maintain local identity metadata, and review assigned foundation roles." + titleExtra={ + + } + > +
+ + + + ); +} diff --git a/wren-ui/src/pages/api/ask_task/streaming_answer.ts b/wren-ui/src/pages/api/ask_task/streaming_answer.ts index c878bbae94..d63a083a59 100644 --- a/wren-ui/src/pages/api/ask_task/streaming_answer.ts +++ b/wren-ui/src/pages/api/ask_task/streaming_answer.ts @@ -29,6 +29,24 @@ class ContentMap { const contentMap = new ContentMap(); +const parseSSEMessages = (chunk: Buffer): string[] => { + return chunk + .toString('utf-8') + .split(/\r?\n/) + .filter((line) => line.startsWith('data: ')) + .map((line) => line.slice('data: '.length).trim()) + .filter(Boolean) + .flatMap((payload) => { + try { + const eventData = JSON.parse(payload); + return eventData?.message ? [String(eventData.message)] : []; + } catch (error) { + console.error(`Failed to parse streaming answer payload: ${payload}`); + return []; + } + }); +}; + export default async function handler( req: NextApiRequest, res: NextApiResponse, @@ -69,29 +87,47 @@ export default async function handler( } const stream = await wrenAIAdaptor.streamTextBasedAnswer(queryId); + let streamEnded = false; stream.on('data', (chunk) => { - // pass the chunk directly to the client - const chunkString = chunk.toString('utf-8'); - let message = ''; - const match = chunkString.match(/data: {"message":"([\s\S]*?)"}/); - if (match && match[1]) { - message = match[1]; - } else { - console.log(`not able to match: ${chunkString}`); + for (const message of parseSSEMessages(chunk)) { + contentMap.appendContent(queryId, message); } - contentMap.appendContent(queryId, message); res.write(chunk); }); stream.on('end', () => { + streamEnded = true; + const finalContent = contentMap.getContent(queryId)?.trim(); res.write(`data: ${JSON.stringify({ done: true })}\n\n`); res.end(); + if (!finalContent) { + askingService + .changeThreadResponseAnswerDetailStatus( + Number(responseId), + ThreadResponseAnswerStatus.FAILED, + ) + .then(() => { + console.error( + `Thread response ${responseId} answer stream ended without content`, + ); + contentMap.remove(queryId); + }) + .catch((error) => { + console.error( + 'Failed to update empty thread response answer detail status', + error, + ); + contentMap.remove(queryId); + }); + return; + } + askingService .changeThreadResponseAnswerDetailStatus( Number(responseId), ThreadResponseAnswerStatus.FINISHED, - contentMap.getContent(queryId), + finalContent, ) .then(() => { console.log( @@ -122,6 +158,9 @@ export default async function handler( // destroy the stream if the client closes the connection req.on('close', () => { + if (streamEnded) { + return; + } stream.destroy(); askingService .changeThreadResponseAnswerDetailStatus( diff --git a/wren-ui/src/pages/api/graphql.ts b/wren-ui/src/pages/api/graphql.ts index 60ed8334f4..bb5cd2e5f8 100644 --- a/wren-ui/src/pages/api/graphql.ts +++ b/wren-ui/src/pages/api/graphql.ts @@ -13,7 +13,7 @@ import { GeneralErrorCodes, } from '@/apollo/server/utils/error'; import { TelemetryEvent } from '@/apollo/server/telemetry/telemetry'; -import { components } from '@/common'; +import { components, initComponents } from '@/common'; const serverConfig = getConfig(); const logger = getLogger('APOLLO'); @@ -28,6 +28,7 @@ export const config: PageConfig = { }; const bootstrapServer = async () => { + const componentGraph = components ?? initComponents(); const { telemetry, @@ -47,6 +48,10 @@ const bootstrapServer = async () => { instructionRepository, apiHistoryRepository, dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, + organizationRepository, // adaptors wrenEngineAdaptor, ibisAdaptor, @@ -62,11 +67,13 @@ const bootstrapServer = async () => { sqlPairService, instructionService, + rbacService, + organizationService, // background trackers projectRecommendQuestionBackgroundTracker, threadRecommendQuestionBackgroundTracker, dashboardCacheBackgroundTracker, - } = components; + } = componentGraph; const modelService = new ModelService({ projectService, @@ -126,42 +133,48 @@ const bootstrapServer = async () => { }, introspection: process.env.NODE_ENV !== 'production', context: (): IContext => ({ - config: serverConfig, - telemetry, - // adaptor - wrenEngineAdaptor, - ibisServerAdaptor: ibisAdaptor, - wrenAIAdaptor, - // services - projectService, - modelService, - mdlService, - deployService, - askingService, - queryService, - dashboardService, - sqlPairService, - instructionService, - // repository - projectRepository, - modelRepository, - modelColumnRepository, - modelNestedColumnRepository, - relationRepository, - viewRepository, - deployRepository: deployLogRepository, - schemaChangeRepository, - learningRepository, - dashboardRepository, - dashboardItemRepository, - sqlPairRepository, - instructionRepository, - apiHistoryRepository, - dashboardItemRefreshJobRepository, - // background trackers - projectRecommendQuestionBackgroundTracker, - threadRecommendQuestionBackgroundTracker, - dashboardCacheBackgroundTracker, + config: serverConfig, + telemetry, + // adaptor + wrenEngineAdaptor, + ibisServerAdaptor: ibisAdaptor, + wrenAIAdaptor, + // services + projectService, + modelService, + mdlService, + deployService, + askingService, + queryService, + dashboardService, + sqlPairService, + instructionService, + rbacService, + organizationService, + // repository + projectRepository, + modelRepository, + modelColumnRepository, + modelNestedColumnRepository, + relationRepository, + viewRepository, + deployRepository: deployLogRepository, + schemaChangeRepository, + learningRepository, + dashboardRepository, + dashboardItemRepository, + sqlPairRepository, + instructionRepository, + apiHistoryRepository, + dashboardItemRefreshJobRepository, + roleRepository, + userRepository, + userRoleRepository, + organizationRepository, + // background trackers + projectRecommendQuestionBackgroundTracker, + threadRecommendQuestionBackgroundTracker, + dashboardCacheBackgroundTracker, }), }); await apolloServer.start(); diff --git a/wren-ui/src/pages/api/v1/ask.ts b/wren-ui/src/pages/api/v1/ask.ts index d2e98d0eb2..b95cd1b783 100644 --- a/wren-ui/src/pages/api/v1/ask.ts +++ b/wren-ui/src/pages/api/v1/ask.ts @@ -30,6 +30,7 @@ const { apiHistoryRepository, projectService, deployService, + mdlService, wrenAIAdaptor, queryService, } = components; @@ -62,7 +63,6 @@ export default async function handler( throw new ApiError('Question is required', 400); } - // Get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); if (!lastDeploy) { throw new ApiError( @@ -71,6 +71,11 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const { manifest } = await mdlService.makeModelMDL(project); + const deployId = await deployService.ensureDeploymentPrepared( + project.id, + manifest, + ); // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); @@ -83,7 +88,8 @@ export default async function handler( // Step 1: Generate SQL const askTask = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: @@ -195,7 +201,7 @@ export default async function handler( const queryResult = await queryService.preview(sql, { project, limit: sampleSize || 500, - manifest: lastDeploy.manifest, + manifest, modelingOnly: false, }); sqlData = queryResult; diff --git a/wren-ui/src/pages/api/v1/generate_sql.ts b/wren-ui/src/pages/api/v1/generate_sql.ts index fa5b3859e1..7fd6904db6 100644 --- a/wren-ui/src/pages/api/v1/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/generate_sql.ts @@ -23,6 +23,7 @@ const { apiHistoryRepository, projectService, deployService, + mdlService, wrenAIAdaptor, wrenEngineAdaptor, ibisAdaptor, @@ -61,9 +62,7 @@ export default async function handler( throw new ApiError('Question is required', 400); } - // get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); - if (!lastDeploy) { throw new ApiError( 'No deployment found, please deploy a model first', @@ -71,6 +70,11 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const { manifest } = await mdlService.makeModelMDL(project); + const deployId = await deployService.ensureDeploymentPrepared( + project.id, + manifest, + ); // ask AI service to generate SQL const histories = threadId @@ -78,7 +82,8 @@ export default async function handler( : undefined; const task = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: @@ -120,14 +125,14 @@ export default async function handler( let nativeSql: string; if (project.type === DataSourceName.DUCKDB) { nativeSql = await wrenEngineAdaptor.getNativeSQL(sql, { - manifest: lastDeploy.manifest, + manifest, modelingOnly: false, }); } else { nativeSql = await ibisAdaptor.getNativeSql({ dataSource: project.type, sql, - mdl: lastDeploy.manifest, + mdl: manifest, }); } diff --git a/wren-ui/src/pages/api/v1/generate_summary.ts b/wren-ui/src/pages/api/v1/generate_summary.ts index 2ca538ce51..a7b6b103a0 100644 --- a/wren-ui/src/pages/api/v1/generate_summary.ts +++ b/wren-ui/src/pages/api/v1/generate_summary.ts @@ -21,8 +21,13 @@ import { getLogger } from '@server/utils'; const logger = getLogger('API_GENERATE_SUMMARY'); logger.level = 'debug'; -const { projectService, wrenAIAdaptor, deployService, queryService } = - components; +const { + projectService, + wrenAIAdaptor, + deployService, + mdlService, + queryService, +} = components; interface GenerateSummaryRequest { question: string; @@ -67,6 +72,7 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const { manifest } = await mdlService.makeModelMDL(project); // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); @@ -77,7 +83,7 @@ export default async function handler( const queryResult = await queryService.preview(sql, { project, limit: sampleSize || 500, - manifest: lastDeploy.manifest, + manifest, modelingOnly: false, }); sqlData = queryResult; diff --git a/wren-ui/src/pages/api/v1/generate_vega_chart.ts b/wren-ui/src/pages/api/v1/generate_vega_chart.ts index 1d52692b04..fd4de61127 100644 --- a/wren-ui/src/pages/api/v1/generate_vega_chart.ts +++ b/wren-ui/src/pages/api/v1/generate_vega_chart.ts @@ -17,8 +17,13 @@ import { PreviewDataResponse } from '@server/services/queryService'; import { transformToObjects } from '@server/utils/dataUtils'; import { enhanceVegaSpec } from '@/utils/vegaSpecUtils'; -const { projectService, wrenAIAdaptor, deployService, queryService } = - components; +const { + projectService, + wrenAIAdaptor, + deployService, + mdlService, + queryService, +} = components; const MAX_WAIT_TIME = 1000 * 60 * 3; // 3 minutes @@ -97,6 +102,7 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const { manifest } = await mdlService.makeModelMDL(project); // Execute the SQL query to get the data let queryResult: PreviewDataResponse; @@ -104,7 +110,7 @@ export default async function handler( queryResult = (await queryService.preview(sql, { project, limit: sampleSize, - manifest: lastDeploy.manifest, + manifest, modelingOnly: false, })) as PreviewDataResponse; } catch (queryError) { @@ -125,6 +131,7 @@ export default async function handler( const task = await wrenAIAdaptor.generateChart({ query: question, sql, + data: queryResult as unknown as Record, projectId: project.id.toString(), configurations: { language: WrenAILanguage[project.language] || WrenAILanguage.EN, diff --git a/wren-ui/src/pages/api/v1/organizations/[id]/select.ts b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts new file mode 100644 index 0000000000..f97abe3d78 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/[id]/select.ts @@ -0,0 +1,71 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + parseOrganizationId, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_SELECT_ORGANIZATION'); +logger.level = 'debug'; + +const getOrganizationService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['POST']); + const organizationService = getOrganizationService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + const organizationId = parseOrganizationId(req.query.id); + const organization = + await organizationService.selectCurrentOrganization(organizationId); + + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + id: organization.id, + name: organization.name, + identifier: organization.identifier, + description: organization.description, + isCurrent: Boolean(organization.isCurrent), + createdAt: organization.createdAt, + updatedAt: organization.updatedAt, + }, + projectId, + apiType: ApiType.SELECT_ORGANIZATION, + startTime, + requestPayload: { id: organizationId }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: ApiType.SELECT_ORGANIZATION, + requestPayload: { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/current.ts b/wren-ui/src/pages/api/v1/organizations/current.ts new file mode 100644 index 0000000000..5902ea7d47 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/current.ts @@ -0,0 +1,130 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + getCurrentProjectName, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_CURRENT_ORGANIZATION'); +logger.level = 'debug'; + +const getComponentGraph = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph; +}; + +const getOrganizationService = () => { + return getComponentGraph().organizationService; +}; + +const getOrganizationMemberService = () => { + return getComponentGraph().organizationMemberService; +}; + +const serializeOrganization = (organization) => + organization + ? { + id: organization.id, + name: organization.name, + identifier: organization.identifier, + description: organization.description, + isCurrent: Boolean(organization.isCurrent), + createdAt: organization.createdAt, + updatedAt: organization.updatedAt, + } + : null; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'PUT', 'DELETE']); + const organizationService = getOrganizationService(); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + if (req.method === 'PUT') { + const organization = await organizationService.updateCurrentOrganization( + req.body, + ); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: serializeOrganization(organization), + projectId, + apiType: ApiType.UPDATE_CURRENT_ORGANIZATION, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + return; + } + + if (req.method === 'DELETE') { + await organizationMemberService.deleteCurrentOrganization(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.DELETE_CURRENT_ORGANIZATION, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + const [currentOrganization, organizations, currentProjectName] = + await Promise.all([ + organizationService.getCurrentOrganization(), + organizationService.listOrganizations(), + getCurrentProjectName(), + ]); + + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + currentOrganization: serializeOrganization(currentOrganization), + organizations: organizations.map(serializeOrganization), + currentProjectName, + }, + projectId, + apiType: ApiType.GET_CURRENT_ORGANIZATION, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'PUT' + ? ApiType.UPDATE_CURRENT_ORGANIZATION + : req.method === 'DELETE' + ? ApiType.DELETE_CURRENT_ORGANIZATION + : ApiType.GET_CURRENT_ORGANIZATION, + requestPayload: req.method === 'PUT' ? req.body : {}, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/index.ts b/wren-ui/src/pages/api/v1/organizations/index.ts new file mode 100644 index 0000000000..5973b70da1 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/index.ts @@ -0,0 +1,92 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + getCurrentProjectName, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATIONS'); +logger.level = 'debug'; + +const getOrganizationService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationService; +}; + +const serializeOrganization = (organization) => ({ + id: organization.id, + name: organization.name, + identifier: organization.identifier, + description: organization.description, + isCurrent: Boolean(organization.isCurrent), + createdAt: organization.createdAt, + updatedAt: organization.updatedAt, +}); + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'POST']); + const organizationService = getOrganizationService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + if (req.method === 'GET') { + const organizations = await organizationService.listOrganizations(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + organizations: organizations.map(serializeOrganization), + currentProjectName: await getCurrentProjectName(), + }, + projectId, + apiType: ApiType.GET_ORGANIZATIONS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + const organization = await organizationService.createOrganization(req.body); + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: serializeOrganization(organization), + projectId, + apiType: ApiType.CREATE_ORGANIZATION, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'GET' + ? ApiType.GET_ORGANIZATIONS + : ApiType.CREATE_ORGANIZATION, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/invitations/[id].ts b/wren-ui/src/pages/api/v1/organizations/invitations/[id].ts new file mode 100644 index 0000000000..9174e8ecd1 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/invitations/[id].ts @@ -0,0 +1,62 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + parseOrganizationId, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATION_INVITATION'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['DELETE']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + const invitationId = parseOrganizationId(req.query.id); + + await organizationMemberService.removeInvitation(invitationId); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.REMOVE_ORGANIZATION_INVITATION, + startTime, + requestPayload: { id: invitationId }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: ApiType.REMOVE_ORGANIZATION_INVITATION, + requestPayload: { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/invitations/accept/[token].ts b/wren-ui/src/pages/api/v1/organizations/invitations/accept/[token].ts new file mode 100644 index 0000000000..e53a1d81f8 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/invitations/accept/[token].ts @@ -0,0 +1,65 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ACCEPT_ORGANIZATION_INVITATION'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['POST']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + const token = Array.isArray(req.query.token) ? req.query.token[0] : req.query.token; + + if (!token) { + throw new Error('Invitation token is required'); + } + + const member = await organizationMemberService.acceptInvitation(token); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: member, + projectId, + apiType: ApiType.ACCEPT_ORGANIZATION_INVITATION, + startTime, + requestPayload: { token }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: ApiType.ACCEPT_ORGANIZATION_INVITATION, + requestPayload: { token: req.query.token }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/members/[id].ts b/wren-ui/src/pages/api/v1/organizations/members/[id].ts new file mode 100644 index 0000000000..88c6d290c8 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/members/[id].ts @@ -0,0 +1,83 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + parseOrganizationId, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATION_MEMBER'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['PATCH', 'DELETE']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + const memberId = parseOrganizationId(req.query.id); + + if (req.method === 'PATCH') { + const member = await organizationMemberService.updateMember( + memberId, + req.body, + ); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: member, + projectId, + apiType: ApiType.UPDATE_ORGANIZATION_MEMBER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + return; + } + + await organizationMemberService.removeMember(memberId); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.REMOVE_ORGANIZATION_MEMBER, + startTime, + requestPayload: { id: memberId }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'PATCH' + ? ApiType.UPDATE_ORGANIZATION_MEMBER + : ApiType.REMOVE_ORGANIZATION_MEMBER, + requestPayload: req.method === 'PATCH' ? req.body : { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/members/current.ts b/wren-ui/src/pages/api/v1/organizations/members/current.ts new file mode 100644 index 0000000000..2d2ad798db --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/members/current.ts @@ -0,0 +1,60 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_CURRENT_ORGANIZATION_MEMBER'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['DELETE']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + await organizationMemberService.leaveCurrentOrganization(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.LEAVE_ORGANIZATION, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: ApiType.LEAVE_ORGANIZATION, + requestPayload: {}, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/organizations/members/index.ts b/wren-ui/src/pages/api/v1/organizations/members/index.ts new file mode 100644 index 0000000000..5cdc25cb71 --- /dev/null +++ b/wren-ui/src/pages/api/v1/organizations/members/index.ts @@ -0,0 +1,103 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_ORGANIZATION_MEMBERS'); +logger.level = 'debug'; + +const getBaseUrl = (req: NextApiRequest) => { + const protocol = + (req.headers['x-forwarded-proto'] as string | undefined) || 'http'; + const host = + (req.headers['x-forwarded-host'] as string | undefined) || req.headers.host; + return `${protocol}://${host}`; +}; + +const withInviteLinks = ( + payload: any, + baseUrl: string, +) => ({ + ...payload, + invitations: (payload?.invitations || []).map((invitation) => ({ + ...invitation, + inviteLink: `${baseUrl}/organization/invitations/${invitation.token}`, + })), +}); + +const withInviteLink = (payload: T, baseUrl: string) => ({ + ...payload, + inviteLink: `${baseUrl}/organization/invitations/${payload.token}`, +}); + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'POST']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + if (req.method === 'GET') { + const payload = + await organizationMemberService.listCurrentOrganizationMembers(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: withInviteLinks(payload, getBaseUrl(req)), + projectId, + apiType: ApiType.GET_ORGANIZATION_MEMBERS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + const member = await organizationMemberService.inviteMember(req.body); + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: withInviteLink(member, getBaseUrl(req)), + projectId, + apiType: ApiType.INVITE_ORGANIZATION_MEMBER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'GET' + ? ApiType.GET_ORGANIZATION_MEMBERS + : ApiType.INVITE_ORGANIZATION_MEMBER, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/projects/[id]/select.ts b/wren-ui/src/pages/api/v1/projects/[id]/select.ts new file mode 100644 index 0000000000..7aef051872 --- /dev/null +++ b/wren-ui/src/pages/api/v1/projects/[id]/select.ts @@ -0,0 +1,77 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + ApiError, + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_SELECT_PROJECT'); +logger.level = 'debug'; + +const getProjectService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.projectService; +}; + +const parseProjectId = (value: string | string[] | undefined) => { + const rawValue = Array.isArray(value) ? value[0] : value; + if (!rawValue || !/^\d+$/.test(rawValue)) { + throw new ApiError('Invalid project id', 400); + } + return rawValue; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + if (req.method !== 'POST') { + throw new ApiError('Method not allowed', 405); + } + + const projectId = parseProjectId(req.query.id); + const projectService = getProjectService(); + logger.debug(`API select project request for ${projectId}`); + const project = await projectService.selectProject(projectId); + logger.debug( + `API selected project ${String(project.id)} (${project.type || 'unknown'})`, + ); + + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + project: { + id: project.id, + displayName: project.displayName, + projectType: project.projectType || 'CLASSIC', + }, + }, + projectId: Number.isSafeInteger(Number(projectId)) ? Number(projectId) : 0, + apiType: ApiType.SELECT_PROJECT, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: 0, + apiType: ApiType.SELECT_PROJECT, + requestPayload: {}, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/projects/access/[id].ts b/wren-ui/src/pages/api/v1/projects/access/[id].ts new file mode 100644 index 0000000000..ce3244638e --- /dev/null +++ b/wren-ui/src/pages/api/v1/projects/access/[id].ts @@ -0,0 +1,94 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, + parseOrganizationId, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_PROJECT_ACCESS_MEMBER'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components, initComponents } = require('@/common'); + let componentGraph = components ?? globalThis.__wrenComponents; + + if ( + !componentGraph?.organizationMemberService || + typeof componentGraph.organizationMemberService.listCurrentProjectAccess !== + 'function' + ) { + componentGraph = initComponents(); + globalThis.__wrenComponents = componentGraph; + } + + if (!componentGraph?.organizationMemberService) { + throw new Error('Organization member service is not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['PATCH', 'DELETE']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + const memberId = parseOrganizationId(req.query.id); + + if (req.method === 'PATCH') { + const member = + await organizationMemberService.updateProjectMemberPermission( + memberId, + req.body, + ); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: member, + projectId, + apiType: ApiType.UPDATE_PROJECT_MEMBER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + return; + } + + await organizationMemberService.removeProjectMember(memberId); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.REMOVE_PROJECT_MEMBER, + startTime, + requestPayload: { id: memberId }, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'PATCH' + ? ApiType.UPDATE_PROJECT_MEMBER + : ApiType.REMOVE_PROJECT_MEMBER, + requestPayload: req.method === 'PATCH' ? req.body : { id: req.query.id }, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/projects/access/current.ts b/wren-ui/src/pages/api/v1/projects/access/current.ts new file mode 100644 index 0000000000..ba5915b875 --- /dev/null +++ b/wren-ui/src/pages/api/v1/projects/access/current.ts @@ -0,0 +1,88 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_PROJECT_ACCESS'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components, initComponents } = require('@/common'); + let componentGraph = components ?? globalThis.__wrenComponents; + + if ( + !componentGraph?.organizationMemberService || + typeof componentGraph.organizationMemberService.listCurrentProjectAccess !== + 'function' + ) { + componentGraph = initComponents(); + globalThis.__wrenComponents = componentGraph; + } + + if (!componentGraph?.organizationMemberService) { + throw new Error('Organization member service is not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'POST']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + if (req.method === 'GET') { + const payload = await organizationMemberService.listCurrentProjectAccess(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: payload, + projectId, + apiType: ApiType.GET_PROJECT_ACCESS, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + const member = await organizationMemberService.addProjectMember(req.body); + await respondWithSimple({ + res, + statusCode: 201, + responsePayload: member, + projectId, + apiType: ApiType.ADD_PROJECT_MEMBER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'GET' + ? ApiType.GET_PROJECT_ACCESS + : ApiType.ADD_PROJECT_MEMBER, + requestPayload: req.method === 'GET' ? {} : req.body, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/projects/current.ts b/wren-ui/src/pages/api/v1/projects/current.ts new file mode 100644 index 0000000000..852f028440 --- /dev/null +++ b/wren-ui/src/pages/api/v1/projects/current.ts @@ -0,0 +1,111 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { coerceBoolean } from '@server/repositories/baseRepository'; +import { + ApiError, + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; + +const logger = getLogger('API_CURRENT_PROJECT'); +logger.level = 'debug'; + +const getProjectService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.projectService; +}; + +const serializeProject = (project) => + project + ? { + id: String(project.id), + displayName: project.displayName, + projectType: project.projectType || 'CLASSIC', + isCurrent: coerceBoolean(project.isCurrent), + hasDataSource: Boolean(project.type), + type: project.type || null, + createdAt: project.createdAt, + updatedAt: project.updatedAt, + } + : null; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + if (req.method !== 'GET') { + throw new ApiError('Method not allowed', 405); + } + + const projectService = getProjectService(); + let currentProject = null; + try { + currentProject = await projectService.getCurrentProject(); + logger.debug( + `Resolved current project ${String(currentProject?.id)} (${currentProject?.type || 'unknown'})`, + ); + } catch (error) { + logger.warn( + `Failed to resolve current project through project service: ${ + error instanceof Error ? error.message : String(error) + }`, + ); + currentProject = null; + } + const projects = await projectService.listProjects(); + logger.debug( + `Loaded ${projects.length} project(s) for current project API: ${projects + .map((project) => `${String(project.id)}:${project.type || 'unknown'}:${coerceBoolean(project.isCurrent)}`) + .join(', ') || 'none'}`, + ); + const serializedCurrentProject = + currentProject && + projects.some( + (project) => String(project.id) === String(currentProject.id), + ) + ? currentProject + : projects.find((project) => coerceBoolean(project.isCurrent)) || null; + logger.debug( + `Current project API returning ${ + serializedCurrentProject + ? `${String(serializedCurrentProject.id)} (${serializedCurrentProject.type || 'unknown'})` + : 'no current project' + }`, + ); + + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { + currentProject: serializeProject(serializedCurrentProject), + projects: projects.map(serializeProject), + }, + projectId: Number.isSafeInteger(Number(serializedCurrentProject?.id)) + ? Number(serializedCurrentProject?.id) + : 0, + apiType: ApiType.GET_CURRENT_PROJECT, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: 0, + apiType: ApiType.GET_CURRENT_PROJECT, + requestPayload: {}, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/api/v1/run_sql.ts b/wren-ui/src/pages/api/v1/run_sql.ts index 0154d05d16..dcdb3cb63d 100644 --- a/wren-ui/src/pages/api/v1/run_sql.ts +++ b/wren-ui/src/pages/api/v1/run_sql.ts @@ -15,7 +15,7 @@ import { transformToObjects } from '@server/utils/dataUtils'; const logger = getLogger('API_RUN_SQL'); logger.level = 'debug'; -const { projectService, queryService, deployService } = components; +const { projectService, queryService, deployService, mdlService } = components; /** * Validates the SQL result and ensures it has the expected format @@ -69,7 +69,7 @@ export default async function handler( ); } - const manifest = deployment.manifest; + const { manifest } = await mdlService.makeModelMDL(project); // Execute the SQL query try { diff --git a/wren-ui/src/pages/api/v1/stream/ask.ts b/wren-ui/src/pages/api/v1/stream/ask.ts index 12e53eb4db..afcfab5e76 100644 --- a/wren-ui/src/pages/api/v1/stream/ask.ts +++ b/wren-ui/src/pages/api/v1/stream/ask.ts @@ -44,6 +44,7 @@ const { apiHistoryRepository, projectService, deployService, + mdlService, wrenAIAdaptor, queryService, } = components; @@ -123,7 +124,6 @@ export default async function handler( // Send message start event sendMessageStart(res); - // Get current project's last deployment const lastDeploy = await deployService.getLastDeployment(project.id); if (!lastDeploy) { throw new ApiError( @@ -132,6 +132,11 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const { manifest } = await mdlService.makeModelMDL(project); + const deployId = await deployService.ensureDeploymentPrepared( + project.id, + manifest, + ); // Create a new thread if it's a new question const newThreadId = threadId || uuidv4(); @@ -150,7 +155,8 @@ export default async function handler( }); const askTask = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: @@ -309,7 +315,7 @@ export default async function handler( const queryResult = await queryService.preview(sql, { project, limit: sampleSize || 500, - manifest: lastDeploy.manifest, + manifest, modelingOnly: false, }); sqlData = queryResult; diff --git a/wren-ui/src/pages/api/v1/stream/generate_sql.ts b/wren-ui/src/pages/api/v1/stream/generate_sql.ts index 6102592657..0c18bf4c22 100644 --- a/wren-ui/src/pages/api/v1/stream/generate_sql.ts +++ b/wren-ui/src/pages/api/v1/stream/generate_sql.ts @@ -29,8 +29,13 @@ import { const logger = getLogger('API_STREAM_GENERATE_SQL'); logger.level = 'debug'; -const { apiHistoryRepository, projectService, deployService, wrenAIAdaptor } = - components; +const { + apiHistoryRepository, + projectService, + deployService, + mdlService, + wrenAIAdaptor, +} = components; export default async function handler( req: NextApiRequest, @@ -74,6 +79,11 @@ export default async function handler( Errors.GeneralErrorCodes.NO_DEPLOYMENT_FOUND, ); } + const { manifest } = await mdlService.makeModelMDL(project); + const deployId = await deployService.ensureDeploymentPrepared( + project.id, + manifest, + ); // Get conversation history if threadId is provided const histories = threadId @@ -90,7 +100,8 @@ export default async function handler( const askTask = await wrenAIAdaptor.ask({ query: question, - deployId: lastDeploy.hash, + deployId, + projectId: project.id.toString(), histories: transformHistoryInput(histories) as any, configurations: { language: diff --git a/wren-ui/src/pages/api/v1/users/current.ts b/wren-ui/src/pages/api/v1/users/current.ts new file mode 100644 index 0000000000..6e9346d2d0 --- /dev/null +++ b/wren-ui/src/pages/api/v1/users/current.ts @@ -0,0 +1,97 @@ +import { NextApiRequest, NextApiResponse } from 'next'; +import { ApiType } from '@server/repositories/apiHistoryRepository'; +import { + handleApiError, + respondWithSimple, +} from '@/apollo/server/utils/apiUtils'; +import { getLogger } from '@server/utils'; +import { + assertAllowedMethods, + getCurrentProjectContext, +} from '@/apollo/server/middlewares/organizationApi'; + +const logger = getLogger('API_CURRENT_USER'); +logger.level = 'debug'; + +const getOrganizationMemberService = () => { + const { components } = require('@/common'); + const componentGraph = components ?? globalThis.__wrenComponents; + if (!componentGraph) { + throw new Error('Components are not initialized'); + } + return componentGraph.organizationMemberService; +}; + +export default async function handler( + req: NextApiRequest, + res: NextApiResponse, +) { + const startTime = Date.now(); + + try { + assertAllowedMethods(req, ['GET', 'PUT', 'DELETE']); + const organizationMemberService = getOrganizationMemberService(); + const projectContext = await getCurrentProjectContext(); + const projectId = projectContext.id ?? 0; + + if (req.method === 'DELETE') { + await organizationMemberService.deleteCurrentUserAccount(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: { success: true }, + projectId, + apiType: ApiType.DELETE_CURRENT_USER, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + return; + } + + if (req.method === 'PUT') { + const user = await organizationMemberService.updateCurrentUserProfile( + req.body, + ); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: user, + projectId, + apiType: ApiType.UPDATE_CURRENT_USER, + startTime, + requestPayload: req.body, + headers: req.headers as Record, + }); + return; + } + + const user = await organizationMemberService.getCurrentUserProfile(); + await respondWithSimple({ + res, + statusCode: 200, + responsePayload: user, + projectId, + apiType: ApiType.GET_CURRENT_USER, + startTime, + requestPayload: {}, + headers: req.headers as Record, + }); + } catch (error) { + await handleApiError({ + error, + res, + projectId: (await getCurrentProjectContext()).id ?? 0, + apiType: + req.method === 'DELETE' + ? ApiType.DELETE_CURRENT_USER + : req.method === 'PUT' + ? ApiType.UPDATE_CURRENT_USER + : ApiType.GET_CURRENT_USER, + requestPayload: req.method === 'PUT' ? req.body : {}, + headers: req.headers as Record, + startTime, + logger, + }); + } +} diff --git a/wren-ui/src/pages/home/[id].tsx b/wren-ui/src/pages/home/[id].tsx index e27fe1f8b2..1d107a397b 100644 --- a/wren-ui/src/pages/home/[id].tsx +++ b/wren-ui/src/pages/home/[id].tsx @@ -2,10 +2,10 @@ import { useRouter } from 'next/router'; import { useParams } from 'next/navigation'; import { ComponentRef, + useRef, useCallback, useEffect, useMemo, - useRef, useState, } from 'react'; import { isEmpty } from 'lodash'; @@ -70,6 +70,11 @@ const getThreadResponseIsFinished = (threadResponse: ThreadResponse) => { return isAnswerFinished !== false && isChartFinished !== false; }; +const THREAD_RESPONSE_POLL_INTERVAL_MS = 1000; +const THREAD_RECOMMENDATION_POLL_INTERVAL_MS = 1000; +const THREAD_RESPONSE_POLL_MAX_INTERVAL_MS = 10000; +const THREAD_RECOMMENDATION_POLL_MAX_INTERVAL_MS = 10000; + export default function HomeThread() { const $prompt = useRef>(null); const router = useRouter(); @@ -93,7 +98,8 @@ export default function HomeThread() { const { data, updateQuery: updateThreadQuery } = useThreadQuery({ variables: { threadId }, - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', skip: threadId === null, onError: () => router.push(Path.Home), }); @@ -101,7 +107,13 @@ export default function HomeThread() { onError: (error) => console.error(error), onCompleted(next) { const nextResponse = next.createThreadResponse; + if (!nextResponse) { + return; + } updateThreadQuery((prev) => { + if (!prev?.thread?.responses) { + return prev; + } return { ...prev, thread: { @@ -123,18 +135,27 @@ export default function HomeThread() { }); const [fetchThreadResponse, threadResponseResult] = useThreadResponseLazyQuery({ - pollInterval: 1000, + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', onCompleted(next) { const nextResponse = next.threadResponse; - updateThreadQuery((prev) => ({ - ...prev, - thread: { - ...prev.thread, - responses: prev.thread.responses.map((response) => - response.id === nextResponse.id ? nextResponse : response, - ), - }, - })); + if (!nextResponse) { + return; + } + updateThreadQuery((prev) => { + if (!prev?.thread?.responses) { + return prev; + } + return { + ...prev, + thread: { + ...prev.thread, + responses: prev.thread.responses.map((response) => + response.id === nextResponse.id ? nextResponse : response, + ), + }, + }; + }); }, }); @@ -147,8 +168,29 @@ export default function HomeThread() { fetchThreadRecommendationQuestions, threadRecommendationQuestionsResult, ] = useGetThreadRecommendationQuestionsLazyQuery({ - pollInterval: 1000, + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); + const threadResponsePollingRef = useRef | null>( + null, + ); + const threadResponsePollingSessionRef = useRef(0); + const threadResponsePollingTargetRef = useRef(null); + const threadResponsePollingRequestRef = useRef | null>(null); + const threadResponsePollingDelayRef = useRef(THREAD_RESPONSE_POLL_INTERVAL_MS); + const lastThreadResponseFingerprintRef = useRef(null); + const threadRecommendationPollingRef = useRef< + ReturnType | null + >(null); + const threadRecommendationPollingSessionRef = useRef(0); + const threadRecommendationPollingTargetRef = useRef(null); + const threadRecommendationPollingRequestRef = useRef | null>( + null, + ); + const threadRecommendationPollingDelayRef = useRef( + THREAD_RECOMMENDATION_POLL_INTERVAL_MS, + ); + const lastThreadRecommendationFingerprintRef = useRef(null); const [generateThreadResponseAnswer] = useGenerateThreadResponseAnswerMutation({ @@ -183,6 +225,150 @@ export default function HomeThread() { [pollingResponse], ); + const stopThreadResponsePolling = useCallback(() => { + threadResponsePollingSessionRef.current += 1; + threadResponsePollingTargetRef.current = null; + if (threadResponsePollingRef.current) { + clearTimeout(threadResponsePollingRef.current); + threadResponsePollingRef.current = null; + } + threadResponsePollingDelayRef.current = THREAD_RESPONSE_POLL_INTERVAL_MS; + }, []); + + const startThreadResponsePolling = useCallback( + async (responseId?: number) => { + if (!responseId) return; + if ( + threadResponsePollingTargetRef.current === responseId && + (threadResponsePollingRequestRef.current || threadResponsePollingRef.current) + ) { + return; + } + + stopThreadResponsePolling(); + threadResponsePollingTargetRef.current = responseId; + const pollingSessionId = threadResponsePollingSessionRef.current; + + const run = async () => { + if (threadResponsePollingSessionRef.current !== pollingSessionId) return; + if (threadResponsePollingRequestRef.current) { + await threadResponsePollingRequestRef.current; + if (threadResponsePollingSessionRef.current !== pollingSessionId) { + return; + } + } + + let shouldContinuePolling = true; + try { + const request = fetchThreadResponse({ + variables: { responseId }, + }); + threadResponsePollingRequestRef.current = request.then( + () => undefined, + ); + const result = await request; + if (getThreadResponseIsFinished(result.data?.threadResponse)) { + shouldContinuePolling = false; + stopThreadResponsePolling(); + setShowRecommendedQuestions(true); + } + } catch (error) { + console.error(error); + } finally { + threadResponsePollingRequestRef.current = null; + if ( + shouldContinuePolling && + threadResponsePollingSessionRef.current === pollingSessionId + ) { + threadResponsePollingRef.current = setTimeout( + run, + threadResponsePollingDelayRef.current, + ); + } + } + }; + + await run(); + }, + [fetchThreadResponse, stopThreadResponsePolling], + ); + + const stopThreadRecommendationPolling = useCallback(() => { + threadRecommendationPollingSessionRef.current += 1; + threadRecommendationPollingTargetRef.current = null; + if (threadRecommendationPollingRef.current) { + clearTimeout(threadRecommendationPollingRef.current); + threadRecommendationPollingRef.current = null; + } + threadRecommendationPollingDelayRef.current = + THREAD_RECOMMENDATION_POLL_INTERVAL_MS; + }, []); + + const startThreadRecommendationPolling = useCallback( + async (nextThreadId?: number) => { + if (!nextThreadId) return; + if ( + threadRecommendationPollingTargetRef.current === nextThreadId && + (threadRecommendationPollingRequestRef.current || + threadRecommendationPollingRef.current) + ) { + return; + } + + stopThreadRecommendationPolling(); + threadRecommendationPollingTargetRef.current = nextThreadId; + const pollingSessionId = threadRecommendationPollingSessionRef.current; + + const run = async () => { + if ( + threadRecommendationPollingSessionRef.current !== pollingSessionId + ) { + return; + } + if (threadRecommendationPollingRequestRef.current) { + await threadRecommendationPollingRequestRef.current; + if ( + threadRecommendationPollingSessionRef.current !== pollingSessionId + ) { + return; + } + } + + let shouldContinuePolling = true; + try { + const request = fetchThreadRecommendationQuestions({ + variables: { threadId: nextThreadId }, + }); + threadRecommendationPollingRequestRef.current = request.then( + () => undefined, + ); + const result = await request; + const task = result.data?.getThreadRecommendationQuestions; + if (!task || isRecommendedFinished(task.status)) { + shouldContinuePolling = false; + stopThreadRecommendationPolling(); + } + } catch (error) { + console.error(error); + } finally { + threadRecommendationPollingRequestRef.current = null; + if ( + shouldContinuePolling && + threadRecommendationPollingSessionRef.current === pollingSessionId + ) { + threadRecommendationPollingRef.current = setTimeout( + run, + threadRecommendationPollingDelayRef.current, + ); + } + } + }; + + await run(); + }, + [fetchThreadRecommendationQuestions, stopThreadRecommendationPolling], + ); + const onFixSQLStatement = async (responseId: number, sql: string) => { await updateThreadResponse({ variables: { where: { id: responseId }, data: { sql } }, @@ -190,57 +376,65 @@ export default function HomeThread() { }; const onGenerateThreadResponseAnswer = async (responseId: number) => { + if (!responseId) return; + await generateThreadResponseAnswer({ variables: { responseId } }); - fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; const onGenerateThreadResponseChart = async (responseId: number) => { + if (!responseId) return; + await generateThreadResponseChart({ variables: { responseId } }); - fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; const onAdjustThreadResponseChart = async ( responseId: number, data: AdjustThreadResponseChartInput, ) => { + if (!responseId) return; + await adjustThreadResponseChart({ variables: { responseId, data }, }); - fetchThreadResponse({ variables: { responseId } }); + await startThreadResponsePolling(responseId); }; const onGenerateThreadRecommendedQuestions = async () => { + if (!threadId) return; + await generateThreadRecommendationQuestions({ variables: { threadId } }); - fetchThreadRecommendationQuestions({ variables: { threadId } }); + await startThreadRecommendationPolling(threadId); }; const handleUnfinishedTasks = useCallback( (responses: ThreadResponse[]) => { + const latestResponse = [...(responses || [])].sort( + (a, b) => Number(b?.id || 0) - Number(a?.id || 0), + )[0]; + if (!latestResponse) return; + // unfinished asking task - const unfinishedAskingResponse = (responses || []).find( - (response) => - response?.askingTask && !getIsFinished(response?.askingTask?.status), - ); - if (unfinishedAskingResponse) { - askPrompt.onFetching(unfinishedAskingResponse?.askingTask?.queryId); + const unfinishedTaskId = + latestResponse?.askingTask && + !getIsFinished(latestResponse.askingTask.status) + ? latestResponse.askingTask.queryId + : null; + if (unfinishedTaskId) { + askPrompt.onFetching(unfinishedTaskId); return; } // unfinished thread response - const unfinishedThreadResponse = (responses || []).find( - (response) => !getThreadResponseIsFinished(response), - ); - if ( - canFetchThreadResponse(unfinishedThreadResponse?.askingTask) && - unfinishedThreadResponse + !getThreadResponseIsFinished(latestResponse) && + canFetchThreadResponse(latestResponse?.askingTask) ) { - fetchThreadResponse({ - variables: { responseId: unfinishedThreadResponse.id }, - }); + startThreadResponsePolling(latestResponse.id); } }, - [askPrompt, fetchThreadResponse], + [askPrompt, startThreadResponsePolling], ); // store thread questions for instant recommended questions @@ -255,13 +449,26 @@ export default function HomeThread() { // stop all requests when change thread useEffect(() => { if (threadId !== null) { - fetchThreadRecommendationQuestions({ variables: { threadId } }); setShowRecommendedQuestions(true); + void (async () => { + try { + const result = await fetchThreadRecommendationQuestions({ + variables: { threadId }, + }); + const status = + result.data?.getThreadRecommendationQuestions?.status || null; + if (status && !isRecommendedFinished(status)) { + await startThreadRecommendationPolling(threadId); + } + } catch (error) { + console.error(error); + } + })(); } return () => { askPrompt.onStopPolling(); - threadResponseResult.stopPolling(); - threadRecommendationQuestionsResult.stopPolling(); + stopThreadResponsePolling(); + stopThreadRecommendationPolling(); $prompt.current?.close(); }; }, [threadId]); @@ -275,11 +482,43 @@ export default function HomeThread() { useEffect(() => { if (isPollingResponseFinished) { - threadResponseResult.stopPolling(); + stopThreadResponsePolling(); setShowRecommendedQuestions(true); } }, [isPollingResponseFinished]); + useEffect(() => { + const fingerprint = JSON.stringify({ + id: pollingResponse?.id || null, + askingStatus: pollingResponse?.askingTask?.status || null, + askingType: pollingResponse?.askingTask?.type || null, + answerStatus: pollingResponse?.answerDetail?.status || null, + chartStatus: pollingResponse?.chartDetail?.status || null, + breakdownStatus: pollingResponse?.breakdownDetail?.status || null, + adjustmentStatus: pollingResponse?.adjustmentTask?.status || null, + sql: pollingResponse?.sql || null, + }); + + if (lastThreadResponseFingerprintRef.current === fingerprint) { + threadResponsePollingDelayRef.current = Math.min( + threadResponsePollingDelayRef.current * 2, + THREAD_RESPONSE_POLL_MAX_INTERVAL_MS, + ); + } else { + threadResponsePollingDelayRef.current = THREAD_RESPONSE_POLL_INTERVAL_MS; + lastThreadResponseFingerprintRef.current = fingerprint; + } + }, [ + pollingResponse?.id, + pollingResponse?.askingTask?.status, + pollingResponse?.askingTask?.type, + pollingResponse?.answerDetail?.status, + pollingResponse?.chartDetail?.status, + pollingResponse?.breakdownDetail?.status, + pollingResponse?.adjustmentTask?.status, + pollingResponse?.sql, + ]); + const recommendedQuestions = useMemo( () => threadRecommendationQuestionsResult.data @@ -289,15 +528,40 @@ export default function HomeThread() { useEffect(() => { if (isRecommendedFinished(recommendedQuestions?.status)) { - threadRecommendationQuestionsResult.stopPolling(); + stopThreadRecommendationPolling(); } }, [recommendedQuestions]); + useEffect(() => { + const fingerprint = JSON.stringify({ + status: recommendedQuestions?.status || null, + count: recommendedQuestions?.questions?.length || 0, + errorCode: recommendedQuestions?.error?.code || null, + }); + + if (lastThreadRecommendationFingerprintRef.current === fingerprint) { + threadRecommendationPollingDelayRef.current = Math.min( + threadRecommendationPollingDelayRef.current * 2, + THREAD_RECOMMENDATION_POLL_MAX_INTERVAL_MS, + ); + } else { + threadRecommendationPollingDelayRef.current = + THREAD_RECOMMENDATION_POLL_INTERVAL_MS; + lastThreadRecommendationFingerprintRef.current = fingerprint; + } + }, [ + recommendedQuestions?.status, + recommendedQuestions?.questions?.length, + recommendedQuestions?.error?.code, + ]); + const onCreateResponse = async (payload: CreateThreadResponseInput) => { try { askPrompt.onStopPolling(); - const threadId = thread.id; + const threadId = thread?.id; + if (!threadId) return; + await createThreadResponse({ variables: { threadId, data: payload }, }); diff --git a/wren-ui/src/pages/home/dashboard.tsx b/wren-ui/src/pages/home/dashboard.tsx index eda1cc9ee2..b6cc935c45 100644 --- a/wren-ui/src/pages/home/dashboard.tsx +++ b/wren-ui/src/pages/home/dashboard.tsx @@ -81,7 +81,7 @@ export default function Dashboard() { }, }); - const onRemoveDashboardItemFromQueryCache = (id: number) => { + const onRemoveDashboardItemFromQueryCache = (id: string) => { updateDashboardQuery((prev) => { return { ...prev, @@ -99,7 +99,7 @@ export default function Dashboard() { } }; - const onDelete = async (id: number) => { + const onDelete = async (id: string) => { await deleteDashboardItem({ variables: { where: { id } } }); }; diff --git a/wren-ui/src/pages/home/index.tsx b/wren-ui/src/pages/home/index.tsx index afff9b5abd..14e33ccde9 100644 --- a/wren-ui/src/pages/home/index.tsx +++ b/wren-ui/src/pages/home/index.tsx @@ -94,14 +94,16 @@ export default function Home() { const askPrompt = useAskPrompt(); const { data: suggestedQuestionsData } = useSuggestedQuestionsQuery({ - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const [createThread, { loading: threadCreating }] = useCreateThreadMutation({ onError: (error) => console.error(error), onCompleted: () => homeSidebar.refetch(), }); const [preloadThread] = useThreadLazyQuery({ - fetchPolicy: 'cache-and-network', + fetchPolicy: 'network-only', + nextFetchPolicy: 'network-only', }); const { data: settingsResult } = useGetSettingsQuery(); diff --git a/wren-ui/src/pages/modeling.tsx b/wren-ui/src/pages/modeling.tsx index 46eb8f11ef..a9470a5a76 100644 --- a/wren-ui/src/pages/modeling.tsx +++ b/wren-ui/src/pages/modeling.tsx @@ -1,11 +1,33 @@ import dynamic from 'next/dynamic'; import { useRouter } from 'next/router'; import { useSearchParams } from 'next/navigation'; -import { forwardRef, useEffect, useMemo, useRef } from 'react'; -import { message } from 'antd'; +import React, { forwardRef, useEffect, useMemo, useRef, useState } from 'react'; +import { + Alert, + Button, + Collapse, + Dropdown, + Input, + Menu, + Select, + Space, + Spin, + Table, + Tag, + message, +} from 'antd'; +import { + DeleteOutlined, + EditOutlined, + RobotOutlined, + SaveOutlined, + TableOutlined, +} from '@ant-design/icons'; +import { gql, useApolloClient, useMutation } from '@apollo/client'; import styled from 'styled-components'; import { MORE_ACTION, NODE_TYPE } from '@/utils/enum'; import { editCalculatedField } from '@/utils/modelingHelper'; +import SimpleLayout from '@/components/layouts/SimpleLayout'; import SiderLayout from '@/components/layouts/SiderLayout'; import MetadataDrawer from '@/components/pages/modeling/MetadataDrawer'; import EditMetadataModal from '@/components/pages/modeling/EditMetadataModal'; @@ -46,21 +68,281 @@ import { } from '@/apollo/client/graphql/relationship.generated'; import * as events from '@/utils/events'; +const GENERATE_MODELING_SEMANTICS = gql` + mutation GenerateModelingSemantics($data: GenerateModelingSemanticsInput!) { + generateModelingSemantics(data: $data) + } +`; + +const MODELING_SEMANTICS_RESULT = gql` + query ModelingSemanticsResult($queryId: String!) { + modelingSemanticsResult(queryId: $queryId) + } +`; + +const GENERATE_MODELING_RELATIONSHIPS = gql` + mutation GenerateModelingRelationships { + generateModelingRelationships + } +`; + +const MODELING_RELATIONSHIPS_RESULT = gql` + query ModelingRelationshipsResult($queryId: String!) { + modelingRelationshipsResult(queryId: $queryId) + } +`; + +const SAVE_MODELING_RELATIONSHIPS = gql` + mutation SaveModelingRelationships($data: [ModelingRelationshipInput!]!) { + saveModelingRelationships(data: $data) + } +`; + +const SAVE_MODELING_SEMANTICS = gql` + mutation SaveModelingSemantics($data: [SaveModelingSemanticInput!]!) { + saveModelingSemantics(data: $data) + } +`; + const Diagram = dynamic(() => import('@/components/diagram'), { ssr: false }); // https://github.com/vercel/next.js/issues/4957#issuecomment-413841689 const ForwardDiagram = forwardRef(function ForwardDiagram(props: any, ref) { return ; }); +const semanticText = (value: any): string => { + if (Array.isArray(value)) { + return value.map(semanticText).filter(Boolean).join(', '); + } + return typeof value === 'string' ? value.trim() : ''; +}; + +const firstSemanticText = (...values: any[]): string => + values.map(semanticText).find(Boolean) || ''; + +const addSemanticText = >( + payload: T, + key: string, + value: any, +): T => { + const text = semanticText(value); + return text ? { ...payload, [key]: text } : payload; +}; + const DiagramWrapper = styled.div` position: relative; height: 100%; `; +const AssistantAction = styled.div` + position: absolute; + top: 16px; + right: 16px; + z-index: 10; +`; + +const AssistantPage = styled.div` + min-height: calc(100vh - 48px); + background: #f5f5f5; + padding: 72px 24px; +`; + +const AssistantCard = styled.div` + max-width: 1060px; + margin: 0 auto; + padding: 34px 72px; + background: #fff; + border: 1px solid #e5e7eb; + border-radius: 2px; +`; + +const AssistantBack = styled.button` + display: block; + max-width: 1060px; + margin: 0 auto 8px; + padding: 0; + border: 0; + background: transparent; + color: #5f6368; + cursor: pointer; +`; + +const AssistantTitle = styled.h1` + margin: 0 0 14px; + color: #3c4043; + font-size: 40px; + line-height: 1.2; +`; + +const AssistantDescription = styled.p` + margin: 0 0 26px; + color: #5f6368; + line-height: 1.6; +`; + +const AssistantFooter = styled.div` + display: flex; + justify-content: space-between; + margin-top: 28px; +`; + +const SemanticReviewCard = styled.div` + padding: 18px 16px 28px; + border-bottom: 1px solid #e5e7eb; + + &:last-child { + border-bottom: 0; + } +`; + +const RelationshipGroup = styled.div` + margin-top: 24px; + border: 1px solid #e5e7eb; + border-radius: 4px; + overflow: hidden; +`; + +const RelationshipGroupTitle = styled.div` + display: flex; + align-items: center; + gap: 8px; + padding: 14px 16px; + border-bottom: 1px solid #e5e7eb; +`; + +const AssistantCenter = styled.div` + display: flex; + flex-direction: column; + align-items: center; + justify-content: center; + min-height: 170px; + color: #0f3bff; +`; + +type AssistantRelationship = { + clientId: string; + fromModel: string; + fromColumn: string; + toModel: string; + toColumn: string; + type: string; + reason: string; +}; + +const SEMANTIC_EXAMPLE_PROMPTS = [ + { + label: 'College', + text: 'The purpose of this dataset is to monitor academic performance by tracking student enrollments, grades, and GPA calculations, and to identify areas for student support.', + }, + { + label: 'E-commerce', + text: 'This dataset includes historical pricing information, discount rates, and promotional activities. It supports dynamic pricing strategies, promotion effectiveness analysis, and competitive pricing assessments.', + }, + { + label: 'Human Resources', + text: 'This dataset tracks job postings, applicant details, interview processes, and hiring outcomes. It supports recruitment strategy optimization, time-to-hire analysis, and candidate sourcing effectiveness.', + }, +]; + +const RELATIONSHIP_TYPES = [ + { label: 'Many-to-one', value: 'MANY_TO_ONE' }, + { label: 'One-to-many', value: 'ONE_TO_MANY' }, + { label: 'One-to-one', value: 'ONE_TO_ONE' }, +]; + +const relationshipTypeLabel = (type: string) => + RELATIONSHIP_TYPES.find((item) => item.value === type)?.label || type; + +const normalizeRelationshipType = (type: string) => { + const normalized = String(type || '') + .trim() + .replace(/[\s-]+/g, '_') + .toUpperCase(); + if (normalized === 'MANY_TO_ONE') return 'MANY_TO_ONE'; + if (normalized === 'ONE_TO_MANY') return 'ONE_TO_MANY'; + if (normalized === 'ONE_TO_ONE') return 'ONE_TO_ONE'; + return type; +}; + +const getRelationshipFieldValue = (model = '', column = '') => + `${model}.${column}`; + +const resolveRelationshipField = ( + value = '', + models: Array<{ referenceName?: string; fields?: any[] }> = [], +) => { + const fieldValue = String(value || ''); + for (const model of models) { + for (const field of model.fields || []) { + if ( + fieldValue === + getRelationshipFieldValue(model.referenceName, field.referenceName) + ) { + return { + model: model.referenceName || '', + column: field.referenceName || '', + }; + } + } + } + + return { model: '', column: '' }; +}; + +const resolveRelationshipFieldParts = ( + model = '', + column = '', + fallbackValue = '', + models: Array<{ referenceName?: string; fields?: any[] }> = [], +) => { + const resolvedField = + model && column + ? resolveRelationshipField( + getRelationshipFieldValue(model, column), + models, + ) + : { model: '', column: '' }; + + return resolvedField.model && resolvedField.column + ? resolvedField + : resolveRelationshipField(fallbackValue, models); +}; + +const renderIcon = (IconComponent) => React.createElement(IconComponent as any); +const ASSISTANT_CANCELLED = 'ASSISTANT_CANCELLED'; +const ASSISTANT_SAVE_MESSAGE_KEY = 'modeling-ai-assistant-save'; +const ASSISTANT_INITIAL_POLL_INTERVAL_MS = 1000; +const ASSISTANT_MAX_POLL_INTERVAL_MS = 5000; +const ASSISTANT_MAX_POLL_ATTEMPTS = 1800; + export default function Modeling() { const router = useRouter(); const searchParams = useSearchParams(); + const apolloClient = useApolloClient(); const diagramRef = useRef(null); + const assistantRunIdRef = useRef(0); + const assistantSavingRef = useRef(false); + const [assistantMode, setAssistantMode] = useState< + 'semantics' | 'relationships' | null + >(null); + const [assistantLoading, setAssistantLoading] = useState(false); + const [assistantSaving, setAssistantSaving] = useState(false); + const [assistantError, setAssistantError] = useState(null); + const [selectedModels, setSelectedModels] = useState([]); + const [semanticPrompt, setSemanticPrompt] = useState(''); + const [semanticStep, setSemanticStep] = useState<'pick' | 'generate'>('pick'); + const [semanticSearch, setSemanticSearch] = useState(''); + const [semanticResult, setSemanticResult] = useState([]); + const [relationshipResult, setRelationshipResult] = useState< + AssistantRelationship[] + >([]); + const [originalRelationshipResult, setOriginalRelationshipResult] = useState< + AssistantRelationship[] + >([]); + const [relationshipAutoStarted, setRelationshipAutoStarted] = useState(false); + const [editingRelationshipKey, setEditingRelationshipKey] = useState< + string | null + >(null); const { data } = useDiagramQuery({ fetchPolicy: 'cache-and-network', @@ -83,6 +365,7 @@ export default function Modeling() { awaitRefetchQueries: true, ...options, onCompleted: () => { + window.dispatchEvent(new Event('wren:modeling-changed')); // refetch to get latest deploy status deployStatusQueryResult.refetch(); @@ -173,7 +456,6 @@ export default function Modeling() { }, }), ); - const [deleteRelationshipMutation] = useDeleteRelationshipMutation( getBaseOptions({ onCompleted: () => { @@ -199,6 +481,12 @@ export default function Modeling() { }, }), ); + const [generateModelingSemantics] = useMutation(GENERATE_MODELING_SEMANTICS); + const [generateModelingRelationships] = useMutation( + GENERATE_MODELING_RELATIONSHIPS, + ); + const [saveModelingSemantics] = useMutation(SAVE_MODELING_SEMANTICS); + const [saveModelingRelationships] = useMutation(SAVE_MODELING_RELATIONSHIPS); const diagramData = useMemo(() => { if (!data) return null; @@ -276,11 +564,16 @@ export default function Modeling() { const onSelect = (selectKeys) => { if (diagramRef.current) { const { getNodes, fitBounds } = diagramRef.current; - const node = getNodes().find((node) => node.id === selectKeys[0]); + const selectedKey = selectKeys?.[0]; + if (!selectedKey) return; + + const node = getNodes().find((node) => node.id === selectedKey); + if (!node?.position) return; + const position = { ...node.position, - width: node.width, - height: node.height, + width: node.width ?? 1, + height: node.height ?? 1, }; fitBounds(position); } @@ -377,6 +670,994 @@ export default function Modeling() { const modelLoading = modelCreating || modelUpdating; const relationshipLoading = relationshipUpdating || relationshipCreating; + const waitForAssistantResult = async ( + queryId: string, + query: any, + fieldName: string, + runId: number, + ) => { + if (!queryId) { + throw new Error('AI assistant did not return a task id.'); + } + + for (let attempt = 0; attempt < ASSISTANT_MAX_POLL_ATTEMPTS; attempt += 1) { + if (assistantRunIdRef.current !== runId) { + throw new Error(ASSISTANT_CANCELLED); + } + const res = await apolloClient.query({ + query, + variables: { queryId }, + fetchPolicy: 'network-only', + }); + if (assistantRunIdRef.current !== runId) { + throw new Error(ASSISTANT_CANCELLED); + } + const payload = res.data?.[fieldName]; + const status = String(payload?.status || '').toLowerCase(); + if (status === 'finished') return payload.response || []; + if (status === 'failed') { + throw new Error(payload.error?.message || 'AI assistant failed.'); + } + const pollInterval = Math.min( + ASSISTANT_INITIAL_POLL_INTERVAL_MS * Math.max(1, attempt + 1), + ASSISTANT_MAX_POLL_INTERVAL_MS, + ); + await new Promise((resolve) => setTimeout(resolve, pollInterval)); + } + throw new Error('AI assistant timed out.'); + }; + + const normalizeSemanticModel = (name: string, value: any): any => { + const modelName = semanticText(value?.name) || semanticText(name); + return { + name: modelName, + displayName: firstSemanticText( + value?.displayName, + value?.alias, + value?.properties?.displayName, + value?.properties?.alias, + ), + description: firstSemanticText( + value?.description, + value?.properties?.description, + ), + columns: (value?.columns || []).map((column) => { + const columnName = semanticText(column?.name); + return { + name: columnName, + type: column?.type, + displayName: firstSemanticText( + column?.displayName, + column?.alias, + column?.properties?.displayName, + column?.properties?.alias, + ), + description: firstSemanticText( + column?.description, + column?.properties?.description, + ), + }; + }), + }; + }; + + const normalizeSemanticResult = (result: any): any[] => { + if (Array.isArray(result)) { + return result.map((model) => normalizeSemanticModel(model?.name, model)); + } + if (Array.isArray(result?.models)) { + return result.models.map((model) => + normalizeSemanticModel(model?.name, model), + ); + } + if (Array.isArray(result?.semantics)) { + return result.semantics.map((model) => + normalizeSemanticModel(model?.name, model), + ); + } + if (Array.isArray(result?.descriptions)) { + return result.descriptions.map((model) => + normalizeSemanticModel(model?.name, model), + ); + } + if (result && typeof result === 'object') { + return Object.entries(result).map(([name, value]) => + normalizeSemanticModel(name, value), + ); + } + return []; + }; + + const normalizeRelationshipResult = ( + result: any, + ): AssistantRelationship[] => { + const relationships = Array.isArray(result) + ? result + : Array.isArray(result?.relationships) + ? result.relationships + : Array.isArray(result?.response?.relationships) + ? result.response.relationships + : []; + + return relationships.map((relationship, index) => { + const availableModels = diagramData?.models || []; + const from = resolveRelationshipFieldParts( + relationship.fromModel || '', + relationship.fromColumn || '', + relationship.from || relationship.fromField || '', + availableModels, + ); + const to = resolveRelationshipFieldParts( + relationship.toModel || '', + relationship.toColumn || '', + relationship.to || relationship.toField || '', + availableModels, + ); + const fromModel = from.model; + const fromColumn = from.column; + const toModel = to.model; + const toColumn = to.column; + + return { + clientId: + relationship.clientId || + [ + fromModel, + fromColumn, + toModel, + toColumn, + relationship.type, + index, + ].join(':'), + fromModel, + fromColumn, + toModel, + toColumn, + type: normalizeRelationshipType(relationship.type), + reason: relationship.reason || relationship.description || '', + }; + }); + }; + + const openAssistant = (mode: 'semantics' | 'relationships') => { + assistantRunIdRef.current += 1; + setAssistantMode(mode); + setAssistantError(null); + setAssistantSaving(false); + setSemanticStep('pick'); + setSemanticSearch(''); + setSelectedModels( + diagramData?.models?.map((model) => model.referenceName) || [], + ); + setSemanticResult([]); + setRelationshipResult([]); + setOriginalRelationshipResult([]); + setRelationshipAutoStarted(false); + setEditingRelationshipKey(null); + }; + + const runAssistant = async () => { + const runId = assistantRunIdRef.current + 1; + assistantRunIdRef.current = runId; + try { + setAssistantLoading(true); + setAssistantError(null); + if (assistantMode === 'semantics') { + if (!selectedModels.length) { + throw new Error('Select at least one model.'); + } + const res = await generateModelingSemantics({ + variables: { + data: { + selectedModels, + userPrompt: + semanticPrompt || 'Describe this dataset for analytics.', + }, + }, + }); + const queryId = res.data?.generateModelingSemantics?.queryId; + const result = await waitForAssistantResult( + queryId, + MODELING_SEMANTICS_RESULT, + 'modelingSemanticsResult', + runId, + ); + if (assistantRunIdRef.current !== runId) return; + const normalizedResult = normalizeSemanticResult(result); + if (!normalizedResult.length) { + throw new Error('AI assistant returned no semantic descriptions.'); + } + setSemanticResult(normalizedResult); + setSemanticStep('generate'); + } + if (assistantMode === 'relationships') { + if (!diagramData?.models || diagramData.models.length < 2) { + throw new Error('At least two models are required.'); + } + const res = await generateModelingRelationships(); + const queryId = res.data?.generateModelingRelationships?.queryId; + const result = await waitForAssistantResult( + queryId, + MODELING_RELATIONSHIPS_RESULT, + 'modelingRelationshipsResult', + runId, + ); + if (assistantRunIdRef.current !== runId) return; + const normalizedResult = normalizeRelationshipResult(result); + setRelationshipResult(normalizedResult); + setOriginalRelationshipResult(normalizedResult); + } + } catch (error: any) { + if (error.message !== ASSISTANT_CANCELLED) { + const errorMessage = + error.message || 'Failed to run Modeling AI Assistant.'; + setAssistantError(errorMessage); + message.error(errorMessage); + } + } finally { + if (assistantRunIdRef.current === runId) { + setAssistantLoading(false); + } + } + }; + + const updateSemanticModelDescription = ( + modelName: string, + description: string, + ) => { + setSemanticResult((models) => + models.map((model) => + model.name === modelName ? { ...model, description } : model, + ), + ); + }; + + const updateSemanticModelDisplayName = ( + modelName: string, + displayName: string, + ) => { + setSemanticResult((models) => + models.map((model) => + model.name === modelName ? { ...model, displayName } : model, + ), + ); + }; + + const updateSemanticColumnDisplayName = ( + modelName: string, + columnName: string, + displayName: string, + ) => { + setSemanticResult((models) => + models.map((model) => + model.name === modelName + ? { + ...model, + columns: (model.columns || []).map((column) => + column.name === columnName + ? { ...column, displayName } + : column, + ), + } + : model, + ), + ); + }; + + const updateSemanticColumnDescription = ( + modelName: string, + columnName: string, + description: string, + ) => { + setSemanticResult((models) => + models.map((model) => + model.name === modelName + ? { + ...model, + columns: (model.columns || []).map((column) => + column.name === columnName + ? { ...column, description } + : column, + ), + } + : model, + ), + ); + }; + + const updateRelationship = ( + clientId: string, + changes: Partial, + ) => { + setRelationshipResult((relationships) => + relationships.map((relationship) => + relationship.clientId === clientId + ? { ...relationship, ...changes } + : relationship, + ), + ); + }; + + const updateRelationshipField = ( + clientId: string, + side: 'from' | 'to', + value: string, + ) => { + const field = resolveRelationshipField(value, diagramData?.models || []); + updateRelationship( + clientId, + side === 'from' + ? { fromModel: field.model, fromColumn: field.column } + : { toModel: field.model, toColumn: field.column }, + ); + }; + + const deleteSuggestedRelationship = (clientId: string) => { + setRelationshipResult((relationships) => + relationships.filter( + (relationship) => relationship.clientId !== clientId, + ), + ); + if (editingRelationshipKey === clientId) { + setEditingRelationshipKey(null); + } + }; + + const closeAssistant = () => { + assistantRunIdRef.current += 1; + setAssistantMode(null); + setAssistantError(null); + setAssistantLoading(false); + setAssistantSaving(false); + setSemanticStep('pick'); + setSemanticSearch(''); + setRelationshipAutoStarted(false); + setEditingRelationshipKey(null); + }; + + const saveAssistantResult = async () => { + if (assistantSavingRef.current) return; + assistantSavingRef.current = true; + try { + if (!diagramData) return; + setAssistantSaving(true); + if (assistantMode === 'semantics') { + const latestDiagram = await apolloClient.query({ + query: DIAGRAM, + fetchPolicy: 'network-only', + }); + const currentModels = latestDiagram.data?.diagram?.models || []; + const data = semanticResult.flatMap((model) => { + const diagramModel = currentModels.find( + (item) => item.referenceName === model.name, + ); + if (!diagramModel) return []; + const modelPayload = { + modelId: diagramModel.modelId, + referenceName: diagramModel.referenceName, + columns: (model.columns || []) + .map((column) => { + const field = diagramModel.fields.find( + (item) => item.referenceName === column.name, + ); + return field + ? addSemanticText( + addSemanticText( + { + id: field.columnId, + referenceName: field.referenceName, + }, + 'displayName', + column.displayName, + ), + 'description', + column.description, + ) + : null; + }) + .filter(Boolean), + }; + + return addSemanticText( + addSemanticText(modelPayload, 'displayName', model.displayName), + 'description', + model.description, + ); + }); + + if (!data.length) { + throw new Error('No semantic descriptions to save.'); + } + + await saveModelingSemantics({ + variables: { data }, + refetchQueries, + awaitRefetchQueries: true, + }); + } + if (assistantMode === 'relationships') { + const res = await saveModelingRelationships({ + variables: { + data: relationshipResult.map((relationship) => ({ + fromModel: relationship.fromModel, + fromColumn: relationship.fromColumn, + toModel: relationship.toModel, + toColumn: relationship.toColumn, + type: relationship.type, + description: relationship.reason, + })), + }, + refetchQueries, + awaitRefetchQueries: true, + }); + const createdCount = + res.data?.saveModelingRelationships?.createdCount || 0; + const skippedCount = + res.data?.saveModelingRelationships?.skippedCount || 0; + + if (!createdCount) { + throw new Error( + skippedCount + ? 'No valid new relationships could be saved for the current models.' + : 'No relationship suggestions to save.', + ); + } + + if (skippedCount) { + message.warning( + `${createdCount} relationship(s) saved. ${skippedCount} invalid or duplicate suggestion(s) skipped.`, + ); + } else { + message.success({ + key: ASSISTANT_SAVE_MESSAGE_KEY, + content: 'Saved Modeling AI Assistant suggestions.', + }); + } + } + closeAssistant(); + if (assistantMode !== 'relationships') { + message.success({ + key: ASSISTANT_SAVE_MESSAGE_KEY, + content: 'Saved Modeling AI Assistant suggestions.', + }); + } + } catch (error: any) { + message.error(error.message || 'Failed to save assistant suggestions.'); + } finally { + assistantSavingRef.current = false; + setAssistantSaving(false); + } + }; + + const semanticModelOptions = (diagramData?.models || []).filter((model) => { + const keyword = semanticSearch.trim().toLowerCase(); + if (!keyword) return true; + return [model.displayName, model.referenceName] + .filter(Boolean) + .some((value) => String(value).toLowerCase().includes(keyword)); + }); + + const relationshipFieldOptions = (diagramData?.models || []).flatMap( + (model) => + (model.fields || []).map((field) => ({ + label: `${model.referenceName}.${field.referenceName}`, + value: getRelationshipFieldValue( + model.referenceName, + field.referenceName, + ), + })), + ); + + const relationshipGroups = relationshipResult.reduce< + Record + >((groups, relationship) => { + const key = relationship.fromModel || 'Unknown model'; + groups[key] = [...(groups[key] || []), relationship]; + return groups; + }, {}); + + const isRelationshipGenerating = + assistantMode === 'relationships' && + (!relationshipAutoStarted || assistantLoading) && + !relationshipResult.length; + + useEffect(() => { + if (assistantMode !== 'relationships') return; + if (relationshipAutoStarted) return; + setRelationshipAutoStarted(true); + runAssistant(); + }, [assistantMode, relationshipAutoStarted]); + + useEffect( + () => () => { + assistantRunIdRef.current += 1; + }, + [], + ); + + if (assistantMode === 'semantics') { + return ( + + + + + ← Back to modeling + + + {semanticStep === 'pick' && ( + <> + Pick models + + + Good semantics improve how AI understands and queries your + data. + {' '} + Select models to generate semantics with AI. Modeling AI + Assistant will help you create semantics that improve how AI + understands and queries your data. + +
+ {selectedModels.length}/{diagramData?.models?.length || 0}{' '} + model(s) +
+ setSemanticSearch(event.target.value)} + /> +
setSelectedModels(keys as string[]), + }} + columns={[ + { + title: 'Model name', + render: (_value, model) => + model.displayName || model.referenceName, + }, + ]} + /> + + + + + + )} + + {semanticStep === 'generate' && ( + <> + Generate semantics +

User Prompt

+ + Help AI better understand your data by providing a brief + description of your dataset's purpose. Modeling AI + Assistant will use this context to generate more relevant + semantics. + +
+ + setSemanticPrompt(event.target.value) + } + placeholder="Describe what this dataset represents and how it is used." + /> + +
+ {assistantLoading && ( + + +
+ Generating semantic descriptions... +
+ + )} + {assistantError && ( + + )} + + +
+ Following, we provide some example prompts based on some + real world datasets. +
+ + {SEMANTIC_EXAMPLE_PROMPTS.map((example) => ( +
+ {example.label} +
+ {example.text} +
+
+ ))} + + + + {semanticResult.length ? ( +
+
+ Generated semantics +
+ Review the semantics generated by AI. +
+
+ {semanticResult.map((model) => ( + +
+ {model.name} + + {(model.columns || []).length} column(s) + +
+
Description
+ + updateSemanticModelDescription( + model.name, + event.target.value, + ) + } + /> +
Alias / synonyms
+ + updateSemanticModelDisplayName( + model.name, + event.target.value, + ) + } + /> +
( + + updateSemanticColumnDisplayName( + model.name, + column.name, + event.target.value, + ) + } + /> + ), + }, + { title: 'Type', dataIndex: 'type', width: 140 }, + { + title: 'Description', + render: (_value, column) => ( + + updateSemanticColumnDescription( + model.name, + column.name, + event.target.value, + ) + } + /> + ), + }, + ]} + /> + + ))} + + ) : null} + + + + + + )} + + + + + ); + } + + if (assistantMode === 'relationships') { + return ( + + + + + ← Back to modeling + + + Generate relationships + + Modeling AI Assistant will use AI to discover potential + connections between your models. +
+ Review the suggested relationships and adjust them before saving + to your data models. +
+ Learn more:{' '} + + Modeling AI Assistant / Generate relationships + + + + {isRelationshipGenerating ? ( + + +
Generating...
+ + ) : ( + <> + {assistantError ? ( + + + + ) : !relationshipResult.length ? ( + +
+ No relationship suggestions were generated. +
+ + ) : null} + + {Object.entries(relationshipGroups).map( + ([modelName, relationships]) => ( + + + {renderIcon(TableOutlined)} + {modelName} + +
+ editingRelationshipKey === record.clientId ? ( + + updateRelationshipField( + record.clientId, + 'to', + value, + ) + } + /> + ) : ( + getRelationshipFieldValue( + record.toModel, + record.toColumn, + ) + ), + }, + { + title: 'Type', + width: 170, + render: (_value, record) => + editingRelationshipKey === record.clientId ? ( + + + + + + + + + + + + + + ); +} diff --git a/wren-ui/src/pages/organization/index.tsx b/wren-ui/src/pages/organization/index.tsx new file mode 100644 index 0000000000..da6625737f --- /dev/null +++ b/wren-ui/src/pages/organization/index.tsx @@ -0,0 +1,14 @@ +import { useEffect } from 'react'; +import { useRouter } from 'next/router'; +import PageLoading from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; + +export default function OrganizationIndex() { + const router = useRouter(); + + useEffect(() => { + router.replace(Path.OrganizationGeneral); + }, [router]); + + return ; +} diff --git a/wren-ui/src/pages/organization/invitations/[token].tsx b/wren-ui/src/pages/organization/invitations/[token].tsx new file mode 100644 index 0000000000..03a03b7aba --- /dev/null +++ b/wren-ui/src/pages/organization/invitations/[token].tsx @@ -0,0 +1,64 @@ +import { useEffect, useState } from 'react'; +import { Button, Result, Spin } from 'antd'; +import { useRouter } from 'next/router'; +import SimpleLayout from '@/components/layouts/SimpleLayout'; + +export default function AcceptOrganizationInvitationPage() { + const router = useRouter(); + const token = Array.isArray(router.query.token) + ? router.query.token[0] + : router.query.token; + const [state, setState] = useState< + 'loading' | 'success' | 'error' + >('loading'); + const [message, setMessage] = useState('Accepting invitation...'); + + useEffect(() => { + if (!router.isReady || !token) return; + + const acceptInvitation = async () => { + try { + const response = await fetch( + `/api/v1/organizations/invitations/accept/${token}`, + { + method: 'POST', + }, + ); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to accept invitation'); + } + setState('success'); + setMessage('Invitation accepted successfully.'); + } catch (error: any) { + setState('error'); + setMessage(error.message || 'Failed to accept invitation'); + } + }; + + void acceptInvitation(); + }, [router.isReady, token]); + + return ( + +
+ {state === 'loading' ? ( +
+ +
{message}
+
+ ) : ( + void router.push('/organization/members')}> + Go to members + + } + /> + )} +
+ + ); +} diff --git a/wren-ui/src/pages/organization/members.tsx b/wren-ui/src/pages/organization/members.tsx new file mode 100644 index 0000000000..e563c20691 --- /dev/null +++ b/wren-ui/src/pages/organization/members.tsx @@ -0,0 +1,762 @@ +import { useEffect, useMemo, useState } from 'react'; +import { + Button, + Form, + Input, + Modal, + message, +} from 'antd'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; +import { getRelativeTime } from '@/utils/time'; + +type OrganizationRole = 'Admin' | 'Member'; +type ProjectPermissionRole = 'Owner' | 'Editor' | 'Viewer'; +type InvitationStatus = 'Pending' | 'Accepted' | 'Expired'; + +interface ProjectOption { + id: number; + displayName: string; +} + +interface MemberProject { + projectId: number; + displayName: string; + permission: ProjectPermissionRole; +} + +interface MemberRecord { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationRole; + projects: MemberProject[]; +} + +interface PendingInviteRecord { + id: number; + email: string; + organizationRole: OrganizationRole; + status: InvitationStatus; + token: string; + inviteLink: string; + createdAt: string; + expiresAt: string; + acceptedAt?: string | null; + projects: MemberProject[]; +} + +interface MembersResponse { + members: MemberRecord[]; + invitations: PendingInviteRecord[]; + projects: ProjectOption[]; + currentUserId: number | null; + error?: string; +} + +const ROLE_OPTIONS = ['Admin', 'Member']; +const PROJECT_PERMISSION_OPTIONS = ['Owner', 'Editor', 'Viewer']; + +const MembersCard = styled.div` + margin-top: 16px; + border: 1px solid var(--gray-4); + border-radius: 4px; + background: white; + overflow: hidden; +`; + +const MembersHeaderRow = styled.div` + display: grid; + grid-template-columns: minmax(0, 1.2fr) minmax(0, 1fr) 180px 120px; + gap: 16px; + padding: 14px 20px; + border-bottom: 1px solid var(--gray-4); + font-weight: 600; + color: var(--gray-8); +`; + +const MembersRow = styled.div` + display: grid; + grid-template-columns: minmax(0, 1.2fr) minmax(0, 1fr) 180px 120px; + gap: 16px; + padding: 16px 20px; + align-items: center; + + &:not(:last-child) { + border-bottom: 1px solid var(--gray-4); + } +`; + +const ProjectsBox = styled.div` + border: 1px solid var(--gray-4); + border-radius: 4px; + background: white; + overflow: hidden; +`; + +const ProjectsToolbar = styled.div` + padding: 12px 12px 0; +`; + +const ProjectRow = styled.div` + display: grid; + grid-template-columns: 28px 1fr 140px; + gap: 12px; + align-items: center; + padding: 10px 12px; + border-top: 1px solid var(--gray-4); +`; + +const ProjectHeader = styled(ProjectRow)` + background: var(--gray-3); + font-weight: 600; +`; + +const EmptyProjectsRow = styled.div` + padding: 16px 12px; + border-top: 1px solid var(--gray-4); + color: var(--gray-7); +`; + +const ProjectCheckbox = styled.input` + width: 16px; + height: 16px; + cursor: pointer; +`; + +const NativeSelect = styled.select` + width: 100%; + min-height: 40px; + border: 1px solid var(--gray-4); + border-radius: 6px; + padding: 8px 12px; + background: white; + + &:disabled { + background: var(--gray-3); + color: var(--gray-6); + cursor: not-allowed; + } +`; + +const ModalLabel = styled.div` + margin-bottom: 8px; + font-weight: 500; + color: var(--gray-8); +`; + +const SectionTitle = styled.h4` + margin: 16px 0 0; + color: var(--gray-8); +`; + +const PendingHeaderRow = styled.div` + display: grid; + grid-template-columns: minmax(0, 1.2fr) 120px 140px minmax(0, 1.4fr) 80px; + gap: 16px; + padding: 14px 20px; + border-bottom: 1px solid var(--gray-4); + font-weight: 600; + color: var(--gray-8); +`; + +const PendingRow = styled.div` + display: grid; + grid-template-columns: minmax(0, 1.2fr) 120px 140px minmax(0, 1.4fr) 80px; + gap: 16px; + padding: 16px 20px; + align-items: center; + + &:not(:last-child) { + border-bottom: 1px solid var(--gray-4); + } +`; + +const NameCell = styled.div` + display: flex; + align-items: center; + gap: 12px; +`; + +const InitialsBadge = styled.div` + width: 40px; + height: 40px; + border-radius: 999px; + display: inline-flex; + align-items: center; + justify-content: center; + background: #c27ba0; + color: white; + font-weight: 600; +`; + +const MemberName = styled.div` + display: flex; + flex-direction: column; +`; + +const MemberEmail = styled.span` + color: var(--gray-7); +`; + +const InviteLinkRow = styled.div` + display: flex; + align-items: center; + gap: 8px; +`; + +const InviteLinkInput = styled.input` + width: 100%; + min-height: 32px; + border: 1px solid var(--gray-4); + border-radius: 6px; + padding: 6px 10px; + background: white; +`; + +const StatusBadge = styled.span<{ $status: InvitationStatus }>` + display: inline-flex; + align-items: center; + justify-content: center; + min-width: 72px; + padding: 4px 10px; + border-radius: 999px; + font-size: 12px; + font-weight: 600; + color: ${(props) => + props.$status === 'Pending' + ? '#4e9f2d' + : props.$status === 'Accepted' + ? '#1d70b8' + : '#c2410c'}; + background: ${(props) => + props.$status === 'Pending' + ? '#f0f9e8' + : props.$status === 'Accepted' + ? '#eaf4ff' + : '#fff3e8'}; + border: 1px solid + ${(props) => + props.$status === 'Pending' + ? '#9ed270' + : props.$status === 'Accepted' + ? '#9cc9f5' + : '#fdba74'}; +`; + +const ActionArea = styled.div` + display: flex; + justify-content: flex-end; +`; + +const InlineConfirm = styled.div` + display: flex; + align-items: center; + gap: 8px; + justify-content: flex-end; +`; + +const getInitials = (name: string) => + (name || 'U') + .trim() + .split(/\s+/) + .slice(0, 2) + .map((part) => part[0]) + .join('') + .toUpperCase(); + +const buildProjectSelectionMap = ( + projects: ProjectOption[], + organizationRole: OrganizationRole, +) => { + if (organizationRole === 'Admin') { + return projects.reduce>((acc, project) => { + acc[project.id] = 'Owner'; + return acc; + }, {}); + } + return {}; +}; + +export default function OrganizationMembersPage() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [updatingMemberId, setUpdatingMemberId] = useState(null); + const [visible, setVisible] = useState(false); + const [members, setMembers] = useState([]); + const [invitations, setInvitations] = useState([]); + const [projects, setProjects] = useState([]); + const [currentUserId, setCurrentUserId] = useState(null); + const [confirmingMemberId, setConfirmingMemberId] = useState(null); + const [confirmingInviteId, setConfirmingInviteId] = useState(null); + const [projectSelections, setProjectSelections] = useState< + Record + >({}); + const organizationRole = Form.useWatch( + 'organizationRole', + form, + ) as OrganizationRole | undefined; + const projectSearch = Form.useWatch('projectSearch', form) as string | undefined; + + const loadMembers = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/organizations/members'); + const payload = (await response.json()) as MembersResponse; + if (!response.ok) { + throw new Error(payload.error || 'Failed to load members'); + } + setMembers(payload.members || []); + setInvitations(payload.invitations || []); + setProjects(payload.projects || []); + setCurrentUserId(payload.currentUserId ?? null); + } catch (error: any) { + message.error(error.message || 'Failed to load members'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadMembers(); + }, []); + + useEffect(() => { + if (!visible) return; + if (organizationRole === 'Admin') { + setProjectSelections(buildProjectSelectionMap(projects, organizationRole)); + return; + } + setProjectSelections({}); + }, [organizationRole, visible, projects]); + + const openInviteModal = () => { + form.resetFields(); + form.setFieldsValue({ + organizationRole: 'Admin', + projectSearch: '', + }); + setProjectSelections(buildProjectSelectionMap(projects, 'Admin')); + setVisible(true); + }; + + const filteredProjects = useMemo(() => { + const search = `${projectSearch || ''}`.trim().toLowerCase(); + return projects.filter((project) => + !search + ? true + : project.displayName.toLowerCase().includes(search), + ); + }, [projects, projectSearch]); + + const toggleProject = (projectId: number, checked: boolean) => { + setProjectSelections((prev) => { + const next = { ...prev }; + if (checked) { + next[projectId] = next[projectId] || 'Owner'; + } else { + delete next[projectId]; + } + return next; + }); + }; + + const updateProjectPermission = ( + projectId: number, + permission: ProjectPermissionRole, + ) => { + setProjectSelections((prev) => ({ + ...prev, + [projectId]: permission, + })); + }; + + const inviteMember = async () => { + try { + const values = await form.validateFields(); + const response = await fetch('/api/v1/organizations/members', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + email: values.email, + organizationRole: values.organizationRole, + projects: Object.entries(projectSelections).map( + ([projectId, permission]) => ({ + projectId: Number(projectId), + permission, + }), + ), + }), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to invite member'); + } + message.success('Member invited successfully.'); + setVisible(false); + form.resetFields(); + setProjectSelections({}); + await loadMembers(); + } catch (error: any) { + if (error?.errorFields) return; + message.error(error.message || 'Failed to invite member'); + } finally { + setSaving(false); + } + }; + + const updateMemberRole = async ( + memberId: number, + organizationRole: OrganizationRole, + ) => { + try { + setUpdatingMemberId(memberId); + const response = await fetch(`/api/v1/organizations/members/${memberId}`, { + method: 'PATCH', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ organizationRole }), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to update member'); + } + setMembers((prev) => + prev.map((member) => + member.id === memberId + ? { ...member, organizationRole: payload.organizationRole } + : member, + ), + ); + message.success('Member role updated successfully.'); + } catch (error: any) { + message.error(error.message || 'Failed to update member'); + } finally { + setUpdatingMemberId(null); + } + }; + + const removeMember = async (memberId: number) => { + try { + setUpdatingMemberId(memberId); + const response = await fetch(`/api/v1/organizations/members/${memberId}`, { + method: 'DELETE', + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to remove member'); + } + setMembers((prev) => prev.filter((member) => member.id !== memberId)); + message.success('Member removed successfully.'); + } catch (error: any) { + message.error(error.message || 'Failed to remove member'); + } finally { + setUpdatingMemberId(null); + setConfirmingMemberId(null); + } + }; + + const removeInvitation = async (invitationId: number) => { + try { + const response = await fetch( + `/api/v1/organizations/invitations/${invitationId}`, + { + method: 'DELETE', + }, + ); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to remove invitation'); + } + setInvitations((prev) => + prev.filter((invitation) => invitation.id !== invitationId), + ); + message.success('Invitation removed successfully.'); + } catch (error: any) { + message.error(error.message || 'Failed to remove invitation'); + } finally { + setConfirmingInviteId(null); + } + }; + + const copyInviteLink = async (inviteLink: string) => { + try { + await navigator.clipboard.writeText(inviteLink); + message.success('Invitation link copied.'); + } catch { + message.error('Failed to copy invitation link'); + } + }; + + return ( + + Invite people + + } + > + +
+ Pending invites + + +
Invitee
+
Status
+
Created
+
Invite link
+
+ + {invitations.length ? ( + invitations.map((invitation) => ( + + + + {getInitials(invitation.email.split('@')[0])} + +
{invitation.email}
+ + + {invitation.status} + +
{getRelativeTime(invitation.createdAt)}
+ + + + + + {confirmingInviteId === invitation.id ? ( + + + + + ) : ( + + )} + + + )) + ) : ( + +
No pending invitations
+
+
+
+
+ + )} + + + Organization members + + +
Name
+
Email
+
Role
+
Actions
+ + {members.map((member) => ( + + + {getInitials(member.name)} + + + {member.name} + {member.userId === currentUserId ? ' (me)' : ''} + + + + {member.email} + + void updateMemberRole( + member.id, + event.target.value as OrganizationRole, + ) + } + > + {ROLE_OPTIONS.map((role) => ( + + ))} + + + {member.userId === currentUserId ? null : confirmingMemberId === member.id ? ( + + + + + ) : ( + + )} + + + ))} + +
+ + + { + setVisible(false); + setProjectSelections({}); + form.resetFields(); + }} + onOk={() => { + setSaving(true); + void inviteMember(); + }} + confirmLoading={saving} + okText="Invite" + destroyOnClose + width={720} + > +
+ + + + + + + {ROLE_OPTIONS.map((role) => ( + + ))} + + + + + + + + {Object.keys(projectSelections).length}/{projects.length} project(s) + + + + + + + +
+
Project name
+
Permission
+ + + {filteredProjects.map((project) => { + const checked = Boolean(projectSelections[project.id]); + const permission = projectSelections[project.id] || 'Owner'; + return ( + + + toggleProject(project.id, event.target.checked) + } + /> +
{project.displayName}
+ + updateProjectPermission( + project.id, + event.target.value as ProjectPermissionRole, + ) + } + > + {PROJECT_PERMISSION_OPTIONS.map((role) => ( + + ))} + + + ); + })} + {!filteredProjects.length && ( + + {projects.length + ? 'No projects match your search.' + : 'No projects are available for this organization yet.'} + + )} + + + + + + ); +} diff --git a/wren-ui/src/pages/project/access-control.tsx b/wren-ui/src/pages/project/access-control.tsx new file mode 100644 index 0000000000..79a80d9f55 --- /dev/null +++ b/wren-ui/src/pages/project/access-control.tsx @@ -0,0 +1,506 @@ +import { useEffect, useMemo, useState } from 'react'; +import { + Button, + Form, + Input, + Modal, + Typography, + message, +} from 'antd'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; + +type OrganizationRole = 'Admin' | 'Member'; +type ProjectPermission = 'Owner' | 'Contributor' | 'Viewer'; + +interface ProjectAccessMember { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationRole; + permission: ProjectPermission; + isCurrentUser: boolean; + canEditPermission: boolean; + canRemove: boolean; +} + +interface AvailableProjectMember { + id: number; + userId: number; + name: string; + email: string; + organizationRole: OrganizationRole; +} + +interface ProjectAccessResponse { + members: ProjectAccessMember[]; + availableMembers: AvailableProjectMember[]; + currentUserId: number | null; + canManageAccess: boolean; + error?: string; +} + +const PERMISSION_OPTIONS: ProjectPermission[] = [ + 'Owner', + 'Contributor', + 'Viewer', +]; + +const MembersCard = styled.div` + margin-top: 16px; + border: 1px solid var(--gray-4); + border-radius: 4px; + background: white; + overflow: hidden; +`; + +const MembersHeaderRow = styled.div` + display: grid; + grid-template-columns: minmax(0, 1.2fr) minmax(0, 1fr) 120px 160px 120px; + gap: 16px; + padding: 14px 20px; + border-bottom: 1px solid var(--gray-4); + font-weight: 600; + color: var(--gray-8); +`; + +const MembersRow = styled.div` + display: grid; + grid-template-columns: minmax(0, 1.2fr) minmax(0, 1fr) 120px 160px 120px; + gap: 16px; + padding: 16px 20px; + align-items: center; + + &:not(:last-child) { + border-bottom: 1px solid var(--gray-4); + } +`; + +const SectionTitle = styled.h4` + margin: 16px 0 0; + color: var(--gray-8); +`; + +const NameCell = styled.div` + display: flex; + align-items: center; + gap: 12px; +`; + +const InitialsBadge = styled.div` + width: 40px; + height: 40px; + border-radius: 999px; + display: inline-flex; + align-items: center; + justify-content: center; + background: #e58dc4; + color: white; + font-weight: 600; +`; + +const MemberName = styled.div` + display: flex; + flex-direction: column; +`; + +const MemberEmail = styled.span` + color: var(--gray-7); +`; + +const RoleBadge = styled.span<{ $role: OrganizationRole }>` + display: inline-flex; + align-items: center; + justify-content: center; + min-width: 52px; + padding: 3px 10px; + border-radius: 6px; + font-size: 12px; + font-weight: 600; + color: ${(props) => (props.$role === 'Admin' ? '#6f3ff5' : '#2f6cf6')}; + border: 1px solid + ${(props) => (props.$role === 'Admin' ? '#8b5cf6' : '#6b8cff')}; + background: ${(props) => (props.$role === 'Admin' ? '#f3ebff' : '#eef4ff')}; +`; + +const NativeSelect = styled.select` + width: 100%; + min-height: 40px; + border: 1px solid var(--gray-4); + border-radius: 6px; + padding: 8px 12px; + background: white; + + &:disabled { + background: var(--gray-3); + color: var(--gray-6); + cursor: not-allowed; + } +`; + +const ActionArea = styled.div` + display: flex; + justify-content: flex-end; +`; + +const InlineConfirm = styled.div` + display: flex; + align-items: center; + gap: 8px; + justify-content: flex-end; +`; + +const SearchResults = styled.div` + border: 1px solid var(--gray-4); + border-radius: 6px; + background: white; + overflow: hidden; + margin-top: 8px; +`; + +const SearchResultRow = styled.button<{ $selected?: boolean }>` + width: 100%; + display: flex; + align-items: center; + gap: 12px; + padding: 10px 12px; + border: none; + background: ${(props) => (props.$selected ? 'var(--gray-3)' : 'white')}; + text-align: left; + cursor: pointer; + + &:hover { + background: var(--gray-3); + } +`; + +const SearchPlaceholder = styled.div` + padding: 14px 12px; + color: var(--gray-7); +`; + +const getInitials = (name: string) => + (name || 'U') + .trim() + .split(/\s+/) + .slice(0, 2) + .map((part) => part[0]) + .join('') + .toUpperCase(); + +export default function ProjectAccessControlPage() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [saving, setSaving] = useState(false); + const [visible, setVisible] = useState(false); + const [members, setMembers] = useState([]); + const [availableMembers, setAvailableMembers] = useState< + AvailableProjectMember[] + >([]); + const [canManageAccess, setCanManageAccess] = useState(false); + const [confirmingMemberId, setConfirmingMemberId] = useState( + null, + ); + + const memberSearch = Form.useWatch('memberSearch', form) as string | undefined; + const selectedMemberId = Form.useWatch('organizationMemberId', form) as + | number + | undefined; + + const loadAccess = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/projects/access/current'); + const payload = (await response.json()) as ProjectAccessResponse; + if (!response.ok) { + throw new Error(payload.error || 'Failed to load project access'); + } + setMembers(payload.members || []); + setAvailableMembers(payload.availableMembers || []); + setCanManageAccess(Boolean(payload.canManageAccess)); + } catch (error: any) { + message.error(error.message || 'Failed to load project access'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadAccess(); + }, []); + + const filteredAvailableMembers = useMemo(() => { + const keyword = `${memberSearch || ''}`.trim().toLowerCase(); + if (!keyword) { + return availableMembers; + } + return availableMembers.filter((member) => + `${member.name} ${member.email}`.toLowerCase().includes(keyword), + ); + }, [availableMembers, memberSearch]); + + const openAddMemberModal = () => { + form.resetFields(); + form.setFieldsValue({ + permission: 'Owner', + memberSearch: '', + }); + setVisible(true); + }; + + const addMember = async () => { + try { + const values = await form.validateFields(); + const response = await fetch('/api/v1/projects/access/current', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + organizationMemberId: values.organizationMemberId, + permission: values.permission, + }), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to add member'); + } + message.success('Project member added successfully.'); + setVisible(false); + form.resetFields(); + await loadAccess(); + } catch (error: any) { + if (error?.errorFields) return; + message.error(error.message || 'Failed to add member'); + } finally { + setSaving(false); + } + }; + + const updatePermission = async ( + memberId: number, + permission: ProjectPermission, + ) => { + try { + const response = await fetch(`/api/v1/projects/access/${memberId}`, { + method: 'PATCH', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ permission }), + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to update permission'); + } + setMembers((prev) => + prev.map((member) => (member.id === memberId ? payload : member)), + ); + message.success('Project permission updated successfully.'); + } catch (error: any) { + message.error(error.message || 'Failed to update permission'); + } + }; + + const removeMember = async (memberId: number) => { + try { + const response = await fetch(`/api/v1/projects/access/${memberId}`, { + method: 'DELETE', + }); + const payload = await response.json(); + if (!response.ok) { + throw new Error(payload.error || 'Failed to remove member'); + } + setMembers((prev) => prev.filter((member) => member.id !== memberId)); + message.success('Project member removed successfully.'); + } catch (error: any) { + message.error(error.message || 'Failed to remove member'); + } finally { + setConfirmingMemberId(null); + } + }; + + return ( + + Add member + + } + > + +
+ Manage access + + +
Name
+
Email
+
Organization role
+
Permission
+
+ + {members.map((member) => ( + + + {getInitials(member.name)} + + + {member.name} + {member.isCurrentUser ? ' (me)' : ''} + + + + {member.email} + + {member.organizationRole} + + + void updatePermission( + member.id, + event.target.value as ProjectPermission, + ) + } + > + {PERMISSION_OPTIONS.map((permission) => ( + + ))} + + + {!member.canRemove || !canManageAccess ? null : confirmingMemberId === member.id ? ( + + + Are you sure? + + + + + ) : ( + + )} + + + ))} + +
+ + + { + setVisible(false); + form.resetFields(); + }} + onOk={() => { + setSaving(true); + void addMember(); + }} + confirmLoading={saving} + okText="Add" + destroyOnClose + > +
+ + + + member.id === selectedMemberId) + ?.name || '' + : '' + } + readOnly + /> + + + + + + + + {filteredAvailableMembers.length ? ( + filteredAvailableMembers.map((member) => ( + form.setFieldsValue({ organizationMemberId: member.id })} + > + {getInitials(member.name)} + + {member.name} + {member.email} + + + )) + ) : ( + + No organization members are available to add. + + )} + + + + + {PERMISSION_OPTIONS.map((permission) => ( + + ))} + + + + + + ); +} diff --git a/wren-ui/src/pages/project/danger-zone.tsx b/wren-ui/src/pages/project/danger-zone.tsx new file mode 100644 index 0000000000..5ad315d2de --- /dev/null +++ b/wren-ui/src/pages/project/danger-zone.tsx @@ -0,0 +1,201 @@ +import { useEffect, useMemo, useState } from 'react'; +import { Button, Modal, Typography, message } from 'antd'; +import { useRouter } from 'next/router'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; +import { Path } from '@/utils/enum'; +import { useResetCurrentProjectMutation } from '@/apollo/client/graphql/settings.generated'; + +type OrganizationRole = 'Admin' | 'Member'; + +interface ProjectAccessMember { + userId: number; + organizationRole: OrganizationRole; + isCurrentUser: boolean; +} + +interface ProjectAccessResponse { + members: ProjectAccessMember[]; + currentUserId: number | null; + error?: string; +} + +const IntroText = styled(Typography.Text)` + display: block; + margin-top: 12px; + color: var(--gray-7); +`; + +const WarningText = styled(Typography.Text)` + display: block; + margin-top: 12px; + color: var(--red-6); +`; + +const DangerPanel = styled.div` + margin-top: 24px; + border: 1px solid var(--red-5); + border-radius: 4px; + background: white; +`; + +const DangerRow = styled.div` + display: flex; + align-items: center; + justify-content: space-between; + gap: 24px; + padding: 20px 24px; + + & + & { + border-top: 1px solid var(--gray-4); + } +`; + +const DangerCopy = styled.div` + min-width: 0; +`; + +const DangerTitle = styled.div` + color: var(--gray-8); + font-weight: 600; + margin-bottom: 6px; +`; + +const DangerDescription = styled(Typography.Text)` + color: var(--gray-6); +`; + +const readJsonResponse = async (response: Response): Promise => { + const text = await response.text(); + return (text ? JSON.parse(text) : {}) as T; +}; + +export default function ProjectDangerZonePage() { + const router = useRouter(); + const [loading, setLoading] = useState(true); + const [access, setAccess] = useState(null); + const [resetCurrentProject, { loading: resetting, client }] = + useResetCurrentProjectMutation({ + onError: (error) => + message.error(error.message || 'Failed to reset project'), + }); + + const loadProjectAccess = async () => { + setLoading(true); + try { + const response = await fetch('/api/v1/projects/access/current', { + cache: 'no-store', + headers: { + 'Cache-Control': 'no-cache', + }, + }); + const payload = await readJsonResponse(response); + if (!response.ok) { + throw new Error(payload.error || 'Failed to load project access'); + } + setAccess(payload); + } catch (error: any) { + message.error(error.message || 'Failed to load project access'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadProjectAccess(); + }, []); + + const currentMember = useMemo( + () => + access?.members.find( + (member) => + member.isCurrentUser || member.userId === access.currentUserId, + ) || null, + [access], + ); + const canDeleteProject = currentMember?.organizationRole === 'Admin'; + + const resetProject = async () => { + await resetCurrentProject(); + await client.clearStore(); + message.success('Project reset successfully.'); + await router.push(Path.OnboardingConnection); + }; + + const deleteProject = async () => { + await resetCurrentProject(); + await client.clearStore(); + message.success('Project deleted successfully.'); + await router.push(Path.OrganizationGeneral); + }; + + return ( + + + Use Danger Zone for destructive project actions. + Actions in this section cannot be undone. + + + + + Reset project + + Resetting a project removes its current settings and records, + including data connection information, Modeling page + information, and Home page threads. + + + + + + + + Delete project + + Only organization admins can delete a project. Deleting a + project permanently removes access to its ask records for all + owners and members. + + + + + + + + ); +} diff --git a/wren-ui/src/pages/project/general.tsx b/wren-ui/src/pages/project/general.tsx new file mode 100644 index 0000000000..4f9fd7b09b --- /dev/null +++ b/wren-ui/src/pages/project/general.tsx @@ -0,0 +1,279 @@ +import { useEffect, useMemo, useState } from 'react'; +import { Button, Form, Input, Select, Typography, message } from 'antd'; +import CopyOutlined from '@ant-design/icons/CopyOutlined'; +import styled from 'styled-components'; +import OrganizationSettingsLayout from '@/components/organization/SettingsLayout'; +import { LoadingWrapper } from '@/components/PageLoading'; +import { + useGetSettingsQuery, + useUpdateCurrentProjectMutation, +} from '@/apollo/client/graphql/settings.generated'; +import { + ProjectLanguage, + WorkspaceProjectType, +} from '@/apollo/client/graphql/__types__'; +import { getLanguageText } from '@/utils/language'; + +interface CurrentProjectRecord { + id: string; + displayName: string; + projectType: WorkspaceProjectType; + isCurrent: boolean; + hasDataSource: boolean; + type?: string | null; + createdAt?: string; + updatedAt?: string; +} + +interface CurrentProjectResponse { + currentProject: CurrentProjectRecord | null; + projects: CurrentProjectRecord[]; + error?: string; +} + +const SettingsCard = styled.div` + margin-top: 16px; + border: 1px solid var(--gray-4); + border-radius: 4px; + padding: 20px 28px 28px; + background: white; +`; + +const DetailsTitle = styled(Typography.Title)` + && { + margin-top: 20px; + margin-bottom: 0; + color: var(--gray-8); + } +`; + +const InlineRow = styled.div` + display: flex; + align-items: flex-start; + gap: 16px; + max-width: 860px; + margin-bottom: 20px; +`; + +const LabelCell = styled.div` + width: 160px; + text-align: right; + color: var(--gray-8); + flex-shrink: 0; + padding-top: 8px; +`; + +const FieldCell = styled.div` + flex: 1; +`; + +const Actions = styled.div` + margin-left: 176px; + display: flex; + gap: 8px; +`; + +const HelperText = styled(Typography.Text)` + display: block; + margin-top: 6px; + color: var(--gray-6); +`; + +const languageOptions = Object.keys(ProjectLanguage).map((key) => ({ + label: getLanguageText(key as ProjectLanguage), + value: key, +})); + +export default function ProjectGeneralPage() { + const [form] = Form.useForm(); + const [loading, setLoading] = useState(true); + const [project, setProject] = useState(null); + const { data: settingsData } = useGetSettingsQuery(); + const [initialValues, setInitialValues] = useState({ + displayName: '', + language: ProjectLanguage.EN, + }); + + const [updateCurrentProject, { loading: saving }] = + useUpdateCurrentProjectMutation({ + onError: (error) => + message.error(error.message || 'Failed to update project'), + onCompleted: () => { + message.success('Project updated successfully.'); + }, + }); + + const loadProject = async () => { + setLoading(true); + try { + const currentProjectResponse = await fetch('/api/v1/projects/current'); + const currentProjectPayload = + (await currentProjectResponse.json()) as CurrentProjectResponse; + if (!currentProjectResponse.ok) { + throw new Error( + currentProjectPayload.error || 'Failed to load current project', + ); + } + + const currentProject = currentProjectPayload.currentProject; + if (!currentProject) { + throw new Error('No current project found'); + } + + const values = { + displayName: currentProject.displayName || '', + language: settingsData?.settings?.language || ProjectLanguage.EN, + }; + + setProject(currentProject); + setInitialValues(values); + form.setFieldsValue(values); + } catch (error: any) { + message.error(error.message || 'Failed to load project settings'); + } finally { + setLoading(false); + } + }; + + useEffect(() => { + void loadProject(); + }, [settingsData?.settings?.language]); + + const currentDisplayName = Form.useWatch('displayName', form); + const currentLanguage = Form.useWatch('language', form); + const hasChanges = useMemo( + () => + (currentDisplayName || '').trim() !== initialValues.displayName || + currentLanguage !== initialValues.language, + [currentDisplayName, currentLanguage, initialValues], + ); + + const resetChanges = () => { + form.setFieldsValue(initialValues); + }; + + const copyProjectId = async () => { + if (!project?.id) return; + try { + await navigator.clipboard.writeText(String(project.id)); + message.success('Project ID copied.'); + } catch { + message.error('Failed to copy project ID'); + } + }; + + const saveChanges = async () => { + try { + const values = await form.validateFields(); + await updateCurrentProject({ + variables: { + data: { + displayName: values.displayName.trim(), + language: values.language, + }, + }, + }); + const updatedValues = { + displayName: values.displayName.trim(), + language: values.language, + }; + setInitialValues(updatedValues); + setProject((current) => + current + ? { + ...current, + displayName: updatedValues.displayName, + } + : current, + ); + form.setFieldsValue(updatedValues); + } catch (error: any) { + if (error?.errorFields) return; + message.error(error.message || 'Failed to update project'); + } + }; + + return ( + + +
+ Details + +
+ + Project ID: + + } + onClick={() => void copyProjectId()} + /> + } + /> + + + + + Project name: + + + + + + + + + Project language: + + + + + + + + Email : + + + + + + + + + + + + + + + Account Verification + + + To ensure security, account verification is handled by the configured + identity provider or invitation flow for your organization. + + + + Change password + + + + Password : + + + + You will receive a confirmation link via email when password + reset delivery is configured. + + + + + + + ); +} diff --git a/wren-ui/src/utils/enum/menu.ts b/wren-ui/src/utils/enum/menu.ts index 2d5e55c354..4ab0941256 100644 --- a/wren-ui/src/utils/enum/menu.ts +++ b/wren-ui/src/utils/enum/menu.ts @@ -3,4 +3,7 @@ export enum MENU_KEY { INSTRUCTIONS = 'instructions', API_HISTORY = 'api-history', API_REFERENCE = 'api-reference', + ADMIN_USERS = 'admin-users', + ADMIN_ROLES = 'admin-roles', + ADMIN_ASSIGNMENTS = 'admin-assignments', } diff --git a/wren-ui/src/utils/enum/path.ts b/wren-ui/src/utils/enum/path.ts index 10f5fa8091..2bf5289333 100644 --- a/wren-ui/src/utils/enum/path.ts +++ b/wren-ui/src/utils/enum/path.ts @@ -12,4 +12,18 @@ export enum Path { KnowledgeInstructions = '/knowledge/instructions', APIManagement = '/api-management', APIManagementHistory = '/api-management/history', + Administration = '/administration', + AdministrationUsers = '/administration/users', + AdministrationRoles = '/administration/roles', + AdministrationAssignments = '/administration/assignments', + Organization = '/organization', + OrganizationGeneral = '/organization/general', + OrganizationMembers = '/organization/members', + OrganizationDangerZone = '/organization/danger-zone', + ProjectGeneral = '/project/general', + ProjectAccessControl = '/project/access-control', + ProjectDangerZone = '/project/danger-zone', + ProjectCreate = '/projects/create', + UserProfile = '/user/profile', + UserDangerZone = '/user/danger-zone', } diff --git a/wren-ui/src/utils/errorHandler.tsx b/wren-ui/src/utils/errorHandler.tsx index 06e9070092..0254870c80 100644 --- a/wren-ui/src/utils/errorHandler.tsx +++ b/wren-ui/src/utils/errorHandler.tsx @@ -226,7 +226,10 @@ class CreateRelationshipErrorHandler extends ErrorHandler { public getErrorMessage(error: GraphQLError) { switch (error.extensions?.code) { default: - return 'Failed to create relationship.'; + return replaceMessage( + 'Failed to create %{relationship}.', + error.message, + ); } } } diff --git a/wren-ui/tools/knex.js b/wren-ui/tools/knex.js index 0324ea66f4..93e1a72d60 100644 --- a/wren-ui/tools/knex.js +++ b/wren-ui/tools/knex.js @@ -1,10 +1,110 @@ +const fs = require('fs'); +const path = require('path'); + const DB_TYPE = process.env.DB_TYPE; // export DB_TYPE=pg const PG_URL = process.env.PG_URL; const DEBUG = process.env.DEBUG === 'true'; // export DEBUG=true const SQLITE_FILE = process.env.SQLITE_FILE; // export SQLITE_FILE=./db.sqlite3 -const getKnex = () => { - if (DB_TYPE === 'pg') { +const APP_TABLE_ORDER = [ + 'project', + 'roles', + 'users', + 'user_roles', + 'organization', + 'organization_members', + 'organization_member_projects', + 'organization_invitations', + 'organization_invitation_projects', + 'model', + 'model_column', + 'model_nested_column', + 'relation', + 'metric', + 'metric_measure', + 'view', + 'deploy_log', + 'thread', + 'thread_response', + 'schema_change', + 'learning', + 'dashboard', + 'dashboard_item', + 'sql_pair', + 'instruction', + 'dashboard_item_refresh_job', + 'asking_task', + 'api_history', +]; + +const normalizeDbType = (dbType) => + (dbType || 'mssql').trim().toLowerCase().replace(/[-_ ]/g, ''); + +const ensureMssqlConfig = () => { + if (process.env.MSSQL_URL) { + return; + } + + const missingFields = [ + ['MSSQL_HOST', process.env.MSSQL_HOST], + ['MSSQL_DATABASE', process.env.MSSQL_DATABASE], + ['MSSQL_USER', process.env.MSSQL_USER], + ['MSSQL_PASSWORD', process.env.MSSQL_PASSWORD], + ].filter(([, value]) => !value); + + if (missingFields.length > 0) { + throw new Error( + `MSSQL is the required Wren UI application database. Missing configuration: ${missingFields + .map(([key]) => key) + .join(', ')}`, + ); + } +}; + +const parseBooleanUrlParam = (searchParams, key, fallback) => { + const value = searchParams.get(key); + if (value === null) return fallback; + return value.toLowerCase() === 'true'; +}; + +const getMssqlConnection = () => { + if (process.env.MSSQL_URL) { + const url = new URL(process.env.MSSQL_URL); + return { + server: url.hostname, + port: url.port ? parseInt(url.port) : 1433, + database: decodeURIComponent(url.pathname.replace(/^\//, '')), + user: decodeURIComponent(url.username), + password: decodeURIComponent(url.password), + options: { + encrypt: parseBooleanUrlParam(url.searchParams, 'encrypt', false), + trustServerCertificate: parseBooleanUrlParam( + url.searchParams, + 'trustServerCertificate', + true, + ), + }, + }; + } + + return { + server: process.env.MSSQL_HOST || 'localhost', + port: process.env.MSSQL_PORT ? parseInt(process.env.MSSQL_PORT) : 1433, + database: process.env.MSSQL_DATABASE || 'wren_ui', + user: process.env.MSSQL_USER, + password: process.env.MSSQL_PASSWORD, + options: { + encrypt: process.env.MSSQL_ENCRYPT === 'true', + trustServerCertificate: + process.env.MSSQL_TRUST_SERVER_CERTIFICATE !== 'false', + }, + }; +}; + +const getKnex = (options = {}) => { + const dbType = normalizeDbType(options.dbType || DB_TYPE); + + if (dbType === 'pg' || dbType === 'postgres' || dbType === 'postgresql') { console.log('using pg'); /* eslint-disable @typescript-eslint/no-var-requires */ return require('knex')({ @@ -13,20 +113,245 @@ const getKnex = () => { debug: DEBUG, pool: { min: 2, max: 10 }, }); - } else { - console.log('using sqlite'); + } + + if (dbType === 'mssql' || dbType === 'sqlserver') { + console.log('using mssql'); + ensureMssqlConfig(); /* eslint-disable @typescript-eslint/no-var-requires */ return require('knex')({ - client: 'better-sqlite3', - connection: { - filename: SQLITE_FILE, - }, - useNullAsDefault: true, + client: 'mssql', + connection: getMssqlConnection(), + debug: DEBUG, + pool: { min: 2, max: 10 }, }); } + + throw new Error( + `Unsupported DB_TYPE "${options.dbType || DB_TYPE || ''}". Wren UI application storage now requires MSSQL.`, + ); +}; + +const getSqliteFile = () => { + const appRoot = path.resolve(__dirname, '..'); + const candidates = [ + SQLITE_FILE, + path.join(appRoot, 'data', 'db.sqlite3'), + path.join(appRoot, 'db.sqlite3'), + ].filter(Boolean); + + const sqliteFile = candidates.find((candidate) => fs.existsSync(candidate)); + if (!sqliteFile) { + throw new Error( + `SQLite source database not found. Set SQLITE_FILE explicitly. Checked: ${candidates.join( + ', ', + )}`, + ); + } + return sqliteFile; +}; + +const getSourceTables = async (sourceDb) => { + const rows = await sourceDb.raw(` + SELECT name + FROM sqlite_master + WHERE type = 'table' + AND name NOT LIKE 'sqlite_%' + AND name NOT IN ('knex_migrations', 'knex_migrations_lock') + `); + const tableNames = rows.map((row) => row.name); + const orderedTables = APP_TABLE_ORDER.filter((table) => + tableNames.includes(table), + ); + const remainingTables = tableNames + .filter((table) => !orderedTables.includes(table)) + .sort(); + return [...orderedTables, ...remainingTables]; +}; + +const getTargetColumns = async (targetDb, tableName) => { + const rows = await targetDb('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ TABLE_SCHEMA: 'dbo', TABLE_NAME: tableName }) + .orderBy('ORDINAL_POSITION', 'asc'); + return rows.map((row) => row.COLUMN_NAME); +}; + +const getTargetIdentityColumns = async (targetDb, tableName) => { + const rows = await targetDb('INFORMATION_SCHEMA.COLUMNS') + .select('COLUMN_NAME') + .where({ TABLE_SCHEMA: 'dbo', TABLE_NAME: tableName }) + .whereRaw( + "COLUMNPROPERTY(OBJECT_ID(TABLE_SCHEMA + '.' + TABLE_NAME), COLUMN_NAME, 'IsIdentity') = 1", + ); + return rows.map((row) => row.COLUMN_NAME); +}; + +const parseCount = (row) => Number(row.count || row.Count || row[''] || 0); + +const getTableCount = async (db, tableName) => { + const [row] = await db(tableName).count({ count: '*' }); + return parseCount(row); +}; + +const ensureTargetIsEmpty = async (targetDb, tableNames) => { + const nonEmptyTables = []; + for (const tableName of tableNames) { + if (!(await targetDb.schema.hasTable(tableName))) { + continue; + } + const count = await getTableCount(targetDb, tableName); + if (count > 0) { + nonEmptyTables.push(`${tableName} (${count})`); + } + } + + if (nonEmptyTables.length > 0 && process.env.MIGRATE_OVERWRITE !== 'true') { + throw new Error( + `Refusing to copy into a non-empty MSSQL database. Non-empty tables: ${nonEmptyTables.join( + ', ', + )}. Set MIGRATE_OVERWRITE=true to delete target Wren UI rows first.`, + ); + } +}; + +const setForeignKeysEnabled = async (targetDb, enabled) => { + const rows = await targetDb + .select( + targetDb.raw( + "QUOTENAME(SCHEMA_NAME(schema_id)) + '.' + QUOTENAME(name) AS full_name", + ), + ) + .from('sys.tables') + .where({ is_ms_shipped: 0 }); + + for (const row of rows) { + const fullName = row.full_name || row.fullName; + if (enabled) { + await targetDb.raw( + `ALTER TABLE ${fullName} WITH CHECK CHECK CONSTRAINT ALL`, + ); + } else { + await targetDb.raw(`ALTER TABLE ${fullName} NOCHECK CONSTRAINT ALL`); + } + } +}; + +const clearTargetTables = async (targetDb, tableNames) => { + for (const tableName of [...tableNames].reverse()) { + if (await targetDb.schema.hasTable(tableName)) { + await targetDb(tableName).delete(); + } + } +}; + +const copyTable = async (sourceDb, targetDb, tableName) => { + if (!(await targetDb.schema.hasTable(tableName))) { + console.log(`Skipping ${tableName}: target table does not exist`); + return { tableName, sourceCount: 0, targetCount: 0 }; + } + + const sourceRows = await sourceDb(tableName).select('*'); + if (sourceRows.length === 0) { + console.log(`Copied ${tableName}: 0 rows`); + return { + tableName, + sourceCount: 0, + targetCount: await getTableCount(targetDb, tableName), + }; + } + + const targetColumns = await getTargetColumns(targetDb, tableName); + const commonColumns = targetColumns.filter((column) => + Object.prototype.hasOwnProperty.call(sourceRows[0], column), + ); + const rows = sourceRows.map((row) => + Object.fromEntries(commonColumns.map((column) => [column, row[column]])), + ); + const chunkSize = Math.max( + 1, + Math.floor(1800 / Math.max(commonColumns.length, 1)), + ); + const identityColumns = await getTargetIdentityColumns(targetDb, tableName); + const hasIdentityId = + commonColumns.includes('id') && identityColumns.includes('id'); + + if (hasIdentityId) { + await targetDb.raw(`SET IDENTITY_INSERT [dbo].[${tableName}] ON`); + } + + try { + for (let index = 0; index < rows.length; index += chunkSize) { + await targetDb(tableName).insert(rows.slice(index, index + chunkSize)); + } + } finally { + if (hasIdentityId) { + await targetDb.raw(`SET IDENTITY_INSERT [dbo].[${tableName}] OFF`); + } + } + + const targetCount = await getTableCount(targetDb, tableName); + console.log(`Copied ${tableName}: ${sourceRows.length} rows`); + return { tableName, sourceCount: sourceRows.length, targetCount }; +}; + +const migrateSqliteToMssql = async () => { + const sqliteFile = getSqliteFile(); + const sourceDb = getKnex({ dbType: 'sqlite', sqliteFile }); + const targetDb = getKnex({ dbType: 'mssql' }); + + try { + const migrationsDir = path.resolve(__dirname, '..', 'migrations'); + console.log(`Migrating Wren UI application tables from ${sqliteFile}`); + console.log('Running MSSQL schema migrations'); + await targetDb.migrate.latest({ directory: migrationsDir }); + + const tableNames = await getSourceTables(sourceDb); + await ensureTargetIsEmpty(targetDb, tableNames); + + await targetDb.transaction(async (trx) => { + await setForeignKeysEnabled(trx, false); + try { + if (process.env.MIGRATE_OVERWRITE === 'true') { + await clearTargetTables(trx, tableNames); + } + + const verification = []; + for (const tableName of tableNames) { + verification.push(await copyTable(sourceDb, trx, tableName)); + } + + const mismatches = verification.filter( + ({ sourceCount, targetCount }) => sourceCount !== targetCount, + ); + if (mismatches.length > 0) { + throw new Error( + `Row-count verification failed: ${mismatches + .map( + ({ tableName, sourceCount, targetCount }) => + `${tableName} sqlite=${sourceCount} mssql=${targetCount}`, + ) + .join(', ')}`, + ); + } + } finally { + await setForeignKeysEnabled(trx, true); + } + }); + + console.log('SQLite to MSSQL migration completed successfully.'); + } finally { + await sourceDb.destroy(); + await targetDb.destroy(); + } }; const main = async () => { + if (process.env.MIGRATE_SQLITE_TO_MSSQL === 'true') { + await migrateSqliteToMssql(); + return; + } + const knex = getKnex(); const query = knex.queryBuilder(); @@ -36,7 +361,12 @@ const main = async () => { .whereIn('id', [7, 8]); console.log(projects); - process.exit(0); + await knex.destroy(); }; -main(); +main() + .then(() => process.exit(0)) + .catch((error) => { + console.error(error); + process.exit(1); + });