diff --git a/src/astro/databases/snowflake.py b/src/astro/databases/snowflake.py index da288a8a5..7eb15937b 100644 --- a/src/astro/databases/snowflake.py +++ b/src/astro/databases/snowflake.py @@ -12,6 +12,7 @@ from snowflake.connector import pandas_tools from snowflake.connector.errors import ProgrammingError +from astro import settings from astro.constants import ( DEFAULT_CHUNK_SIZE, FileLocation, @@ -37,6 +38,11 @@ FileType.PARQUET: "MATCH_BY_COLUMN_NAME=CASE_INSENSITIVE", } +DEFAULT_STORAGE_INTEGRATION = { + FileLocation.S3: settings.SNOWFLAKE_STORAGE_INTEGRATION_AMAZON, + FileLocation.GS: settings.SNOWFLAKE_STORAGE_INTEGRATION_GOOGLE, +} + NATIVE_LOAD_SUPPORTED_FILE_TYPES = (FileType.CSV, FileType.NDJSON, FileType.PARQUET) NATIVE_LOAD_SUPPORTED_FILE_LOCATIONS = (FileLocation.GS, FileLocation.S3) @@ -142,7 +148,6 @@ class SnowflakeDatabase(BaseDatabase): """ def __init__(self, conn_id: str = DEFAULT_CONN_ID): - self.storage_integration: Optional[str] = None super().__init__(conn_id) @property @@ -196,7 +201,9 @@ def _create_stage_auth_sub_statement( :param storage_integration: Previously created Snowflake storage integration :return: String containing line to be used for authentication on the remote storage """ - + storage_integration = storage_integration or DEFAULT_STORAGE_INTEGRATION.get( + file.location.location_type + ) if storage_integration is not None: auth = f"storage_integration = {storage_integration};" else: diff --git a/tests/benchmark/Dockerfile b/tests/benchmark/Dockerfile index ece9eb23d..1d2d241bd 100644 --- a/tests/benchmark/Dockerfile +++ b/tests/benchmark/Dockerfile @@ -11,6 +11,7 @@ ENV AIRFLOW_HOME=/opt/app/ ENV PYTHONPATH=/opt/app/ ENV ASTRO_PUBLISH_BENCHMARK_DATA=True ENV GCP_BUCKET=dag-authoring +ENV AIRFLOW__ASTRO_SDK__SNOWFLAKE_STORAGE_INTEGRATION_GOOGLE=gcs_int_python_sdk # Debian Bullseye is shipped with Python 3.9 # Upgrade built-in pip