From fbc5c5028431640f281e7509e157835aade9cd73 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?= Date: Wed, 18 Dec 2024 16:07:19 -0600 Subject: [PATCH 1/3] fix: Handle type-mapping with `SQLConnector.jsonschema_to_sql` --- pyproject.toml | 2 -- target_snowflake/connector.py | 28 +++++++++++++++------------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index db57ad8..b9e788d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,8 +40,6 @@ target-version = "py39" [tool.ruff.lint] ignore = [ - "ANN101", # missing-type-self - "ANN102", # missing-type-cls "ANN201", "TD", "D", diff --git a/target_snowflake/connector.py b/target_snowflake/connector.py index 3c76f6c..712fd4f 100644 --- a/target_snowflake/connector.py +++ b/target_snowflake/connector.py @@ -3,7 +3,7 @@ import urllib.parse from enum import Enum from functools import cached_property -from operator import contains, eq +from operator import eq from pathlib import Path from typing import TYPE_CHECKING, Any, cast @@ -13,7 +13,7 @@ from cryptography.hazmat.primitives import serialization from singer_sdk import typing as th from singer_sdk.connectors import SQLConnector -from singer_sdk.connectors.sql import FullyQualifiedName +from singer_sdk.connectors.sql import FullyQualifiedName, JSONSchemaToSQL from singer_sdk.exceptions import ConfigValidationError from snowflake.sqlalchemy import URL from snowflake.sqlalchemy.base import SnowflakeIdentifierPreparer @@ -27,6 +27,7 @@ from sqlalchemy.engine import Engine +# TODO: Remove this when when JSON schema to SQL is stable SNOWFLAKE_MAX_STRING_LENGTH = 16777216 @@ -89,6 +90,8 @@ class SnowflakeConnector(SQLConnector): allow_merge_upsert: bool = False # Whether MERGE UPSERT is supported. allow_temp_tables: bool = True # Whether temp tables are supported. + max_varchar_length = 16_777_216 + def __init__(self, *args: Any, **kwargs: Any) -> None: self.table_cache: dict = {} self.schema_cache: dict = {} @@ -317,6 +320,16 @@ def _conform_max_length(jsonschema_type): # noqa: ANN205, ANN001 jsonschema_type["maxLength"] = SNOWFLAKE_MAX_STRING_LENGTH return jsonschema_type + @cached_property + def jsonschema_to_sql(self) -> JSONSchemaToSQL: + to_sql = super().jsonschema_to_sql + to_sql.register_type_handler("integer", NUMBER) + to_sql.register_type_handler("object", VARIANT) + to_sql.register_type_handler("array", VARIANT) + to_sql.register_type_handler("number", sct.DOUBLE) + to_sql.register_format_handler("date-time", TIMESTAMP_NTZ) + return to_sql + def to_sql_type(self, jsonschema_type: dict) -> sqlalchemy.types.TypeEngine: """Return a JSON Schema representation of the provided type. @@ -336,23 +349,12 @@ def to_sql_type(self, jsonschema_type: dict) -> sqlalchemy.types.TypeEngine: maxlength = jsonschema_type.get("maxLength", SNOWFLAKE_MAX_STRING_LENGTH) # define type maps string_submaps = [ - TypeMap(eq, TIMESTAMP_NTZ(), "date-time"), - TypeMap(contains, sqlalchemy.types.TIME(), "time"), - TypeMap(eq, sqlalchemy.types.DATE(), "date"), TypeMap(eq, sqlalchemy.types.VARCHAR(maxlength), None), ] - type_maps = [ - TypeMap(th._jsonschema_type_check, NUMBER(), ("integer",)), # noqa: SLF001 - TypeMap(th._jsonschema_type_check, VARIANT(), ("object",)), # noqa: SLF001 - TypeMap(th._jsonschema_type_check, VARIANT(), ("array",)), # noqa: SLF001 - TypeMap(th._jsonschema_type_check, sct.DOUBLE(), ("number",)), # noqa: SLF001 - ] # apply type maps if th._jsonschema_type_check(jsonschema_type, ("string",)): # noqa: SLF001 datelike_type = th.get_datelike_property_type(jsonschema_type) target_type = evaluate_typemaps(string_submaps, datelike_type, target_type) - else: - target_type = evaluate_typemaps(type_maps, jsonschema_type, target_type) return cast(sqlalchemy.types.TypeEngine, target_type) From 18c449a6b6c5bc5546cbaf57ae2a3a824aee85d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?= Date: Wed, 18 Dec 2024 18:48:34 -0600 Subject: [PATCH 2/3] Rely entirely on Singer SDK JSON Schema to SQL type conversion --- target_snowflake/connector.py | 66 ++--------------------------------- 1 file changed, 2 insertions(+), 64 deletions(-) diff --git a/target_snowflake/connector.py b/target_snowflake/connector.py index 712fd4f..9128565 100644 --- a/target_snowflake/connector.py +++ b/target_snowflake/connector.py @@ -3,15 +3,13 @@ import urllib.parse from enum import Enum from functools import cached_property -from operator import eq from pathlib import Path -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any import snowflake.sqlalchemy.custom_types as sct import sqlalchemy from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization -from singer_sdk import typing as th from singer_sdk.connectors import SQLConnector from singer_sdk.connectors.sql import FullyQualifiedName, JSONSchemaToSQL from singer_sdk.exceptions import ConfigValidationError @@ -27,31 +25,6 @@ from sqlalchemy.engine import Engine -# TODO: Remove this when when JSON schema to SQL is stable -SNOWFLAKE_MAX_STRING_LENGTH = 16777216 - - -class TypeMap: - def __init__(self, operator, map_value, match_value=None) -> None: # noqa: ANN001 - self.operator = operator - self.map_value = map_value - self.match_value = match_value - - def match(self, compare_value): # noqa: ANN001 - try: - if self.match_value: - return self.operator(compare_value, self.match_value) - return self.operator(compare_value) - except TypeError: - return False - - -def evaluate_typemaps(type_maps, compare_value, unmatched_value): # noqa: ANN001 - for type_map in type_maps: - if type_map.match(compare_value): - return type_map.map_value - return unmatched_value - class SnowflakeFullyQualifiedName(FullyQualifiedName): def __init__( @@ -312,16 +285,9 @@ def get_column_alter_ddl( }, ) - @staticmethod - def _conform_max_length(jsonschema_type): # noqa: ANN205, ANN001 - """Alter jsonschema representations to limit max length to Snowflake's VARCHAR length.""" - max_length = jsonschema_type.get("maxLength") - if max_length and max_length > SNOWFLAKE_MAX_STRING_LENGTH: - jsonschema_type["maxLength"] = SNOWFLAKE_MAX_STRING_LENGTH - return jsonschema_type - @cached_property def jsonschema_to_sql(self) -> JSONSchemaToSQL: + # https://docs.snowflake.com/en/sql-reference/intro-summary-data-types.html to_sql = super().jsonschema_to_sql to_sql.register_type_handler("integer", NUMBER) to_sql.register_type_handler("object", VARIANT) @@ -330,34 +296,6 @@ def jsonschema_to_sql(self) -> JSONSchemaToSQL: to_sql.register_format_handler("date-time", TIMESTAMP_NTZ) return to_sql - def to_sql_type(self, jsonschema_type: dict) -> sqlalchemy.types.TypeEngine: - """Return a JSON Schema representation of the provided type. - - Uses custom Snowflake types from [snowflake-sqlalchemy](https://github.com/snowflakedb/snowflake-sqlalchemy/blob/main/src/snowflake/sqlalchemy/custom_types.py) - - Args: - jsonschema_type: The JSON Schema representation of the source type. - - Returns: - The SQLAlchemy type representation of the data type. - """ - # start with default implementation - jsonschema_type = SnowflakeConnector._conform_max_length(jsonschema_type) - target_type = super().to_sql_type(jsonschema_type) - # snowflake max and default varchar length - # https://docs.snowflake.com/en/sql-reference/intro-summary-data-types.html - maxlength = jsonschema_type.get("maxLength", SNOWFLAKE_MAX_STRING_LENGTH) - # define type maps - string_submaps = [ - TypeMap(eq, sqlalchemy.types.VARCHAR(maxlength), None), - ] - # apply type maps - if th._jsonschema_type_check(jsonschema_type, ("string",)): # noqa: SLF001 - datelike_type = th.get_datelike_property_type(jsonschema_type) - target_type = evaluate_typemaps(string_submaps, datelike_type, target_type) - - return cast(sqlalchemy.types.TypeEngine, target_type) - def schema_exists(self, schema_name: str) -> bool: if schema_name in self.schema_cache: return True From 1882efba071d2759a419cc750a58a8f7276b32e5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Edgar=20Ram=C3=ADrez=20Mondrag=C3=B3n?= Date: Wed, 18 Dec 2024 19:41:34 -0600 Subject: [PATCH 3/3] test: Re-enable schema update tests --- tests/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/core.py b/tests/core.py index 79dc433..e190c8f 100644 --- a/tests/core.py +++ b/tests/core.py @@ -586,7 +586,7 @@ def singer_filepath(self) -> Path: SnowflakeTargetRecordMissingKeyProperty, SnowflakeTargetRecordMissingRequiredProperty, SnowflakeTargetSchemaNoProperties, - # SnowflakeTargetSchemaUpdates, + SnowflakeTargetSchemaUpdates, TargetSpecialCharsInAttributes, # Implicitly asserts special chars handled SnowflakeTargetReservedWords, SnowflakeTargetReservedWordsNoKeyProps,