feat(athena): track Spark Connect calc IDs and DPU stats, unify response field, tidy comments, and add functional test for the Spark 3.5 path

Signed-off-by: Daisuke Taniwaki <daisuketaniwaki@gmail.com>

Author: dtaniwaki

dbt-athena/hatch.toml

@@ -17,9 +17,11 @@ dependencies = [
     "moto~=5.0.13",
     "pre-commit==3.7.0",
     "pyparsing~=3.1.4",
+    "pyspark[connect]>=3.5.6,<3.6",
     "pytest>=7.0,<8.0",
     "pytest-dotenv~=0.5",
     "pytest-xdist~=3.6",
+    "setuptools",
 ]
 [envs.default.scripts]
 setup = [
@@ -65,8 +67,10 @@ dependencies = [
     "ddtrace==2.3.0",
     "moto~=5.0.13",
     "pyparsing~=3.1.4",
+    "pyspark[connect]>=3.5.6,<3.6",
     "pytest>=7.0,<8.0",
     "pytest-xdist~=3.6",
+    "setuptools",
 ]
 [envs.ci.scripts]
 unit-tests = "python -m pytest tests/unit --ddtrace"
@@ -79,8 +83,10 @@ dependencies = [
     "ddtrace==2.3.0",
     "moto~=5.0.13",
     "pyparsing~=3.1.4",
+    "pyspark[connect]>=3.5.6,<3.6",
     "pytest>=7.0,<8.0",
     "pytest-xdist~=3.6",
+    "setuptools",
 ]
 [envs.cd.scripts]
 unit-tests = "python -m pytest tests/unit --ddtrace"

dbt-athena/pyproject.toml

@@ -38,6 +38,9 @@ dependencies=[
     "pydantic>=1.10,<3.0",
     "tenacity>=8.2,<10.0",
 ]
+[project.optional-dependencies]
+spark_connect = ["pyspark[connect]>=3.5.6,<3.6", "setuptools"]
+
 [project.urls]
 Homepage = "https://github.com/dbt-labs/dbt-adapters/tree/main/dbt-athena"
 Documentation = "https://docs.getdbt.com"

dbt-athena/src/dbt/adapters/athena/connections.py

@@ -48,7 +48,9 @@
 @dataclass
 class AthenaAdapterResponse(AdapterResponse):
     data_scanned_in_bytes: Optional[int] = None
+    dpu_execution_in_millis: Optional[int] = None
     spark_session_id: Optional[str] = None
+    spark_calculation_execution_ids: Optional[List[str]] = None
 
 
 @dataclass
@@ -80,26 +82,20 @@ class AthenaCredentials(Credentials):
     lf_tags_database: Optional[Dict[str, str]] = None
 
     def __post_init__(self) -> None:
-        # Surface mis-configured Spark Connect integer fields at profile-load
-        # time rather than waiting until a python model runs — the misconfig
-        # would otherwise only manifest once a Spark 3.5 model is submitted,
-        # potentially minutes into a long dbt run.
-        self._validate_positive_int_field("spark_connect_max_sessions")
-        self._validate_positive_int_field("spark_connect_session_concurrency")
-
-    def _validate_positive_int_field(self, field_name: str) -> None:
-        raw = getattr(self, field_name)
-        if raw is None:
-            return
-        try:
-            value = int(raw)
-        except (TypeError, ValueError) as e:
-            raise DbtRuntimeError(f"{field_name} must be an integer (got {raw!r}).") from e
-        if value < 1:
-            raise DbtRuntimeError(
-                f"{field_name} must be >= 1 (got {value}). Omit the field to use the default."
-            )
-        setattr(self, field_name, value)
+        # Fail fast at profile load so misconfig doesn't surface mid-run.
+        for field_name in ("spark_connect_max_sessions", "spark_connect_session_concurrency"):
+            raw = getattr(self, field_name)
+            if raw is None:
+                continue
+            try:
+                value = int(raw)
+            except (TypeError, ValueError) as e:
+                raise DbtRuntimeError(f"{field_name} must be an integer (got {raw!r}).") from e
+            if value < 1:
+                raise DbtRuntimeError(
+                    f"{field_name} must be >= 1 (got {value}). Omit the field to use the default."
+                )
+            setattr(self, field_name, value)
 
     @property
     def type(self) -> str:
@@ -362,21 +358,10 @@ def process_query_stats(cursor: AthenaCursor) -> Tuple[int, int]:
         return cursor.rowcount, cursor.data_scanned_in_bytes
 
     def cleanup_all(self) -> None:
-        # Terminate Spark Connect sessions owned by THIS invocation so DPUs
-        # are released immediately instead of waiting for idle timeout.
-        # Scoping by invocation id prevents multi-invocation hosts (dbt Cloud
-        # workers, test harnesses) from killing sessions belonging to other
-        # live invocations that share the singleton pool.
-        #
-        # ``spark_connect.session`` is imported lazily because it pulls in the
-        # pyspark runtime lookup path; importing at module load time would
-        # force every user of the Athena connection manager (including pure
-        # SQL workflows) to pay for Spark imports.
+        # Lazy import: avoid pulling pyspark for SQL-only workflows.
         from dbt_common.invocation import get_invocation_id
 
-        from dbt.adapters.athena.spark_connect.session import (
-            SparkConnectSessionPool,
-        )
+        from dbt.adapters.athena.spark_connect.session import SparkConnectSessionPool
 
         SparkConnectSessionPool().terminate_by_invocation(get_invocation_id())
         super().cleanup_all()

dbt-athena/src/dbt/adapters/athena/impl.py

@@ -9,7 +9,18 @@
 from functools import lru_cache
 from textwrap import dedent
 from threading import Lock
-from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple, Type
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Dict,
+    FrozenSet,
+    Iterable,
+    List,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+)
 from urllib.parse import urlparse
 from uuid import uuid4
 
