Repository: TobikoData/sqlmesh
Pull request #2067 (2024-02-01 00:13:09)

Feat: Add support for MotherDuck

izeigerman: @eakmanrq I'm not quite sure how multi-catalog works with MotherDuck, so I'd rather defer implementation until there's actual demand for it.

Patch:
diff --git a/docs/guides/configuration.md b/docs/guides/configuration.md
index df25fb8f..b816d3f6 100644
--- a/docs/guides/configuration.md
+++ b/docs/guides/configuration.md
@@ -309,6 +309,8 @@ If no match is found, the catalog defined in the model or the default catalog de
 * [BigQuery](../integrations/engines/bigquery.md)
 * [Databricks](../integrations/engines/databricks.md)
 * [DuckDB](../integrations/engines/duckdb.md)
+* [MotherDuck](../integrations/engines/motherduck.md)
+* [MySQL](../integrations/engines/mysql.md)
 * [MSSQL](../integrations/engines/mssql.md)
 * [Snowflake](../integrations/engines/snowflake.md)
 * [Spark](../integrations/engines/spark.md)
@@ -473,6 +475,7 @@ These pages describe the connection configuration options for each execution eng
 * [BigQuery](../integrations/engines/bigquery.md)
 * [Databricks](../integrations/engines/databricks.md)
 * [DuckDB](../integrations/engines/duckdb.md)
+* [MotherDuck](../integrations/engines/motherduck.md)
 * [MySQL](../integrations/engines/mysql.md)
 * [MSSQL](../integrations/engines/mssql.md)
 * [Postgres](../integrations/engines/postgres.md)
