feat(athena): add Apache Spark 3.5 support via Spark Connect

[dtaniwaki]

resolves #1854
docs N/A

Problem

Athena workgroups configured with Apache Spark engine version 3.5 cannot run dbt Python models because the Calculation API (StartCalculationExecution) is not supported in Spark 3.5. Additionally, Spark 3.5 no longer accepts CoordinatorDpuSize, DefaultExecutorDpuSize, or SparkProperties in EngineConfiguration, requiring a different configuration format (Classifications).

Apache Spark 3.5 on Athena upgrades the runtime from Spark 3.2.1 to 3.5.6, bringing performance improvements such as avoiding unnecessary shuffles in Storage-Partitioned Joins when partition keys mismatch but join expressions are compatible (release notes).

Solution

Adds an alternative code execution path using Spark Connect via GetSessionEndpoint, which is available in Spark 3.5. Users opt in by setting spark_engine_version: "3.5" in their model or profile config.

Key changes:

• Spark Connect execution path: Submits code via Spark Connect instead of the Calculations API, with a ChannelBuilder that auto-refreshes the Athena AuthToken before expiry.
• Spark Connect session pool: Thread-safe singleton keyed by (invocation_id, fingerprint) that reuses sessions across models within an invocation. Required because Spark Connect binds a persistent gRPC channel to a single session, so the existing AthenaSparkSessionManager can't be reused.
• Engine configuration: For Spark 3.5, drops the unsupported keys and moves Spark properties into Classifications (spark-defaults).
• Transient error retry: Retries transient Spark Connect errors (credential/region loading, gRPC pool shutdown, inactive sessions, session-quota exhaustion, retriable gRPC status codes) up to spark_connect_max_retries times with exponential backoff, terminating the failed session each time.
• DPU budget tracking: Reserves DPUs per session (min(MaxConcurrentDpus, maxExecutors + 1)) against an account-wide budget (default 60, the L-E20AD6B8 quota) so the pool throttles before AWS rejects with Maximum allowed sessions. Region-level required capacity not being available errors fall into the same backoff path.

How it works

dbt's own thread pool drives parallelism; the Spark Connect session pool just decides whether each thread gets a fresh session, a warm one, or has to wait.

There are two communication channels with Athena: a boto3 control plane that the pool uses to create / look up / terminate sessions, and a gRPC data plane that PySpark clients use to actually run code against the Spark runtime.

flowchart TB
    subgraph Client["dbt host process"]
        direction TB
        subgraph Threads["dbt thread pool (profile: threads: N)"]
            direction LR
            T1[dbt thread 1]
            T2[dbt thread 2]
            T3[dbt thread N]
            T1 ~~~ T2 ~~~ T3
        end
        Pool["Spark Connect session pool"]
        subgraph Clients["PySpark clients (up to max_sessions per fingerprint)"]
            direction LR
            PS1["SparkSession 1"]
            PS2["SparkSession M"]
        end
        Threads -->|"acquire / release<br/>(≤ session_concurrency<br/>models per session)"| Pool
        Pool --> PS1 & PS2
    end

    subgraph AWS["Athena (AWS-managed)"]
        direction TB
        Ctrl["Athena control plane"]
        subgraph WG["Athena Spark workgroup (engine version 3.5)"]
            direction TB
            subgraph S1["session 1"]
                direction TB
                SC1["Spark Connect server<br/>session 1 endpoint"]
                Spark1["Spark runtime<br/>Apache Spark 3.5"]
                SC1 --> Spark1
            end
            subgraph S2["session M"]
                direction TB
                SC2["Spark Connect server<br/>session M endpoint"]
                Spark2["Spark runtime<br/>Apache Spark 3.5"]
                SC2 --> Spark2
            end
        end
    end

    Pool -. "boto3<br/>(create session,<br/>fetch endpoint URL +<br/>initial auth token)" .-> Ctrl
    PS1 & PS2 -. "boto3<br/>(refresh auth token<br/>before expiry)" .-> Ctrl
    PS1 -->|gRPC + current auth token| SC1
    PS2 -->|gRPC + current auth token| SC2

