diff --git a/.env.example b/.env.example index 1193904..ffdc314 100644 --- a/.env.example +++ b/.env.example @@ -9,3 +9,13 @@ EMBEDDING_PROVIDER=local EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2 ENABLE_SCHEMA_RAG=false ENABLE_GOLDEN_RECORDS=false +ENABLE_EXTERNAL_CONNECTIONS=true +ALLOW_PRIVATE_DATABASE_HOSTS=false +CONNECTION_ENCRYPTION_KEY= +AUTH_SIGNING_KEY= +DATABASE_CONNECTION_TIMEOUT_SECONDS=5 +DATABASE_STATEMENT_TIMEOUT_MS=10000 +DATABASE_MAX_RESULT_ROWS=500 +DATABASE_MAX_SQL_LENGTH=10000 +SCHEMA_FULL_CONTEXT_TABLE_LIMIT=20 +SCHEMA_RETRIEVAL_TABLE_LIMIT=8 diff --git a/README.md b/README.md index fe3adf3..1eb0d71 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,11 @@ # QueryMindAI > Ask questions, not SQL. -QueryMindAI is an open-source natural-language analytics platform that converts business questions into safe, explainable SQL queries over structured data. +QueryMindAI is an open-source AI data workspace that lets users connect a read-only database and query it safely using natural language. ## Current status -QueryMindAI is an early production-minded release. Implemented today: PostgreSQL schema introspection, optional schema-vector retrieval, provider-based SQL generation, parser validation, read-only execution, verified-example retrieval, query history, demo data, Docker Compose, and Render configuration. Authentication, multi-tenant credential storage, a polished connection manager, and comparative model benchmarks are roadmap items. Some retained dashboard views are explicitly demo UI. +QueryMindAI is an early production-minded release. Implemented today: encrypted saved PostgreSQL connections, public-host SSRF checks, SSL connections, catalog snapshots, structured SQL generation, explicit approval drafts, parser/table validation, read-only execution, verified examples, history, audit events, Docker Compose, and Render configuration. Signed anonymous workspace sessions provide ownership isolation for the public demo; account login and enterprise identity integration remain roadmap work. ## Screenshots @@ -24,27 +24,25 @@ QueryMindAI is an early production-minded release. Implemented today: PostgreSQL ```mermaid flowchart LR - U[Next.js web] --> A[FastAPI API] - A --> P[Query orchestrator] - P --> S[Schema introspection / optional RAG] - P --> V[Verified examples] - P --> L[LLM provider client] - P --> G[sqlglot safety gate] - G --> D[(Read-only PostgreSQL)] + B[Browser] --> W[QueryMindAI Web] --> A[QueryMindAI API] + A --> C[Connection Manager] --> S[Schema Catalog] + S --> L[LLM Query Planner] --> V[SQL Validator] + V --> E[Execution Engine] --> D[(Customer PostgreSQL Database)] ``` ### Query lifecycle ```mermaid flowchart LR - Q[Validate question] --> S[Retrieve schema] --> E[Optional verified example] - E --> L[Generate SQL] --> C[Clean] --> V[Parse + validate] - V --> X[Read-only execute] --> R[Structured response] + Q[Validate question] --> S[Retrieve bounded schema] --> L[Generate structured SQL] + L --> V[Parse, table-check, cap LIMIT] --> D[Store expiring draft] + D --> H[Human review and approval] --> R[Revalidate SQL] + R --> X[Read-only execute with timeout] --> O[Results and history] ``` ## Tech stack -FastAPI, SQLAlchemy, Alembic, PostgreSQL/pgvector, sqlglot, OpenAI-compatible provider API, Next.js 15, React 19, TypeScript, Tailwind CSS, Docker Compose, GitHub Actions, and Render. +FastAPI, SQLAlchemy, Alembic, PostgreSQL, sqlglot, OpenAI-compatible provider API, Next.js 15, React 19, TypeScript, Tailwind CSS, Docker Compose, GitHub Actions, and Render. ## Repository structure @@ -59,7 +57,7 @@ scripts Operator helpers ## Quickstart -Prerequisites: Python 3.12, Node 20, PostgreSQL 16 with pgvector, and optionally Ollama. +Prerequisites: Python 3.12, Node 20, PostgreSQL 16+, and optionally Ollama. ```bash cp apps/api/.env.example apps/api/.env @@ -105,8 +103,10 @@ For the default Compose stack, optional vector features are off so schema intros | `LLM_MODEL`, `LLM_FALLBACK_MODEL` | Primary and retry fallback generation models | | `EMBEDDING_PROVIDER`, `EMBEDDING_MODEL` | Local semantic-retrieval configuration | | `ENABLE_SCHEMA_RAG`, `ENABLE_GOLDEN_RECORDS` | Optional embedding features | -| `ENABLE_EXTERNAL_CONNECTIONS` | Raw connection endpoint; defaults false | -| `MAX_QUERY_ROWS`, `STATEMENT_TIMEOUT_MS` | Execution safeguards | +| `ENABLE_EXTERNAL_CONNECTIONS`, `ALLOW_PRIVATE_DATABASE_HOSTS` | BYOD and SSRF policy flags | +| `CONNECTION_ENCRYPTION_KEY`, `AUTH_SIGNING_KEY` | Backend-only Fernet and session signing secrets | +| `DATABASE_*` | Connect timeout, statement timeout, SQL length, and result caps | +| `SCHEMA_*_TABLE_LIMIT` | Full-context and retrieval bounds | | `NEXT_PUBLIC_API_URL` | Browser-visible API `/api/v1` URL | See the component `.env.example` files for all defaults. @@ -120,6 +120,10 @@ Semantic retrieval uses lazy, process-cached `sentence-transformers/all-MiniLM-L ## API overview - `GET /health`, `GET /ready` +- `POST /api/v1/auth/session` +- `POST/GET/DELETE /api/v1/connections[...]` +- `POST /api/v1/queries/generate`, `POST /api/v1/queries/execute` +- `GET /api/v1/queries/history`, `GET /api/v1/queries/{id}` - `GET /api/v1/schema`, `GET /api/v1/query-history` - `POST /api/v1/query` - `POST /api/v1/verified-examples` @@ -141,9 +145,9 @@ The API validates question size, limits generated SQL length, rejects comments/m ## Limitations and roadmap -Current limitations include PostgreSQL-only execution, provider-dependent SQL quality, unavailable verified-example saving when embeddings are disabled, no auth/multi-tenancy, no encrypted external credential vault, in-memory external schema cache, and demo-only management screens. +Current limitations include PostgreSQL-only execution, publicly reachable hosts only, provider-dependent SQL quality, application-level encryption rather than a managed vault, anonymous browser-bound sessions rather than account authentication, no rate limiter, and no guaranteed detection of inherited write privileges. -Roadmap: authenticated workspaces, secure connection management, API-backed connection pages, richer history review, execution-backed evaluations, additional SQL dialects, observability exports, and reviewed embedding lifecycle tools. +Roadmap: OIDC accounts and organizations, MySQL, richer history review, execution-backed evaluations, observability exports, and a customer-side connector for private databases: `Customer Database → QueryMind Connector inside customer network → outbound TLS → QueryMindAI Cloud`. The connector is documented only and is not implemented. ## Contributing and license diff --git a/apps/api/.env.example b/apps/api/.env.example index 8c75f1c..291ae1f 100644 --- a/apps/api/.env.example +++ b/apps/api/.env.example @@ -13,6 +13,15 @@ EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2 ENABLE_SCHEMA_RAG=false ENABLE_GOLDEN_RECORDS=false ENABLE_EXTERNAL_CONNECTIONS=false +ALLOW_PRIVATE_DATABASE_HOSTS=false +CONNECTION_ENCRYPTION_KEY= +AUTH_SIGNING_KEY= +DATABASE_CONNECTION_TIMEOUT_SECONDS=5 +DATABASE_STATEMENT_TIMEOUT_MS=10000 +DATABASE_MAX_RESULT_ROWS=500 +DATABASE_MAX_SQL_LENGTH=10000 +SCHEMA_FULL_CONTEXT_TABLE_LIMIT=20 +SCHEMA_RETRIEVAL_TABLE_LIMIT=8 MAX_QUERY_ROWS=100 STATEMENT_TIMEOUT_MS=10000 MAX_SQL_LENGTH=20000 diff --git a/apps/api/README.md b/apps/api/README.md index cac35c8..96d315c 100644 --- a/apps/api/README.md +++ b/apps/api/README.md @@ -8,7 +8,7 @@ The modular FastAPI service turns validated questions into schema-aware PostgreS ## Endpoints -`GET /health` is liveness. `GET /ready` checks PostgreSQL. Under `/api/v1`: `GET /schema`, `GET /query-history`, `POST /query`, `POST /verified-examples`; `/golden-record` is a deprecated compatibility alias. `/connect-and-query` returns 403 unless explicitly enabled. +`GET /health` is liveness and `GET /ready` checks the application PostgreSQL database. BYOD endpoints under `/api/v1` include signed session creation; connection create/list/get/test/refresh/schema/delete; query generate/execute/history/get; and saving an executed query as verified. All BYOD resources are scoped to the verified session subject. Legacy demo endpoints remain for compatibility. ## Configuration and providers @@ -28,6 +28,8 @@ uvicorn app.main:app --reload --port 8000 Production startup never creates tables. Alembic creates application tables and pgvector. The demo seed is idempotent. Use `../../scripts/create_readonly_role.sql` as a reviewed template and run query traffic with that role. +Saved customer credentials require `CONNECTION_ENCRYPTION_KEY` (a Fernet key). Signed workspace sessions require a separate `AUTH_SIGNING_KEY`. Generate them with the commands in `../../docs/deployment.md`. BYOD accepts URL or structured PostgreSQL input, requires SSL, blocks unsafe networks by default, and stores normalized schema snapshots without row data. + ## Safety model `sqlglot` parses PostgreSQL and enforces one read-only query, accepts SELECT/CTE SELECT, rejects comments, multiple/write/admin statements, caps LIMIT, and limits SQL size. Execution uses `SET LOCAL TRANSACTION READ ONLY` plus statement timeout. Database permissions remain the primary boundary. diff --git a/apps/api/app/api/v1/endpoints/byod.py b/apps/api/app/api/v1/endpoints/byod.py new file mode 100644 index 0000000..4b9fb03 --- /dev/null +++ b/apps/api/app/api/v1/endpoints/byod.py @@ -0,0 +1,150 @@ +from typing import Annotated + +from fastapi import APIRouter, Depends, status +from sqlalchemy import select, text +from sqlalchemy.orm import Session + +from app.audit.service import record_audit +from app.auth.dependencies import get_current_user_id, issue_session +from app.auth.schemas import SessionResponse +from app.connections.schemas import (ConnectionCreate, ConnectionResponse, ConnectionTestResponse, + SchemaResponse) +from app.connections.security import decrypt_config, normalize_connection_input +from app.connections.service import (create_saved_connection, inspect_and_check, latest_snapshot, masked, + owned_connection, refresh_owned_schema) +from app.core.config import settings +from app.core.exceptions import DependencyError, FeatureDisabledError, NotFoundError, QueryMindError +from app.db.session import get_db +from app.models.byod import DatabaseConnection +from app.query_planner.schemas import (QueryDraftResponse, QueryExecuteRequest, QueryExecutionResponse, + QueryGenerateRequest) +from app.query_planner.service import generate_draft +from app.execution.service import execute_draft +from app.schemas.query_schema import VerifiedExampleRequest +from app.services.golden_record_service import save_golden_record + +router = APIRouter() +UserId = Annotated[str, Depends(get_current_user_id)] +Db = Annotated[Session, Depends(get_db)] + + +def enabled(): + if not settings.ENABLE_EXTERNAL_CONNECTIONS: + raise FeatureDisabledError("Saved external database connections are disabled") + + +@router.post("/auth/session", response_model=SessionResponse) +def create_workspace_session(): + enabled(); token, _, expires_at = issue_session() + return {"access_token": token, "expires_at": expires_at} + + +@router.post("/connections", response_model=ConnectionResponse, status_code=status.HTTP_201_CREATED) +def create_connection(request: ConnectionCreate, user_id: UserId, db: Db): + enabled() + config = normalize_connection_input(request) + try: + connection, warnings = create_saved_connection(db, user_id, request.name, config) + return masked(connection, warnings) + except QueryMindError: + db.rollback(); raise + except Exception as exc: + db.rollback() + raise DependencyError("The database connection could not be established") from exc + + +@router.get("/connections", response_model=list[ConnectionResponse]) +def list_connections(user_id: UserId, db: Db): + enabled() + items = db.scalars(select(DatabaseConnection).where(DatabaseConnection.user_id == user_id) + .order_by(DatabaseConnection.created_at.desc())).all() + return [masked(item) for item in items] + + +@router.get("/connections/{connection_id}", response_model=ConnectionResponse) +def get_connection(connection_id: str, user_id: UserId, db: Db): + enabled(); return masked(owned_connection(db, connection_id, user_id)) + + +@router.post("/connections/{connection_id}/test", response_model=ConnectionTestResponse) +def test_connection(connection_id: str, user_id: UserId, db: Db): + enabled(); saved = owned_connection(db, connection_id, user_id) + try: + _, read_only, warnings = inspect_and_check(decrypt_config(saved.encrypted_connection_data)) + except Exception as exc: + saved.status = "failed"; record_audit(db, user_id, "connection_tested", saved.id, {"status": "failed"}); db.commit() + raise DependencyError("The database connection test failed") from exc + saved.status = "connected"; record_audit(db, user_id, "connection_tested", saved.id, {"status": "connected"}); db.commit() + return {"status": "connected", "read_only": read_only, "warnings": warnings} + + +@router.post("/connections/{connection_id}/refresh-schema", response_model=SchemaResponse) +def refresh_schema(connection_id: str, user_id: UserId, db: Db): + enabled(); saved = owned_connection(db, connection_id, user_id) + snapshot = refresh_owned_schema(db, saved, user_id) + return {"connection_id": saved.id, "schema_hash": snapshot.schema_hash, + "metadata": snapshot.normalized_metadata, "updated_at": snapshot.updated_at} + + +@router.get("/connections/{connection_id}/schema", response_model=SchemaResponse) +def get_schema(connection_id: str, user_id: UserId, db: Db): + enabled(); saved = owned_connection(db, connection_id, user_id); snapshot = latest_snapshot(db, saved.id) + return {"connection_id": saved.id, "schema_hash": snapshot.schema_hash, + "metadata": snapshot.normalized_metadata, "updated_at": snapshot.updated_at} + + +@router.delete("/connections/{connection_id}", status_code=status.HTTP_204_NO_CONTENT) +def delete_connection(connection_id: str, user_id: UserId, db: Db): + enabled(); saved = owned_connection(db, connection_id, user_id) + record_audit(db, user_id, "connection_deleted", saved.id, {"name": saved.name}) + db.execute(text("DELETE FROM golden_records WHERE connection_key = :key"), {"key": saved.id}) + db.execute(text("DELETE FROM schema_embeddings WHERE connection_key = :key"), {"key": saved.id}) + db.delete(saved); db.commit() + + +@router.post("/queries/generate", response_model=QueryDraftResponse) +async def generate_query(request: QueryGenerateRequest, user_id: UserId, db: Db): + enabled(); draft = await generate_draft(db, user_id, request.connection_id, request.question) + return {"draft_id": draft.id, "connection_id": draft.connection_id, "sql": draft.sql, + "explanation": draft.explanation, "assumptions": draft.assumptions, "warnings": draft.warnings, + "confidence": float(draft.confidence), "expires_at": draft.expires_at} + + +@router.post("/queries/execute", response_model=QueryExecutionResponse) +def run_query(request: QueryExecuteRequest, user_id: UserId, db: Db): + enabled(); return execute_draft(db, user_id, request.draft_id) + + +@router.get("/queries/history") +def byod_history(user_id: UserId, db: Db, limit: int = 50): + enabled() + rows = db.execute(text(""" + SELECT id, connection_id, question, generated_sql AS sql, execution_status, row_count, + duration_ms, warnings, created_at FROM query_history + WHERE user_id = :user_id ORDER BY created_at DESC LIMIT :limit + """), {"user_id": user_id, "limit": min(max(limit, 1), 100)}).mappings().all() + return {"items": [dict(row) for row in rows]} + + +@router.post("/queries/{query_id}/verified-example", status_code=status.HTTP_201_CREATED) +async def save_query_as_verified(query_id: int, user_id: UserId, db: Db): + enabled() + row = db.execute(text("""SELECT connection_id, question, generated_sql FROM query_history + WHERE id = :id AND user_id = :user_id AND execution_status = 'success'"""), + {"id": query_id, "user_id": user_id}).mappings().first() + if not row: raise NotFoundError("Successful query history item was not found") + owned_connection(db, row["connection_id"], user_id) + saved = await save_golden_record(VerifiedExampleRequest(question=row["question"], sql=row["generated_sql"], + connection_key=row["connection_id"], source="approved_execution"), db) + return {"id": saved.id, "status": "saved", "message": "Verified example saved."} + + +@router.get("/queries/{query_id}") +def get_query(query_id: int, user_id: UserId, db: Db): + enabled() + row = db.execute(text(""" + SELECT id, connection_id, question, generated_sql AS sql, execution_status, row_count, + duration_ms, warnings, created_at FROM query_history WHERE id = :id AND user_id = :user_id + """), {"id": query_id, "user_id": user_id}).mappings().first() + if not row: raise NotFoundError("Query history item was not found") + return dict(row) diff --git a/apps/api/app/api/v1/endpoints/query.py b/apps/api/app/api/v1/endpoints/query.py index 1b96e53..4efceab 100644 --- a/apps/api/app/api/v1/endpoints/query.py +++ b/apps/api/app/api/v1/endpoints/query.py @@ -4,20 +4,17 @@ from app.core.config import settings from app.core.exceptions import FeatureDisabledError -from app.db.dynamic_session import get_dynamic_session from app.db.session import get_db -from app.schemas.query_schema import (ConnectQueryRequest, QueryRequest, QueryResponse, +from app.schemas.query_schema import (QueryRequest, QueryResponse, VerifiedExampleRequest, VerifiedExampleResponse) from app.services.golden_record_service import save_golden_record -from app.services.query_service import execute_query -from app.services.schema_introspect_service import get_dynamic_schema router = APIRouter() @router.post("/query", response_model=QueryResponse) async def query_db(request: QueryRequest, db: Session = Depends(get_db)): - return await execute_query(db, request.question) + raise FeatureDisabledError("Combined generation and execution is disabled; use /queries/generate then /queries/execute") @router.get("/schema") @@ -43,17 +40,6 @@ def query_history(limit: int = 20, db: Session = Depends(get_db)): return {"items": [dict(row) for row in rows], "source": "real_api_data"} -@router.post("/connect-and-query", response_model=QueryResponse) -async def connect_and_query(request: ConnectQueryRequest): - if not settings.ENABLE_EXTERNAL_CONNECTIONS: - raise FeatureDisabledError("External database connections are disabled") - db = get_dynamic_session(request.db_config.model_dump()) - try: - return await execute_query(db, request.question, schema_override=get_dynamic_schema(db)) - finally: - db.close() - - @router.post("/verified-examples", response_model=VerifiedExampleResponse, status_code=status.HTTP_201_CREATED) async def add_verified_example(request: VerifiedExampleRequest, db: Session = Depends(get_db)): row = await save_golden_record(request, db) diff --git a/apps/api/app/audit/__init__.py b/apps/api/app/audit/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/audit/service.py b/apps/api/app/audit/service.py new file mode 100644 index 0000000..99440e6 --- /dev/null +++ b/apps/api/app/audit/service.py @@ -0,0 +1,8 @@ +from sqlalchemy.orm import Session + +from app.models.byod import AuditLog + + +def record_audit(db: Session, user_id: str, action: str, connection_id: str | None = None, metadata: dict | None = None): + # Metadata must contain identifiers/status only—never credentials, URLs, SQL results, or decrypted config. + db.add(AuditLog(user_id=user_id, connection_id=connection_id, action=action, event_metadata=metadata or {})) diff --git a/apps/api/app/auth/__init__.py b/apps/api/app/auth/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/auth/dependencies.py b/apps/api/app/auth/dependencies.py new file mode 100644 index 0000000..bbf572c --- /dev/null +++ b/apps/api/app/auth/dependencies.py @@ -0,0 +1,42 @@ +import base64 +import hashlib +import hmac +import json +import time +import uuid +from typing import Annotated + +from fastapi import Header + +from app.core.config import settings +from app.core.exceptions import AuthenticationError, DependencyError + + +def _key() -> bytes: + if not settings.AUTH_SIGNING_KEY: + raise DependencyError("AUTH_SIGNING_KEY is required for saved connection sessions") + return settings.AUTH_SIGNING_KEY.encode() + + +def issue_session() -> tuple[str, str, int]: + user_id = f"usr_{uuid.uuid4().hex}" + expires_at = int(time.time()) + settings.AUTH_SESSION_TTL_SECONDS + body = base64.urlsafe_b64encode(json.dumps({"sub": user_id, "exp": expires_at}, separators=(",", ":")).encode()).decode().rstrip("=") + signature = base64.urlsafe_b64encode(hmac.new(_key(), body.encode(), hashlib.sha256).digest()).decode().rstrip("=") + return f"{body}.{signature}", user_id, expires_at + + +def get_current_user_id(authorization: Annotated[str | None, Header(max_length=4096)] = None) -> str: + if not authorization or not authorization.startswith("Bearer "): + raise AuthenticationError("A signed workspace session is required") + try: + body, supplied = authorization[7:].split(".", 1) + expected = base64.urlsafe_b64encode(hmac.new(_key(), body.encode(), hashlib.sha256).digest()).decode().rstrip("=") + if not hmac.compare_digest(supplied, expected): + raise ValueError + payload = json.loads(base64.urlsafe_b64decode(body + "=" * (-len(body) % 4))) + if int(payload["exp"]) <= int(time.time()) or not str(payload["sub"]).startswith("usr_"): + raise ValueError + return str(payload["sub"]) + except (ValueError, KeyError, json.JSONDecodeError) as exc: + raise AuthenticationError("The workspace session is invalid or expired") from exc diff --git a/apps/api/app/auth/schemas.py b/apps/api/app/auth/schemas.py new file mode 100644 index 0000000..b4bc994 --- /dev/null +++ b/apps/api/app/auth/schemas.py @@ -0,0 +1,7 @@ +from pydantic import BaseModel + + +class SessionResponse(BaseModel): + access_token: str + token_type: str = "bearer" + expires_at: int diff --git a/apps/api/app/catalog/__init__.py b/apps/api/app/catalog/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/catalog/service.py b/apps/api/app/catalog/service.py new file mode 100644 index 0000000..e9c4996 --- /dev/null +++ b/apps/api/app/catalog/service.py @@ -0,0 +1,72 @@ +import hashlib +import json + +from sqlalchemy import inspect + + +SYSTEM_SCHEMAS = {"information_schema", "pg_catalog", "pg_toast"} + + +def introspect_postgresql(connection) -> dict: + inspector = inspect(connection) + schemas = [] + for schema_name in sorted(name for name in inspector.get_schema_names() if name not in SYSTEM_SCHEMAS and not name.startswith("pg_")): + objects = [] + table_names = [(name, "table") for name in inspector.get_table_names(schema=schema_name)] + view_names = [(name, "view") for name in inspector.get_view_names(schema=schema_name)] + for object_name, object_type in sorted(table_names + view_names): + columns = [{"name": column["name"], "type": str(column["type"]), "nullable": bool(column["nullable"])} + for column in inspector.get_columns(object_name, schema=schema_name)] + pk = inspector.get_pk_constraint(object_name, schema=schema_name) or {} + fks = [] + for fk in inspector.get_foreign_keys(object_name, schema=schema_name): + fks.append({"columns": fk.get("constrained_columns", []), "referred_schema": fk.get("referred_schema") or schema_name, + "referred_table": fk.get("referred_table"), "referred_columns": fk.get("referred_columns", [])}) + indexes = [{"name": index.get("name"), "columns": index.get("column_names", []), "unique": bool(index.get("unique"))} + for index in inspector.get_indexes(object_name, schema=schema_name)] if object_type == "table" else [] + objects.append({"name": object_name, "type": object_type, "columns": columns, + "primary_key": pk.get("constrained_columns", []), "foreign_keys": fks, "indexes": indexes}) + schemas.append({"name": schema_name, "objects": objects}) + return {"database_type": "postgresql", "schemas": schemas} + + +def schema_hash(metadata: dict) -> str: + encoded = json.dumps(metadata, sort_keys=True, separators=(",", ":")).encode() + return hashlib.sha256(encoded).hexdigest() + + +def qualified_tables(metadata: dict) -> set[str]: + result = set() + for schema in metadata.get("schemas", []): + for obj in schema.get("objects", []): + result.add(obj["name"].lower()) + result.add(f"{schema['name']}.{obj['name']}".lower()) + return result + + +def schema_context(metadata: dict, question: str, full_limit: int, retrieval_limit: int) -> str: + objects = [(schema["name"], obj) for schema in metadata.get("schemas", []) for obj in schema.get("objects", [])] + if len(objects) > full_limit: + terms = {term.lower() for term in question.replace("_", " ").split() if len(term) > 2} + scored = [] + for schema_name, obj in objects: + haystack = " ".join([obj["name"], *(col["name"] for col in obj["columns"])]).lower() + scored.append((sum(term in haystack for term in terms), schema_name, obj)) + selected = [(schema, obj) for _, schema, obj in sorted(scored, key=lambda item: (-item[0], item[1], item[2]["name"]))[:retrieval_limit]] + selected_names = {(schema, obj["name"]) for schema, obj in selected} + # Add directly related tables while respecting the same bounded context budget. + for schema, obj in list(selected): + for fk in obj.get("foreign_keys", []): + target = (fk["referred_schema"], fk["referred_table"]) + if target not in selected_names: + match = next(((s, candidate) for s, candidate in objects if (s, candidate["name"]) == target), None) + if match and len(selected) < retrieval_limit: + selected.append(match); selected_names.add(target) + objects = selected + lines = [] + for schema, obj in objects: + cols = ", ".join(f"{col['name']} {col['type']}" for col in obj["columns"]) + lines.append(f"{schema}.{obj['name']} ({cols})") + for fk in obj.get("foreign_keys", []): + lines.append(f"FK {schema}.{obj['name']}.{','.join(fk['columns'])} -> {fk['referred_schema']}.{fk['referred_table']}.{','.join(fk['referred_columns'])}") + return "\n".join(lines) diff --git a/apps/api/app/connections/__init__.py b/apps/api/app/connections/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/connections/schemas.py b/apps/api/app/connections/schemas.py new file mode 100644 index 0000000..09c969f --- /dev/null +++ b/apps/api/app/connections/schemas.py @@ -0,0 +1,52 @@ +from datetime import datetime +from typing import Literal + +from pydantic import BaseModel, Field, model_validator + + +SSLMode = Literal["require", "verify-ca", "verify-full"] + + +class ConnectionCreate(BaseModel): + name: str = Field(min_length=1, max_length=128) + connection_string: str | None = Field(default=None, min_length=1, max_length=2048) + host: str | None = Field(default=None, max_length=253) + port: int = Field(default=5432, ge=1, le=65535) + database: str | None = Field(default=None, min_length=1, max_length=128) + username: str | None = Field(default=None, min_length=1, max_length=128) + password: str | None = Field(default=None, min_length=1, max_length=512) + ssl_mode: SSLMode = "require" + + @model_validator(mode="after") + def one_input_mode(self): + structured = all([self.host, self.database, self.username, self.password]) + if bool(self.connection_string) == bool(structured): + raise ValueError("Provide either a connection string or all structured connection fields") + return self + + +class ConnectionResponse(BaseModel): + id: str + name: str + database_type: str + host: str + database: str + username: str + ssl_mode: str + status: str + last_connected_at: datetime | None + created_at: datetime + warnings: list[str] = [] + + +class ConnectionTestResponse(BaseModel): + status: str + read_only: bool | None + warnings: list[str] + + +class SchemaResponse(BaseModel): + connection_id: str + schema_hash: str + metadata: dict + updated_at: datetime diff --git a/apps/api/app/connections/security.py b/apps/api/app/connections/security.py new file mode 100644 index 0000000..6d2faa4 --- /dev/null +++ b/apps/api/app/connections/security.py @@ -0,0 +1,83 @@ +import ipaddress +import json +import socket +from functools import lru_cache + +from cryptography.fernet import Fernet, InvalidToken +from sqlalchemy.engine import URL, make_url + +from app.core.config import settings +from app.core.exceptions import ConnectionSecurityError, DependencyError + + +def normalize_connection_input(request) -> dict: + if request.connection_string: + try: + url = make_url(request.connection_string) + except Exception as exc: + raise ConnectionSecurityError("Invalid PostgreSQL connection string") from exc + if url.drivername not in {"postgresql", "postgres", "postgresql+psycopg2"}: + raise ConnectionSecurityError("Only PostgreSQL connection strings are supported") + if not all([url.host, url.database, url.username, url.password]): + raise ConnectionSecurityError("Connection string must include host, database, username, and password") + query = dict(url.query) + ssl_mode = query.get("sslmode", request.ssl_mode) + config = {"host": url.host, "port": url.port or 5432, "database": url.database, + "username": url.username, "password": url.password, "ssl_mode": ssl_mode} + else: + config = {"host": request.host, "port": request.port, "database": request.database, + "username": request.username, "password": request.password, "ssl_mode": request.ssl_mode} + if config["port"] != 5432: + raise ConnectionSecurityError("Only the standard PostgreSQL port 5432 is supported") + if config["ssl_mode"] not in {"require", "verify-ca", "verify-full"}: + raise ConnectionSecurityError("SSL mode must be require, verify-ca, or verify-full") + return config + + +def _blocked(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: + return ip.is_loopback or ip.is_private or ip.is_link_local or ip.is_unspecified or ip.is_reserved or ip.is_multicast + + +def validate_public_host(host: str) -> list[str]: + if not host or host.rstrip(".").lower() == "localhost": + raise ConnectionSecurityError("Localhost database hosts are not allowed") + try: + literal = ipaddress.ip_address(host.strip("[]")) + addresses = [literal] + except ValueError: + try: + addresses = list({ipaddress.ip_address(item[4][0]) for item in socket.getaddrinfo(host, 5432, type=socket.SOCK_STREAM)}) + except (socket.gaierror, ValueError) as exc: + raise ConnectionSecurityError("Database hostname could not be resolved") from exc + if not addresses: + raise ConnectionSecurityError("Database hostname did not resolve") + if not settings.ALLOW_PRIVATE_DATABASE_HOSTS and any(_blocked(ip) for ip in addresses): + raise ConnectionSecurityError("Private, local, link-local, and reserved database hosts are blocked") + return [str(ip) for ip in addresses] + + +@lru_cache +def _fernet() -> Fernet: + if not settings.CONNECTION_ENCRYPTION_KEY: + raise DependencyError("CONNECTION_ENCRYPTION_KEY is required for saved connections") + try: + return Fernet(settings.CONNECTION_ENCRYPTION_KEY.encode()) + except Exception as exc: + raise DependencyError("CONNECTION_ENCRYPTION_KEY is invalid") from exc + + +def encrypt_config(config: dict) -> str: + return _fernet().encrypt(json.dumps(config, separators=(",", ":")).encode()).decode() + + +def decrypt_config(value: str) -> dict: + try: + return json.loads(_fernet().decrypt(value.encode()).decode()) + except (InvalidToken, ValueError, json.JSONDecodeError) as exc: + raise DependencyError("Stored connection credentials could not be decrypted") from exc + + +def sqlalchemy_url(config: dict) -> URL: + return URL.create("postgresql+psycopg2", username=config["username"], password=config["password"], + host=config["host"], port=config["port"], database=config["database"], + query={"sslmode": config["ssl_mode"]}) diff --git a/apps/api/app/connections/service.py b/apps/api/app/connections/service.py new file mode 100644 index 0000000..5bb98f9 --- /dev/null +++ b/apps/api/app/connections/service.py @@ -0,0 +1,105 @@ +from datetime import datetime, timezone + +from sqlalchemy import create_engine, select, text +from sqlalchemy.orm import Session +from sqlalchemy.pool import NullPool + +from app.audit.service import record_audit +from app.catalog.service import introspect_postgresql, schema_hash +from app.connections.security import decrypt_config, encrypt_config, sqlalchemy_url, validate_public_host +from app.core.config import settings +from app.core.exceptions import DependencyError, NotFoundError +from app.models.byod import DatabaseConnection, SchemaSnapshot + + +def owned_connection(db: Session, connection_id: str, user_id: str) -> DatabaseConnection: + connection = db.scalar(select(DatabaseConnection).where(DatabaseConnection.id == connection_id, + DatabaseConnection.user_id == user_id)) + if not connection: + raise NotFoundError("Database connection was not found") + return connection + + +def masked(connection: DatabaseConnection, warnings: list[str] | None = None) -> dict: + return {"id": connection.id, "name": connection.name, "database_type": connection.database_type, + "host": connection.masked_host, "database": connection.database_name, "username": connection.username, + "ssl_mode": connection.ssl_mode, "status": connection.status, + "last_connected_at": connection.last_connected_at, "created_at": connection.created_at, + "warnings": warnings or []} + + +def external_engine(config: dict): + validate_public_host(config["host"]) # DNS rebinding defense: re-check immediately before connect. + return create_engine(sqlalchemy_url(config), poolclass=NullPool, pool_pre_ping=True, + connect_args={"connect_timeout": settings.DATABASE_CONNECTION_TIMEOUT_SECONDS}) + + +def inspect_and_check(config: dict) -> tuple[dict, bool | None, list[str]]: + engine = external_engine(config) + warnings = [] + try: + with engine.connect() as conn: + metadata = introspect_postgresql(conn) + write_privilege = conn.execute(text(""" + SELECT EXISTS ( + SELECT 1 FROM information_schema.table_privileges + WHERE grantee = current_user AND privilege_type IN ('INSERT','UPDATE','DELETE','TRUNCATE','TRIGGER','REFERENCES') + ) + """)).scalar() + default_read_only = str(conn.execute(text("SHOW default_transaction_read_only")).scalar()).lower() == "on" + if write_privilege: + read_only = False + warnings.append("The database role appears to have write privileges; use a dedicated read-only role.") + elif default_read_only: + read_only = True + else: + read_only = None + warnings.append("No direct table write grants were found, but read-only status could not be conclusively verified.") + if not metadata["schemas"]: + warnings.append("No accessible application schemas were found.") + return metadata, read_only, warnings + finally: + engine.dispose() + + +def save_snapshot(db: Session, connection_id: str, metadata: dict) -> SchemaSnapshot: + digest = schema_hash(metadata) + existing = db.scalar(select(SchemaSnapshot).where(SchemaSnapshot.connection_id == connection_id, + SchemaSnapshot.schema_hash == digest)) + if existing: + return existing + snapshot = SchemaSnapshot(connection_id=connection_id, normalized_metadata=metadata, schema_hash=digest) + db.add(snapshot); db.flush() + return snapshot + + +def latest_snapshot(db: Session, connection_id: str) -> SchemaSnapshot: + snapshot = db.scalar(select(SchemaSnapshot).where(SchemaSnapshot.connection_id == connection_id) + .order_by(SchemaSnapshot.created_at.desc())) + if not snapshot: + raise NotFoundError("No schema snapshot is available; refresh the schema first") + return snapshot + + +def create_saved_connection(db: Session, user_id: str, name: str, config: dict) -> tuple[DatabaseConnection, list[str]]: + validate_public_host(config["host"]) + metadata, _, warnings = inspect_and_check(config) + connection = DatabaseConnection(user_id=user_id, name=name, encrypted_connection_data=encrypt_config(config), + masked_host=config["host"], database_name=config["database"], username=config["username"], + ssl_mode=config["ssl_mode"], status="connected", last_connected_at=datetime.now(timezone.utc)) + db.add(connection); db.flush() + save_snapshot(db, connection.id, metadata) + record_audit(db, user_id, "connection_created", connection.id, {"database_type": "postgresql"}) + record_audit(db, user_id, "connection_tested", connection.id, {"status": "connected"}) + db.commit(); db.refresh(connection) + return connection, warnings + + +def refresh_owned_schema(db: Session, connection: DatabaseConnection, user_id: str) -> SchemaSnapshot: + config = decrypt_config(connection.encrypted_connection_data) + metadata, _, _ = inspect_and_check(config) + snapshot = save_snapshot(db, connection.id, metadata) + connection.status = "connected"; connection.last_connected_at = datetime.now(timezone.utc) + record_audit(db, user_id, "schema_refreshed", connection.id, {"schema_hash": snapshot.schema_hash}) + db.commit(); db.refresh(snapshot) + return snapshot diff --git a/apps/api/app/core/config.py b/apps/api/app/core/config.py index f4292c2..d1c31b4 100644 --- a/apps/api/app/core/config.py +++ b/apps/api/app/core/config.py @@ -21,6 +21,17 @@ class Settings(BaseSettings): ENABLE_SCHEMA_RAG: bool = False ENABLE_GOLDEN_RECORDS: bool = False ENABLE_EXTERNAL_CONNECTIONS: bool = False + ALLOW_PRIVATE_DATABASE_HOSTS: bool = False + CONNECTION_ENCRYPTION_KEY: str | None = None + AUTH_SIGNING_KEY: str | None = None + AUTH_SESSION_TTL_SECONDS: int = Field(2592000, ge=300, le=31536000) + DATABASE_CONNECTION_TIMEOUT_SECONDS: int = Field(5, ge=1, le=30) + DATABASE_STATEMENT_TIMEOUT_MS: int = Field(10000, ge=100, le=300000) + DATABASE_MAX_RESULT_ROWS: int = Field(500, ge=1, le=10000) + DATABASE_MAX_SQL_LENGTH: int = Field(10000, ge=100, le=100000) + SCHEMA_FULL_CONTEXT_TABLE_LIMIT: int = Field(20, ge=1, le=500) + SCHEMA_RETRIEVAL_TABLE_LIMIT: int = Field(8, ge=1, le=100) + QUERY_DRAFT_TTL_SECONDS: int = Field(900, ge=60, le=86400) MAX_QUESTION_LENGTH: int = Field(1000, ge=1, le=10000) MAX_SQL_LENGTH: int = Field(20000, ge=100, le=100000) MAX_QUERY_ROWS: int = Field(100, ge=1, le=10000) diff --git a/apps/api/app/core/exceptions.py b/apps/api/app/core/exceptions.py index 2d3ec30..50dce9d 100644 --- a/apps/api/app/core/exceptions.py +++ b/apps/api/app/core/exceptions.py @@ -30,3 +30,18 @@ class DependencyError(QueryMindError): class FeatureDisabledError(QueryMindError): status_code = 403 code = "feature_disabled" + + +class NotFoundError(QueryMindError): + status_code = 404 + code = "not_found" + + +class AuthenticationError(QueryMindError): + status_code = 401 + code = "authentication_required" + + +class ConnectionSecurityError(QueryMindError): + status_code = 400 + code = "connection_rejected" diff --git a/apps/api/app/db/base.py b/apps/api/app/db/base.py index e69de29..fa2b68a 100644 --- a/apps/api/app/db/base.py +++ b/apps/api/app/db/base.py @@ -0,0 +1,5 @@ +from sqlalchemy.orm import DeclarativeBase + + +class Base(DeclarativeBase): + pass diff --git a/apps/api/app/db/dynamic_session.py b/apps/api/app/db/dynamic_session.py deleted file mode 100644 index 6f5b3d4..0000000 --- a/apps/api/app/db/dynamic_session.py +++ /dev/null @@ -1,11 +0,0 @@ -from sqlalchemy import URL, create_engine -from sqlalchemy.orm import sessionmaker - -from app.core.config import settings - - -def get_dynamic_session(db_config): - url = URL.create("postgresql+psycopg2", username=db_config["user"], password=db_config["password"], - host=db_config["host"], port=db_config["port"], database=db_config["database"]) - engine = create_engine(url, connect_args={"connect_timeout": settings.DATABASE_CONNECT_TIMEOUT_SECONDS}, pool_pre_ping=True) - return sessionmaker(bind=engine)() diff --git a/apps/api/app/execution/__init__.py b/apps/api/app/execution/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/execution/service.py b/apps/api/app/execution/service.py new file mode 100644 index 0000000..751c400 --- /dev/null +++ b/apps/api/app/execution/service.py @@ -0,0 +1,56 @@ +import time + +from sqlalchemy import text +from sqlalchemy.orm import Session + +from app.audit.service import record_audit +from app.catalog.service import qualified_tables +from app.connections.security import decrypt_config +from app.connections.service import external_engine, latest_snapshot, owned_connection +from app.core.config import settings +from app.core.exceptions import DependencyError +from app.query_planner.service import owned_draft +from app.services.validation_service import validate_sql + + +def execute_draft(db: Session, user_id: str, draft_id: str) -> dict: + draft = owned_draft(db, draft_id, user_id) + saved = owned_connection(db, draft.connection_id, user_id) + snapshot = latest_snapshot(db, saved.id) + safe_sql = validate_sql(draft.sql, max_rows=settings.DATABASE_MAX_RESULT_ROWS, + allowed_tables=qualified_tables(snapshot.normalized_metadata), + max_sql_length=settings.DATABASE_MAX_SQL_LENGTH) + config = decrypt_config(saved.encrypted_connection_data) + engine = external_engine(config) + started = time.perf_counter() + try: + with engine.connect() as conn: + transaction = conn.begin() + try: + conn.execute(text("SET TRANSACTION READ ONLY")) + conn.execute(text("SELECT set_config('statement_timeout', :timeout, true)"), + {"timeout": f"{settings.DATABASE_STATEMENT_TIMEOUT_MS}ms"}) + # Safety invariant: safe_sql was revalidated immediately before this call. + result = conn.execute(text(safe_sql)) + columns = list(result.keys()) + rows = [dict(row._mapping) for row in result.fetchmany(settings.DATABASE_MAX_RESULT_ROWS)] + finally: + transaction.rollback() + except Exception as exc: + raise DependencyError("The customer database query failed or timed out") from exc + finally: + engine.dispose() + duration = round((time.perf_counter() - started) * 1000) + history = db.execute(text(""" + INSERT INTO query_history + (connection_key, connection_id, user_id, question, generated_sql, execution_status, + row_count, duration_ms, warnings, status, verified_example_used) + VALUES (:key, :connection_id, :user_id, :question, :sql, 'success', :rows, :duration, + CAST(:warnings AS jsonb), 'success', false) + RETURNING id + """), {"key": saved.id, "connection_id": saved.id, "user_id": user_id, "question": draft.question, + "sql": safe_sql, "rows": len(rows), "duration": duration, "warnings": "[]"}).scalar_one() + record_audit(db, user_id, "query_executed", saved.id, {"query_id": history, "row_count": len(rows), "duration_ms": duration}) + db.commit() + return {"query_id": history, "sql": safe_sql, "columns": columns, "rows": rows, + "row_count": len(rows), "duration_ms": duration, "warnings": draft.warnings or []} diff --git a/apps/api/app/history/__init__.py b/apps/api/app/history/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/llm/__init__.py b/apps/api/app/llm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/main.py b/apps/api/app/main.py index 8e5bcd6..88aa707 100644 --- a/apps/api/app/main.py +++ b/apps/api/app/main.py @@ -8,6 +8,7 @@ from sqlalchemy import text from app.api.v1.endpoints.query import router as query_router +from app.api.v1.endpoints.byod import router as byod_router from app.core.config import settings from app.core.exceptions import QueryMindError from app.core.logger import configure_logging, request_id_context @@ -19,6 +20,7 @@ app.add_middleware(CORSMiddleware, allow_origins=settings.CORS_ALLOW_ORIGINS, allow_credentials=True, allow_methods=["*"], allow_headers=["*"]) app.include_router(query_router, prefix="/api/v1", tags=["analytics"]) +app.include_router(byod_router, prefix="/api/v1", tags=["connections"]) @app.middleware("http") diff --git a/apps/api/app/models/__init__.py b/apps/api/app/models/__init__.py new file mode 100644 index 0000000..657b6c9 --- /dev/null +++ b/apps/api/app/models/__init__.py @@ -0,0 +1,3 @@ +from app.models.byod import AuditLog, DatabaseConnection, QueryDraft, SchemaSnapshot + +__all__ = ["AuditLog", "DatabaseConnection", "QueryDraft", "SchemaSnapshot"] diff --git a/apps/api/app/models/byod.py b/apps/api/app/models/byod.py new file mode 100644 index 0000000..9ab30bf --- /dev/null +++ b/apps/api/app/models/byod.py @@ -0,0 +1,70 @@ +import uuid +from datetime import datetime, timezone + +from sqlalchemy import JSON, DateTime, ForeignKey, Integer, String, Text, UniqueConstraint +from sqlalchemy.orm import Mapped, mapped_column + +from app.db.base import Base + + +def new_id(prefix: str) -> str: + return f"{prefix}_{uuid.uuid4().hex}" + + +def now() -> datetime: + return datetime.now(timezone.utc) + + +class DatabaseConnection(Base): + __tablename__ = "database_connections" + id: Mapped[str] = mapped_column(String(64), primary_key=True, default=lambda: new_id("conn")) + user_id: Mapped[str] = mapped_column(String(128), index=True) + name: Mapped[str] = mapped_column(String(128)) + database_type: Mapped[str] = mapped_column(String(32), default="postgresql") + encrypted_connection_data: Mapped[str] = mapped_column(Text) + masked_host: Mapped[str] = mapped_column(String(255)) + database_name: Mapped[str] = mapped_column(String(128)) + username: Mapped[str] = mapped_column(String(128)) + ssl_mode: Mapped[str] = mapped_column(String(32)) + status: Mapped[str] = mapped_column(String(32), default="pending") + last_connected_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True)) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now) + updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now, onupdate=now) + + +class SchemaSnapshot(Base): + __tablename__ = "schema_snapshots" + __table_args__ = (UniqueConstraint("connection_id", "schema_hash", name="uq_connection_schema_hash"),) + id: Mapped[str] = mapped_column(String(64), primary_key=True, default=lambda: new_id("schema")) + connection_id: Mapped[str] = mapped_column(ForeignKey("database_connections.id", ondelete="CASCADE"), index=True) + schema_name: Mapped[str] = mapped_column(String(128), default="all") + normalized_metadata: Mapped[dict] = mapped_column(JSON) + schema_hash: Mapped[str] = mapped_column(String(64)) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now) + updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now, onupdate=now) + + +class QueryDraft(Base): + __tablename__ = "query_drafts" + id: Mapped[str] = mapped_column(String(64), primary_key=True, default=lambda: new_id("draft")) + user_id: Mapped[str] = mapped_column(String(128), index=True) + connection_id: Mapped[str] = mapped_column(ForeignKey("database_connections.id", ondelete="CASCADE"), index=True) + question: Mapped[str] = mapped_column(Text) + sql: Mapped[str] = mapped_column(Text) + explanation: Mapped[str] = mapped_column(Text) + assumptions: Mapped[list] = mapped_column(JSON, default=list) + warnings: Mapped[list] = mapped_column(JSON, default=list) + confidence: Mapped[str] = mapped_column(String(16)) + schema_hash: Mapped[str] = mapped_column(String(64)) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now) + expires_at: Mapped[datetime] = mapped_column(DateTime(timezone=True)) + + +class AuditLog(Base): + __tablename__ = "audit_logs" + id: Mapped[str] = mapped_column(String(64), primary_key=True, default=lambda: new_id("audit")) + user_id: Mapped[str] = mapped_column(String(128), index=True) + connection_id: Mapped[str | None] = mapped_column(String(64), index=True) + action: Mapped[str] = mapped_column(String(64), index=True) + event_metadata: Mapped[dict] = mapped_column("metadata", JSON, default=dict) + created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), default=now) diff --git a/apps/api/app/organizations/__init__.py b/apps/api/app/organizations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/query_planner/__init__.py b/apps/api/app/query_planner/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/app/query_planner/schemas.py b/apps/api/app/query_planner/schemas.py new file mode 100644 index 0000000..d10bfcc --- /dev/null +++ b/apps/api/app/query_planner/schemas.py @@ -0,0 +1,48 @@ +from datetime import datetime +from typing import Any + +from pydantic import BaseModel, Field + +from app.core.config import settings + + +class QueryGenerateRequest(BaseModel): + connection_id: str = Field(min_length=1, max_length=64) + question: str = Field(min_length=1, max_length=settings.MAX_QUESTION_LENGTH) + + +class QueryDraftResponse(BaseModel): + draft_id: str + connection_id: str + sql: str + explanation: str + assumptions: list[str] + warnings: list[str] + confidence: float + expires_at: datetime + + +class QueryExecuteRequest(BaseModel): + draft_id: str = Field(min_length=1, max_length=64) + + +class QueryExecutionResponse(BaseModel): + query_id: int + sql: str + columns: list[str] + rows: list[dict[str, Any]] + row_count: int + duration_ms: int + warnings: list[str] + + +class QueryHistoryItem(BaseModel): + id: int + connection_id: str | None + question: str + sql: str + execution_status: str | None + row_count: int + duration_ms: int + warnings: list[str] + created_at: datetime diff --git a/apps/api/app/query_planner/service.py b/apps/api/app/query_planner/service.py new file mode 100644 index 0000000..76b9b4c --- /dev/null +++ b/apps/api/app/query_planner/service.py @@ -0,0 +1,45 @@ +from datetime import datetime, timedelta, timezone + +from sqlalchemy import select +from sqlalchemy.orm import Session + +from app.audit.service import record_audit +from app.catalog.service import qualified_tables, schema_context +from app.connections.service import latest_snapshot, owned_connection +from app.core.config import settings +from app.core.exceptions import NotFoundError +from app.models.byod import QueryDraft +from app.services.llm_service import generate_sql +from app.services.validation_service import validate_sql + + +async def generate_draft(db: Session, user_id: str, connection_id: str, question: str) -> QueryDraft: + owned_connection(db, connection_id, user_id) + snapshot = latest_snapshot(db, connection_id) + context = schema_context(snapshot.normalized_metadata, question, settings.SCHEMA_FULL_CONTEXT_TABLE_LIMIT, + settings.SCHEMA_RETRIEVAL_TABLE_LIMIT) + generated = await generate_sql(question, context) + safe_sql = validate_sql(generated.sql, max_rows=settings.DATABASE_MAX_RESULT_ROWS, + allowed_tables=qualified_tables(snapshot.normalized_metadata), + max_sql_length=settings.DATABASE_MAX_SQL_LENGTH) + draft = QueryDraft(user_id=user_id, connection_id=connection_id, question=question.strip(), sql=safe_sql, + explanation=generated.explanation, assumptions=generated.assumptions, warnings=[], + confidence=str(generated.confidence), schema_hash=snapshot.schema_hash, + expires_at=datetime.now(timezone.utc) + timedelta(seconds=settings.QUERY_DRAFT_TTL_SECONDS)) + db.add(draft); db.flush() + record_audit(db, user_id, "query_generated", connection_id, {"draft_id": draft.id}) + db.commit(); db.refresh(draft) + return draft + + +def owned_draft(db: Session, draft_id: str, user_id: str) -> QueryDraft: + draft = db.scalar(select(QueryDraft).where(QueryDraft.id == draft_id, QueryDraft.user_id == user_id)) + if not draft: + raise NotFoundError("Query draft was not found") + expires = draft.expires_at + if expires.tzinfo is None: + expires = expires.replace(tzinfo=timezone.utc) + if expires <= datetime.now(timezone.utc): + raise NotFoundError("Query draft has expired; generate SQL again") + owned_connection(db, draft.connection_id, user_id) + return draft diff --git a/apps/api/app/schemas/query_schema.py b/apps/api/app/schemas/query_schema.py index 419704e..238088b 100644 --- a/apps/api/app/schemas/query_schema.py +++ b/apps/api/app/schemas/query_schema.py @@ -35,19 +35,6 @@ class QueryResponse(BaseModel): metadata: QueryMetadata -class DBConfig(BaseModel): - host: str = Field(min_length=1, max_length=253, pattern=r"^[A-Za-z0-9._:-]+$") - port: int = Field(ge=1, le=65535) - database: str = Field(min_length=1, max_length=63, pattern=r"^[A-Za-z0-9_-]+$") - user: str = Field(min_length=1, max_length=63) - password: str = Field(min_length=1, max_length=512) - - -class ConnectQueryRequest(BaseModel): - db_config: DBConfig - question: str = Field(min_length=1, max_length=settings.MAX_QUESTION_LENGTH) - - class VerifiedExampleRequest(BaseModel): question: str = Field(min_length=1, max_length=settings.MAX_QUESTION_LENGTH) sql: str = Field(min_length=1, max_length=settings.MAX_SQL_LENGTH) diff --git a/apps/api/app/services/golden_record_service.py b/apps/api/app/services/golden_record_service.py index 4555bac..ce9aca5 100644 --- a/apps/api/app/services/golden_record_service.py +++ b/apps/api/app/services/golden_record_service.py @@ -1,11 +1,13 @@ import logging +import json +import math from sqlalchemy import text from sqlalchemy.orm import Session from app.core.config import settings from app.core.exceptions import DependencyError -from app.core.embedding_provider import as_pgvector, embed_text +from app.core.embedding_provider import embed_text from app.schemas.query_schema import VerifiedExampleRequest from app.services.validation_service import validate_sql @@ -13,7 +15,7 @@ async def _embedding(value: str) -> str: - return as_pgvector(await embed_text(value)) + return await embed_text(value) async def get_golden_record(question: str, db: Session, connection_key: str | None = None) -> dict | None: @@ -21,13 +23,13 @@ async def get_golden_record(question: str, db: Session, connection_key: str | No return None try: vector = await _embedding(question) - row = db.execute(text(""" - SELECT question, sql_query, 1 - (embedding <=> CAST(:vector AS vector)) AS similarity - FROM golden_records WHERE connection_key = :key - ORDER BY embedding <=> CAST(:vector AS vector) LIMIT 1 - """), {"key": connection_key or settings.DEFAULT_CONNECTION_KEY, "vector": vector}).fetchone() - if row and float(row.similarity) >= settings.GOLDEN_RECORD_SIMILARITY_THRESHOLD: - return {"question": row.question, "sql": row.sql_query, "similarity": float(row.similarity)} + rows = db.execute(text("SELECT question, sql_query, embedding FROM golden_records WHERE connection_key = :key AND embedding IS NOT NULL"), + {"key": connection_key or settings.DEFAULT_CONNECTION_KEY}).fetchall() + def cosine(other): + dot=sum(a*b for a,b in zip(vector,other)); denom=math.sqrt(sum(a*a for a in vector))*math.sqrt(sum(b*b for b in other)); return dot/denom if denom else 0 + best = max(((cosine(row.embedding), row) for row in rows), default=None, key=lambda item:item[0]) + if best and best[0] >= settings.GOLDEN_RECORD_SIMILARITY_THRESHOLD: + return {"question": best[1].question, "sql": best[1].sql_query, "similarity": best[0]} except Exception: db.rollback() logger.warning("Verified-example retrieval unavailable", extra={"stage": "verified_example_retrieval"}) @@ -35,21 +37,19 @@ async def get_golden_record(question: str, db: Session, connection_key: str | No async def save_golden_record(request: VerifiedExampleRequest, db: Session): - if not settings.ENABLE_GOLDEN_RECORDS: - raise DependencyError("Verified examples are disabled") safe_sql = validate_sql(request.sql) try: - vector = await _embedding(request.question) + vector = await _embedding(request.question) if settings.ENABLE_GOLDEN_RECORDS else None row = db.execute(text(""" INSERT INTO golden_records (connection_key, question, sql_query, embedding, reviewer, source) - VALUES (:key, :question, :sql, CAST(:vector AS vector), :reviewer, :source) + VALUES (:key, :question, :sql, CAST(:vector AS jsonb), :reviewer, :source) ON CONFLICT (connection_key, question, sql_query) DO UPDATE SET reviewer = COALESCE(EXCLUDED.reviewer, golden_records.reviewer), source = COALESCE(EXCLUDED.source, golden_records.source), updated_at = NOW() RETURNING id, created_at """), {"key": request.connection_key, "question": request.question.strip(), "sql": safe_sql, - "vector": vector, "reviewer": request.reviewer, "source": request.source}).fetchone() + "vector": json.dumps(vector) if vector is not None else None, "reviewer": request.reviewer, "source": request.source}).fetchone() db.commit() return row except Exception as exc: diff --git a/apps/api/app/services/schema_service.py b/apps/api/app/services/schema_service.py index 9ed301d..fcd6fd6 100644 --- a/apps/api/app/services/schema_service.py +++ b/apps/api/app/services/schema_service.py @@ -1,10 +1,11 @@ import logging +import math from sqlalchemy import text from sqlalchemy.orm import Session from app.core.config import settings -from app.core.embedding_provider import as_pgvector, embed_text +from app.core.embedding_provider import embed_text from app.services.schema_introspect_service import get_dynamic_schema logger = logging.getLogger(__name__) @@ -15,14 +16,12 @@ async def get_relevant_schema(question: str, db: Session, connection_key: str | if not settings.ENABLE_SCHEMA_RAG: return fallback try: - vector = as_pgvector(await embed_text(question)) - rows = db.execute(text(""" - SELECT schema_ddl FROM schema_embeddings - WHERE connection_key = :key - ORDER BY embedding <=> CAST(:vector AS vector) - LIMIT :limit - """), {"key": connection_key or settings.DEFAULT_CONNECTION_KEY, "vector": vector, "limit": top_k}) - schemas = [row[0] for row in rows] + vector = await embed_text(question) + rows = db.execute(text("SELECT schema_ddl, embedding FROM schema_embeddings WHERE connection_key = :key AND embedding IS NOT NULL"), + {"key": connection_key or settings.DEFAULT_CONNECTION_KEY}).fetchall() + def cosine(other): + dot=sum(a*b for a,b in zip(vector,other)); denom=math.sqrt(sum(a*a for a in vector))*math.sqrt(sum(b*b for b in other)); return dot/denom if denom else 0 + schemas = [row.schema_ddl for row in sorted(rows, key=lambda row: cosine(row.embedding), reverse=True)[:top_k]] return "\n".join(schemas) if schemas else fallback except Exception: db.rollback() diff --git a/apps/api/app/services/validation_service.py b/apps/api/app/services/validation_service.py index abc840f..0fd032a 100644 --- a/apps/api/app/services/validation_service.py +++ b/apps/api/app/services/validation_service.py @@ -12,10 +12,11 @@ ) -def validate_sql(sql: str, max_rows: int | None = None) -> str: +def validate_sql(sql: str, max_rows: int | None = None, allowed_tables: set[str] | None = None, + max_sql_length: int | None = None) -> str: if not sql or not sql.strip(): raise UnsafeSQLError("The provider returned empty SQL") - if len(sql) > settings.MAX_SQL_LENGTH: + if len(sql) > (max_sql_length or settings.MAX_SQL_LENGTH): raise UnsafeSQLError("Generated SQL exceeds the configured maximum length") if "/*" in sql or "*/" in sql or re.search(r"--[^\n\r]*", sql): raise UnsafeSQLError("SQL comments are not allowed") @@ -32,6 +33,22 @@ def validate_sql(sql: str, max_rows: int | None = None) -> str: raise UnsafeSQLError("Write and administrative SQL statements are not allowed") if statement.find(exp.Select) is None: raise UnsafeSQLError("Only SELECT queries and WITH ... SELECT queries are allowed") + if statement.find(exp.Anonymous) is not None and any( + function.name.upper() in {"CALL", "DO", "PG_READ_FILE", "PG_LS_DIR", "DBLINK", "LO_IMPORT", "LO_EXPORT"} + for function in statement.find_all(exp.Anonymous) + ): + raise UnsafeSQLError("Stored procedures and unsafe functions are not allowed") + if allowed_tables is not None: + cte_names = {cte.alias_or_name.lower() for cte in statement.find_all(exp.CTE)} + referenced = set() + for table in statement.find_all(exp.Table): + if table.name.lower() in cte_names: + continue + qualified = f"{table.db}.{table.name}".lower() if table.db else table.name.lower() + referenced.add(qualified) + unknown = sorted(table for table in referenced if table not in allowed_tables) + if unknown: + raise UnsafeSQLError(f"SQL references tables outside the selected schema: {', '.join(unknown)}") limit = max_rows or settings.MAX_QUERY_ROWS existing = statement.args.get("limit") diff --git a/apps/api/app/sql_validation/__init__.py b/apps/api/app/sql_validation/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/apps/api/migrations/versions/0001_initial.py b/apps/api/migrations/versions/0001_initial.py index 3a825f3..e483cb1 100644 --- a/apps/api/migrations/versions/0001_initial.py +++ b/apps/api/migrations/versions/0001_initial.py @@ -8,18 +8,17 @@ def upgrade(): - op.execute("CREATE EXTENSION IF NOT EXISTS vector") op.execute(""" CREATE TABLE golden_records ( id BIGSERIAL PRIMARY KEY, connection_key VARCHAR(128) NOT NULL, - question TEXT NOT NULL, sql_query TEXT NOT NULL, embedding vector, + question TEXT NOT NULL, sql_query TEXT NOT NULL, embedding JSONB, reviewer VARCHAR(128), source VARCHAR(128), created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), CONSTRAINT uq_verified_example UNIQUE (connection_key, question, sql_query) ); CREATE TABLE schema_embeddings ( id BIGSERIAL PRIMARY KEY, connection_key VARCHAR(128) NOT NULL, - table_name VARCHAR(128) NOT NULL, schema_ddl TEXT NOT NULL, embedding vector, + table_name VARCHAR(128) NOT NULL, schema_ddl TEXT NOT NULL, embedding JSONB, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), CONSTRAINT uq_schema_embedding UNIQUE (connection_key, table_name) ); diff --git a/apps/api/migrations/versions/0002_byod_connections.py b/apps/api/migrations/versions/0002_byod_connections.py new file mode 100644 index 0000000..e5160c9 --- /dev/null +++ b/apps/api/migrations/versions/0002_byod_connections.py @@ -0,0 +1,61 @@ +"""Saved BYOD connections, catalog snapshots, drafts, history ownership, and audit.""" +from alembic import op + +revision = "0002_byod_connections" +down_revision = "0001_initial" +branch_labels = None +depends_on = None + + +def upgrade(): + op.execute(""" + ALTER TABLE golden_records ALTER COLUMN embedding TYPE JSONB USING embedding::text::jsonb; + ALTER TABLE schema_embeddings ALTER COLUMN embedding TYPE JSONB USING embedding::text::jsonb; + CREATE TABLE database_connections ( + id VARCHAR(64) PRIMARY KEY, user_id VARCHAR(128) NOT NULL, name VARCHAR(128) NOT NULL, + database_type VARCHAR(32) NOT NULL, encrypted_connection_data TEXT NOT NULL, + masked_host VARCHAR(255) NOT NULL, database_name VARCHAR(128) NOT NULL, + username VARCHAR(128) NOT NULL, ssl_mode VARCHAR(32) NOT NULL, status VARCHAR(32) NOT NULL, + last_connected_at TIMESTAMPTZ, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ); + CREATE INDEX ix_database_connections_user_id ON database_connections(user_id); + CREATE TABLE schema_snapshots ( + id VARCHAR(64) PRIMARY KEY, connection_id VARCHAR(64) NOT NULL REFERENCES database_connections(id) ON DELETE CASCADE, + schema_name VARCHAR(128) NOT NULL, normalized_metadata JSONB NOT NULL, schema_hash VARCHAR(64) NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + CONSTRAINT uq_connection_schema_hash UNIQUE(connection_id, schema_hash) + ); + CREATE INDEX ix_schema_snapshots_connection_id ON schema_snapshots(connection_id); + CREATE TABLE query_drafts ( + id VARCHAR(64) PRIMARY KEY, user_id VARCHAR(128) NOT NULL, + connection_id VARCHAR(64) NOT NULL REFERENCES database_connections(id) ON DELETE CASCADE, + question TEXT NOT NULL, sql TEXT NOT NULL, explanation TEXT NOT NULL, + assumptions JSONB NOT NULL DEFAULT '[]', warnings JSONB NOT NULL DEFAULT '[]', confidence VARCHAR(16) NOT NULL, + schema_hash VARCHAR(64) NOT NULL, created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), expires_at TIMESTAMPTZ NOT NULL + ); + CREATE INDEX ix_query_drafts_user_id ON query_drafts(user_id); + CREATE INDEX ix_query_drafts_connection_id ON query_drafts(connection_id); + CREATE TABLE audit_logs ( + id VARCHAR(64) PRIMARY KEY, user_id VARCHAR(128) NOT NULL, connection_id VARCHAR(64), + action VARCHAR(64) NOT NULL, metadata JSONB NOT NULL DEFAULT '{}', created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() + ); + CREATE INDEX ix_audit_logs_user_id ON audit_logs(user_id); + CREATE INDEX ix_audit_logs_connection_id ON audit_logs(connection_id); + ALTER TABLE query_history ADD COLUMN IF NOT EXISTS connection_id VARCHAR(64) REFERENCES database_connections(id) ON DELETE CASCADE; + ALTER TABLE query_history ADD COLUMN IF NOT EXISTS user_id VARCHAR(128); + ALTER TABLE query_history ADD COLUMN IF NOT EXISTS execution_status VARCHAR(32); + ALTER TABLE query_history ADD COLUMN IF NOT EXISTS warnings JSONB NOT NULL DEFAULT '[]'; + ALTER TABLE query_history ADD COLUMN IF NOT EXISTS error_code VARCHAR(64); + """) + + +def downgrade(): + op.execute(""" + ALTER TABLE query_history DROP COLUMN IF EXISTS error_code; + ALTER TABLE query_history DROP COLUMN IF EXISTS warnings; + ALTER TABLE query_history DROP COLUMN IF EXISTS execution_status; + ALTER TABLE query_history DROP COLUMN IF EXISTS user_id; + ALTER TABLE query_history DROP COLUMN IF EXISTS connection_id; + DROP TABLE IF EXISTS audit_logs, query_drafts, schema_snapshots, database_connections; + """) diff --git a/apps/api/requirements.txt b/apps/api/requirements.txt index 0530bee..c486a6e 100644 --- a/apps/api/requirements.txt +++ b/apps/api/requirements.txt @@ -1,8 +1,8 @@ alembic==1.16.5 +cryptography==45.0.7 fastapi==0.135.1 httpx==0.28.1 openai==2.29.0 -pgvector==0.4.1 psycopg2-binary==2.9.11 pydantic-settings==2.10.1 pytest==8.4.2 diff --git a/apps/api/tests/test_api_errors.py b/apps/api/tests/test_api_errors.py index a8af5b1..d08344f 100644 --- a/apps/api/tests/test_api_errors.py +++ b/apps/api/tests/test_api_errors.py @@ -1,8 +1,3 @@ -from unittest.mock import AsyncMock, patch - -from app.core.exceptions import ProviderError, UnsafeSQLError - - def test_question_validation(client): response = client.post("/api/v1/query", json={"question": " "}) assert response.status_code == 422 @@ -12,16 +7,9 @@ def test_external_connections_disabled(client): response = client.post("/api/v1/connect-and-query", json={"question": "count rows", "db_config": { "host": "example.test", "port": 5432, "database": "demo", "user": "reader", "password": "secret" }}) - assert response.status_code == 403 + assert response.status_code == 404 -def test_provider_error_status(client): - with patch("app.api.v1.endpoints.query.execute_query", new=AsyncMock(side_effect=ProviderError("provider unavailable"))): - response = client.post("/api/v1/query", json={"question": "Show orders"}) - assert response.status_code == 502 - - -def test_unsafe_sql_status(client): - with patch("app.api.v1.endpoints.query.execute_query", new=AsyncMock(side_effect=UnsafeSQLError("rejected"))): - response = client.post("/api/v1/query", json={"question": "Show orders"}) - assert response.status_code == 422 +def test_legacy_combined_query_is_disabled(client): + response = client.post("/api/v1/query", json={"question": "Show orders"}) + assert response.status_code == 403 diff --git a/apps/api/tests/test_auth.py b/apps/api/tests/test_auth.py new file mode 100644 index 0000000..a898697 --- /dev/null +++ b/apps/api/tests/test_auth.py @@ -0,0 +1,12 @@ +import pytest + +from app.auth.dependencies import get_current_user_id, issue_session +from app.core.config import settings +from app.core.exceptions import AuthenticationError + + +def test_signed_session_subject_cannot_be_forged(monkeypatch): + monkeypatch.setattr(settings,"AUTH_SIGNING_KEY","test-signing-key-not-for-production") + token,user_id,_=issue_session() + assert get_current_user_id(f"Bearer {token}")==user_id + with pytest.raises(AuthenticationError): get_current_user_id(f"Bearer {token}tampered") diff --git a/apps/api/tests/test_byod_services.py b/apps/api/tests/test_byod_services.py new file mode 100644 index 0000000..79c33c2 --- /dev/null +++ b/apps/api/tests/test_byod_services.py @@ -0,0 +1,76 @@ +from types import SimpleNamespace +from unittest.mock import AsyncMock, patch + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import Session + +from app.catalog.service import introspect_postgresql, qualified_tables, schema_hash +from app.connections.service import masked, owned_connection +from app.core.config import settings +from app.core.exceptions import NotFoundError, UnsafeSQLError +from app.db.base import Base +from app.models.byod import DatabaseConnection, SchemaSnapshot +from app.query_planner.service import generate_draft +from app.schemas.llm_schema import SQLGeneration +from app.services.validation_service import validate_sql + + +METADATA = {"database_type":"postgresql","schemas":[{"name":"public","objects":[ + {"name":"orders","type":"table","columns":[{"name":"id","type":"INTEGER","nullable":False}], + "primary_key":["id"],"foreign_keys":[],"indexes":[]} +]}]} + + +@pytest.fixture +def byod_db(): + engine=create_engine("sqlite+pysqlite:///:memory:"); Base.metadata.create_all(engine) + with Session(engine) as session: yield session + + +def add_connection(db, user="usr_owner"): + value=DatabaseConnection(user_id=user,name="Analytics",encrypted_connection_data="ciphertext",masked_host="db.example.com", + database_name="analytics",username="reader",ssl_mode="require",status="connected") + db.add(value);db.flush();return value + + +def test_masked_output_never_contains_encrypted_credentials(byod_db): + value=add_connection(byod_db); output=masked(value) + assert "encrypted_connection_data" not in output and "password" not in output and "ciphertext" not in str(output) + + +def test_ownership_enforced(byod_db): + value=add_connection(byod_db) + assert owned_connection(byod_db,value.id,"usr_owner") is value + with pytest.raises(NotFoundError): owned_connection(byod_db,value.id,"usr_attacker") + + +def test_schema_hash_is_stable_and_tables_normalized(): + assert schema_hash(METADATA)==schema_hash({"schemas":METADATA["schemas"],"database_type":"postgresql"}) + assert qualified_tables(METADATA)=={"orders","public.orders"} + + +def test_schema_introspection_normalization(): + inspector=SimpleNamespace( + get_schema_names=lambda:["public","pg_catalog"], get_table_names=lambda schema:["orders"], + get_view_names=lambda schema:[], get_columns=lambda name,schema:[{"name":"id","type":"INTEGER","nullable":False}], + get_pk_constraint=lambda name,schema:{"constrained_columns":["id"]}, get_foreign_keys=lambda name,schema:[], + get_indexes=lambda name,schema:[{"name":"ix_orders_id","column_names":["id"],"unique":True}]) + with patch("app.catalog.service.inspect",return_value=inspector): metadata=introspect_postgresql(object()) + assert metadata["schemas"][0]["objects"][0]["primary_key"]==["id"] + + +def test_allowed_table_validation(): + assert "LIMIT" in validate_sql("SELECT * FROM public.orders",allowed_tables={"orders","public.orders"}) + with pytest.raises(UnsafeSQLError): validate_sql("SELECT * FROM secrets",allowed_tables={"orders","public.orders"}) + + +@pytest.mark.asyncio +async def test_generation_creates_draft_without_database_execution(byod_db, monkeypatch): + monkeypatch.setattr(settings,"DATABASE_MAX_RESULT_ROWS",50) + connection=add_connection(byod_db); snapshot=SchemaSnapshot(connection_id=connection.id,normalized_metadata=METADATA,schema_hash=schema_hash(METADATA));byod_db.add(snapshot);byod_db.commit() + generated=SQLGeneration(sql="SELECT id FROM orders",explanation="Reads orders",tables_used=["orders"],assumptions=[],confidence=.9) + with patch("app.query_planner.service.generate_sql",new=AsyncMock(return_value=generated)), \ + patch("app.connections.service.external_engine") as external: + draft=await generate_draft(byod_db,"usr_owner",connection.id,"Show orders") + assert draft.sql.endswith("LIMIT 50") and external.call_count==0 diff --git a/apps/api/tests/test_connection_security.py b/apps/api/tests/test_connection_security.py new file mode 100644 index 0000000..a7106f5 --- /dev/null +++ b/apps/api/tests/test_connection_security.py @@ -0,0 +1,54 @@ +import socket +from unittest.mock import patch + +import pytest +from cryptography.fernet import Fernet + +from app.connections.schemas import ConnectionCreate +from app.connections.security import (_fernet, decrypt_config, encrypt_config, normalize_connection_input, + validate_public_host) +from app.core.config import settings +from app.core.exceptions import ConnectionSecurityError, DependencyError + + +def test_structured_connection_input(): + request = ConnectionCreate(name="Analytics", host="db.example.com", database="analytics", + username="reader", password="secret", ssl_mode="require") + config = normalize_connection_input(request) + assert config == {"host":"db.example.com","port":5432,"database":"analytics","username":"reader","password":"secret","ssl_mode":"require"} + + +def test_postgresql_connection_string_parsing(): + request = ConnectionCreate(name="Analytics", connection_string="postgresql://reader:secret@db.example.com:5432/analytics?sslmode=verify-full") + config = normalize_connection_input(request) + assert config["host"] == "db.example.com" and config["ssl_mode"] == "verify-full" + + +@pytest.mark.parametrize("url", ["mysql://u:p@db.example.com/db", "sqlite:///tmp/db"]) +def test_non_postgresql_scheme_rejected(url): + with pytest.raises(ConnectionSecurityError): normalize_connection_input(ConnectionCreate(name="x", connection_string=url)) + + +@pytest.mark.parametrize("host", ["localhost", "127.0.0.1", "10.1.2.3", "172.16.1.2", "192.168.1.5", "169.254.169.254", "::1", "fc00::1"]) +def test_private_local_and_link_local_hosts_rejected(host, monkeypatch): + monkeypatch.setattr(settings, "ALLOW_PRIVATE_DATABASE_HOSTS", False) + with pytest.raises(ConnectionSecurityError): validate_public_host(host) + + +def test_public_dns_host_accepted(monkeypatch): + monkeypatch.setattr(settings, "ALLOW_PRIVATE_DATABASE_HOSTS", False) + answer = [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("203.0.113.10", 5432))] + # TEST-NET is reserved and correctly blocked; use a real globally routable-shaped address. + answer[0] = (socket.AF_INET, socket.SOCK_STREAM, 6, "", ("8.8.8.8", 5432)) + with patch("app.connections.security.socket.getaddrinfo", return_value=answer): + assert validate_public_host("db.example.com") == ["8.8.8.8"] + + +def test_encryption_round_trip_and_wrong_key_failure(monkeypatch): + first = Fernet.generate_key().decode(); second = Fernet.generate_key().decode() + monkeypatch.setattr(settings, "CONNECTION_ENCRYPTION_KEY", first); _fernet.cache_clear() + encrypted = encrypt_config({"host":"db.example.com","password":"secret"}) + assert "secret" not in encrypted and decrypt_config(encrypted)["password"] == "secret" + monkeypatch.setattr(settings, "CONNECTION_ENCRYPTION_KEY", second); _fernet.cache_clear() + with pytest.raises(DependencyError): decrypt_config(encrypted) + _fernet.cache_clear() diff --git a/apps/api/tests/test_execution.py b/apps/api/tests/test_execution.py new file mode 100644 index 0000000..6f6b186 --- /dev/null +++ b/apps/api/tests/test_execution.py @@ -0,0 +1,70 @@ +from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import Session + +from app.core.config import settings +from app.core.exceptions import DependencyError, NotFoundError +from app.execution.service import execute_draft +from app.db.base import Base +from app.models.byod import QueryDraft, SchemaSnapshot +from app.query_planner.service import owned_draft +from test_byod_services import METADATA, add_connection, schema_hash + + +@pytest.fixture +def byod_db(): + engine=create_engine("sqlite+pysqlite:///:memory:"); Base.metadata.create_all(engine) + with Session(engine) as session: yield session + + +def add_draft(db, connection, user="usr_owner"): + draft=QueryDraft(user_id=user,connection_id=connection.id,question="Show orders",sql="SELECT id FROM orders LIMIT 500", + explanation="Reads orders",assumptions=[],warnings=[],confidence="0.9",schema_hash=schema_hash(METADATA), + expires_at=datetime.now(timezone.utc)+timedelta(minutes=5)) + db.add(draft);db.add(SchemaSnapshot(connection_id=connection.id,normalized_metadata=METADATA,schema_hash=schema_hash(METADATA)));db.commit();return draft + + +def test_execution_requires_owned_valid_draft(byod_db): + connection=add_connection(byod_db);draft=add_draft(byod_db,connection) + with pytest.raises(NotFoundError): owned_draft(byod_db,draft.id,"usr_attacker") + + +def test_execution_timeout_is_safely_wrapped(byod_db): + connection=add_connection(byod_db);draft=add_draft(byod_db,connection) + class Engine: + def connect(self): raise TimeoutError("connect timeout") + def dispose(self): pass + with patch("app.execution.service.decrypt_config",return_value={}),patch("app.execution.service.external_engine",return_value=Engine()),pytest.raises(DependencyError,match="failed or timed out"): + execute_draft(byod_db,"usr_owner",draft.id) + + +def test_execution_revalidates_and_caps_rows(byod_db,monkeypatch): + monkeypatch.setattr(settings,"DATABASE_MAX_RESULT_ROWS",2) + connection=add_connection(byod_db);draft=add_draft(byod_db,connection) + class Result: + def keys(self): return ["id"] + def fetchmany(self,size): assert size==2; return [SimpleNamespace(_mapping={"id":1}),SimpleNamespace(_mapping={"id":2})] + class Tx: + def rollback(self): pass + class Conn: + def __enter__(self): return self + def __exit__(self,*args): pass + def begin(self): return Tx() + def execute(self,statement,params=None): + if str(statement).startswith("SELECT id"): return Result() + return SimpleNamespace() + class Engine: + def connect(self): return Conn() + def dispose(self): pass + original=byod_db.execute + def execute(statement,params=None,*args,**kwargs): + if "INSERT INTO query_history" in str(statement): return SimpleNamespace(scalar_one=lambda:7) + return original(statement,params,*args,**kwargs) + byod_db.execute=execute + with patch("app.execution.service.decrypt_config",return_value={}),patch("app.execution.service.external_engine",return_value=Engine()): + output=execute_draft(byod_db,"usr_owner",draft.id) + assert output["row_count"]==2 and output["sql"].endswith("LIMIT 2") diff --git a/apps/api/tests/test_validation.py b/apps/api/tests/test_validation.py index 9f16e65..91377dd 100644 --- a/apps/api/tests/test_validation.py +++ b/apps/api/tests/test_validation.py @@ -33,3 +33,8 @@ def test_rejects_write_or_admin(sql): def test_rejects_multiple_statements_or_comments(sql): with pytest.raises(UnsafeSQLError): validate_sql(sql) + + +def test_rejects_unknown_table_against_catalog(): + with pytest.raises(UnsafeSQLError): + validate_sql("SELECT * FROM payroll", allowed_tables={"orders", "public.orders"}) diff --git a/apps/web/README.md b/apps/web/README.md index 4499bff..b4c0b6c 100644 --- a/apps/web/README.md +++ b/apps/web/README.md @@ -4,7 +4,7 @@ Next.js 15, React 19, TypeScript, and Tailwind UI for QueryMindAI. ## Architecture and implemented pages -`src/app` contains routes and retained page components, `src/lib/api.ts` is the typed fetch boundary, and `src/types` contains contracts. The AI Query Assistant is API-backed and implemented: question entry, loading/error states, Generated SQL, rows, explanation, warnings, and verified-example approval. Schema/history backend APIs exist; the retained Browser, History, Dashboard, Connections, and login designs contain demo data or coming-soon interactions and must not be treated as operational telemetry/authentication. +`src/app` contains routes, `src/lib/api.ts` is the typed authenticated fetch boundary, and `src/types` contains contracts. Data Connections creates/tests/deletes encrypted PostgreSQL connections; Schema Explorer renders normalized live catalog snapshots; AI Query Workspace separates generation from explicit execution; Query History lists approved executions without stored rows. The client stores a signed anonymous workspace token in browser local storage; this is ownership isolation, not account login. ## Configuration and API integration diff --git a/apps/web/src/app/dashboard-and-connection/components/DashboardContent.tsx b/apps/web/src/app/dashboard-and-connection/components/DashboardContent.tsx index 6daba67..5b3dd01 100644 --- a/apps/web/src/app/dashboard-and-connection/components/DashboardContent.tsx +++ b/apps/web/src/app/dashboard-and-connection/components/DashboardContent.tsx @@ -1,413 +1,28 @@ 'use client'; -import React, { useState } from 'react'; -import { Search, TrendingUp, Zap, BarChart2, Plus, CheckCircle, AlertCircle, Sparkles, ArrowRight, } from 'lucide-react'; - -// ─── Data ──────────────────────────────────────────────────────────────────── - -const metricCards = [ - { - id: 'total-queries', - label: 'Total Queries', - value: '14,208', - sub: '12% increase this week', - subIcon: TrendingUp, - subColor: 'text-green-500', - bg: 'bg-blue-600', - textColor: 'text-white', - subTextColor: 'text-blue-200', - valueColor: 'text-white', - labelColor: 'text-blue-200', - wide: true, - }, - { - id: 'successful-conversions', - label: 'Successful Conversions', - value: '8,590', - sub: 'AI Generated', - subIcon: Zap, - subColor: 'text-red-500', - bg: 'bg-white', - textColor: 'text-gray-900', - subTextColor: 'text-red-500', - valueColor: 'text-gray-900', - labelColor: 'text-gray-500', - wide: false, - }, - { - id: 'avg-latency', - label: 'Average Latency', - value: '24ms', - sub: 'Real-time monitoring active', - subIcon: BarChart2, - subColor: 'text-blue-500', - bg: 'bg-blue-50', - textColor: 'text-gray-900', - subTextColor: 'text-gray-500', - valueColor: 'text-gray-900', - labelColor: 'text-gray-500', - wide: false, - }, -]; - -const databases = [ - { - id: 'pg-prod', - name: 'PostgreSQL Production', - host: 'db-prod-cluster-01.aws-east.com', - status: 'LIVE', - statusColor: 'text-red-500', - statusBg: 'bg-red-50', - statusDot: 'bg-red-500', - meta: '342 TABLES', - avatars: ['AC', 'SR', '+4'], - icon: '≡', - iconBg: 'bg-blue-100', - iconColor: 'text-blue-600', - }, - { - id: 'marketing-lake', - name: 'Marketing Data Lake', - host: 'bq-analytics-marketing-v2', - status: 'IDLE', - statusColor: 'text-gray-500', - statusBg: 'bg-gray-100', - statusDot: 'bg-gray-400', - meta: '2.4 TB STORED', - avatars: ['PM'], - icon: '▦', - iconBg: 'bg-blue-100', - iconColor: 'text-blue-600', - }, -]; - -type QueryStatus = 'success' | 'ai' | 'error'; - -interface QueryRow { - id: string; - status: QueryStatus; - snippet: string; - meta: string; - connection: string; - connectionDot: string; - execTime: string | null; - execHighlight?: boolean; -} - -const queryRows: QueryRow[] = [ - { - id: 'q1', - status: 'success', - snippet: 'SELECT user_id, count(*) FROM transactions...', - meta: '2 minutes ago by Alex Chen', - connection: 'PostgreSQL Production', - connectionDot: 'bg-blue-500', - execTime: null, - execHighlight: true, - }, - { - id: 'q2', - status: 'ai', - snippet: '"Find all customers who bought item X in last 30 days"', - meta: 'AI GENERATED · 15 MINUTES AGO', - connection: 'PostgreSQL Production', - connectionDot: 'bg-blue-500', - execTime: '0.8s', - }, - { - id: 'q3', - status: 'error', - snippet: "DELETE FROM logs WHERE created_at < '2020-01-01'", - meta: 'Permission Denied · 1 hour ago', - connection: 'Marketing Data Lake', - connectionDot: 'bg-yellow-500', - execTime: '--', - }, - { - id: 'q4', - status: 'success', - snippet: 'SELECT sum(revenue) FROM monthly_reports...', - meta: '3 hours ago by Sarah Smith', - connection: 'Marketing Data Lake', - connectionDot: 'bg-yellow-500', - execTime: '4.5s', - }, -]; - -// ─── Sub-components ─────────────────────────────────────────────────────────── - -function StatusIcon({ status }: { status: QueryStatus }) { - if (status === 'success') { - return ( -
- -
- ); - } - if (status === 'ai') { - return ( -
- -
- ); - } - return ( -
- -
- ); -} - -// ─── Main Component ─────────────────────────────────────────────────────────── +import { FormEvent, useCallback, useEffect, useState } from 'react'; +import { Database, Plus, RefreshCw, ShieldCheck, Trash2 } from 'lucide-react'; +import { api, ApiError } from '@/lib/api'; +import type { SavedConnection } from '@/types/api'; export default function DashboardContent() { - const [activeTab, setActiveTab] = useState<'connections' | 'documentation'>('connections'); - const [queryFilter, setQueryFilter] = useState<'All' | 'AI Assisted' | 'Errors'>('All'); - - const filteredQueries = queryRows.filter((q) => { - if (queryFilter === 'All') return true; - if (queryFilter === 'AI Assisted') return q.status === 'ai'; - if (queryFilter === 'Errors') return q.status === 'error'; - return true; - }); - - return ( -
- {/* Top Header Bar */} -
-
- Query.AI -
- - Search queries, tables, or connections... -
-
-
- - -
-
- - {/* Main Content */} -
- {/* Workspace Overview Header */} -
-
-

