Repository: TobikoData/sqlmesh
Pull request #1814 (2023-12-07 07:46:51)

Toby/test results


Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 1409dae3..ab3f66eb 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -84,16 +84,12 @@ Options:
 
 ## dag
 ```
-Usage: sqlmesh dag [OPTIONS]
+Usage: sqlmesh dag [OPTIONS] FILE
 
-  Renders the dag using graphviz.
-
-  This command requires a manual install of both the python and system
-  graphviz package.
+  Renders the dag as an html file.
 
 Options:
-  --file TEXT  The file to which the dag image should be written.
-  --help       Show this message and exit.
+  --help  Show this message and exit.
 ```
 
 ## diff
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index cb7de270..720b5220 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -349,14 +349,13 @@ def invalidate(ctx: click.Context, environment: str) -> None:
 
 
 @cli.command("dag")
-@opt.file
+@click.argument("file", required=True)
 @click.pass_context
 @error_handler
 def dag(ctx: click.Context, file: str) -> None:
     """
-    Renders the dag using graphviz.
+    Renders the dag as an html file.
 
-    This command requires a manual install of both the python and system graphviz package.
     """
     rendered_dag_path = ctx.obj.render_dag(file)
     if rendered_dag_path:
diff --git a/sqlmesh/cli/options.py b/sqlmesh/cli/options.py
index a9c0cf64..869cd46e 100644
--- a/sqlmesh/cli/options.py
+++ b/sqlmesh/cli/options.py
@@ -42,11 +42,6 @@ expand = click.option(
     help="Whether or not to expand materialized models (defaults to False). If True, all referenced models are expanded as raw queries. Multiple model names can also be specified, in which case only they will be expanded as raw queries.",
 )
 
