Improve sqlmesh as lib tooling (#4)

* Upgrade to latest sqlmesh

* WIP

* Refactor sqlmesh controller

* bump

* remove signals

* fix

* update

* remove type keyword

* Remove unused context.py

Author: ravenac95

dagster_sqlmesh/asset.py

@@ -1,115 +1,31 @@
-from dataclasses import dataclass, field
-from typing import (
-    Union,
-    Iterable,
-    Dict,
-    Mapping,
-    Set,
-    Any,
-    Optional,
-)
+import typing as t
 import logging
 
-import sqlglot
-from sqlglot import exp
-from sqlmesh.core.context import Context
-from sqlmesh.core.console import Console
-from sqlmesh.core.model import Model
 from dagster import (
-    AssetDep,
     multi_asset,
-    AssetCheckResult,
-    AssetMaterialization,
-    AssetOut,
-    AssetKey,
     RetryPolicy,
 )
-from dagster._core.definitions.asset_dep import CoercibleToAssetDep
+
+from dagster_sqlmesh.controller.dagster import DagsterSQLMeshController
+from dagster_sqlmesh.translator import SQLMeshDagsterTranslator
 
 from .config import SQLMeshContextConfig
-from .console import EventConsole, ConsoleEventHandler, DebugEventConsole
-from .utils import sqlmesh_model_name_to_key
-from .context import DagsterSQLMeshContext
 
 logger = logging.getLogger(__name__)
 
 
-MultiAssetResponse = Iterable[Union[AssetCheckResult, AssetMaterialization]]
-
-
-@dataclass(kw_only=True)
-class SQLMeshParsedFQN:
-    catalog: str
-    schema: str
-    view_name: str
-
-
-def parse_fqn(fqn: str):
-    split_fqn = fqn.split(".")
-
-    # Remove any quotes
-    split_fqn = list(map(lambda a: a.strip("'\""), split_fqn))
-    return SQLMeshParsedFQN(
-        catalog=split_fqn[0], schema=split_fqn[1], view_name=split_fqn[2]
-    )
-
-
-@dataclass(kw_only=True)
-class SQLMeshModelDep:
-    fqn: str
-    model: Optional[Model] = None
-
-    def parse_fqn(self):
-        return parse_fqn(self.fqn)
-
-
-@dataclass(kw_only=True)
-class SQLMeshMultiAssetOptions:
-    outs: Dict[str, AssetOut] = field(default_factory=lambda: {})
-    deps: Iterable[CoercibleToAssetDep] = field(default_factory=lambda: {})
-    internal_asset_deps: Dict[str, Set[AssetKey]] = field(default_factory=lambda: {})
-
-
-class SQLMeshDagsterTranslator:
-    def get_asset_key_from_model(self, context: Context, model: Model) -> AssetKey:
-        return AssetKey(model.view_name)
-
-    def get_asset_key_fqn(self, context: Context, fqn: str) -> AssetKey:
-        table = self.get_fqn_to_table(context, fqn)
-        return AssetKey(table.name)
-
-    def get_fqn_to_table(self, context: Context, fqn: str) -> exp.Table:
-        dialect = self.get_context_dialect(context)
-        return sqlglot.to_table(fqn, dialect=dialect)
-
-    def get_context_dialect(self, context: Context) -> str:
-        return context.engine_adapter.dialect
-
-    # def get_asset_deps(
-    #     self, context: Context, model: Model, deps: List[SQLMeshModelDep]
-    # ) -> List[AssetKey]:
-    #     asset_keys: List[AssetKey] = []
-    #     for dep in deps:
-    #         if dep.model:
-    #             asset_keys.append(AssetKey(dep.model.view_name))
-    #         else:
-    #             parsed_fqn = dep.parse_fqn()
-    #             asset_keys.append(AssetKey([parsed_fqn.view_name]))
-    #     return asset_keys
-
-
 # Define a SQLMesh Asset
 def sqlmesh_assets(
     *,
     config: SQLMeshContextConfig,
-    name: Optional[str] = None,
-    dagster_sqlmesh_translator: Optional[SQLMeshDagsterTranslator] = None,
+    name: t.Optional[str] = None,
+    dagster_sqlmesh_translator: t.Optional[SQLMeshDagsterTranslator] = None,
     compute_kind: str = "sqlmesh",
-    op_tags: Optional[Mapping[str, Any]] = None,
-    required_resource_keys: Optional[Set[str]] = None,
-    retry_policy: Optional[RetryPolicy] = None,
+    op_tags: t.Optional[t.Mapping[str, t.Any]] = None,
+    required_resource_keys: t.Optional[t.Set[str]] = None,
+    retry_policy: t.Optional[RetryPolicy] = None,
 ):
-    controller = setup_sqlmesh_controller(config)
+    controller = DagsterSQLMeshController.setup(config)
     if not dagster_sqlmesh_translator:
         dagster_sqlmesh_translator = SQLMeshDagsterTranslator()
     conversion = controller.to_asset_outs(dagster_sqlmesh_translator)
@@ -123,79 +39,4 @@ def sqlmesh_assets(
         compute_kind=compute_kind,
         retry_policy=retry_policy,
         required_resource_keys=required_resource_keys,
-        # can_subset=True,
-    )
-
-
-@dataclass
-class SQLMeshController:
-    console: EventConsole
-    context: DagsterSQLMeshContext
-
-    def add_event_handler(self, handler: ConsoleEventHandler):
-        return self.console.add_handler(handler)
-
-    def remove_event_handler(self, handler_id: str):
-        return self.console.remove_handler(handler_id)
-
-    def to_asset_outs(
-        self, translator: SQLMeshDagsterTranslator
-    ) -> SQLMeshMultiAssetOptions:
-        context = self.context
-        dag = context.dag
-        output = SQLMeshMultiAssetOptions()
-        depsMap: Dict[str, CoercibleToAssetDep] = {}
-
-        for model_fqn, deps in dag.graph.items():
-            logger.debug(f"model found: {model_fqn}")
-            model = context.get_model(model_fqn)
-            if not model:
-                # If no model is returned this seems to be an asset dependency
-                continue
-            asset_out = translator.get_asset_key_from_model(
-                context,
-                model,
-            )
-            model_deps = [
-                SQLMeshModelDep(fqn=dep, model=context.get_model(dep)) for dep in deps
-            ]
-            internal_asset_deps: Set[AssetKey] = set()
-            for dep in model_deps:
-                if dep.model:
-                    internal_asset_deps.add(
-                        translator.get_asset_key_from_model(context, dep.model)
-                    )
-                else:
-                    table = translator.get_fqn_to_table(context, dep.fqn)
-                    key = translator.get_asset_key_fqn(context, dep.fqn)
-                    internal_asset_deps.add(key)
-                    # create an external dep
-                    depsMap[table.name] = AssetDep(key)
-            model_key = sqlmesh_model_name_to_key(model.name)
-            output.outs[model_key] = AssetOut(key=asset_out, is_required=False)
-            output.internal_asset_deps[model_key] = internal_asset_deps
-
-        output.deps = list(depsMap.values())
-        return output
-
-
-def setup_sqlmesh_controller(
-    config: SQLMeshContextConfig,
-    debug_console: Optional[Console] = None,
-    log_override: Optional[logging.Logger] = None,
-):
-    console = EventConsole(log_override=log_override)
-    if debug_console:
-        console = DebugEventConsole(debug_console)
-    options: Dict[str, Any] = dict(
-        paths=config.path,
-        gateway=config.gateway,
-        console=console,
-    )
-    if config.sqlmesh_config:
-        options["config"] = config.sqlmesh_config
-    context = DagsterSQLMeshContext(**options)
-    return SQLMeshController(
-        console=console,
-        context=context,
     )

dagster_sqlmesh/conftest.py

@@ -4,13 +4,12 @@
 import shutil
 import os
 from dataclasses import dataclass
-from typing import cast, List, Optional, Any, Dict
+import typing as t
 
 import pytest
 import duckdb
 import polars
 from sqlmesh.utils.date import TimeLike
-from sqlmesh.core.plan.builder import PlanBuilder
 from sqlmesh.core.console import get_console
 from sqlmesh.core.config import (
     Config as SQLMeshConfig,
@@ -20,8 +19,9 @@
 )
 
 from dagster_sqlmesh.config import SQLMeshContextConfig
-from dagster_sqlmesh.events import ConsoleRecorder, show_plan_summary
-from dagster_sqlmesh.asset import setup_sqlmesh_controller
+from dagster_sqlmesh.events import ConsoleRecorder
+from dagster_sqlmesh.controller.base import PlanOptions, RunOptions
+from dagster_sqlmesh.controller.dagster import DagsterSQLMeshController
 
 logger = logging.getLogger(__name__)
 
@@ -58,7 +58,9 @@ def create_controller(self, enable_debug_console: bool = False):
         console = None
         if enable_debug_console:
             console = get_console()
-        return setup_sqlmesh_controller(self.context_config, debug_console=console)
+        return DagsterSQLMeshController.setup(
+            self.context_config, debug_console=console
+        )
 
     def query(self, *args, **kwargs):
         conn = duckdb.connect(self.db_path)
@@ -99,21 +101,19 @@ def run(
         *,
         environment: str,
         apply: bool = False,
-        execution_time: Optional[TimeLike] = None,
+        execution_time: t.Optional[TimeLike] = None,
         enable_debug_console: bool = False,
-        start: Optional[TimeLike] = None,
-        end: Optional[TimeLike] = None,
-        restate_models: Optional[List[str]] = None,
+        start: t.Optional[TimeLike] = None,
+        end: t.Optional[TimeLike] = None,
+        restate_models: t.Optional[t.List[str]] = None,
     ):
         controller = self.create_controller(enable_debug_console=enable_debug_console)
-        controller.add_event_handler(ConsoleRecorder())
-        plan_options: Dict[str, Any] = dict(
-            environment=environment,
+        recorder = ConsoleRecorder()
+        # controller.add_event_handler(ConsoleRecorder())
+        plan_options = PlanOptions(
             enable_preview=True,
         )
-        run_options: Dict[str, Any] = dict(
-            environment=environment,
-        )
+        run_options = RunOptions()
         if execution_time:
             plan_options["execution_time"] = execution_time
             run_options["execution_time"] = execution_time
@@ -126,19 +126,12 @@ def run(
             plan_options["end"] = end
             run_options["end"] = end
 
-        builder = cast(
-            PlanBuilder,
-            controller.context.plan_builder(**plan_options),
-        )
-        if apply:
-            logger.debug("making plan")
-            plan = builder.build()
-            show_plan_summary(logger, plan, lambda x: x.is_model)
-            logger.debug("applying plan")
-            controller.context.apply(plan)
-            logger.debug("running through the scheduler")
-            controller.context.run(**run_options)
-        controller.context.close()
+        for _context, event in controller.plan_and_run(
+            environment,
+            plan_options=plan_options,
+            run_options=run_options,
+        ):
+            recorder(event)
 
 
 @pytest.fixture