dvt-core 0.59.0a51__py3-none-any.whl

dbt/compute/strategies/base.py

@@ -0,0 +1,165 @@
+"""
+Base Connection Strategy for Spark Engines
+
+Defines the abstract interface for different Spark connection strategies.
+Uses composition over inheritance for flexible platform support.
+
+v0.5.98: Added JAR provisioning and connectivity testing methods.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Set, Tuple
+
+try:
+    from pyspark.sql import SparkSession
+
+    PYSPARK_AVAILABLE = True
+except ImportError:
+    PYSPARK_AVAILABLE = False
+    SparkSession = None
+
+
+class BaseConnectionStrategy(ABC):
+    """
+    Abstract base class for Spark connection strategies.
+
+    Different strategies implement different ways to connect to Spark:
+    - LocalStrategy: Embedded PySpark (in-process)
+    - DatabricksStrategy: Databricks Connect (remote cluster)
+    - EMRStrategy: AWS EMR cluster
+    - DataprocStrategy: GCP Dataproc
+    - StandaloneStrategy: Self-managed Spark clusters
+    """
+
+    def __init__(self, config: Dict[str, Any], app_name: str = "DVT-Compute"):
+        """
+        Initialize connection strategy.
+
+        :param config: Strategy-specific configuration
+        :param app_name: Spark application name
+        """
+        self.config = config
+        self.app_name = app_name
+
+    @abstractmethod
+    def get_spark_session(self) -> SparkSession:
+        """
+        Create and return a SparkSession.
+
+        :returns: Initialized SparkSession
+        :raises DbtRuntimeError: If session creation fails
+        """
+        pass
+
+    @abstractmethod
+    def validate_config(self) -> None:
+        """
+        Validate strategy-specific configuration.
+
+        :raises DbtRuntimeError: If configuration is invalid
+        """
+        pass
+
+    def estimate_cost(self, duration_minutes: float) -> float:
+        """
+        Estimate cost for running on this platform.
+
+        Default implementation returns 0.0 (free). Override for cloud platforms.
+
+        :param duration_minutes: Estimated query duration in minutes
+        :returns: Estimated cost in USD
+        """
+        return 0.0
+
+    @abstractmethod
+    def close(self, spark: Optional[SparkSession]) -> None:
+        """
+        Clean up Spark session.
+
+        :param spark: SparkSession to clean up (may be None)
+        """
+        pass
+
+    def get_platform_name(self) -> str:
+        """
+        Get human-readable platform name.
+
+        :returns: Platform name (e.g., "local", "databricks", "emr")
+        """
+        return self.__class__.__name__.replace("Strategy", "").lower()
+
+    def get_jar_provisioning_config(self, adapter_types: Set[str]) -> Dict[str, str]:
+        """
+        Get Spark configuration for JDBC JAR provisioning.
+
+        Default implementation returns empty dict. Override in subclasses
+        to provide platform-specific JAR configuration.
+
+        Local platforms use spark.jars (local file paths).
+        Remote platforms use spark.jars.packages (Maven coordinates).
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Dictionary of Spark config keys/values (e.g., {"spark.jars": "..."})
+        """
+        return {}
+
+    def test_connectivity(self) -> Tuple[bool, str]:
+        """
+        Test basic connectivity to the Spark cluster.
+
+        Creates a session, runs a simple query, and returns status.
+        Override for platform-specific connectivity testing.
+
+        :returns: Tuple of (success, message)
+        """
+        try:
+            spark = self.get_spark_session()
+            # Run a simple SQL query to verify connectivity
+            spark.sql("SELECT 1 AS test").collect()
+            return (True, "Session created and SQL test passed")
+        except Exception as e:
+            return (False, str(e))
+
+    def test_jdbc_connectivity(
+        self,
+        jdbc_url: str,
+        properties: Dict[str, str],
+        table_or_query: str = "(SELECT 1 AS test) AS t",
+    ) -> Tuple[bool, str]:
+        """
+        Test JDBC connectivity through the Spark cluster.
+
+        Creates a session and attempts to read from a JDBC source.
+        This verifies that JDBC drivers are properly configured.
+
+        :param jdbc_url: JDBC connection URL
+        :param properties: JDBC connection properties (user, password, driver)
+        :param table_or_query: Table name or SQL query wrapped in parentheses
+        :returns: Tuple of (success, message)
+        """
+        try:
+            spark = self.get_spark_session()
+
+            # Attempt JDBC read
+            df = (
+                spark.read.format("jdbc")
+                .option("url", jdbc_url)
+                .option("dbtable", table_or_query)
+                .options(**properties)
+                .load()
+            )
+
+            # Force evaluation
+            row_count = df.count()
+            return (True, f"JDBC read successful ({row_count} rows)")
+        except Exception as e:
+            error_msg = str(e)
+            # Provide helpful error messages for common issues
+            if "ClassNotFoundException" in error_msg:
+                return (False, f"JDBC driver not found: {error_msg}")
+            elif "No suitable driver" in error_msg:
+                return (False, f"JDBC driver not loaded: {error_msg}")
+            elif "Authentication" in error_msg.lower() or "password" in error_msg.lower():
+                return (False, f"Authentication failed: {error_msg}")
+            else:
+                return (False, f"JDBC test failed: {error_msg}")

dbt/compute/strategies/dataproc.py

@@ -0,0 +1,207 @@
+"""
+GCP Dataproc Spark Connection Strategy
+
+Provides connection to Google Cloud Dataproc Spark clusters.
+
+v0.5.98: New strategy for GCP Dataproc clusters with Maven-based JAR provisioning.
+
+Configuration:
+{
+    "project": "my-gcp-project",         # Required: GCP project ID
+    "region": "us-central1",             # Required: Dataproc region
+    "cluster": "my-dataproc-cluster",    # Required: Cluster name
+    "spark.driver.memory": "4g",         # Optional: driver memory
+    "spark.executor.memory": "8g",       # Optional: executor memory
+}
+
+Requirements:
+- GCP Dataproc cluster must be running
+- gcloud SDK configured (gcloud auth login)
+- Dataproc connector or direct YARN access
+
+Cost Estimate:
+- Typical 5-node Dataproc cluster: ~$1.00/hr (n1-standard-4 instances)
+- Dataproc pricing includes Spark/Hadoop runtime at no extra cost
+"""
+
+from typing import Any, Dict, Optional, Set, Tuple
+
+from dbt.compute.strategies.base import BaseConnectionStrategy
+from dbt_common.exceptions import DbtRuntimeError
+
+try:
+    from pyspark.sql import SparkSession
+
+    PYSPARK_AVAILABLE = True
+except ImportError:
+    PYSPARK_AVAILABLE = False
+    SparkSession = None
+
+
+class DataprocStrategy(BaseConnectionStrategy):
+    """
+    GCP Dataproc Spark cluster connection strategy.
+
+    Connects to Dataproc clusters using YARN as the resource manager.
+    Uses spark.jars.packages for JDBC JAR provisioning.
+    """
+
+    def validate_config(self) -> None:
+        """
+        Validate Dataproc strategy configuration.
+
+        Required:
+        - project: GCP project ID
+        - region: Dataproc region
+        - cluster: Cluster name
+
+        :raises DbtRuntimeError: If configuration is invalid
+        """
+        if not isinstance(self.config, dict):
+            raise DbtRuntimeError(
+                f"Dataproc config must be a dictionary, got {type(self.config)}"
+            )
+
+        # Check required fields
+        required_fields = ["project", "region", "cluster"]
+        missing = [f for f in required_fields if f not in self.config]
+        if missing:
+            raise DbtRuntimeError(
+                f"Dataproc config missing required fields: {', '.join(missing)}"
+            )
+
+    def get_spark_session(self, adapter_types: Optional[Set[str]] = None) -> SparkSession:
+        """
+        Create Spark session connected to Dataproc cluster.
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Initialized SparkSession connected to Dataproc
+        :raises DbtRuntimeError: If session creation fails
+        """
+        if not PYSPARK_AVAILABLE:
+            raise DbtRuntimeError("PySpark is not available. Install it with: pip install pyspark")
+
+        try:
+            builder = SparkSession.builder.appName(self.app_name)
+
+            # Set YARN master for Dataproc
+            builder = builder.master("yarn")
+
+            # Get JDBC JAR config
+            if adapter_types is None:
+                from dbt.compute.jar_provisioning import get_required_adapter_types
+                adapter_types = get_required_adapter_types()
+
+            if adapter_types:
+                jar_config = self.get_jar_provisioning_config(adapter_types)
+                for key, value in jar_config.items():
+                    builder = builder.config(key, value)
+
+            # Apply user-provided configs
+            for key, value in self.config.items():
+                if key not in ("project", "region", "cluster"):
+                    builder = builder.config(key, value)
+
+            # Default Dataproc optimizations
+            default_configs = {
+                "spark.submit.deployMode": "client",
+                "spark.dynamicAllocation.enabled": "true",
+                "spark.sql.execution.arrow.pyspark.enabled": "true",
+            }
+            for key, value in default_configs.items():
+                if key not in self.config:
+                    builder = builder.config(key, value)
+
+            # Create session
+            spark = builder.getOrCreate()
+            spark.sparkContext.setLogLevel("WARN")
+
+            return spark
+
+        except Exception as e:
+            error_msg = str(e)
+            if "Connection refused" in error_msg:
+                raise DbtRuntimeError(
+                    f"Cannot connect to Dataproc cluster '{self.config.get('cluster')}'. "
+                    f"Ensure the cluster is running. Error: {error_msg}"
+                ) from e
+            raise DbtRuntimeError(f"Failed to create Dataproc Spark session: {error_msg}") from e
+
+    def close(self, spark: Optional[SparkSession]) -> None:
+        """
+        Clean up Spark session.
+
+        For Dataproc, we stop the application but the cluster continues running.
+
+        :param spark: SparkSession to clean up
+        """
+        if spark:
+            try:
+                spark.stop()
+            except Exception:
+                pass  # Best effort cleanup
+
+    def estimate_cost(self, duration_minutes: float) -> float:
+        """
+        Estimate cost for Dataproc execution.
+
+        Based on typical 5-node Dataproc cluster with n1-standard-4 instances.
+
+        :param duration_minutes: Estimated query duration in minutes
+        :returns: Estimated cost in USD
+        """
+        # Typical Dataproc cluster: 5x n1-standard-4 @ ~$0.19/hr each = ~$0.95/hr total
+        # Plus Dataproc fee: $0.01/vCPU/hr = ~$0.20/hr for 20 vCPUs
+        hourly_cost = 1.15
+        hours = duration_minutes / 60.0
+        return round(hourly_cost * hours, 2)
+
+    def get_platform_name(self) -> str:
+        """Get platform name."""
+        return "dataproc"
+
+    def get_jar_provisioning_config(self, adapter_types: Set[str]) -> Dict[str, str]:
+        """
+        Get Spark config for JDBC JAR provisioning using Maven coordinates.
+
+        Dataproc clusters download JDBC drivers from Maven Central at session startup.
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Dictionary with spark.jars.packages config
+        """
+        from dbt.compute.jar_provisioning import RemoteJARProvisioning
+
+        provisioning = RemoteJARProvisioning()
+        return provisioning.get_spark_config(adapter_types)
+
+    def test_connectivity(self) -> Tuple[bool, str]:
+        """
+        Test connectivity to Dataproc cluster.
+
+        :returns: Tuple of (success, message)
+        """
+        if not PYSPARK_AVAILABLE:
+            return (False, "PySpark not installed")
+
+        try:
+            spark = self.get_spark_session()
+            spark.sql("SELECT 1 AS test").collect()
+            return (True, "Dataproc session created and SQL test passed")
+        except Exception as e:
+            error_msg = str(e)
+            if "Connection refused" in error_msg:
+                return (False, "Cannot connect to Dataproc cluster (connection refused)")
+            return (False, f"Dataproc connection failed: {e}")
+
+    def get_cluster_info(self) -> Dict[str, Any]:
+        """
+        Get information about the Dataproc configuration.
+
+        :returns: Dictionary with cluster metadata
+        """
+        return {
+            "platform": "dataproc",
+            "project": self.config.get("project", "unknown"),
+            "region": self.config.get("region", "unknown"),
+            "cluster": self.config.get("cluster", "unknown"),
+        }

dbt/compute/strategies/emr.py

@@ -0,0 +1,203 @@
+"""
+AWS EMR (Elastic MapReduce) Spark Connection Strategy
+
+Provides connection to AWS EMR Spark clusters via YARN.
+
+v0.5.98: New strategy for AWS EMR clusters with Maven-based JAR provisioning.
+
+Configuration:
+{
+    "master": "yarn",                    # Required: YARN resource manager
+    "spark.submit.deployMode": "client", # Optional: client or cluster
+    "spark.driver.memory": "4g",         # Optional: driver memory
+    "spark.executor.memory": "8g",       # Optional: executor memory
+    "spark.executor.instances": "4",     # Optional: number of executors
+}
+
+Requirements:
+- AWS EMR cluster must be running
+- AWS credentials configured (aws configure or IAM role)
+- Spark must be accessible from client machine (e.g., via SSH tunnel or VPN)
+
+Cost Estimate:
+- Typical 5-node EMR cluster: ~$1.20/hr (m5.xlarge instances)
+- On-demand pricing varies by instance type and region
+"""
+
+from typing import Any, Dict, Optional, Set, Tuple
+
+from dbt.compute.strategies.base import BaseConnectionStrategy
+from dbt_common.exceptions import DbtRuntimeError
+
+try:
+    from pyspark.sql import SparkSession
+
+    PYSPARK_AVAILABLE = True
+except ImportError:
+    PYSPARK_AVAILABLE = False
+    SparkSession = None
+
+
+class EMRStrategy(BaseConnectionStrategy):
+    """
+    AWS EMR Spark cluster connection strategy.
+
+    Connects to EMR clusters using YARN as the resource manager.
+    Uses spark.jars.packages for JDBC JAR provisioning.
+    """
+
+    def validate_config(self) -> None:
+        """
+        Validate EMR strategy configuration.
+
+        Required:
+        - master: Must be "yarn" for EMR
+
+        :raises DbtRuntimeError: If configuration is invalid
+        """
+        if not isinstance(self.config, dict):
+            raise DbtRuntimeError(
+                f"EMR config must be a dictionary, got {type(self.config)}"
+            )
+
+        # Check master is yarn
+        master = self.config.get("master", "")
+        if master.lower() != "yarn":
+            raise DbtRuntimeError(
+                f"EMR config requires master='yarn', got: {master}"
+            )
+
+    def get_spark_session(self, adapter_types: Optional[Set[str]] = None) -> SparkSession:
+        """
+        Create Spark session connected to EMR cluster via YARN.
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Initialized SparkSession connected to EMR
+        :raises DbtRuntimeError: If session creation fails
+        """
+        if not PYSPARK_AVAILABLE:
+            raise DbtRuntimeError("PySpark is not available. Install it with: pip install pyspark")
+
+        try:
+            builder = SparkSession.builder.appName(self.app_name)
+
+            # Set YARN master
+            builder = builder.master("yarn")
+
+            # Get JDBC JAR config
+            if adapter_types is None:
+                from dbt.compute.jar_provisioning import get_required_adapter_types
+                adapter_types = get_required_adapter_types()
+
+            if adapter_types:
+                jar_config = self.get_jar_provisioning_config(adapter_types)
+                for key, value in jar_config.items():
+                    builder = builder.config(key, value)
+
+            # Apply user-provided configs
+            for key, value in self.config.items():
+                if key != "master":  # master already set
+                    builder = builder.config(key, value)
+
+            # Default EMR optimizations
+            default_configs = {
+                "spark.submit.deployMode": "client",
+                "spark.dynamicAllocation.enabled": "true",
+                "spark.sql.execution.arrow.pyspark.enabled": "true",
+            }
+            for key, value in default_configs.items():
+                if key not in self.config:
+                    builder = builder.config(key, value)
+
+            # Create session
+            spark = builder.getOrCreate()
+            spark.sparkContext.setLogLevel("WARN")
+
+            return spark
+
+        except Exception as e:
+            error_msg = str(e)
+            if "Connection refused" in error_msg:
+                raise DbtRuntimeError(
+                    f"Cannot connect to EMR cluster. Ensure the cluster is running "
+                    f"and accessible from this machine. Error: {error_msg}"
+                ) from e
+            raise DbtRuntimeError(f"Failed to create EMR Spark session: {error_msg}") from e
+
+    def close(self, spark: Optional[SparkSession]) -> None:
+        """
+        Clean up Spark session.
+
+        For EMR, we stop the application but the cluster continues running.
+
+        :param spark: SparkSession to clean up
+        """
+        if spark:
+            try:
+                spark.stop()
+            except Exception:
+                pass  # Best effort cleanup
+
+    def estimate_cost(self, duration_minutes: float) -> float:
+        """
+        Estimate cost for EMR execution.
+
+        Based on typical 5-node EMR cluster with m5.xlarge instances.
+
+        :param duration_minutes: Estimated query duration in minutes
+        :returns: Estimated cost in USD
+        """
+        # Typical EMR cluster: 5x m5.xlarge @ ~$0.24/hr each = ~$1.20/hr total
+        hourly_cost = 1.20
+        hours = duration_minutes / 60.0
+        return round(hourly_cost * hours, 2)
+
+    def get_platform_name(self) -> str:
+        """Get platform name."""
+        return "emr"
+
+    def get_jar_provisioning_config(self, adapter_types: Set[str]) -> Dict[str, str]:
+        """
+        Get Spark config for JDBC JAR provisioning using Maven coordinates.
+
+        EMR clusters download JDBC drivers from Maven Central at session startup.
+
+        :param adapter_types: Set of adapter types that need JDBC drivers
+        :returns: Dictionary with spark.jars.packages config
+        """
+        from dbt.compute.jar_provisioning import RemoteJARProvisioning
+
+        provisioning = RemoteJARProvisioning()
+        return provisioning.get_spark_config(adapter_types)
+
+    def test_connectivity(self) -> Tuple[bool, str]:
+        """
+        Test connectivity to EMR cluster.
+
+        :returns: Tuple of (success, message)
+        """
+        if not PYSPARK_AVAILABLE:
+            return (False, "PySpark not installed")
+
+        try:
+            spark = self.get_spark_session()
+            spark.sql("SELECT 1 AS test").collect()
+            return (True, "EMR session created and SQL test passed")
+        except Exception as e:
+            error_msg = str(e)
+            if "Connection refused" in error_msg:
+                return (False, "Cannot connect to EMR cluster (connection refused)")
+            return (False, f"EMR connection failed: {e}")
+
+    def get_cluster_info(self) -> Dict[str, Any]:
+        """
+        Get information about the EMR configuration.
+
+        :returns: Dictionary with cluster metadata
+        """
+        return {
+            "platform": "emr",
+            "master": self.config.get("master", "yarn"),
+            "deploy_mode": self.config.get("spark.submit.deployMode", "client"),
+            "executor_instances": self.config.get("spark.executor.instances", "dynamic"),
+        }