sqlbenchdag 0.1.0__py3-none-any.whl

sql_benchmarks/harness.py

@@ -0,0 +1,84 @@
+import os
+import shutil
+import tempfile
+import time
+from .constants import ROOT_DIR, DATA_DIR, RESULTS_DIR, VIOLATIONS_DIR, REPORTS_DIR
+from .utils.integrity_monitor import IntegrityMonitor
+
+class IsolationHarness:
+    """
+    Provides a 'Clean Room' for experiment execution.
+    Redirects all file outputs to a scratchpad to prevent state contamination.
+    """
+    
+    def __init__(self, experiment_id: str):
+        self.experiment_id = experiment_id
+        self.scratchpad_root = None
+        self._monitor = None
+        
+    def provision(self) -> dict:
+        """Creates the scratchpad and returns the redirection environment."""
+        self.scratchpad_root = tempfile.mkdtemp(prefix=f"sb_{self.experiment_id}_")
+        
+        # Create standard layout inside scratchpad
+        paths = {
+            "SB_DATA_DIR": os.path.join(self.scratchpad_root, "data"),
+            "SB_RESULTS_DIR": os.path.join(self.scratchpad_root, "results"),
+            "SB_VIOLATIONS_DIR": os.path.join(self.scratchpad_root, "violations"),
+            "SB_REPORTS_DIR": os.path.join(self.scratchpad_root, "reports"),
+            "SCRATCHPAD_ROOT": self.scratchpad_root
+        }
+        
+        for p in paths.values():
+            os.makedirs(p, exist_ok=True)
+            
+        # Ensure critical subdirectories exist
+        os.makedirs(os.path.join(paths["SB_DATA_DIR"], "staging"), exist_ok=True)
+        os.makedirs(os.path.join(paths["SB_DATA_DIR"], "duckdb"), exist_ok=True)
+        os.makedirs(os.path.join(paths["SB_RESULTS_DIR"], "fragments"), exist_ok=True)
+
+        # Snapshot the package source NOW: if any code changes while the
+        # experiment runs, the Experiment ID no longer describes what ran.
+        from .constants import PACKAGE_DIR
+        self._monitor = IntegrityMonitor(PACKAGE_DIR)
+
+        return paths
+
+    def check_integrity(self) -> list:
+        """
+        Compares the package source against the snapshot taken at provision().
+        Returns drift entries (MODIFIED/ADDED/DELETED) — any code change during
+        an experiment invalidates the run, because the Experiment ID was
+        computed from the code as it stood at submission.
+
+        History note: a previous implementation here only inspected a
+        sentinel file named secure_pill.txt (planted by its own test) and
+        could never detect real tampering. Replaced with the IntegrityMonitor
+        snapshot mechanism this docstring's promises always implied.
+        """
+        if self._monitor is None:
+            return []
+        return self._monitor.check_drift()
+
+    def cleanup(self):
+        """Removes the scratchpad — and ONLY a scratchpad.
+
+        Guarded after a historical incident: a miscomputed path turned a
+        cleanup into a recursive deletion of real work ("brilliant
+        amnesia"). rmtree here refuses any path that is not a genuine
+        sb_-prefixed directory under the system temp dir. Refusal is loud;
+        a leaked scratchpad is recoverable, a deleted repo is not.
+        """
+        if not self.scratchpad_root:
+            return
+        real = os.path.realpath(self.scratchpad_root)
+        tmp_root = os.path.realpath(tempfile.gettempdir())
+        if not (real.startswith(tmp_root + os.sep)
+                and os.path.basename(real).startswith("sb_")):
+            print(f"[CRITICAL] cleanup() refused: '{real}' is not a temp scratchpad. "
+                  f"Leaving it untouched.")
+            return
+        if os.path.exists(real):
+            # Wait a beat for any lagging file handles
+            time.sleep(0.1)
+            shutil.rmtree(real)

sql_benchmarks/jobs.py

@@ -0,0 +1,8 @@
+from dagster import define_asset_job, AssetSelection
+
+# Defines a job that selects ALL assets (Ingestion -> Benchmarks)
+# This is what the Sensor triggers when it sees a new YAML file.
+benchmark_job = define_asset_job(
+    name="benchmark_job",
+    selection=AssetSelection.all()
+)

sql_benchmarks/partitions.py

