feat(dagster): instrument resource loading times (#4658)

* feat(dagster): create a resource registry to lazily load resources

* fix(dagster): fixes sql factory

* fix(dagster): remove ecosystems data

* fix(dagster): fix instrumented defintions loading

* fix(dagster): typing fixes for dlt factory

* chore(ops): set to deploy with dagster eagerly loading tables

* fix

* Update warehouse/oso_dagster/definitions/resources.py

Co-authored-by: Copilot <[email protected]>

* update

* fix(dagster): move some things around

* fix docs

* more docs fixes

---------

Co-authored-by: Copilot <[email protected]>

Author: ravenac95

lib/oso-core/oso_core/logging/decorators.py

@@ -2,13 +2,15 @@
 Some potentially useful logging decorators
 """
 
+import functools
 import logging
 
 
 def time_function(logger: logging.Logger):
     """Decorator to time a function and log the duration."""
 
     def decorator(func):
+        @functools.wraps(func)
         def wrapper(*args, **kwargs):
             import time
 

ops/helm-charts/oso-dagster/Chart.yaml

@@ -3,7 +3,7 @@ name: oso-dagster
 description: Extension of the dagster template
 
 type: application
-version: 0.18.0
+version: 0.19.0
 appVersion: "1.0.0"
 dependencies:
 - name: dagster

ops/helm-charts/oso-dagster/values.yaml

@@ -13,6 +13,7 @@ dagster:
 
 osoDagster:
   verboseLogs: 0
+  eagerlyLoadSqlTables: 0
 
 global:
   dagsterInternalName: "dagster"

ops/k8s-apps/base/dagster/dagster.yaml

@@ -44,6 +44,7 @@ spec:
         # Include internal python libraries in the dagster logs.
         managedPythonLoggers:
           - metrics_tools
+          - oso_core
           - oso_dagster
       nameOverride: "dagster"
       ingress:

ops/k8s-apps/production/dagster/custom-helm-values.yaml

@@ -4,6 +4,10 @@ metadata:
   name: production-dagster
 spec:
   values:
+    # We should eventually set this to 0, but we are setting it to 1 to gather
+    # timing data
+    osoDagster:
+      eagerlyLoadSqlTables: 1
     pg:
       host: production-cloudsql-proxy-gcloud-sqlproxy.production-cloudsql-proxy.svc.cluster.local
       port: "5432"

warehouse/oso_dagster/assets/ecosystems.py

@@ -19,7 +19,7 @@
 class SQLMeshRunConfig(dg.Config):
     # Set this to True to restate the selected models
     restate_models: list[str] | None = None
-    
+
     # Set this to True to dynamically identify models to restate based on entity_category tags
     restate_by_entity_category: bool = False
     # List of entity categories to restate (e.g., ["project", "collection"])
@@ -79,30 +79,34 @@ async def sqlmesh_project(
         trino: TrinoResource,
         config: SQLMeshRunConfig,
     ):
-
         restate_models = config.restate_models[:] if config.restate_models else []
-        
-        # Ensure that both trino and the mcs are available
+
         async with multiple_async_contexts(
             trino=trino.ensure_available(log_override=context.log),
         ):
             # If restate_by_entity_category is True, dynamically identify models based on entity categories
             if config.restate_by_entity_category:
                 # Ensure we have entity categories to filter by
                 if not config.restate_entity_categories:
-                    context.log.info("restate_by_entity_category is True but no entity categories specified. Using both 'project' and 'collection'.")
+                    context.log.info(
+                        "restate_by_entity_category is True but no entity categories specified. Using both 'project' and 'collection'."
+                    )
                     entity_categories = ["project", "collection"]
                 else:
                     entity_categories = config.restate_entity_categories
-                    context.log.info(f"Filtering models by entity categories: {entity_categories}")
+                    context.log.info(
+                        f"Filtering models by entity categories: {entity_categories}"
+                    )
 
                 for category in entity_categories:
                     restate_models.append(f"tag:entity_category={category}")
 
             plan_options: PlanOptions = {"skip_tests": config.skip_tests}
             if config.allow_destructive_models:
-                plan_options["allow_destructive_models"] = config.allow_destructive_models
-            
+                plan_options["allow_destructive_models"] = (
+                    config.allow_destructive_models
+                )
+
             # If we specify a dev_environment, we will first plan it for safety
             if dev_environment:
                 context.log.info("Planning dev environment")
@@ -130,7 +134,7 @@ async def sqlmesh_project(
                 yield result
 
     all_assets_selection = AssetSelection.assets(sqlmesh_project)
-    
+
     return AssetFactoryResponse(
         assets=[sqlmesh_project],
         jobs=[

warehouse/oso_dagster/assets/sqlmesh.py

@@ -83,6 +83,16 @@ class DagsterConfig(BaseSettings):
     mcs_k8s_deployment_name: str = ""
     mcs_connect_timeout: int = 240
 
+    # This is a bit of a legacy configuration that we need to remove
+    cbt_search_paths: list[str] = Field(
+        default_factory=lambda: [os.path.join(os.path.dirname(__file__), "models")]
+    )
+
+    clickhouse_importer_secret_group_name: str = "clickhouse_importer"
+    clickhouse_secret_group_name: str = "clickhouse"
+
+    eagerly_load_sql_tables: bool = False
+
     @model_validator(mode="after")
     def handle_generated_config(self):
         """Handles any configurations that can be generated from other configuration values"""