@@ -950,4 +953,4 @@ Example enabling debug mode for the CLI command `sqlmesh plan`:
     ```cmd
     C:\> set SQLMESH_DEBUG=1
     C:\> sqlmesh plan
-    ```
\ No newline at end of file
+    ```
diff --git a/docs/guides/connections.md b/docs/guides/connections.md
index 2fe626a0..764bc7fd 100644
--- a/docs/guides/connections.md
+++ b/docs/guides/connections.md
@@ -84,6 +84,7 @@ default_gateway: local_db
 * [BigQuery](../integrations/engines/bigquery.md)
 * [Databricks](../integrations/engines/databricks.md)
 * [DuckDB](../integrations/engines/duckdb.md)
+* [MotherDuck](../integrations/engines/motherduck.md)
 * [MySQL](../integrations/engines/mysql.md)
 * [MSSQL](../integrations/engines/mssql.md)
 * [Postgres](../integrations/engines/postgres.md)
diff --git a/docs/integrations/engines/motherduck.md b/docs/integrations/engines/motherduck.md
new file mode 100644
index 00000000..04759cc8
--- /dev/null
+++ b/docs/integrations/engines/motherduck.md
@@ -0,0 +1,14 @@
+# MotherDuck
+
+## Local/Built-in Scheduler
+**Engine Adapter Type**: `motherduck`
+
+### Connection options
+
+| Option             | Description                                                                                                 | Type   | Required |
+|--------------------|-------------------------------------------------------------------------------------------------------------|:------:|:--------:|
+| `type`             | Engine type name - must be `motherduck`                                                                     | string | Y        |
+| `database`         | The database name.                                                                                          | string | Y        |
+| `token`            | The optional MotherDuck token. If not specified, the user will be prompted to login with their web browser. | string | N        |
+| `extensions`       | Extension to load into duckdb. Only autoloadable extensions are supported.                                  | list   | N        |
+| `connector_config` | Configuration to pass into the duckdb connector.                                                            | dict   | N        |
diff --git a/docs/integrations/overview.md b/docs/integrations/overview.md
index bd0e91a0..7200fc6c 100644
--- a/docs/integrations/overview.md
+++ b/docs/integrations/overview.md
@@ -14,6 +14,7 @@ SQLMesh supports the following execution engines for running SQLMesh projects:
 * [BigQuery](./engines/bigquery.md)
 * [Databricks](./engines/databricks.md)
 * [DuckDB](./engines/duckdb.md)
+* [MotherDuck](./engines/motherduck.md)
 * [MySQL](./engines/mysql.md)
 * [MSSQL](./engines/mssql.md)
 * [Postgres](./engines/postgres.md)
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index f23e11b4..40c2f64a 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -123,6 +123,7 @@ These pages describe the connection configuration options for each execution eng
 * [BigQuery](../integrations/engines/bigquery.md)
 * [Databricks](../integrations/engines/databricks.md)
 * [DuckDB](../integrations/engines/duckdb.md)
+* [MotherDuck](../integrations/engines/motherduck.md)
 * [MySQL](../integrations/engines/mysql.md)
 * [MSSQL](../integrations/engines/mssql.md)
 * [Postgres](../integrations/engines/postgres.md)
diff --git a/sqlmesh/core/config/__init__.py b/sqlmesh/core/config/__init__.py
index e8427d3a..d7b85a77 100644
--- a/sqlmesh/core/config/__init__.py
+++ b/sqlmesh/core/config/__init__.py
@@ -6,6 +6,7 @@ from sqlmesh.core.config.connection import (
     DatabricksConnectionConfig,
     DuckDBConnectionConfig,
     GCPPostgresConnectionConfig,
+    MotherDuckConnectionConfig,
     MSSQLConnectionConfig,
     MySQLConnectionConfig,
     PostgresConnectionConfig,
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index b66558f7..e804f2a8 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -104,49 +104,26 @@ class ConnectionConfig(abc.ABC, BaseConfig):
         return None
 
 
-class DuckDBConnectionConfig(ConnectionConfig):
-    """Configuration for the DuckDB connection.
+class BaseDuckDBConnectionConfig(ConnectionConfig):
+    """Common configuration for the DuckDB-based connections.
 
     Args:
-        database: The optional database name. If not specified, the in-memory database will be used.
-        catalogs: Key is the name of the catalog and value is the path.
         extensions: A list of autoloadable extensions to load.
         connector_config: A dictionary of configuration to pass into the duckdb connector.
         concurrent_tasks: The maximum number of tasks that can use this connection concurrently.
         register_comments: Whether or not to register model comments with the SQL engine.
     """
 
-    database: t.Optional[str] = None
-    catalogs: t.Optional[t.Dict[str, str]] = None
     extensions: t.List[str] = []
     connector_config: t.Dict[str, t.Any] = {}
 
     concurrent_tasks: Literal[1] = 1
     register_comments: bool = True
 
-    type_: Literal["duckdb"] = Field(alias="type", default="duckdb")
-
-    _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
-
-    @model_validator(mode="before")
-    @model_validator_v1_args
-    def _validate_database_catalogs(
-        cls, values: t.Dict[str, t.Optional[str]]
-    ) -> t.Dict[str, t.Optional[str]]:
-        if values.get("database") and values.get("catalogs"):
-            raise ConfigError(
-                "Cannot specify both `database` and `catalogs`. Define all your catalogs in `catalogs` and have the first entry be the default catalog"
-            )
-        return values
-
     @property
     def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
         return {"config": self.connector_config} if self.connector_config else {}
 
-    @property
-    def _connection_kwargs_keys(self) -> t.Set[str]:
-        return {"database"}
-
     @property
     def _engine_adapter(self) -> t.Type[EngineAdapter]:
         return engine_adapter.DuckDBEngineAdapter
@@ -171,7 +148,7 @@ class DuckDBConnectionConfig(ConnectionConfig):
                 except Exception as e:
                     raise ConfigError(f"Failed to load extension {extension}: {e}")
 
-            for i, (alias, path) in enumerate((self.catalogs or {}).items()):
+            for i, (alias, path) in enumerate((getattr(self, "catalogs", None) or {}).items()):
                 try:
                     cursor.execute(f'''ATTACH '{path}' AS "{alias}"''')
                 except BinderException as e:
@@ -183,11 +160,68 @@ class DuckDBConnectionConfig(ConnectionConfig):
                         and path == ":memory:"
                     ):
                         raise e
-                if i == 0 and not self.database:
+                if i == 0 and not getattr(self, "database", None):
                     cursor.execute(f'USE "{alias}"')
 
         return init
 
+
+class MotherDuckConnectionConfig(BaseDuckDBConnectionConfig):
+    """Configuration for the MotherDuck connection.
+
+    Args:
+        database: The database name.
+        token: The optional MotherDuck token. If not specified, the user will be prompted to login with their web browser.
+    """
+
+    database: str
+    token: t.Optional[str] = None
+
+    type_: Literal["motherduck"] = Field(alias="type", default="motherduck")
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        return set()
+
+    @property
+    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
+        """kwargs that are for execution config only"""
+        connection_str = f"md:{self.database}"
+        if self.token:
+            connection_str += f"?motherduck_token={self.token}"
+        return {"database": connection_str}
+
+
+class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
+    """Configuration for the DuckDB connection.
+
+    Args:
+        database: The optional database name. If not specified, the in-memory database will be used.
+        catalogs: Key is the name of the catalog and value is the path.
+    """
+
+    database: t.Optional[str] = None
+    catalogs: t.Optional[t.Dict[str, str]] = None
+
+    type_: Literal["duckdb"] = Field(alias="type", default="duckdb")
+
+    _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
+
+    @model_validator(mode="before")
+    @model_validator_v1_args
+    def _validate_database_catalogs(
+        cls, values: t.Dict[str, t.Optional[str]]
+    ) -> t.Dict[str, t.Optional[str]]:
+        if values.get("database") and values.get("catalogs"):
+            raise ConfigError(
+                "Cannot specify both `database` and `catalogs`. Define all your catalogs in `catalogs` and have the first entry be the default catalog"
+            )
+        return values
+
+    @property
+    def _connection_kwargs_keys(self) -> t.Set[str]:
+        return {"database"}
+
     def create_engine_adapter(self, register_comments_override: bool = False) -> EngineAdapter:
         """Checks if another engine adapter has already been created that shares a catalog that points to the same data
         file. If so, it uses that same adapter instead of creating a new one. As a result, any additional configuration
@@ -1198,7 +1232,9 @@ class TrinoConnectionConfig(ConnectionConfig):
 CONNECTION_CONFIG_TO_TYPE = {
     # Map all subclasses of ConnectionConfig to the value of their `type_` field.
     tpe.all_field_infos()["type_"].default: tpe
-    for tpe in subclasses(__name__, ConnectionConfig, exclude=(ConnectionConfig,))
+    for tpe in subclasses(
+        __name__, ConnectionConfig, exclude=(ConnectionConfig, BaseDuckDBConnectionConfig)
+    )
 }
 
 
diff --git a/sqlmesh/core/plan/builder.py b/sqlmesh/core/plan/builder.py
index 595d5653..ef4a9481 100644
--- a/sqlmesh/core/plan/builder.py
+++ b/sqlmesh/core/plan/builder.py
@@ -338,7 +338,11 @@ class PlanBuilder:
         return {
             self._context_diff.snapshots[s_id].name
             for s_id in dag.subdag(
-                *[self._model_fqn_to_snapshot[m].snapshot_id for m in self._backfill_models]
+                *[
+                    self._model_fqn_to_snapshot[m].snapshot_id
+                    for m in self._backfill_models
+                    if m in self._model_fqn_to_snapshot
+                ]
             ).sorted
         }
 
diff --git a/sqlmesh/core/selector.py b/sqlmesh/core/selector.py
index f816146b..1ddbd9a8 100644
--- a/sqlmesh/core/selector.py
+++ b/sqlmesh/core/selector.py
@@ -134,11 +134,16 @@ class Selector:
         return value, include_upstream, include_downstream
 
     def _get_models(
-        self, model_name: str, include_upstream: bool, include_downstream: bool
+        self,
+        model_name: str,
+        include_upstream: bool,
+        include_downstream: bool,
+        models: t.Optional[t.Dict[str, Model]] = None,
     ) -> t.Set[str]:
         result = {model_name}
         if include_upstream:
-            result.update(self._dag.upstream(model_name))
+            models = models or self._models
+            result.update([u for u in self._dag.upstream(model_name) if u in models])
         if include_downstream:
             result.update(self._dag.downstream(model_name))
         return result
@@ -222,6 +227,8 @@ class Selector:
                 logger.warning(f"Expression '{selection}' doesn't match any models.")
 
             for model_fqn in matched_models:
-                results.update(self._get_models(model_fqn, include_upstream, include_downstream))
+                results.update(
+                    self._get_models(model_fqn, include_upstream, include_downstream, models=models)
+                )
 
         return results


Test patch:
diff --git a/pytest.ini b/pytest.ini
index 56845e1b..d02eb892 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -21,6 +21,7 @@ markers =
     bigquery: test for BigQuery
     databricks: test for Databricks
     duckdb: test for DuckDB
+    motherduck: test for MotherDuck
     mssql: test for MSSQL
     mysql: test for MySQL
     postgres: test for Postgres
diff --git a/tests/core/engine_adapter/test_integration.py b/tests/core/engine_adapter/test_integration.py
index f55ac11d..c5f0a9b6 100644
--- a/tests/core/engine_adapter/test_integration.py
+++ b/tests/core/engine_adapter/test_integration.py
@@ -455,6 +455,14 @@ def config() -> Config:
                 pytest.mark.remote,
             ],
         ),
+        pytest.param(
+            "motherduck",
+            marks=[
+                pytest.mark.motherduck,
+                pytest.mark.engine,
+                pytest.mark.remote,
+            ],
+        ),
         pytest.param(
             "redshift",
             marks=[
@@ -526,6 +534,12 @@ def test_catalog_operations(ctx: TestContext):
         ctx.engine_adapter.cursor.connection.autocommit(False)
     elif ctx.dialect == "snowflake":
         ctx.engine_adapter.execute(f'CREATE DATABASE IF NOT EXISTS "{catalog_name}"')
+    elif ctx.dialect == "duckdb":
+        try:
+            # Only applies to MotherDuck
+            ctx.engine_adapter.execute(f'CREATE DATABASE IF NOT EXISTS "{catalog_name}"')
+        except Exception:
+            pass
     current_catalog = ctx.engine_adapter.get_current_catalog()
     ctx.engine_adapter.set_current_catalog(catalog_name)
     assert ctx.engine_adapter.get_current_catalog() == catalog_name
@@ -588,6 +602,12 @@ def test_drop_schema_catalog(ctx: TestContext):
         ctx.engine_adapter.cursor.connection.autocommit(False)
     elif ctx.dialect == "snowflake":
         ctx.engine_adapter.execute(f'CREATE DATABASE IF NOT EXISTS "{catalog_name}"')
+    elif ctx.dialect == "duckdb":
+        try:
+            # Only applies to MotherDuck
+            ctx.engine_adapter.execute(f'CREATE DATABASE IF NOT EXISTS "{catalog_name}"')
+        except Exception:
+            pass
     elif ctx.dialect == "bigquery":
         catalog_name = "tobiko-test"
 
diff --git a/tests/core/test_selector.py b/tests/core/test_selector.py
index d769b695..032fcca6 100644
--- a/tests/core/test_selector.py
+++ b/tests/core/test_selector.py
@@ -349,6 +349,36 @@ def test_expand_model_tags(mocker: MockerFixture, make_snapshot, tags, tag_selec
     assert selector.expand_model_selections(tag_selections) == output
 
 
+def test_select_models_with_external_parent(mocker: MockerFixture):
+    default_catalog = "test_catalog"
+    added_model = SqlModel(
+        name="db.added_model",
+        query=d.parse_one("SELECT 1 AS a FROM external"),
+        default_catalog=default_catalog,
+        tags=["tag1"],
+    )
+
+    env_name = "test_env"
+
+    state_reader_mock = mocker.Mock()
+    state_reader_mock.get_environment.return_value = Environment(
+        name=env_name,
+        snapshots=[],
+        start_at="2023-01-01",
+        end_at="2023-02-01",
+        plan_id="test_plan_id",
+    )
+    state_reader_mock.get_snapshots.return_value = {}
+
+    local_models: UniqueKeyDict[str, Model] = UniqueKeyDict("models")
+    local_models[added_model.fqn] = added_model
+
+    selector = Selector(state_reader_mock, local_models, default_catalog=default_catalog)
+
+    expanded_selections = selector.expand_model_selections(["+*added_model*"])
+    assert expanded_selections == {added_model.fqn}
+
+
 def _assert_models_equal(actual: t.Dict[str, Model], expected: t.Dict[str, Model]) -> None:
     assert set(actual) == set(expected)
     for name, model in actual.items():