Repository: TobikoData/sqlmesh
Pull request #2566 (2024-05-06 00:32:16)

chore: add documentation for new duckdb catalog options
Docs for this feature: https://github.com/TobikoData/sqlmesh/pull/2565


Patch:
diff --git a/docs/integrations/engines/duckdb.md b/docs/integrations/engines/duckdb.md
index 4aafb5ef..f3cdf011 100644
--- a/docs/integrations/engines/duckdb.md
+++ b/docs/integrations/engines/duckdb.md
@@ -5,15 +5,15 @@
 
 ### Connection options
 
-| Option     | Description                                                                                                                                                      |  Type  | Required |
-|--------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------|:------:|:--------:|
-| `type`             | Engine type name - must be `duckdb`                                                                                                                      | string |    Y     |
-| `database`         | The optional database name. If not specified, the in-memory database is used. Cannot be defined if using `catalogs`.                                     | string |    N     |
-| `catalogs`         | Mapping to define multiple catalogs. Name is the alias and value is the path. First entry is the default catalog. Cannot be defined if using `database`. |  dict  |    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     |
+| Option             | Description                                                                                                                                                                                                                                     |  Type  | Required |
+|--------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:------:|:--------:|
+| `type`             | Engine type name - must be `duckdb`                                                                                                                                                                                                             | string |    Y     |
+| `database`         | The optional database name. If not specified, the in-memory database is used. Cannot be defined if using `catalogs`.                                                                                                                            | string |    N     |
+| `catalogs`         | Mapping to define multiple catalogs. Can [attach DuckDB catalogs](#duckdb-catalogs-example) or [catalogs for other connections](#other-connection-catalogs-example). First entry is the default catalog. Cannot be defined if using `database`. |  dict  |    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     |
 
-#### Catalogs Example
+#### DuckDB Catalogs Example
 
 This example specifies two catalogs. The first catalog is named "persistent" and maps to the DuckDB file database `local.duckdb`. The second catalog is named "ephemeral" and maps to the DuckDB in-memory database.
 
@@ -58,6 +58,76 @@ SQLMesh will place models with the explicit catalog "ephemeral", such as `epheme
     )
     ```
 
+#### Other Connection Catalogs Example
+
+Catalogs can also be defined to connect to anything that [DuckDB can be attached to](https://duckdb.org/docs/sql/statements/attach.html).
+
+Below are examples of connecting to a SQLite database and a PostgreSQL database. 
+The SQLite database is read-write, while the PostgreSQL database is read-only.
+
+=== "YAML"
+
+    ```yaml linenums="1"
+    gateways:
+        my_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
+    ```
+
+=== "Python"
+
+    ```python linenums="1"
+    from sqlmesh.core.config import (
+        Config,
+        ModelDefaultsConfig,
+        GatewayConfig,
+        DuckDBConnectionConfig
+    )
+    from sqlmesh.core.config.connection import DuckDBAttachOptions
+
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect=<dialect>),
+        gateways={
+            "my_gateway": GatewayConfig(
+                connection=DuckDBConnectionConfig(
+                    catalogs={
+                        "memory": ":memory:",
+                        "sqlite": DuckDBAttachOptions(
+                            type="sqlite", 
+                            path="test.db"
+                        ),
+                        "postgres": DuckDBAttachOptions(
+                            type="postgres", 
+                            path="dbname=postgres user=postgres host=127.0.0.1", 
+                            read_only=True
+                        ),
+                    }
+                )
+            ),
+        }
+    )
+    ```
+
+##### Connectors without schemas
+
+Some connections, like SQLite, do not support schema names and therefore objects will be attached under the default schema name of `main`.
+
+Example: mounting a SQLite database with the name `sqlite` that has a table `example_table` will be accessible as `sqlite.main.example_table`.
+
+##### Sensitive fields in paths
+
+If a connector, like Postgres, requires sensitive information in the path, it might support defining environment variables instead. 
+[See DuckDB Documentation for more information](https://duckdb.org/docs/extensions/postgres#configuring-via-environment-variables).
+
 #### Cloud service authentication
 
 DuckDB can read data directly from cloud services via extensions (e.g., [httpfs](https://duckdb.org/docs/extensions/httpfs/s3api), [azure](https://duckdb.org/docs/extensions/azure)).
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",