Loading

Tuning rules of thumb:

• Throughput: keep max_sessions × session_concurrency ≥ dbt threads for the dominant fingerprint, otherwise dbt threads will sit idle waiting for the pool.
• DPU isolation vs sharing: a session's DPU budget is fixed by its engine config and is shared by all models running on it. Set session_concurrency: 1 (and scale max_sessions up) when each model needs the full DPU budget to itself. Raise session_concurrency when the per-model DPU footprint is small and you'd rather amortize session-startup cost across more models.
• DPU budget: defaults to the account-wide Athena Spark 3.5 DPU quota (60, L-E20AD6B8). Lower spark_connect_dpu_budget when multiple dbt processes share one AWS account so they don't fight for the same quota; raise it if AWS has raised your account quota.

New options

Option	Scope	Default	Description
spark_engine_version	profile / model config	"3"	Set to "3.5" to opt into Spark Connect execution.
spark_connect_max_sessions	profile	4	Maximum concurrent Spark Connect sessions per fingerprint (engine config + workgroup + engine version). Raises throughput when multiple models share a fingerprint but would exceed per-session concurrency.
spark_connect_session_concurrency	profile	1	Maximum in-flight models sharing a single Spark Connect session. 1 isolates each model; larger values trade isolation for lower session-startup overhead.
spark_connect_dpu_budget	profile	60	Account-wide DPU budget. The pool waits to start a session whose reserved DPUs would push the live total above this, throttling before AWS rejects with Maximum allowed sessions.
spark_connect_pool_acquire_timeout	profile	21600	Maximum cumulative seconds (across retries) a model will wait to acquire a Spark Connect session from the pool before failing. Default is 6 hours; acts as a safety net for DPU exhaustion or a stuck pool.
spark_connect_max_retries	profile	3	Maximum number of retries for transient Spark Connect errors (each retry uses a fresh session). 0 disables retries (single attempt).

The existing per-model timeout config (default 43200s = 12h) continues to apply on the Spark Connect path as the per-attempt Spark execution budget, so model-level execution timeouts do not need a new knob.

Example configuration:

# profiles.yml
my_athena_profile:
  target: dev
  outputs:
    dev:
      type: athena
      spark_work_group: my-spark-workgroup
      spark_engine_version: "3.5"
      spark_connect_max_sessions: 4               # optional, default 4
      spark_connect_session_concurrency: 1        # optional, default 1
      spark_connect_dpu_budget: 60                # optional, default 60 (account quota)
      spark_connect_pool_acquire_timeout: 21600   # optional, default 21600 (6h)
      spark_connect_max_retries: 3                # optional, default 3
      # ... other settings

# dbt_project.yml or model config
models:
  my_project:
    python_models:
      +spark_engine_version: "3.5"
      +timeout: 3600   # optional, existing per-model execution timeout (seconds), default 43200

Optional dependency for Spark Connect:

pip install "dbt-athena[spark_connect]"

Test infrastructure (Terraform)

The functional test (test_spark_connect_python_submissions.py) is gated on DBT_TEST_ATHENA_SPARK_WORK_GROUP and requires a workgroup with engine version Apache Spark version 3.5 plus a dedicated execution role. The Terraform below provisions the minimal AWS resources needed to run the test against a real Athena account.

Terraform (click to expand)

terraform {
  required_version = ">= 1.5"

  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = "~> 5.0"
    }
  }
}

provider "aws" {
  profile = "<your-aws-profile>"
  region  = "<your-region>"
}

locals {
  account_id = "<your-account-id>"
  region     = "<your-region>"
  name       = "<your-spark-test-name>"
  bucket     = "<your-test-output-bucket>"
  kms_key_id = "<your-kms-key-id>"

  workgroup_arn = "arn:aws:athena:${local.region}:${local.account_id}:workgroup/${local.name}"
}

