dvt-core 0.52.2__cp310-cp310-macosx_10_9_x86_64.whl

dbt/compute/filter_pushdown.py

@@ -0,0 +1,273 @@
+"""
+Filter pushdown optimization for federated queries.
+
+Extracts filters (WHERE, LIMIT, ORDER BY) from compiled SQL and rewrites them
+to be pushed down to source databases in their native SQL dialects.
+"""
+
+import re
+import sys
+import sqlparse
+from sqlparse.sql import Statement, Token, TokenList, Identifier, Where, Comparison
+from sqlparse.tokens import Keyword, Whitespace
+from typing import Dict, List, Optional, Any
+
+
+class FilterPushdownOptimizer:
+    """
+    Optimizes federated queries by pushing filters down to source databases.
+
+    Strategy:
+    1. Parse compiled SQL to extract filters per source table
+    2. Rewrite filters in each source adapter's SQL dialect
+    3. Return subqueries for JDBC reads instead of plain table names
+
+    Example:
+        Input SQL:
+            SELECT * FROM snowflake_table WHERE date > '2024-01-01' LIMIT 10
+
+        Output:
+            JDBC subquery: (SELECT * FROM snowflake_table WHERE date > '2024-01-01' LIMIT 10)
+    """
+
+    def __init__(self, compiled_sql: str, source_tables: List[Any]):
+        """
+        Initialize optimizer with compiled SQL and source table metadata.
+
+        Args:
+            compiled_sql: The fully compiled SQL from the model
+            source_tables: List of SourceTableMetadata objects
+        """
+        self.compiled_sql = compiled_sql
+        self.source_tables = source_tables
+        self.parsed = sqlparse.parse(compiled_sql)[0] if compiled_sql else None
+
+    def extract_limit(self) -> Optional[int]:
+        """
+        Extract LIMIT clause from SQL.
+
+        Returns:
+            Limit value as integer, or None if no LIMIT clause
+        """
+        if not self.parsed:
+            return None
+
+        # Simple regex approach for LIMIT (works for most cases)
+        limit_match = re.search(r'\bLIMIT\s+(\d+)\b', self.compiled_sql, re.IGNORECASE)
+        if limit_match:
+            return int(limit_match.group(1))
+
+        return None
+
+    def extract_sample_clause(self) -> Optional[Dict[str, Any]]:
+        """
+        Extract SAMPLE/TABLESAMPLE clause from SQL (Snowflake-specific sampling).
+
+        Snowflake supports several SAMPLE methods:
+        - SAMPLE (N) or SAMPLE (N ROWS) - Row-count sampling
+        - SAMPLE SYSTEM (P) - System/block sampling with P% probability
+        - SAMPLE BERNOULLI (P) - Bernoulli/row-level sampling with P% probability
+        - SAMPLE BLOCK (P) - Alias for SYSTEM
+        - TABLESAMPLE ... - Alternative syntax
+        - REPEATABLE(seed) or SEED(seed) - Reproducible sampling
+
+        Returns:
+            Dict with keys:
+                - 'method': 'ROWS', 'SYSTEM', 'BERNOULLI', 'BLOCK'
+                - 'value': int (row count or percentage)
+                - 'seed': Optional int for reproducible sampling
+                - 'full_clause': The complete SAMPLE clause to push down
+            or None if no SAMPLE clause
+        """
+        if not self.parsed:
+            return None
+
+        # Try to match complete SAMPLE/TABLESAMPLE clause
+        # Pattern: (TABLE)?SAMPLE <method>? (value) (REPEATABLE|SEED)?(seed)?
+
+        # Match: SAMPLE (N) or SAMPLE (N ROWS) with optional seed
+        sample_rows = re.search(
+            r'\b(?:TABLE)?SAMPLE\s*\(\s*(\d+)(?:\s+ROWS)?\s*\)'
+            r'(?:\s+(?:REPEATABLE|SEED)\s*\(\s*(\d+)\s*\))?',
+            self.compiled_sql,
+            re.IGNORECASE
+        )
+        if sample_rows:
+            result = {
+                'method': 'ROWS',
+                'value': int(sample_rows.group(1)),
+                'seed': int(sample_rows.group(2)) if sample_rows.group(2) else None
+            }
+            # Build full clause
+            clause = f"SAMPLE ({result['value']})"
+            if result['seed']:
+                clause += f" REPEATABLE ({result['seed']})"
+            result['full_clause'] = clause
+            return result
+
+        # Match: SAMPLE SYSTEM|BERNOULLI|BLOCK (P) with optional seed
+        sample_method = re.search(
+            r'\b(?:TABLE)?SAMPLE\s+(SYSTEM|BERNOULLI|BLOCK)\s*\(\s*(\d+(?:\.\d+)?)\s*\)'
+            r'(?:\s+(?:REPEATABLE|SEED)\s*\(\s*(\d+)\s*\))?',
+            self.compiled_sql,
+            re.IGNORECASE
+        )
+        if sample_method:
+            method = sample_method.group(1).upper()
+            # BLOCK is an alias for SYSTEM
+            if method == 'BLOCK':
+                method = 'SYSTEM'
+
+            result = {
+                'method': method,
+                'value': float(sample_method.group(2)),
+                'seed': int(sample_method.group(3)) if sample_method.group(3) else None
+            }
+            # Build full clause
+            clause = f"SAMPLE {result['method']} ({result['value']})"
+            if result['seed']:
+                clause += f" REPEATABLE ({result['seed']})"
+            result['full_clause'] = clause
+            return result
+
+        return None
+
+    def extract_where_clauses(self) -> Dict[str, List[str]]:
+        """
+        Extract WHERE clauses that apply to specific source tables.
+
+        Returns:
+            Dict mapping table name/alias to list of WHERE conditions
+
+        Example:
+            {
+                'snowflake_table': ['date > \'2024-01-01\'', 'status = \'active\''],
+                'postgres_table': ['id > 100']
+            }
+        """
+        # TODO: Implement WHERE clause extraction using sqlparse
+        # For now, return empty dict - LIMIT pushdown is the priority
+        return {}
+
+    def build_pushdown_subquery(
+        self,
+        source_table: Any,
+        adapter_type: str
+    ) -> Optional[str]:
+        """
+        Build a subquery with pushed-down filters for a specific source table.
+
+        Args:
+            source_table: SourceTableMetadata object
+            adapter_type: Adapter type (postgres, snowflake, etc.)
+
+        Returns:
+            SQL subquery with filters, or None if no pushdown possible
+
+        Example:
+            Input: table="schema.table", LIMIT 10
+            Output: "(SELECT * FROM schema.table LIMIT 10)"
+        """
+        limit = self.extract_limit()
+        sample_clause = self.extract_sample_clause()
+        where_clauses = self.extract_where_clauses()
+
+        # DVT v0.4.7: Suppressed debug output for clean console
+        # Debug info: LIMIT={limit}, SAMPLE={sample_clause}, WHERE={where_clauses}
+
+        # If no filters to push down, return None (read full table)
+        if not limit and not sample_clause and not where_clauses:
+            return None
+
+        # Build subquery
+        qualified_name = source_table.qualified_name
+        subquery_parts = [f"SELECT * FROM {qualified_name}"]
+
+        # Add SAMPLE clause (Snowflake-specific, goes right after FROM)
+        if sample_clause and adapter_type.lower() == 'snowflake':
+            # Use the pre-built full_clause from extract_sample_clause
+            # This includes all sampling options: method, value, and seed
+            subquery_parts.append(sample_clause['full_clause'])
+
+        # Add WHERE clauses (if any)
+        table_key = source_table.identifier  # or qualified_name
+        if table_key in where_clauses:
+            conditions = " AND ".join(where_clauses[table_key])
+            subquery_parts.append(f"WHERE {conditions}")
+
+        # Add LIMIT (if present and no SAMPLE used)
+        # Note: SAMPLE takes precedence over LIMIT for Snowflake
+        if limit and not (sample_clause and adapter_type.lower() == 'snowflake'):
+            # Rewrite LIMIT in adapter's dialect
+            limit_clause = self._rewrite_limit_for_adapter(limit, adapter_type)
+            if limit_clause:
+                subquery_parts.append(limit_clause)
+
+        subquery = " ".join(subquery_parts)
+        # DVT v0.4.7: Suppressed debug output
+        return f"({subquery})"
+
+    def _rewrite_limit_for_adapter(self, limit: int, adapter_type: str) -> Optional[str]:
+        """
+        Rewrite LIMIT clause for specific adapter's SQL dialect.
+
+        Args:
+            limit: Limit value
+            adapter_type: Adapter type (postgres, snowflake, redshift, etc.)
+
+        Returns:
+            LIMIT clause in adapter's dialect
+        """
+        # Most adapters support standard LIMIT syntax
+        standard_adapters = [
+            'postgres', 'postgresql',
+            'snowflake',
+            'redshift',
+            'mysql',
+            'sqlite',
+            'bigquery'
+        ]
+
+        if adapter_type.lower() in standard_adapters:
+            return f"LIMIT {limit}"
+
+        # SQL Server / TSQL uses TOP
+        if adapter_type.lower() in ['sqlserver', 'mssql', 'tsql']:
+            # Note: This should go in SELECT clause, not at the end
+            # For now, return None - we'll handle this in a future iteration
+            return None
+
+        # Oracle uses ROWNUM or FETCH FIRST (12c+)
+        if adapter_type.lower() == 'oracle':
+            return f"FETCH FIRST {limit} ROWS ONLY"
+
+        # Default: standard LIMIT
+        return f"LIMIT {limit}"
+
+
+def optimize_jdbc_table_read(
+    source_table: Any,
+    compiled_sql: str,
+    source_tables: List[Any],
+    adapter_type: str
+) -> str:
+    """
+    Optimize JDBC table read by pushing down filters.
+
+    Args:
+        source_table: SourceTableMetadata for this table
+        compiled_sql: Compiled SQL from the model
+        source_tables: All source tables in the query
+        adapter_type: Source adapter type
+
+    Returns:
+        Table identifier (plain name or subquery with filters)
+    """
+    optimizer = FilterPushdownOptimizer(compiled_sql, source_tables)
+    subquery = optimizer.build_pushdown_subquery(source_table, adapter_type)
+
+    if subquery:
+        return subquery
+    else:
+        # No filters to push down - read full table
+        return source_table.qualified_name

