Repository: TobikoData/sqlmesh
Pull request #2039 (2024-01-29 03:19:01)

fix: share duckdb adapter is overlapping data files
Part of fix for: https://github.com/TobikoData/sqlmesh/issues/2037

Result of this change: https://github.com/TobikoData/sqlmesh/pull/1975

Prior to this change we would allow multiple adapters to share a common data file for DuckDB. This actually works fine until multiple adapters are created again pointed to the already created file and then the second adapter (typically state sync) will not read any data. This is why testing before of having multiple adapter run a plan and check the result did not show any issues.

The solution is to check when creating a new adapter for DuckDB if the new adapter is going to point to file that overlaps a file already used by another adapter. If so, we just reuse the past adapter. This has a downside that any additional configuration associated with the new adapter is ignored. One could consider merging the previous and new catalogs together if this overlap is detected but then it doesn't address other config. For example, what if the new one has other connection config set? How do we merge those? I'm wondering if this naive approach will work in most cases to where we can handle the more complex situations later. For example we will eventually want to support Motherduck and that could be a good time to rework things with a more robust use-case in mind.


Patch:
diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py
index 73e5f26e..b66558f7 100644
--- a/sqlmesh/core/config/connection.py
+++ b/sqlmesh/core/config/connection.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import abc
 import base64
+import logging
 import os
 import pathlib
 import sys
@@ -33,6 +34,9 @@ else:
     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 @@ class DuckDBConnectionConfig(ConnectionConfig):
 
         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


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index 499fadb0..99dfa9e5 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -16,6 +16,7 @@ from pytest_mock.plugin import MockerFixture
 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()
diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py
index 64937fb1..702be84e 100644
--- a/tests/core/test_connection_config.py
+++ b/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]