handle case insensitive db name in schema collection (#19384)

* handle case insensitive db name in schema collection

* add unit test

* add changelog

* update log

Author: lu-zhengda

sqlserver/changelog.d/19384.fixed

@@ -0,0 +1 @@
+Fix KeyError in SQL Server schema collection caused by case-insensitive database name mismatches.

sqlserver/datadog_checks/sqlserver/connection.py

@@ -8,7 +8,7 @@
 from datadog_checks.base import AgentCheck, ConfigurationError
 from datadog_checks.base.log import get_check_logger
 from datadog_checks.sqlserver.cursor import CommenterCursorWrapper
-from datadog_checks.sqlserver.utils import construct_use_statement
+from datadog_checks.sqlserver.utils import construct_use_statement, is_collation_case_insensitive
 
 try:
     import adodbapi
@@ -361,7 +361,7 @@ def _check_db_exists(self):
                 cursor.execute(DATABASE_EXISTS_QUERY)
                 for row in cursor.fetchall():
                     # collation_name can be NULL if db offline, in that case assume its case_insensitive
-                    case_insensitive = not row.collation_name or 'CI' in row.collation_name
+                    case_insensitive = is_collation_case_insensitive(row.collation_name)
                     self.existing_databases[row.name.lower()] = (
                         case_insensitive,
                         row.name,

sqlserver/datadog_checks/sqlserver/schemas.py

@@ -31,7 +31,13 @@
     SCHEMA_QUERY,
     TABLES_IN_SCHEMA_QUERY,
 )
-from datadog_checks.sqlserver.utils import convert_to_bool, execute_query, get_list_chunks, is_azure_sql_database
+from datadog_checks.sqlserver.utils import (
+    convert_to_bool,
+    execute_query,
+    get_list_chunks,
+    is_azure_sql_database,
+    is_collation_case_insensitive,
+)
 
 
 class SubmitData:
@@ -57,9 +63,23 @@ def reset(self):
         self.db_to_schemas.clear()
         self.db_info.clear()
 
-    def store_db_infos(self, db_infos):
+    def store_db_infos(self, db_infos, databases):
+        dbs = set(databases)
         for db_info in db_infos:
-            self.db_info[db_info['name']] = db_info
+            case_insensitive = is_collation_case_insensitive(db_info.get('collation'))
+            db_name = db_info['name']
+            db_name_lower = db_name.lower()
+            if db_name not in dbs:
+                if db_name.lower() in dbs and case_insensitive:
+                    db_name = db_name_lower
+                else:
+                    self._log.debug(
+                        "Skipping db {} as it is not in the databases list {} or collation is case sensitive".format(
+                            db_name, dbs
+                        )
+                    )
+                    continue
+            self.db_info[db_name] = db_info
 
     def store(self, db_name, schema, tables, columns_count):
         self._columns_count += columns_count
@@ -277,7 +297,7 @@ def _collect_schemas_data(self):
 
         databases = self._check.get_databases()
         db_infos = self._query_db_information(databases)
-        self._data_submitter.store_db_infos(db_infos)
+        self._data_submitter.store_db_infos(db_infos, databases)
         self._fetch_for_databases()
         self._data_submitter.submit()
         self._log.debug("Finished collect_schemas_data")

sqlserver/datadog_checks/sqlserver/utils.py

@@ -215,3 +215,12 @@ def convert_to_bool(value):
         return bool(value)
     else:
         return value
+
+
+def is_collation_case_insensitive(collation):
+    """
+    Checks if the collation is case insensitive
+    :param collation: The collation string
+    :return: bool
+    """
+    return not collation or 'CI' in collation.upper()

sqlserver/tests/test_unit.py

@@ -754,7 +754,9 @@ def test_submit_data():
 
     dataSubmitter, submitted_data = set_up_submitter_unit_test()
 
-    dataSubmitter.store_db_infos([{"id": 3, "name": "test_db1"}, {"id": 4, "name": "test_db2"}])
+    dataSubmitter.store_db_infos(
+        [{"id": 3, "name": "test_db1"}, {"id": 4, "name": "test_db2"}], ["test_db1", "test_db2"]
+    )
     schema1 = {"id": "1"}
     schema2 = {"id": "2"}
     schema3 = {"id": "3"}
@@ -787,6 +789,46 @@ def test_submit_data():
     assert deep_compare(data, expected_data)
 
 
+@pytest.mark.parametrize(
+    "db_infos, databases, expected_dbs",
+    [
+        pytest.param(
+            [
+                {"id": 3, "name": "test_db1", "collation": "SQL_Latin1_General_CP1_CI_AS"},
+                {"id": 4, "name": "TEST_DB2", "collation": "SQL_Latin1_General_CP1_CI_AS"},
+            ],
+            ["test_db1", "test_db2"],
+            ["test_db1", "test_db2"],
+            id="case_insensitive",
+        ),
+        pytest.param(
+            [{"id": 3, "name": "test_db1", "collation": "SQL_Latin1_General_CP1_CS_AS"}],
+            ["TEST_DB1"],
+            [],
+            id="case_sensitive",
+        ),
+        pytest.param(
+            [{"id": 3, "name": "test_db1", "collation": "SQL_Latin1_General_CP1_CS_AS"}],
+            ["test_db1"],
+            ["test_db1"],
+            id="case_sensitive_lowercase",
+        ),
+        pytest.param(
+            [{"id": 3, "name": "TEST_DB1", "collation": "SQL_Latin1_General_CP1_CS_AS"}],
+            ["TEST_DB1"],
+            ["TEST_DB1"],
+            id="case_sensitive_uppercase",
+        ),
+    ],
+)
+def test_store_db_infos_case_sensitive(db_infos, databases, expected_dbs):
+    dataSubmitter, _ = set_up_submitter_unit_test()
+    dataSubmitter.db_info.clear()
+
+    dataSubmitter.store_db_infos(db_infos, databases)
+    assert list(dataSubmitter.db_info.keys()) == expected_dbs
+
+
 def test_fetch_throws(instance_docker):
     check = SQLServer(CHECK_NAME, {}, [instance_docker])
     schemas = Schemas(check, check._config)