Fix snowflake optimisation retrieval of envvar

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:

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