feat(oso_dagster): asset caching for sqlmesh related assets

ravenac95 · ravenac95 · commit 8fd146873590 · 2025-08-01T15:14:48.000Z
diff --git a/warehouse/oso_dagster/assets/sqlmesh.py b/warehouse/oso_dagster/assets/sqlmesh.py
@@ -1,14 +1,22 @@
 import copy
 
 import dagster as dg
-from dagster import AssetExecutionContext, AssetSelection, RunConfig, define_asset_job
+from dagster import (
+    AssetExecutionContext,
+    AssetSelection,
+    ResourceParam,
+    RunConfig,
+    define_asset_job,
+)
 from dagster_sqlmesh import (
+    DagsterSQLMeshCacheOptions,
     PlanOptions,
     SQLMeshContextConfig,
     SQLMeshDagsterTranslator,
     SQLMeshResource,
     sqlmesh_assets,
 )
+from oso_dagster.config import DagsterConfig
 from oso_dagster.factories.common import AssetFactoryResponse
 from oso_dagster.resources.trino import TrinoResource
 from oso_dagster.utils.asynctools import multiple_async_contexts
@@ -60,29 +68,38 @@ class SQLMeshRunConfig(dg.Config):
 @early_resources_asset_factory()
 def sqlmesh_factory(
     sqlmesh_infra_config: dict,
-    sqlmesh_config: SQLMeshContextConfig,
+    sqlmesh_context_config: SQLMeshContextConfig,
     sqlmesh_translator: SQLMeshDagsterTranslator,
+    sqlmesh_cache_options: DagsterSQLMeshCacheOptions,
 ):
     dev_environment = sqlmesh_infra_config["dev_environment"]
     environment = sqlmesh_infra_config["environment"]
 
+    print(sqlmesh_cache_options)
+
     @sqlmesh_assets(
-        config=sqlmesh_config,
+        config=sqlmesh_context_config,
         environment=environment,
         dagster_sqlmesh_translator=sqlmesh_translator,
         enabled_subsetting=False,
         op_tags=op_tags,
+        cache_options=sqlmesh_cache_options,
     )
     async def sqlmesh_project(
         context: AssetExecutionContext,
+        global_config: ResourceParam[DagsterConfig],
         sqlmesh: SQLMeshResource,
         trino: TrinoResource,
         config: SQLMeshRunConfig,
     ):
         restate_models = config.restate_models[:] if config.restate_models else []
 
