Fixes some duckdb running issues (#2863)

opensource-observer · Jan 25, 2025 · ccdb48e · ccdb48e
1 parent 436db50
commit ccdb48e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/warehouse/metrics_mesh/macros/oso_source.py b/warehouse/metrics_mesh/macros/oso_source.py
@@ -46,8 +46,9 @@ def oso_source(evaluator: MacroEvaluator, table_name: exp.Expression):
         select * from source_public.table
         ```
     """
-    if evaluator.runtime_stage in ["loading", "creating"]:
+    if evaluator.runtime_stage in ["loading"]:
         return table_name
+
     if evaluator.engine_adapter.dialect == "duckdb":
         # We hardcode the rewrite rules for duckdb here to 1) keep things consistent
         # and 2) ensure tests can run even on production when using duckdb

diff --git a/warehouse/metrics_tools/local/utils.py b/warehouse/metrics_tools/local/utils.py
@@ -35,7 +35,7 @@
     "opensource-observer.oso_playground.stg_farcaster__profiles": "bigquery.oso.stg_farcaster__profiles",
     "opensource-observer.oso_playground.stg_lens__owners": "bigquery.oso.stg_lens__owners",
     "opensource-observer.oso_playground.stg_lens__profiles": "bigquery.oso.stg_lens__profiles",
-    "opensource-observer.oso_playground.stg_ossd__current_collections": "bigquery.stg_ossd__current_collections",
+    "opensource-observer.oso_playground.stg_ossd__current_collections": "bigquery.oso.stg_ossd__current_collections",
     "opensource-observer.oso_playground.stg_ossd__current_projects": "bigquery.oso.stg_ossd__current_projects",
     "opensource-observer.oso_playground.stg_ossd__current_repositories": "bigquery.oso.stg_ossd__current_repositories",
     "opensource-observer.oso_playground.timeseries_events_aux_issues_by_artifact_v0": "bigquery.oso.timeseries_events_aux_issues_by_artifact_v0",

diff --git a/warehouse/metrics_tools/source/rewrite.py b/warehouse/metrics_tools/source/rewrite.py
@@ -16,8 +16,16 @@
 
 
 def oso_source_for_pymodel(context: ExecutionContext, table_name: str) -> exp.Table:
+    # Rewrite rules can be injected into the sqlmesh context via the sqlmesh
+    # config. However, in general we don't want to modify the source unless it's
+    # running on duckdb locally. So this allows rewriting as an override but
+    # defaults to using duckdb if we're currently running duckdb.
+    oso_source_rewrite_config_default: t.List[dict] = []
+    if context.engine_adapter.dialect == "duckdb":
+        oso_source_rewrite_config_default = DUCKDB_REWRITE_RULES
     oso_source_rewrite_config = t.cast(
-        t.List[dict], context.var("oso_source_rewrite", [])
+        t.List[dict],
+        context.var("oso_source_rewrite", oso_source_rewrite_config_default),
     )
     return oso_source_rewrite(oso_source_rewrite_config, table_name)