make amps work better (#49)

dagster-io · Oct 24, 2023 · f911a36 · f911a36
1 parent 59e5f4b
commit f911a36
Show file tree

Hide file tree

Showing 7 changed files with 38 additions and 25 deletions.
diff --git a/dbt_project/models/ANALYTICS/company_stats.sql b/dbt_project/models/ANALYTICS/company_stats.sql
@@ -1,8 +1,3 @@
-{{
-        config(
-                dagster_auto_materialize_policy={"type":"lazy"}
-        )
-}}
 select
         order_date,
         company,

diff --git a/dbt_project/models/ANALYTICS/order_stats.sql b/dbt_project/models/ANALYTICS/order_stats.sql
@@ -1,8 +1,3 @@
-{{
-        config(
-                dagster_auto_materialize_policy={"type":"eager"}
-        )
-}}
 select
         {{ date_trunc("day", "order_date") }} as order_date,
         count(*) as n_orders,

diff --git a/dbt_project/models/ANALYTICS/sku_stats.sql b/dbt_project/models/ANALYTICS/sku_stats.sql
@@ -1,8 +1,3 @@
-{{
-        config(
-                dagster_auto_materialize_policy={"type":"eager"}
-        )
-}}
 select
         order_date,
         sku,

diff --git a/dbt_project/models/ANALYTICS/weekly_order_summary.sql b/dbt_project/models/ANALYTICS/weekly_order_summary.sql
@@ -1,7 +1,6 @@
 
 {{
         config(
-                dagster_auto_materialize_policy={"type":"eager"},
                 dagster_freshness_policy={"cron_schedule": "0 9 * * MON", "maximum_lag_minutes": (24+9)*60}
         )
 }}

diff --git a/dbt_project/models/MARKETING/company_perf.sql b/dbt_project/models/MARKETING/company_perf.sql
@@ -1,8 +1,3 @@
-{{
-        config(
-                dagster_auto_materialize_policy={"type":"lazy"}
-        )
-}}
 select
         company,
         sum(n_orders) as n_orders,

diff --git a/hooli_data_eng/assets/dbt_assets.py b/hooli_data_eng/assets/dbt_assets.py
@@ -7,6 +7,8 @@
 )
 from dagster_dbt.asset_decorator import dbt_assets
 from dagster import (
+    AutoMaterializePolicy,
+    AutoMaterializeRule,
     AssetKey,
     DailyPartitionsDefinition,
     WeeklyPartitionsDefinition,
@@ -35,6 +37,15 @@
     file_relative_path(__file__, "../../dbt_project/target/manifest.json")
 )
 
+allow_outdated_parents_policy = AutoMaterializePolicy.eager().without_rules(
+    AutoMaterializeRule.skip_on_parent_outdated()
+)
+
+allow_outdated_and_missing_parents_policy = AutoMaterializePolicy.eager().without_rules(
+    AutoMaterializeRule.skip_on_parent_outdated(), 
+    AutoMaterializeRule.skip_on_parent_missing() # non-partitioned assets should run even if some upstream partitions are missing
+)
+
 
 class CustomDagsterDbtTranslator(DagsterDbtTranslator):
     def get_description(self, dbt_resource_props: Mapping[str, Any]) -> str:
@@ -51,6 +62,9 @@ def get_asset_key(self, dbt_resource_props: Mapping[str, Any]) -> AssetKey:
         if node_path == "models/sources.yml":
             prefix = "RAW_DATA"
 
+        if node_path == "MARKETING/company_perf.sql":
+            prefix = "ANALYTICS"
+
         return AssetKey([prefix, dbt_resource_props["name"]])
 
     def get_group_name(self, dbt_resource_props: Mapping[str, Any]):
@@ -60,6 +74,9 @@ def get_group_name(self, dbt_resource_props: Mapping[str, Any]):
         if node_path == "models/sources.yml":
             prefix = "RAW_DATA"
 
+        if node_path == "MARKETING/company_perf.sql":
+            prefix = "ANALYTICS"
+
         return prefix
 
     def get_metadata(self, dbt_resource_props: Mapping[str, Any]) -> Mapping[str, Any]:
@@ -71,10 +88,24 @@ def get_metadata(self, dbt_resource_props: Mapping[str, Any]) -> Mapping[str, An
         if dbt_resource_props["name"] == "users_cleaned":
             metadata = {"partition_expr": "created_at"}
 
+        if dbt_resource_props["name"] in ["company_perf", "sku_stats", "company_stats"]:
+            metadata = {}
+
         default_metadata = default_metadata_from_dbt_resource_props(dbt_resource_props)
 
         return {**default_metadata, **metadata}
 
+    def get_auto_materialize_policy(
+        self, dbt_resource_props: Mapping[str, Any]
+    ):
+        return allow_outdated_parents_policy
+
+
+class CustomDagsterDbtTranslatorForViews(CustomDagsterDbtTranslator):
+    def get_auto_materialize_policy(
+        self, dbt_resource_props: Mapping[str, Any]
+    ):
+        return allow_outdated_and_missing_parents_policy
 
 def _process_partitioned_dbt_assets(context: OpExecutionContext, dbt2: DbtCliResource):
     # map partition key range to dbt vars
@@ -145,8 +176,9 @@ def weekly_dbt_assets(context: OpExecutionContext, dbt2: DbtCliResource):
 dbt_views = load_assets_from_dbt_project(
     DBT_PROJECT_DIR,
     DBT_PROFILES_DIR,
-    key_prefix=["ANALYTICS"],
-    source_key_prefix="ANALYTICS",
+    #key_prefix=["ANALYTICS"],
+    #source_key_prefix="ANALYTICS",
     select="company_perf sku_stats company_stats",
-    node_info_to_group_fn=lambda x: "ANALYTICS",
+    #node_info_to_group_fn=lambda x: "ANALYTICS",
+    dagster_dbt_translator=CustomDagsterDbtTranslatorForViews()
 )
diff --git a/hooli_data_eng/assets/marketing/__init__.py b/hooli_data_eng/assets/marketing/__init__.py
@@ -11,14 +11,16 @@
 )
 import pandas as pd
 
+from hooli_data_eng.assets.dbt_assets import allow_outdated_parents_policy
+
 # These assets take data from a SQL table managed by 
 # dbt and create summaries using pandas 
 # The assets are updated via freshness policies 
 # and an associated reconciliation sensor
 @asset(
     key_prefix="MARKETING", 
     freshness_policy=FreshnessPolicy(maximum_lag_minutes=24*60), 
-    auto_materialize_policy=AutoMaterializePolicy.lazy(),
+    auto_materialize_policy=allow_outdated_parents_policy,
     compute_kind="pandas",
     op_tags={"owner": "[email protected]"},
     ins={"company_perf": AssetIn(key_prefix=["ANALYTICS"])}