dbt/compute/jar_provisioning.py

@@ -0,0 +1,255 @@
+"""
+JAR Provisioning Module
+
+Centralized JDBC JAR provisioning for Spark compute engines.
+
+v0.5.98: Supports two provisioning strategies:
+- LocalJARProvisioning: Uses spark.jars with local file paths (fast startup)
+- RemoteJARProvisioning: Uses spark.jars.packages with Maven coordinates (remote clusters)
+
+Local Spark uses local JARs from .dvt/jdbc_jars/ for instant startup.
+Remote clusters (Databricks, EMR, Dataproc, Standalone) use Maven coordinates
+so Spark workers can download JARs directly from Maven Central.
+"""
+
+import glob
+import os
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+
+
+# Maven coordinates for JDBC drivers - used by remote clusters
+# Format: groupId:artifactId:version
+# These are the same JARs as downloaded by `dvt target sync`, but expressed as Maven coordinates
+JDBC_MAVEN_COORDINATES = {
+    # Official dbt-labs adapters - JDBC drivers only
+    "postgres": "org.postgresql:postgresql:42.7.4",
+    "snowflake": "net.snowflake:snowflake-jdbc:3.16.1",
+    "bigquery": "com.google.cloud.bigdataoss:gcs-connector:hadoop3-2.2.22",
+    "redshift": "com.amazon.redshift:redshift-jdbc42:2.1.0.32",
+    "spark": "",  # Native, no JDBC needed
+    "databricks": "com.databricks:databricks-jdbc:2.6.36",
+    "trino": "io.trino:trino-jdbc:443",
+    "duckdb": "org.duckdb:duckdb_jdbc:1.1.3",
+    # Community adapters - JDBC drivers only (verified on Maven)
+    "mysql": "com.mysql:mysql-connector-j:9.1.0",
+    "sqlserver": "com.microsoft.sqlserver:mssql-jdbc:12.8.1.jre11",
+    "synapse": "com.microsoft.sqlserver:mssql-jdbc:12.8.1.jre11",
+    "fabric": "com.microsoft.sqlserver:mssql-jdbc:12.8.1.jre11",
+    "oracle": "com.oracle.database.jdbc:ojdbc11:23.6.0.24.10",
+    "teradata": "com.teradata.jdbc:terajdbc:20.00.00.20",
+    "clickhouse": "com.clickhouse:clickhouse-jdbc:0.6.5",
+    "greenplum": "org.postgresql:postgresql:42.7.4",  # PostgreSQL compatible
+    "vertica": "com.vertica.jdbc:vertica-jdbc:24.3.0-0",
+    "sqlite": "org.xerial:sqlite-jdbc:3.47.1.0",
+    "mariadb": "org.mariadb.jdbc:mariadb-java-client:3.4.1",
+    "exasol": "com.exasol:exasol-jdbc:24.2.0",
+    "db2": "com.ibm.db2:jcc:11.5.9.0",
+    "presto": "io.prestosql:presto-jdbc:350",
+    "hive": "org.apache.hive:hive-jdbc:3.1.3",
+    "singlestore": "com.singlestore:singlestore-jdbc-client:1.2.9",
+    "starrocks": "com.mysql:mysql-connector-j:9.1.0",  # MySQL wire protocol
+    "doris": "com.mysql:mysql-connector-j:9.1.0",  # MySQL wire protocol
+    "materialize": "org.postgresql:postgresql:42.7.4",  # PostgreSQL wire protocol
+    "neo4j": "org.neo4j:neo4j-jdbc-driver:4.0.10",
+    "timescaledb": "org.postgresql:postgresql:42.7.4",  # PostgreSQL extension
+    "questdb": "org.postgresql:postgresql:42.7.4",  # PostgreSQL wire protocol
+}
+
+
+class JARProvisioning(ABC):
+    """Abstract base class for JAR provisioning strategies."""
+
+    @abstractmethod
+    def get_spark_config(self, adapter_types: Set[str]) -> Dict[str, str]:
+        """
+        Get Spark configuration for JDBC JARs.
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Dictionary of Spark config keys/values
+        """
+        pass
+
+    @abstractmethod
+    def get_provisioning_type(self) -> str:
+        """
+        Get the type of JAR provisioning.
+
+        :returns: 'local' or 'maven'
+        """
+        pass
+
+
+class LocalJARProvisioning(JARProvisioning):
+    """
+    Local JAR provisioning using spark.jars with local file paths.
+
+    Best for local Spark (spark-local) where JARs are pre-downloaded
+    to .dvt/jdbc_jars/ directory via `dvt target sync`.
+
+    Advantages:
+    - Instant startup (no JAR download at runtime)
+    - Works offline
+    - Consistent JAR versions
+
+    Disadvantages:
+    - Only works for local Spark (JARs must be on local filesystem)
+    - Requires running `dvt target sync` first
+    """
+
+    def __init__(self, project_dir: Optional[str] = None):
+        """
+        Initialize local JAR provisioning.
+
+        :param project_dir: Path to project root directory (defaults to cwd)
+        """
+        self.project_dir = project_dir or os.getcwd()
+        self.jdbc_jars_dir = os.path.join(self.project_dir, ".dvt", "jdbc_jars")
+
+    def get_jar_paths(self) -> List[str]:
+        """
+        Discover all JDBC JAR files from project cache.
+
+        :returns: List of absolute JAR file paths
+        """
+        if not os.path.exists(self.jdbc_jars_dir):
+            return []
+
+        jar_pattern = os.path.join(self.jdbc_jars_dir, "*.jar")
+        return sorted(glob.glob(jar_pattern))
+
+    def get_spark_config(self, adapter_types: Set[str]) -> Dict[str, str]:
+        """
+        Get Spark config using local JAR paths.
+
+        :param adapter_types: Set of adapter types (ignored - uses all JARs found)
+        :returns: Dictionary with spark.jars config
+        """
+        jar_paths = self.get_jar_paths()
+
+        if jar_paths:
+            return {"spark.jars": ",".join(jar_paths)}
+        return {}
+
+    def get_provisioning_type(self) -> str:
+        """Get provisioning type."""
+        return "local"
+
+
+class RemoteJARProvisioning(JARProvisioning):
+    """
+    Remote JAR provisioning using spark.jars.packages with Maven coordinates.
+
+    Best for remote Spark clusters (Databricks, EMR, Dataproc, Standalone)
+    where Spark workers need to download JARs from Maven Central.
+
+    Advantages:
+    - Works with any remote Spark cluster
+    - No need to pre-install JARs on cluster
+    - Spark handles dependency resolution
+
+    Disadvantages:
+    - Requires network access to Maven Central
+    - First query may be slower (JAR download)
+    - May need to configure Maven repositories for private JARs
+    """
+
+    def __init__(self, profiles_dir: Optional[str] = None):
+        """
+        Initialize remote JAR provisioning.
+
+        :param profiles_dir: Path to DVT profiles directory (for scanning profiles.yml)
+        """
+        self.profiles_dir = profiles_dir or str(Path.home() / ".dvt")
+
+    def get_maven_coordinates(self, adapter_types: Set[str]) -> List[str]:
+        """
+        Get Maven coordinates for the specified adapter types.
+
+        :param adapter_types: Set of adapter types
+        :returns: List of Maven coordinates (group:artifact:version)
+        """
+        coordinates = []
+        seen = set()  # Avoid duplicates (e.g., postgres and timescaledb both use postgresql)
+
+        for adapter_type in adapter_types:
+            coord = JDBC_MAVEN_COORDINATES.get(adapter_type.lower(), "")
+            if coord and coord not in seen:
+                coordinates.append(coord)
+                seen.add(coord)
+
+        return sorted(coordinates)
+
+    def get_spark_config(self, adapter_types: Set[str]) -> Dict[str, str]:
+        """
+        Get Spark config using Maven coordinates.
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Dictionary with spark.jars.packages config
+        """
+        coordinates = self.get_maven_coordinates(adapter_types)
+
+        if coordinates:
+            return {"spark.jars.packages": ",".join(coordinates)}
+        return {}
+
+    def get_provisioning_type(self) -> str:
+        """Get provisioning type."""
+        return "maven"
+
+
+def get_required_adapter_types(profiles_dir: Optional[str] = None) -> Set[str]:
+    """
+    Scan profiles.yml and return the set of adapter types needed.
+
+    :param profiles_dir: Path to DVT profiles directory
+    :returns: Set of adapter type names (e.g., {'postgres', 'snowflake'})
+    """
+    from dbt.clients.yaml_helper import load_yaml_text
+
+    if profiles_dir is None:
+        profiles_dir = str(Path.home() / ".dvt")
+
+    profiles_path = Path(profiles_dir) / "profiles.yml"
+    if not profiles_path.exists():
+        return set()
+
+    try:
+        content = profiles_path.read_text()
+        profiles = load_yaml_text(content) or {}
+
+        adapter_types = set()
+        for profile_name, profile_data in profiles.items():
+            if isinstance(profile_data, dict):
+                outputs = profile_data.get("outputs", {})
+                for target_name, target_config in outputs.items():
+                    if isinstance(target_config, dict):
+                        adapter_type = target_config.get("type")
+                        if adapter_type:
+                            adapter_types.add(adapter_type.lower())
+
+        return adapter_types
+
+    except Exception:
+        return set()
+
+
+def get_provisioning_for_platform(
+    platform: str,
+    project_dir: Optional[str] = None,
+    profiles_dir: Optional[str] = None,
+) -> JARProvisioning:
+    """
+    Factory function to get the appropriate JAR provisioning strategy.
+
+    :param platform: Spark platform ('local', 'databricks', 'emr', 'dataproc', 'standalone')
+    :param project_dir: Project directory (for local provisioning)
+    :param profiles_dir: Profiles directory (for remote provisioning)
+    :returns: JARProvisioning instance
+    """
+    if platform.lower() == "local":
+        return LocalJARProvisioning(project_dir=project_dir)
+    else:
+        # All remote platforms use Maven coordinates
+        return RemoteJARProvisioning(profiles_dir=profiles_dir)