@datarecce/ui 0.1.29 → 0.1.31

package/recce-source/recce/adapter/dbt_adapter/dbt_version.py

@@ -0,0 +1,42 @@
+class DbtVersion:
+
+    def __init__(self):
+        from dbt import version as dbt_version
+
+        dbt_version = self.parse(dbt_version.__version__)
+        if dbt_version.is_prerelease:
+            dbt_version = self.parse(dbt_version.base_version)
+        self.dbt_version = dbt_version
+
+    @staticmethod
+    def parse(version: str):
+        from packaging import version as v
+
+        return v.parse(version)
+
+    def as_version(self, other):
+        from packaging.version import Version
+
+        if isinstance(other, Version):
+            return other
+        if isinstance(other, str):
+            return self.parse(other)
+        return self.parse(str(other))
+
+    def __ge__(self, other):
+        return self.dbt_version >= self.as_version(other)
+
+    def __gt__(self, other):
+        return self.dbt_version > self.as_version(other)
+
+    def __lt__(self, other):
+        return self.dbt_version < self.as_version(other)
+
+    def __le__(self, other):
+        return self.dbt_version <= self.as_version(other)
+
+    def __eq__(self, other):
+        return self.dbt_version.release[:2] == self.as_version(other).release[:2]
+
+    def __str__(self):
+        return ".".join([str(x) for x in list(self.dbt_version.release)])

package/recce-source/recce/adapter/sqlmesh_adapter.py

@@ -0,0 +1,141 @@
+import typing as t
+from dataclasses import dataclass
+from typing import Dict, Optional, Type
+
+import pandas as pd
+from sqlglot import Expression, parse_one, select
+from sqlglot.optimizer import traverse_scope
+from sqlmesh.core.context import Context as SqlmeshContext
+from sqlmesh.core.environment import Environment
+from sqlmesh.core.state_sync import StateReader
+
+from recce.adapter.base import BaseAdapter
+from recce.models import RunType
+from recce.tasks import QueryDiffTask, QueryTask, RowCountDiffTask, Task
+
+sqlmesh_supported_registry: Dict[RunType, Type[Task]] = {
+    RunType.QUERY: QueryTask,
+    RunType.QUERY_DIFF: QueryDiffTask,
+    RunType.ROW_COUNT_DIFF: RowCountDiffTask,
+}
+
+
+@dataclass
+class SqlmeshAdapter(BaseAdapter):
+    context: SqlmeshContext
+    base_env: Environment
+    curr_env: Environment
+
+    def support_tasks(self):
+        return {run_type.value: True for run_type in sqlmesh_supported_registry}
+
+    def get_lineage(self, base: Optional[bool] = False):
+        state_reader: StateReader = self.context.state_reader
+        env = self.base_env if base else self.curr_env
+        nodes = {}
+        parent_map = {}
+
+        for snapshot in state_reader.get_snapshots(env.snapshots, hydrate_seeds=True).values():
+
+            if snapshot.node_type.lower() != "model":
+                continue
+
+            model = snapshot.model
+            if not model:
+                continue
+
+            node = {
+                "unique_id": model.name,
+                "name": model.name,
+                "resource_type": snapshot.node_type.lower(),
+                "checksum": {"checksum": snapshot.fingerprint.data_hash},
+            }
+
+            columns = {}
+            for column, type in model.columns_to_types.items():
+                columns[column] = {"name": column, "type": str(type)}
+            node["columns"] = columns
+
+            nodes[snapshot.name] = node
+            parents = [snapshotId.name for snapshotId in snapshot.parents]
+            parent_map[snapshot.name] = parents
+
+        manifest_metadata = {}
+        catalog_metadata = {}
+
+        return dict(
+            parent_map=parent_map,
+            nodes=nodes,
+            manifest_metadata=manifest_metadata,
+            catalog_metadata=catalog_metadata,
+        )
+
+    def get_model(self, model_id: str, base=False):
+        raise NotImplementedError()
+
+    def get_node_name_by_id(self, unique_id):
+        model = self.context.models.get(unique_id)
+        if model is None:
+            return None
+        return model.name
+
+    @classmethod
+    def load(cls, **kwargs):
+        sqlmesh_envs = kwargs.get("sqlmesh_envs")
+        if sqlmesh_envs is None:
+            raise Exception("'--sqlmesh-envs SOURCE:TARGET' is required")
+
+        envs = sqlmesh_envs.split(":")
+        if len(envs) != 2:
+            raise Exception('sqlmesh_envs must be in the format of "SOURCE:TARGET"')
+
+        sqlmesh_config = kwargs.get("sqlmesh_config", None)
+        context = SqlmeshContext(config=sqlmesh_config)
+        base_env = context.state_reader.get_environment(envs[0])
+        curr_env = context.state_reader.get_environment(envs[1])
+        if base_env is None:
+            raise Exception(f"Source environment {envs[0]} not found")
+        if curr_env is None:
+            raise Exception(f"Target environment {envs[1]} not found")
+
+        return cls(context=context, base_env=base_env, curr_env=curr_env)
+
+    def replace_virtual_tables(self, sql: t.Union[Expression, str], base: bool = None) -> Expression:
+        """
+        Replace virtual tables based on the env name.
+
+        Args:
+            sql: SQL expression to replace virtual tables
+            base: True: replace virtual tables with base env, False: replace virtual tables with current env, None: no replacement
+        """
+        if isinstance(sql, str):
+            expression = parse_one(sql, dialect=self.context.default_dialect)
+        else:
+            expression = sql
+
+        if base is not None:
+            env = self.base_env if base else self.curr_env
+            if env.name != "prod":
+                model_names = [model.name for model in self.context.models.values()]
+                for scope in traverse_scope(expression):
+                    for table in scope.tables:
+                        if f"{table.db}.{table.name}" in model_names:
+                            table.args["db"] = f"{table.args['db']}__{env.name}"
+
+        return expression
+
+    def fetchdf_with_limit(
+        self, sql: t.Union[Expression, str], base: Optional[bool] = None, limit: Optional[int] = None
+    ) -> (pd.DataFrame, bool):
+        expression = self.replace_virtual_tables(sql, base=base)
+        if limit:
+            expression = (
+                select("*").from_("__QUERY").with_("__QUERY", as_=expression).limit(limit + 1 if limit else None)
+            )
+        df = self.context.fetchdf(expression)
+        if limit and len(df) > limit:
+            df = df.head(limit)
+            more = True
+        else:
+            more = False
+        return df, more

