Repository: TobikoData/sqlmesh
Pull request #2565 (2024-05-05 17:49:03)

 Feat: Full support for Duckdb attach options
This PR adds additional config surface to allow for DuckDB's attach options to be leveraged, similar to duckdb-dbt.

Ref: #2336

cc @eakmanrq 

CLAassistant: [![CLA assistant check](https://cla-assistant.io/pull/badge/not_signed)](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2565) <br/>Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our [Contributor License Agreement](https://cla-assistant.io/TobikoData/sqlmesh?pullRequest=2565) before we can accept your contribution.<br/><sub>You have signed the CLA already but the status is still pending? Let us [recheck](https://cla-assistant.io/check/TobikoData/sqlmesh?pullRequest=2565) it.</sub>
wseaton: There is some formatting churn here caused by my editor setup that shouldn't be happening anymore. Will fixup, but wanted to note that the duckdb integration tests are passing w/ the following patch:

```patch
diff --git a/tests/core/engine_adapter/config.yaml b/tests/core/engine_adapter/config.yaml
index 66be470e..ca07c66a 100644
--- a/tests/core/engine_adapter/config.yaml
+++ b/tests/core/engine_adapter/config.yaml
@@ -4,7 +4,10 @@ gateways:
       type: duckdb
       catalogs:
         memory: ':memory:'
-        testing: 'testing.duckdb'
+        testing:
+          type: duckdb
+          path: 'testing.duckdb'
+          read_only: false
   inttest_trino:
     connection:
       type: trino
diff --git a/tests/core/engine_adapter/test_integration.py b/tests/core/engine_adapter/test_integration.py
index 6bbc746a..64e60d10 100644
--- a/tests/core/engine_adapter/test_integration.py
+++ b/tests/core/engine_adapter/test_integration.py
@@ -526,8 +526,9 @@ def engine_adapter(mark_gateway: t.Tuple[str, str], config) -> EngineAdapter:
         engine_adapter.DEFAULT_BATCH_SIZE = 1
     # Clear our any local db files that may have been left over from previous runs
     if mark == "duckdb":
-        for raw_path in (connection_config.catalogs or {}).values():
-            pathlib.Path(raw_path).unlink(missing_ok=True)
+        for options in (connection_config.catalogs or {}).values():
+            path = options if isinstance(options, str) else options.path
+            pathlib.Path(path).unlink(missing_ok=True)
     return engine_adapter
```

Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index aeeaf275..3f56a63b 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -165,21 +165,28 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
                 except Exception as e:
                     raise ConfigError(f"Failed to set connector config {field} to {setting}: {e}")
 
-            for i, (alias, path) in enumerate((getattr(self, "catalogs", None) or {}).items()):
+            for i, (alias, path_options) in enumerate(
+                (getattr(self, "catalogs", None) or {}).items()
+            ):
                 # we parse_identifier and generate to ensure that `alias` has exactly one set of quotes
                 # regardless of whether it comes in quoted or not
                 alias = exp.parse_identifier(alias, dialect="duckdb").sql(
                     identify=True, dialect="duckdb"
                 )
                 try:
-                    cursor.execute(f"""ATTACH '{path}' AS {alias}""")
+                    query = (
+                        path_options.to_sql(alias)
+                        if isinstance(path_options, DuckDBAttachOptions)
+                        else f"ATTACH '{path_options}' AS {alias}"
+                    )
+                    cursor.execute(query)
                 except BinderException as e:
                     # If a user tries to create a catalog pointing at `:memory:` and with the name `memory`
                     # then we don't want to raise since this happens by default. They are just doing this to
                     # set it as the default catalog.
                     if not (
                         'database with name "memory" already exists' in str(e)
-                        and path == ":memory:"
+                        and path_options == ":memory:"
                     ):
                         raise e
                 if i == 0 and not getattr(self, "database", None):
@@ -214,6 +221,23 @@ class MotherDuckConnectionConfig(BaseDuckDBConnectionConfig):
         return {"database": connection_str}
 
 
+class DuckDBAttachOptions(BaseConfig):
+    type: str
+    path: str
+    read_only: bool = False
+
+    def to_sql(self, alias: str) -> str:
+        options = []
+        # 'duckdb' is actually not a supported type, but we'd like to allow it for
+        # fully qualified attach options or integration testing, similar to duckdb-dbt
+        if self.type != "duckdb":
+            options.append(f"TYPE {self.type.upper()}")
+        if self.read_only:
+            options.append("READ_ONLY")
+        options_sql = f" ({', '.join(options)})" if options else ""
+        return f"ATTACH '{self.path}' AS {alias}{options_sql}"
+
+
 class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
     """Configuration for the DuckDB connection.
 
@@ -223,7 +247,7 @@ class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
     """
 
     database: t.Optional[str] = None
-    catalogs: t.Optional[t.Dict[str, str]] = None
+    catalogs: t.Optional[t.Dict[str, t.Union[str, DuckDBAttachOptions]]] = None
 
     type_: Literal["duckdb"] = Field(alias="type", default="duckdb")
 
@@ -253,18 +277,19 @@ class DuckDBConnectionConfig(BaseDuckDBConnectionConfig):
             data_files.add(self.database)
         data_files.discard(":memory:")
         for data_file in data_files:
-            if adapter := DuckDBConnectionConfig._data_file_to_adapter.get(data_file):
-                logger.info(
-                    f"Using existing DuckDB adapter due to overlapping data file: {data_file}"
-                )
+            key = data_file if isinstance(data_file, str) else data_file.path
+            if adapter := DuckDBConnectionConfig._data_file_to_adapter.get(key):
+                logger.info(f"Using existing DuckDB adapter due to overlapping data file: {key}")
                 return adapter
+
         if data_files:
             logger.info(f"Creating new DuckDB adapter for data files: {data_files}")
         else:
             logger.info("Creating new DuckDB adapter for in-memory database")
         adapter = super().create_engine_adapter(register_comments_override)
         for data_file in data_files:
-            DuckDBConnectionConfig._data_file_to_adapter[data_file] = adapter
+            key = data_file if isinstance(data_file, str) else data_file.path
+            DuckDBConnectionConfig._data_file_to_adapter[key] = adapter
         return adapter
 
     def get_catalog(self) -> t.Optional[str]:


Test patch:
diff --git a/tests/core/test_config.py b/tests/core/test_config.py
index e308d1b6..ffd13162 100644
--- a/tests/core/test_config.py
+++ b/tests/core/test_config.py
@@ -13,6 +13,7 @@ from sqlmesh.core.config import (
     GatewayConfig,
     ModelDefaultsConfig,
 )
+from sqlmesh.core.config.connection import DuckDBAttachOptions
 from sqlmesh.core.config.feature_flag import DbtFeatureFlag, FeatureFlag
 from sqlmesh.core.config.loader import (
     load_config_from_env,
@@ -486,3 +487,48 @@ def test_variables():
         ConfigError, match="Unsupported variable value type: <class 'sqlglot.expressions.Column'>"
     ):
         Config(variables={"invalid_var": exp.column("sqlglot_expr")})
+
+
+def test_load_duckdb_attach_config(tmp_path):
+    config_path = tmp_path / "config_duckdb_attach.yaml"
+    with open(config_path, "w", encoding="utf-8") as fd:
+        fd.write(
+            """
+gateways:
+    another_gateway:
+        connection:
+            type: duckdb
+            catalogs:
+              memory: ':memory:'
+              sqlite:
+                type: 'sqlite'
+                path: 'test.db'
+              postgres:
+                type: 'postgres'
+                path: 'dbname=postgres user=postgres host=127.0.0.1'
+                read_only: true
+model_defaults:
+    dialect: ''
+        """
+        )
+
+    config = load_config_from_paths(
+        Config,
+        project_paths=[config_path],
+    )
+
+    assert config.gateways["another_gateway"].connection.catalogs.get("memory") == ":memory:"
+
+    attach_config_1 = config.gateways["another_gateway"].connection.catalogs.get("sqlite")
+
+    assert isinstance(attach_config_1, DuckDBAttachOptions)
+    assert attach_config_1.type == "sqlite"
+    assert attach_config_1.path == "test.db"
+    assert attach_config_1.read_only is False
+
+    attach_config_2 = config.gateways["another_gateway"].connection.catalogs.get("postgres")
+
+    assert isinstance(attach_config_2, DuckDBAttachOptions)
+    assert attach_config_2.type == "postgres"
+    assert attach_config_2.path == "dbname=postgres user=postgres host=127.0.0.1"
+    assert attach_config_2.read_only is True
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 8ecb90ca..45db2689 100644
--- a/tests/core/test_connection_config.py
+++ b/tests/core/test_connection_config.py
@@ -7,6 +7,7 @@ from _pytest.fixtures import FixtureRequest
 from sqlmesh.core.config.connection import (
     BigQueryConnectionConfig,
     ConnectionConfig,
+    DuckDBAttachOptions,
     DuckDBConnectionConfig,
     GCPPostgresConnectionConfig,
     MySQLConnectionConfig,
@@ -525,6 +526,41 @@ def test_duckdb_shared(make_config, caplog, kwargs1, kwargs2, shared_adapter):
         assert "Creating new DuckDB adapter" in caplog.messages[1]
 
 
+def test_duckdb_attach_catalog(make_config):
+    config = make_config(
+        type="duckdb",
+        catalogs={
+            "test1": "test1.duckdb",
+            "test2": DuckDBAttachOptions(
+                type="duckdb",
+                path="test2.duckdb",
+            ),
+        },
+    )
+    assert isinstance(config, DuckDBConnectionConfig)
+    assert config.get_catalog() == "test1"
+
+    assert config.catalogs.get("test2").read_only is False
+    assert config.catalogs.get("test2").path == "test2.duckdb"
+    assert config.is_recommended_for_state_sync is True
+
+
+def test_duckdb_attach_options():
+
+    options = DuckDBAttachOptions(
+        type="postgres", path="dbname=postgres user=postgres host=127.0.0.1", read_only=True
+    )
+
+    assert (
+        options.to_sql(alias="db")
+        == "ATTACH 'dbname=postgres user=postgres host=127.0.0.1' AS db (TYPE POSTGRES, READ_ONLY)"
+    )
+
+    options = DuckDBAttachOptions(type="duckdb", path="test.db", read_only=False)
+
+    assert options.to_sql(alias="db") == "ATTACH 'test.db' AS db"
+
+
 def test_bigquery(make_config):
     config = make_config(
         type="bigquery",