sqlbenchdag 0.1.0__py3-none-any.whl

sql_benchmarks/resources/postgres.py

@@ -0,0 +1,188 @@
+# sql_benchmarks_dagster/resources/postgres.py (Configuration and Delegation)
+import time
+import os
+from dagster import ConfigurableResource
+from typing import Dict, Any, Optional
+import socket 
+import docker
+from docker.errors import NotFound, APIError
+from .base_engine import IBenchmarkEngine # We import it for type hinting, but don't inherit
+from .postgres_client import PostgresClient 
+from pydantic import ConfigDict, PrivateAttr
+from sqlalchemy import create_engine, text
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import OperationalError
+from ..utils.system import thrash_os_cache
+from ..constants import DATA_DIR
+
+# Inheritance is simplified to prevent MRO conflicts. It satisfies IBenchmarkEngine via Protocol.
+class PostgresEngine(ConfigurableResource): 
+    
+    # --- CONFIGURATION (Immutable) ---
+    connection_string: str
+    container_name: str = "benchmark_postgres"
+    docker_image: str = "postgres:15" 
+    model_config = ConfigDict(extra='forbid')
+    
+    _runtime_connection_string: Optional[str] = PrivateAttr(default=None)
+
+    # --- FACTORY METHOD ---
+    def _get_client(self) -> PostgresClient:
+        # Use runtime string if set (dynamic port), else config default
+        target_conn = self._runtime_connection_string or self.connection_string
+        return PostgresClient(target_conn)
+
+    # --- IBenchmarkEngine Implementation (Delegation) ---
+    def clear_cache(self, settings: dict = None):
+        """
+        Enforces a multi-layer Cold Cache:
+        1. Host Layer: Floods host RAM to evict OS Page Cache.
+        2. Database Layer: Recreates the container to wipe DBMS buffer pools.
+        """
+        thrash_os_cache()
+        self.setup_docker(settings)
+        self._wait_for_ready()
+
+    def run_query(self, sql: str, partition_key: str, engine_params: Dict[str, Any] = None) -> Optional[float]:
+        # engine_params is the 'postgres' namespace: session settings (work_mem, ...)
+        self.clear_cache(engine_params)
+        client = self._get_client()
+        return client.run_query(sql=sql, pg_settings=engine_params)
+
+    def bulk_load(self, filepath: str, table_name: str, partition_key: str, table_def: dict = None) -> None:
+        self.setup_docker()
+        self._wait_for_ready()
+        client = self._get_client()
+        client.bulk_load(filepath, table_name, partition_key, table_def)
+
+    def get_engine_name(self) -> str:
+        return "postgres"
+    
+    def get_engine(self):
+        target_conn = self._runtime_connection_string or self.connection_string
+        return create_engine(target_conn)
+
+    # --- EXTERNAL/SYSTEM/CONFIG HELPERS ---
+    def _get_port_from_url(self) -> int:
+        try:
+            target_conn = self._runtime_connection_string or self.connection_string
+            url = make_url(target_conn)
+            return url.port or 5432
+        except Exception:
+            return 5432
+
+    def _check_port_available(self, port: int) -> bool:
+        """Returns True if port is free (connect returns non-zero)."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            return s.connect_ex(('localhost', port)) != 0
+
+    def _find_free_port(self, start_port: int = 5432) -> int:
+        port = start_port
+        while True:
+            if self._check_port_available(port):
+                return port
+            port += 1
+
+    def _wait_for_ready(self, timeout=30):
+        start = time.time()
+        while time.time() - start < timeout:
+            try:
+                # Use dynamic engine creation
+                engine = self.get_engine()
+                with engine.connect() as conn:
+                    conn.execute(text("SELECT 1"))
+                return
+            except Exception:
+                time.sleep(1)
+        raise TimeoutError("Postgres container timed out during restart.")    
+
+    def _kill_zombie_container(self):
+        """
+        Forcefully removes the container using Docker SDK.
+        """
+        client = docker.from_env()
+        try:
+            container = client.containers.get(self.container_name)
+            container.remove(force=True)
+        except NotFound:
+            pass # Already gone
+        except APIError as e:
+            raise RuntimeError(f"Failed to cleanup container {self.container_name}: {e}")
+
+    def setup_docker(self, settings: dict = None):
+        """
+        Robust Provisioning using Docker SDK
+        """
+        # 1. CLEANUP
+        self._kill_zombie_container()
+
+        target_port = self._get_port_from_url()
+
+        # 2. DYNAMIC PORT ALLOCATION (Auto-Resolve Conflict)
+        # Check if the desired port is free. If not, find a new one.
+        is_free = False
+        for _ in range(5):
+             if self._check_port_available(target_port):
+                 is_free = True
+                 break
+             time.sleep(1)
+
+        if not is_free:
+            print(f"[WARN] Port {target_port} is busy. searching for free port...")
+            new_port = self._find_free_port(start_port=target_port + 1)
+            print(f"[INFO] Switched to available port: {new_port}")
+            target_port = new_port
+            
+            # CRITICAL: Update connection_string so clients connect to the new port
+            url = make_url(self.connection_string)
+            new_url = url.set(port=new_port)
+            self._runtime_connection_string = str(new_url)
+            
+        # 3. PREPARE FILESYSTEM
+        # A dedicated home for the DB files prevents "Directory not empty" errors.
+        db_storage_path = os.path.join(DATA_DIR, "postgres_db")
+        os.makedirs(db_storage_path, exist_ok=True)
+
+        # 4. BUILD & RUN CONTAINER
+        client = docker.from_env()
+        
+        # Construct config commands
+        # Postgres entrypoint interprets arguments as config flags if they start with -c
+        command_args = []
+        if settings:
+            for key, val in settings.items():
+                command_args.extend(["-c", f"{key}={val}"])
+        
+        # Retry logic for "Port is already allocated" race condition
+        max_retries = 5
+        for attempt in range(max_retries):
+            try:
+                client.containers.run(
+                    image=self.docker_image,
+                    name=self.container_name,
+                    detach=True,
+                    ports={f'5432/tcp': target_port},
+                    volumes={
+                        'pg_bench_data': {'bind': '/var/lib/postgresql/data', 'mode': 'rw'},
+                        DATA_DIR: {'bind': '/mnt/data', 'mode': 'rw'}
+                    },
+                    environment={
+                        "POSTGRES_PASSWORD": make_url(self.connection_string).password or "password",
+                        "POSTGRES_HOST_AUTH_METHOD": "trust",
+                        "POSTGRES_DB": make_url(self.connection_string).database
+                    },
+                    shm_size="2gb",
+                    command=command_args
+                )
+                break # Success
+            except APIError as e:
+                # If port is busy or conflict, wait and retry
+                if "port is already allocated" in str(e) or "Conflict" in str(e):
+                    if attempt < max_retries - 1:
+                        time.sleep(2)
+                        # Try cleanup again just in case
+                        self._kill_zombie_container()
+                        continue
+                raise RuntimeError(f"Postgres failed to start via Docker SDK: {e}")
+
+        # No wait here. Orchestration layer handles the wait.

sql_benchmarks/resources/postgres_client.py

@@ -0,0 +1,184 @@
+import io
+import time
+import polars as pl
+import pyarrow.parquet as pq
+from sqlalchemy import create_engine, text
+from typing import Dict, Any, Optional
+
+# Module-level constant so callers can identify pg-setting dimension keys
+# without importing the class or reaching into its internals.
+PG_SETTING_KEYS = frozenset({
+    "work_mem",
+    "random_page_cost",
+    "enable_hashjoin",
+    "enable_nestloop",
+    "enable_seqscan",
+    "enable_sort",
+    "enable_mergejoin",
+    "effective_cache_size",
+    "max_parallel_workers_per_gather",
+})
+
+
+class PostgresClient:
+    """
+    STATEFUL CLIENT: Restored original logic using Polars + PyArrow for robust ingestion.
+    """
+    
+    def __init__(self, connection_string: str):
+        self.connection_string = connection_string
+        self.engine = create_engine(connection_string) 
+
+    # Allowlist of Postgres settings the benchmark harness is permitted to set.
+    # Prevents arbitrary SQL injection via pg_settings YAML keys.
+    # Also exported as a module-level constant for callers that need to
+    # identify which dimension keys are pg settings (without coupling to this class).
+    _ALLOWED_PG_SETTINGS = PG_SETTING_KEYS
+
+    def run_query(self, sql: str, partition_key: str = None, pg_settings: Dict[str, Any] = None) -> Optional[float]:
+        """Execute a benchmark query, optionally applying Postgres session settings first.
+
+        partition_key is accepted but unused — present for interface symmetry with DuckDBClient.
+        pg_settings must only contain keys from PG_SETTING_KEYS (enforced below).
+        """
+        pg_settings = pg_settings or {}
+
+        with self.engine.connect() as conn:
+            for key, val in pg_settings.items():
+                if key not in PG_SETTING_KEYS:
+                    raise ValueError(f"pg_setting '{key}' is not in the allowlist.")
+                # Key is allowlisted (safe to interpolate); value is parameterized.
+                conn.execute(text(f"SET {key} = :val"), {"val": str(val)})
+
+            start = time.time()
+            conn.execute(text(sql))
+            conn.commit()
+            return time.time() - start
+
+    # Updated signature to accept partition_key (passed from factory), even if unused logic-wise
+    def bulk_load(self, file_path: str, table_name: str, partition_key: str = None, table_def: dict = None):
+        print(f"Streaming {file_path} to {table_name}...")
+
+        if file_path.endswith(".parquet"):
+            # 1. Infer Schema using Polars (Fixes UndefinedTable error)
+            df_schema = pl.scan_parquet(file_path).limit(1).collect()
+            self._create_schema(table_name, df_schema)
+            # 2. Stream Data
+            self._stream_parquet(file_path, table_name)
+
+        elif file_path.endswith(".json"):
+             self._stream_json(file_path, table_name)
+        elif file_path.endswith(".csv"):
+             self._stream_csv(file_path, table_name)
+
+        self._execute_internal(f"ANALYZE {table_name};")
+
+        # Apply declared schema constraints (PK, indexes, FKs) AFTER the bulk
+        # load — building them post-load is correct and faster than per-row.
+        # This is one-time setup cost, outside the timed query loop.
+        if table_def:
+            self._apply_constraints(table_def, table_name, partition_key)
+
+    def _apply_constraints(self, table_def: dict, table_name: str, partition_key: str):
+        from ..utils.ddl import PostgresDDLGenerator
+        gen = PostgresDDLGenerator(table_def, table_name, partition_key)
+        statements = []
+        pk_sql = gen.generate_pk_sql()
+        if pk_sql:
+            statements.append(pk_sql)
+        statements.extend(gen.generate_index_sqls())
+        statements.extend(gen.generate_fk_sqls())
+        for sql in statements:
+            self._execute_internal(sql)
+        if statements:
+            # Refresh stats so the planner can choose an index scan.
+            self._execute_internal(f"ANALYZE {table_name};")
+
+    def _create_schema(self, table_name: str, sample_df: pl.DataFrame):
+        """
+        Creates the table schema dynamically using a robust Type Map.
+        """
+        # Improved Type Map as requested
+        type_map = {
+            pl.Int8: "SMALLINT",
+            pl.Int16: "SMALLINT",
+            pl.Int32: "INTEGER",
+            pl.Int64: "BIGINT",
+            pl.Float32: "REAL",
+            pl.Float64: "DOUBLE PRECISION",
+            pl.String: "TEXT",
+            pl.Utf8: "TEXT",
+            pl.Boolean: "BOOLEAN",
+            pl.Date: "DATE",
+            pl.Datetime: "TIMESTAMP",
+            pl.Object: "TEXT"
+        }
+
+        cols = []
+        for name, dtype in sample_df.schema.items():
+            # Default to TEXT if not found, or use the mapped type
+            pg_type = type_map.get(dtype, "TEXT")
+            # Handle datetime subclasses if necessary
+            if isinstance(dtype, pl.Datetime):
+                pg_type = "TIMESTAMP"
+                
+            cols.append(f'"{name}" {pg_type}')
+
+        # Clean slate: Drop and Create
+        ddl = f"DROP TABLE IF EXISTS {table_name}; CREATE TABLE {table_name} ({', '.join(cols)});"
+        self._execute_internal(ddl)
+
+    def _stream_parquet(self, path: str, table_name: str, batch_size: int = 500_000):
+        # Use the engine's raw connection (Restored your try/finally logic)
+        raw_conn = self.engine.raw_connection()
+        try:
+            with raw_conn.cursor() as cur:
+                parquet_file = pq.ParquetFile(path)
+
+                for batch in parquet_file.iter_batches(batch_size=batch_size):
+                    df = pl.from_arrow(batch)
+                    self._copy_buffer(cur, df, table_name)
+
+            raw_conn.commit()
+        finally:
+            raw_conn.close()
+
+    def _stream_json(self, path: str, table_name: str):
+        raw_conn = self.engine.raw_connection()
+        try:
+            with raw_conn.cursor() as cur:
+                df = pl.read_ndjson(path)
+                self._create_schema(table_name, df)
+                self._copy_buffer(cur, df, table_name)
+            raw_conn.commit()
+        finally:
+            raw_conn.close()
+
+    def _stream_csv(self, path: str, table_name: str):
+        raw_conn = self.engine.raw_connection()
+        try:
+            with open(path, 'r') as f:
+                with raw_conn.cursor() as cur:
+                     cur.copy_expert(f"COPY {table_name} FROM STDIN WITH (FORMAT CSV, HEADER)", f)
+            raw_conn.commit()
+        finally:
+            raw_conn.close()
+
+    def _copy_buffer(self, cursor, df: pl.DataFrame, table_name: str):
+        """
+        High-performance buffer write using Polars CSV serializer.
+        """
+        csv_buffer = io.StringIO()
+        # write_csv is extremely fast in Polars
+        df.write_csv(csv_buffer, include_header=False)
+        csv_buffer.seek(0)
+        cursor.copy_expert(
+            f"COPY {table_name} FROM STDIN WITH (FORMAT CSV)",
+            csv_buffer
+        )
+
+    def _execute_internal(self, sql: str):
+        """Helper to execute DDL/Utility queries"""
+        with self.engine.connect() as conn:
+            conn.execute(text(sql))
+            conn.commit()

sql_benchmarks/resources/quack.py

@@ -0,0 +1,56 @@
+import os
+from typing import Dict, Any, Optional
+
+from dagster import ConfigurableResource
+from pydantic import ConfigDict
+
+from ..utils.system import thrash_os_cache
+from .quack_client import QuackClient
+
+
+class QuackEngine(ConfigurableResource):
+    """
+    DuckDB-over-Quack: the same storage and SQL dialect as the duckdb engine,
+    measured through the Quack client-server protocol (HTTP) instead of
+    in-process calls. Comparing 'duckdb' vs 'quack' on identical scenarios
+    isolates the protocol cost.
+
+    Uses its OWN data_folder so its database files never contend with the
+    in-process duckdb engine's files within the same experiment.
+    """
+    data_folder: str
+    port: int = 9494
+    token: str = "sb-local-quack-token"
+    # attach mode (False): client-side planning via USE remote.
+    # pushdown (True): SQL text shipped via remote.query() — server-side execution.
+    pushdown: bool = False
+    model_config = ConfigDict(extra='forbid')
+
+    def _get_client(self) -> QuackClient:
+        return QuackClient(data_folder=self.data_folder, port=self.port,
+                           token=self.token, pushdown=self.pushdown)
+
+    # --- IBenchmarkEngine Implementation (Delegation) ---
+
+    def get_engine_name(self) -> str:
+        return "quack"
+
+    def bulk_load(self, filepath: str, table_name: str, partition_key: str, table_def: dict = None) -> None:
+        client = self._get_client()
+        client.bulk_load(filepath, table_name, partition_key, table_def)
+
+    def run_query(self,
+                  sql: str,
+                  partition_key: str,
+                  engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """engine_params is the 'quack' namespace — reserved, not yet applied."""
+        self.clear_cache()
+        client = self._get_client()
+        return client.run_query(sql=sql,
+                                partition_key=partition_key,
+                                engine_params=engine_params)
+
+    def clear_cache(self):
+        # OS page cache flush; the per-query server restart in QuackClient
+        # provides the engine-level cold start.
+        thrash_os_cache()

sql_benchmarks/resources/quack_client.py

@@ -0,0 +1,169 @@
+"""
+QuackClient: benchmark client for DuckDB's Quack client-server protocol
+(beta since DuckDB 1.5.3).
+
+Topology per measurement (mirrors the Postgres cold-cache methodology):
+  1. Any previous server on this port is stopped (cold start guarantee).
+  2. A fresh server subprocess opens the partition's .duckdb file and serves
+     it via CALL quack_serve over HTTP, bound to localhost only.
+  3. The measurement connection ATTACHes with 'quack:localhost:<port>',
+     switches default catalog with USE, and executes the SAME SQL files as
+     the in-process duckdb engine — measured time therefore includes Quack
+     protocol serialization and HTTP transport, which is the point.
+  4. The server is stopped after the measurement.
+
+Server startup and ATTACH handshake happen BEFORE the timer starts; only
+query execution + result materialisation are measured.
+"""
+import os
+import re
+import subprocess
+import sys
+import time
+from typing import Any, Dict, Optional
+
+import duckdb
+
+from .duckdb_client import DuckDBClient
+
+# Token is interpolated into SQL on both server and client side; restrict the
+# alphabet so a misconfigured env var fails loudly instead of injecting SQL.
+_SAFE_TOKEN = re.compile(r"^[A-Za-z0-9_-]{8,}$")
+
+# Module-level server registry keyed by port (facades create a fresh client
+# per call — instance state would orphan server processes).
+_SERVERS: Dict[int, subprocess.Popen] = {}
+
+# Runs in a subprocess: open the partition db, serve it, block forever.
+# db_path and port arrive via argv; the token via env (kept out of `ps`).
+_SERVER_SCRIPT = """
+import os, sys, time
+import duckdb
+db_path, port = sys.argv[1], sys.argv[2]
+token = os.environ["SB_QUACK_TOKEN"]
+con = duckdb.connect(db_path)
+con.execute("LOAD quack;")
+con.execute("CALL quack_serve('quack:localhost:%s', token := '%s')" % (port, token))
+print("READY", flush=True)
+while True:
+    time.sleep(3600)
+"""
+
+
+class QuackClient:
+    """STATEFUL CLIENT: owns server lifecycle + measurement I/O for Quack."""
+
+    def __init__(self, data_folder: str, port: int = 9494, token: str = "",
+                 pushdown: bool = False):
+        if not _SAFE_TOKEN.match(token):
+            raise ValueError(
+                "Quack token must be >=8 chars of [A-Za-z0-9_-] "
+                "(it is interpolated into quack_serve/ATTACH statements)."
+            )
+        self.data_folder = data_folder
+        self.port = port
+        self.token = token
+        # Execution mode. attach (default): the CLIENT plans the query against
+        # the remote catalog (USE remote) — table data may stream over HTTP.
+        # pushdown: the SQL text is shipped via remote.query('...') and
+        # executes fully SERVER-side; only the result set crosses the wire.
+        # Comparing the two isolates where Quack spends its time.
+        self.pushdown = pushdown
+        # File-level operations (db path layout, parquet bulk load) are
+        # identical to the in-process duckdb engine — compose, don't copy.
+        self._duck = DuckDBClient(data_folder=data_folder)
+
+    # --- IBenchmarkEngine-facing operations -------------------------------
+
+    def bulk_load(self, filepath: str, target_table_name: str, partition_key: str, table_def: dict = None) -> None:
+        # The server holds the db file open; loading requires exclusive access.
+        self.stop_server()
+        self._duck.bulk_load(filepath, target_table_name, partition_key, table_def)
+
+    def run_query(self,
+                  sql: str,
+                  partition_key: str,
+                  engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """Cold-start a server for the partition, measure the query over Quack.
+
+        engine_params is the 'quack' namespace — reserved for server-side
+        settings; not yet applied.
+        """
+        db_path = self._duck._get_db_path(partition_key)
+        self.stop_server()
+        proc = self._start_server(db_path)
+        con = duckdb.connect()
+        try:
+            con.execute("LOAD quack;")
+            self._attach_with_retry(con, proc)
+
+            if self.pushdown:
+                # Ship the SQL text; execution happens entirely server-side.
+                # Single-quote doubling is the only escaping SQL strings need.
+                wrapped = "FROM remote.query('{}')".format(sql.replace("'", "''"))
+                start = time.time()
+                con.sql(wrapped).fetchall()
+                end = time.time()
+            else:
+                con.execute("USE remote")
+                start = time.time()
+                con.sql(sql).fetchall()
+                end = time.time()
+            return end - start
+        except duckdb.NotImplementedException as e:
+            # A protocol capability gap, not a config error: e.g. attach mode
+            # cannot run multi-table joins in the Quack beta ("Multiple
+            # streaming scans ... not currently supported"). Record as DNF —
+            # the limitation IS the measurement. Anything else still raises.
+            print(f"[Quack] DNF — protocol limitation: {e}")
+            return None
+        finally:
+            con.close()
+            self.stop_server()
+
+    # --- Server lifecycle --------------------------------------------------
+
+    def stop_server(self) -> None:
+        proc = _SERVERS.pop(self.port, None)
+        if proc is None:
+            return
+        proc.terminate()
+        try:
+            proc.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+            proc.wait()
+
+    def _start_server(self, db_path: str) -> subprocess.Popen:
+        env = {**os.environ, "SB_QUACK_TOKEN": self.token}
+        proc = subprocess.Popen(
+            [sys.executable, "-c", _SERVER_SCRIPT, db_path, str(self.port)],
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+        )
+        _SERVERS[self.port] = proc
+        return proc
+
+    def _attach_with_retry(self, con, proc: subprocess.Popen, timeout_s: float = 15.0) -> None:
+        """ATTACH doubles as the readiness probe. Fails loudly with server stderr."""
+        uri = f"quack:localhost:{self.port}"
+        deadline = time.time() + timeout_s
+        last_err = None
+        while time.time() < deadline:
+            if proc.poll() is not None:
+                _, stderr = proc.communicate()
+                raise RuntimeError(
+                    f"Quack server exited (code {proc.returncode}) before serving "
+                    f"{uri}. Stderr:\n{stderr}"
+                )
+            try:
+                con.execute(
+                    f"ATTACH '{uri}' AS remote (TOKEN '{self.token}', DISABLE_SSL true)"
+                )
+                return
+            except duckdb.Error as e:
+                last_err = e
+                time.sleep(0.2)
+        raise RuntimeError(f"Quack server at {uri} not ready after {timeout_s}s: {last_err}")