data "aws_iam_policy_document" "trust" {
  statement {
    effect  = "Allow"
    actions = ["sts:AssumeRole"]

    principals {
      type        = "Service"
      identifiers = ["athena.amazonaws.com"]
    }

    condition {
      test     = "StringEquals"
      variable = "aws:SourceAccount"
      values   = [local.account_id]
    }

    condition {
      test     = "ArnEquals"
      variable = "aws:SourceArn"
      values   = [local.workgroup_arn]
    }
  }
}

data "aws_iam_policy_document" "permissions" {
  statement {
    sid     = "GlueCatalogFullAccess"
    effect  = "Allow"
    actions = ["glue:*"]
    resources = [
      "arn:aws:glue:${local.region}:${local.account_id}:catalog",
      "arn:aws:glue:${local.region}:${local.account_id}:database/*",
      "arn:aws:glue:${local.region}:${local.account_id}:table/*",
      "arn:aws:glue:${local.region}:${local.account_id}:userDefinedFunction/*",
    ]
  }

  statement {
    sid     = "S3OutputBucket"
    effect  = "Allow"
    actions = ["s3:*"]
    resources = [
      "arn:aws:s3:::${local.bucket}",
      "arn:aws:s3:::${local.bucket}/*",
    ]
  }

  statement {
    sid       = "LakeFormation"
    effect    = "Allow"
    actions   = ["lakeformation:GetDataAccess"]
    resources = ["*"]
  }

  statement {
    sid    = "AthenaSelf"
    effect = "Allow"
    actions = [
      "athena:GetWorkGroup",
      "athena:StartQueryExecution",
      "athena:GetQueryExecution",
      "athena:GetQueryResults",
      "athena:StopQueryExecution",
      "athena:ListWorkGroups",
    ]
    resources = ["*"]
  }

  statement {
    sid    = "KmsForOutputBucket"
    effect = "Allow"
    actions = [
      "kms:Decrypt",
      "kms:Encrypt",
      "kms:GenerateDataKey",
      "kms:DescribeKey",
    ]
    resources = ["arn:aws:kms:${local.region}:${local.account_id}:key/${local.kms_key_id}"]
  }
}

resource "aws_iam_role" "spark_test" {
  name               = local.name
  description        = "Spark Connect functional test execution role for dbt-athena (Apache Spark 3.5)"
  assume_role_policy = data.aws_iam_policy_document.trust.json
}

resource "aws_iam_role_policy" "spark_test_inline" {
  name   = "${local.name}-inline"
  role   = aws_iam_role.spark_test.name
  policy = data.aws_iam_policy_document.permissions.json
}

resource "aws_athena_workgroup" "spark_test" {
  name        = local.name
  description = "Spark Connect functional test workgroup for dbt-athena (Apache Spark 3.5)"

  configuration {
    enforce_workgroup_configuration = false
    execution_role                  = aws_iam_role.spark_test.arn

    engine_version {
      selected_engine_version = "Apache Spark version 3.5"
    }
  }
}

output "execution_role_arn" {
  value = aws_iam_role.spark_test.arn
}

output "workgroup_name" {
  value = aws_athena_workgroup.spark_test.name
}

Checklist

• I have read the contributing guide and understand what's expected of me
• I have run this code in development and it appears to resolve the stated issue
• This PR includes tests, or tests are not required/relevant for this PR
• This PR has no interface changes (e.g. macros, cli, logs, json artifacts, config files, adapter interface, etc) or this PR has already received feedback and approval from Product or DX

[Copilot]

Pull request overview

Adds an opt-in execution path for Athena Spark engine version 3.5 Python models by switching from the unsupported Calculations API to Spark Connect (GetSessionEndpoint), while updating engine configuration handling to match Spark 3.5 requirements.

Changes:

• Add Spark Connect submission path (endpoint polling + auth token refresh) when spark_engine_version: "3.5".
• Update Spark session EngineConfiguration generation to use Classifications for Spark 3.5 and adjust expected keys.
• Extend adapter response to include DPU execution timing, plus add/adjust unit tests and macros for Spark Connect DataFrame typing and Iceberg write behavior.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 8 comments.

Show a summary per file

