airbyte-lib snowflake integration

.github/workflows/airbyte-ci-tests.yml

@@ -79,7 +79,7 @@ jobs:
               - 'airbyte-ci/connectors/metadata_service/orchestrator/**'
               - '!**/*.md'
             airbyte_lib:
-              - 'airbyte_lib/**'
+              - 'airbyte-lib/**'
               - '!**/*.md'
 
       - name: Run airbyte-ci/connectors/connector_ops tests
@@ -162,7 +162,8 @@ jobs:
           docker_hub_password: ${{ secrets.DOCKER_HUB_PASSWORD }}
           docker_hub_username: ${{ secrets.DOCKER_HUB_USERNAME }}
           gcs_credentials: ${{ secrets.METADATA_SERVICE_PROD_GCS_CREDENTIALS }}
+          gcp_gsm_credentials: ${{ secrets.GCP_GSM_CREDENTIALS }}
           sentry_dsn: ${{ secrets.SENTRY_AIRBYTE_CI_DSN }}
           github_token: ${{ secrets.GH_PAT_MAINTENANCE_OCTAVIA }}
-          subcommand: "test airbyte-lib"
+          subcommand: "test airbyte-lib --pass-env-var=GCP_GSM_CREDENTIALS --poetry-run-command='pytest'"
           tailscale_auth_key: ${{ secrets.TAILSCALE_AUTH_KEY }}

airbyte-lib/airbyte_lib/caches/__init__.py

@@ -3,6 +3,7 @@
 from airbyte_lib.caches.base import SQLCacheBase
 from airbyte_lib.caches.duckdb import DuckDBCache, DuckDBCacheConfig
 from airbyte_lib.caches.postgres import PostgresCache, PostgresCacheConfig
+from airbyte_lib.caches.snowflake import SnowflakeCacheConfig, SnowflakeSQLCache
 
 
 # We export these classes for easy access: `airbyte_lib.caches...`
@@ -12,4 +13,6 @@
     "PostgresCache",
     "PostgresCacheConfig",
     "SQLCacheBase",
+    "SnowflakeCacheConfig",
+    "SnowflakeSQLCache",
 ]

airbyte-lib/airbyte_lib/caches/base.py

@@ -15,8 +15,9 @@
 import sqlalchemy
 import ulid
 from overrides import overrides
-from sqlalchemy import CursorResult, Executable, TextClause, create_engine, text
+from sqlalchemy import create_engine, text
 from sqlalchemy.pool import StaticPool
+from sqlalchemy.sql.elements import TextClause
 
 from airbyte_lib._file_writers.base import FileWriterBase, FileWriterBatchHandle
 from airbyte_lib._processors import BatchHandle, RecordProcessor
@@ -28,7 +29,9 @@
     from pathlib import Path
 
     from sqlalchemy.engine import Connection, Engine
+    from sqlalchemy.engine.cursor import CursorResult
     from sqlalchemy.engine.reflection import Inspector
+    from sqlalchemy.sql.base import Executable
 
     from airbyte_protocol.models import ConfiguredAirbyteStream
 
@@ -593,7 +596,7 @@ def _write_files_to_new_table(
                     schema=self.config.schema_name,
                     if_exists="append",
                     index=False,
-                    dtype=self._get_sql_column_definitions(stream_name),  # type: ignore[arg-type]
+                    dtype=self._get_sql_column_definitions(stream_name),
                 )
         return temp_table_name
 

airbyte-lib/airbyte_lib/caches/postgres.py

@@ -28,7 +28,9 @@ class PostgresCacheConfig(SQLCacheConfigBase, ParquetWriterConfig):
     @overrides
     def get_sql_alchemy_url(self) -> str:
         """Return the SQLAlchemy URL to use."""
-        return f"postgresql+psycopg://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
+        return (
+            f"postgresql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}"
+        )
 
     def get_database_name(self) -> str:
         """Return the name of the database."""

airbyte-lib/airbyte_lib/caches/snowflake.py

@@ -10,9 +10,10 @@
 from typing import TYPE_CHECKING
 
 from overrides import overrides
+from snowflake.sqlalchemy import URL
 
 from airbyte_lib._file_writers import ParquetWriter, ParquetWriterConfig
-from airbyte_lib.caches.base import SQLCacheBase, SQLCacheConfigBase
+from airbyte_lib.caches.base import RecordDedupeMode, SQLCacheBase, SQLCacheConfigBase
 from airbyte_lib.telemetry import CacheTelemetryInfo
 
 
@@ -31,16 +32,25 @@ class SnowflakeCacheConfig(SQLCacheConfigBase, ParquetWriterConfig):
     password: str
     warehouse: str
     database: str
+    role: str
+
+    dedupe_mode = RecordDedupeMode.APPEND
 
     # Already defined in base class:
     # schema_name: str
 
     @overrides
     def get_sql_alchemy_url(self) -> str:
         """Return the SQLAlchemy URL to use."""
-        return (
-            f"snowflake://{self.username}:{self.password}@{self.account}/"
-            f"?warehouse={self.warehouse}&database={self.database}&schema={self.schema_name}"
+        return str(
+            URL(
+                account=self.account,
+                user=self.username,
+                password=self.password,
+                database=self.database,
+                warehouse=self.warehouse,
+                role=self.role,
+            )
         )
 
     def get_database_name(self) -> str:
@@ -67,6 +77,7 @@ def _write_files_to_new_table(
         """Write a file(s) to a new table.
 
         TODO: Override the base implementation to use the COPY command.
+        TODO: Make sure this works for all data types.
         """
         return super()._write_files_to_new_table(files, stream_name, batch_id)
 

airbyte-lib/airbyte_lib/results.py

@@ -1,6 +1,6 @@
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 
-from sqlalchemy import Engine
+from sqlalchemy.engine import Engine
 
 from airbyte_lib.caches import SQLCacheBase
 from airbyte_lib.datasets import CachedDataset