feat: rudimentary scheduler (#5621)

* wip: working scheduler

* feat: add udm models asset

* clean up

* fix: move things to resource dependency injection

* fix: update to latest dagster-sqlmesh

* Update warehouse/oso_dagster/resources/udm_state.py

Co-authored-by: Raymond Cheng <[email protected]>

* Update warehouse/oso_dagster/resources/udm_engine_adapter.py

Co-authored-by: Raymond Cheng <[email protected]>

* remove prototyping code

* fix: make engine adapter loading lazy

* fix: updated lint rules for new sqlmesh version

* fix: various build fixes for updated dagster-sqlmesh

* fix: dagster run issues

* Update warehouse/oso_dagster/resources/udm_engine_adapter.py

---------

Co-authored-by: Raymond Cheng <[email protected]>

Author: ravenac95

pyproject.toml

@@ -67,7 +67,7 @@ dependencies = [
     "kr8s==0.20.9",
     "structlog>=25.4.0",
     "pandas-gbq>=0.29.2",
-    "dagster-sqlmesh>=0.20.0",
+    "dagster-sqlmesh>=0.22.0",
     "oso-core",
     "pyoso",
     "metrics-service",
@@ -110,6 +110,7 @@ pyoso = { workspace = true }
 oso-mcp = { workspace = true }
 oso-core = { workspace = true }
 metrics-service = { workspace = true }
+scheduler = { workspace = true }
 
 [tool.uv.workspace]
 members = [
@@ -119,7 +120,8 @@ members = [
     "warehouse/pynessie-gc",
     "warehouse/metrics-service",
     "lib/oso-core",
-    "warehouse/oso_semantic"
+    "warehouse/oso_semantic",
+    "warehouse/scheduler",
 ]
 
 [build-system]

uv.lock

@@ -19,7 +19,9 @@
 
 logger = logging.getLogger(__name__)
 
-CallableAliasList = t.List[t.Callable[..., t.Any] | t.Tuple[t.Callable[..., t.Any], t.List[str]]]
+CallableAliasList = t.List[
+    t.Callable[..., t.Any] | t.Tuple[t.Callable[..., t.Any], t.List[str]]
+]
 
 
 class MacroOverridingModel(model):
@@ -41,12 +43,12 @@ def model(self, *args, **kwargs):
         if len(self._additional_macros) > 0:
             macros = MacroRegistry(f"macros_for_{self.name}")
             system_macros = kwargs.get("macros", macros)
-            macros.update(system_macros) 
+            macros.update(system_macros)
             macros.update(macro.get_registry())
 
             additional_macros = create_macro_registry_from_list(self._additional_macros)
             macros.update(additional_macros)
-            if not macros.get('datetrunc'):
+            if not macros.get("datetrunc"):
                 raise ValueError("FUCL")
             kwargs["macros"] = macros
 
@@ -68,12 +70,14 @@ def escape_triple_quotes(input_string: str) -> str:
 
 
 def create_unregistered_macro_registry(
-    macros: t.List[t.Callable[..., t.Any] | t.Tuple[t.Callable[..., t.Any], t.List[str]]]
+    macros: t.List[
+        t.Callable[..., t.Any] | t.Tuple[t.Callable[..., t.Any], t.List[str]]
+    ],
 ):
     registry = MacroRegistry(f"macro_registry_{uuid.uuid4().hex}")
     for additional_macro in macros:
         if isinstance(additional_macro, tuple):
-            registry.update(create_unregistered_wrapped_macro(*additional_macro)) # type: ignore
+            registry.update(create_unregistered_wrapped_macro(*additional_macro))  # type: ignore
         else:
             registry.update(create_unregistered_wrapped_macro(additional_macro))
     return registry
@@ -218,11 +222,31 @@ def create_basic_python_env(
     return serialized
 
 
+def _dict_sort(obj: t.Any) -> str:
+    """Copied from sqlmesh for now"""
+    try:
+        if isinstance(obj, dict):
+            obj = dict(sorted(obj.items(), key=lambda x: str(x[0])))
+    except Exception:
+        logger.warning("Failed to sort non-recursive dict", exc_info=True)
+    return repr(obj)
+
+
 class PrettyExecutable(Executable):
     @classmethod
-    def value(cls, v: t.Any, is_metadata: t.Optional[bool] = None) -> Executable:
-        pretty_v = json.dumps(v, indent=1)
-        return cls(payload=pretty_v, kind=ExecutableKind.VALUE)
+    def value(
+        cls,
+        v: t.Any,
+        is_metadata: t.Optional[bool] = None,
+        sort_root_dict: bool = False,
+    ) -> Executable:
+        payload = _dict_sort(v) if sort_root_dict else repr(v)
+        pretty_v = json.dumps(payload, indent=1)
+        return Executable(
+            payload=pretty_v,
+            kind=ExecutableKind.VALUE,
+            is_metadata=is_metadata or None,
+        )
 
 
 def create_import_call_env(
@@ -291,7 +315,7 @@ def create_macro_registry_from_list(macro_list: CallableAliasList):
     registry = MacroRegistry("macros")
     for additional_macro in macro_list:
         if isinstance(additional_macro, tuple):
-            registry.update(create_unregistered_wrapped_macro(*additional_macro)) # type: ignore
+            registry.update(create_unregistered_wrapped_macro(*additional_macro))  # type: ignore
         else:
             registry.update(create_unregistered_wrapped_macro(additional_macro))
     return registry

warehouse/metrics_tools/models/tools.py

@@ -0,0 +1,40 @@
+import functools
+
+import dagster as dg
+from oso_dagster.factories.common import (
+    AssetFactoryResponse,
+    early_resources_asset_factory,
+)
+from oso_dagster.resources.udm_engine_adapter import (
+    UserDefinedModelEngineAdapterResource,
+)
+from oso_dagster.resources.udm_state import UserDefinedModelStateResource
+from scheduler.evaluator import UserDefinedModelEvaluator
+
+
+@early_resources_asset_factory()
+def udm_models() -> AssetFactoryResponse:
+    @dg.asset
+    async def udm_models(
+        context: dg.AssetExecutionContext,
+        udm_engine_adapter: dg.ResourceParam[UserDefinedModelEngineAdapterResource],
+        udm_state: dg.ResourceParam[UserDefinedModelStateResource],
+    ) -> None:
+        async with udm_state.get_client() as udm_client:
+            evaluator = UserDefinedModelEvaluator.prepare(udm_client)
+            await evaluator.evaluate(
+                functools.partial(
+                    udm_engine_adapter.get_adapter, log_override=context.log
+                )
+            )
+
+    job = dg.define_asset_job(
+        name="udm_models_job", selection=dg.AssetSelection.assets(udm_models)
+    )
+
+    return AssetFactoryResponse(
+        assets=[udm_models],
+        jobs=[job],
+        # Run this job every 2 hours to evaluate UDM models
+        schedules=[dg.ScheduleDefinition(job=job, cron_schedule="0 */2 * * *")],
+    )

warehouse/oso_dagster/assets/default/udm.py

@@ -13,8 +13,6 @@
     IntermediateAssetDep,
     IntermediateAssetOut,
     PlanOptions,
-    SQLMeshContextConfig,
-    SQLMeshDagsterTranslator,
     SQLMeshMultiAssetOptions,
     SQLMeshResource,
     sqlmesh_asset_from_multi_asset_options,
@@ -24,6 +22,7 @@
 from oso_dagster.config import DagsterConfig
 from oso_dagster.factories import AssetFactoryResponse, cacheable_asset_factory
 from oso_dagster.factories.common import CacheableDagsterContext
+from oso_dagster.resources.sqlmesh import PrefixedSQLMeshContextConfig
 from oso_dagster.resources.trino import TrinoResource
 from oso_dagster.utils.asynctools import multiple_async_contexts
 from pydantic import BaseModel
@@ -133,16 +132,15 @@ def sqlmesh_factory(
     def cacheable_sqlmesh_multi_asset_options(
         *,
         sqlmesh_infra_config: dict,
-        sqlmesh_context_config: SQLMeshContextConfig,
-        sqlmesh_translator: SQLMeshDagsterTranslator,
+        sqlmesh_context_config: ResourceParam[PrefixedSQLMeshContextConfig],
     ) -> CacheableSQLMeshMultiAssetOptions:
         environment = sqlmesh_infra_config["environment"]
 
         return CacheableSQLMeshMultiAssetOptions.from_sqlmesh_multi_asset_options(
             sqlmesh_to_multi_asset_options(
                 config=sqlmesh_context_config,
                 environment=environment,
-                dagster_sqlmesh_translator=sqlmesh_translator,
+                # dagster_sqlmesh_translator=sqlmesh_translator,
             )
         )
 
@@ -163,6 +161,7 @@ async def sqlmesh_project(
             context: AssetExecutionContext,
             global_config: ResourceParam[DagsterConfig],
             sqlmesh: SQLMeshResource,
+            sqlmesh_context_config: ResourceParam[PrefixedSQLMeshContextConfig],
             trino: TrinoResource,
             config: SQLMeshRunConfig,
         ):
@@ -210,6 +209,7 @@ def run_sqlmesh(
                     all(
                         sqlmesh.run(
                             context,
+                            config=sqlmesh_context_config,
                             environment=dev_environment,
                             plan_options=copy.deepcopy(plan_options),
                             start=config.start,
@@ -223,6 +223,7 @@ def run_sqlmesh(
                 context.log.info("Starting to process prod environment")
                 for result in sqlmesh.run(
                     context,
+                    config=sqlmesh_context_config,
                     environment=environment,
                     plan_options=copy.deepcopy(plan_options),
                     start=config.start,

warehouse/oso_dagster/assets/sqlmesh/sqlmesh.py

@@ -8,7 +8,7 @@
 from oso_dagster.config import DagsterConfig
 from oso_dagster.factories import AssetFactoryResponse, cacheable_asset_factory
 from oso_dagster.factories.common import CacheableDagsterContext
-from oso_dagster.resources import PrefixedSQLMeshTranslator, SQLMeshExporter
+from oso_dagster.resources import SQLMeshExporter
 from oso_dagster.resources.sqlmesh import SQLMeshExportedAssetDefinition
 from pydantic import BaseModel
 from sqlmesh.core.model import Model
@@ -38,8 +38,7 @@ def sqlmesh_export_factory(
     )
     def cacheable_exported_assets_defs(
         sqlmesh_infra_config: dict,
-        sqlmesh_context_config: SQLMeshContextConfig,
-        sqlmesh_translator: PrefixedSQLMeshTranslator,
+        sqlmesh_context_config: ResourceParam[SQLMeshContextConfig],
         sqlmesh_exporters: ResourceParam[t.List[SQLMeshExporter]],
     ) -> SQLMeshExportedAssetsCollection:
         environment = sqlmesh_infra_config["environment"]
@@ -49,6 +48,7 @@ def cacheable_exported_assets_defs(
             context_factory=DEFAULT_CONTEXT_FACTORY,
         )
         assets_map: dict[str, SQLMeshExportedAssetDefinition] = {}
+        sqlmesh_translator = sqlmesh_context_config.get_translator()
 
         with controller.instance(environment) as mesh:
             models = mesh.models()

warehouse/oso_dagster/assets/sqlmesh/sqlmesh_export.py

@@ -6,6 +6,7 @@
 from dagster_sqlmesh import SQLMeshResource
 from dagster_sqlmesh.controller.base import DEFAULT_CONTEXT_FACTORY
 from oso_dagster.factories import dlt_factory
+from oso_dagster.resources.sqlmesh import PrefixedSQLMeshContextConfig
 from sqlglot import exp
 
 
@@ -17,6 +18,7 @@
 def get_rendered_models(
     context: AssetExecutionContext,
     sqlmesh: SQLMeshResource,
+    sqlmesh_context_config: PrefixedSQLMeshContextConfig,
     environment: str,
 ):
     """
@@ -31,7 +33,9 @@ def get_rendered_models(
         Dict: Information about each rendered model
     """
     controller = sqlmesh.get_controller(
-        context_factory=DEFAULT_CONTEXT_FACTORY, log_override=context.log
+        config=sqlmesh_context_config,
+        context_factory=DEFAULT_CONTEXT_FACTORY,
+        log_override=context.log,
     )
 
     with controller.instance(environment, "model_renderer") as mesh:
@@ -71,6 +75,7 @@ def get_rendered_models(
 def sqlmesh_render_models(
     context: AssetExecutionContext,
     sqlmesh: SQLMeshResource,
+    sqlmesh_context_config: PrefixedSQLMeshContextConfig,
     sqlmesh_infra_config: dict,
 ):
     """
@@ -90,4 +95,5 @@ def sqlmesh_render_models(
         context=context,
         sqlmesh=sqlmesh,
         environment=environment,
+        sqlmesh_context_config=sqlmesh_context_config,
     )