@@ -265,7 +276,10 @@ def is_work_group_output_location_enforced(self) -> bool:
             return False
 
     def _s3_table_prefix(
-        self, s3_data_dir: Optional[str], s3_tmp_table_dir: Optional[str], is_temporary_table: bool
+        self,
+        s3_data_dir: Optional[str],
+        s3_tmp_table_dir: Optional[str],
+        is_temporary_table: bool,
     ) -> str:
         """
         Returns the root location for storing tables in S3.
@@ -320,15 +334,25 @@ def generate_s3_location(
 
         mapping = {
             S3DataNaming.UNIQUE: path.join(table_prefix, str(uuid4())),
-            S3DataNaming.TABLE: path.join(table_prefix, s3_path_table_part),  # type:ignore
+            S3DataNaming.TABLE: path.join(
+                table_prefix,
+                s3_path_table_part,  # type:ignore[arg-type]
+            ),
             S3DataNaming.TABLE_UNIQUE: path.join(
-                table_prefix, s3_path_table_part, str(uuid4())  # type:ignore
+                table_prefix,
+                s3_path_table_part,  # type:ignore[arg-type]
+                str(uuid4()),
             ),
             S3DataNaming.SCHEMA_TABLE: path.join(
-                table_prefix, schema_name, s3_path_table_part  # type:ignore
+                table_prefix,
+                schema_name,  # type:ignore[arg-type]
+                s3_path_table_part,  # type:ignore[arg-type]
             ),
             S3DataNaming.SCHEMA_TABLE_UNIQUE: path.join(
-                table_prefix, schema_name, s3_path_table_part, str(uuid4())  # type:ignore
+                table_prefix,
+                schema_name,  # type:ignore[arg-type]
+                s3_path_table_part,  # type:ignore[arg-type]
+                str(uuid4()),
             ),
         }
 
@@ -355,7 +379,9 @@ def get_glue_table(self, relation: AthenaRelation) -> Optional[GetTableResponseT
 
         try:
             table = glue_client.get_table(
-                CatalogId=catalog_id, DatabaseName=relation.schema, Name=relation.identifier
+                CatalogId=catalog_id,
+                DatabaseName=relation.schema,
+                Name=relation.identifier,
             )
         except ClientError as e:
             if e.response["Error"]["Code"] == "EntityNotFoundException":
@@ -449,7 +475,10 @@ def quote(self, identifier: str) -> str:  # type:ignore
 
     @available
     def quote_seed_column(
-        self, column: str, quote_config: Optional[bool], quote_character: Optional[str] = None
+        self,
+        column: str,
+        quote_config: Optional[bool],
+        quote_character: Optional[str] = None,
     ) -> str:
         if quote_character:
             old_value = self.quote_character
@@ -656,7 +685,8 @@ def _get_one_catalog(
                     for table in page["TableList"]:
                         catalog.extend(
                             self._get_one_table_for_catalog(
-                                table, information_schema.database  # type:ignore
+                                table,
+                                information_schema.database,  # type:ignore
                             )
                         )
             table = agate.Table.from_object(catalog)
@@ -679,7 +709,9 @@ def _get_one_catalog(
                     for table in page["TableMetadataList"]:
                         catalog.extend(
                             self._get_one_table_for_non_glue_catalog(
-                                table, schema, information_schema.database  # type:ignore
+                                table,
+                                schema,
+                                information_schema.database,  # type:ignore
                             )
                         )
             table = agate.Table.from_object(catalog)
@@ -714,7 +746,11 @@ def _get_data_catalog(self, database: str) -> Optional[DataCatalogTypeDef]:
                         config=get_boto3_config(num_retries=creds.effective_num_retries),
                     )
                 catalog_id = sts.get_caller_identity()["Account"]
-                return {"Name": database, "Type": "GLUE", "Parameters": {"catalog-id": catalog_id}}
+                return {
+                    "Name": database,
+                    "Type": "GLUE",
+                    "Parameters": {"catalog-id": catalog_id},
+                }
             with boto3_client_lock:
                 athena = client.session.client(
                     "athena",
@@ -800,7 +836,8 @@ def _get_one_catalog_by_relations(
             glue_table_definition = self.get_glue_table(_rel)
             if glue_table_definition:
                 _table_definition = self._get_one_table_for_catalog(
-                    glue_table_definition["Table"], _rel.database  # type:ignore
+                    glue_table_definition["Table"],
+                    _rel.database,  # type:ignore
                 )
                 _table_definitions.extend(_table_definition)
         table = agate.Table.from_object(_table_definitions)
@@ -1137,13 +1174,23 @@ def persist_docs_to_glue(
             )
 
     def generate_python_submission_response(self, submission_result: Any) -> AthenaAdapterResponse:
+        # ``code`` is intentionally left unset: AdapterResponse.code is
+        # reserved for driver-level status codes (SQLSTATE-style).  We
+        # surface the success/failure signal via ``_message`` instead.
         if not submission_result:
             return AthenaAdapterResponse(_message="ERROR")
         result = submission_result if isinstance(submission_result, dict) else {}
+        statistics = result.get("Statistics") or {}
+        dpu_execution_in_millis = (
+            statistics.get("DpuExecutionInMillis") if isinstance(statistics, dict) else None
+        )
         spark_session_id = result.get("SparkSessionId")
+        spark_calculation_execution_ids = result.get("SparkCalculationExecutionIds")
         return AthenaAdapterResponse(
             _message="OK",
+            dpu_execution_in_millis=dpu_execution_in_millis,
             spark_session_id=spark_session_id,
+            spark_calculation_execution_ids=spark_calculation_execution_ids,
         )
 
     @property

dbt-athena/src/dbt/adapters/athena/python_submissions.py

@@ -140,10 +140,6 @@ def submit(self, compiled_code: str) -> Any:
                 relation_name=self.relation_name,
             ).submit(compiled_code)
 
-        # Seeing an empty calculation along with main python model code calculation is submitted for almost every model
-        # Also, if not returning the result json, we are getting green ERROR messages instead of OK messages.
-        # And with this handling, the run model code in target folder every model under run folder seems to be empty
-        # Need to fix this work around solution
         if compiled_code.strip():
             while True:
                 try:
@@ -194,22 +190,23 @@ def submit(self, compiled_code: str) -> Any:
                         CalculationExecutionId=calculation_execution_id
                     )
                     result = execution_response.get("Result") or {}
+                    statistics = execution_response.get("Statistics")
+                    if statistics is not None:
+                        result["Statistics"] = statistics
                     result["SparkSessionId"] = self.session_id
+                    result["SparkCalculationExecutionIds"] = [calculation_execution_id]
                 except Exception as e:
                     LOGGER.error(f"Unable to retrieve results: Got: {e}")
-                    result = {"SparkSessionId": self.session_id}
+                    # Preserve identifiers so CloudWatch / Athena console can
+                    # still be used to debug the failed fetch.
+                    result = {
+                        "SparkSessionId": self.session_id,
+                        "SparkCalculationExecutionIds": [calculation_execution_id],
+                    }
             return result
         else:
-            # dbt submits an empty "ghost" calculation alongside every python
-            # model to keep the adapter response shape consistent.  This
-            # branch returns placeholder data without hitting Athena.
-            return {
-                "ResultS3Uri": "string",
-                "ResultType": "string",
-                "StdErrorS3Uri": "string",
-                "StdOutS3Uri": "string",
-                "SparkSessionId": self.session_id,
-            }
+            # dbt-core wraps every python model with an empty submit; skip Athena to save DPUs.
+            return {"SparkSessionId": self.session_id}
 
     def poll_until_session_idle(self) -> None:
         """