sqlbenchdag 0.1.0__py3-none-any.whl

sql_benchmarks/api/models/experiments.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel
+from typing import Optional
+
+
+class ExperimentSubmitRequest(BaseModel):
+    config_yaml: str
+
+
+class ExperimentSubmitResponse(BaseModel):
+    experiment_id: str
+    status: str  # "queued" | "duplicate" | "rejected"
+    detail: Optional[str] = None
+
+
+class ExperimentStatus(BaseModel):
+    experiment_id: str
+    status: str  # "queued" | "running" | "complete" | "not_found"
+    has_results: bool
+    fragment_count: int
+    has_csv: bool

sql_benchmarks/api/models/recommend.py

@@ -0,0 +1,11 @@
+from pydantic import BaseModel
+from typing import Dict, List
+
+
+class RecommendResponse(BaseModel):
+    recommended_engine: str
+    confidence: str  # "high" | "medium" | "low"
+    reasoning: str
+    supporting_experiments: List[str]
+    engine_scores: Dict[str, float]  # {engine: mean_duration_seconds}
+    caveats: List[str]

sql_benchmarks/api/models/results.py

@@ -0,0 +1,66 @@
+from pydantic import BaseModel
+from typing import Any, Dict, List, Optional
+
+
+class FragmentMeta(BaseModel):
+    timestamp: str
+    experiment_id: str
+    dagster_run_id: str
+    engine: str
+    asset: str
+    partition: str
+
+
+class FragmentMetrics(BaseModel):
+    duration_seconds: float
+    replication_factor: int
+    # Raw per-replication measurements; None for fragments written before
+    # raw capture (or DNF sentinels, which carry an empty list).
+    durations_raw: Optional[List[float]] = None
+
+
+class Fragment(BaseModel):
+    meta: FragmentMeta
+    metrics: FragmentMetrics
+    parameters: Dict[str, Any]
+
+
+class ExperimentSummary(BaseModel):
+    experiment_id: str
+    suite: Optional[str] = None
+    engines: List[str]
+    partition_count: int
+    fragment_count: int
+    has_csv: bool
+    has_dashboard: bool
+    created_at: Optional[float] = None  # unix timestamp from metadata
+
+
+class ExperimentResult(BaseModel):
+    experiment_id: str
+    config: Optional[Dict[str, Any]] = None
+    summary: ExperimentSummary
+    fragments: List[Fragment]
+
+
+class EngineRanking(BaseModel):
+    engine: str
+    mean_duration_seconds: float
+    median_duration_seconds: float
+    p95_duration_seconds: float
+    sample_count: int
+    rank: int
+
+
+class CompareResult(BaseModel):
+    experiment_id: str
+    suite: Optional[str] = None
+    partition: Optional[str] = None
+    rankings: List[EngineRanking]
+    winner: str
+    speedup_vs_slowest: float
+
+
+class ResultsListResponse(BaseModel):
+    experiments: List[ExperimentSummary]
+    total: int

sql_benchmarks/api/routers/catalog.py

@@ -0,0 +1,19 @@
+from fastapi import APIRouter
+
+from ..data.catalog_reader import CatalogReader
+from ..models.catalog import CatalogEnginesResponse, CatalogSuitesResponse
+
+router = APIRouter(prefix="/v1/catalog", tags=["catalog"])
+_reader = CatalogReader()
+
+
+@router.get("/engines", response_model=CatalogEnginesResponse)
+def list_engines():
+    """List available database engines and the test suites they support."""
+    return _reader.get_engines_response()
+
+
+@router.get("/suites", response_model=CatalogSuitesResponse)
+def list_suites():
+    """List all benchmark test suites with their SQL content per engine."""
+    return _reader.get_suites_response()

sql_benchmarks/api/routers/experiments.py

