Repository: TobikoData/sqlmesh
Pull request #1962 (2024-01-12 00:47:09)

Fix: Include elk worker in output bundle


Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 38bcd866..cc2fe2c3 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -93,7 +93,8 @@ Usage: sqlmesh dag [OPTIONS] FILE
   Render the DAG as an html file.
 
 Options:
-  --help  Show this message and exit.
+  --help               Show this message and exit.
+  --select-model TEXT  Select specific models to include in the dag.
 ```
 
 ## diff
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 3964d015..d9fa5ceb 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -355,11 +355,17 @@ def invalidate(ctx: click.Context, environment: str) -> None:
 
 @cli.command("dag")
 @click.argument("file", required=True)
+@click.option(
+    "--select-model",
+    type=str,
+    multiple=True,
+    help="Select specific models to include in the dag.",
+)
 @click.pass_context
 @error_handler
-def dag(ctx: click.Context, file: str) -> None:
+def dag(ctx: click.Context, file: str, select_model: t.List[str]) -> None:
     """Render the DAG as an html file."""
-    rendered_dag_path = ctx.obj.render_dag(file)
+    rendered_dag_path = ctx.obj.render_dag(file, select_model)
     if rendered_dag_path:
         ctx.obj.console.log_success(f"Generated the dag to {rendered_dag_path}")
 
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index cd21fcf7..3dafe860 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -933,13 +933,7 @@ class Context(BaseContext):
         if include_unmodified is None:
             include_unmodified = self.config.include_unmodified
 
-        model_selector = Selector(
-            self.state_reader,
-            self._models,
-            context_path=self.path,
-            default_catalog=self.default_catalog,
-            dialect=self.config.dialect,
-        )
+        model_selector = self._new_selector()
 
         if backfill_models:
             backfill_models = model_selector.expand_model_selections(backfill_models)
@@ -1139,18 +1133,32 @@ class Context(BaseContext):
             self.console.show_row_diff(table_diff.row_diff(), show_sample=show_sample)
         return table_diff
 
-    def get_dag(self, **options: t.Any) -> GraphHTML:
-        """Gets an HTML object representation of the DAG."""
+    def get_dag(
+        self, select_models: t.Optional[t.Collection[str]] = None, **options: t.Any
+    ) -> GraphHTML:
+        """Gets an HTML object representation of the DAG.
+
+        Args:
+            select_models: A list of model selection strings that should be included in the dag.
+        Returns:
+            An html object that renders the dag.
+        """
+        dag = (
+            self.dag.prune(*self._new_selector().expand_model_selections(select_models))
+            if select_models
+            else self.dag
+        )
+
         nodes = {}
         edges: t.List[t.Dict] = []
 
-        for node, deps in self.dag.graph.items():
+        for node, deps in dag.graph.items():
             nodes[node] = {
                 "id": node,
                 "label": node.split(".")[-1],
                 "title": f"<span>{node}</span>",
             }
-            edges.extend({"from": node, "to": d} for d in deps)
+            edges.extend({"from": d, "to": node} for d in deps)
 
         return GraphHTML(
             nodes,
@@ -1173,15 +1181,16 @@ class Context(BaseContext):
             },
         )
 
-    def render_dag(self, path: str) -> None:
+    def render_dag(self, path: str, select_models: t.Optional[t.Collection[str]] = None) -> None:
         """Render the dag as HTML and save it to a file.
 
         Args:
             path: filename to save the dag html to
+            select_models: A list of model selection strings that should be included in the dag.
         """
 
         with open(path, "w", encoding="utf-8") as file:
-            file.write(str(self.get_dag()))
+            file.write(str(self.get_dag(select_models)))
 
     def create_test(
         self,
@@ -1644,6 +1653,15 @@ class Context(BaseContext):
     def _new_state_sync(self) -> StateSync:
         return self._provided_state_sync or self._scheduler.create_state_sync(self)
 
+    def _new_selector(self) -> Selector:
+        return Selector(
+            self.state_reader,
+            self._models,
+            context_path=self.path,
+            default_catalog=self.default_catalog,
+            dialect=self.config.dialect,
+        )
+
     def _register_notification_targets(self) -> None:
         event_notifications = collections.defaultdict(set)
         for target in self.notification_targets:
diff --git a/sqlmesh/core/plan/definition.py b/sqlmesh/core/plan/definition.py
index 66dc955c..24f61636 100644
--- a/sqlmesh/core/plan/definition.py
+++ b/sqlmesh/core/plan/definition.py
@@ -50,8 +50,6 @@ class Plan(PydanticModel, frozen=True):
     execution_time: t.Optional[TimeLike] = None
 
     # Make sure to use @cached_property instead once we stop supporting 3.7.
-    _categorized: t.Optional[t.List[Snapshot]] = None
-    _uncategorized: t.Optional[t.List[Snapshot]] = None
     _all_modified: t.Optional[t.Dict[SnapshotId, Snapshot]] = None
     _snapshots: t.Optional[t.Dict[SnapshotId, Snapshot]] = None
     _environment: t.Optional[Environment] = None
@@ -112,24 +110,20 @@ class Plan(PydanticModel, frozen=True):
     @property
     def categorized(self) -> t.List[Snapshot]:
         """Returns the already categorized snapshots."""
-        if self._categorized is None:
-            self._categorized = [
-                self.context_diff.snapshots[s_id]
-                for s_id in sorted(self.directly_modified)
-                if self.context_diff.snapshots[s_id].version
-            ]
-        return self._categorized
+        return [
+            self.context_diff.snapshots[s_id]
+            for s_id in sorted(self.directly_modified)
+            if self.context_diff.snapshots[s_id].version
+        ]
 
     @property
     def uncategorized(self) -> t.List[Snapshot]:
         """Returns the uncategorized snapshots."""
-        if self._uncategorized is None:
-            self._uncategorized = [
-                self.context_diff.snapshots[s_id]
-                for s_id in sorted(self.directly_modified)
-                if not self.context_diff.snapshots[s_id].version
-            ]
-        return self._uncategorized
+        return [
+            self.context_diff.snapshots[s_id]
+            for s_id in sorted(self.directly_modified)
+            if not self.context_diff.snapshots[s_id].version
+        ]
 
     @property
     def snapshots(self) -> t.Dict[SnapshotId, Snapshot]:
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index a0353f8c..a0a100ec 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -497,12 +497,18 @@ class SQLMeshMagics(Magics):
 
     @magic_arguments()
     @argument("--file", "-f", type=str, help="An optional file path to write the HTML output to.")
+    @argument(
+        "--select-model",
+        type=str,
+        nargs="*",
+        help="Select specific models to include in the dag.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def dag(self, context: Context, line: str) -> None:
         """Displays the HTML DAG."""
         args = parse_argstring(self.dag, line)
-        dag = context.get_dag()
+        dag = context.get_dag(args.select_model)
         if args.file:
             with open(args.file, "w") as file:
                 file.write(str(dag))
diff --git a/sqlmesh/utils/dag.py b/sqlmesh/utils/dag.py
index b8f6a57a..3c04ee81 100644
--- a/sqlmesh/utils/dag.py
+++ b/sqlmesh/utils/dag.py
@@ -56,7 +56,7 @@ class DAG(t.Generic[T]):
             nodes: The nodes of the new subdag.
 
         Returns:
-            A new dag consisting of the specified nodes.
+            A new dag consisting of the specified nodes and upstream.
         """
         queue = set(nodes)
         graph = {}
@@ -69,6 +69,23 @@ class DAG(t.Generic[T]):
 
         return DAG(graph)
 
+    def prune(self, *nodes: T) -> DAG[T]:
+        """Create a dag keeping only the included nodes.
+
+        Args:
+            nodes: The nodes of the new pruned dag.
+
+        Returns:
+            A new dag consisting of the specified nodes.
+        """
+        graph = {}
+
+        for node, deps in self._dag.items():
+            if node in nodes:
+                graph[node] = {dep for dep in deps if dep in nodes}
+
+        return DAG(graph)
+
     def upstream(self, node: T) -> t.List[T]:
         """Returns all upstream dependencies in topologically sorted order."""
         return self.subdag(node).sorted[:-1]
diff --git a/web/client/src/library/components/graph/help.ts b/web/client/src/library/components/graph/help.ts
index 87d231fa..7b52f86e 100644
--- a/web/client/src/library/components/graph/help.ts
+++ b/web/client/src/library/components/graph/help.ts
@@ -52,11 +52,14 @@ function createGraphLayout({
 } {
   // https://eclipse.dev/elk/reference/options.html
   const elk: any = new ELK({
-    workerUrl: '/node_modules/elkjs/lib/elk-worker.min.js',
+    workerUrl: new URL(
+      '/node_modules/elkjs/lib/elk-worker.min.js?worker',
+      import.meta.url,
+    ).href,
   })
 
   return {
-    terminate: () => elk.worker.worker.terminate(),
+    terminate: () => elk.worker.terminate(),
     create: async () =>
       new Promise((resolve, reject) => {
         elk
diff --git a/web/client/src/workers/index.ts b/web/client/src/workers/index.ts
index 7a5b2748..a2eafdfe 100644
--- a/web/client/src/workers/index.ts
+++ b/web/client/src/workers/index.ts
@@ -1,9 +1,9 @@
 const sqlglotWorker = new Worker(
-  new URL('./sqlglot/sqlglot.ts', import.meta.url),
+  new URL('./sqlglot/sqlglot.ts?worker', import.meta.url),
 )
 
 function createLineageWorker(): Worker {
-  return new Worker(new URL('./lineage.ts', import.meta.url), {
+  return new Worker(new URL('./lineage.ts?worker', import.meta.url), {
     type: 'module',
   })
 }


Test patch:
diff --git a/tests/utils/test_dag.py b/tests/utils/test_dag.py
index cb46f6bb..58b53c4c 100644
--- a/tests/utils/test_dag.py
+++ b/tests/utils/test_dag.py
@@ -92,3 +92,10 @@ def test_reversed_graph():
         "c": set(),
         "d": set(),
     }
+
+
+def test_prune():
+    assert DAG({"a": {"b", "d"}, "b": {"d", "e"}}).prune("a", "d").graph == {
+        "a": {"d"},
+        "d": set(),
+    }