dagster-snowflake 0.26.18rc1__tar.gz → 0.27.14__tar.gz

{dagster_snowflake-0.26.18rc1/dagster_snowflake.egg-info → dagster_snowflake-0.27.14}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-snowflake
-Version: 0.26.18rc1
+Version: 0.27.14
 Summary: Package for Snowflake Dagster framework components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-snowflake
 Author: Dagster Labs
@@ -10,11 +10,12 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.9,<3.13
+Requires-Python: >=3.9,<3.14
 License-File: LICENSE
-Requires-Dist: dagster==1.10.18rc1
+Requires-Dist: dagster==1.11.14
 Requires-Dist: snowflake-connector-python>=3.4.0
 Requires-Dist: pyOpenSSL>=22.1.0
 Provides-Extra: snowflake-sqlalchemy

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14}/dagster_snowflake/__init__.py

@@ -1,5 +1,8 @@
 from dagster_shared.libraries import DagsterLibraryRegistry
 
+from dagster_snowflake.components import (
+    SnowflakeConnectionComponent as SnowflakeConnectionComponent,
+)
 from dagster_snowflake.ops import snowflake_op_for_query as snowflake_op_for_query
 from dagster_snowflake.resources import (
     SnowflakeConnection as SnowflakeConnection,

dagster_snowflake-0.27.14/dagster_snowflake/components/__init__.py

@@ -0,0 +1,5 @@
+from dagster_snowflake.components.sql_component.component import SnowflakeConnectionComponent
+
+__all__ = [
+    "SnowflakeConnectionComponent",
+]

dagster_snowflake-0.27.14/dagster_snowflake/components/sql_component/component.py

@@ -0,0 +1,61 @@
+from functools import cached_property
+from typing import Any, cast
+
+import dagster as dg
+from dagster._annotations import preview, public
+from dagster._core.definitions.definitions_class import Definitions
+from dagster.components.core.context import ComponentLoadContext
+from dagster.components.lib.sql_component.sql_client import SQLClient
+from pydantic import BaseModel, create_model
+
+from dagster_snowflake.resources import SnowflakeResource
+
+
+@public
+@preview
+class SnowflakeConnectionComponentBase(dg.Component, dg.Resolvable, dg.Model, SQLClient):
+    """A component that represents a Snowflake connection. Use this component if you are
+    also using the TemplatedSqlComponent to execute SQL queries, and need to connect to Snowflake.
+    """
+
+    @cached_property
+    def _snowflake_resource(self) -> SnowflakeResource:
+        return SnowflakeResource(
+            **{
+                (field.alias or field_name): getattr(self, field_name)
+                for field_name, field in self.__class__.model_fields.items()
+            }
+        )
+
+    def connect_and_execute(self, sql: str) -> None:
+        """Connect to the SQL database and execute the SQL query."""
+        return self._snowflake_resource.connect_and_execute(sql)
+
+    def build_defs(self, context: ComponentLoadContext) -> Definitions:
+        return Definitions()
+
+
+def _copy_fields_to_model(
+    copy_from: type[BaseModel], copy_to: type[BaseModel], new_model_cls_name: str
+) -> None:
+    """Given two models, creates a copy of the second model with the fields of the first model."""
+    field_definitions: dict[str, tuple[type, Any]] = {
+        field_name: (cast("type", field.annotation), field)
+        for field_name, field in copy_from.model_fields.items()
+    }
+
+    return create_model(
+        new_model_cls_name,
+        __base__=copy_to,
+        __doc__=copy_to.__doc__,
+        **field_definitions,  # type: ignore
+    )
+
+
+SnowflakeConnectionComponent = public(preview)(
+    _copy_fields_to_model(
+        copy_from=SnowflakeResource,
+        copy_to=SnowflakeConnectionComponentBase,
+        new_model_cls_name="SnowflakeConnectionComponent",
+    )
+)

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14}/dagster_snowflake/resources.py

