Reuse legacy schema detail connection (DataDog#23544)

Keep the legacy schema detail connection open for the duration of each database schema collection instead of reconnecting once per table.

Author: lu-zhengda

sqlserver/changelog.d/23544.fixed

@@ -0,0 +1 @@
+Reuse the auxiliary SQL Server schema collection connection for legacy table detail queries.

sqlserver/datadog_checks/sqlserver/schemas.py

@@ -86,6 +86,7 @@ def __init__(self, check: SQLServer):
         config.max_tables = check._config.schema_config.get('max_tables', 300)
         self._is_2016_or_earlier = None
         self._database_compatibility_levels: dict[str, int] = {}
+        self._pre_2017_cursor = None
         super().__init__(check, config)
 
     @property
@@ -106,14 +107,26 @@ def _get_databases(self) -> list[DatabaseInfo]:
 
     @contextlib.contextmanager
     def _get_cursor(self, database_name):
-        with self._check.connection.open_managed_default_connection(KEY_PREFIX):
-            with self._check.connection.get_managed_cursor(KEY_PREFIX) as cursor:
+        with contextlib.ExitStack() as stack:
+            try:
+                stack.enter_context(self._check.connection.open_managed_default_connection(KEY_PREFIX))
+                cursor = stack.enter_context(self._check.connection.get_managed_cursor(KEY_PREFIX))
                 switch_db_statement = construct_use_statement(database_name)
                 cursor.execute(switch_db_statement)
                 self._is_2016_or_earlier = self._should_use_legacy_schema_query(database_name)
+
+                if self._is_2016_or_earlier:
+                    stack.enter_context(self._check.connection.open_managed_default_connection(KEY_PREFIX_PRE_2017))
+                    self._pre_2017_cursor = stack.enter_context(
+                        self._check.connection.get_managed_cursor(KEY_PREFIX_PRE_2017)
+                    )
+                    self._pre_2017_cursor.execute(switch_db_statement)
+
                 query = self._get_tables_query()
                 cursor.execute(query)
                 yield cursor
+            finally:
+                self._pre_2017_cursor = None
 
     def _record_database_compatibility_levels(self, databases: list[DatabaseInfo]) -> None:
         self._database_compatibility_levels = {}
@@ -205,25 +218,25 @@ def _map_row(self, database: DatabaseInfo, cursor_row) -> DatabaseObject:
         object = super()._map_row(database, cursor_row)
         if self._is_2016_or_earlier:
             # We need to fetch the related data for each table
-            # Use a key_prefix to get a separate connection to avoid conflicts with the main connection
-            with self._check.connection.open_managed_default_connection(KEY_PREFIX_PRE_2017):
-                with self._check.connection.get_managed_cursor(KEY_PREFIX_PRE_2017) as cursor:
-                    switch_db_statement = construct_use_statement(database.get("name"))
-                    cursor.execute(switch_db_statement)
-                    table_id = str(cursor_row.get("table_id"))
-                    columns_query = COLUMN_QUERY.replace("schema_tables.table_id", table_id)
-                    cursor.execute(columns_query)
-                    columns = cursor.fetchall_dict()
-                    indexes_query = INDEX_QUERY_PRE_2017.replace("schema_tables.table_id", table_id)
-                    cursor.execute(indexes_query)
-                    indexes = cursor.fetchall_dict()
-                    foreign_keys_query = FOREIGN_KEY_QUERY_PRE_2017.replace("schema_tables.table_id", table_id)
-                    cursor.execute(foreign_keys_query)
-                    foreign_keys = cursor.fetchall_dict()
-                    partitions_query = PARTITIONS_QUERY.replace("schema_tables.table_id", table_id)
-                    cursor.execute(partitions_query)
-                    partition_row = cursor.fetchone_dict()
-                    partition_count = partition_row.get("partition_count") if partition_row else None
+            # Use a separate connection to avoid conflicts with the main cursor while it streams table rows.
+            cursor = self._pre_2017_cursor
+            if cursor is None:
+                raise RuntimeError("pre-2017 schema cursor is not initialized")
+
+            table_id = str(cursor_row.get("table_id"))
+            columns_query = COLUMN_QUERY.replace("schema_tables.table_id", table_id)
+            cursor.execute(columns_query)
+            columns = cursor.fetchall_dict()
+            indexes_query = INDEX_QUERY_PRE_2017.replace("schema_tables.table_id", table_id)
+            cursor.execute(indexes_query)
+            indexes = cursor.fetchall_dict()
+            foreign_keys_query = FOREIGN_KEY_QUERY_PRE_2017.replace("schema_tables.table_id", table_id)
+            cursor.execute(foreign_keys_query)
+            foreign_keys = cursor.fetchall_dict()
+            partitions_query = PARTITIONS_QUERY.replace("schema_tables.table_id", table_id)
+            cursor.execute(partitions_query)
+            partition_row = cursor.fetchone_dict()
+            partition_count = partition_row.get("partition_count") if partition_row else None
         else:
             columns = json.loads(cursor_row.get("columns") or "[]")
             indexes = json.loads(cursor_row.get("indexes") or "[]")

sqlserver/tests/test_unit.py

@@ -27,7 +27,7 @@
     STATIC_INFO_VERSION,
 )
 from datadog_checks.sqlserver.metrics import SqlFractionMetric
