|
| 1 | +import os |
| 2 | + |
| 3 | +import pytest |
| 4 | + |
| 5 | +from dbt.tests.adapter.python_model.test_python_model import ( |
| 6 | + BasePythonModelTests, |
| 7 | + basic_sql, |
| 8 | + schema_yml, |
| 9 | + second_sql, |
| 10 | +) |
| 11 | + |
| 12 | +# spark_engine_version="3.5" routes execution through Spark Connect. |
| 13 | +# DBT_TEST_ATHENA_SPARK_WORK_GROUP must point to a workgroup whose |
| 14 | +# engine version is set to Apache Spark 3.5. |
| 15 | +spark_connect_python = """ |
| 16 | +def model(dbt, _): |
| 17 | + dbt.config( |
| 18 | + materialized='table', |
| 19 | + spark_engine_version='3.5', |
| 20 | + ) |
| 21 | + df = dbt.ref("my_sql_model") |
| 22 | + df2 = dbt.ref("my_versioned_sql_model", v=1) |
| 23 | + df3 = dbt.ref("my_versioned_sql_model", version=1) |
| 24 | + df4 = dbt.ref("test", "my_versioned_sql_model", v=1) |
| 25 | + df5 = dbt.ref("test", "my_versioned_sql_model", version=1) |
| 26 | + df6 = dbt.source("test_source", "test_table") |
| 27 | + df = df.limit(2) |
| 28 | + return df |
| 29 | +""" |
| 30 | + |
| 31 | + |
| 32 | +@pytest.mark.skipif( |
| 33 | + not os.getenv("DBT_TEST_ATHENA_SPARK_WORK_GROUP"), |
| 34 | + reason="DBT_TEST_ATHENA_SPARK_WORK_GROUP must point to a Spark 3.5 workgroup.", |
| 35 | +) |
| 36 | +class TestSparkConnectPythonModel(BasePythonModelTests): |
| 37 | + @pytest.fixture(scope="class") |
| 38 | + def project_config_update(self): |
| 39 | + # Athena defaults SQL models to views, but Spark's HiveExternalCatalog |
| 40 | + # cannot resolve Athena views (no S3 location), so dbt.ref() blows up |
| 41 | + # with "Can not create a Path from an empty string". Force tables. |
| 42 | + return {"models": {"+materialized": "table"}} |
| 43 | + |
| 44 | + @pytest.fixture(scope="class") |
| 45 | + def models(self): |
| 46 | + return { |
| 47 | + "schema.yml": schema_yml, |
| 48 | + "my_sql_model.sql": basic_sql, |
| 49 | + "my_versioned_sql_model_v1.sql": basic_sql, |
| 50 | + "my_python_model.py": spark_connect_python, |
| 51 | + "second_sql_model.sql": second_sql, |
| 52 | + } |
0 commit comments