Repository: TobikoData/sqlmesh
Pull request #2938 (2024-07-24 15:21:59)

Chore: add persistent duckdb gateway to sushi and wursthall examples
These changes allow wursthall to use a persistent duckdb database with the CLI `--gateway duckdb_persistent` option.

They also modify sushi so its persistent duckdb config is specified with the same `--gateway duckdb_persistent` option (rather than the current `--config local_config` option).


Patch:
diff --git a/examples/sushi/config.py b/examples/sushi/config.py
index afa8a83d..abd99d12 100644
--- a/examples/sushi/config.py
+++ b/examples/sushi/config.py
@@ -29,9 +29,17 @@ model_defaults = ModelDefaultsConfig(**defaults)
 model_defaults_iceberg = ModelDefaultsConfig(**defaults, storage_format="iceberg")
 
 
-# An in memory DuckDB config.
+# A DuckDB config, in-memory by default.
 config = Config(
-    default_connection=DuckDBConnectionConfig(),
+    gateways={
+        "duckdb": GatewayConfig(
+            connection=DuckDBConnectionConfig(),
+        ),
+        "duckdb_persistent": GatewayConfig(
+            connection=DuckDBConnectionConfig(database=f"{DATA_DIR}/duckdb.db"),
+        ),
+    },
+    default_gateway="duckdb",
     model_defaults=model_defaults,
 )
 
@@ -54,12 +62,6 @@ test_config = Config(
     model_defaults=model_defaults,
 )
 
-# A stateful DuckDB config.
-local_config = Config(
-    default_connection=DuckDBConnectionConfig(database=f"{DATA_DIR}/local.duckdb"),
-    model_defaults=model_defaults,
-)
-
 airflow_config = Config(
     default_scheduler=AirflowSchedulerConfig(),
     gateways=GatewayConfig(
diff --git a/examples/wursthall/config.yaml b/examples/wursthall/config.yaml
index 299a2e51..d4a21061 100644
--- a/examples/wursthall/config.yaml
+++ b/examples/wursthall/config.yaml
@@ -3,6 +3,10 @@ gateways:
   duckdb:
     connection:
       type: duckdb
+  duckdb_persistent:
+    connection:
+      type: duckdb
+      database: duckdb.db
 
 model_defaults:
   dialect: 'duckdb'
diff --git a/sqlmesh/core/config/__init__.py b/sqlmesh/core/config/__init__.py
index efdaa1c7..b5b66dde 100644
--- a/sqlmesh/core/config/__init__.py
+++ b/sqlmesh/core/config/__init__.py
@@ -4,6 +4,7 @@ from sqlmesh.core.config.categorizer import (
 )
 from sqlmesh.core.config.common import EnvironmentSuffixTarget as EnvironmentSuffixTarget
 from sqlmesh.core.config.connection import (
+    BaseDuckDBConnectionConfig as BaseDuckDBConnectionConfig,
     BigQueryConnectionConfig as BigQueryConnectionConfig,
     ConnectionConfig as ConnectionConfig,
     DatabricksConnectionConfig as DatabricksConnectionConfig,
@@ -16,6 +17,7 @@ from sqlmesh.core.config.connection import (
     RedshiftConnectionConfig as RedshiftConnectionConfig,
     SnowflakeConnectionConfig as SnowflakeConnectionConfig,
     SparkConnectionConfig as SparkConnectionConfig,
+    TrinoConnectionConfig as TrinoConnectionConfig,
     parse_connection_config as parse_connection_config,
 )
 from sqlmesh.core.config.gateway import GatewayConfig as GatewayConfig


Test patch:
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index 6444dc12..71537f99 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -47,7 +47,7 @@ def test_global_config(copy_to_temp_path: t.Callable):
 
 
 def test_named_config(copy_to_temp_path: t.Callable):
-    context = Context(paths=copy_to_temp_path("examples/sushi"), config="local_config")
+    context = Context(paths=copy_to_temp_path("examples/sushi"), config="test_config")
     assert len(context.config.gateways) == 1
 
 
@@ -85,7 +85,7 @@ def test_config_not_found(copy_to_temp_path: t.Callable):
         ConfigError,
         match=r".*config could not be found.*",
     ):
-        Context(paths="nonexistent/directory", config="local_config")
+        Context(paths="nonexistent/directory", config="config")
 
 
 def test_custom_macros(sushi_context):
@@ -594,14 +594,14 @@ def test_plan_default_end(sushi_context_pre_scheduling: Context):
 def test_plan_start_ahead_of_end(copy_to_temp_path):
     path = copy_to_temp_path("examples/sushi")
     with freezegun.freeze_time("2024-01-02 00:00:00"):
-        context = Context(paths=path, config="local_config")
+        context = Context(paths=path, gateway="duckdb_persistent")
         context.plan("prod", no_prompts=True, auto_apply=True)
         assert context.state_sync.max_interval_end_for_environment("prod") == to_timestamp(
             "2024-01-02"
         )
         context.close()
     with freezegun.freeze_time("2024-01-03 00:00:00"):
-        context = Context(paths=path, config="local_config")
+        context = Context(paths=path, gateway="duckdb_persistent")
         expression = d.parse(
             """
                 MODEL(