diff --git a/dlt/common/destination/capabilities.py b/dlt/common/destination/capabilities.py index 54d9ac1a6e..3a8972c3b8 100644 --- a/dlt/common/destination/capabilities.py +++ b/dlt/common/destination/capabilities.py @@ -180,6 +180,7 @@ class DestinationCapabilitiesContext(ContainerInjectableContext): # use naming convention in the schema naming_convention: TNamingConventionReferenceArg = None alter_add_multi_column: bool = True + supports_create_schema_if_not_exists: bool = False supports_create_table_if_not_exists: bool = True supports_truncate_command: bool = True schema_supports_numeric_precision: bool = True diff --git a/dlt/destinations/impl/snowflake/factory.py b/dlt/destinations/impl/snowflake/factory.py index f6901978f3..613700c803 100644 --- a/dlt/destinations/impl/snowflake/factory.py +++ b/dlt/destinations/impl/snowflake/factory.py @@ -145,6 +145,7 @@ def _raw_capabilities(self) -> DestinationCapabilitiesContext: caps.is_max_text_data_type_length_in_bytes = True caps.supports_ddl_transactions = True caps.alter_add_multi_column = True + caps.supports_create_schema_if_not_exists = True caps.supports_clone_table = True caps.supported_merge_strategies = ["delete-insert", "upsert", "scd2", "insert-only"] caps.supported_replace_strategies = [ diff --git a/dlt/destinations/sql_client.py b/dlt/destinations/sql_client.py index e51fb89d64..e506941eae 100644 --- a/dlt/destinations/sql_client.py +++ b/dlt/destinations/sql_client.py @@ -133,7 +133,12 @@ def has_dataset(self) -> bool: return len(rows) > 0 def create_dataset(self) -> None: - self.execute_sql("CREATE SCHEMA %s" % self.fully_qualified_dataset_name()) + not_exists_clause = ( + " IF NOT EXISTS" if self.capabilities.supports_create_schema_if_not_exists else "" + ) + self.execute_sql( + "CREATE SCHEMA%s %s" % (not_exists_clause, self.fully_qualified_dataset_name()) + ) def drop_dataset(self) -> None: # assert self.fully_qualified_dataset_name() != "None" diff --git a/tests/load/snowflake/test_snowflake_table_builder.py b/tests/load/snowflake/test_snowflake_table_builder.py index 84e3db924b..fce797be24 100644 --- a/tests/load/snowflake/test_snowflake_table_builder.py +++ b/tests/load/snowflake/test_snowflake_table_builder.py @@ -21,6 +21,7 @@ SnowflakeClientConfiguration, SnowflakeCredentials, ) +from dlt.destinations.impl.snowflake.sql_client import SnowflakeSqlClient from dlt.common.destination.typing import PreparedTableSchema from dlt.common.schema.typing import TColumnSchema @@ -57,6 +58,31 @@ def create_client(schema: Schema, use_decfloat: bool = False) -> SnowflakeClient ) +def test_create_dataset_uses_idempotent_schema_sql( + monkeypatch: pytest.MonkeyPatch, +) -> None: + credentials = SnowflakeCredentials() + credentials.database = "test_db" + sql_client = SnowflakeSqlClient( + "test_dataset", + "test_dataset_staging", + credentials, + snowflake().capabilities(), + ) + executed_sql: list[str] = [] + + def capture_sql(sql: str, *args: object, **kwargs: object) -> None: + executed_sql.append(sql) + + monkeypatch.setattr(sql_client, "execute_sql", capture_sql) + + sql_client.create_dataset() + + assert executed_sql == [ + "CREATE SCHEMA IF NOT EXISTS %s" % sql_client.fully_qualified_dataset_name() + ] + + def test_create_table(snowflake_client: SnowflakeClient) -> None: # make sure we are in case insensitive mode assert snowflake_client.capabilities.generates_case_sensitive_identifiers() is False diff --git a/tests/load/test_sql_client.py b/tests/load/test_sql_client.py index a47048a829..a7687123c8 100644 --- a/tests/load/test_sql_client.py +++ b/tests/load/test_sql_client.py @@ -1,11 +1,13 @@ import os import pytest import datetime # noqa: I251 -from typing import Iterator, Any, Tuple, Type, Union +from contextlib import contextmanager +from typing import ContextManager, Iterator, Any, Optional, Sequence, Tuple, Type, Union from threading import Thread, Event from time import sleep from dlt.common import pendulum, Decimal +from dlt.common.destination import DestinationCapabilitiesContext from dlt.common.destination.exceptions import IdentifierTooLongException from dlt.common.schema.typing import LOADS_TABLE_NAME from dlt.common.storages import FileStorage @@ -40,11 +42,61 @@ ) +class RecordingSqlClient(SqlClientBase[Any]): + def __init__(self, capabilities: DestinationCapabilitiesContext) -> None: + super().__init__(None, "dataset", "dataset_staging", capabilities) + self.statements: list[str] = [] + + def open_connection(self) -> None: + return None + + def close_connection(self) -> None: + return None + + @contextmanager + def begin_transaction(self) -> Iterator[None]: + yield None + + @property + def native_connection(self) -> None: + return None + + def execute_sql(self, sql: str, *args: Any, **kwargs: Any) -> Optional[Sequence[Sequence[Any]]]: + self.statements.append(sql) + return None + + def execute_query(self, query: str, *args: Any, **kwargs: Any) -> ContextManager[DBApiCursor]: + raise NotImplementedError + + @staticmethod + def _make_database_exception(ex: Exception) -> Exception: + return ex + + @pytest.fixture def file_storage() -> FileStorage: return FileStorage(get_test_storage_root(), file_type="b", makedirs=True) +@pytest.mark.parametrize( + ("supports_create_schema_if_not_exists", "expected_sql"), + [ + (False, "CREATE SCHEMA dataset"), + (True, "CREATE SCHEMA IF NOT EXISTS dataset"), + ], +) +def test_create_dataset_uses_schema_idempotency_capability( + supports_create_schema_if_not_exists: bool, expected_sql: str +) -> None: + capabilities = DestinationCapabilitiesContext.generic_capabilities() + capabilities.supports_create_schema_if_not_exists = supports_create_schema_if_not_exists + sql_client = RecordingSqlClient(capabilities) + + sql_client.create_dataset() + + assert sql_client.statements == [expected_sql] + + @pytest.mark.parametrize( "client", destinations_configs(