From d2cebdcb515f422aee4db3ce81bdfdedf8d68573 Mon Sep 17 00:00:00 2001 From: Eliraz Date: Mon, 6 Jul 2026 17:22:59 +0300 Subject: [PATCH] add temporal service --- agent/uv.lock | 4 + ...56b789_add_is_partial_to_table_profiles.py | 29 + backend/app/config.py | 3 + backend/app/infra_init.py | 3 +- backend/app/routers/profiling.py | 117 +++- backend/e2e_test.py | 2 +- backend/pyproject.toml | 1 + backend/uv.lock | 42 ++ core/pyproject.toml | 2 + core/src/core/config.py | 6 + core/src/core/models/models.py | 2 + core/src/core/services/__init__.py | 0 .../src/core}/services/profiling_engine.py | 534 +++++++++++++----- docker-compose.yml | 64 +++ frontend/src/api/client.ts | 6 +- .../src/components/tables/ProfilingTab.tsx | 81 ++- frontend/src/types/index.ts | 1 + worker/Dockerfile | 27 + worker/app/config.py | 18 + worker/app/main.py | 66 +++ worker/app/workflows/__init__.py | 1 + worker/app/workflows/profiling_activities.py | 304 ++++++++++ worker/app/workflows/profiling_workflow.py | 212 +++++++ worker/pyproject.toml | 19 + 24 files changed, 1359 insertions(+), 185 deletions(-) create mode 100644 backend/alembic/versions/a123e456b789_add_is_partial_to_table_profiles.py create mode 100644 core/src/core/services/__init__.py rename {backend/app => core/src/core}/services/profiling_engine.py (62%) create mode 100644 worker/Dockerfile create mode 100644 worker/app/config.py create mode 100644 worker/app/main.py create mode 100644 worker/app/workflows/__init__.py create mode 100644 worker/app/workflows/profiling_activities.py create mode 100644 worker/app/workflows/profiling_workflow.py create mode 100644 worker/pyproject.toml diff --git a/agent/uv.lock b/agent/uv.lock index 9698e93..85f958a 100644 --- a/agent/uv.lock +++ b/agent/uv.lock @@ -199,6 +199,8 @@ version = "0.1.0" source = { directory = "../core" } dependencies = [ { name = "alembic" }, + { name = "langchain" }, + { name = "langchain-openai" }, { name = "pgvector" }, { name = "psycopg2-binary" }, { name = "pydantic" }, @@ -211,6 +213,8 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "alembic", specifier = "==1.14.0" }, + { name = "langchain" }, + { name = "langchain-openai" }, { name = "pgvector", specifier = ">=0.2.0" }, { name = "psycopg2-binary", specifier = "==2.9.9" }, { name = "pydantic", specifier = "==2.10.4" }, diff --git a/backend/alembic/versions/a123e456b789_add_is_partial_to_table_profiles.py b/backend/alembic/versions/a123e456b789_add_is_partial_to_table_profiles.py new file mode 100644 index 0000000..d50fbfe --- /dev/null +++ b/backend/alembic/versions/a123e456b789_add_is_partial_to_table_profiles.py @@ -0,0 +1,29 @@ +"""add is_partial to table_profiles + +Revision ID: a123e456b789 +Revises: f9a3d1c8e205 +Create Date: 2026-07-05 12:00:00.000000 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "a123e456b789" +down_revision: Union[str, None] = "f9a3d1c8e205" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Add is_partial column to table_profiles table + op.add_column( + "table_profiles", + sa.Column("is_partial", sa.Boolean(), nullable=False, server_default="false") + ) + + +def downgrade() -> None: + op.drop_column("table_profiles", "is_partial") diff --git a/backend/app/config.py b/backend/app/config.py index 22c58f9..eec6162 100644 --- a/backend/app/config.py +++ b/backend/app/config.py @@ -94,7 +94,10 @@ def override_trino_settings(self) -> "Settings": LANGFUSE_WAIT_MAX_ATTEMPTS: int = 20 LANGFUSE_WAIT_INITIAL_DELAY_SECS: float = 0.5 LANGFUSE_WAIT_BACKOFF_FACTOR: float = 1.5 + # Concurrency and Trino PROFILER_MAX_CONCURRENT_QUERIES: int = 10 + PROFILING_CHUNK_SIZE: int = 5 + TEMPORAL_HOST: str = "localhost:7233" # JWT Config JWT_SECRET: str = "dev-secret-change-in-production" diff --git a/backend/app/infra_init.py b/backend/app/infra_init.py index 69574bb..674e2ed 100644 --- a/backend/app/infra_init.py +++ b/backend/app/infra_init.py @@ -1090,7 +1090,7 @@ def _ensure_airlines_registered() -> None: ) from sqlmodel import Session, select - from app.services.profiling_engine import run_table_profiling + from core.services.profiling_engine import run_table_profiling logger.info("[InfraInit] Registering airlines Snowflake tables...") @@ -1192,6 +1192,7 @@ def _run_profile( if result.success else ProfilingStatus.failed ) + profile.is_partial = not result.success or bool(result.errors) profile.version = result.version profile.row_count = result.row_count profile.sample_size = result.sample_size diff --git a/backend/app/routers/profiling.py b/backend/app/routers/profiling.py index 34cc4ea..fa9b620 100644 --- a/backend/app/routers/profiling.py +++ b/backend/app/routers/profiling.py @@ -9,6 +9,10 @@ import traceback from datetime import datetime, timedelta from typing import Any +import anyio +import anyio.to_thread +from temporalio.client import Client, WorkflowExecutionStatus +from app.config import settings from core.db.engine import engine, get_session from core.models.models import ( @@ -26,7 +30,7 @@ from sqlmodel import Session, select from app.services.join_detection import discover_joins_for_table -from app.services.profiling_engine import ( +from core.services.profiling_engine import ( build_context_for_llm, generate_table_summary, run_table_profiling, @@ -74,10 +78,36 @@ def _upsert_ai_summary(session: Session, table_id: str, summary: str) -> None: # ── Background worker ────────────────────────────────────────────────────────── -def _run_profile_job(table_id: str): +# ── Background worker ────────────────────────────────────────────────────────── +async def trigger_temporal_profiling_workflow(table_id: str, resume_from_partial: bool = False) -> bool: + """ + Attempts to trigger the profiling workflow via Temporal. + Returns True if started successfully, False otherwise. + """ + try: + logger.info("[Profiling] Connecting to Temporal client at %s", settings.TEMPORAL_HOST) + client = await Client.connect(settings.TEMPORAL_HOST) + await client.start_workflow( + "TableProfilingWorkflow", + id=f"profile-{table_id}", + task_queue="profiling-tasks", + args=[table_id, resume_from_partial], + ) + logger.info("[Profiling] Successfully started Temporal workflow for table %s", table_id) + return True + except Exception as e: + if "WorkflowAlreadyStartedError" in str(type(e)) or "WorkflowExecutionAlreadyStartedError" in str(type(e)): + logger.info("[Profiling] Profiling workflow is already running for table %s", table_id) + return True + logger.error("[Profiling] Failed to start Temporal workflow for table %s: %s", table_id, e) + return False + + +def _run_profile_job_local(table_id: str): """ Background task: runs real Trino profiling via profiling_engine, then upserts results into table_profiles + column_profiles. + Fallback when Temporal is not available. """ with Session(engine) as session: table = session.get(Table, table_id) @@ -104,10 +134,6 @@ def _run_profile_job(table_id: str): logger.error(f"[Profiling] Engine failed for {table_id}: {exc}") return - if not result.success: - logger.error(f"[Profiling] Engine unsuccessful for {table_id}") - return - # Persist results (Upsert by table_id) with Session(engine) as session: profile = session.exec( @@ -118,6 +144,7 @@ def _run_profile_job(table_id: str): session.add(profile) profile.status = ProfilingStatus.completed + profile.is_partial = not result.success or bool(result.errors) profile.row_count = result.row_count profile.sample_size = result.sample_size profile.column_count = result.column_count @@ -178,9 +205,16 @@ def _run_profile_job(table_id: str): logger.warning("[Profiling] AI summary step failed for %s: %s", table_id, exc) +async def _trigger_profiling(table_id: str, resume_from_partial: bool = False): + success = await trigger_temporal_profiling_workflow(table_id, resume_from_partial) + if not success: + logger.info("Falling back to local FastAPI background task profiling for %s", table_id) + await anyio.to_thread.run_sync(_run_profile_job_local, table_id) + + # ── GET /tables/{id}/profile ─────────────────────────────────────────────────── @router.get("/tables/{table_id}/profile", response_model=TableProfileRead) -def get_table_profile(table_id: str, session: Session = Depends(get_session)): +async def get_table_profile(table_id: str, session: Session = Depends(get_session)): table = session.get(Table, table_id) if not table: raise HTTPException(status_code=404, detail="Table not found") @@ -194,6 +228,31 @@ def get_table_profile(table_id: str, session: Session = Depends(get_session)): status_code=404, detail="No profile found. Run POST /profile/run first." ) + # Sync state from Temporal if stuck + if profile.status in (ProfilingStatus.running, ProfilingStatus.pending): + try: + client = await Client.connect(settings.TEMPORAL_HOST) + handle = client.get_workflow_handle(f"profile-{table_id}") + desc = await handle.describe() + if desc.status in ( + WorkflowExecutionStatus.FAILED, + WorkflowExecutionStatus.TERMINATED, + WorkflowExecutionStatus.TIMED_OUT, + WorkflowExecutionStatus.CANCELED + ): + profile.status = ProfilingStatus.failed + session.add(profile) + session.commit() + session.refresh(profile) + except Exception as e: + if "NotFound" in str(type(e)): + profile.status = ProfilingStatus.failed + session.add(profile) + session.commit() + session.refresh(profile) + else: + logger.warning("[Profiling] Failed to sync temporal status for %s: %s", table_id, e) + return profile @@ -211,8 +270,7 @@ def run_all_profiles( tables = session.exec(select(Table)).all() count = 0 for table in tables: - # If not force, check for running profile - background_tasks.add_task(_run_profile_job, str(table.id)) + background_tasks.add_task(_trigger_profiling, table.id) count += 1 logger.info( @@ -233,6 +291,7 @@ def run_table_profile( table_id: str, background_tasks: BackgroundTasks, force: bool = False, + resume_from_partial: bool = False, session: Session = Depends(get_session), ): """ @@ -248,16 +307,54 @@ def run_table_profile( select(TableProfile).where(TableProfile.table_id == table_id) ).first() - background_tasks.add_task(_run_profile_job, table_id) + background_tasks.add_task(_trigger_profiling, table_id, resume_from_partial) logger.info(f"[Profiling] Queued profiling job: table={table_id}") if profile: + # Optimistically update to pending + profile.status = ProfilingStatus.pending + session.add(profile) + session.commit() + session.refresh(profile) return profile # Return a dummy pending profile just to satisfy response_model if not exists yet return TableProfile(table_id=table_id, status=ProfilingStatus.pending) +# ── POST /tables/{id}/profile/terminate ─────────────────────────────────────── +@router.post("/tables/{table_id}/profile/terminate", status_code=200) +async def terminate_table_profile(table_id: str, session: Session = Depends(get_session)): + """ + Terminates a running Temporal profiling workflow for the given table. + """ + table = session.get(Table, table_id) + profile = session.exec( + select(TableProfile).where(TableProfile.table_id == table_id) + ).first() + + if not table or not profile or profile.status not in (ProfilingStatus.running, ProfilingStatus.pending): + raise HTTPException(status_code=404, detail="Table not found") + + try: + client = await Client.connect(settings.TEMPORAL_HOST) + handle = client.get_workflow_handle(f"profile-{table_id}") + await handle.cancel() + except Exception as e: + logger.warning(f"Failed to cancel temporal workflow for {table_id}: {e}") + + profile = session.exec( + select(TableProfile).where(TableProfile.table_id == table_id) + ).first() + + if profile: + profile.status = ProfilingStatus.failed + session.add(profile) + session.commit() + + return {"message": "Profiling job terminated", "status": "failed"} + + # ── GET /tables/{id}/profile/columns ────────────────────────────────────────── @router.get( "/tables/{table_id}/profile/columns", response_model=list[ColumnProfileRead] diff --git a/backend/e2e_test.py b/backend/e2e_test.py index a9137c8..cc5f118 100644 --- a/backend/e2e_test.py +++ b/backend/e2e_test.py @@ -11,7 +11,7 @@ from sqlmodel import Session, select from app.routers.tables import create_table -from app.services.profiling_engine import run_table_profiling +from core.services.profiling_engine import run_table_profiling logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s" diff --git a/backend/pyproject.toml b/backend/pyproject.toml index 226e152..4215ada 100644 --- a/backend/pyproject.toml +++ b/backend/pyproject.toml @@ -25,6 +25,7 @@ dependencies = [ "minio>=7.2.0", "core", "mcp>=1.2.0", + "temporalio>=1.6.0", "python-core-utils[keycloak] @ git+ssh://git@github.com/matzpen-agency/python-core-utils.git@1.2.0#egg=python-core-utils", ] diff --git a/backend/uv.lock b/backend/uv.lock index 9c4330b..b483de3 100644 --- a/backend/uv.lock +++ b/backend/uv.lock @@ -1122,6 +1122,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/79/7b/2c79738432f5c924bef5071f933bcc9efd0473bac3b4aa584a6f7c1c8df8/mypy_extensions-1.1.0-py3-none-any.whl", hash = "sha256:1be4cccdb0f2482337c4743e60421de3a356cd97508abadd57d47403e94f5505", size = 4963, upload-time = "2025-04-22T14:54:22.983Z" }, ] +[[package]] +name = "nexus-rpc" +version = "1.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/35/d5/cd1ffb202b76ebc1b33c1332a3416e55a39929006982adc2b1eb069aaa9b/nexus_rpc-1.4.0.tar.gz", hash = "sha256:3b8b373d4865671789cc43623e3dc0bcbf192562e40e13727e17f1c149050fba", size = 82367, upload-time = "2026-02-25T22:01:34.053Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/11/52/6327a5f4fda01207205038a106a99848a41c83e933cd23ea2cab3d2ebc6c/nexus_rpc-1.4.0-py3-none-any.whl", hash = "sha256:14c953d3519113f8ccec533a9efdb6b10c28afef75d11cdd6d422640c40b3a49", size = 29645, upload-time = "2026-02-25T22:01:33.122Z" }, +] + [[package]] name = "numpy" version = "2.4.6" @@ -2000,6 +2012,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/96/00/2b325970b3060c7cecebab6d295afe763365822b1306a12eeab198f74323/starlette-0.41.3-py3-none-any.whl", hash = "sha256:44cedb2b7c77a9de33a8b74b2b90e9f50d11fcf25d8270ea525ad71a25374ff7", size = 73225, upload-time = "2024-11-18T19:45:02.027Z" }, ] +[[package]] +name = "temporalio" +version = "1.30.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nexus-rpc" }, + { name = "protobuf" }, + { name = "types-protobuf" }, + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/00/b0/ad8fc3cd7425c6551a637bf23c798e8fdd8eb7a3ec4fee4f46f7678ba8d2/temporalio-1.30.0.tar.gz", hash = "sha256:7c025919511bb465392d547e48ccb85fd560a995db4ebcc82fdb43cddf088e6f", size = 2686876, upload-time = "2026-07-02T21:04:46.713Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/39/842fdffe93388dd30ac12a53a698f71cbfb68b3bc938f30f3e5d6a36d4ad/temporalio-1.30.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:6773a6b708dee7675fcbb681bf28e48337ce43b8467ceba8f903e78ae68909f8", size = 14520026, upload-time = "2026-07-02T21:04:31.384Z" }, + { url = "https://files.pythonhosted.org/packages/40/f3/a2237d5265eb29de591abeac7610a48616b590a1b923b4919f60ee81adfa/temporalio-1.30.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:4038c3ce2d9acc12fef31dd16ef9be8cc7f721672da5662aa05e3942d0b5c9d1", size = 14018523, upload-time = "2026-07-02T21:04:34.405Z" }, + { url = "https://files.pythonhosted.org/packages/e6/57/dc648d812f4c688bd246a616f0d65c5f03b33675df2effb1b480e1df6d21/temporalio-1.30.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:108d1a56e174eabc18add58316084cdead230e239d3df22bbe999d6954986591", size = 14330502, upload-time = "2026-07-02T21:04:37.39Z" }, + { url = "https://files.pythonhosted.org/packages/1c/e1/dbd57de0f5090891850c2ee5490319834a12b704076b11f32a4a149998d4/temporalio-1.30.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47b156138de30c2cd6723d5bfc06a30abcf680344c5481abb31f2a98a9b1f809", size = 14833364, upload-time = "2026-07-02T21:04:40.454Z" }, + { url = "https://files.pythonhosted.org/packages/30/2a/6d41289c11465ba276a8b417f31d2e469f0fe4b8afacd61d5244151c5fce/temporalio-1.30.0-cp310-abi3-win_amd64.whl", hash = "sha256:3adee28d5ec47bd6309a5eeef7b00126373f119bc5c1b058a34de098542f4da7", size = 15181893, upload-time = "2026-07-02T21:04:43.609Z" }, +] + [[package]] name = "text2sql-backend" version = "0.1.0" @@ -2025,6 +2056,7 @@ dependencies = [ { name = "python-multipart" }, { name = "requests" }, { name = "sqlmodel" }, + { name = "temporalio" }, { name = "trino" }, { name = "uvicorn", extra = ["standard"] }, ] @@ -2064,6 +2096,7 @@ requires-dist = [ { name = "python-multipart", specifier = "==0.0.20" }, { name = "requests", specifier = "==2.32.3" }, { name = "sqlmodel", specifier = "==0.0.22" }, + { name = "temporalio", specifier = ">=1.6.0" }, { name = "trino", specifier = "==0.328.0" }, { name = "uvicorn", extras = ["standard"], specifier = "==0.32.1" }, ] @@ -2114,6 +2147,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/14/6a/e9fc6a5b8f9a380a4a56b9f1e4dba5c6899561868017b17f6de382808b6f/types_passlib-1.7.7.20260211-py3-none-any.whl", hash = "sha256:c0f1ad440c513a6c07f333b28249530686056fd54a7b3ac6128ae31fd46305d3", size = 40457, upload-time = "2026-02-10T15:11:58.647Z" }, ] +[[package]] +name = "types-protobuf" +version = "7.34.1.20260518" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/29/59/e2b13b499d15e6720150c4b1a8d91e31fcacf716b432397475b3151ff7e4/types_protobuf-7.34.1.20260518.tar.gz", hash = "sha256:28cfaded25889cb83ebfb63cfb0a43628f0b6f3785767bec17287dc6468795f2", size = 68936, upload-time = "2026-05-18T06:01:47.332Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/1f/ec5caf72c2e3b688ca3927e0979a04ddad19e1afc4bf1c199bd743e0f419/types_protobuf-7.34.1.20260518-py3-none-any.whl", hash = "sha256:a0a5337413347166439c0e07cbc26c6164d091401c6f01b1dfd8cdb966c4dd8f", size = 85992, upload-time = "2026-05-18T06:01:45.696Z" }, +] + [[package]] name = "types-pyasn1" version = "0.6.0.20260408" diff --git a/core/pyproject.toml b/core/pyproject.toml index 98652ec..38d7141 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -12,6 +12,8 @@ dependencies = [ "trino>=0.328.0", "pgvector>=0.2.0", "psycopg2-binary==2.9.9", + "langchain", + "langchain-openai", "python-core-utils[postgres,redis] @ git+ssh://git@github.com/matzpen-agency/python-core-utils.git@1.2.0", ] diff --git a/core/src/core/config.py b/core/src/core/config.py index e757850..5d2e3b7 100644 --- a/core/src/core/config.py +++ b/core/src/core/config.py @@ -24,6 +24,12 @@ class CoreSettings(BaseSettings): CATALOG_VALID_TTL: int = 300 + PROFILER_MAX_CONCURRENT_QUERIES: int = 20 + PROFILING_CHUNK_SIZE: int = 15 + LLM_BASE_URL: str | None = None + LLM_MODEL: str | None = None + LLM_API_KEY: str | None = None + # Starburst Galaxy Configuration USE_GALAXY: bool = False GALAXY_HOST: str = "" diff --git a/core/src/core/models/models.py b/core/src/core/models/models.py index f82ccae..1ff6d2c 100644 --- a/core/src/core/models/models.py +++ b/core/src/core/models/models.py @@ -445,6 +445,7 @@ class TableProfile(SQLModel, table=True): index=True, ) status: ProfilingStatus = Field(default=ProfilingStatus.pending) + is_partial: bool = Field(default=False) row_count: int | None = None sample_size: int | None = None # rows returned by TABLESAMPLE column_count: int | None = None @@ -465,6 +466,7 @@ class TableProfileRead(SQLModel): id: str table_id: str status: ProfilingStatus + is_partial: bool | None = None row_count: int | None sample_size: int | None column_count: int | None diff --git a/core/src/core/services/__init__.py b/core/src/core/services/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/app/services/profiling_engine.py b/core/src/core/services/profiling_engine.py similarity index 62% rename from backend/app/services/profiling_engine.py rename to core/src/core/services/profiling_engine.py index b60fce1..c0aef2b 100644 --- a/backend/app/services/profiling_engine.py +++ b/core/src/core/services/profiling_engine.py @@ -6,28 +6,45 @@ ready for PostgreSQL persistence and LLM context injection. """ -from langfuse.api.unstable.commons.types import array_options_evaluation_rule_filter import concurrent.futures import logging from dataclasses import dataclass, field from datetime import date, datetime from decimal import Decimal -from typing import Any +from typing import Any, Optional, TypedDict + +class PrecomputedMetrics(TypedDict, total=False): + null_count: int + null_rate: float + distinct_count: int + min: Any + max: Any + avg: Any + quants: list[Any] + stddev: Any + min_unix: Any + max_unix: Any + +from langchain_openai import ChatOpenAI +from langchain_core.messages import SystemMessage, HumanMessage +from pydantic import BaseModel, ConfigDict, Field +import os + +from core import TrinoExecutionResult from core.trino import execute_query_sync -from app.config import settings +import os logger = logging.getLogger(__name__) +from core.config import settings + # ── Thresholds ───────────────────────────────────────────────────────────────── SAMPLE_PERCENT = 10 # TABLESAMPLE BERNOULLI(10) SAMPLE_LIMIT = 10_000 TOP_VALUES_LIMIT = 50 MIN_EXAMPLES = 3 -QUERY_TIMEOUT_SECONDS = ( - settings.TRINO_REQUEST_TIMEOUT -) # Hard per-query timeout in seconds CATEGORICAL_DISTINCT_THRESHOLD = 50 # Max unique values for standard string categories NUMERIC_CATEGORICAL_THRESHOLD = 15 # Strict limit for integers (e.g., status codes) @@ -93,6 +110,42 @@ # ── Data structures ──────────────────────────────────────────────────────────── +class HistogramBucket(BaseModel): + lo: Any + hi: Any + count: int + label: str + +class BaseStats(BaseModel): + model_config = ConfigDict(coerce_numbers_to_str=False) + +class CategoricalStats(BaseStats): + values: list[Any] = Field(default_factory=list) + frequencies: dict[str, int] = Field(default_factory=dict) + top_values_raw: list[dict] = Field(default_factory=list) + +class ContinuousStats(BaseStats): + min: float | int | str | None = None + max: float | int | str | None = None + avg: float | None = None + stddev: float | None = None + q25: float | None = None + median: float | None = None + q75: float | None = None + histogram: list[HistogramBucket] | None = None + examples: list[str] | None = None + +class TimeStats(BaseStats): + min: float | int | str | None = None + max: float | int | str | None = None + stddev: float | None = None + q25: float | None = None + median: float | None = None + q75: float | None = None + histogram: list[HistogramBucket] | None = None + examples: list[str] | None = None + + @dataclass class ColumnStats: column_name: str @@ -105,11 +158,11 @@ class ColumnStats: is_large_categorical: bool = False is_free_text: bool = False is_geo: bool = False - is_temporal: bool = False + is_time: bool = False is_boolean: bool = False - is_continuous: bool = False + is_continuous: bool = False semantic_type: str = "continuous" # categorical | continuous | temporal | geo | large_categorical | free_text | boolean - + top_values: list[dict] = field(default_factory=list) value_frequencies: dict[str, int] = field(default_factory=dict) min_value: str | None = None @@ -184,9 +237,9 @@ def build_numeric_stats_query(fqn: str, col: str) -> str: f"SELECT " f' MIN({col}) AS min_val, ' f' MAX({col}) AS max_val, ' - f' AVG(CAST({col} AS DOUBLE)) AS avg_val, ' - f' APPROX_PERCENTILE(CAST({col} AS DOUBLE), ARRAY[0.25, 0.5, 0.75]) AS quants, ' - f' STDDEV_POP(CAST({col} AS DOUBLE)) AS std_val ' + f' AVG(TRY_CAST({col} AS DOUBLE)) AS avg_val, ' + f' APPROX_PERCENTILE(TRY_CAST({col} AS DOUBLE), ARRAY[0.25, 0.5, 0.75]) AS quants, ' + f' STDDEV_POP(TRY_CAST({col} AS DOUBLE)) AS std_val ' f'FROM {fqn} WHERE {col} IS NOT NULL' ) @@ -221,19 +274,106 @@ def build_generic_histogram_query( """ -def build_time_stats_query(fqn: str, col: str) -> str: - # Safely extract unix epoch stats for temporal fields +def build_time_stats_query(fqn: str, col: str, data_type: str = "") -> str: + dtype = data_type.lower() + if any(t in dtype for t in ["int", "double", "float", "real", "decimal", "numeric"]): + return ( + f"SELECT " + f' MIN({col}) AS min_val, ' + f' MAX({col}) AS max_val, ' + f' APPROX_PERCENTILE(TRY_CAST({col} AS DOUBLE), ARRAY[0.25, 0.5, 0.75]) AS quants, ' + f' STDDEV_POP(TRY_CAST({col} AS DOUBLE)) AS std_val, ' + f' MIN(TRY_CAST({col} AS DOUBLE)) AS min_unix, ' + f' MAX(TRY_CAST({col} AS DOUBLE)) AS max_unix ' + f'FROM {fqn} WHERE {col} IS NOT NULL' + ) return ( f"SELECT " f' MIN({col}) AS min_val, ' f' MAX({col}) AS max_val, ' - f' APPROX_PERCENTILE(to_unixtime(CAST({col} AS TIMESTAMP)), ARRAY[0.25, 0.5, 0.75]) AS quants, ' - f' STDDEV_POP(to_unixtime(CAST({col} AS TIMESTAMP))) AS std_val, ' - f' MIN(to_unixtime(CAST({col} AS TIMESTAMP))) AS min_unix, ' - f' MAX(to_unixtime(CAST({col} AS TIMESTAMP))) AS max_unix ' + f' APPROX_PERCENTILE(to_unixtime(TRY_CAST({col} AS TIMESTAMP)), ARRAY[0.25, 0.5, 0.75]) AS quants, ' + f' STDDEV_POP(to_unixtime(TRY_CAST({col} AS TIMESTAMP))) AS std_val, ' + f' MIN(to_unixtime(TRY_CAST({col} AS TIMESTAMP))) AS min_unix, ' + f' MAX(to_unixtime(TRY_CAST({col} AS TIMESTAMP))) AS max_unix ' f'FROM {fqn} WHERE {col} IS NOT NULL' ) +def build_combined_profiling_query(fqn: str, columns_meta: list[tuple]) -> str | None: + select_parts = ["COUNT(*) AS __row_count"] + has_any = False + + for col in columns_meta: + col_name, data_type = col[0], col[1] + if _is_row_type(data_type) or _is_complex_type(data_type): + continue + + safe_col = f'"{col_name}"' + select_parts.append(f'COUNT({safe_col}) AS "{col_name}__non_null"') + select_parts.append(f'APPROX_DISTINCT({safe_col}) AS "{col_name}__distinct"') + has_any = True + + dtype_lower = data_type.lower() + if any(t in dtype_lower for t in NUMERIC_TYPES): + select_parts.append(f'MIN({safe_col}) AS "{col_name}__min"') + select_parts.append(f'MAX({safe_col}) AS "{col_name}__max"') + select_parts.append(f'AVG(TRY_CAST({safe_col} AS DOUBLE)) AS "{col_name}__avg"') + select_parts.append(f'APPROX_PERCENTILE(TRY_CAST({safe_col} AS DOUBLE), ARRAY[0.25, 0.5, 0.75]) AS "{col_name}__quants"') + select_parts.append(f'STDDEV_POP(TRY_CAST({safe_col} AS DOUBLE)) AS "{col_name}__std"') + elif any(t in dtype_lower for t in TIME_TYPES): + select_parts.append(f'MIN({safe_col}) AS "{col_name}__min"') + select_parts.append(f'MAX({safe_col}) AS "{col_name}__max"') + select_parts.append(f'APPROX_PERCENTILE(to_unixtime(TRY_CAST({safe_col} AS TIMESTAMP)), ARRAY[0.25, 0.5, 0.75]) AS "{col_name}__quants"') + select_parts.append(f'STDDEV_POP(to_unixtime(TRY_CAST({safe_col} AS TIMESTAMP))) AS "{col_name}__std"') + select_parts.append(f'MIN(to_unixtime(TRY_CAST({safe_col} AS TIMESTAMP))) AS "{col_name}__min_unix"') + select_parts.append(f'MAX(to_unixtime(TRY_CAST({safe_col} AS TIMESTAMP))) AS "{col_name}__max_unix"') + + if not has_any: + return None + + return f"SELECT {', '.join(select_parts)} FROM {fqn}" + + +def parse_combined_result(columns_meta: list[tuple], row_dict: dict[str, Any], row_count: int) -> dict[str, PrecomputedMetrics]: + col_metrics = {} + for col in columns_meta: + col_name, data_type = col[0], col[1] + if _is_row_type(data_type) or _is_complex_type(data_type): + continue + + col_key = col_name.lower() + non_null_key = f"{col_key}__non_null" + distinct_key = f"{col_key}__distinct" + if non_null_key not in row_dict: + continue + + non_nulls = int(row_dict[non_null_key] or 0) + nulls = row_count - non_nulls + distinct_count = int(row_dict[distinct_key] or 0) + + metrics: PrecomputedMetrics = { + "null_count": nulls, + "null_rate": round(nulls / max(row_count, 1), 4), + "distinct_count": distinct_count, + } + + dtype_lower = data_type.lower() + if any(t in dtype_lower for t in NUMERIC_TYPES): + metrics["min"] = row_dict.get(f"{col_key}__min") + metrics["max"] = row_dict.get(f"{col_key}__max") + metrics["avg"] = row_dict.get(f"{col_key}__avg") + metrics["quants"] = row_dict.get(f"{col_key}__quants") # type: ignore + metrics["stddev"] = row_dict.get(f"{col_key}__std") + elif any(t in dtype_lower for t in TIME_TYPES): + metrics["min"] = row_dict.get(f"{col_key}__min") + metrics["max"] = row_dict.get(f"{col_key}__max") + metrics["quants"] = row_dict.get(f"{col_key}__quants") # type: ignore + metrics["stddev"] = row_dict.get(f"{col_key}__std") + metrics["min_unix"] = row_dict.get(f"{col_key}__min_unix") + metrics["max_unix"] = row_dict.get(f"{col_key}__max_unix") + + col_metrics[col_name] = metrics + return col_metrics + # ── Helpers ──────────────────────────────────────────────────────────────────── def _fqn(catalog: str, schema: str, table: str) -> str: return f'"{catalog}"."{schema}"."{table}"' @@ -259,11 +399,7 @@ def _make_json_safe(obj: Any) -> Any: return obj -def _safe_float(val) -> float | None: - try: - return float(val) if val is not None else None - except (TypeError, ValueError): - return None + def _is_complex_type(dtype: str) -> bool: @@ -277,39 +413,50 @@ def _is_row_type(dtype: str) -> bool: return dtype.lower().strip().startswith("row(") -_global_trino_executor = concurrent.futures.ThreadPoolExecutor( - max_workers=settings.PROFILER_MAX_CONCURRENT_QUERIES -) +_global_trino_executor = None + +def _get_executor(): + global _global_trino_executor + if _global_trino_executor is None: + _global_trino_executor = concurrent.futures.ThreadPoolExecutor( + max_workers=settings.PROFILER_MAX_CONCURRENT_QUERIES + ) + return _global_trino_executor def execute_with_timeout(query: str, table_id: str): """Runs execute_query_sync in a global thread pool and enforces a hard wall-clock timeout.""" - future = _global_trino_executor.submit(execute_query_sync, query, table_id) + request_timeout = settings.TRINO_REQUEST_TIMEOUT + future = _get_executor().submit(execute_query_sync, query, table_id) try: - return future.result(timeout=QUERY_TIMEOUT_SECONDS) + return future.result(timeout=request_timeout) except concurrent.futures.TimeoutError: logger.warning( - f"[TrinoClient] Query timed out after {QUERY_TIMEOUT_SECONDS}s: {query[:120]}" + f"[TrinoClient] Query timed out after {request_timeout}s: {query[:120]}" ) - from core import TrinoExecutionResult - return TrinoExecutionResult( success=False, - error_message=f"Query timed out after {QUERY_TIMEOUT_SECONDS}s", + columns=[], + rows=[], + error_message=f"Query timed out after {request_timeout}s" ) def _fetch_one(query: str, table_id: str, default: Any = None) -> Any: """Helper to fetch a single row and handle boilerplates.""" res = execute_with_timeout(query, table_id) - if res.success and res.rows: + if not res.success: + raise Exception(f"Query failed: {res.error_message}") + if res.rows: return res.rows[0] return default def _fetch_all(query: str, table_id: str) -> list[tuple]: """Helper to fetch all rows safely.""" res = execute_with_timeout(query, table_id) - return res.rows if res.success and res.rows else [] + if not res.success: + raise Exception(f"Query failed: {res.error_message}") + return res.rows if res.rows else [] def _classify_semantic_type( @@ -323,18 +470,18 @@ def _classify_semantic_type( """Classifies column data dynamically based on limits and statistical ratios.""" dtype_lower = data_type.lower() col_lower = col_name.lower() - + # 1. Base Type Flags is_geo = col_lower in GEO_HINTS or "geo" in col_lower or "coord" in col_lower - is_temporal = dtype_lower in TIME_TYPES or any(h in col_lower for h in TIME_HINTS) + is_time = dtype_lower in TIME_TYPES or any(h in col_lower for h in TIME_HINTS) is_string = any(t in dtype_lower for t in STRING_TYPES) is_numeric = any(t in dtype_lower for t in NUMERIC_TYPES) - + # 2. Math Setup N = max(row_count - null_count, 1) # Non-null rows - U = distinct_count # Unique values + U = distinct_count # Unique values R = U / N if N > 0 else 0 # Uniqueness Ratio - + # ========================================== # LAYER 1: Explicit Database Types & Fast Pass # ========================================== @@ -342,8 +489,8 @@ def _classify_semantic_type( return "boolean" if is_geo: return "geo" - if is_temporal: - return "temporal" + if is_time: + return "time" # ========================================== # LAYER 2: Constants & The Boolean Heuristic @@ -359,18 +506,18 @@ def _classify_semantic_type( if is_numeric: if U <= NUMERIC_CATEGORICAL_THRESHOLD: return "categorical" - return "continuous" + return "continuous" # ========================================== # LAYER 4: High Uniqueness (Identifiers) # ========================================== - if R >= CARDINALITY_RATIO_IDENTIFIER: + if R >= CARDINALITY_RATIO_IDENTIFIER: return "free_text" # ========================================== # LAYER 5: The Anchor Category (Strict Lists) # ========================================== - if U <= CATEGORICAL_DISTINCT_THRESHOLD: + if U <= CATEGORICAL_DISTINCT_THRESHOLD: return "categorical" # ========================================== @@ -380,13 +527,13 @@ def _classify_semantic_type( # Stage 1: We don't have the T10 math yet. Tell the engine to go fetch it. if t10_coverage is None: return "requires_t10_check" - + # Stage 2: We have the T10 math. Apply the dispersion trap! - # If the data is moderately dispersed (R >= 0.50) AND the top 10 values + # If the data is moderately dispersed (R >= 0.50) AND the top 10 values # barely cover any ground (<= 10%), it is noisy free text. if R >= CARDINALITY_RATIO_DISPERSION and t10_coverage <= TOP_10_COVERAGE_DISPERSION: return "free_text" - + # Otherwise, it has enough repetition to be a searchable category (like Company Name). return "large_categorical" @@ -396,8 +543,12 @@ def _classify_semantic_type( # ── Stats Extraction Helpers ─────────────────────────────────────────────────── -def _build_histogram_data(fqn: str, field_path: str, table_id: str, min_val: float, max_val: float, is_temporal: bool = False) -> list[dict]: - cast_expr = f'to_unixtime(CAST({field_path} AS TIMESTAMP))' if is_temporal else f'CAST({field_path} AS DOUBLE)' +def _build_histogram_data(fqn: str, field_path: str, table_id: str, min_val: float, max_val: float, is_time: bool = False, data_type: str = "") -> list[dict]: + dtype = data_type.lower() + if is_time and not any(t in dtype for t in ["int", "double", "float", "real", "decimal", "numeric"]): + cast_expr = f'to_unixtime(TRY_CAST({field_path} AS TIMESTAMP))' + else: + cast_expr = f'TRY_CAST({field_path} AS DOUBLE)' rows = _fetch_all(build_generic_histogram_query(fqn, field_path, cast_expr, min_val, max_val, 8), table_id) if not rows: return [] @@ -408,7 +559,7 @@ def _build_histogram_data(fqn: str, field_path: str, table_id: str, min_val: flo for i in range(1, 9): cnt = bucket_counts.get(i, 0) lo, hi = min_val + (i - 1) * step, min_val + i * step - if is_temporal: + if is_time: hist_data.append({ "lo": datetime.fromtimestamp(lo).isoformat() if max_val > min_val else datetime.fromtimestamp(min_val).isoformat(), "hi": datetime.fromtimestamp(hi).isoformat() if max_val > min_val else datetime.fromtimestamp(max_val).isoformat(), @@ -423,73 +574,145 @@ def _build_histogram_data(fqn: str, field_path: str, table_id: str, min_val: flo "label": f"{lo:g}", }) if max_val <= min_val: break # Constant value, single bucket needed - + if "null" in bucket_counts: hist_data.append({"lo": None, "hi": None, "count": bucket_counts["null"], "label": "Null / Unknown"}) return hist_data -def _extract_continuous_stats(fqn: str, field_path: str, table_id: str) -> dict: - row = _fetch_one(build_numeric_stats_query(fqn, field_path), table_id) - if not row or row[0] is None: return {} - - stats = { - "min": _make_json_safe(row[0]), "max": _make_json_safe(row[1]), - "avg": _safe_float(row[2]), "stddev": _safe_float(row[4]) - } - quants = row[3] - if isinstance(quants, list) and len(quants) >= 3: - stats.update({"q25": _safe_float(quants[0]), "median": _safe_float(quants[1]), "q75": _safe_float(quants[2])}) +def _extract_continuous_stats(fqn: str, field_path: str, table_id: str, precomputed: PrecomputedMetrics | None = None) -> ContinuousStats | None: + stats = None + if precomputed: + quants = precomputed.get("quants") + stats = ContinuousStats( + min=precomputed.get("min"), + max=precomputed.get("max"), + avg=precomputed.get("avg"), + stddev=precomputed.get("stddev"), + q25=quants[0] if isinstance(quants, list) and len(quants) >= 3 else None, + median=quants[1] if isinstance(quants, list) and len(quants) >= 3 else None, + q75=quants[2] if isinstance(quants, list) and len(quants) >= 3 else None, + ) + + if not stats or stats.min is None or stats.max is None: + row = _fetch_one(build_numeric_stats_query(fqn, field_path), table_id) + if not row or row[0] is None: return None + quants = row[3] + stats = ContinuousStats( + min=row[0], + max=row[1], + avg=row[2], + stddev=row[4], + q25=quants[0] if isinstance(quants, list) and len(quants) >= 3 else None, + median=quants[1] if isinstance(quants, list) and len(quants) >= 3 else None, + q75=quants[2] if isinstance(quants, list) and len(quants) >= 3 else None, + ) + + if stats.min is not None and stats.max is not None: + stats.histogram = _build_histogram_data(fqn, field_path, table_id, float(stats.min), float(stats.max)) - if stats.get("min") is not None and stats.get("max") is not None: - stats["histogram"] = _build_histogram_data(fqn, field_path, table_id, float(stats["min"]), float(stats["max"])) return stats -def _extract_temporal_stats(fqn: str, field_path: str, table_id: str) -> dict: - row = _fetch_one(build_time_stats_query(fqn, field_path), table_id) - if not row or row[0] is None: return {} - - stats = {"min": _make_json_safe(row[0]), "max": _make_json_safe(row[1]), "stddev": _safe_float(row[3])} - quants, min_unix, max_unix = row[2], _safe_float(row[4]), _safe_float(row[5]) +def _extract_time_stats(fqn: str, field_path: str, data_type: str, table_id: str, precomputed: PrecomputedMetrics | None = None, sample_data: list[dict] | None = None) -> TimeStats | None: + stats = None + min_unix = max_unix = None - if isinstance(quants, list) and len(quants) >= 3: - stats.update({"q25": _safe_float(quants[0]), "median": _safe_float(quants[1]), "q75": _safe_float(quants[2])}) - + if precomputed: + quants = precomputed.get("quants") + stats = TimeStats( + min=precomputed.get("min"), + max=precomputed.get("max"), + stddev=precomputed.get("stddev"), + q25=quants[0] if isinstance(quants, list) and len(quants) >= 3 else None, + median=quants[1] if isinstance(quants, list) and len(quants) >= 3 else None, + q75=quants[2] if isinstance(quants, list) and len(quants) >= 3 else None, + ) + min_unix = precomputed.get("min_unix") + max_unix = precomputed.get("max_unix") + if stats.min is None: + stats = None + + if not stats: + row = _fetch_one(build_time_stats_query(fqn, field_path, data_type), table_id) + if not row or row[0] is None: return None + quants = row[2] + stats = TimeStats( + min=row[0], + max=row[1], + stddev=row[3], + q25=quants[0] if isinstance(quants, list) and len(quants) >= 3 else None, + median=quants[1] if isinstance(quants, list) and len(quants) >= 3 else None, + q75=quants[2] if isinstance(quants, list) and len(quants) >= 3 else None, + ) + min_unix, max_unix = row[4], row[5] + + dtype = data_type.lower() + + try: + min_unix = float(min_unix) if min_unix is not None else None + max_unix = float(max_unix) if max_unix is not None else None + except (TypeError, ValueError): + min_unix = max_unix = None + if min_unix is not None and max_unix is not None: - stats["histogram"] = _build_histogram_data(fqn, field_path, table_id, min_unix, max_unix, is_temporal=True) + stats.histogram = _build_histogram_data(fqn, field_path, table_id, min_unix, max_unix, is_time=True, data_type=data_type) + + # Extract examples from sample_data to avoid extra query + examples = [] + col_name = field_path.strip('"') + if sample_data: + for r in sample_data: + val = r.get(col_name) + if val is not None and val != "": + val_str = str(val) + if val_str not in examples: + examples.append(val_str) + if len(examples) >= 3: + break + if not examples: + ex_rows = _fetch_all(f"SELECT DISTINCT {field_path} FROM {fqn} WHERE {field_path} IS NOT NULL LIMIT 3", table_id) + examples = [str(r[0]) for r in ex_rows] - ex_rows = _fetch_all(f"SELECT DISTINCT {field_path} FROM {fqn} WHERE {field_path} IS NOT NULL LIMIT 3", table_id) - stats["examples"] = [str(r[0]) for r in ex_rows] + stats.examples = examples return stats -def _extract_categorical_stats(fqn: str, field_path: str, table_id: str, limit: int = TOP_VALUES_LIMIT) -> dict: +def _extract_categorical_stats(fqn: str, field_path: str, table_id: str, limit: int = TOP_VALUES_LIMIT) -> CategoricalStats | None: rows = _fetch_all(build_top_values_query(fqn, field_path, limit=limit), table_id) + if not rows: + return None top_vals = [{"value": _make_json_safe(r[0]), "count": int(r[1])} for r in rows] - return { - "values": [v["value"] for v in top_vals], - "frequencies": {v["value"]: v["count"] for v in top_vals}, - "top_values_raw": top_vals # Internal use - } + stats = CategoricalStats( + values=[v["value"] for v in top_vals], + frequencies={str(v["value"]): v["count"] for v in top_vals}, + top_values_raw=top_vals + ) + return stats # ── Extraction Pipelines ─────────────────────────────────────────────────────── -def _process_column_metrics(fqn: str, table_id: str, field_path: str, field_name: str, field_type: str, row_count: int) -> dict: +def _process_column_metrics(fqn: str, table_id: str, field_path: str, field_name: str, field_type: str, row_count: int, precomputed: PrecomputedMetrics | None = None, sample_data: list[dict] | None = None) -> dict: """Shared pipeline for semantic detection and stat extraction for both root columns and row children.""" stats: dict[str, Any] = {"type": field_type.lower()} - - null_row = _fetch_one(build_null_ratio_query(fqn, field_path), table_id, default=(1, 0)) - if null_row is None: - logger.warning(f"Null count query failed for {field_path}. Defaulting to 0.") - null_row = (1, 0) - total, non_nulls = int(null_row[0] or 1), int(null_row[1] or 0) - nulls = total - non_nulls - - dist_row = _fetch_one(build_distinct_count_query(fqn, field_path), table_id, default=(0,)) - if dist_row is None: - logger.warning(f"Distinct count query failed for {field_path}. Defaulting to 0.") - distinct_count = 0 + + if precomputed is None: + null_row = _fetch_one(build_null_ratio_query(fqn, field_path), table_id, default=(1, 0)) + if null_row is None: + logger.warning(f"Null count query failed for {field_path}. Defaulting to 0.") + null_row = (1, 0) + total, non_nulls = int(null_row[0] or 1), int(null_row[1] or 0) + nulls = total - non_nulls + + dist_row = _fetch_one(build_distinct_count_query(fqn, field_path), table_id, default=(0,)) + if dist_row is None: + logger.warning(f"Distinct count query failed for {field_path}. Defaulting to 0.") + distinct_count = 0 + else: + distinct_count = int(dist_row[0] or 0) else: - distinct_count = int(dist_row[0] or 0) + nulls = precomputed["null_count"] + distinct_count = precomputed["distinct_count"] + total = row_count semantic_type = _classify_semantic_type(field_name, field_type, row_count, nulls, distinct_count) + t10_rows = None if semantic_type == "requires_t10_check": t10_rows = _fetch_all(build_top_values_query(fqn, field_path, limit=10), table_id) t10_coverage = sum(int(r[1]) for r in t10_rows) / max(row_count - nulls, 1) if t10_rows else 0.0 @@ -504,19 +727,28 @@ def _process_column_metrics(fqn: str, table_id: str, field_path: str, field_name if semantic_type in ("categorical", "boolean"): cat_stats = _extract_categorical_stats(fqn, field_path, table_id) - stats.update(cat_stats) - if semantic_type == "boolean": - stats["examples"] = cat_stats.get("values", [])[:3] + if cat_stats: + cat_dict = cat_stats.model_dump(exclude_none=True) + stats.update(cat_dict) + if semantic_type == "boolean": + stats["examples"] = cat_dict.get("values", [])[:3] elif semantic_type == "large_categorical": - ex_rows = _fetch_all(build_top_values_query(fqn, field_path, limit=3), table_id) - stats["examples"] = [str(r[0]) for r in ex_rows] + if t10_rows is not None: + stats["examples"] = [str(r[0]) for r in t10_rows[:3]] + else: + ex_rows = _fetch_all(build_top_values_query(fqn, field_path, limit=3), table_id) + stats["examples"] = [str(r[0]) for r in ex_rows] elif semantic_type == "continuous": - stats.update(_extract_continuous_stats(fqn, field_path, table_id)) + c_stats = _extract_continuous_stats(fqn, field_path, table_id, precomputed) + if c_stats: + stats.update(c_stats.model_dump(exclude_none=True)) - elif semantic_type == "temporal": - stats.update(_extract_temporal_stats(fqn, field_path, table_id)) + elif semantic_type == "time": + t_stats = _extract_time_stats(fqn, field_path, field_type, table_id, precomputed, sample_data) + if t_stats: + stats.update(t_stats.model_dump(exclude_none=True)) return stats @@ -542,14 +774,13 @@ def _analyze_row_column( "note": "Max depth reached", "children": [], } - - from app.services.trino_client import _parse_row_fields # Assuming parser is accessible, kept logic intact mentally - + + # Simple inline parser for brevity (matches your original) def parse_fields(rt): inner = rt.strip()[4:-1].strip() if rt.lower().startswith("row(") else rt return [tuple(p.strip().split(None, 1)) if len(p.strip().split(None, 1)) == 2 else (p.strip(), "unknown") for p in inner.split(',')] - + fields = parse_fields(row_type) children: list[dict] = [] @@ -585,7 +816,9 @@ def parse_fields(rt): # ── Per-Column Analysis ──────────────────────────────────────────────────────── def _analyze_column( - fqn: str, table_id: str, col_name: str, data_type: str, row_count: int + fqn: str, table_id: str, col_name: str, data_type: str, row_count: int, + precomputed: PrecomputedMetrics | None = None, + sample_data: list[dict] | None = None, ) -> ColumnStats: stats = ColumnStats(column_name=col_name, data_type=data_type) safe_col = f'"{col_name}"' @@ -605,17 +838,17 @@ def _analyze_column( "note": "Skipped (array/map/json)", } return stats - - metrics = _process_column_metrics(fqn, table_id, safe_col, col_name, data_type, row_count) - + + metrics = _process_column_metrics(fqn, table_id, safe_col, col_name, data_type, row_count, precomputed, sample_data) + # Map back to ColumnStats dataclass stats.semantic_type = metrics.pop("semantic_type") setattr(stats, f"is_{stats.semantic_type}", True) - + stats.null_count = metrics.pop("null_count", 0) stats.null_rate = metrics.pop("null_rate", 0.0) stats.distinct_count = metrics.pop("distinct_count", 0) - + # Populate specific fields based on semantic type if "top_values_raw" in metrics: stats.top_values = metrics.pop("top_values_raw") @@ -655,9 +888,25 @@ def run_table_profiling(table_id: str, catalog: str, schema: str, table: str, ve columns_meta = _fetch_all(build_column_metadata_query(catalog, schema, table), table_id) result.column_count = len(columns_meta) + # Combined stats query execution + precomputed_metrics = {} + if result.row_count > 0 and columns_meta: + combined_query = build_combined_profiling_query(fqn, columns_meta) + if combined_query: + try: + res = execute_with_timeout(combined_query, table_id) + if res.success and res.rows: + row_dict = dict(zip(res.columns, res.rows[0])) + precomputed_metrics = parse_combined_result(columns_meta, row_dict, result.row_count) + logger.info("[ProfilingEngine] Combined profiling query succeeded.") + else: + logger.warning("[ProfilingEngine] Combined profiling query failed or returned no rows: %s", res.error_message) + except Exception as e: + logger.warning("[ProfilingEngine] Combined profiling query failed, falling back to individual: %s", e) + col_stats: list[ColumnStats] = [] with concurrent.futures.ThreadPoolExecutor(max_workers=min(30, len(columns_meta) or 1)) as worker_executor: - futures = {worker_executor.submit(_analyze_column, fqn, table_id, col[0], col[1], result.row_count): col[:2] for col in columns_meta} + futures = {worker_executor.submit(_analyze_column, fqn, table_id, col[0], col[1], result.row_count, precomputed_metrics.get(col[0]), result.sample_data): col[:2] for col in columns_meta} for future in concurrent.futures.as_completed(futures): try: col_name, data_type = futures[future] @@ -675,18 +924,18 @@ def run_table_profiling(table_id: str, catalog: str, schema: str, table: str, ve # Auto Insights... insights = [f"~{result.row_count:,} rows (COUNT(*))."] if result.row_count else [] if result.sample_size: insights.append(f"{result.sample_size:,} rows sampled.") - - for flag, name in [("is_categorical", "categorical"), ("is_temporal", "Time"), ("is_geo", "Geographic")]: + + for flag, name in [("is_categorical", "categorical"), ("is_time", "Time"), ("is_geo", "Geographic")]: cols = [c.column_name for c in col_stats if getattr(c, flag)] if cols: insights.append(f"{name} columns: {', '.join(cols[:5])}.") high_null = [c.column_name for c in col_stats if c.null_rate > 0.20] if high_null: insights.append(f"High null rate (>20%): {', '.join(high_null[:5])}.") - + if result.row_count > 0: pk_candidates = [c.column_name for c in col_stats if c.distinct_count >= result.row_count * 0.95] if pk_candidates: insights.append(f"PK candidates: {', '.join(pk_candidates[:3])}.") - + result.auto_insights = insights # Profile JSON construction @@ -697,7 +946,7 @@ def run_table_profiling(table_id: str, catalog: str, schema: str, table: str, ve "columns": [{ "name": c.column_name, "data_type": c.data_type, "semantic_type": c.semantic_type, "is_categorical": c.is_categorical, "is_large_categorical": c.is_large_categorical, - "is_free_text": c.is_free_text, "is_boolean": c.is_boolean, "is_temporal": c.is_temporal, + "is_free_text": c.is_free_text, "is_boolean": c.is_boolean, "is_time": c.is_time, "is_geo": c.is_geo, "is_continuous": c.is_continuous, "distinct_count": c.distinct_count, "null_rate": c.null_rate, "stats": c.stats_json } for c in col_stats], @@ -733,7 +982,7 @@ def build_context_for_llm( # Whitelist safe keys directly into the context object to prevent nesting boilerplate for key in ["values", "examples", "min", "max", "avg", "distinct_count"]: if key in stats: col_ctx[key] = stats[key] - + context_columns.append(col_ctx) return { @@ -748,18 +997,8 @@ def build_context_for_llm( def generate_table_summary(result: "TableProfilingResult") -> str: - """ - Generate a ≤3-sentence plain-English description of a table using the LLM. - Called once at the end of `_run_profile_job` and stored in EnrichmentVersion. - - Uses the same LLM endpoint as the agent (LLM_BASE_URL / LLM_API_KEY / LLM_MODEL). - Returns an empty string on any failure so profiling is never blocked. - """ - import os - - llm_base_url = os.environ.get("LLM_BASE_URL", "http://localhost:11434/v1") - llm_api_key = os.environ.get("LLM_API_KEY", "ollama") - llm_model = os.environ.get("LLM_MODEL", "gemma4:e4b") + """Uses LLM to summarize the entire table based on JSON profile.""" + try: col_lines = [] @@ -782,26 +1021,17 @@ def generate_table_summary(result: "TableProfilingResult") -> str: "for querying. Be specific about what the table contains." ) - import httpx + llm = ChatOpenAI( + model=settings.LLM_MODEL, + base_url=settings.LLM_BASE_URL, + api_key=settings.LLM_API_KEY, + temperature=0.0, + timeout=60.0 + ) + + response = llm.invoke(prompt) + return response.content - payload = { - "model": llm_model, - "messages": [{"role": "user", "content": prompt}], - "temperature": 0.0, - "stream": False, - } - with httpx.Client(timeout=60.0) as client: - resp = client.post( - f"{llm_base_url}/chat/completions", - headers={ - "Authorization": f"Bearer {llm_api_key}", - "Content-Type": "application/json", - }, - json=payload, - ) - resp.raise_for_status() - data = resp.json() - return data["choices"][0]["message"]["content"].strip() except Exception as exc: logger.warning("[ProfilingEngine] generate_table_summary failed: %s", exc) return "" diff --git a/docker-compose.yml b/docker-compose.yml index b7d786a..f72ed5d 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -411,6 +411,69 @@ services: - ingestion-volume-dags:/opt/airflow/dags - ingestion-volume-tmp:/tmp + temporal: + image: temporalio/auto-setup:1.24.2 + ports: + - "7233:7233" + - "8233:8233" + environment: + - DB=postgres12 + - DB_PORT=5432 + - POSTGRES_USER=postgres + - POSTGRES_PWD=postgres + - POSTGRES_SEEDS=db + depends_on: + db: + condition: service_healthy + + temporal-ui: + image: temporalio/ui:latest + ports: + - "8088:8080" + environment: + - TEMPORAL_ADDRESS=temporal:7233 + - TEMPORAL_CORS_ORIGINS=http://localhost:3000 + - TEMPORAL_CSRF_COOKIE_INSECURE=true + depends_on: + temporal: + condition: service_started + + worker: + build: + context: . + dockerfile: ./worker/Dockerfile + secrets: + - deploy_key + command: python -m app.main + env_file: + - ./backend/.env + environment: + - DATABASE_URL=postgresql://postgres:postgres@db:5432/text2sql + - POSTGRES_HOST=db + - POSTGRES_USER=postgres + - POSTGRES_PASSWORD=postgres + - POSTGRES_PORT=5432 + - POSTGRES_DB=text2sql + - TRINO_HOST=trino + - TRINO_PORT=8080 + - TRINO_USER=trino + - TRINO_PASSWORD= + - TRINO_CATALOG=minio + - TRINO_SCHEMA=simple_retail + - TRINO_HTTP_SCHEME=http + - TRINO_REQUEST_TIMEOUT=600.0 + - TRINO_ENABLED=true + - TEMPORAL_HOST=temporal:7233 + - LLM_BASE_URL=http://host.docker.internal:11434/v1 + - LLM_MODEL=gemma4:e4b + - LLM_API_KEY=ollama + - ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS=600 + depends_on: + db: + condition: service_healthy + temporal: + condition: service_started + backend: build: context: . @@ -460,6 +523,7 @@ services: - MINIO_SECRET_KEY=password123 - REDIS_URL=redis://redis:6379 - REDIS_SSL=false + - TEMPORAL_HOST=temporal:7233 depends_on: db: condition: service_healthy diff --git a/frontend/src/api/client.ts b/frontend/src/api/client.ts index d297a1b..dd55c2c 100644 --- a/frontend/src/api/client.ts +++ b/frontend/src/api/client.ts @@ -128,8 +128,10 @@ export const auditApi = { // ── Profiling ───────────────────────────────────────────────────────────────── export const profilingApi = { get: (tableId: string) => api.get(`/tables/${tableId}/profile`).then((r) => r.data), - run: (tableId: string) => - api.post(`/tables/${tableId}/profile/run`).then((r) => r.data), + run: (tableId: string, params?: { force?: boolean; resume_from_partial?: boolean }) => + api.post(`/tables/${tableId}/profile/run`, null, { params }).then((r) => r.data), + terminate: (tableId: string) => + api.post(`/tables/${tableId}/profile/terminate`).then((r) => r.data), getColumns: (tableId: string) => api.get(`/tables/${tableId}/profile/columns`).then((r) => r.data), getCrossProfiles: (tableId: string) => diff --git a/frontend/src/components/tables/ProfilingTab.tsx b/frontend/src/components/tables/ProfilingTab.tsx index b680971..debccdf 100644 --- a/frontend/src/components/tables/ProfilingTab.tsx +++ b/frontend/src/components/tables/ProfilingTab.tsx @@ -10,6 +10,7 @@ import { Link2, Maximize2, RefreshCw, + X, } from 'lucide-react'; import { Area, @@ -1337,18 +1338,25 @@ export function ProfilingTab({ tableId }: { tableId: string }) { }); const runMutation = useMutation({ - mutationFn: () => profilingApi.run(tableId), + mutationFn: (params?: { force?: boolean; resume_from_partial?: boolean }) => profilingApi.run(tableId, params), + onSuccess: () => { + qc.invalidateQueries({ queryKey: ['profile', tableId] }); + }, + }); + + const terminateMutation = useMutation({ + mutationFn: () => profilingApi.terminate(tableId), onSuccess: () => { qc.invalidateQueries({ queryKey: ['profile', tableId] }); }, }); const profile = profileQ.data; - const isRunning = profile?.status === 'running' || runMutation.isPending; + const isRunning = profile?.status === 'running' || profile?.status === 'pending' || runMutation.isPending; useEffect(() => { let poll: ReturnType; - if (profile?.status === 'running') { + if (profile?.status === 'running' || profile?.status === 'pending') { poll = setInterval(async () => { await qc.invalidateQueries({ queryKey: ['profile', tableId] }); await qc.invalidateQueries({ queryKey: ['profile-columns', tableId] }); @@ -1390,25 +1398,60 @@ export function ProfilingTab({ tableId }: { tableId: string }) { )} - + ) : profile?.status === 'failed' || profile?.is_partial ? ( + <> + + + ) : ( - + )} - {isRunning ? 'Running…' : profile ? 'Re-profile' : 'Run Profiling'} - + {!profile && !isRunning && ( diff --git a/frontend/src/types/index.ts b/frontend/src/types/index.ts index 9e3ee3c..aaa27ed 100644 --- a/frontend/src/types/index.ts +++ b/frontend/src/types/index.ts @@ -33,6 +33,7 @@ export type TableProfile = Omit< sample_data?: any[]; profile_json?: any; auto_insights?: any[]; + is_partial?: boolean; }; export type ColumnProfile = Omit & { diff --git a/worker/Dockerfile b/worker/Dockerfile new file mode 100644 index 0000000..30a9211 --- /dev/null +++ b/worker/Dockerfile @@ -0,0 +1,27 @@ +FROM python:3.12-slim + +# Install uv, git and openssh-client +COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv +RUN apt-get update && apt-get install -y --no-install-recommends git openssh-client && rm -rf /var/lib/apt/lists/* + +# Configure git to use ssh for github.com +RUN git config --global url."git@github.com:".insteadOf "https://github.com/" + +WORKDIR /app + +# Copy workspace dependencies +COPY core /app/core +COPY worker /app/worker + +WORKDIR /app/worker + +# Mount SSH deploy key secret and sync dependencies +RUN --mount=type=secret,id=deploy_key,target=/root/.ssh/id_rsa,mode=0600 \ + mkdir -p -m 0700 /root/.ssh && \ + ssh-keyscan github.com >> /root/.ssh/known_hosts && \ + uv sync --no-dev --no-install-project + +# Add virtualenv bin to PATH +ENV PATH="/app/worker/.venv/bin:$PATH" + +CMD ["python", "-m", "app.main"] diff --git a/worker/app/config.py b/worker/app/config.py new file mode 100644 index 0000000..1e7114d --- /dev/null +++ b/worker/app/config.py @@ -0,0 +1,18 @@ +from pydantic_settings import BaseSettings, SettingsConfigDict + +class Settings(BaseSettings): + model_config = SettingsConfigDict(env_file=".env", extra="ignore") + + TEMPORAL_HOST: str = "temporal:7233" + TRINO_REQUEST_TIMEOUT: float = 600.0 + PROFILER_MAX_CONCURRENT_QUERIES: int = 10 + PROFILING_CHUNK_SIZE: int = 5 + + # LLM config for generating table summaries + LLM_BASE_URL: str = "http://host.docker.internal:11434/v1" + LLM_MODEL: str = "gemma4:e4b" + LLM_API_KEY: str | None = "ollama" + + ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS: int = 600 + +settings = Settings() diff --git a/worker/app/main.py b/worker/app/main.py new file mode 100644 index 0000000..1aca070 --- /dev/null +++ b/worker/app/main.py @@ -0,0 +1,66 @@ +import asyncio +import concurrent.futures +import logging + +from temporalio.client import Client +from temporalio.worker import Worker + +from app.config import settings +from app.workflows.profiling_activities import ( + compute_chunk_metrics_activity, + fetch_table_metadata_activity, + generate_ai_summary_activity, + persist_profiling_results_activity, + profile_column_activity, +) +from app.workflows.profiling_workflow import TableProfilingWorkflow + +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s" +) +logger = logging.getLogger(__name__) + + +async def main(): + logger.info("Connecting to Temporal server at %s", settings.TEMPORAL_HOST) + + client = None + for attempt in range(1, 16): + try: + client = await Client.connect(settings.TEMPORAL_HOST) + logger.info("Connected to Temporal server successfully!") + break + except Exception as e: + logger.warning( + "Attempt %d/15: Failed to connect to Temporal server at %s: %s. Retrying in 5 seconds...", + attempt, + settings.TEMPORAL_HOST, + e + ) + await asyncio.sleep(5) + + if not client: + logger.error("Could not connect to Temporal server. Exiting.") + raise ConnectionError("Failed to connect to Temporal server") + + with concurrent.futures.ThreadPoolExecutor(max_workers=100) as activity_executor: + worker = Worker( + client, + task_queue="profiling-tasks", + workflows=[TableProfilingWorkflow], + activities=[ + fetch_table_metadata_activity, + compute_chunk_metrics_activity, + profile_column_activity, + generate_ai_summary_activity, + persist_profiling_results_activity, + ], + activity_executor=activity_executor, + ) + logger.info("Temporal profiling worker is running...") + await worker.run() + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/worker/app/workflows/__init__.py b/worker/app/workflows/__init__.py new file mode 100644 index 0000000..3212a65 --- /dev/null +++ b/worker/app/workflows/__init__.py @@ -0,0 +1 @@ +# Workflows package init diff --git a/worker/app/workflows/profiling_activities.py b/worker/app/workflows/profiling_activities.py new file mode 100644 index 0000000..3d292c4 --- /dev/null +++ b/worker/app/workflows/profiling_activities.py @@ -0,0 +1,304 @@ +import dataclasses +import logging +from datetime import datetime, timedelta + +from core.db.engine import engine +from core.models.models import ColumnProfile, ProfilingStatus, Table, TableProfile +from sqlmodel import Session, select +from temporalio import activity + +from core.services.profiling_engine import ( + ColumnStats, + TableProfilingResult, + _analyze_column, + _fetch_all, + _fetch_one, + _fqn, + _make_json_safe, + build_column_metadata_query, + build_combined_profiling_query, + build_row_count_query, + build_sample_query, + execute_query_sync, + execute_with_timeout, + generate_table_summary, + parse_combined_result, +) + +logger = logging.getLogger(__name__) + + +@activity.defn +def fetch_table_metadata_activity(payload: dict) -> dict: + table_id = payload["table_id"] + resume_from_partial = payload.get("resume_from_partial", False) + logger.info("Starting fetch_table_metadata_activity for table_id: %s (resume=%s)", table_id, resume_from_partial) + with Session(engine) as session: + table = session.get(Table, table_id) + if not table: + raise ValueError(f"Table with ID {table_id} not found") + + catalog = table.catalog + schema_name = table.schema_name + table_name = table.name + fqn = _fqn(catalog, schema_name, table_name) + + profile = session.exec( + select(TableProfile).where(TableProfile.table_id == table_id) + ).first() + if not profile: + profile = TableProfile(table_id=table_id) + + profile.status = ProfilingStatus.running + profile.updated_at = datetime.now() + session.add(profile) + session.commit() + session.refresh(profile) + profile_id = profile.id + + row_count_res = _fetch_one(build_row_count_query(fqn), table_id, default=(0,)) + row_count = int(row_count_res[0] or 0) + + columns_meta = _fetch_all(build_column_metadata_query(catalog, schema_name, table_name), table_id) + + sample_data = [] + sample_size = 0 + sample_res = execute_query_sync(build_sample_query(fqn), table_id) + if sample_res.success and sample_res.rows: + sample_size = len(sample_res.rows) + sample_data = [_make_json_safe(dict(zip(sample_res.columns, row, strict=False))) for row in sample_res.rows[:50]] + + existing_columns = [] + if resume_from_partial: + with Session(engine) as session: + existing_cols_db = session.exec( + select(ColumnProfile).where(ColumnProfile.table_id == table_id) + ).all() + for col in existing_cols_db: + existing_columns.append({ + "column_name": col.column_name, + "data_type": col.data_type, + "null_count": col.null_count, + "null_rate": col.null_rate, + "distinct_count": col.distinct_count, + "min_value": col.min_value, + "max_value": col.max_value, + "avg_value": col.avg_value, + "median_value": col.median_value, + "top_values": col.top_values, + "is_categorical": col.is_categorical, + "is_geo": col.is_geo, + "is_time": col.is_time, + "semantic_type": col.semantic_type, + "stats_json": col.stats_json, + "errors": [], + }) + + return { + "table_id": table_id, + "profile_id": profile_id, + "catalog": catalog, + "schema_name": schema_name, + "table_name": table_name, + "fqn": fqn, + "row_count": row_count, + "columns_meta": columns_meta, + "sample_data": sample_data, + "sample_size": sample_size, + "existing_columns": existing_columns, + } + + +@activity.defn +def compute_chunk_metrics_activity(payload: dict) -> dict: + fqn = payload["fqn"] + table_id = payload["table_id"] + row_count = payload["row_count"] + columns_chunk = payload["columns_chunk"] + + if row_count <= 0 or not columns_chunk: + return {} + + combined_query = build_combined_profiling_query(fqn, columns_chunk) + if not combined_query: + return {} + + logger.info("Executing chunk profiling query for table %s (%d columns)", fqn, len(columns_chunk)) + res = execute_with_timeout(combined_query, table_id) + if res.success and res.rows: + row_dict = dict(zip(res.columns, res.rows[0])) + return parse_combined_result(columns_chunk, row_dict, row_count) + else: + logger.warning("Chunk profiling query unsuccessful: %s", res.error_message) + return {} + + +@activity.defn +def profile_column_activity(payload: dict) -> dict: + col_name = payload["col_name"] + data_type = payload["data_type"] + row_count = payload["row_count"] + precomputed = payload.get("precomputed") + sample_data = payload.get("sample_data") + catalog = payload["catalog"] + schema = payload["schema_name"] + table = payload["table_name"] + table_id = payload["table_id"] + + fqn = _fqn(catalog, schema, table) + logger.info("Profiling column %s (%s)", col_name, data_type) + stats = _analyze_column(fqn, table_id, col_name, data_type, row_count, precomputed, sample_data) + + # Convert dataclass to dict + return dataclasses.asdict(stats) + + +@activity.defn +def generate_ai_summary_activity(payload: dict) -> str: + table_fqn = payload["table_fqn"] + row_count = payload["row_count"] + column_count = payload["column_count"] + column_stats_dicts = payload["column_stats"] + + # Re-construct ColumnStats objects for generate_table_summary + col_stats = [] + for cs in column_stats_dicts: + # Filter fields to match dataclass + fields = {f.name: cs[f.name] for f in dataclasses.fields(ColumnStats) if f.name in cs} + col_stats.append(ColumnStats(**fields)) + + # Re-construct TableProfilingResult + result = TableProfilingResult( + table_id=payload["table_id"], + table_fqn=table_fqn, + version=1, + computed_at=datetime.now(), + row_count=row_count, + column_count=column_count, + column_stats=col_stats, + ) + + logger.info("Generating AI summary for table %s", table_fqn) + return generate_table_summary(result) + + +@activity.defn +def persist_profiling_results_activity(payload: dict) -> None: + table_id = payload["table_id"] + profile_id = payload["profile_id"] + row_count = payload["row_count"] + sample_size = payload["sample_size"] + column_count = payload["column_count"] + sample_data = payload["sample_data"] + column_stats_dicts = payload["column_stats"] + ai_summary = payload.get("ai_summary") + is_partial = payload.get("is_partial", False) + errors = payload.get("errors", []) + + logger.info("Persisting profiling results for table_id: %s (is_partial=%s)", table_id, is_partial) + + col_stats = [] + for cs in column_stats_dicts: + fields = {f.name: cs[f.name] for f in dataclasses.fields(ColumnStats) if f.name in cs} + col_stats.append(ColumnStats(**fields)) + + # Reconstruct profile_json + insights = [f"~{row_count:,} rows (COUNT(*))."] if row_count else [] + if sample_size: + insights.append(f"{sample_size:,} rows sampled.") + + for flag, name in [("is_categorical", "categorical"), ("is_time", "Time"), ("is_geo", "Geographic")]: + cols = [c.column_name for c in col_stats if getattr(c, flag)] + if cols: + insights.append(f"{name} columns: {', '.join(cols[:5])}.") + + high_null = [c.column_name for c in col_stats if c.null_rate > 0.20] + if high_null: + insights.append(f"High null rate (>20%): {', '.join(high_null[:5])}.") + + if row_count > 0: + pk_candidates = [c.column_name for c in col_stats if c.distinct_count >= row_count * 0.95] + if pk_candidates: + insights.append(f"PK candidates: {', '.join(pk_candidates[:3])}.") + + null_rate_avg = round(sum(c.null_rate for c in col_stats) / len(col_stats), 4) if col_stats else 0.0 + + profile_json = _make_json_safe({ + "table": payload["table_fqn"], + "version": 1, + "computed_at": datetime.now().isoformat(), + "row_count": row_count, + "sample_size": sample_size, + "column_count": column_count, + "null_rate_avg": null_rate_avg, + "columns": [{ + "name": c.column_name, "data_type": c.data_type, "semantic_type": c.semantic_type, + "is_categorical": c.is_categorical, "is_large_categorical": c.is_large_categorical, + "is_free_text": c.is_free_text, "is_boolean": c.is_boolean, "is_time": c.is_time, + "is_geo": c.is_geo, "is_continuous": c.is_continuous, "distinct_count": c.distinct_count, + "null_rate": c.null_rate, "stats": c.stats_json + } for c in col_stats], + "insights": insights, + "errors": errors, + "failed_subtasks": payload.get("failed_subtasks", []), + }) + + with Session(engine) as session: + profile = session.get(TableProfile, profile_id) + if not profile: + profile = TableProfile(id=profile_id, table_id=table_id) + session.add(profile) + + profile.status = ProfilingStatus.completed + profile.is_partial = is_partial + profile.row_count = row_count + profile.sample_size = sample_size + profile.column_count = column_count + profile.null_rate_avg = null_rate_avg + profile.auto_insights = insights + profile.sample_data = sample_data + profile.profile_json = profile_json + profile.cached_until = datetime.now() + timedelta(hours=24) + profile.updated_at = datetime.now() + session.add(profile) + session.commit() + + # Clear old column profiles + old_cols = session.exec( + select(ColumnProfile).where(ColumnProfile.table_id == table_id) + ).all() + for old_c in old_cols: + session.delete(old_c) + session.commit() + + # Persist new column profiles + for cs in col_stats: + cp = ColumnProfile( + table_id=table_id, + profile_id=profile_id, + column_name=cs.column_name, + data_type=cs.data_type, + null_count=cs.null_count, + null_rate=cs.null_rate, + distinct_count=cs.distinct_count, + min_value=cs.min_value, + max_value=cs.max_value, + avg_value=cs.avg_value, + median_value=cs.median_value, + top_values=cs.top_values, + is_categorical=cs.is_categorical, + is_geo=cs.is_geo, + is_time=cs.is_time, + semantic_type=cs.semantic_type, + stats_json=_make_json_safe(cs.stats_json), + ) + session.add(cp) + session.commit() + + if ai_summary: + try: + from app.routers.profiling import _upsert_ai_summary + with Session(engine) as session: + _upsert_ai_summary(session, table_id, ai_summary) + except Exception as exc: + logger.warning("[Profiling] AI summary persist step failed: %s", exc) diff --git a/worker/app/workflows/profiling_workflow.py b/worker/app/workflows/profiling_workflow.py new file mode 100644 index 0000000..f358c62 --- /dev/null +++ b/worker/app/workflows/profiling_workflow.py @@ -0,0 +1,212 @@ +import asyncio +import logging +from datetime import timedelta + +from temporalio import workflow +from temporalio.common import RetryPolicy + +# Import activity definitions and config bypassing sandbox restrictions +with workflow.unsafe.imports_passed_through(): + from app.config import settings + from app.workflows.profiling_activities import ( + compute_chunk_metrics_activity, + fetch_table_metadata_activity, + generate_ai_summary_activity, + persist_profiling_results_activity, + profile_column_activity, + ) + +logger = logging.getLogger(__name__) + + +@workflow.defn +class TableProfilingWorkflow: + @workflow.run + async def run(self, table_id: str, resume_from_partial: bool = False) -> None: + logger.info("Running TableProfilingWorkflow for table_id: %s (resume=%s)", table_id, resume_from_partial) + + # 1. Fetch metadata activity + metadata = await workflow.execute_activity( + fetch_table_metadata_activity, + {"table_id": table_id, "resume_from_partial": resume_from_partial}, + start_to_close_timeout=timedelta(seconds=settings.ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS), + retry_policy=RetryPolicy(maximum_attempts=3), + ) + + catalog = metadata["catalog"] + schema_name = metadata["schema_name"] + table_name = metadata["table_name"] + fqn = metadata["fqn"] + row_count = metadata["row_count"] + columns_meta = metadata["columns_meta"] + sample_data = metadata["sample_data"] + sample_size = metadata["sample_size"] + profile_id = metadata["profile_id"] + existing_columns = metadata.get("existing_columns", []) + + column_stats = [] + # Immediately append previously completed columns + if existing_columns: + column_stats.extend(existing_columns) + + errors = [] + failed_subtasks = [] + is_partial = False + + existing_col_names = {c["column_name"] for c in existing_columns} + columns_meta_to_run = [c for c in columns_meta if c[0] not in existing_col_names] + + try: + # If there are no columns or table is empty, we don't profile columns + if row_count > 0 and columns_meta_to_run: + # 2. Get precomputed metrics via chunked queries + precomputed_metrics = {} + try: + chunk_size = settings.PROFILING_CHUNK_SIZE + chunks = [columns_meta_to_run[i:i + chunk_size] for i in range(0, len(columns_meta_to_run), chunk_size)] + + async def run_chunk(chunk): + return await workflow.execute_activity( + compute_chunk_metrics_activity, + { + "fqn": fqn, + "table_id": table_id, + "row_count": row_count, + "columns_chunk": chunk + }, + start_to_close_timeout=timedelta(seconds=settings.ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS), + retry_policy=RetryPolicy(maximum_attempts=3), + ) + + chunk_tasks = [run_chunk(c) for c in chunks] + chunk_results = await asyncio.gather(*chunk_tasks, return_exceptions=True) + + for i, res in enumerate(chunk_results): + if isinstance(res, BaseException): + logger.error("Failed to run profiling chunk %d: %s", i, res) + is_partial = True + errors.append(f"Chunk profiling failed: {res}") + failed_subtasks.append(f"Profiling Chunk {i}") + else: + precomputed_metrics.update(res) + except Exception as exc: + logger.error("Failed to run chunked profiling queries: %s", exc) + is_partial = True + errors.append(f"Chunked profiling setup failed: {exc}") + failed_subtasks.append("Chunked Profiling Setup") + + # Helper to run column profiling activity with error handling and retry policy + async def run_col_profile(col): + col_name, data_type = col[0], col[1] + payload = { + "col_name": col_name, + "data_type": data_type, + "row_count": row_count, + "precomputed": precomputed_metrics.get(col_name), + "sample_data": sample_data, + "catalog": catalog, + "schema_name": schema_name, + "table_name": table_name, + "table_id": table_id, + } + try: + res = await workflow.execute_activity( + profile_column_activity, + payload, + start_to_close_timeout=timedelta(seconds=settings.ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS), + retry_policy=RetryPolicy( + initial_interval=timedelta(seconds=2), + backoff_coefficient=2.0, + maximum_attempts=3, + ), + ) + return res, None + except Exception as exc: + return None, (col_name, str(exc)) + + # 3. Schedule and run all column profiling in parallel + tasks = [run_col_profile(col) for col in columns_meta_to_run] + results = await asyncio.gather(*tasks) + + # Process results of concurrent activities + for res, err in results: + if res is not None: + column_stats.append(res) + if res.get("errors"): + is_partial = True + failed_subtasks.append(f"Profile column with errors: {res.get('column_name')}") + for col_err in res.get("errors"): + errors.append(f"Column {res.get('column_name')} error: {col_err}") + elif err is not None: + col_name, error_msg = err + is_partial = True + failed_subtasks.append(f"Profile column: {col_name}") + errors.append(f"Column {col_name} failed: {error_msg}") + elif row_count == 0: + # Table is empty, build default stats for columns that aren't already existing + for col in columns_meta_to_run: + col_name, data_type = col[0], col[1] + column_stats.append({ + "column_name": col_name, + "data_type": data_type, + "null_count": 0, + "null_rate": 0.0, + "distinct_count": 0, + "semantic_type": "continuous", + "stats_json": {"type": "continuous", "distinct_count": 0, "null_rate": 0.0}, + "errors": ["Table is empty"], + }) + + # 4. Generate AI summary activity (skip if no columns completed or it's a completely failed run) + ai_summary = "" + if column_stats and not (is_partial and len(column_stats) == 0): + summary_payload = { + "table_id": table_id, + "table_fqn": fqn, + "row_count": row_count, + "column_count": len(columns_meta), + "column_stats": column_stats, + } + try: + ai_summary = await workflow.execute_activity( + generate_ai_summary_activity, + summary_payload, + start_to_close_timeout=timedelta(seconds=settings.ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS), + retry_policy=RetryPolicy(maximum_attempts=3), + ) + except Exception as exc: + logger.warning("generate_ai_summary_activity failed: %s", exc) + is_partial = True + failed_subtasks.append("Generate AI Summary") + errors.append(f"AI Summary generation failed: {exc}") + except Exception as exc: + logger.error("Workflow failed with error: %s", exc) + is_partial = True + failed_subtasks.append("Workflow Execution Failed") + errors.append(f"Workflow crashed: {exc!s}") + raise + finally: + # 5. Persist results activity + persist_payload = { + "table_id": table_id, + "profile_id": profile_id, + "table_fqn": fqn, + "row_count": row_count, + "sample_size": sample_size, + "column_count": len(columns_meta), + "sample_data": sample_data, + "column_stats": column_stats, + "ai_summary": ai_summary, + "is_partial": is_partial, + "failed_subtasks": failed_subtasks, + "errors": errors, + } + try: + await workflow.execute_activity( + persist_profiling_results_activity, + persist_payload, + start_to_close_timeout=timedelta(seconds=settings.ACTIVITY_START_TO_CLOSE_TIMEOUT_SECONDS), + retry_policy=RetryPolicy(maximum_attempts=3), + ) + except Exception as p_exc: + logger.error("Failed to persist results in finally block: %s", p_exc) diff --git a/worker/pyproject.toml b/worker/pyproject.toml new file mode 100644 index 0000000..5369d8f --- /dev/null +++ b/worker/pyproject.toml @@ -0,0 +1,19 @@ +[project] +name = "text2sql-worker" +version = "0.1.0" +description = "TextToSQL Studio Temporal Worker" +requires-python = ">=3.12" +dependencies = [ + "temporalio>=1.6.0", + "pydantic==2.10.4", + "pydantic-settings==2.7.0", + "python-dotenv==1.0.1", + "sqlmodel==0.0.22", + "trino==0.328.0", + "psycopg2-binary==2.9.9", + "pandas==2.2.3", + "core", +] + +[tool.uv.sources] +core = { path = "../core", editable = true }