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

metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_basic_iolets.json

@@ -22,6 +22,17 @@
         }
     }
 },
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    }
+},
 {
     "entityType": "dataFlow",
     "entityUrn": "urn:li:dataFlow:(airflow,basic_iolets,prod)",
@@ -85,6 +96,17 @@
         }
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,basic_iolets,prod),run_data_task)",

metadata-ingestion-modules/airflow-plugin/tests/integration/goldens/v1_simple_dag.json

@@ -23,6 +23,17 @@
         }
     }
 },
+{
+    "entityType": "dataFlow",
+    "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    }
+},
 {
     "entityType": "dataFlow",
     "entityUrn": "urn:li:dataFlow:(airflow,simple_dag,prod)",
@@ -86,6 +97,17 @@
         }
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),task_1)",
@@ -540,6 +562,17 @@
         }
     }
 },
+{
+    "entityType": "dataJob",
+    "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)",
+    "changeType": "UPSERT",
+    "aspectName": "status",
+    "aspect": {
+        "json": {
+            "removed": false
+        }
+    }
+},
 {
     "entityType": "dataJob",
     "entityUrn": "urn:li:dataJob:(urn:li:dataFlow:(airflow,simple_dag,prod),run_another_data_task)",

metadata-ingestion-modules/airflow-plugin/tests/unit/test_airflow.py

@@ -316,7 +316,7 @@ def test_lineage_backend(mock_emit, inlets, outlets, capture_executions):
         assert all(map(lambda let: isinstance(let, Dataset), op2.outlets))
 
         # Check that the right things were emitted.
-        assert mock_emitter.emit.call_count == 17 if capture_executions else 9
+        assert mock_emitter.emit.call_count == 19 if capture_executions else 11
 
         # TODO: Replace this with a golden file-based comparison.
         assert mock_emitter.method_calls[0].args[0].aspectName == "dataFlowInfo"
@@ -325,127 +325,139 @@ def test_lineage_backend(mock_emit, inlets, outlets, capture_executions):
             == "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
         )
 
-        assert mock_emitter.method_calls[1].args[0].aspectName == "ownership"
+        assert mock_emitter.method_calls[1].args[0].aspectName == "status"
         assert (
             mock_emitter.method_calls[1].args[0].entityUrn
             == "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
         )
 
-        assert mock_emitter.method_calls[2].args[0].aspectName == "globalTags"
+        assert mock_emitter.method_calls[2].args[0].aspectName == "ownership"
         assert (
             mock_emitter.method_calls[2].args[0].entityUrn
             == "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
         )
 
-        assert mock_emitter.method_calls[3].args[0].aspectName == "dataJobInfo"
+        assert mock_emitter.method_calls[3].args[0].aspectName == "globalTags"
         assert (
             mock_emitter.method_calls[3].args[0].entityUrn
-            == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
+            == "urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod)"
         )
 
-        assert mock_emitter.method_calls[4].args[0].aspectName == "dataJobInputOutput"
+        assert mock_emitter.method_calls[4].args[0].aspectName == "dataJobInfo"
         assert (
             mock_emitter.method_calls[4].args[0].entityUrn
             == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
         )
+
+        assert mock_emitter.method_calls[5].args[0].aspectName == "status"
         assert (
-            mock_emitter.method_calls[4].args[0].aspect.inputDatajobs[0]
+            mock_emitter.method_calls[5].args[0].entityUrn
+            == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
+        )
+
+        assert mock_emitter.method_calls[6].args[0].aspectName == "dataJobInputOutput"
+        assert (
+            mock_emitter.method_calls[6].args[0].entityUrn
+            == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
+        )
+        assert (
+            mock_emitter.method_calls[6].args[0].aspect.inputDatajobs[0]
             == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task1_upstream)"
         )
         assert (
-            mock_emitter.method_calls[4].args[0].aspect.inputDatajobs[1]
+            mock_emitter.method_calls[6].args[0].aspect.inputDatajobs[1]
             == "urn:li:dataJob:(urn:li:dataFlow:(airflow,testDag,PROD),testTask)"
         )
         assert (
-            mock_emitter.method_calls[4].args[0].aspect.inputDatasets[0]
+            mock_emitter.method_calls[6].args[0].aspect.inputDatasets[0]
             == "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
         )
         assert (
-            mock_emitter.method_calls[4].args[0].aspect.outputDatasets[0]
+            mock_emitter.method_calls[6].args[0].aspect.outputDatasets[0]
             == "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
         )
 
-        assert mock_emitter.method_calls[5].args[0].aspectName == "datasetKey"
+        assert mock_emitter.method_calls[7].args[0].aspectName == "datasetKey"
         assert (
-            mock_emitter.method_calls[5].args[0].entityUrn
+            mock_emitter.method_calls[7].args[0].entityUrn
             == "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
         )
 