Workspace Overview

-

- Manage your intelligent data infrastructure and monitor AI-assisted query performance. -

-
-
-
-

Active Connections

-

12

-
-
-

AI Accuracy

-

98.4%

-
-
-
- - {/* Metric Cards */} -
- {metricCards.map((card) => ( -
- {/* Background decoration for blue card */} - {card.id === 'total-queries' && ( -
-
-
-
- )} -

{card.label}

-

{card.value}

-
- - {card.sub} -
-
- ))} -
- - {/* Active Databases */} -
-
-

Active Databases

- -
-
- {databases.map((db) => ( -
-
-
- {db.icon} -
- - - {db.status} - -
-

{db.name}

-

{db.host}

-
- {db.meta} -
- {db.avatars.map((av, i) => ( -
- {av} -
- ))} -
-
-
- ))} - - {/* Add New Node card */} -
-
- -
-

Add New Node

-

Connect MySQL, Redshift, or Snowflake clusters

-
-
-
- - {/* Recent Query Intelligence */} -
-
-

Recent Query Intelligence

-
- {(['All', 'AI Assisted', 'Errors'] as const).map((f) => ( - - ))} -
-
- -
- {/* Table Header */} -
-
Status
-
Query Snippet
-
Connection
-
Execution Time
-
- - {/* Rows */} - {filteredQueries.map((row, idx) => ( -
- {/* Status */} -
- -
- - {/* Query Snippet */} -
-

- {row.snippet} -

-

- {row.meta} -

-
- - {/* Connection */} -
- - {row.connection} -
- - {/* Execution Time */} -
- {row.execHighlight ? ( -
- -
- ) : ( - {row.execTime} - )} -
-
- ))} - - {/* Show More */} -
- -
-
-
-
-
- ); -} \ No newline at end of file + const [items,setItems]=useState([]); const [loading,setLoading]=useState(true); + const [error,setError]=useState(''); const [mode,setMode]=useState<'fields'|'url'>('fields'); const [saving,setSaving]=useState(false); + const load=useCallback(async()=>{setLoading(true);try{setItems(await api.listConnections());setError('')}catch(e){setError(e instanceof ApiError?e.message:'Could not load connections')}finally{setLoading(false)}},[]); + useEffect(()=>{void load()},[load]); + async function submit(event:FormEvent){event.preventDefault();setSaving(true);setError('');const data=new FormData(event.currentTarget); + const body:Record={name:data.get('name'),ssl_mode:data.get('ssl_mode')}; + if(mode==='url')body.connection_string=data.get('connection_string'); else Object.assign(body,{host:data.get('host'),port:Number(data.get('port')),database:data.get('database'),username:data.get('username'),password:data.get('password')}); + try{await api.createConnection(body);event.currentTarget.reset();await load()}catch(e){setError(e instanceof ApiError?e.message:'Connection failed')}finally{setSaving(false)}} + async function remove(id:string){if(!confirm('Delete this connection and its schema, drafts, and query history?'))return;await api.deleteConnection(id);await load()} + return
+

