diff --git a/sqlit/domains/connections/providers/mssql/adapter.py b/sqlit/domains/connections/providers/mssql/adapter.py index dc67ba09..57fa08e5 100644 --- a/sqlit/domains/connections/providers/mssql/adapter.py +++ b/sqlit/domains/connections/providers/mssql/adapter.py @@ -3,6 +3,7 @@ from __future__ import annotations import struct +import weakref from typing import TYPE_CHECKING, Any from sqlit.domains.connections.providers.adapters.base import ( @@ -87,11 +88,24 @@ def _first_actionable_line(text: str) -> str: return "" +def _is_use_unsupported_error(exc: Exception) -> bool: + """True for Azure SQL Database's refusal to run a USE statement.""" + return "USE statement is not supported" in str(exc) + + class SQLServerAdapter(DatabaseAdapter): """Adapter for Microsoft SQL Server using the mssql-python driver.""" def __init__(self) -> None: self._supports_cross_database_queries_override: bool | None = None + # Azure SQL Database rejects USE, so metadata for a database other than + # the one the connection is bound to needs its own connection. The + # config a connection was opened with is kept so we can rebuild the + # connection string with DATABASE= swapped; the extra connections are + # keyed by (id(parent), lowercased database name). + self._configs: dict[int, Any] = {} + self._db_conns: dict[tuple[int, str], Any] = {} + self._conn_refs: dict[int, Any] = {} @property def name(self) -> str: @@ -293,8 +307,42 @@ def connect(self, config: ConnectionConfig) -> Any: conn = mssql_python.connect(conn_str, attrs_before=attrs_before) # Enable autocommit to allow DDL statements like CREATE DATABASE conn.autocommit = True + self._track(conn, config) return conn + def _track(self, conn: Any, config: ConnectionConfig) -> None: + """Remember the config a connection was opened with. + + Entries are keyed by id(), so they have to go before the id can be + handed to a different object: a weakref callback drops them when the + connection is collected without an explicit disconnect(). Drivers whose + connection objects don't support weak references keep the entry until + disconnect() clears it. + """ + key = id(conn) + try: + self._conn_refs[key] = weakref.ref(conn, lambda _ref, key=key: self._forget(key)) + except TypeError: + pass + self._configs[key] = config + + def _forget(self, parent_id: int) -> None: + """Drop the bookkeeping for a connection and its per-database siblings.""" + for key in [k for k in self._db_conns if k[0] == parent_id]: + sibling = self._db_conns.pop(key) + self._forget(id(sibling)) + try: + sibling.close() + except Exception: + pass + self._configs.pop(parent_id, None) + self._conn_refs.pop(parent_id, None) + + def disconnect(self, conn: Any) -> None: + """Close the connection plus any per-database siblings opened for it.""" + self._forget(id(conn)) + super().disconnect(conn) + def _preflight_azure_credentials(self, config: ConnectionConfig) -> str | None: """Acquire a SQL Entra token and return it for direct ODBC attach. @@ -360,11 +408,61 @@ def get_databases(self, conn: Any) -> list[str]: return [row[0] for row in cursor.fetchall()] def _get_cursor_for_database(self, conn: Any, database: str | None) -> Any: - """Get a cursor for the specified database using USE statement.""" - cursor = conn.cursor() - if database: + """Get a cursor scoped to `database`. + + On a normal SQL Server instance that is a `USE`. Azure SQL Database + (EngineEdition 5/6) rejects USE outright — "USE statement is not + supported to switch between databases" — so there we open a second + connection bound to the target database and cache it per parent + connection. The first rejected USE also flips the cross-database + capability off, so later calls skip straight to the sibling path even + when detect_capabilities never ran (the process worker reads + capabilities before it connects). + """ + if not database: + return conn.cursor() + + if self._connection_database(conn) == database.lower(): + return conn.cursor() + + if self.supports_cross_database_queries: + cursor = conn.cursor() + try: + cursor.execute(f"USE [{database}]") + except Exception as exc: + if not _is_use_unsupported_error(exc): + raise + self._supports_cross_database_queries_override = False + else: + return cursor + + target = self._connection_for_database(conn, database) + return target.cursor() + + def _connection_database(self, conn: Any) -> str: + """Lowercased database the connection was opened against, or ''.""" + config = self._configs.get(id(conn)) + endpoint = config.tcp_endpoint if config is not None else None + return (endpoint.database or "").lower() if endpoint is not None else "" + + def _connection_for_database(self, conn: Any, database: str) -> Any: + """Return a connection bound to `database`, opening one if needed.""" + key = (id(conn), database.lower()) + cached = self._db_conns.get(key) + if cached is not None: + return cached + + config = self._configs.get(id(conn)) + if config is None: + # Connection came from somewhere other than our connect() — nothing + # to rebuild a connection string from, so USE is the only option. + cursor = conn.cursor() cursor.execute(f"USE [{database}]") - return cursor + return conn + + sibling = self.connect(self.apply_database_override(config, database)) + self._db_conns[key] = sibling + return sibling def get_tables(self, conn: Any, database: str | None = None) -> list[TableInfo]: """Get list of tables with schema from SQL Server.""" diff --git a/tests/unit/test_mssql_azure_database_switch.py b/tests/unit/test_mssql_azure_database_switch.py new file mode 100644 index 00000000..aea55c85 --- /dev/null +++ b/tests/unit/test_mssql_azure_database_switch.py @@ -0,0 +1,226 @@ +"""Unit tests for switching databases on Azure SQL Database. + +Azure SQL Database (EngineEdition 5/6) rejects USE outright: + + DDBC Error: USE statement is not supported to switch between databases. + Use a new connection to connect to a different database. + +The adapter falls back to opening a second connection bound to the target +database, so the explorer can still browse other databases on the server. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + +USE_NOT_SUPPORTED = ( + "Driver Error: Syntax error or access violation; DDBC Error: USE statement " + "is not supported to switch between databases. Use a new connection to " + "connect to a different database." +) + + +class FakeCursor: + def __init__(self, conn: FakeConnection) -> None: + self.conn = conn + + def execute(self, sql: str, *args: object) -> None: + self.conn.executed.append(sql) + if sql.startswith("USE") and self.conn.use_error is not None: + raise self.conn.use_error() + + def fetchall(self) -> list: + return [] + + def fetchone(self) -> None: + return None + + +class FakeConnection: + """Stand-in for an mssql_python connection. + + `use_error` builds the exception a USE statement should raise, or is None + on a server where USE works. + """ + + def __init__(self, database: str, use_error) -> None: + self.database = database + self.use_error = use_error + self.executed: list[str] = [] + self.closed = False + self.autocommit = False + + def cursor(self) -> FakeCursor: + return FakeCursor(self) + + def close(self) -> None: + self.closed = True + + +def _database_from(conn_str: str) -> str: + for part in conn_str.split(";"): + key, _, value = part.partition("=") + if key.strip().upper() == "DATABASE": + return value.strip() + return "" + + +@pytest.fixture +def driver(): + """Patch mssql_python so connect() hands back FakeConnections. + + Yields a factory: `adapter, opened = driver(use_error=...)`, where `opened` + accumulates every connection the adapter opened, in order. + """ + module = MagicMock() + with patch.dict("sys.modules", {"mssql_python": module}): + + def _factory(use_error): + from sqlit.domains.connections.providers.mssql.adapter import SQLServerAdapter + + opened: list[FakeConnection] = [] + + def fake_connect(conn_str: str, attrs_before=None) -> FakeConnection: + conn = FakeConnection(_database_from(conn_str), use_error) + opened.append(conn) + return conn + + module.connect.side_effect = fake_connect + return SQLServerAdapter(), opened + + yield _factory + + +def azure_driver(driver): + return driver(lambda: RuntimeError(USE_NOT_SUPPORTED)) + + +def sql_server_driver(driver): + return driver(None) + + +@pytest.fixture +def config(): + def _config(database: str): + from sqlit.domains.connections.domain.config import ConnectionConfig, TcpEndpoint + + return ConnectionConfig( + name="test_mssql", + db_type="mssql", + endpoint=TcpEndpoint( + host="server.database.windows.net", + port="1433", + database=database, + username="sa", + password="password", + ), + options={"auth_type": "sql"}, + ) + + return _config + + +class TestAzureDatabaseSwitching: + def test_use_failure_falls_back_to_a_new_connection(self, driver, config): + adapter, opened = azure_driver(driver) + root = adapter.connect(config("master")) + + adapter.get_tables(root, database="TestDB") + + assert root.executed == ["USE [TestDB]"] + assert [conn.database for conn in opened] == ["master", "TestDB"] + assert "INFORMATION_SCHEMA.TABLES" in opened[1].executed[0] + + def test_second_lookup_reuses_the_new_connection(self, driver, config): + adapter, opened = azure_driver(driver) + root = adapter.connect(config("master")) + + adapter.get_tables(root, database="TestDB") + adapter.get_views(root, database="TestDB") + + assert len(opened) == 2, "Expected one extra connection, not one per lookup" + # The rejected USE is probed once; after that the adapter knows better. + assert root.executed == ["USE [TestDB]"] + + def test_each_database_gets_its_own_connection(self, driver, config): + adapter, opened = azure_driver(driver) + root = adapter.connect(config("master")) + + adapter.get_tables(root, database="TestDB") + adapter.get_tables(root, database="OtherDB") + + assert [conn.database for conn in opened] == ["master", "TestDB", "OtherDB"] + + def test_current_database_needs_no_use_or_reconnect(self, driver, config): + adapter, opened = azure_driver(driver) + root = adapter.connect(config("TestDB")) + + adapter.get_tables(root, database="TestDB") + + assert len(opened) == 1 + assert not any(sql.startswith("USE") for sql in root.executed) + + def test_disconnect_closes_the_extra_connections(self, driver, config): + adapter, opened = azure_driver(driver) + root = adapter.connect(config("master")) + + adapter.get_tables(root, database="TestDB") + adapter.get_tables(root, database="OtherDB") + adapter.disconnect(root) + + assert all(conn.closed for conn in opened) + assert adapter._db_conns == {} + assert adapter._configs == {} + + def test_collected_connection_drops_its_bookkeeping(self, driver, config): + """A connection closed by GC rather than disconnect() must not linger. + + The per-connection state is keyed by id(), which Python reuses once an + object is collected - a stale entry would bind a later connection to + the wrong database. + """ + import gc + + adapter, opened = azure_driver(driver) + root = adapter.connect(config("master")) + adapter.get_tables(root, database="TestDB") + sibling = opened[1] + + del root + opened.clear() # the fixture's bookkeeping would otherwise keep it alive + gc.collect() + + assert adapter._configs == {} + assert adapter._db_conns == {} + assert sibling.closed is True + + def test_capability_flips_off_after_a_rejected_use(self, driver, config): + adapter, _opened = azure_driver(driver) + root = adapter.connect(config("master")) + + assert adapter.supports_cross_database_queries is True + adapter.get_tables(root, database="TestDB") + assert adapter.supports_cross_database_queries is False + + +class TestRegularSQLServerUnchanged: + def test_use_is_still_the_happy_path(self, driver, config): + adapter, opened = sql_server_driver(driver) + root = adapter.connect(config("master")) + + adapter.get_tables(root, database="TestDB") + + assert len(opened) == 1, "Regular SQL Server should not open extra connections" + assert root.executed[0] == "USE [TestDB]" + + def test_unrelated_use_errors_are_not_swallowed(self, driver, config): + adapter, opened = driver(lambda: RuntimeError("Login failed for user 'sa'.")) + root = adapter.connect(config("master")) + + with pytest.raises(RuntimeError, match="Login failed"): + adapter.get_tables(root, database="TestDB") + + assert len(opened) == 1, "A non-USE failure must not trigger a reconnect" + assert adapter.supports_cross_database_queries is True