-from datadog_checks.sqlserver.schemas import SQLServerSchemaCollector
+from datadog_checks.sqlserver.schemas import KEY_PREFIX, KEY_PREFIX_PRE_2017, SQLServerSchemaCollector
 from datadog_checks.sqlserver.sqlserver import SQLConnectionError
 from datadog_checks.sqlserver.utils import (
     Database,
@@ -219,16 +219,165 @@ def test_schema_collector_uses_database_compatibility_level_for_schema_query(
         }
     )
     collector._database_compatibility_levels = {"datadog_test": compatibility_level}
-    cursor = mock.Mock()
-    collector._check.connection.open_managed_default_connection.return_value = contextlib.nullcontext()
-    collector._check.connection.get_managed_cursor.return_value = contextlib.nullcontext(cursor)
+    main_cursor = mock.Mock()
+    detail_cursor = mock.Mock()
+    collector._check.connection.open_managed_default_connection.side_effect = [
+        contextlib.nullcontext(),
+        contextlib.nullcontext(),
+    ]
+    collector._check.connection.get_managed_cursor.side_effect = [
+        contextlib.nullcontext(main_cursor),
+        contextlib.nullcontext(detail_cursor),
+    ]
 
     with collector._get_cursor("datadog_test"):
         pass
 
-    tables_query = cursor.execute.call_args_list[1][0][0]
+    tables_query = main_cursor.execute.call_args_list[1][0][0]
     assert collector._is_2016_or_earlier is expected_legacy
     assert ("STRING_AGG" in tables_query) is not expected_legacy