@@ -0,0 +1,79 @@
+import os
+
+import yaml
+from fastapi import APIRouter, BackgroundTasks, HTTPException
+
+from ...constants import CONFIG_ARCHIVE_DIR, EXPERIMENTS_DIR, ROOT_DIR
+from ...coordinator import ExperimentCoordinator
+from ...utils.hasher import generate_experiment_hash
+from ...validator import ExperimentValidator
+from ..data.reader import ResultReader
+from ..models.experiments import ExperimentStatus, ExperimentSubmitRequest, ExperimentSubmitResponse
+
+router = APIRouter(prefix="/v1/experiments", tags=["experiments"])
+_reader = ResultReader()
+
+
+def _run_experiment(yaml_path: str):
+    coordinator = ExperimentCoordinator(yaml_path, headless=True)
+    coordinator.run()
+
+
+@router.post("", response_model=ExperimentSubmitResponse, status_code=202)
+def submit_experiment(body: ExperimentSubmitRequest, background_tasks: BackgroundTasks):
+    """
+    Submit a new benchmark experiment as a YAML config string.
+    Returns immediately with the experiment ID. Use /status to poll progress.
+    """
+    try:
+        config = yaml.safe_load(body.config_yaml)
+    except yaml.YAMLError as e:
+        raise HTTPException(status_code=422, detail=f"Invalid YAML: {e}")
+
+    try:
+        ExperimentValidator.validate(config, source_label="api_submission")
+    except ValueError as e:
+        raise HTTPException(status_code=422, detail=str(e))
+
+    exp_id = generate_experiment_hash(config, ROOT_DIR)
+
+    if os.path.exists(os.path.join(CONFIG_ARCHIVE_DIR, f"config_{exp_id}.yaml")):
+        return ExperimentSubmitResponse(
+            experiment_id=exp_id,
+            status="duplicate",
+            detail="Results already exist for this experiment. Retrieve them at /v1/results/{exp_id}",
+        )
+
+    queue_dir = os.path.join(EXPERIMENTS_DIR, "queue")
+    os.makedirs(queue_dir, exist_ok=True)
+    yaml_path = os.path.join(queue_dir, f"{exp_id}.yaml")
+    with open(yaml_path, "w") as f:
+        yaml.dump(config, f, sort_keys=False)
+
+    background_tasks.add_task(_run_experiment, yaml_path)
+
+    return ExperimentSubmitResponse(experiment_id=exp_id, status="queued")
+
+
+@router.get("/{exp_id}/status", response_model=ExperimentStatus)
+def get_status(exp_id: str):
+    """Check the status of a submitted experiment."""
+    fragments = _reader.get_fragments(exp_id)
+    has_csv = _reader.has_csv(exp_id)
+
+    if _reader.is_complete(exp_id):
+        status = "complete"
+    elif _reader.results_exist(exp_id):
+        status = "running"
+    elif _reader.is_queued(exp_id):
+        status = "queued"
+    else:
+        status = "not_found"
+
+    return ExperimentStatus(
+        experiment_id=exp_id,
+        status=status,
+        has_results=_reader.results_exist(exp_id),
+        fragment_count=len(fragments),
+        has_csv=has_csv,
+    )

sql_benchmarks/api/routers/recommend.py

@@ -0,0 +1,22 @@
+from typing import Optional
+
+from fastapi import APIRouter, Query
+
+from ..data.reader import ResultReader
+from ..logic.ranker import recommend_engine
+from ..models.recommend import RecommendResponse
+
+router = APIRouter(prefix="/v1", tags=["recommend"])
+_reader = ResultReader()
+
+
+@router.get("/recommend", response_model=RecommendResponse)
+def recommend(
+    suite: Optional[str] = Query(None, description="Test suite name (e.g. analytical_wall)"),
+    scale: Optional[str] = Query(None, description="Partition key substring to filter by scale (e.g. 'large')"),
+):
+    """
+    Get an engine recommendation based on pre-computed benchmark data.
+    Returns the fastest engine for the given suite and scale, with confidence and reasoning.
+    """
+    return recommend_engine(_reader, suite=suite, scale=scale)

sql_benchmarks/api/routers/results.py

