fix: share duckdb adapter is overlapping data files (#2039)

Author: eakmanrq

sqlmesh/core/config/connection.py

@@ -2,6 +2,7 @@
 
 import abc
 import base64
+import logging
 import os
 import pathlib
 import sys
@@ -33,6 +34,9 @@
     from typing_extensions import Literal
 
 
+logger = logging.getLogger(__name__)
+
+
 class ConnectionConfig(abc.ABC, BaseConfig):
     concurrent_tasks: int
     register_comments: bool
@@ -122,6 +126,8 @@ class DuckDBConnectionConfig(ConnectionConfig):
 
     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(
@@ -182,6 +188,29 @@ def init(cursor: duckdb.DuckDBPyConnection) -> None:
 
         return init
 
+    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
+        associated with the new adapter will be ignored."""
+        data_files = set((self.catalogs or {}).values())
+        if self.database:
+            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}"
+                )
+                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
+        return adapter
+
     def get_catalog(self) -> t.Optional[str]:
         if self.database:
             # Remove `:` from the database name in order to handle if `:memory:` is passed in

tests/conftest.py

@@ -16,6 +16,7 @@
 from sqlglot import exp, maybe_parse, parse_one
 from sqlglot.helper import ensure_list
 
+from sqlmesh.core.config import DuckDBConnectionConfig
 from sqlmesh.core.context import Context
 from sqlmesh.core.engine_adapter.base import EngineAdapter
 from sqlmesh.core.macros import macro
@@ -176,6 +177,12 @@ def rescope_global_models(request):
     model.set_registry(existing_registry)
 
 
+@pytest.fixture(scope="function", autouse=True)
+def rescope_duckdb_classvar(request):
+    DuckDBConnectionConfig._data_file_to_adapter = {}
+    yield
+
+
 @pytest.fixture
 def duck_conn() -> duckdb.DuckDBPyConnection:
     return duckdb.connect()

tests/core/test_connection_config.py

@@ -366,3 +366,145 @@ def test_duckdb(make_config):
     )
     assert isinstance(config, DuckDBConnectionConfig)
     assert config._static_connection_kwargs == {"config": {"foo": "bar"}}
+
+
+@pytest.mark.parametrize(
+    "kwargs1, kwargs2, shared_adapter",
+    [
+        (
+            {
+                "database": "test.duckdb",
+            },
+            {
+                "database": "test.duckdb",
+            },
+            True,
+        ),
+        (
+            {},
+            {},
+            False,
+        ),
+        (
+            {
+                "database": "test1.duckdb",
+            },
+            {
+                "database": "test2.duckdb",
+            },
+            False,
+        ),
+        (
+            {
+                "database": ":memory:",
+            },
+            {
+                "database": ":memory:",
+            },
+            False,
+        ),
+        (
+            {
+                "database": ":memory:",
+            },
+            {
+                "database": "test.duckdb",
+            },
+            False,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test": "test.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test": "test.duckdb",
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test": ":memory:",
+                }
+            },
+            {
+                "catalogs": {
+                    "test": ":memory:",
+                }
+            },
+            False,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test1": ":memory:",
+                    "test2": "test2.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test1": ":memory:",
+                    "test2": "test2.duckdb",
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test1": ":memory:",
+                    "test2": "test2.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test1": "test2.duckdb",
+                    "test2": ":memory:",
+                }
+            },
+            True,
+        ),
+        (
+            {
+                "catalogs": {
+                    "test1": "test1.duckdb",
+                    "test2": "test2.duckdb",
+                    "test3": "test3.duckdb",
+                }
+            },
+            {
+                "catalogs": {
+                    "test1": "test1_miss.duckdb",
+                    "test2": "test2_miss.duckdb",
+                    "test3": "test3.duckdb",
+                }
+            },
+            True,
+        ),
+    ],
+)
+def test_duckdb_shared(make_config, caplog, kwargs1, kwargs2, shared_adapter):
+    config1 = make_config(
+        type="duckdb",
+        **kwargs1,
+    )
+    config2 = make_config(
+        type="duckdb",
+        **kwargs2,
+    )
+    assert isinstance(config1, DuckDBConnectionConfig)
+    assert isinstance(config2, DuckDBConnectionConfig)
+    adapter1 = config1.create_engine_adapter()
+    adapter2 = config2.create_engine_adapter()
+    if shared_adapter:
+        assert id(adapter1) == id(adapter2)
+        assert "Creating new DuckDB adapter" in caplog.messages[0]
+        assert "Using existing DuckDB adapter" in caplog.messages[1]
+    else:
+        assert id(adapter1) != id(adapter2)
+        assert "Creating new DuckDB adapter" in caplog.messages[0]
+        assert "Creating new DuckDB adapter" in caplog.messages[1]