@@ -19,6 +19,7 @@ from dagster._annotations import public
 from dagster._core.definitions.resource_definition import dagster_maintained_resource
 from dagster._core.storage.event_log.sql_event_log import SqlDbConnection
 from dagster._utils.cached_method import cached_method
+from dagster.components.lib.sql_component.sql_client import SQLClient
 from pydantic import Field, model_validator, validator
 
 from dagster_snowflake.constants import (
@@ -40,7 +41,7 @@ except ImportError:
     raise
 
 
-class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext):
+class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext, SQLClient):
     """A resource for connecting to the Snowflake data warehouse.
 
     If connector configuration is not set, SnowflakeResource.get_connection() will return a
@@ -281,9 +282,9 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
 
     @validator("connector")
     def validate_connector(cls, v: Optional[str]) -> Optional[str]:
-        if v is not None and v != "sqlalchemy":
+        if v is not None and v not in ["sqlalchemy", "adbc"]:
             raise ValueError(
-                "Snowflake Resource: 'connector' configuration value must be None or sqlalchemy."
+                "Snowflake Resource: 'connector' configuration value must be None, sqlalchemy or adbc."
             )
         return v
 
@@ -391,6 +392,81 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
 
         return sqlalchemy_engine_args
 
+    @property
+    @cached_method
+    def _adbc_connection_args(self) -> Mapping[str, Any]:
+        config = self._resolved_config_dict
+        adbc_engine_args = {}
+
+        if config.get("account"):
+            adbc_engine_args["adbc.snowflake.sql.account"] = config["account"]
+        if config.get("user"):
+            adbc_engine_args["username"] = config["user"]
+        if config.get("password"):
+            adbc_engine_args["password"] = config["password"]
+        if config.get("database"):
+            adbc_engine_args["adbc.snowflake.sql.db"] = config["database"]
+        if config.get("schema"):
+            adbc_engine_args["adbc.snowflake.sql.schema"] = config["schema"]
+        if config.get("role"):
+            adbc_engine_args["adbc.snowflake.sql.role"] = config["role"]
+        if config.get("warehouse"):
+            adbc_engine_args["adbc.snowflake.sql.warehouse"] = config["warehouse"]
+
+        if config.get("authenticator"):
+            auth_mapping = {
+                "snowflake": "auth_snowflake",
+                "oauth": "auth_oauth",
+                "externalbrowser": "auth_ext_browser",
+                "okta": "auth_okta",
+                "jwt": "auth_jwt",
+                "snowflake_jwt": "auth_jwt",
+            }
+            auth_type = auth_mapping.get(config["authenticator"].lower(), config["authenticator"])
+            adbc_engine_args["adbc.snowflake.sql.auth_type"] = auth_type
+
+        if config.get("private_key") or config.get("private_key_path"):
+            # ADBC expects the raw private key value as bytes for jwt_private_key_pkcs8_value
+            adbc_engine_args["adbc.snowflake.sql.auth_type"] = "auth_jwt"
+            if config.get("private_key"):
+                adbc_engine_args["adbc.snowflake.sql.client_option.jwt_private_key_pkcs8_value"] = (
+                    config["private_key"]
+                )
+            elif config.get("private_key_path"):
+                adbc_engine_args["adbc.snowflake.sql.client_option.jwt_private_key"] = config[
+                    "private_key_path"
+                ]
+
+            if config.get("private_key_password"):
+                adbc_engine_args[
+                    "adbc.snowflake.sql.client_option.jwt_private_key_pkcs8_password"
+                ] = config["private_key_password"]
+
+        if config.get("login_timeout"):
+            adbc_engine_args["adbc.snowflake.sql.client_option.login_timeout"] = (
+                f"{config['login_timeout']}s"
+            )
+        if config.get("network_timeout"):
+            adbc_engine_args["adbc.snowflake.sql.client_option.request_timeout"] = (
+                f"{config['network_timeout']}s"
+            )
+        if config.get("client_session_keep_alive") is not None:
+            adbc_engine_args["adbc.snowflake.sql.client_option.keep_session_alive"] = str(
+                config["client_session_keep_alive"]
+            ).lower()
+
+        adbc_engine_args["adbc.snowflake.sql.client_option.app_name"] = (
+            SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER
+        )
+
+        if config.get("additional_snowflake_connection_args"):
+            for key, value in config["additional_snowflake_connection_args"].items():
+                # Allow direct ADBC option names to be passed through
+                if key.startswith("adbc.snowflake."):
+                    adbc_engine_args[key] = value  # noqa: PERF403
+
+        return adbc_engine_args
+
     def _snowflake_private_key(self, config) -> bytes:
         # If the user has defined a path to a private key, we will use that.
         if config.get("private_key_path", None) is not None:
@@ -476,6 +552,15 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
             yield conn
             conn.close()
             engine.dispose()
+        elif self.connector == "adbc":
+            import adbc_driver_snowflake.dbapi
+
+            conn = adbc_driver_snowflake.dbapi.connect(
+                db_kwargs=self._adbc_connection_args,  # pyright: ignore[reportArgumentType]
+            )
+
+            yield conn
+            conn.close()
         else:
             conn = snowflake.connector.connect(**self._connection_args)
 
@@ -493,6 +578,10 @@ class SnowflakeResource(ConfigurableResource, IAttachDifferentObjectToOpContext)
             snowflake_connection_resource=self,
         )
 
+    def connect_and_execute(self, sql: str) -> None:
+        with self.get_connection() as conn:
+            conn.cursor().execute(sql)
+
 
 class SnowflakeConnection:
     """A connection to Snowflake that can execute queries. In general this class should not be

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14}/dagster_snowflake/snowflake_io_manager.py