@@ -0,0 +1,51 @@
+from typing import Optional
+
+from fastapi import APIRouter, HTTPException, Query
+
+from ..data.reader import ResultReader
+from ..logic.comparator import compare_experiment
+from ..models.results import CompareResult, ExperimentResult, ResultsListResponse
+
+router = APIRouter(prefix="/v1/results", tags=["results"])
+_reader = ResultReader()
+
+
+@router.get("", response_model=ResultsListResponse)
+def list_results(
+    suite: Optional[str] = Query(None, description="Filter by test suite name"),
+    engine: Optional[str] = Query(None, description="Filter by engine name"),
+    limit: int = Query(50, ge=1, le=500),
+    offset: int = Query(0, ge=0),
+):
+    """List completed benchmark experiments, optionally filtered by suite or engine."""
+    summaries = _reader.filter_experiments(suite=suite, engine=engine)
+    summaries.sort(key=lambda s: s.created_at or 0, reverse=True)
+    page = summaries[offset: offset + limit]
+    return ResultsListResponse(experiments=page, total=len(summaries))
+
+
+@router.get("/{exp_id}", response_model=ExperimentResult)
+def get_result(exp_id: str):
+    """Get full benchmark results for a specific experiment ID."""
+    if not _reader.results_exist(exp_id):
+        raise HTTPException(status_code=404, detail=f"Experiment '{exp_id}' not found")
+    summary = _reader.build_summary(exp_id)
+    fragments = _reader.get_fragments(exp_id)
+    config = _reader.get_config(exp_id)
+    return ExperimentResult(
+        experiment_id=exp_id,
+        config=config,
+        summary=summary,
+        fragments=fragments,
+    )
+
+
+@router.get("/{exp_id}/compare", response_model=CompareResult)
+def compare_result(
+    exp_id: str,
+    partition: Optional[str] = Query(None, description="Filter to a specific partition key"),
+):
+    """Get a ranked cross-engine performance comparison for an experiment."""
+    if not _reader.results_exist(exp_id):
+        raise HTTPException(status_code=404, detail=f"Experiment '{exp_id}' not found")
+    return compare_experiment(exp_id, _reader, partition=partition)

sql_benchmarks/assets/__init__.py

@@ -0,0 +1,4 @@
+from . import data_factory
+from . import ingestion_factory   
+from . import benchmark_factory
+from . import data_quality

sql_benchmarks/assets/benchmark_factory.py

