fix: Set DatabricksConnectionConfig to only use a shared_connection with U2M OAuth (#5859)

Signed-off-by: davem-bis <68955845+davem-bis@users.noreply.github.com>

Author: davem-bis

sqlmesh/core/config/connection.py

@@ -13,8 +13,8 @@
 from sys import version_info
 
 import pydantic
+from pydantic import Field, computed_field
 from packaging import version
-from pydantic import Field
 from pydantic_core import from_json
 from sqlglot import exp
 from sqlglot.errors import ParseError
@@ -110,7 +110,14 @@ class ConnectionConfig(abc.ABC, BaseConfig):
     catalog_type_overrides: t.Optional[t.Dict[str, str]] = None
 
     # Whether to share a  single connection across threads or create a new connection per thread.
-    shared_connection: t.ClassVar[bool] = False
+    #
+    # MyPy throws a "Decorators on top of @property are not supported" error despite this being a
+    # valid decoration, and Pydantic recommend disabling the MyPy hint for this reason - see:
+    # https://pydantic.dev/docs/validation/2.0/usage/computed_fields/
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def shared_connection(self) -> bool:
+        return False
 
     @property
     @abc.abstractmethod
@@ -311,7 +318,10 @@ class BaseDuckDBConnectionConfig(ConnectionConfig):
 
     token: t.Optional[str] = None
 
-    shared_connection: t.ClassVar[bool] = True
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def shared_connection(self) -> bool:
+        return True
 
     _data_file_to_adapter: t.ClassVar[t.Dict[str, EngineAdapter]] = {}
 
@@ -820,11 +830,15 @@ class DatabricksConnectionConfig(ConnectionConfig):
     DISPLAY_NAME: t.ClassVar[t.Literal["Databricks"]] = "Databricks"
     DISPLAY_ORDER: t.ClassVar[t.Literal[3]] = 3
 
-    shared_connection: t.ClassVar[bool] = True
-
     _concurrent_tasks_validator = concurrent_tasks_validator
     _http_headers_validator = http_headers_validator
 
+    @computed_field  # type: ignore[prop-decorator]
+    @property
+    def shared_connection(self) -> bool:
+        """The connection should only be shared if U2M OAuth is being used"""
+        return self.auth_type is not None and self.oauth_client_secret is None
+
     @model_validator(mode="before")
     def _databricks_connect_validator(cls, data: t.Any) -> t.Any:
         # SQLQueryContextLogger will output any error SQL queries even if they are in a try/except block.

sqlmesh/core/config/loader.py

@@ -272,4 +272,4 @@ def convert_config_type(
     config_obj: Config,
     config_type: t.Type[C],
 ) -> C:
-    return config_type.parse_obj(config_obj.dict())
+    return config_type.parse_obj(config_obj.dict(exclude_computed_fields=True))

tests/core/test_config.py

@@ -532,6 +532,7 @@ def test_connection_config_serialization():
         "extensions": [],
         "pre_ping": False,
         "pretty_sql": False,
+        "shared_connection": True,
         "connector_config": {},
         "secrets": [],
         "filesystems": [],
@@ -544,6 +545,7 @@ def test_connection_config_serialization():
         "extensions": [],
         "pre_ping": False,
         "pretty_sql": False,
+        "shared_connection": True,
         "connector_config": {},
         "secrets": [],
         "filesystems": [],

tests/core/test_connection_config.py

@@ -1426,26 +1426,27 @@ def test_databricks(make_config):
         )
 
 
-def test_databricks_shared_connection(make_config):
-    """Databricks should use a shared connection pool to prevent OAuth CSRF races.
+def test_databricks__u2m_oauth__shared_connection_pool(make_config):
+    """Databricks should use a shared connection pool when using OAuth to prevent CSRF races.
 
     When concurrent_tasks > 1, ThreadLocalConnectionPool creates one connection per
     thread. For U2M OAuth, each thread triggers its own browser-based OAuth flow;
     these race on the CSRF state parameter and cause MismatchingStateError.
 