package/recce-source/recce/apis/check_api.py

@@ -0,0 +1,203 @@
+from typing import Optional
+from uuid import UUID
+
+from fastapi import APIRouter, BackgroundTasks, HTTPException
+from pydantic import BaseModel
+
+from recce.apis.check_func import (
+    create_check_from_run,
+    create_check_without_run,
+    export_persistent_state,
+)
+from recce.apis.run_func import submit_run
+from recce.event import log_api_event
+from recce.exceptions import RecceException
+from recce.models import Check, CheckDAO, Run, RunDAO, RunType
+
+check_router = APIRouter(tags=["check"])
+
+
+class CreateCheckIn(BaseModel):
+    name: Optional[str] = None
+    description: str = ""
+    run_id: Optional[str] = None
+    type: Optional[RunType] = (None,)
+    params: Optional[dict] = None
+    view_options: Optional[dict] = None
+    track_props: Optional[dict] = None
+
+
+class CheckOut(BaseModel):
+    check_id: UUID
+    name: str
+    description: str
+    type: str
+    params: Optional[dict] = None
+    view_options: Optional[dict] = None
+    is_checked: bool = False
+    is_preset: bool = False
+    last_run: Optional[Run] = None
+
+    @classmethod
+    def from_check(cls, check: Check):
+        check_related_runs = RunDAO().list_by_check_id(check.check_id)
+        last_run = check_related_runs[-1] if len(check_related_runs) > 0 else None
+        return CheckOut(
+            check_id=check.check_id,
+            name=check.name,
+            description=check.description,
+            type=check.type.value,
+            params=check.params,
+            view_options=check.view_options,
+            is_checked=check.is_checked,
+            is_preset=check.is_preset,
+            last_run=last_run,
+        )
+
+
+@check_router.post("/checks", status_code=201, response_model=CheckOut)
+async def create_check(check_in: CreateCheckIn, background_tasks: BackgroundTasks):
+    try:
+        if check_in.run_id is not None:
+            check = create_check_from_run(
+                check_in.run_id,
+                check_name=check_in.name,
+                check_description=check_in.description,
+                check_view_options=check_in.view_options,
+            )
+        else:
+            check = create_check_without_run(
+                check_name=check_in.name,
+                check_description=check_in.description,
+                check_type=check_in.type,
+                params=check_in.params,
+                check_view_options=check_in.view_options,
+            )
+        log_api_event(
+            "create_check",
+            dict(
+                type=str(check.type),
+                track_props=check_in.track_props,
+            ),
+        )
+    except NameError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    background_tasks.add_task(export_persistent_state)
+    return CheckOut.from_check(check)
+
+
+class RunCheckIn(BaseModel):
+    nowait: Optional[bool] = False
+
+
+@check_router.post("/checks/{check_id}/run", status_code=201, response_model=Run)
+async def run_check_handler(check_id: UUID, input: RunCheckIn):
+    check = CheckDAO().find_check_by_id(check_id)
+    if not check:
+        raise HTTPException(status_code=404, detail=f"Check ID '{check_id}' not found")
+
+    try:
+        log_api_event(
+            "rerun_check",
+            dict(
+                type=str(check.type),
+                rerun=True,
+            ),
+        )
+        run, future = submit_run(check.type, check.params, check_id=check_id)
+    except RecceException as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+    if input.nowait:
+        return run
+    else:
+        run.result = await future
+        return run
+
+
+@check_router.get("/checks", status_code=200, response_model=list[CheckOut], response_model_exclude_none=True)
+async def list_checks_handler():
+    checks = []
+
+    for check in CheckDAO().list():
+        checks.append(CheckOut.from_check(check))
+
+    return checks
+
+
+@check_router.get("/checks/{check_id}", status_code=200, response_model=CheckOut)
+async def get_check_handler(check_id: UUID):
+    check = CheckDAO().find_check_by_id(check_id)
+    if check is None:
+        raise HTTPException(status_code=404, detail="Not Found")
+
+    runs = RunDAO().list_by_check_id(check_id)
+    last_run = runs[-1] if len(runs) > 0 else None
+
+    out = CheckOut.from_check(check)
+    out.last_run = last_run
+    return out
+
+
+class PatchCheckIn(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+    params: Optional[dict] = None
+    view_options: Optional[dict] = None
+    is_checked: Optional[bool] = None
+
+
+@check_router.patch("/checks/{check_id}", status_code=200, response_model=CheckOut, response_model_exclude_none=True)
+async def update_check_handler(check_id: UUID, patch: PatchCheckIn, background_tasks: BackgroundTasks):
+    check = CheckDAO().update_check_by_id(check_id, patch)
+    if check is None:
+        raise HTTPException(status_code=404, detail="Not Found")
+
+    background_tasks.add_task(export_persistent_state)
+    return CheckOut.from_check(check)
+
+
+class DeleteCheckOut(BaseModel):
+    check_id: UUID
+
+
+@check_router.delete("/checks/{check_id}", status_code=200, response_model=DeleteCheckOut)
+async def delete_handler(check_id: UUID, background_tasks: BackgroundTasks):
+    CheckDAO().delete(check_id)
+    background_tasks.add_task(export_persistent_state)
+    return DeleteCheckOut(check_id=check_id)
+
+
+class ReorderChecksIn(BaseModel):
+    source: int
+    destination: int
+
+
+@check_router.post("/checks/reorder", status_code=200)
+async def reorder_handler(order: ReorderChecksIn):
+    try:
+        CheckDAO().reorder(order.source, order.destination)
+    except RecceException as e:
+        raise HTTPException(status_code=400, detail=e.message)
+
+
+@check_router.post("/checks/{check_id}/mark-as-preset", status_code=204)
+async def mark_as_preset_check_handler(check_id: UUID, background_tasks: BackgroundTasks):
+    """
+    Mark an existing check as a preset check (cloud users only).
+
+    This creates a preset check from the specified check.
+    Only available for users with cloud mode enabled.
+
+    Returns:
+        204 No Content: Successfully marked check as preset
+        400 Bad Request: Error with detail message (e.g., not in cloud mode, check not found)
+    """
+    try:
+        CheckDAO().mark_as_preset_check(check_id)
+        background_tasks.add_task(export_persistent_state)
+    except RecceException as e:
+        raise HTTPException(status_code=400, detail=e.message)