-        assert mock_emitter.method_calls[6].args[0].aspectName == "datasetKey"
+        assert mock_emitter.method_calls[8].args[0].aspectName == "datasetKey"
         assert (
-            mock_emitter.method_calls[6].args[0].entityUrn
+            mock_emitter.method_calls[8].args[0].entityUrn
             == "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
         )
 
-        assert mock_emitter.method_calls[7].args[0].aspectName == "ownership"
+        assert mock_emitter.method_calls[9].args[0].aspectName == "ownership"
         assert (
-            mock_emitter.method_calls[7].args[0].entityUrn
+            mock_emitter.method_calls[9].args[0].entityUrn
             == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
         )
 
-        assert mock_emitter.method_calls[8].args[0].aspectName == "globalTags"
+        assert mock_emitter.method_calls[10].args[0].aspectName == "globalTags"
         assert (
-            mock_emitter.method_calls[8].args[0].entityUrn
+            mock_emitter.method_calls[10].args[0].entityUrn
             == "urn:li:dataJob:(urn:li:dataFlow:(airflow,test_lineage_is_sent_to_backend,prod),task2)"
         )
 
         if capture_executions:
             assert (
-                mock_emitter.method_calls[9].args[0].aspectName
+                mock_emitter.method_calls[11].args[0].aspectName
                 == "dataProcessInstanceProperties"
             )
             assert (
-                mock_emitter.method_calls[9].args[0].entityUrn
+                mock_emitter.method_calls[11].args[0].entityUrn
                 == "urn:li:dataProcessInstance:5e274228107f44cc2dd7c9782168cc29"
             )
 
             assert (
-                mock_emitter.method_calls[10].args[0].aspectName
+                mock_emitter.method_calls[12].args[0].aspectName
                 == "dataProcessInstanceRelationships"
             )
             assert (
-                mock_emitter.method_calls[10].args[0].entityUrn
+                mock_emitter.method_calls[12].args[0].entityUrn
                 == "urn:li:dataProcessInstance:5e274228107f44cc2dd7c9782168cc29"
             )
             assert (
-                mock_emitter.method_calls[11].args[0].aspectName
+                mock_emitter.method_calls[13].args[0].aspectName
                 == "dataProcessInstanceInput"
             )
             assert (
-                mock_emitter.method_calls[11].args[0].entityUrn
+                mock_emitter.method_calls[13].args[0].entityUrn
                 == "urn:li:dataProcessInstance:5e274228107f44cc2dd7c9782168cc29"
             )
             assert (
-                mock_emitter.method_calls[12].args[0].aspectName
+                mock_emitter.method_calls[14].args[0].aspectName
                 == "dataProcessInstanceOutput"
             )
             assert (
-                mock_emitter.method_calls[12].args[0].entityUrn
+                mock_emitter.method_calls[14].args[0].entityUrn
                 == "urn:li:dataProcessInstance:5e274228107f44cc2dd7c9782168cc29"
             )
-            assert mock_emitter.method_calls[13].args[0].aspectName == "datasetKey"
+            assert mock_emitter.method_calls[15].args[0].aspectName == "datasetKey"
             assert (
-                mock_emitter.method_calls[13].args[0].entityUrn
+                mock_emitter.method_calls[15].args[0].entityUrn
                 == "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableConsumed,PROD)"
             )
-            assert mock_emitter.method_calls[14].args[0].aspectName == "datasetKey"
+            assert mock_emitter.method_calls[16].args[0].aspectName == "datasetKey"
             assert (
-                mock_emitter.method_calls[14].args[0].entityUrn
+                mock_emitter.method_calls[16].args[0].entityUrn
                 == "urn:li:dataset:(urn:li:dataPlatform:snowflake,mydb.schema.tableProduced,PROD)"
             )
             assert (
-                mock_emitter.method_calls[15].args[0].aspectName
+                mock_emitter.method_calls[17].args[0].aspectName
                 == "dataProcessInstanceRunEvent"
             )
             assert (
-                mock_emitter.method_calls[15].args[0].entityUrn
+                mock_emitter.method_calls[17].args[0].entityUrn
                 == "urn:li:dataProcessInstance:5e274228107f44cc2dd7c9782168cc29"
             )
             assert (
-                mock_emitter.method_calls[16].args[0].aspectName
+                mock_emitter.method_calls[18].args[0].aspectName
                 == "dataProcessInstanceRunEvent"
             )
             assert (
-                mock_emitter.method_calls[16].args[0].entityUrn
+                mock_emitter.method_calls[18].args[0].entityUrn
                 == "urn:li:dataProcessInstance:5e274228107f44cc2dd7c9782168cc29"
             )

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),
+        }
+    )