@@ -0,0 +1,17 @@
+from dagster import StaticPartitionsDefinition
+from .config_loader import ConfigLoader 
+
+# Initialize the Configuration Compiler
+try:
+    COMPILER = ConfigLoader()
+except ValueError as e:
+    # If the compiler raises a structural error, we raise it here to fail hard at load time.
+    raise e
+
+# Extract the compiled artifacts for Dagster
+partitions_def = StaticPartitionsDefinition(COMPILER.partition_keys)
+SCENARIO_CONFIG = COMPILER.scenario_config
+
+# If the config file was missing, we must ensure Dagster boots cleanly
+if not COMPILER.partition_keys:
+    partitions_def = StaticPartitionsDefinition(["init"])

sql_benchmarks/plugins/data_sources/declarative_gen.py

@@ -0,0 +1,211 @@
+import os
+import polars as pl
+import numpy as np
+from dagster import MaterializeResult, MetadataValue
+from ...utils.schema import TableDef  # Strict Pydantic Model
+from ...utils.providers import PROVIDER_REGISTRY
+from ...constants import DEFAULT_CHUNK_SIZE
+
+
+def resolve_params_recursive(obj, params):
+    """Recursively replace string values in obj that match keys in params."""
+    if isinstance(obj, dict):
+        return {k: resolve_params_recursive(v, params) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [resolve_params_recursive(i, params) for i in obj]
+    elif isinstance(obj, str) and obj in params:
+        return params[obj]
+    # Handle Pydantic Models by dumping them first (if passed directly)
+    elif hasattr(obj, "model_dump"):
+        return resolve_params_recursive(obj.model_dump(), params)
+    return obj
+
+def generate(context, params, table_name, target_path, dataset_config):
+    
+    # 1. SCHEMA VALIDATION
+    if 'tables' not in dataset_config:
+        raise ValueError("Missing 'tables' section in dataset config.")
+        
+    dataset_config = resolve_params_recursive(dataset_config, params)
+    
+    tables = dataset_config.get("tables", {})
+    raw_table_def = tables.get(table_name)
+    if not raw_table_def:
+        raise ValueError(f"Table '{table_name}' not defined in dataset config.")
+
+    # Validate using Pydantic (Throws ValidationError if invalid)
+    # We strip unknown fields if strict mode is issue, but Schema is set to 'allow' extra.
+    table_model = TableDef(**tables[table_name])
+
+    # 2. RESOLVE ROW COUNT
+    # Priority: Matrix Params > Table Config
+    base_rows = params.get('rows')
+    if base_rows is None:
+        base_rows = table_model.rows
+    
+    # Handle variable reference (e.g. rows: "small")
+    if isinstance(base_rows, str) and not base_rows.isdigit():
+        var_name = base_rows
+        if var_name in params:
+            base_rows = params[var_name]
+        else:
+            raise ValueError(f"Table '{table_name}' expects param '{var_name}' (from matrix) but it was missing.")
+            
+    if base_rows is None:
+        raise ValueError(f"Could not resolve row count for '{table_name}'.")
+        
+    row_count = int(base_rows)
+
+    # Helper to resolve rows for ANY table (current or target)
+    def _resolve_rows(t_name, t_def):
+        r = params.get('rows')
+        if r is None:
+            r = t_def.get('rows') 
+        
+        # Handle variable references
+        if isinstance(r, str) and not r.isdigit():
+             if r in params:
+                 r = params[r]
+             else:
+                 # Fallback: maybe the target table uses a fixed number?
+                 # If we can't resolve it from params, we can't guess.
+                 pass
+        
+        return int(r) if r else None
+
+    # Batch Generation Strategy
+    
+    # Prepare parent directory
+    parent_dir = os.path.dirname(target_path)
+    if parent_dir:
+        os.makedirs(parent_dir, exist_ok=True)
+
+    # ... (rest of logic) ...
+    
+    chunk_count = (row_count // DEFAULT_CHUNK_SIZE) + (1 if row_count % DEFAULT_CHUNK_SIZE > 0 else 0)
+
+    print(f"[Gen] Generating {row_count} rows in batches of {DEFAULT_CHUNK_SIZE}...")
+
+    # Initialize empty list for temp chunks
+    temp_files = []
+    
+    try:
+        for i in range(chunk_count):
+            offset = i * DEFAULT_CHUNK_SIZE
+            current_size = min(DEFAULT_CHUNK_SIZE, row_count - offset)
+            
+            chunk_path = f"{target_path}.part_{i}"
+            _generate_chunk(offset, current_size, table_model, dataset_config, params, row_count, chunk_path,
+                            base_seed=dataset_config.get("seed", 42))
+            temp_files.append(chunk_path)
+            print(f"[Gen] Batch {i+1} done.")
+            
+        # Merge efficiently using Polars scan
+        # If single chunk, just rename
+        if len(temp_files) == 1:
+            import shutil
+            shutil.move(temp_files[0], target_path)
+        else:
+            pl.scan_parquet(f"{target_path}.part_*").collect().write_parquet(target_path)
+            
+    finally:
+        # Cleanup
+        for f in temp_files:
+            if os.path.exists(f):
+                os.remove(f)
+
+    return MaterializeResult(
+        metadata={
+            "path": MetadataValue.path(target_path),
+            "row_count": MetadataValue.int(row_count),
+            "table": table_name
+        }
+    )
+
+def _generate_chunk(offset, size, table_model, dataset_config, params, total_rows, output_path,
+                    base_seed=42):
+    # Enable reproducible generation. The base seed comes from dataset.seed
+    # in the YAML (default 42, preserving all pre-seed-param experiment IDs'
+    # data). Because the config is hashed, a different seed automatically
+    # yields a different Experiment ID — seed changes are never silent.
+    # Offset ensures different chunks get different sequences, but the same
+    # chunk always gets the same sequence.
+    np.random.seed(base_seed + offset)
+
+    data = {}
+    if not table_model.columns:
+        # Empty table with rows?
+        data["_row_id"] = np.arange(offset, offset + size) # Dummy
+    else:
+        for col_def in table_model.columns:
+            p_name = col_def.provider
+            generator_func = PROVIDER_REGISTRY.get(p_name)
+            
+            if not generator_func:
+                 raise ValueError(f"Unknown provider '{p_name}'")
+
+            kwargs = col_def.model_dump()
+            # Note: kwargs are already resolved at the top level
+            
+            # Dynamic Params substitution - Removed as per instruction
+            # for k, v in kwargs.items():
+            #     if isinstance(v, str) and v in params:
+            #         kwargs[k] = params[v]
+            
+            kwargs['table_name'] = params.get('table_name', 'unknown') 
+            
+            # Contextual Handling
+            # Foreign Key needs target_rows to generate valid IDs
+            if p_name in ["foreign_key", "foreign key"]:
+                target_table = col_def.target_table
+                if target_table:
+                    tables = dataset_config.get("tables", {})
+                    target_def = tables.get(target_table, {})
+                    target_rows = None
+
+                    # Resolve target table's row count
+                    target_row_val = target_def.get('rows') or params.get('rows')
+                    if isinstance(target_row_val, str) and target_row_val in params:
+                        target_rows = int(params[target_row_val])
+                    elif target_row_val is not None:
+                        target_rows = int(target_row_val)
+                    else:
+                        target_rows = total_rows  # Fallback to current table's row count
+
+                    kwargs['target_rows'] = target_rows
+
+            # EXECUTE PROVIDER
+            # Note: Most providers don't care about offset (random). 
+            # Sequence might? 'sequence' provider usually starts at 1.
+            # We need to handle 'sequence' specially or pass offset via kwargs if provider supports it.
+            if p_name == "sequence":
+                kwargs['start'] = offset + 1 
+            
+            # Text Concat needs existing data (in this chunk)
+            if p_name == "text_concat":
+                results = generator_func(size, existing_data=data, **kwargs)
+            else:
+                results = generator_func(size, **kwargs)
+
+            # Multi-column providers return a dict of {col_name: array}.
+            # Expand them directly into data instead of nesting under one key.
+            if isinstance(results, dict):
+                data.update(results)
+            else:
+                data[col_def.name] = results
+            
+    df = pl.DataFrame(data)
+
+    # Apply Null Masks natively
+    for col_def in table_model.columns:
+        # null_probability will be a float here because it was resolved globally
+        if col_def.null_probability > 0:
+            mask = np.random.rand(size) < col_def.null_probability
+            
+            # Use Polars expression to nullify values
+            # This preserves the underlying Arrow type (Int64, Utf8) instead of casting to Object
+            df = df.with_columns(
+                pl.when(pl.lit(mask)).then(None).otherwise(pl.col(col_def.name)).alias(col_def.name)
+            )
+
+    df.write_parquet(output_path)

sql_benchmarks/plugins/data_sources/local_file_loader.py

@@ -0,0 +1,62 @@
+import polars as pl
+import os
+from dagster import MaterializeResult, MetadataValue
+
+# CHANGED: 'output_dir' -> 'target_path'
+def generate(context, params, table_name, target_path, dataset_config):
+    """
+    Normalizes REAL data (CSV/Parquet/JSON) into standard Parquet for the platform.
+    """
+    partition_key = context.partition_key
+    
+    # 1. GET PATH
+    paths = dataset_config.get("paths", {})
+    source_template = paths.get(table_name)
+    
+    if not source_template:
+        raise ValueError(f"CRITICAL: No path defined for table '{table_name}' in dataset.paths")
+    
+    # INTERPOLATE
+    # Support {topology} or other vars from the partition params
+    source_file = source_template.format(**params)
+
+    if not os.path.exists(source_file):
+        raise FileNotFoundError(f"CRITICAL: Real data file not found: {source_file}")
+
+    context.log.info(f"Normalizing real data from {source_file}...")
+
+    # 2. STREAMING CONVERT (Lazy Evaluation)
+    # Ensure directory exists
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+
+    file_ext = os.path.splitext(source_file)[1].lower()
+
+    try:
+        # Create a LazyFrame (Scan) based on extension
+        if file_ext == ".csv":
+            q = pl.scan_csv(source_file, infer_schema_length=10000)
+        elif file_ext == ".parquet":
+            q = pl.scan_parquet(source_file)
+        elif file_ext in [".json", ".ndjson"]:
+            q = pl.scan_ndjson(source_file)
+        else:
+            raise ValueError(f"Unsupported file type: {file_ext}")
+
+        # 3. SINK TO PARQUET
+        # Use the explicit target_path provided by the factory
+        q.sink_parquet(target_path)
+        
+        # 4. METADATA
+        row_count = pl.scan_parquet(target_path).select(pl.len()).collect().item()
+            
+    except Exception as e:
+        raise RuntimeError(f"Failed to process {source_file}: {e}")
+
+    return MaterializeResult(
+        metadata={
+            "path": MetadataValue.path(target_path),
+            "row_count": MetadataValue.int(row_count),
+            "source_file": MetadataValue.path(source_file),
+            "loader_backend": "polars_streaming"
+        }
+    )

sql_benchmarks/plugins/data_sources/tpc_h.py

@@ -0,0 +1,61 @@
+import duckdb
+import os
+from dagster import MaterializeResult, MetadataValue
+
+# CHANGED: 'output_dir' -> 'target_path'
+def generate(context, params, table_name, target_path, dataset_config):
+    partition_key = context.partition_key
+    
+    if 'tables' not in dataset_config:
+        raise ValueError("Critical: 'dataset.tables' missing in active.yaml")
+        
+    if not dataset_config['tables'].get(table_name):
+        raise ValueError(f"Table '{table_name}' requested but not defined in YAML.")
+
+    scale_factor = params.get('scale_factor', 0.01)
+    
+    context.log.info(f"Generating TPC-H table '{table_name}' with SF={scale_factor}...")
+
+    # 1. Setup DuckDB
+    con = duckdb.connect(database=':memory:')
+    con.install_extension("tpch")
+    con.load_extension("tpch")
+    
+    # 2. Generate Data
+    con.execute(f"CALL dbgen(sf={scale_factor})")
+    
+    # 3. Export with Type Casting
+    # Ensure directory exists (Defensive)
+    os.makedirs(os.path.dirname(target_path), exist_ok=True)
+    
+    # Get columns for the table
+    cols_info = con.execute(f"DESCRIBE {table_name}").fetchall()
+    
+    # Build a SELECT statement that casts Decimals to Doubles
+    select_parts = []
+    for col_name, col_type, _, _, _, _ in cols_info:
+        if "DECIMAL" in col_type:
+            select_parts.append(f"CAST({col_name} AS DOUBLE) AS {col_name}")
+        else:
+            select_parts.append(col_name)
+            
+    select_query = ", ".join(select_parts)
+    
+    # USE THE PASSED TARGET_PATH
+    con.execute(f"""
+        COPY (SELECT {select_query} FROM {table_name}) 
+        TO '{target_path}' 
+        (FORMAT PARQUET, COMPRESSION 'SNAPPY')
+    """)
+    
+    # 4. Metadata
+    row_count = con.execute(f"SELECT count(*) FROM {table_name}").fetchone()[0]
+    
+    return MaterializeResult(
+        metadata={
+            "path": MetadataValue.path(target_path),
+            "row_count": MetadataValue.int(row_count),
+            "scale_factor": MetadataValue.float(scale_factor),
+            "source": "duckdb_tpch_extension"
+        }
+    )

sql_benchmarks/resources/actian.py

@@ -0,0 +1,234 @@
+import os
+import time
+import socket
+from dagster import ConfigurableResource
+from typing import Dict, Any, Optional
+from pydantic import ConfigDict, PrivateAttr
+from .base_engine import IBenchmarkEngine
+from .actian_client import ActianClient
+
+
+class ActianEngine(ConfigurableResource):
+    """
+    Dagster Resource for Actian Vector on EC2.
+
+    Manages SSH tunnel lifecycle and provides cold-cache enforcement
+    via remote service restart.
+    """
+
+    # --- CONFIGURATION ---
+    # Local connection (through tunnel)
+    local_port: int = 27832
+
+    # Remote EC2 configuration
+    ec2_host: str = os.getenv("ACTIAN_EC2_HOST", "")
+    ec2_user: str = os.getenv("ACTIAN_EC2_USER", "ingres")
+    ssh_key_path: str = os.getenv("ACTIAN_SSH_KEY", os.path.expanduser("~/.ssh/benchmark-key-pair.pem"))
+
+    # Actian database configuration
+    database: str = os.getenv("ACTIAN_DATABASE", "benchmark_db")
+    db_user: str = os.getenv("ACTIAN_USER", "ingres")
+    db_password: str = os.getenv("ACTIAN_PASSWORD", "AntiGravity")
+
+    # Remote paths
+    remote_data_dir: str = "/tmp/benchmark_data"
+    vwload_path: str = "/opt/Actian/VectorVW/ingres/bin/vwload"
+    actian_sql_path: str = "/opt/Actian/VectorVW/ingres/bin/sql"
+
+    model_config = ConfigDict(extra='forbid')
+
+    # Private attributes for runtime state
+    _tunnel: Any = PrivateAttr(default=None)
+    _ssh_client: Any = PrivateAttr(default=None)
+
+    def _ensure_tunnel(self):
+        """Establishes SSH tunnel if not already active."""
+        if self._tunnel is not None and self._tunnel.is_active:
+            return
+
+        if not self.ec2_host:
+            raise ValueError("ACTIAN_EC2_HOST environment variable must be set")
+
+        from sshtunnel import SSHTunnelForwarder
+
+        self._tunnel = SSHTunnelForwarder(
+            (self.ec2_host, 22),
+            ssh_username=self.ec2_user,
+            ssh_pkey=self.ssh_key_path,
+            remote_bind_address=('localhost', 27832),
+            local_bind_address=('localhost', self.local_port),
+        )
+        self._tunnel.start()
+        print(f"[Actian] SSH tunnel established: localhost:{self._tunnel.local_bind_port} -> {self.ec2_host}:27832")
+
+        # Give the tunnel a moment to stabilize
+        time.sleep(1)
+
+    def _ensure_ssh(self):
+        """Establishes SSH client for remote commands."""
+        if self._ssh_client is not None:
+            # Check if connection is still alive
+            try:
+                self._ssh_client.exec_command("echo ok", timeout=5)
+                return
+            except Exception:
+                self._ssh_client = None
+
+        if not self.ec2_host:
+            raise ValueError("ACTIAN_EC2_HOST environment variable must be set")
+
+        import paramiko
+    # 1. Expand the path (handling both the env var and the default)
+        full_key_path = os.path.expanduser(self.ssh_key_path)
+    
+        # 2. Architect Fix: Explicitly load as RSAKey. 
+        # This bypasses the logic where Paramiko scans for DSSKey and crashes.
+        try:
+            pkey = paramiko.RSAKey.from_private_key_file(full_key_path)
+        except Exception as e:
+            raise RuntimeError(f"Could not load RSA key from {full_key_path}: {e}")
+    
+        # 3. Connect using 'pkey' instead of 'key_filename'
+        self._ssh_client = paramiko.SSHClient()
+        self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        
+        self._ssh_client.connect(
+            hostname=self.ec2_host,
+            username=self.ec2_user,
+            pkey=pkey, # Use the loaded object here
+            timeout=30
+        )
+        # self._ssh_client = paramiko.SSHClient()
+        # self._ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+        # self._ssh_client.connect(
+        #     hostname=self.ec2_host,
+        #     username=self.ec2_user,
+        #     key_filename=self.ssh_key_path,
+        #     timeout=30
+        # )
+        print(f"[Actian] SSH connection established to {self.ec2_host}")
+
+    def _ssh_exec(self, command: str, timeout: int = 300) -> tuple[str, str, int]:
+        """Execute a command on the remote EC2 instance."""
+        self._ensure_ssh()
+
+        stdin, stdout, stderr = self._ssh_client.exec_command(command, timeout=timeout)
+        exit_code = stdout.channel.recv_exit_status()
+
+        return stdout.read().decode(), stderr.read().decode(), exit_code
+
+    def _get_client(self) -> ActianClient:
+        """Factory method for creating the Actian client."""
+        self._ensure_tunnel()
+
+        connection_params = {
+            "host": "localhost",
+            "port": self._tunnel.local_bind_port,
+            "user": self.db_user,
+            "password": self.db_password,
+            "database": self.database
+        }
+
+        # Pass SSH details for vwload operations
+        ssh_params = {
+            "ec2_host": self.ec2_host,
+            "ec2_user": self.ec2_user,
+            "ssh_key_path": self.ssh_key_path,
+            "remote_data_dir": self.remote_data_dir,
+            "vwload_path": self.vwload_path,
+            "actian_sql_path": self.actian_sql_path,
+        }
+
+        return ActianClient(connection_params, ssh_params)
+
+    # --- IBenchmarkEngine Implementation ---
+
+    def get_engine_name(self) -> str:
+        return "actian"
+
+    def clear_cache(self, settings: dict = None):
+        """
+        Enforces cold cache by restarting Actian Vector service on EC2.
+
+        This clears both:
+        1. X100 engine memory segments
+        2. Linux page cache (via service restart)
+        """
+        print("[Actian] Clearing cache via service restart...")
+
+        # Restart Actian Vector service
+        stdout, stderr, exit_code = self._ssh_exec(
+            "sudo systemctl restart actian-vectorVW",
+            timeout=120
+        )
+
+        if exit_code != 0:
+            raise RuntimeError(f"Failed to restart Actian service: {stderr}")
+
+        # Wait for service to be ready
+        self._wait_for_ready()
+        print("[Actian] Service restarted, cache cleared")
+
+    def run_query(self, sql: str, partition_key: str, engine_params: Dict[str, Any] = None) -> Optional[float]:
+        """Execute a benchmark query with cold cache.
+
+        engine_params is the 'actian' namespace — accepted but not yet applied
+        (Vector session tuning is a future capability).
+        """
+        self.clear_cache()
+        client = self._get_client()
+        return client.run_query(sql, {})
+
+    def bulk_load(self, filepath: str, table_name: str, partition_key: str) -> None:
+        """
+        Bulk load data into Actian Vector.
+
+        1. SCP the file to EC2
+        2. Convert to CSV if needed (vwload prefers CSV)
+        3. Run vwload on EC2
+        """
+        self._ensure_tunnel()
+        self._wait_for_ready()
+
+        client = self._get_client()
+        client.bulk_load(filepath, table_name, partition_key)
+
+    def _wait_for_ready(self, timeout: int = 120):
+        """Wait for Actian to be ready to accept connections."""
+        self._ensure_tunnel()
+
+        start = time.time()
+        while time.time() - start < timeout:
+            try:
+                with socket.create_connection(
+                    ('localhost', self._tunnel.local_bind_port),
+                    timeout=2
+                ):
+                    # Port is open, but let's verify the service is actually responding
+                    time.sleep(2)
+                    return
+            except (socket.timeout, ConnectionRefusedError, OSError):
+                time.sleep(2)
+
+        raise TimeoutError(
+            f"Actian Vector on {self.ec2_host} failed to respond within {timeout}s"
+        )
+
+    def cleanup(self):
+        """Clean up SSH tunnel and connections."""
+        if self._ssh_client:
+            self._ssh_client.close()
+            self._ssh_client = None
+
+        if self._tunnel:
+            self._tunnel.stop()
+            self._tunnel = None
+
+        print("[Actian] Connections closed")
+
+    def __del__(self):
+        """Ensure cleanup on garbage collection."""
+        try:
+            self.cleanup()
+        except Exception:
+            pass