Fix jupyter test

Author: treysp

test.py

@@ -0,0 +1,210 @@
+from unittest import skip
+from sqlmesh.utils import enable_debug_mode
+from sqlmesh.cli.example_project import init_example_project
+from sqlmesh.core.config.loader import load_config_from_paths
+from sqlmesh.core.context import Context
+from sqlglot import parse_one, transpile, exp
+from pathlib import Path
+from sqlmesh.core.config import Config, DuckDBConnectionConfig, ModelDefaultsConfig
+import duckdb
+from sqlmesh.dbt.context import DbtContext
+
+#load_config_from_paths(Path("./sqlmesh-example/config.yaml"))
+
+#### sqlmesh example
+
+project_dir = "sqlmesh-example-fresh"
+gateway = "inttest_databricks"
+state_file = "duckdb.db"
+
+import os
+import shutil
+#os.environ["SQLMESH_DEBUG"] = "1"
+
+os.chdir(f"/Users/trey/tobiko/sqlmesh/{project_dir}")
+#os.chdir("/users/trey/desktop/temp")
+
+if os.path.exists(".cache"):
+    shutil.rmtree(".cache")
+
+if os.path.exists(state_file):
+    os.remove(state_file)
+if os.path.exists(state_file + ".wal"):
+    os.remove(state_file + ".wal")
+
+temp = Context(paths=f"/Users/trey/tobiko/sqlmesh/{project_dir}")
+#temp.print_info()
+#temp.run(ignore_cron=True)
+#temp.test()
+temp.plan(no_prompts=True, auto_apply=True, skip_tests=True)
+
+# #### sushi dbt
+
+# import os
+# os.chdir("./examples/sushi_dbt")
+# #enable_debug_mode()
+# temp = Context(paths="/users/trey/tobiko/sqlmesh/examples/sushi_dbt")
+
+#### sushi
+
+# import os
+# os.chdir("./examples/sushi")
+# #enable_debug_mode()
+# temp = Context(paths="/users/trey/tobiko/sqlmesh/examples/sushi", gateway="inttest_snowflake")
+
+# temp.plan(no_prompts=True, auto_apply=True, skip_tests=True)
+
+#case_2 = duckdb.connect("/users/trey/tobiko/sqlmesh/duckdb_attach/case_2.duckdb")
+
+# temp.render("sqlmesh_example.full_model")
+
+# temp.plan(include_unmodified=True, auto_apply=True, no_prompts=True, skip_tests=True)
+
+# #temp.evaluate("sqlmesh_example.select_external", start = "2020-01-01", end="today", execution_time="now")
+
+# #temp.table_diff(source="prod", target="dev", model_or_snapshot="sqlmesh_example.incremental_model")
+
+# #from sqlglot import exp
+
+# #exp.func("IF", exp.and_(*((exp.column(c, "s").eq(exp.column(c, "t")), exp.column(c, "s").not_().is_(exp.Null()), exp.column(c, "t").not_().is_(exp.Null())) for c in list(set([a.name for a in s_index]))),), 1, 0).sql()
+
+# #exp.func("IF", exp.and_(*(exp.column(c, "s").eq(exp.column(c, "t")), exp.column(c, "s").not_().is_(exp.Null()), exp.column(c, "t").not_().is_(exp.Null()) for c in list(set([a.name for a in s_index]))),), 1, 0).sql()
+# #exp.func("IF", exp.and_(*(exp.and_(exp.column(c, "s").eq(exp.column(c, "t")), exp.column(c, "s").not_().is_(exp.Null()), exp.column(c, "t").not_().is_(exp.Null())) for c in list(set([a.name for a in s_index]))),), 1, 0).sql()
+
+# #-------------
+
+# ################
+# #  MS SQL
+# ################
+
+# import pymssql
+# import pandas as pd
+
+# conn = pymssql.connect("localhost", "SA", "1StrongPwd@@", as_dict=True)
+
+# cur = conn.cursor()
+
+# # tables
+# cur.execute("select * from information_schema.tables where table_schema like '%sqlmesh%'")
+# tables = cur.fetchall()
+# to_delete = [f"{tbl['TABLE_SCHEMA']}.{tbl['TABLE_NAME']}" for tbl in tables]
+
+# for tbl in to_delete:
+#     cur.execute(f"drop table {tbl};")
+
+# # views
+# cur.execute("select * from information_schema.views where table_schema like '%sqlmesh%'")
+# tables = cur.fetchall()
+# to_delete = [f"{tbl['TABLE_SCHEMA']}.{tbl['TABLE_NAME']}" for tbl in tables]
+
+# for tbl in to_delete:
+#     cur.execute(f"drop view {tbl};")
+
+# # schemas
+# cur.execute("select * from information_schema.schemata where schema_name like '%sqlmesh%'")
+# schemas = cur.fetchall()
+# to_delete = [schema['SCHEMA_NAME'] for schema in schemas]
+
+# for schema in to_delete:
+#     cur.execute(f"drop schema {schema};")
+
+# cur.execute("select * from information_schema.schemata where schema_name like '%sqlmesh'")
+# pd.DataFrame(cur.fetchall())
+
+# cur.execute("select * from information_schema.tables where table_schema like '%sqlmesh%'")
+# pd.DataFrame(cur.fetchall())
+
+# cur.close()
+# conn.close()
+# exit()
+
+# #######################
+# #   METRICS
+# #######################
+
+# from sqlmesh.core import Context
+
+# context = Context(paths="/users/trey/tobiko/sqlmesh/examples/sushi")
+
+# query = context.rewrite(
+#     """
+#     SELECT
+#         ds,
+#         METRIC(total_orders_from_active_customers),
+#     FROM
+#         __semantic.__table  -- special table for simple metric queries
+#     GROUP BY
+#         ds
+#     """,
+# )
+
+# from time import sleep
+# from rich.console import Console
+# from rich.live import Live
+# from rich.table import Table
+# from rich.progress import (
+#     Progress,
+#     TextColumn,
+#     BarColumn,
+#     SpinnerColumn,
+#     TimeElapsedColumn,
+# )
+
+# PROGRESS_BAR_WIDTH = 40
+
+# def make_overall_progress(console: Console) -> Progress:
+#     return Progress(
+#         TextColumn("[bold blue]Overall:", justify="right"),
+#         BarColumn(bar_width=PROGRESS_BAR_WIDTH),
+#         "[progress.percentage]{task.percentage:>3.1f}%",
+#         TimeElapsedColumn(),
+#         console=console,
+#         transient=False,
+#     )
+
+# def make_model_progress(console: Console) -> Progress:
+#     return Progress(
+#         TextColumn("{task.fields[model_name]}", justify="right"),
+#         SpinnerColumn(spinner_name="dots"),
+#         console=console,
+#         transient=False,
+#     )
+
+# def demo_table_stateless(models: dict[str, int], delay: float = 0.5):
+#     console = Console()
+#     total_batches = sum(models.values())
+#     completed_total = 0
+#     completed_msgs: list[str] = []  # store completed model messages
+
+#     with Live(console=console, refresh_per_second=10) as live:
+#         for model_name, batch_count in models.items():
+#             for batch_index in range(1, batch_count + 1):
+#                 # compute new totals
+#                 completed_total += 1
+
+#                 # recreate Progress widgets with absolute values
+#                 overall = make_overall_progress(console)
+#                 overall.add_task("", total=total_batches, completed=completed_total)
+
+#                 model = make_model_progress(console)
+#                 task_id = model.add_task(
+#                     "", total=batch_count, completed=batch_index, model_name=model_name
+#                 )
+#                 # remove the per‑model bar once it's done
+#                 if batch_index >= batch_count:
+#                     model.remove_task(task_id)
+#                     completed_msgs.append(model_name)
+
+#                 # assemble grid and re‑render
+#                 grid = Table.grid()
+#                 for msg in completed_msgs:
+#                     grid.add_row(msg)
+#                 grid.add_row(overall)
+#                 grid.add_row(model)
+#                 live.update(grid)
+
+#                 sleep(delay)
+
+# if __name__ == "__main__":
+#     # simulate three “models” with 5, 3, and 4 batches
+#     demo_table_stateless({"Model A": 5, "Model B": 3, "Model C": 4})

tests/integrations/jupyter/test_magics.py

@@ -289,9 +289,6 @@ def test_plan(
     with capture_output() as output:
         notebook.run_line_magic(magic_name="plan", line="--no-prompts --auto-apply")
 
-    # TODO: Should this be going to stdout? This is printing the status updates for when each batch finishes for
-    # the models and how long it took
-    assert len(output.stdout.strip().split("\n")) == 46
     assert not output.stderr
     assert len(output.outputs) == 4
     text_output = convert_all_html_output_to_text(output)
@@ -324,8 +321,15 @@ def test_run_dag(
         "'Evaluating models ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 100.0% • 18/18"
     )
     assert not output.stderr
-    assert len(output.outputs) == 6
+    assert len(output.outputs) == 10
     assert convert_all_html_output_to_text(output) == [
+        "[2K",
+        "Executing model batches ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╸━━ 93.8% • "
+        "15/16 • 0:00:01\n"
+        "sushi.waiter_as_customer_by_day",
+        "[1/1] sushi.waiter_as_customer_by_day   [insert 2023-01-01 - 2023-01-02, "
+        "audits ✔2]   0.15s",
+        "",
         "✔ Model batches executed",
         "Run finished for environment 'prod'",
     ]