-        async with multiple_async_contexts(
-            trino=trino.ensure_available(log_override=context.log),
+        # We use this helper function so we can run sqlmesh both locally and in
+        # a k8s environment
+        def run_sqlmesh(
+            context: AssetExecutionContext,
+            sqlmesh: SQLMeshResource,
+            config: SQLMeshRunConfig,
         ):
             # If restate_by_entity_category is True, dynamically identify models based on entity categories
             if config.restate_by_entity_category:
@@ -133,6 +150,18 @@ async def sqlmesh_project(
             ):
                 yield result
 
+        # Trino can either be `local-trino` or `trino`
+        if "trino" in global_config.sqlmesh_gateway:
+            async with multiple_async_contexts(
+                trino=trino.ensure_available(log_override=context.log),
+            ):
+                for result in run_sqlmesh(context, sqlmesh, config):
+                    yield result
+        else:
+            # If we are not running trino we are using duckdb
+            for result in run_sqlmesh(context, sqlmesh, config):
+                yield result
+
     all_assets_selection = AssetSelection.assets(sqlmesh_project)
 
     return AssetFactoryResponse(
diff --git a/warehouse/oso_dagster/assets/sqlmesh_export.py b/warehouse/oso_dagster/assets/sqlmesh_export.py
@@ -1,53 +1,110 @@
-import logging
 import typing as t
+from pathlib import Path
+from traceback import print_exception
 
-from dagster import AssetKey, ResourceParam
-from dagster_sqlmesh import DagsterSQLMeshController, SQLMeshContextConfig
+import structlog
+from dagster import AssetKey, AssetsDefinition, ResourceParam
+from dagster_sqlmesh import (
+    DagsterSQLMeshCacheOptions,
+    DagsterSQLMeshController,
+    SQLMeshContextConfig,
+)
 from dagster_sqlmesh.controller.base import DEFAULT_CONTEXT_FACTORY
+from oso_dagster.resources.sqlmesh import SQLMeshExportedAssetDefinition
 from sqlmesh.core.model import Model
 
 from ..factories import AssetFactoryResponse, early_resources_asset_factory
 from ..resources import PrefixedSQLMeshTranslator, SQLMeshExporter
 
-logger = logging.getLogger(__name__)
+logger = structlog.get_logger(__name__)
 
 
 @early_resources_asset_factory()
 def sqlmesh_export_factory(
     sqlmesh_infra_config: dict,
-    sqlmesh_config: SQLMeshContextConfig,
+    sqlmesh_context_config: SQLMeshContextConfig,
     sqlmesh_translator: PrefixedSQLMeshTranslator,
     sqlmesh_exporters: ResourceParam[t.List[SQLMeshExporter]],
+    sqlmesh_cache_options: DagsterSQLMeshCacheOptions,
 ):
     environment = sqlmesh_infra_config["environment"]
 
     controller = DagsterSQLMeshController.setup_with_config(
-        config=sqlmesh_config,
+        config=sqlmesh_context_config,
         context_factory=DEFAULT_CONTEXT_FACTORY,
     )
-    assets = []
-
-    with controller.instance(environment) as mesh:
-        models = mesh.models()
-        models_to_export: t.List[t.Tuple[Model, AssetKey]] = []
-        for name, model in models.items():
-            if "export" not in model.tags:
-                continue
-            models_to_export.append(
-                (
-                    model,
-                    sqlmesh_translator.get_asset_key(mesh.context, model.fqn),
-                )
-            )
+    assets: list[AssetsDefinition] = []
+
+    missing_exporters = set()
+
+    # Hack for now to cache all of the sqlmesh export assets
+    if sqlmesh_cache_options.enabled:
+        logger.info("SQLMesh export cache is enabled, caching assets")
+
+        cache_dir = Path(sqlmesh_cache_options.cache_dir)
+        exporter_assets_cache_dir = cache_dir.joinpath("sqlmesh_export_assets")
+        # Ensure the cache directory exists
+        exporter_assets_cache_dir.mkdir(parents=True, exist_ok=True)
 
-        # Create a export assets for this
         for exporter in sqlmesh_exporters:
-            asset_def = exporter.create_export_asset(
-                mesh,
-                sqlmesh_translator,
-                to_export=models_to_export,
+            # Load the definition from the exporter
+            cache_file = exporter_assets_cache_dir / f"{exporter.name()}.json"
+            if cache_file.exists():
+                logger.debug(f"Loading cached asset definition for {exporter.name()}")
+                exporter_asset_def = SQLMeshExportedAssetDefinition.model_validate_json(
+                    cache_file.read_text()
+                )
+                assets.append(exporter.asset_from_definition(exporter_asset_def))
+            else:
+                missing_exporters.add(exporter.name())
+    else:
+        missing_exporters = set(exporter.name() for exporter in sqlmesh_exporters)
+
+    if len(missing_exporters) > 0:
+        try:
+            with controller.instance(environment) as mesh:
+                models = mesh.models()
+                models_to_export: t.List[t.Tuple[Model, AssetKey]] = []
+                for name, model in models.items():
+                    if "export" not in model.tags:
+                        continue
+                    models_to_export.append(
+                        (
+                            model,
+                            sqlmesh_translator.get_asset_key(mesh.context, model.fqn),
+                        )
+                    )
+
+                # Create a export assets for this
+                for exporter in sqlmesh_exporters:
+                    asset_def = exporter.create_export_asset(
+                        mesh,
+                        sqlmesh_translator,
+                        to_export=models_to_export,
+                    )
+                    if sqlmesh_cache_options.enabled:
+                        cache_dir = Path(sqlmesh_cache_options.cache_dir)
+                        exporter_assets_cache_dir = cache_dir.joinpath(
+                            "sqlmesh_export_assets"
+                        )
+
+                        # Save the asset definition to the cache
+                        cache_file = (
+                            exporter_assets_cache_dir / f"{exporter.name()}.json"
+                        )
+                        logger.debug(
+                            f"Caching asset definition for {exporter.name()} at {cache_file}"
+                        )
+                        cache_file.write_text(asset_def.model_dump_json())
+
+                    assets.append(exporter.asset_from_definition(asset_def))
+                    logger.debug(f"exporting for {exporter.name()}")
+        except Exception as e:
+            logger.exception(
+                "Failed to create SQLMesh export assets",
+                error=str(e),
             )
-            assets.append(asset_def)
-            logger.debug(f"exporting for {exporter.__class__.__name__}")
+            print_exception(e)
+            raise e
 
     return AssetFactoryResponse(assets=assets)
diff --git a/warehouse/oso_dagster/config.py b/warehouse/oso_dagster/config.py
@@ -60,6 +60,10 @@ class DagsterConfig(BaseSettings):
     sqlmesh_catalog: str = "iceberg"
     sqlmesh_schema: str = "oso"
     sqlmesh_bq_export_dataset_id: str = "oso"
+    sqlmesh_dagster_asset_cache_enabled: bool = False
+    sqlmesh_dagster_asset_cache_dir: str = ""
+    sqlmesh_dagster_asset_enable_ttl: bool = True
+    sqlmesh_dagster_asset_ttl_seconds: int = 60 * 5
 
     enable_k8s_executor: bool = False
 
diff --git a/warehouse/oso_dagster/definitions/legacy.py b/warehouse/oso_dagster/definitions/legacy.py
@@ -1,12 +1,16 @@
 import logging
 import typing as t
 
+from dotenv import load_dotenv
 from oso_core.logging import setup_module_logging
 from oso_core.logging.decorators import time_function
 
+load_dotenv()
+
 logger = logging.getLogger(__name__)
 
 setup_module_logging("oso_dagster")
+setup_module_logging("dagster_sqlmesh")
 
 
 @time_function(logger, override_name="legacy_main")
@@ -68,7 +72,7 @@ def load_definitions(
         dlt_warehouse_destination: Destination,
         dlt: DagsterDltResource,
         alert_manager: AlertManager,
-        sqlmesh_config: SQLMeshContextConfig,
+        sqlmesh_context_config: SQLMeshContextConfig,
         sqlmesh_infra_config: t.Dict[str, str],
         sqlmesh: SQLMeshResource,
         k8s: K8sResource,
@@ -118,7 +122,7 @@ def load_definitions(
             "dlt_warehouse_destination": dlt_warehouse_destination,
             "project_id": global_config.project_id,
             "alert_manager": alert_manager,
-            "sqlmesh_config": sqlmesh_config,
+            "sqlmesh_config": sqlmesh_context_config,
             "sqlmesh_infra_config": sqlmesh_infra_config,
             "k8s": k8s,
             "trino": trino,
diff --git a/warehouse/oso_dagster/definitions/resources.py b/warehouse/oso_dagster/definitions/resources.py
@@ -1,10 +1,15 @@
+import os
 import typing as t
 
 import structlog
 from dagster import ConfigurableIOManagerFactory
 from dagster_dlt import DagsterDltResource
 from dagster_gcp import BigQueryResource, GCSResource
-from dagster_sqlmesh import SQLMeshContextConfig, SQLMeshResource
+from dagster_sqlmesh import (
+    DagsterSQLMeshCacheOptions,
+    SQLMeshContextConfig,
+    SQLMeshResource,
+)
 from dlt.common.destination import Destination
 from oso_core.logging.decorators import time_function
 from oso_dagster.cbt.cbt import CBTResource
@@ -41,6 +46,7 @@
     LogAlertManager,
 )
 from oso_dagster.utils.secrets import SecretResolver
+from sqlmesh.core.config.connection import DuckDBConnectionConfig, TrinoConnectionConfig
 
 from ..config import DagsterConfig
 from ..utils import GCPSecretResolver, LocalSecretResolver
@@ -103,15 +109,55 @@ def gcs_resource_factory(global_config: DagsterConfig) -> GCSResource:
     return GCSResource(project=global_config.project_id)
 
 
-@resource_factory("sqlmesh_translator")
+@resource_factory("sqlmesh_cache_options")
 @time_function(logger)
-def sqlmesh_translator_factory():
-    return PrefixedSQLMeshTranslator("sqlmesh")
+def sqlmesh_cache_options_factory(
+    global_config: DagsterConfig,
+) -> DagsterSQLMeshCacheOptions:
+    """Factory function to create SQLMesh cache options."""
+
+    return DagsterSQLMeshCacheOptions(
+        enabled=global_config.sqlmesh_dagster_asset_cache_enabled,
+        cache_dir=global_config.sqlmesh_dagster_asset_cache_dir,
+        enable_ttl=global_config.sqlmesh_dagster_asset_enable_ttl,
+        ttl_seconds=global_config.sqlmesh_dagster_asset_ttl_seconds,
+    )
 
 
-@resource_factory("sqlmesh_config")
+@resource_factory("sqlmesh_translator")
+@time_function(logger)
+def sqlmesh_translator_factory(
+    global_config: DagsterConfig, sqlmesh_context_config: SQLMeshContextConfig
+) -> PrefixedSQLMeshTranslator:
+    sqlmesh_config = sqlmesh_context_config.sqlmesh_config
+    if sqlmesh_config is None:
+        raise ValueError("SQLMesh configuration is not set in the context config.")
+    gateway = sqlmesh_config.gateways[global_config.sqlmesh_gateway]
+    connection = gateway.connection
+    if connection is None:
+        raise ValueError("SQLMesh gateway connection is not set in the context config.")
+    default_catalog = ""
+    match connection:
+        case TrinoConnectionConfig(catalog=catalog):
+            default_catalog = catalog
+        case DuckDBConnectionConfig(database=database):
+            if not database:
+                raise ValueError(
+                    "DuckDB database path is not set in the context config."
+                )
+            default_catalog = os.path.basename(database).split(".")[0]
+        case _:
+            raise ValueError(
+                f"Unsupported SQLMesh connection type: {type(connection).__name__}"
+            )
+    return PrefixedSQLMeshTranslator("sqlmesh", default_catalog)
+
+
+@resource_factory("sqlmesh_context_config")
 @time_function(logger)
-def sqlmesh_config_factory(global_config: DagsterConfig) -> SQLMeshContextConfig:
+def sqlmesh_context_config_factory(
+    global_config: DagsterConfig,
+) -> SQLMeshContextConfig:
     return SQLMeshContextConfig(
         path=global_config.sqlmesh_dir,
         gateway=global_config.sqlmesh_gateway,
@@ -145,11 +191,11 @@ def trino_resource_factory(
 
 @resource_factory("sqlmesh")
 def sqlmesh_resource_factory(
-    sqlmesh_config: SQLMeshContextConfig,
+    sqlmesh_context_config: SQLMeshContextConfig,
 ) -> SQLMeshResource:
     """Factory function to create a SQLMesh resource."""
     return SQLMeshResource(
-        config=sqlmesh_config,
+        config=sqlmesh_context_config,
     )
 
 
@@ -347,7 +393,8 @@ def default_resource_registry():
     registry.add(gcs_resource_factory)
     registry.add(sqlmesh_resource_factory)
     registry.add(sqlmesh_translator_factory)
-    registry.add(sqlmesh_config_factory)
+    registry.add(sqlmesh_context_config_factory)
+    registry.add(sqlmesh_cache_options_factory)
     registry.add(k8s_resource_factory)
     registry.add(trino_resource_factory)
     registry.add(sqlmesh_exporter_factory)
diff --git a/warehouse/oso_dagster/resources/sqlmesh.py b/warehouse/oso_dagster/resources/sqlmesh.py
