From 92a3fcff4a8000fe96bf58dd9586d6a387bb09b2 Mon Sep 17 00:00:00 2001 From: Sebastian Stan <73468824+Lindblomsebastian@users.noreply.github.com> Date: Thu, 11 Apr 2024 20:03:04 +0200 Subject: [PATCH] Add external_access_integration and secret to python model impl (#955) * Add external_access_integration and secret to python model impl * adds changie entry * Don't format external_access_integrations * Comment out SQL queries that creates network rules, external access integration and secrets * fixes test this is plural * Runs secrets by test user Turns our secrets need to live in the schema in which they will be used and don't need to be run by ACCOUNTADMIN * Fixes example code httpx has no method called stats, it's status * Fix quoting for external access integrations * Include external access integration in secret test * re-add code to create rules changed permissions of CI user * get_generic_secret_string expects a string, not object --------- Co-authored-by: Ernesto Ongaro --- .../unreleased/Features-20240402-131330.yaml | 6 ++ dbt/adapters/snowflake/impl.py | 17 ++++- tests/functional/adapter/test_python_model.py | 75 +++++++++++++++++++ 3 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 .changes/unreleased/Features-20240402-131330.yaml diff --git a/.changes/unreleased/Features-20240402-131330.yaml b/.changes/unreleased/Features-20240402-131330.yaml new file mode 100644 index 000000000..9176cca48 --- /dev/null +++ b/.changes/unreleased/Features-20240402-131330.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Add external_access_integration and secret to python models +time: 2024-04-02T13:13:30.952425+01:00 +custom: + Author: Lindblomsebastian + Issue: "955" diff --git a/dbt/adapters/snowflake/impl.py b/dbt/adapters/snowflake/impl.py index 3583c888c..923583758 100644 --- a/dbt/adapters/snowflake/impl.py +++ b/dbt/adapters/snowflake/impl.py @@ -205,6 +205,10 @@ def submit_python_job(self, parsed_model: dict, compiled_code: str): packages = parsed_model["config"].get("packages", []) imports = parsed_model["config"].get("imports", []) + external_access_integrations = parsed_model["config"].get( + "external_access_integrations", [] + ) + secrets = parsed_model["config"].get("secrets", {}) # adding default packages we need to make python model work default_packages = ["snowflake-snowpark-python"] package_names = [package.split("==")[0] for package in packages] @@ -213,9 +217,18 @@ def submit_python_job(self, parsed_model: dict, compiled_code: str): packages.append(default_package) packages = "', '".join(packages) imports = "', '".join(imports) - # we can't pass empty imports clause to snowflake + external_access_integrations = ", ".join(external_access_integrations) + secrets = ", ".join(f"'{key}' = {value}" for key, value in secrets.items()) + + # we can't pass empty imports, external_access_integrations or secrets clause to snowflake if imports: imports = f"IMPORTS = ('{imports}')" + if external_access_integrations: + # Black is trying to make this a tuple. + # fmt: off + external_access_integrations = f"EXTERNAL_ACCESS_INTEGRATIONS = ({external_access_integrations})" + if secrets: + secrets = f"SECRETS = ({secrets})" if self.config.args.SEND_ANONYMOUS_USAGE_STATS: snowpark_telemetry_string = "dbtLabs_dbtPython" @@ -230,6 +243,8 @@ def submit_python_job(self, parsed_model: dict, compiled_code: str): LANGUAGE PYTHON RUNTIME_VERSION = '{python_version}' PACKAGES = ('{packages}') +{external_access_integrations} +{secrets} {imports} HANDLER = 'main' EXECUTE AS CALLER diff --git a/tests/functional/adapter/test_python_model.py b/tests/functional/adapter/test_python_model.py index ea67e6c1c..6bf0678c7 100644 --- a/tests/functional/adapter/test_python_model.py +++ b/tests/functional/adapter/test_python_model.py @@ -138,3 +138,78 @@ def teardown_method(self, project): def test_custom_target(self, project): results = run_dbt() assert results[0].node.schema == f"{project.test_schema}_MY_CUSTOM_SCHEMA" + + +EXTERNAL_ACCESS_INTEGRATION_MODE = """ +import pandas +import snowflake.snowpark as snowpark + +def model(dbt, session: snowpark.Session): + dbt.config( + materialized="table", + external_access_integrations=["test_external_access_integration"], + packages=["httpx==0.26.0"] + ) + import httpx + return session.create_dataframe( + pandas.DataFrame( + [{"result": httpx.get(url="https://www.google.com").status_code}] + ) + ) +""" + + +class TestExternalAccessIntegration: + @pytest.fixture(scope="class") + def models(self): + return {"external_access_integration_python_model.py": EXTERNAL_ACCESS_INTEGRATION_MODE} + + def test_external_access_integration(self, project): + project.run_sql( + "create or replace network rule test_network_rule type = host_port mode = egress value_list= ('www.google.com:443');" + ) + project.run_sql( + "create or replace external access integration test_external_access_integration allowed_network_rules = (test_network_rule) enabled = true;" + ) + run_dbt(["run"]) + + +SECRETS_MODE = """ +import pandas +import snowflake.snowpark as snowpark + +def model(dbt, session: snowpark.Session): + dbt.config( + materialized="table", + secrets={"secret_variable_name": "test_secret"}, + external_access_integrations=["test_external_access_integration"], + ) + import _snowflake + return session.create_dataframe( + pandas.DataFrame( + [{"secret_value": _snowflake.get_generic_secret_string('secret_variable_name')}] + ) + ) +""" + + +class TestSecrets: + @pytest.fixture(scope="class") + def models(self): + return {"secret_python_model.py": SECRETS_MODE} + + def test_secrets(self, project): + project.run_sql( + "create or replace secret test_secret type = generic_string secret_string='secret value';" + ) + + # The secrets you specify as values must also be specified in the external access integration. + # See https://docs.snowflake.com/en/developer-guide/external-network-access/creating-using-external-network-access#using-the-external-access-integration-in-a-function-or-procedure + + project.run_sql( + "create or replace network rule test_network_rule type = host_port mode = egress value_list= ('www.google.com:443');" + ) + project.run_sql( + "create or replace external access integration test_external_access_integration allowed_network_rules = (test_network_rule) allowed_authentication_secrets = (test_secret) enabled = true;" + ) + run_dbt(["run"])