-    Setting shared_connection = True causes ThreadLocalSharedConnectionPool to be
-    used instead: a single connection is created (behind a lock) and each thread
-    gets its own cursor, so only one OAuth flow is ever initiated.
+    For non-U2M OAuth authentication types (e.g. access_token and M2M OAuth) then
+    ThreadLocalConnectionPool should still be used.
 
-    See: https://github.com/tobymao/sqlmesh/issues/5646
+    See:
+        https://github.com/tobymao/sqlmesh/issues/5646
+        https://github.com/SQLMesh/sqlmesh/issues/5858
     """
     from sqlmesh.utils.connection_pool import ThreadLocalSharedConnectionPool
 
     config = make_config(
         type="databricks",
         server_hostname="dbc-test.cloud.databricks.com",
         http_path="sql/test/foo",
-        access_token="test-token",
+        auth_type="databricks-oauth",
         concurrent_tasks=4,
     )
     assert isinstance(config, DatabricksConnectionConfig)
@@ -1455,6 +1456,43 @@ def test_databricks_shared_connection(make_config):
     assert isinstance(adapter._connection_pool, ThreadLocalSharedConnectionPool)
 
 
+@patch.object(DatabricksConnectionConfig, "_connection_factory_with_kwargs")
+def test_databricks__m2m_oauth__connection_pool(mock_connection_factory_with_kwargs, make_config):
+    from sqlmesh.utils.connection_pool import ThreadLocalConnectionPool
+
+    config = make_config(
+        type="databricks",
+        server_hostname="dbc-test.cloud.databricks.com",
+        http_path="sql/test/foo",
+        auth_type="databricks-oauth",
+        oauth_client_id="oauth_client_id",
+        oauth_client_secret="oauth_client_secret",
+        concurrent_tasks=4,
+    )
+    assert isinstance(config, DatabricksConnectionConfig)
+    assert config.shared_connection is False
+
+    adapter = config.create_engine_adapter()
+    assert isinstance(adapter._connection_pool, ThreadLocalConnectionPool)
+
+
+def test_databricks__access_token__connection_pool(make_config):
+    from sqlmesh.utils.connection_pool import ThreadLocalConnectionPool
+
+    config = make_config(
+        type="databricks",
+        server_hostname="dbc-test.cloud.databricks.com",
+        http_path="sql/test/foo",
+        access_token="any-token",
+        concurrent_tasks=4,
+    )
+    assert isinstance(config, DatabricksConnectionConfig)
+    assert config.shared_connection is False
+
+    adapter = config.create_engine_adapter()
+    assert isinstance(adapter._connection_pool, ThreadLocalConnectionPool)
+
+
 def test_engine_import_validator():
     with pytest.raises(
         ConfigError,

tests/integrations/jupyter/test_magics.py

@@ -745,16 +745,16 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
         "Models: 20",
         "Macros: 8",
         "",
-        "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []",
-        "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []",
+        "Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []\n  shared_connection: true",
+        "Test Connection:\n  type: duckdb\n  concurrent_tasks: 1\n  register_comments: true\n  pre_ping: false\n  pretty_sql: false\n  extensions: []\n  connector_config: {}\n  secrets: None\n  filesystems: []\n  shared_connection: true",
         "Data warehouse connection succeeded",
     ]
     assert get_all_html_output(output) == [
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Models: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">20</span></pre>",
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Macros: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">8</span></pre>",
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>",
-        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span></pre>',
-        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span></pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span>  shared_connection: true</pre>',
+        '<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection:  type: duckdb  concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span>  register_comments: true  pre_ping: false  pretty_sql: false  extensions: <span style="font-weight: bold">[]</span>  connector_config: <span style="font-weight: bold">{}</span>  secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span>  filesystems: <span style="font-weight: bold">[]</span>  shared_connection: true</pre>',
         "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Data warehouse connection <span style=\"color: #008000; text-decoration-color: #008000\">succeeded</span></pre>",
     ]