@@ -0,0 +1,222 @@
+import json
+import datetime
+import os
+import glob
+import time
+import jinja2
+import statistics
+from dagster import asset, AssetExecutionContext, MaterializeResult, MetadataValue
+from ..partitions import partitions_def, SCENARIO_CONFIG
+from ..constants import RESULTS_DIR
+from ..utils.common import load_context, get_tables_used_in_sql, get_target_sql_dir, infer_metadata_from_sql, get_engine_asset_prefix, get_engine_sql_dialect, get_engine_benchmark_group, get_scoped_asset_name
+
+CTX = load_context()
+ACTIVE_ENGINES = CTX['engines']
+EXPERIMENT_META = CTX['meta']
+EXP_ID = EXPERIMENT_META.get("experiment_id", "unknown")
+VALID_TABLES = set(CTX['tables'])
+FULL_CONFIG = CTX['full_config']
+REPLICATION_FACTOR = FULL_CONFIG.get("execution", {}).get("replication", 1)
+
+def _smart_cast(val):
+    if isinstance(val, (int, float, bool)): return val
+    return str(val)
+
+def write_benchmark_fragment(experiment_id, run_id, engine, asset_name, pk, durations, params):
+    """
+    Writes the atomic result fragment to disk. 
+    Isolates the 'Scientific Proof' logic from the Dagster asset.
+    """
+    # Results are isolated by experiment_id in RESULTS_DIR
+    fragment_path = os.path.join(
+        RESULTS_DIR, 
+        experiment_id,
+        "fragments",
+        f"{asset_name}__{pk}.json"
+    )
+    
+    os.makedirs(os.path.dirname(fragment_path), exist_ok=True)
+    
+    payload = {
+        "meta": {
+            "timestamp": datetime.datetime.now().isoformat(),
+            "experiment_id": experiment_id,
+            "dagster_run_id": run_id,
+            "engine": engine,
+            "asset": asset_name,
+            "partition": pk
+        },
+        "metrics": {
+            "duration_seconds": statistics.mean(durations),
+            # Raw per-replication measurements: the spread under identical
+            # conditions is itself evidence (thermal drift, scheduler noise,
+            # cold-cache approximation). Mean alone hides it.
+            "durations_raw": durations,
+            "replication_factor": REPLICATION_FACTOR
+        },
+        "parameters": params # The "Jagged" Context
+    }
+    
+    with open(fragment_path, "w") as f:
+        json.dump(payload, f, default=str, indent=2)
+        
+    return fragment_path
+
+def make_benchmark_asset(name, engine, used_tables, raw_template, static_meta, extra_context=None):
+    prefix = get_engine_asset_prefix(engine)
+    deps = [get_scoped_asset_name(f"{prefix}{t}_table", EXP_ID) for t in used_tables]
+    asset_base_name = f"{prefix}benchmark_{name}"
+    asset_scoped_name = get_scoped_asset_name(asset_base_name, EXP_ID)
+    tags = {}
+
+    # Single-Lane Limit for engines with exclusive infrastructure:
+    # postgres restarts a shared Docker container; quack variants bind fixed ports.
+    if engine in ("postgres", "quack", "quack_pushdown"):
+        tags["dagster/concurrency_key"] = f"{engine}_exclusive"
+    tags["experiment_scope"] = "partitioned"
+
+    @asset(
+        name=asset_scoped_name,
+        partitions_def=partitions_def,
+        deps=deps,
+        group_name=get_engine_benchmark_group(engine),
+        required_resource_keys={engine},
+        op_tags=tags
+    )
+    def _benchmark(context: AssetExecutionContext):
+        db = getattr(context.resources, engine)
+        pk = context.partition_key
+        params = SCENARIO_CONFIG.get(pk, {})
+        # Each engine receives ONLY its own namespace of engine_params
+        # (assembled by config_loader from execution.engine_params + namespaced
+        # matrix dimensions like 'postgres.work_mem').
+        engine_params = params.get("engine_params", {}).get(engine, {})
+        dims = {k: v for k, v in params.items() if k != "engine_params"}
+
+        # SQL render — dims feeds template variables; engine_params stays out of it
+        render_ctx = {f"{t}_table": f"{t}_{pk}" for t in used_tables}
+        render_ctx.update(dims)
+        sql = jinja2.Template(raw_template).render(render_ctx)
+
+        # 3. Execution (Replicated)
+        # None return means the engine signalled a non-fatal failure (e.g. TypeDB
+        # stack overflow on recursive queries).  We stop after the first None and
+        # record the run as DNF rather than crashing the entire Dagster step.
+        durations = []
+        dnf = False
+        for _ in range(REPLICATION_FACTOR):
+            duration = db.run_query(sql=sql, partition_key=pk, engine_params=engine_params)
+            if duration is None:
+                dnf = True
+                break
+            durations.append(duration)
+
+        if dnf:
+            context.log.warning(
+                f"Engine '{engine}' returned None for partition '{pk}' — "
+                f"recording as DNF (did-not-finish). "
+                f"Likely cause: server crash / stack overflow / OOM during query evaluation."
+            )
+            # Write a sentinel fragment directly (bypass write_benchmark_fragment
+            # which requires at least one duration measurement)
+            experiment_id = EXPERIMENT_META.get("experiment_id", "unknown")
+            fragment_path = os.path.join(
+                RESULTS_DIR, experiment_id, "fragments",
+                f"{asset_scoped_name}__{pk}.json"
+            )
+            os.makedirs(os.path.dirname(fragment_path), exist_ok=True)
+            import json as _json, datetime as _dt
+            _json.dump({
+                "meta": {
+                    "timestamp": _dt.datetime.now().isoformat(),
+                    "experiment_id": experiment_id,
+                    "dagster_run_id": context.run.run_id,
+                    "engine": engine,
+                    "asset": asset_scoped_name,
+                    "partition": pk,
+                },
+                "metrics": {"duration_seconds": None, "durations_raw": [], "replication_factor": 0, "dnf": True},
+                # dnf lives in metrics; duplicating it into parameters leaked a
+                # redundant lowercase 'dnf' column into the flattened CSV.
+                "parameters": params,
+            }, open(fragment_path, "w"), default=str, indent=2)
+            return MaterializeResult(metadata={
+                "dnf": MetadataValue.bool(True),
+                "engine": MetadataValue.text(engine),
+                "partition": MetadataValue.text(pk),
+            })
+
+        # 4. Write Fragment (The Isolated Call)
+        experiment_id = EXPERIMENT_META.get("experiment_id", "unknown")
+
+        write_benchmark_fragment(
+            experiment_id=experiment_id,
+            run_id=context.run.run_id,
+            engine=engine,
+            asset_name=asset_scoped_name,
+            pk=pk,
+            durations=durations,
+            params=dims,
+        )
+
+        meta = {
+            "duration": MetadataValue.float(statistics.mean(durations)),
+            "sql": MetadataValue.md(f"```sql\n{sql}\n```"),
+            "experiment_id": experiment_id,
+            "config_engine": engine,
+            **{k: _smart_cast(v) for k, v in static_meta.items()},
+            **{f"dim_{k}": _smart_cast(v) for k, v in dims.items()},
+        }
+        return MaterializeResult(metadata=meta)
+
+    return _benchmark
+
+def get_benchmark_assets():
+    assets = []
+    
+    if not CTX: 
+        return [] 
+
+    ACTIVE_ENGINES = CTX.get('engines')
+    if not ACTIVE_ENGINES:
+        return []
+        
+    FULL_CONFIG = CTX.get('full_config', {})
+    dataset_cfg = CTX.get('dataset_config', {})
+    VALID_TABLES = CTX.get('tables', set())
+    
+    target_dir = get_target_sql_dir(FULL_CONFIG)
+
+    for engine in ACTIVE_ENGINES:
+        path = os.path.join(target_dir, get_engine_sql_dialect(engine))
+        if not os.path.exists(path): continue
+
+        for f in glob.glob(os.path.join(path, "*.sql")):
+            if os.path.getsize(f) == 0:
+                print(f"[WARN] Skipping empty benchmark file: {f}")
+                continue
+
+            base = os.path.basename(f).replace(".sql", "")
+            tables, raw = get_tables_used_in_sql(f, VALID_TABLES)
+            static_meta = infer_metadata_from_sql(raw, dataset_cfg)
+
+            # Context Injection: Enable {{ column_name }} in SQL
+            col_ctx = {}
+            if dataset_cfg and 'tables' in dataset_cfg:
+                for t in tables:
+                     t_def = dataset_cfg['tables'].get(t, {})
+                     for c in t_def.get('columns', []):
+                         col_ctx[c['name']] = c['name']
+
+            asset_wrapper = make_benchmark_asset(base, engine, tables, raw, static_meta, extra_context=col_ctx)
+            
+            try:
+                asset_obj = asset_wrapper.to_asset_def()
+            except AttributeError:
+                asset_obj = asset_wrapper
+                
+            assets.append(asset_obj)
+            
+    return assets
+
+benchmark_assets = get_benchmark_assets()