+    assert (detail_cursor.execute.call_count == 1) is expected_legacy
+
+
+def test_schema_collector_reuses_pre_2017_connection_for_table_detail_queries() -> None:
+    collector = create_schema_collector(
+        {
+            STATIC_INFO_ENGINE_EDITION: ENGINE_EDITION_STANDARD,
+            STATIC_INFO_MAJOR_VERSION: 13,
+        }
+    )
+    main_cursor = mock.Mock()
+    detail_cursor = mock.Mock()
+    detail_cursor.fetchall_dict.side_effect = [
+        [{"name": "id"}],
+        [{"name": "pk_table_1"}],
+        [{"foreign_key_name": "fk_table_1"}],
+        [{"name": "id"}],
+        [{"name": "pk_table_2"}],
+        [{"foreign_key_name": "fk_table_2"}],
+    ]
+    detail_cursor.fetchone_dict.side_effect = [{"partition_count": 1}, {"partition_count": 2}]
+    collector._check.connection.open_managed_default_connection.side_effect = [
+        contextlib.nullcontext(),
+        contextlib.nullcontext(),
+    ]
+    collector._check.connection.get_managed_cursor.side_effect = [
+        contextlib.nullcontext(main_cursor),
+        contextlib.nullcontext(detail_cursor),
+    ]
+    database = {"name": "datadog_test", "id": "1", "collation": "SQL_Latin1_General_CP1_CI_AS", "owner": "dbo"}
+    rows = [
+        {"schema_name": "test_schema", "schema_id": 1, "owner_name": "dbo", "table_name": "table_1", "table_id": 101},
+        {"schema_name": "test_schema", "schema_id": 1, "owner_name": "dbo", "table_name": "table_2", "table_id": 102},
+    ]
+
+    with collector._get_cursor("datadog_test"):
+        mapped_rows = [collector._map_row(database, row) for row in rows]
+
+    assert collector._check.connection.open_managed_default_connection.call_args_list == [
+        mock.call(KEY_PREFIX),
+        mock.call(KEY_PREFIX_PRE_2017),
+    ]
+    assert collector._check.connection.get_managed_cursor.call_args_list == [
+        mock.call(KEY_PREFIX),
+        mock.call(KEY_PREFIX_PRE_2017),
+    ]
+    assert detail_cursor.execute.call_args_list[0] == mock.call("USE [datadog_test];")
+    assert detail_cursor.execute.call_args_list.count(mock.call("USE [datadog_test];")) == 1
+    assert detail_cursor.execute.call_count == 9
+    assert collector._pre_2017_cursor is None
+    assert [row["schemas"][0]["tables"][0]["partitions"]["partition_count"] for row in mapped_rows] == [1, 2]
+
+
+def configure_pre_2017_detail_cursor(cursor: mock.Mock, partition_count: int = 1) -> None:
+    cursor.fetchall_dict.side_effect = [
+        [{"name": "id"}],
+        [{"name": "pk_table"}],
+        [{"foreign_key_name": "fk_table"}],
+    ]
+    cursor.fetchone_dict.return_value = {"partition_count": partition_count}
+
+
+def test_schema_collector_uses_new_pre_2017_connection_for_each_database() -> None:
+    collector = create_schema_collector(
+        {
+            STATIC_INFO_ENGINE_EDITION: ENGINE_EDITION_STANDARD,
+            STATIC_INFO_MAJOR_VERSION: 13,
+        }
+    )
+    main_cursor_1 = mock.Mock()
+    detail_cursor_1 = mock.Mock()
+    configure_pre_2017_detail_cursor(detail_cursor_1, partition_count=1)
+    main_cursor_2 = mock.Mock()
+    detail_cursor_2 = mock.Mock()
+    configure_pre_2017_detail_cursor(detail_cursor_2, partition_count=2)
+    collector._check.connection.open_managed_default_connection.side_effect = [
+        contextlib.nullcontext(),
+        contextlib.nullcontext(),
+        contextlib.nullcontext(),
+        contextlib.nullcontext(),
+    ]
+    collector._check.connection.get_managed_cursor.side_effect = [
+        contextlib.nullcontext(main_cursor_1),
+        contextlib.nullcontext(detail_cursor_1),
+        contextlib.nullcontext(main_cursor_2),
+        contextlib.nullcontext(detail_cursor_2),
+    ]
+    row = {"schema_name": "test_schema", "schema_id": 1, "owner_name": "dbo", "table_name": "table", "table_id": 101}
+
+    with collector._get_cursor("datadog_test_1"):
+        collector._map_row(
+            {"name": "datadog_test_1", "id": "1", "collation": "SQL_Latin1_General_CP1_CI_AS", "owner": "dbo"},
+            row,
+        )
+    with collector._get_cursor("datadog_test_2"):
+        collector._map_row(
+            {"name": "datadog_test_2", "id": "2", "collation": "SQL_Latin1_General_CP1_CI_AS", "owner": "dbo"},
+            row,
+        )
+
+    assert collector._check.connection.get_managed_cursor.call_args_list == [
+        mock.call(KEY_PREFIX),
+        mock.call(KEY_PREFIX_PRE_2017),
+        mock.call(KEY_PREFIX),
+        mock.call(KEY_PREFIX_PRE_2017),
+    ]
+    assert detail_cursor_1.execute.call_args_list[0] == mock.call("USE [datadog_test_1];")
+    assert detail_cursor_2.execute.call_args_list[0] == mock.call("USE [datadog_test_2];")
+    assert collector._pre_2017_cursor is None
+
+
+def test_schema_collector_does_not_open_pre_2017_connection_for_modern_query() -> None:
+    collector = create_schema_collector(
+        {
+            STATIC_INFO_ENGINE_EDITION: ENGINE_EDITION_STANDARD,
+            STATIC_INFO_MAJOR_VERSION: 14,
+        }
+    )
+    collector._database_compatibility_levels = {"datadog_test": 130}
+    cursor = mock.Mock()
+    collector._check.connection.open_managed_default_connection.return_value = contextlib.nullcontext()
+    collector._check.connection.get_managed_cursor.return_value = contextlib.nullcontext(cursor)
+    database = {"name": "datadog_test", "id": "1", "collation": "SQL_Latin1_General_CP1_CI_AS", "owner": "dbo"}
+    row = {
+        "schema_name": "test_schema",
+        "schema_id": 1,
+        "owner_name": "dbo",
+        "table_name": "table",
+        "table_id": 101,
+        "columns": '[{"name": "id"}]',
+        "indexes": '[{"name": "pk_table"}]',
+        "foreign_keys": '[{"foreign_key_name": "fk_table"}]',
+        "partition_count": 1,
+    }
+
+    with collector._get_cursor("datadog_test"):
+        mapped_row = collector._map_row(database, row)
+
+    assert collector._check.connection.open_managed_default_connection.call_args_list == [mock.call(KEY_PREFIX)]
+    assert collector._check.connection.get_managed_cursor.call_args_list == [mock.call(KEY_PREFIX)]
+    assert collector._pre_2017_cursor is None
+    assert mapped_row["schemas"][0]["tables"][0]["columns"] == [{"name": "id"}]
 
 
 def test_get_cursor(instance_docker):