Merge branch 'main' into fix/clickhouse-primary-key-paren

Author: mday-io

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))

sqlmesh/core/engine_adapter/databricks.py

@@ -154,7 +154,7 @@ def _set_spark_engine_adapter_if_needed(self) -> None:
             host=self._extra_config["databricks_connect_server_hostname"],
             token=self._extra_config.get("databricks_connect_access_token"),
         )
-        if "databricks_connect_use_serverless" in self._extra_config:
+        if self._extra_config.get("databricks_connect_use_serverless"):
             connect_kwargs["serverless"] = True
         else:
             connect_kwargs["cluster_id"] = self._extra_config["databricks_connect_cluster_id"]

tests/core/engine_adapter/test_databricks.py

@@ -732,3 +732,108 @@ def test_columns(mocker: MockFixture, make_mocked_engine_adapter: t.Callable):
     adapter.cursor.execute.assert_called_once_with(
         """SELECT columns.column_name, columns.full_data_type FROM system.information_schema.columns WHERE table_name = 'test_table' AND table_schema = 'test_db' AND table_catalog = 'test_catalog' ORDER BY ordinal_position ASC"""
     )
+
+
+def _make_databricks_connect_adapter(
+    mocker: MockFixture,
+    make_mocked_engine_adapter: t.Callable,
+    extra_config: t.Dict[str, t.Any],
+) -> t.Tuple[DatabricksEngineAdapter, t.Any]:
+    """Helper that creates a DatabricksEngineAdapter with Databricks Connect mocked out."""
+    import sys
+    import types
+
+    mock_session = mocker.MagicMock()
+    mock_builder = mocker.MagicMock()
+    mock_builder.remote.return_value = mock_builder
+    mock_builder.userAgent.return_value = mock_builder
+    mock_builder.getOrCreate.return_value = mock_session
+    mock_databricks_session_cls = mocker.MagicMock()
+    mock_databricks_session_cls.builder = mock_builder
+
+    # databricks.connect is a local import inside the method, so inject via sys.modules
+    mock_connect_module = types.ModuleType("databricks.connect")
+    mock_connect_module.DatabricksSession = mock_databricks_session_cls  # type: ignore
+    mock_databricks_module = types.ModuleType("databricks")
+    mocker.patch.dict(
+        sys.modules,
+        {"databricks": mock_databricks_module, "databricks.connect": mock_connect_module},
+    )
+
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.can_access_spark_session",
+        return_value=False,
+    )
+    mocker.patch(
+        "sqlmesh.core.engine_adapter.databricks.DatabricksEngineAdapter.can_access_databricks_connect",
+        return_value=True,
+    )
+
+    adapter = make_mocked_engine_adapter(
+        DatabricksEngineAdapter,
+        default_catalog="test_catalog",
+        **extra_config,
+    )
+    return adapter, mock_builder
+
+
+def test_databricks_connect_routes_to_cluster_id(
+    mocker: MockFixture, make_mocked_engine_adapter: t.Callable
+) -> None:
+    """cluster_id is used when databricks_connect_use_serverless is absent."""
+    extra_config = {
+        "databricks_connect_server_hostname": "myhost.azuredatabricks.net",
+        "databricks_connect_access_token": "mytoken",
+        "databricks_connect_cluster_id": "0123-456789-mycluster",
+    }
+    _, mock_builder = _make_databricks_connect_adapter(
+        mocker, make_mocked_engine_adapter, extra_config
+    )
+
+    mock_builder.remote.assert_called_once_with(
+        host="myhost.azuredatabricks.net",
+        token="mytoken",
+        cluster_id="0123-456789-mycluster",
+    )
+
+
+def test_databricks_connect_routes_to_serverless(
+    mocker: MockFixture, make_mocked_engine_adapter: t.Callable
+) -> None:
+    """serverless=True is used when databricks_connect_use_serverless is truthy."""
+    extra_config = {
+        "databricks_connect_server_hostname": "myhost.azuredatabricks.net",
+        "databricks_connect_access_token": "mytoken",
+        "databricks_connect_cluster_id": "0123-456789-mycluster",
+        "databricks_connect_use_serverless": True,
+    }
+    _, mock_builder = _make_databricks_connect_adapter(
+        mocker, make_mocked_engine_adapter, extra_config
+    )
+
+    mock_builder.remote.assert_called_once_with(
+        host="myhost.azuredatabricks.net",
+        token="mytoken",
+        serverless=True,
+    )
+
+
+def test_databricks_connect_cluster_id_not_overridden_by_falsy_serverless(
+    mocker: MockFixture, make_mocked_engine_adapter: t.Callable
+) -> None:
+    """cluster_id is used when databricks_connect_use_serverless is present but False."""
+    extra_config = {
+        "databricks_connect_server_hostname": "myhost.azuredatabricks.net",
+        "databricks_connect_access_token": "mytoken",
+        "databricks_connect_cluster_id": "0123-456789-mycluster",
+        "databricks_connect_use_serverless": False,
+    }
+    _, mock_builder = _make_databricks_connect_adapter(
+        mocker, make_mocked_engine_adapter, extra_config
+    )
+
+    mock_builder.remote.assert_called_once_with(
+        host="myhost.azuredatabricks.net",
+        token="mytoken",
+        cluster_id="0123-456789-mycluster",
+    )

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>",
     ]