sql_benchmarks/assets/data_factory.py

@@ -0,0 +1,54 @@
+import importlib
+import os
+from dagster import asset, AssetExecutionContext
+from ..partitions import partitions_def, SCENARIO_CONFIG
+from ..constants import DATA_DIR
+from ..utils.common import load_context, get_data_dependencies, get_scoped_asset_name
+
+from ..plugins.data_sources import declarative_gen
+
+# 1. LOAD CONTEXT
+CTX = load_context()
+EXP_ID = CTX['meta'].get("experiment_id", "unknown")
+
+OUTPUT_DIR = os.path.join(DATA_DIR, "staging")
+
+def make_data_asset(table_name):
+
+    raw_deps = get_data_dependencies(table_name, CTX['table_defs'])
+    asset_deps = [get_scoped_asset_name(f"{d}_parquet", EXP_ID) for d in raw_deps]
+    scoped_name = get_scoped_asset_name(f"{table_name}_parquet", EXP_ID)
+
+    @asset(
+        name=scoped_name,
+        partitions_def=partitions_def,
+        group_name="data_generation",
+        deps=asset_deps,
+        description=f"Generates {table_name}"
+    )
+    def _generator(context: AssetExecutionContext):
+        partition_key = context.partition_key
+        params = SCENARIO_CONFIG[partition_key]
+        
+        # Load Plugin from Config
+        plugin_name = CTX['dataset_config']['source']
+        module = importlib.import_module(plugin_name)
+        
+        # 1. Construct the path
+        file_name = f"{table_name}_{partition_key}.parquet"
+        target_path = os.path.join(OUTPUT_DIR, file_name)
+        
+        # 2. Ensure folder exists
+        os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+        # 3. Call Plugin
+        return module.generate(
+            context=context, 
+            params=params, 
+            table_name=table_name, 
+            target_path=target_path, 
+            dataset_config=CTX['dataset_config']
+        )
+    return _generator
+
+data_assets = [make_data_asset(t) for t in CTX['tables']]

