dagster-snowflake 0.13.3rc0__py3-none-any.whl → 0.28.1__py3-none-any.whl

dagster_snowflake/__init__.py

@@ -1,9 +1,19 @@
-from dagster.core.utils import check_dagster_package_version
+from dagster_shared.libraries import DagsterLibraryRegistry
 
-from .resources import SnowflakeConnection, snowflake_resource
-from .solids import snowflake_solid_for_query
-from .version import __version__
+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,
+    SnowflakeResource as SnowflakeResource,
+    fetch_last_updated_timestamps as fetch_last_updated_timestamps,
+    snowflake_resource as snowflake_resource,
+)
+from dagster_snowflake.snowflake_io_manager import (
+    SnowflakeIOManager as SnowflakeIOManager,
+    build_snowflake_io_manager as build_snowflake_io_manager,
+)
+from dagster_snowflake.version import __version__
 
-check_dagster_package_version("dagster-snowflake", __version__)
-
-__all__ = ["snowflake_solid_for_query", "snowflake_resource", "SnowflakeConnection"]
+DagsterLibraryRegistry.register("dagster-snowflake", __version__)

dagster_snowflake/components/__init__.py

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

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/constants.py

@@ -0,0 +1,7 @@
+# Description: This file contains the Snowflake connection identifiers for the Snowflake partner account.
+# The connection identifiers are used to identify the partner account when connecting to Snowflake.
+
+# We use different connection identifiers for different connection code paths to ensure that each is
+# working as expected.
+SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER = "DagsterLabs_Dagster"
+SNOWFLAKE_PARTNER_CONNECTION_IDENTIFIER_SQLALCHEMY = "DagsterLabs_Dagster_SqlAlchemy"

dagster_snowflake/{solids.py → ops.py}

@@ -1,4 +1,10 @@
-from dagster import InputDefinition, Nothing, check, op, solid
+from dagster import (
+    Nothing,
+    _check as check,
+    op,
+)
+from dagster._core.definitions.input import In
+from dagster._core.storage.tags import COMPUTE_KIND_TAG
 
 
 def _core_create_snowflake_command(dagster_decorator, decorator_name, sql, parameters=None):
@@ -7,9 +13,9 @@ def _core_create_snowflake_command(dagster_decorator, decorator_name, sql, param
 
     @dagster_decorator(
         name=f"snowflake_{decorator_name}",
-        input_defs=[InputDefinition("start", Nothing)],
+        ins={"start": In(Nothing)},
         required_resource_keys={"snowflake"},
-        tags={"kind": "sql", "sql": sql},
+        tags={COMPUTE_KIND_TAG: "sql", "sql": sql},
     )
     def snowflake_fn(context):
         context.resources.snowflake.execute_query(sql=sql, parameters=parameters)
@@ -21,8 +27,8 @@ def snowflake_solid_for_query(sql, parameters=None):
     """This function is a solid factory that constructs solids to execute a snowflake query.
 
     Note that you can only use `snowflake_solid_for_query` if you know the query you'd like to
-    execute at pipeline construction time. If you'd like to execute queries dynamically during
-    pipeline execution, you should manually execute those queries in your custom solid using the
+    execute at job construction time. If you'd like to execute queries dynamically during
+    job execution, you should manually execute those queries in your custom solid using the
     snowflake resource.
 
     Args:
@@ -32,15 +38,15 @@ def snowflake_solid_for_query(sql, parameters=None):
     Returns:
         SolidDefinition: Returns the constructed solid definition.
     """
-    return _core_create_snowflake_command(solid, "solid", sql, parameters)
+    return _core_create_snowflake_command(op, "solid", sql, parameters)
 
 
 def snowflake_op_for_query(sql, parameters=None):
     """This function is an op factory that constructs an op to execute a snowflake query.
 
     Note that you can only use `snowflake_op_for_query` if you know the query you'd like to
-    execute at pipeline construction time. If you'd like to execute queries dynamically during
-    pipeline execution, you should manually execute those queries in your custom op using the
+    execute at graph construction time. If you'd like to execute queries dynamically during
+    job execution, you should manually execute those queries in your custom op using the
     snowflake resource.
 
     Args:

dagster_snowflake/py.typed

@@ -0,0 +1 @@
+partial