feat(athena): add use_iceberg_write_to config for Iceberg Python models

[dtaniwaki]

Thank you for maintaining dbt-athena — I'd appreciate your review on this change.

resolves #1882
docs N/A

Problem

Iceberg Python (Spark) models currently have two write paths, both with limitations:

1. saveAsTable does not support Iceberg-native partition transforms (bucket(), day(), etc.) via the DataFrameWriter API.
2. The spark_ctas SQL path doesn't leverage Spark's DataFrameWriterV2 capabilities like tableProperty().

Additionally, for Iceberg tables the table materialization always uses the __ha intermediate table flow (CTAS to __ha → rename), regardless of the ha config. When a previous Spark run fails, the leftover __ha table causes TABLE_OR_VIEW_ALREADY_EXISTS on retry.

Solution

Add a use_iceberg_write_to model config that uses writeTo().createOrReplace() (DataFrameWriterV2 API). This enables Iceberg-native partition transforms and atomic replacement without the __ha intermediate table.

In table.sql, when use_iceberg_write_to is enabled, the model writes directly to target_relation, skipping the __ha → rename flow. Leftover __ha / __bkp tables from previous HA-flow failures are cleaned up.

{{ config(
    materialized='table',
    table_type='iceberg',
    use_iceberg_write_to=True,
    partitioned_by=['day(created_at)', 'bucket(user_id, 256)'],
) }}

Tests

• Unit (tests/unit/test_py_save_table_as.py): renders athena__py_save_table_as end-to-end with jinja2.FileSystemLoader (same pattern as test_get_partition_batches.py) and asserts on the generated Python — writeTo/createOrReplace branch selection, tojson escaping of extra_table_properties and partition expressions, fall-through to spark_ctas and saveAsTable. Also exec's the inline _parse_iceberg_partition against a stub pyspark.sql.functions to cover transform dispatch and bucket/truncate arity validation.
• Functional (tests/functional/adapter/test_use_iceberg_write_to.py, gated on DBT_TEST_ATHENA_SPARK_WORK_GROUP): runs Iceberg Python models end-to-end against a real Athena Spark workgroup. Covers partitioned + unpartitioned writes, idempotent createOrReplace, table_properties propagation, and the compiler-error path when use_iceberg_write_to=True is combined with table_type != 'iceberg'.

Test infrastructure (Terraform)

The functional test requires an Athena Spark workgroup with Iceberg support. This PR predates the Spark 3.5 / Spark Connect work in #1874, so the workgroup uses Athena's PySpark engine version 3 (Spark 3.2.1) — the engine the Calculations API targets. The Terraform below provisions the minimal AWS resources.

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 functional test execution role for dbt-athena"
  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 functional test workgroup for dbt-athena"

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

    engine_version {
      selected_engine_version = "PySpark engine version 3"
    }

    result_configuration {
      output_location = "s3://${local.bucket}/spark-output/"

      encryption_configuration {
        encryption_option = "SSE_KMS"
        kms_key_arn       = "arn:aws:kms:${local.region}:${local.account_id}:key/${local.kms_key_id}"
      }
    }
  }
}

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 a new config-driven write path for Iceberg Python table models that uses Spark DataFrameWriterV2 writeTo().createOrReplace() to support Iceberg-native partition transforms and avoid HA __ha retry failures.

Changes:

• Introduces use_iceberg_write_to config and threads it through the Python submission path.
• Adds a Python write branch that uses writeTo().createOrReplace() and parses Iceberg partition transform expressions.
• Updates Iceberg table materialization to skip the __ha intermediate table when use_iceberg_write_to is enabled and cleans up leftover __ha/__bkp relations.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 5 comments.

File	Description
dbt-athena/src/dbt/include/athena/macros/materializations/models/table/table.sql	Skips Iceberg HA swap flow for Python models when use_iceberg_write_to is enabled; drops leftover temp/backup relations.
dbt-athena/src/dbt/include/athena/macros/materializations/models/table/create_table_as.sql	Passes use_iceberg_write_to, table_type, and extra_table_properties into Python submission; disables Spark CTAS when using WriterV2.
dbt-athena/src/dbt/include/athena/macros/adapters/python_submissions.sql	Implements the WriterV2 writeTo().createOrReplace() path with transform parsing and table property support.

---

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

[dtaniwaki]

@iconara @colin-rogers-dbt This improvement also helps avoid creating unnecessary temporary tables in iceberg, in addition to #1830. Would you review it?

[Copilot]

Pull request overview

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

---

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

[dtaniwaki]

Fixed in commit ddf8f8b. Both prop_name and prop_value now go through | string | tojson (the | string coerces non-string values like ints/booleans before | tojson handles JSON escaping). location is also normalized to {{ (location ~ "/") | tojson }} for consistency.

[dtaniwaki]

Fixed in commit ddf8f8b. bucket and truncate now share a branch that validates len(args) == 2 and raises ValueError("... requires 2 arguments (column, n), got: <expr>") when the user supplies the wrong arity, instead of falling through to an IndexError on args[1]. Coverage: tests/unit/test_py_save_table_as.py::TestParseIcebergPartition::test_bucket_with_missing_arg_raises_clear_error (and the truncate counterpart) added in d32370c.