feat(ingest/dagster): Add automatic snowflake_pandas_io_manager asset capture

metadata-ingestion-modules/dagster-plugin/examples/iris.py

@@ -0,0 +1,112 @@
+import pandas as pd
+from dagster import MaterializeResult, MetadataValue, asset
+from dagster_aws.redshift import RedshiftClientResource
+from dagster_snowflake import SnowflakeResource
+
+
+@asset(
+    metadata={"schema": "public"},
+    key_prefix=["prod", "snowflake", "test_db", "public"],
+    group_name="iris",
+    io_manager_key="snowflake_io_manager",
+)
+def iris_dataset() -> pd.DataFrame:
+    return pd.read_csv(
+        "https://docs.dagster.io/assets/iris.csv",
+        names=[
+            "sepal_length_cm",
+            "sepal_width_cm",
+            "petal_length_cm",
+            "petal_width_cm",
+            "species",
+        ],
+    )
+
+
+@asset(
+    metadata={"schema": "public"},
+    key_prefix=["prod", "snowflake", "test_db", "public"],
+    group_name="iris",
+    io_manager_key="snowflake_io_manager",
+)
+def iris_cleaned(iris_dataset: pd.DataFrame) -> pd.DataFrame:
+    return iris_dataset.dropna().drop_duplicates()
+
+
+@asset(
+    metadata={"schema": "public"},
+    key_prefix=["prod", "snowflake", "test_db", "public"],
+    group_name="iris",
+    deps=[iris_dataset],
+)
+def iris_setosa(snowflake: SnowflakeResource) -> MaterializeResult:
+    query = """
+        create or replace table TEST_DB.public.iris_setosa as (
+            SELECT *
+            FROM TEST_DB.public.iris_cleaned
+            WHERE species = 'Iris-setosa'
+        );
+    """
+
+    with snowflake.get_connection() as connection:
+        with connection.cursor() as cursor:
+            cursor.execute(query)
+
+    return MaterializeResult(
+        metadata={
+            "Query": MetadataValue.text(query),
+        }
+    )
+
+
+@asset(
+    key_prefix=[
+        "prod",
+        "snowflake",
+        "db_name",
+        "schema_name",
+    ],  # the fqdn asset name to be able identify platform and make sure asset is unique
+    group_name="iris",
+    deps=[iris_dataset],
+)
+def my_asset_table_a(snowflake: SnowflakeResource) -> MaterializeResult:
+    query = """
+        create or replace table db_name.schema_name.my_asset_table_a as (
+            SELECT *
+            FROM db_name.schema_name.my_asset_table_b
+        );
+    """
+
+    with snowflake.get_connection() as connection:
+        with connection.cursor() as cursor:
+            cursor.execute(query)
+
+    return MaterializeResult(  # Adding query to metadata to use it getting lineage from it with sql parser
+        metadata={
+            "Query": MetadataValue.text(query),
+        }
+    )
+
+
+@asset(
+    key_prefix=[
+        "prod",
+        "redshift",
+        "dev",
+        "public",
+        "blood_storage_count",
+    ],  # the fqdn asset name to be able identify platform and make sure asset is unique
+    group_name="blood_storage",
+)
+def blood_storage_cleaned(redshift: RedshiftClientResource) -> MaterializeResult:
+    query = """
+        select count(*) from public.blood_storage;
+    """
+    client = redshift.get_client()
+    client.execute_query(query)
+
+    return MaterializeResult(  # Adding query to metadata to use it getting lineage from it with sql parser
+        metadata={
+            "Query": MetadataValue.text(query),
+        }
+    )

metadata-ingestion-modules/dagster-plugin/setup.py

@@ -15,6 +15,12 @@ def get_long_description():
 
 rest_common = {"requests", "requests_file"}
 
+sqlglot_lib = {
+    # Using an Acryl fork of sqlglot.
+    # https://github.com/tobymao/sqlglot/compare/main...hsheth2:sqlglot:main?expand=1
+    "acryl-sqlglot[rs]==24.0.1.dev7",
+}
+
 _version: str = package_metadata["__version__"]
 _self_pin = (
     f"=={_version}" if not (_version.endswith("dev0") or "docker" in _version) else ""
@@ -28,6 +34,7 @@ def get_long_description():
     # Ignoring the dependency below because it causes issues with the vercel built wheel install
     # f"acryl-datahub[datahub-rest]{_self_pin}",
     "acryl-datahub[datahub-rest]",
+    *sqlglot_lib,
 }
 
 mypy_stubs = {
@@ -51,6 +58,9 @@ def get_long_description():
 base_dev_requirements = {
     *base_requirements,
     *mypy_stubs,
+    "dagster-aws >= 0.11.0",
+    "dagster-snowflake >= 0.11.0",
+    "dagster-snowflake-pandas >= 0.11.0",
     "black==22.12.0",
     "coverage>=5.1",
     "flake8>=6.0.0",