@@ -5,7 +5,7 @@ from typing import Any, Optional, cast
 
 from dagster import IOManagerDefinition, OutputContext, io_manager
 from dagster._config.pythonic_config import ConfigurableIOManagerFactory
-from dagster._core.definitions.time_window_partitions import TimeWindow
+from dagster._core.definitions.partitions.utils import TimeWindow
 from dagster._core.storage.db_io_manager import (
     DbClient,
     DbIOManager,
@@ -15,7 +15,6 @@ from dagster._core.storage.db_io_manager import (
 )
 from dagster._core.storage.io_manager import dagster_maintained_io_manager
 from pydantic import Field
-from snowflake.connector.errors import ProgrammingError
 
 from dagster_snowflake.resources import SnowflakeResource
 
@@ -30,7 +29,7 @@ def build_snowflake_io_manager(
     Args:
         type_handlers (Sequence[DbTypeHandler]): Each handler defines how to translate between
             slices of Snowflake tables and an in-memory type - e.g. a Pandas DataFrame. If only
-            one DbTypeHandler is provided, it will be used as teh default_load_type.
+            one DbTypeHandler is provided, it will be used as the default_load_type.
         default_load_type (Type): When an input has no type annotation, load it as this type.
 
     Returns:
@@ -59,7 +58,7 @@ def build_snowflake_io_manager(
 
             snowflake_io_manager = build_snowflake_io_manager([SnowflakePandasTypeHandler(), SnowflakePySparkTypeHandler()])
 
-            defs = Definitions(
+            Definitions(
                 assets=[my_table, my_second_table],
                 resources={
                     "io_manager": snowflake_io_manager.configured({
@@ -75,7 +74,7 @@ def build_snowflake_io_manager(
 
         .. code-block:: python
 
-            defs = Definitions(
+            Definitions(
                 assets=[my_table]
                 resources={"io_manager" snowflake_io_manager.configured(
                     {"database": "my_database", "schema": "my_schema", ...} # will be used as the schema
@@ -364,20 +363,22 @@ class SnowflakeDbClient(DbClient):
 
     @staticmethod
     def ensure_schema_exists(context: OutputContext, table_slice: TableSlice, connection) -> None:
-        schemas = (
-            connection.cursor()
-            .execute(f"show schemas like '{table_slice.schema}' in database {table_slice.database}")
-            .fetchall()
-        )
+        with connection.cursor() as cursor:
+            cursor.execute(
+                f"show schemas like '{table_slice.schema}' in database {table_slice.database}"
+            )
+            schemas = cursor.fetchall()
+
         if len(schemas) == 0:
-            connection.cursor().execute(f"create schema {table_slice.schema};")
+            with connection.cursor() as cursor:
+                cursor.execute(f"create schema {table_slice.schema};")
 
     @staticmethod
     def delete_table_slice(context: OutputContext, table_slice: TableSlice, connection) -> None:
         try:
             connection.cursor().execute(_get_cleanup_statement(table_slice))
-        except ProgrammingError as e:
-            if "does not exist" in e.msg:  # type: ignore
+        except Exception as e:
+            if "does not exist or not authorized" in str(e):
                 # table doesn't exist yet, so ignore the error
                 return
             else:

dagster_snowflake-0.27.14/dagster_snowflake/version.py

@@ -0,0 +1 @@
+__version__ = "0.27.14"

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14/dagster_snowflake.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: dagster-snowflake
-Version: 0.26.18rc1
+Version: 0.27.14
 Summary: Package for Snowflake Dagster framework components.
 Home-page: https://github.com/dagster-io/dagster/tree/master/python_modules/libraries/dagster-snowflake
 Author: Dagster Labs
@@ -10,11 +10,12 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
-Requires-Python: >=3.9,<3.13
+Requires-Python: >=3.9,<3.14
 License-File: LICENSE
-Requires-Dist: dagster==1.10.18rc1
+Requires-Dist: dagster==1.11.14
 Requires-Dist: snowflake-connector-python>=3.4.0
 Requires-Dist: pyOpenSSL>=22.1.0
 Provides-Extra: snowflake-sqlalchemy

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14}/dagster_snowflake.egg-info/SOURCES.txt

@@ -13,6 +13,10 @@ dagster_snowflake/version.py
 dagster_snowflake.egg-info/PKG-INFO
 dagster_snowflake.egg-info/SOURCES.txt
 dagster_snowflake.egg-info/dependency_links.txt
+dagster_snowflake.egg-info/entry_points.txt
 dagster_snowflake.egg-info/not-zip-safe
 dagster_snowflake.egg-info/requires.txt
-dagster_snowflake.egg-info/top_level.txt
+dagster_snowflake.egg-info/top_level.txt
+dagster_snowflake/components/__init__.py
+dagster_snowflake/components/sql_component/__init__.py
+dagster_snowflake/components/sql_component/component.py

dagster_snowflake-0.27.14/dagster_snowflake.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[dagster_dg_cli.registry_modules]
+dagster_snowflake = dagster_snowflake

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14}/dagster_snowflake.egg-info/requires.txt

@@ -1,4 +1,4 @@
-dagster==1.10.18rc1
+dagster==1.11.14
 snowflake-connector-python>=3.4.0
 pyOpenSSL>=22.1.0
 

{dagster_snowflake-0.26.18rc1 → dagster_snowflake-0.27.14}/setup.py

@@ -27,14 +27,15 @@ setup(
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
         "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
         "License :: OSI Approved :: Apache Software License",
         "Operating System :: OS Independent",
     ],
     packages=find_packages(exclude=["dagster_snowflake_tests*"]),
     include_package_data=True,
-    python_requires=">=3.9,<3.13",
+    python_requires=">=3.9,<3.14",
     install_requires=[
-        "dagster==1.10.18rc1",
+        "dagster==1.11.14",
         "snowflake-connector-python>=3.4.0",
         # Workaround for incorrect pin in the snowflake-connector-python package
         # See https://github.com/snowflakedb/snowflake-connector-python/issues/2109
@@ -51,4 +52,9 @@ setup(
         ],
     },
     zip_safe=False,
+    entry_points={
+        "dagster_dg_cli.registry_modules": [
+            "dagster_snowflake = dagster_snowflake",
+        ],
+    },
 )

dagster_snowflake-0.26.18rc1/dagster_snowflake/version.py

@@ -1 +0,0 @@
-__version__ = "0.26.18rc1"