File	Description
dbt-athena/tests/unit/test_session.py	Updates tests for session IDs now being str instead of UUID.
dbt-athena/tests/unit/test_python_submissions.py	Adds routing and Spark Connect unit tests; adjusts expectations for statistics in results.
dbt-athena/tests/unit/test_config.py	Adds Spark 3.5-specific tests and expected keys (including Classifications).
dbt-athena/src/dbt/include/athena/macros/materializations/models/table/create_table_as.sql	Passes additional optional args to python materialization (incl. spark_engine_version) and gates Iceberg CTAS path.
dbt-athena/src/dbt/include/athena/macros/adapters/python_submissions.sql	Uses Spark Connect DataFrame type for 3.5 and adds optional Iceberg writeTo() materialization branch.
dbt-athena/src/dbt/adapters/athena/session.py	Switches global session tracking maps to use string session IDs; adds spark_managed_logging param.
dbt-athena/src/dbt/adapters/athena/python_submissions.py	Implements Spark Connect submission path, endpoint polling, and auth token refresh ChannelBuilder.
dbt-athena/src/dbt/adapters/athena/impl.py	Returns AthenaAdapterResponse for python submissions and attempts to surface DPU execution millis.
dbt-athena/src/dbt/adapters/athena/connections.py	Extends AthenaAdapterResponse with dpu_execution_in_millis.
dbt-athena/src/dbt/adapters/athena/config.py	Adds Spark 3.5 EngineConfiguration behavior (no DPU sizes/SparkProperties; use Classifications).
dbt-athena/pyproject.toml	Adds optional dependency group spark_connect for pyspark[connect].

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[dtaniwaki]

I'm making this PR a draft because it's broken by commits which I added recently 🙇

[Copilot]

Pull request overview

Copilot reviewed 25 out of 26 changed files in this pull request and generated 1 comment.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

_wait_for_endpoint() uses deadline_seconds = min(self.timeout, _ENDPOINT_READY_TIMEOUT_SECONDS), which ignores time already spent in the overall attempt (session acquisition, prior retries/backoff, etc.). This can cause the endpoint wait to exceed the remaining timeout budget, violating the contract that self.timeout covers endpoint readiness + execution. Consider passing the per-attempt remaining budget into _wait_for_endpoint() (or computing an absolute deadline from start_time) and using min(remaining, _ENDPOINT_READY_TIMEOUT_SECONDS) for the endpoint wait loop.

[dtaniwaki]

@iconara @colin-rogers-dbt @nicor88 I think the Spark 3.5 support is ready. I'm running tens of spark models in parallel and it works fine. I also put a terraform example so you can try the feature with the functional tests added in this PR. Would you review the PR?

[dtaniwaki]

I made some refactoring, but the logic is not changed.

[dtaniwaki]

I updated the code with many improvements and squashed the commits. I believe the PR is enough ready for other developer's review.

[dtaniwaki]

Added support for assume_role_arn on Spark Connect Python models: the model body is exec()-ed client-side, so a bare boto3.client(...) in the model previously ran under the caller identity instead of the assumed role (the classic calculation path already runs remotely under the workgroup ExecutionRole). The submitter now installs the assumed session as boto3's process default before exec, so model-level boto3 calls run under assume_role_arn.

[colin-k-rogers]

@dtaniwaki I'm no longer a maintainer/employee so can't help you here

@aahel would you have bandwidth to take a look?

[dtaniwaki]

Thanks for letting me know, and thanks for all your work! No problem at all.

@aahel Thanks in advance if you have time to take a look!

[GrumpyDBA]

Thank you for taking on this task @dtaniwaki

@aahel If it is relevant - it's actually a real and present issue for me currently. AWS have not deployed all Athena versions in all regions. We cannot use Python with DBT at all in some regions. e.g. London:

Here's a sample of available versions London vs Frankfurt

REGION	SELECTED_VERSION
eu-west-2	Apache Spark version 3.5
eu-west-2	AUTO
eu-west-2	Athena engine version 3
eu-central-1	Apache Spark version 3.5
eu-central-1	AUTO
eu-central-1	Athena engine version 3
eu-central-1	PySpark engine version 3

[skillicinski]

Same issue for us in eu-north-1!