sql_benchmarks/assets/data_quality.py

@@ -0,0 +1,91 @@
+from dagster import asset, AssetExecutionContext, MetadataValue, Output
+from typing import List
+import polars as pl
+import json
+import os
+
+from ..constants import DATA_DIR, RESULTS_DIR
+from ..utils.common import load_context, get_scoped_asset_name
+from ..partitions import partitions_def
+
+# Resolved inside assets to stay dynamic
+
+CTX = load_context()
+EXP_ID = CTX['meta'].get("experiment_id", "unknown")
+
+quality_assets: List[object] = []
+
+def make_quality_asset(table_name):
+    """
+    Creates a validation asset for a specific table.
+    Reads the staging Parquet file and verifies stats.
+    """
+    scoped_name = get_scoped_asset_name(f"{table_name}_quality", EXP_ID)
+    deps = [get_scoped_asset_name(f"{table_name}_parquet", EXP_ID)]
+
+    @asset(
+        name=scoped_name,
+        group_name="data_generation",
+        partitions_def=partitions_def,
+        deps=deps,
+        description=f"Validates {table_name} distribution."
+    )
+    def _validate(context: AssetExecutionContext):
+        partition_key = context.partition_key
+        
+        # Resolve EXP_ID inside the function for dynamic behavior (essential for testing)
+        ctx = load_context()
+        current_exp_id = ctx['meta'].get("experiment_id", "unknown")
+        
+        # 2. LOAD DATA
+        filename = f"{table_name}_{partition_key}.parquet"
+        parquet_path = os.path.join(DATA_DIR, "staging", filename)
+        
+        if not os.path.exists(parquet_path):
+            raise FileNotFoundError(f"Staging file not found: {parquet_path}")
+
+        # Read Data
+        df = pl.read_parquet(parquet_path)
+        row_count = df.height
+        
+        stats = {
+            "rows": row_count,
+            "columns": {}
+        }
+        
+        # Compute Stats
+        for col in df.columns:
+            null_count = df[col].null_count()
+            col_stats = {
+                "null_count": null_count,
+                "null_percent": null_count / row_count if row_count > 0 else 0,
+                "cardinality": df[col].n_unique(),
+                "dtype": str(df[col].dtype)
+            }
+            stats["columns"][col] = col_stats
+            
+        if row_count == 0:
+             raise ValueError(f"Table {table_name} is empty!")
+
+        # Results are isolated by current_exp_id in RESULTS_DIR
+        target_path = os.path.join(RESULTS_DIR, current_exp_id, "data_stats", f"{table_name}_{partition_key}.stats.json")
+        stats_dir = os.path.dirname(target_path)
+        os.makedirs(stats_dir, exist_ok=True)
+        
+        with open(target_path, "w") as f:
+            json.dump(stats, f, indent=2)
+
+        return Output(
+            value=target_path,
+            metadata={
+                "stats": MetadataValue.json(stats),
+                "profile_path": MetadataValue.path(target_path),
+                "archive_location": MetadataValue.path(stats_dir)
+            }
+        )
+    return _validate
+
+# Generate for all tables
+if CTX:
+    for t in CTX.get('tables', []):
+        quality_assets.append(make_quality_asset(t))