-file = click.option(
-    "--file",
-    help="The file to which the dag image should be written.",
-)
-
 match_pattern = click.option(
     "-k",
     multiple=True,
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 9e47346e..63c74e73 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -47,6 +47,7 @@ from types import MappingProxyType
 
 import pandas as pd
 from sqlglot import exp
+from sqlglot.lineage import GraphHTML
 
 from sqlmesh.core import constants as c
 from sqlmesh.core.audit import Audit, StandaloneAudit
@@ -102,7 +103,6 @@ from sqlmesh.utils.date import TimeLike, now_ds, to_date
 from sqlmesh.utils.errors import (
     CircuitBreakerError,
     ConfigError,
-    MissingDependencyError,
     PlanError,
     SQLMeshError,
     UncategorizedPlanError,
@@ -110,7 +110,6 @@ from sqlmesh.utils.errors import (
 from sqlmesh.utils.jinja import JinjaMacroRegistry
 
 if t.TYPE_CHECKING:
-    import graphviz
     from typing_extensions import Literal
 
     from sqlmesh.core.engine_adapter._typing import DF, PySparkDataFrame, PySparkSession
@@ -974,60 +973,49 @@ class Context(BaseContext):
             self.console.show_row_diff(table_diff.row_diff(), show_sample=show_sample)
         return table_diff
 
-    def get_dag(self, format: str = "svg") -> graphviz.Digraph:
-        """Gets a graphviz dag.
+    def get_dag(self, **options: t.Any) -> GraphHTML:
+        """Gets an HTML object representation of the DAG."""
+        nodes = {}
+        edges: t.List[t.Dict] = []
 
-        This method requires installing the graphviz base library through your package manager
-        and the python graphviz library.
-
-        To display within Databricks:
-        displayHTML(context.get_dag().pipe(encoding='utf-8'))
-
-        Args:
-            format: The desired format to use for representing the graph
-        """
-        from sqlmesh import runtime_env
-
-        try:
-            import graphviz  # type: ignore
-        except ModuleNotFoundError as e:
-            if runtime_env.is_databricks:
-                raise MissingDependencyError(
-                    "Rendering a dag requires graphviz. Run `pip install graphviz` and then `sudo apt-get install -y python3-dev graphviz libgraphviz-dev pkg-config`"
-                )
-            raise MissingDependencyError(
-                "Rendering a dag requires a manual install of graphviz. Run `pip install graphviz` and then install graphviz library: https://graphviz.org/download/."
-            ) from e
-
-        graph = graphviz.Digraph(node_attr={"shape": "box"}, format=format)
-
-        for name, upstream in self.dag.graph.items():
-            graph.node(name)
-            for u in upstream:
-                graph.edge(u, name)
-        return graph
-
-    def render_dag(self, path: str, format: str = "jpeg") -> str:
-        """Render the dag using graphviz.
+        for node, deps in self.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)
+
+        return GraphHTML(
+            nodes,
+            edges,
+            options={
+                "height": "100%",
+                "width": "100%",
+                "interaction": {},
+                "layout": {
+                    "hierarchical": {
+                        "enabled": True,
+                        "nodeSpacing": 200,
+                        "sortMethod": "directed",
+                    },
+                },
+                "nodes": {
+                    "shape": "box",
+                },
+                **options,
+            },
+        )
 
-        This method requires installing the graphviz base library through your package manager
-        and the python graphviz library.
+    def render_dag(self, path: str) -> None:
+        """Render the dag as HTML and save it to a file.
 
         Args:
-            path: filename to save the dag to
-            format: The desired format to use when rending the dag
+            path: filename to save the dag html to
         """
-        graph = self.get_dag(format=format)
-        # We know graphviz is installed because the command above would have failed if it was not. This allows
-        # us to then catch the specific error that occurs when the system install is missing.
-        import graphviz  # type: ignore
 
-        try:
-            return graph.render(path, format=format)
-        except graphviz.backend.execute.ExecutableNotFound as e:
-            raise MissingDependencyError(
-                "Graphviz is pip-installed but the system install is missing. Instructions: https://graphviz.org/download/"
-            ) from e
+        with open(path, "w", encoding="utf-8") as file:
+            file.write(str(self.get_dag()))
 
     def create_test(
         self,
@@ -1348,7 +1336,10 @@ class Context(BaseContext):
     ) -> t.Tuple[t.Optional[unittest.result.TestResult], t.Optional[str]]:
         if self._test_engine_adapter and not skip_tests:
             result, test_output = self._run_tests()
-            self.console.log_test_results(result, test_output, self._test_engine_adapter.dialect)
+            if result.testsRun > 0:
+                self.console.log_test_results(
+                    result, test_output, self._test_engine_adapter.dialect
+                )
             if not result.wasSuccessful():
                 raise PlanError(
                     "Cannot generate plan due to failing test(s). Fix test(s) and run again"
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index 599a6f8d..931ba069 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -486,29 +486,21 @@ def format_model_expressions(
         return expressions[0].sql(pretty=True, dialect=dialect)
 
     *statements, query = expressions
-    query = query.copy()
-    selects = []
-
-    for expression in query.expressions:
-        column = None
-        comments = expression.comments
-        expression.comments = None
-
-        if not isinstance(expression, exp.Alias) and expression.output_name not in ("", "*"):
-            expression = expression.replace(exp.alias_(expression, expression.output_name))
 
-        column = column or expression
-        expression = expression.this
+    def cast_to_colon(node: exp.Expression) -> exp.Expression:
+        if isinstance(node, exp.Cast) and not node.args.get("format"):
+            this = node.this
 
-        if isinstance(expression, exp.Cast):
-            this = expression.this
             if not isinstance(this, (exp.Binary, exp.Unary)) or isinstance(this, exp.Paren):
-                expression.replace(DColonCast(this=this, to=expression.to))
+                cast = DColonCast(this=this, to=node.to)
+                cast.comments = node.comments
+                node = cast
 
-        column.comments = comments
-        selects.append(column)
+        exp.replace_children(node, cast_to_colon)
+        return node
 
-    query.set("expressions", selects)
+    query = query.copy()
+    exp.replace_children(query, cast_to_colon)
 
     return ";\n\n".join(
         [
diff --git a/sqlmesh/core/engine_adapter/base_postgres.py b/sqlmesh/core/engine_adapter/base_postgres.py
index d319c8e2..c18ecd16 100644
--- a/sqlmesh/core/engine_adapter/base_postgres.py
+++ b/sqlmesh/core/engine_adapter/base_postgres.py
@@ -15,7 +15,7 @@ if t.TYPE_CHECKING:
 
 
 class BasePostgresEngineAdapter(EngineAdapter):
-    DEFAULT_BATCH_SIZE = 1000
+    DEFAULT_BATCH_SIZE = 400
     COLUMNS_TABLE = "information_schema.columns"
     CATALOG_SUPPORT = CatalogSupport.SINGLE_CATALOG_ONLY
 
diff --git a/sqlmesh/integrations/github/cicd/controller.py b/sqlmesh/integrations/github/cicd/controller.py
index 4855404e..940bdb37 100644
--- a/sqlmesh/integrations/github/cicd/controller.py
+++ b/sqlmesh/integrations/github/cicd/controller.py
@@ -480,8 +480,8 @@ class GithubController:
         query = f"""{{
             repository(owner: "{self._event.pull_request_info.owner}", name: "{self._event.pull_request_info.repo}") {{
                 pullRequest(number: {self._event.pull_request_info.pr_number}) {{
-                    title 
-                    state 
+                    title
+                    state
                     mergeStateStatus
                 }}
             }}
@@ -654,6 +654,7 @@ class GithubController:
         ) -> t.Tuple[GithubCheckConclusion, str, t.Optional[str]]:
             if not result:
                 return GithubCheckConclusion.SKIPPED, "Skipped Tests", None
+
             # Clear out console
             self._console.consume_captured_output()
             self._console.log_test_results(
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index 03fde312..a667e116 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -457,8 +457,7 @@ class SQLMeshMagics(Magics):
     def dag(self, line: str) -> None:
         """Displays the dag"""
         self._context.refresh()
-        dag = self._context.get_dag()
-        self.display(HTML(dag.pipe().decode("utf-8")))
+        self.display(self._context.get_dag())
 
     @magic_arguments()
     @line_magic


Test patch:
diff --git a/tests/core/test_dialect.py b/tests/core/test_dialect.py
index d8a76a88..acbc79d7 100644
--- a/tests/core/test_dialect.py
+++ b/tests/core/test_dialect.py
@@ -46,6 +46,7 @@ def test_format_model_expressions():
     sum(n + 1)::int as n, -- n
     o,
     p + 1,
+    CAST(x as int)::int,
     """
         )
     )
@@ -61,26 +62,27 @@ def test_format_model_expressions():
 
 SELECT
   *,
-  a::INT AS a,
+  a::INT,
   b::INT AS b,
   CAST(c + 1 AS INT) AS c,
-  d::INT AS d,
+  d::INT,
   e::INT AS e,
   (
     f + 1
   )::INT AS f,
   SUM(g + 1)::INT AS g,
-  h::INT AS h, /* h */
+  h::INT, /* h */
   i::INT AS i, /* i */
   CAST(j + 1 AS INT) AS j, /* j */
-  k::INT AS k, /* k */
+  k::INT, /* k */
   l::INT AS l, /* l */
   (
     m + 1
   )::INT AS m, /* m */
   SUM(n + 1)::INT AS n, /* n */
-  o AS o,
-  p + 1"""
+  o,
+  p + 1,
+  x::INT::INT"""
     )
 
     x = format_model_expressions(
@@ -130,7 +132,7 @@ def test_format_body_macros():
     1
 )
 SELECT
-  col::INT AS col
+  col::INT
 FROM foo
 @ORDER_BY(@include_order_by)
   @EACH(@columns, item -> @'@iteaoeuatnoehutoenahuoanteuhonateuhaoenthuaoentuhaeotnhaoem'),