Data Connections

Connect a publicly reachable PostgreSQL database with SSL and dedicated read-only credentials.

+
QueryMindAI does not need write access. Credentials are encrypted in the API; connection details and row data are never sent to the LLM.
+ {error&&
{error}
} +

Add PostgreSQL connection

+
+ {mode==='url'?:<>} +
+
+

Saved connections

{loading?

Loading connections…

:items.length===0?

No saved connections yet.

:items.map(item=>

{item.name}

{item.username}@{item.host}/{item.database} · {item.ssl_mode}

{item.status}

)}
+} diff --git a/apps/web/src/app/database-browser/components/DatabaseBrowserContent.tsx b/apps/web/src/app/database-browser/components/DatabaseBrowserContent.tsx index 5e43354..05a363d 100644 --- a/apps/web/src/app/database-browser/components/DatabaseBrowserContent.tsx +++ b/apps/web/src/app/database-browser/components/DatabaseBrowserContent.tsx @@ -1,495 +1,13 @@ 'use client'; -import React, { useState } from 'react'; -import { RefreshCw, Download, Play, ChevronRight, Database, Key, Mail, DollarSign, Grid3X3, Code2, SlidersHorizontal, ChevronLeft, Sparkles, } from 'lucide-react'; - -interface SchemaItem { - name: string; - type: 'table' | 'view'; - selected?: boolean; -} - -interface Schema { - name: string; - expanded: boolean; - items?: SchemaItem[]; -} - -const initialSchemas: Schema[] = [ - { - name: 'production_db', - expanded: true, - items: [ - { name: 'orders_v2', type: 'table', selected: true }, - { name: 'customer_profiles', type: 'table' }, - { name: 'revenue_daily_view', type: 'view' }, - { name: 'inventory_tracking', type: 'table' }, - ], - }, - { name: 'staging_analytics', expanded: false }, - { name: 'raw_events', expanded: false }, -]; - -const columnData = [ - { - name: 'order_id', - type: 'UUID', - badge: 'PK', - badgeColor: 'bg-gray-100 text-gray-600', - icon: , - detail: 'Unique constraint enforced', - fill: 100, - fillLabel: '100% Non-null', - }, - { - name: 'customer_email', - type: 'VARCHAR(255)', - badge: null, - icon: , - detail: 'Email pattern validation', - fill: 98.2, - fillLabel: '98.2% Filled', - }, - { - name: 'amount_usd', - type: 'DECIMAL(12,2)', - badge: null, - icon: , - detail: 'Range: $0.00 - $14k', - fill: 100, - fillLabel: '100% Non-null', - }, -]; - -const previewRows = [ - { - id: 1, - orderId: '550e8400-e29...', - createdAt: '2023-11-21 09:12:04', - customerId: 'USR_8829', - status: 'COMPLETED', - amountUsd: '$129.40', - }, - { - id: 2, - orderId: '7e21124a-f38...', - createdAt: '2023-11-21 09:15:22', - customerId: 'USR_1042', - status: 'SHIPPED', - amountUsd: '$2,450.00', - }, - { - id: 3, - orderId: 'bf12c40c-372...', - createdAt: '2023-11-21 10:01:05', - customerId: 'USR_9001', - status: 'PENDING', - amountUsd: '$89.12', - }, - { - id: 4, - orderId: 'a1288c12-f38...', - createdAt: '2023-11-21 10:04:44', - customerId: 'USR_4211', - status: 'COMPLETED', - amountUsd: '$552.00', - }, - { - id: 5, - orderId: 'd4422e1a-c71...', - createdAt: '2023-11-21 10:12:11', - customerId: 'USR_3321', - status: 'CANCELLED', - amountUsd: '$34.50', - }, -]; - -const statusConfig: Record = { - COMPLETED: { bg: 'bg-green-100', text: 'text-green-700' }, - SHIPPED: { bg: 'bg-blue-100', text: 'text-blue-700' }, - PENDING: { bg: 'bg-yellow-100', text: 'text-yellow-700' }, - CANCELLED: { bg: 'bg-red-100', text: 'text-red-700' }, -}; - -const tabs = ['Columns & Types', 'Sample Data', 'Dependencies', 'Permissions']; - -export default function DatabaseBrowserContent() { - const [schemas, setSchemas] = useState(initialSchemas); - const [activeTab, setActiveTab] = useState('Columns & Types'); - const [previewMode, setPreviewMode] = useState<'Grid' | 'JSON'>('Grid'); - const [promptValue, setPromptValue] = useState(''); - - const toggleSchema = (name: string) => { - setSchemas((prev) => - prev.map((s) => (s.name === name ? { ...s, expanded: !s.expanded } : s)) - ); - }; - - return ( -
- {/* Top Header */} -
- {/* Top Bar */} -
-
- Query.AI - - V4.2 - -
-
-
- - - - Search tables, schemas... -
- -
- - -
- -
-
- - {/* Body */} -
- {/* Schemas Panel */} -
-
- Schemas - -
- -
- {schemas.map((schema) => ( -
- - - {schema.expanded && schema.items && ( -
- {schema.items.map((item) => ( - - ))} -
- )} -
- ))} -
-
- - {/* Main Panel */} -
- {/* Breadcrumb */} -
- production_db - - orders_v2 -
- - {/* Title Row */} -
-
-

orders_v2

-
- - - PostgreSQL 14.2 - - - - - - Updated 4m ago - - - - - - 1.4 GB - -
-
-
- - -
-
- - {/* AI Prompt Section */} -
-
-
- -
-
-

Prompt to Query this Table

-

Let Query.AI write the SQL for you based on natural language.

-
-
-
- setPromptValue(e.target.value)} - placeholder="e.g., Show me total revenue by month for the last year..." - className="flex-1 px-4 py-2.5 rounded-lg border text-sm text-gray-700 placeholder-gray-400 focus:outline-none focus:ring-2 focus:ring-blue-500" - style={{ borderColor: '#d1d5db', background: '#ffffff' }} - /> - -
-
- - {/* Tabs */} -
- {tabs.map((tab) => ( - - ))} -
- - {/* Column Cards */} - {activeTab === 'Columns & Types' && ( - <> -
- {columnData.map((col) => ( -
-
- {col.name} - {col.badge && ( - - {col.badge} - - )} -
-

{col.type}

-
- {col.icon} - {col.detail} -
-
-
-
-

{col.fillLabel}

-
- ))} -
- - {/* Preview Table */} -
-
-

- Preview (Top 50 Rows) -

-
-
- - -
- -
-
- -
- - - - - - - - - - - - - {previewRows.map((row) => ( - - - - - - - - - ))} - -
#order_idcreated_atcustomer_idstatusamount_usd
{row.id} - - {row.orderId} - - {row.createdAt}{row.customerId} - - {row.status} - - - {row.amountUsd} -
- - {/* Pagination */} -
- Showing 5 of 8,241 results -
- - Page 1 of 1,648 - -
-
-
-
- - )} - - {activeTab !== 'Columns & Types' && ( -
- {activeTab} content coming soon -
- )} -
-
-
-
- ); -} +import { useCallback, useEffect, useState } from 'react'; +import { KeyRound, RefreshCw } from 'lucide-react'; +import { api, ApiError } from '@/lib/api'; +import type { CatalogResponse, SavedConnection } from '@/types/api'; + +export default function DatabaseBrowserContent(){const [connections,setConnections]=useState([]);const [selected,setSelected]=useState('');const [catalog,setCatalog]=useState(null);const [error,setError]=useState('');const [loading,setLoading]=useState(false); + const loadConnections=useCallback(async()=>{try{const list=await api.listConnections();setConnections(list);if(list[0])setSelected(v=>v||list[0].id)}catch(e){setError(e instanceof ApiError?e.message:'Could not load connections')}},[]);useEffect(()=>{void loadConnections()},[loadConnections]); + const load=useCallback(async()=>{if(!selected)return;setLoading(true);try{setCatalog(await api.connectionSchema(selected));setError('')}catch(e){setError(e instanceof ApiError?e.message:'Could not load schema')}finally{setLoading(false)}},[selected]);useEffect(()=>{void load()},[load]); + async function refresh(){if(!selected)return;setLoading(true);try{setCatalog(await api.refreshSchema(selected));setError('')}catch(e){setError(e instanceof ApiError?e.message:'Refresh failed')}finally{setLoading(false)}} + return

Schema Explorer

Live API data from the latest encrypted connection snapshot.

+ {error&&
{error}
}{loading?

Introspecting schema…

:!catalog?

Create and select a connection first.

:
{catalog.metadata.schemas.map(schema=>

{schema.name}

{schema.objects.map(obj=>

{obj.name}

{obj.type}
{obj.columns.map(col=>
{obj.primary_key.includes(col.name)&&}{col.name}{col.type}{col.nullable?' · nullable':''}
)}
{obj.foreign_keys.length>0&&
{obj.foreign_keys.map((fk,i)=>

FK {fk.columns.join(', ')} → {fk.referred_schema}.{fk.referred_table}({fk.referred_columns.join(', ')})

)}
}
)}
)}
}
} diff --git a/apps/web/src/app/query-history/components/QueryLibraryContent.tsx b/apps/web/src/app/query-history/components/QueryLibraryContent.tsx index 46495dc..9f97afc 100644 --- a/apps/web/src/app/query-history/components/QueryLibraryContent.tsx +++ b/apps/web/src/app/query-history/components/QueryLibraryContent.tsx @@ -1,333 +1,6 @@ 'use client'; -import React, { useState } from 'react'; -import { - Search, - Bell, - Cloud, - Database, - Calendar, - Tag, - ChevronDown, - Clock, - BarChart2, - Bookmark, - Share2, - Play, - ExternalLink, - AlertTriangle, -} from 'lucide-react'; +import { useEffect, useState } from 'react'; +import { api, ApiError } from '@/lib/api'; -interface QueryCard { - id: string; - timestamp: string; - prompt: string; - sql: string; - tags: string[]; - executionTime?: string; - dataScanned?: string; - status: 'success' | 'failed'; - bookmarked?: boolean; - shared?: boolean; - showRunAgain?: boolean; -} - -const queryCards: QueryCard[] = [ - { - id: '1', - timestamp: '2 mins ago', - prompt: '"Find the top 10 customers by revenue growth in Q3 2023 comparing to Q3 2022, join with regional office locations."', - sql: `SELECT c.name, SUM(o.total_price) as growth_revenue, r.office_name -FROM customers c -JOIN orders o ON c.id = o.customer_id -JOIN regions r ON c.region_id = r.id -WHERE o.order_date BETWEEN '2023-07-01' AND '2023-09-30' -GROUP BY 1, 3 -ORDER BY 2 DESC LIMIT 10;`, - tags: ['#REPORTING', '#QUARTERLY-SYNC'], - executionTime: '1.2s execution', - dataScanned: '4.2 MB scanned', - status: 'success', - showRunAgain: true, - }, - { - id: '2', - timestamp: '4 hours ago', - prompt: '"Debug why the latest migration failed for user_profiles table. Show me the last 5 transactions."', - sql: `SELECT * FROM pg_stat_activity -WHERE query ILIKE '%user_profiles%' -ORDER BY query_start DESC -LIMIT 5;`, - tags: ['#DEBUG', '#OPS'], - status: 'failed', - bookmarked: true, - shared: true, - }, - { - id: '3', - timestamp: 'Yesterday', - prompt: '"Analyze average session duration per device type for mobile users in APAC."', - sql: `SELECT device_type, AVG(duration_ms)/1000/60 as avg_mins -FROM telemetry.sessions -WHERE region = 'APAC' AND platform = 'mobile' -GROUP BY 1;`, - tags: ['#ANALYTICS', '#MOBILE'], - status: 'success', - }, -]; - -function SqlBlock({ code }: { code: string }) { - const keywords = ['SELECT', 'FROM', 'JOIN', 'WHERE', 'GROUP BY', 'ORDER BY', 'LIMIT', 'ON', 'AND', 'AS', 'BETWEEN', 'DESC', 'ILIKE']; - - const highlightLine = (line: string) => { - let result = line; - keywords.forEach((kw) => { - result = result.replace( - new RegExp(`\\b${kw}\\b`, 'g'), - `${kw}` - ); - }); - // Highlight strings - result = result.replace(/'([^']*)'/g, `'$1'`); - // Highlight numbers - result = result.replace(/\b(\d+)\b/g, `$1`); - return result; - }; - - return ( -
- {code.split('\n').map((line, i) => ( -
- ))} -
- ); -} - -export default function QueryLibraryContent() { - const [activeTab, setActiveTab] = useState<'all' | 'shared' | 'favorites'>('all'); - - return ( -
- {/* Top Header Bar */} -
- Query.AI -
- - -
-
- - -
-
- - - -
-
- - {/* Main Content */} -
- {/* Page Header + Tabs */} -
-
-

Query Library

-

- Manage, audit, and reuse your team's computational intelligence -
- across all connected clusters. -

-
- {/* Tabs */} -
- {[ - { key: 'all', label: 'All History' }, - { key: 'shared', label: 'Shared' }, - { key: 'favorites', label: 'Favorites' }, - ].map((tab) => ( - - ))} -
-
- - {/* Filter Bar */} -
- } label="Production Cluster" /> - } label="Last 30 Days" /> - } label="Any Tag" /> -
- Sort by: - -
-
- - {/* Query Cards */} -
- {queryCards.map((card) => ( - - ))} -
- - {/* Footer */} -
- Showing 42 historical queries •{' '} - -
-
-
- ); -} - -function FilterDropdown({ icon, label }: { icon: React.ReactNode; label: string }) { - return ( - - ); -} - -function QueryCardItem({ card }: { card: QueryCard }) { - return ( -
-
- {/* Left content */} -
- {/* Prompt header */} -
- - PROMPT - - · - {card.timestamp} - {(card.bookmarked || card.shared) && ( -
- {card.bookmarked && ( - - )} - {card.shared && ( - - )} -
- )} -
- - {/* Prompt text */} -

- {card.prompt} -

- - {/* SQL block */} -
- -
- - {/* Tags + Stats */} -
- {card.tags.map((tag) => ( - - {tag} - - ))} - - {card.status === 'success' && card.executionTime && ( - <> - - - {card.executionTime} - - - - {card.dataScanned} - - - )} - - {card.status === 'failed' && ( - - - Failed to complete - - )} -
-
- - {/* Right action buttons */} -
- - {card.showRunAgain && ( - - )} -
-
-
- ); -} +export default function QueryLibraryContent(){const [items,setItems]=useState>>([]);const [error,setError]=useState('');useEffect(()=>{api.history().then(r=>setItems(r.items)).catch(e=>setError(e instanceof ApiError?e.message:'Could not load history'))},[]);return

Query History

Successful approved executions. Result row data is not stored.

{error&&
{error}
}
{items.length===0?

No approved queries yet.

:items.map(item=>

{String(item.question)}

{String(item.row_count)} rows · {String(item.duration_ms)}ms
{String(item.sql)}

{new Date(String(item.created_at)).toLocaleString()}

)}
} diff --git a/apps/web/src/app/sql-editor-and-ai-assistant/components/SqlEditorContent.tsx b/apps/web/src/app/sql-editor-and-ai-assistant/components/SqlEditorContent.tsx index 08a8ea6..79118f6 100644 --- a/apps/web/src/app/sql-editor-and-ai-assistant/components/SqlEditorContent.tsx +++ b/apps/web/src/app/sql-editor-and-ai-assistant/components/SqlEditorContent.tsx @@ -1,313 +1,19 @@ 'use client'; -import React, { useState } from 'react'; -import { - Sparkles, - Copy, - CheckCircle, - Download, - ChevronLeft, - ChevronRight, - Zap, - ArrowRight, - BrainCircuit -} from 'lucide-react'; +import { useCallback, useEffect, useState } from 'react'; +import { Download, Play, Save, Sparkles } from 'lucide-react'; import { api, ApiError } from '@/lib/api'; - -const keywords = ['SELECT', 'FROM', 'WHERE', 'JOIN', 'LEFT', 'INNER', 'ON', 'GROUP', 'BY', 'ORDER', 'LIMIT', 'AS', 'AND', 'OR', 'NOT', 'IN', 'IS', 'NULL', 'CASE', 'WHEN', 'THEN', 'END', 'DISTINCT', 'COUNT', 'SUM', 'AVG', 'MAX', 'MIN', 'ROUND', 'DATE_TRUNC', 'INTERVAL', 'NULLIF', 'NOW', 'BETWEEN', 'WITH', 'INSERT', 'UPDATE', 'DELETE', 'HAVING', 'UNION', 'ALL', 'CURRENT_DATE']; - -function syntaxHighlight(sql: string): React.ReactNode[] { - const lines = sql.split('\n'); - return lines.map((line, li) => { - const parts = line.split(/(\s+|[(),;])/); - return ( -
- - {li + 1} - - - {parts.map((part, pi) => { - if (keywords.includes(part.toUpperCase())) { - return {part}; - } - if (/^'[^']*'$/.test(part)) { - return {part}; - } - if (/^\d+(\.\d+)?$/.test(part)) { - return {part}; - } - return {part}; - })} - -
- ); - }); -} - -export default function SQLEditorContent() { - // Input State - const [prompt, setPrompt] = useState(''); - const [lastQuestion, setLastQuestion] = useState(''); - - // Pipeline State - const [isRunning, setIsRunning] = useState(false); - const [isSavingExample, setIsSavingExample] = useState(false); - const [hasResults, setHasResults] = useState(false); - - // Data State - const [generatedSql, setGeneratedSql] = useState('-- Ask the AI to generate SQL...'); - const [queryResults, setQueryResults] = useState>>([]); - const [explanation, setExplanation] = useState('Enter a natural language prompt to generate a query, retrieve data, and see the explanation here.'); - const [warnings, setWarnings] = useState([]); - const [errorMessage, setErrorMessage] = useState(''); - - // Pagination State - const [currentPage, setCurrentPage] = useState(1); - const pageSize = 5; - const totalRows = queryResults.length; - const totalPages = Math.max(1, Math.ceil(totalRows / pageSize)); - - // Handle Pipeline Execution - const handleGenerateAndRun = async () => { - if (!prompt.trim()) return; - setIsRunning(true); - setLastQuestion(prompt); - setErrorMessage(''); - setWarnings([]); - - try { - const data = await api.executeQuery(prompt); - setGeneratedSql(data.sql); - setQueryResults(data.rows || []); - setExplanation(data.explanation || 'Query generated successfully.'); - setWarnings(data.warnings || []); - setHasResults(true); - setCurrentPage(1); - } catch (error) { - setGeneratedSql('-- Error generating SQL. Check connection.'); - const message = error instanceof ApiError ? error.message : 'The pipeline failed to generate or execute the query.'; - setExplanation(message); - setErrorMessage(message); - setQueryResults([]); - setHasResults(false); - } finally { - setIsRunning(false); - } - }; - - const handleSaveVerifiedExample = async () => { - if (!lastQuestion || !generatedSql) { - alert("Please generate a query first."); - return; - } - setIsSavingExample(true); - try { - await api.saveVerifiedExample(lastQuestion, generatedSql); - alert('Verified example saved for semantic retrieval.'); - } catch (error) { - alert(error instanceof ApiError ? error.message : 'Failed to save verified example.'); - } finally { - setIsSavingExample(false); - } - }; - - // Pagination slice - const paginatedResults = queryResults.slice( - (currentPage - 1) * pageSize, - currentPage * pageSize - ); - - return ( -
- {/* Header */} -
-

AI Query Assistant

-

Ask questions, not SQL. Generated queries are validated before execution.

-
- {errorMessage &&
{errorMessage}
} - - {/* AI Prompt Bar */} -
-
- - setPrompt(e.target.value)} - onKeyDown={(e) => e.key === 'Enter' && handleGenerateAndRun()} - placeholder="e.g., Show me the top 10 customers by revenue last month" - className="flex-1 text-sm text-gray-600 placeholder-gray-400 bg-transparent outline-none" - /> - -
-
- - {/* Main two-column area */} -
- {/* SQL Editor Panel */} -
- {/* Editor header */} -
-
- Generated SQL · PostgreSQL -
- - demo_ecom_db -
-
-
- - - {/* NEW TRAIN AI BUTTON */} - -
-
- - {/* Code area - Made editable */} -
-