[prototype] update to match context API #3

hooli_data_eng/assets/forecasting/__init__.py

@@ -99,7 +99,7 @@ def model_stats_by_month(
 ) -> Output[pd.DataFrame]:
     """Model errors by month"""
     a, b = order_forecast_model
-    target_date = pd.to_datetime(context.asset_partition_key_for_output())
+    target_date = pd.to_datetime(context.partition_context.partition_key)
     target_month = target_date.month
     weekly_order_summary["order_date"] = pd.to_datetime(
         weekly_order_summary["order_date"]

hooli_data_eng/assets/marketing/__init__.py

@@ -13,13 +13,13 @@
 
 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 
+# 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), 
+    key_prefix="MARKETING",
+    freshness_policy=FreshnessPolicy(maximum_lag_minutes=24*60),
     auto_materialize_policy=allow_outdated_parents_policy,
     compute_kind="pandas",
     op_tags={"owner": "[email protected]"},
@@ -29,7 +29,7 @@ def avg_orders(context: AssetExecutionContext, company_perf: pd.DataFrame) -> pd
     """ Computes avg order KPI, must be updated regularly for exec dashboard """
 
     return pd.DataFrame({
-        "avg_order": company_perf['total_revenue'] / company_perf['n_orders'] 
+        "avg_order": company_perf['total_revenue'] / company_perf['n_orders']
     })
 
 @asset_check(
@@ -39,14 +39,14 @@ def avg_orders(context: AssetExecutionContext, company_perf: pd.DataFrame) -> pd
 def check_avg_orders(context, avg_orders: pd.DataFrame):
     avg = avg_orders['avg_order'][0]
     return AssetCheckResult(
-        passed= True if (avg < 50) else False, 
+        passed= True if (avg < 50) else False,
         metadata={"actual average": avg, "threshold": 50}
     )
 
 @asset(
-    key_prefix="MARKETING", 
-    freshness_policy=FreshnessPolicy(maximum_lag_minutes=24*60), 
-    compute_kind="snowflake", 
+    key_prefix="MARKETING",
+    freshness_policy=FreshnessPolicy(maximum_lag_minutes=24*60),
+    compute_kind="snowflake",
     metadata={
         "owner": "[email protected]"
     },
@@ -76,7 +76,7 @@ def min_order(context, company_perf: pd.DataFrame) -> pd.DataFrame:
 )
 def key_product_deepdive(context, sku_stats):
     """ Creates a file for a BI tool based on the current quarters top product, represented as a dynamic partition """
-    key_sku = context.partition_key
+    key_sku = context.partition_context.partition_key
     sku = sku_stats[sku_stats['sku'] == key_sku]
     context.add_output_metadata(
         {

hooli_data_eng/assets/raw_data/__init__.py

@@ -27,7 +27,7 @@ def _daily_partition_seq(start, end):
     start = pd.to_datetime(start)
     end = pd.to_datetime(end)
     daily_diffs = int((end - start) / timedelta(hours=24))
-    
+
     return [str(start + timedelta(hours=i)) for i in range(daily_diffs)]
 
 
@@ -41,7 +41,7 @@ def users(context, api: RawDataAPI) -> pd.DataFrame:
     """A table containing all users data"""
     # during a backfill the partition range will span multiple hours
     # during a single run the partition range will be for a single hour
-    first_partition, last_partition = context.asset_partitions_time_window_for_output()
+    first_partition, last_partition = context.partition_context.partition_time_window
     partition_seq = _daily_partition_seq(first_partition, last_partition)
     all_users = []
     for partition in partition_seq:
@@ -69,23 +69,23 @@ def check_users(context, users: pd.DataFrame):
     partitions_def=daily_partitions,
     metadata={"partition_expr": "DT"},
     retry_policy=RetryPolicy(
-        max_retries=3, 
-        delay=1, 
+        max_retries=3,
+        delay=1,
         backoff=Backoff.LINEAR,
         jitter=Jitter.FULL
     ),
     backfill_policy=BackfillPolicy.single_run()
 )
 def orders(context, api: RawDataAPI) -> pd.DataFrame:
     """A table containing all orders that have been placed"""
-    first_partition, last_partition = context.asset_partitions_time_window_for_output()
+    first_partition, last_partition = context.partition_context.partition_time_window
     partition_seq = _daily_partition_seq(first_partition, last_partition)
     all_orders = []
     for partition in partition_seq:
         resp = api.get_orders(partition)
         users = pd.read_json(resp.json())
         all_orders.append(users)
-    
+
     all_orders_df = pd.concat(all_orders)
     all_orders_df['dt'] = pd.to_datetime(all_orders_df['dt'], unit = "ms")
     return all_orders_df