Require key_prefix for SQLServer connections (DataDog#22557)

* Require key_prefix for SQLServer connections

* Changelog

* Rename

* Lint

* WIP

* Lint

* Fix

Author: sethsamuel

sqlserver/changelog.d/22557.added

@@ -0,0 +1 @@
+Require key_prefix for SQLServer connections to avoid unsafely sharing connections across threads

sqlserver/datadog_checks/sqlserver/activity.py

@@ -481,8 +481,8 @@ def collect_activity(self):
 
         # re-use the check's conn module, but set extra_key=dbm-activity- to ensure we get our own
         # raw connection. adodbapi and pyodbc modules are thread safe, but connections are not.
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 connections = self._get_active_connections(cursor)
                 request_cols = self._get_exec_requests_cols_cached(cursor, DM_EXEC_REQUESTS_COLS)
                 input_buffer_columns, input_buffer_join = self._get_input_buffer_columns_and_join()

sqlserver/datadog_checks/sqlserver/agent_history.py

@@ -178,8 +178,8 @@ def collect_agent_history(self):
         Collects all current agent activity for the SQLServer intance.
         :return:
         """
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 history_rows = self._get_new_agent_job_history(cursor)
                 history_event = self._create_agent_jobs_history_event(history_rows)
                 payload = json.dumps(history_event, default=default_json_event_encoding)

sqlserver/datadog_checks/sqlserver/connection.py

@@ -177,7 +177,7 @@ def __init__(self, init_config, instance_config, service_check_handler):
         self.log.debug('Connection initialized.')
 
     @contextmanager
-    def get_managed_cursor(self, key_prefix=None):
+    def get_managed_cursor(self, key_prefix):
         cursor = self.get_cursor(self.DEFAULT_DB_KEY, key_prefix=key_prefix)
         try:
             yield cursor
@@ -226,7 +226,7 @@ def open_managed_default_database(self):
             yield
 
     @contextmanager
-    def open_managed_default_connection(self, key_prefix=None):
+    def open_managed_default_connection(self, key_prefix):
         with self._open_managed_db_connections(self.DEFAULT_DB_KEY, key_prefix=key_prefix):
             yield
 
@@ -692,28 +692,28 @@ def test_network_connectivity(self):
 
         return None
 
-    def _get_current_database_context(self):
+    def _get_current_database_context(self, key_prefix):
         """
         Get the current database name.
         """
-        with self.get_managed_cursor() as cursor:
+        with self.get_managed_cursor(key_prefix) as cursor:
             cursor.execute('select DB_NAME()')
             data = cursor.fetchall()
             return data[0][0]
 
     @contextmanager
-    def restore_current_database_context(self):
+    def restore_current_database_context(self, key_prefix):
         """
         Restores the default database after executing use statements.
         """
-        current_db = self._get_current_database_context()
+        current_db = self._get_current_database_context(key_prefix)
         try:
             yield
         finally:
             if current_db:
                 try:
                     self.log.debug("Restoring the original database context %s", current_db)
-                    with self.get_managed_cursor() as cursor:
+                    with self.get_managed_cursor(key_prefix) as cursor:
                         cursor.execute(construct_use_statement(current_db))
                 except Exception as e:
                     self.log.error("Failed to switch back to the original database context %s: %s", current_db, e)

sqlserver/datadog_checks/sqlserver/deadlocks.py

@@ -133,8 +133,8 @@ def _get_connector(self):
         return self._check.connection.connector
 
     def _set_xe_session_name(self):
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 if self._xe_session_name is None:
                     cursor.execute(
                         get_xe_sessions_query(dm_xe_targets=self._dm_xe_targets, dm_xe_sessions=self._dm_xe_sessions)
@@ -180,8 +180,8 @@ def _query_deadlocks(self):
                 f'Using XE session [{self._xe_session_name}], target [{self._xe_session_target}] to collect deadlocks'
             )
 
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 convert_xml_to_str = False
                 if self._force_convert_xml_to_str or self._get_connector() == "adodbapi":
                     convert_xml_to_str = True

sqlserver/datadog_checks/sqlserver/metadata.py

@@ -139,8 +139,8 @@ def _load_settings_rows(self, cursor):
 
     @tracked_method(agent_check_getter=agent_check_getter)
     def report_sqlserver_metadata(self):
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 settings_rows = self._load_settings_rows(cursor)
                 event = {
                     "host": self._check.reported_hostname,

sqlserver/datadog_checks/sqlserver/schemas.py

@@ -98,15 +98,15 @@ def kind(self):
 
     def _get_databases(self):
         database_names = self._check.get_databases()
-        with self._check.connection.open_managed_default_connection(key_prefix=KEY_PREFIX):
-            with self._check.connection.get_managed_cursor(key_prefix=KEY_PREFIX) as cursor:
+        with self._check.connection.open_managed_default_connection(KEY_PREFIX):
+            with self._check.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                 db_names_formatted = ",".join(["'{}'".format(t) for t in database_names])
                 return execute_query(DB_QUERY.format(db_names_formatted), cursor, convert_results_to_str=True)
 
     @contextlib.contextmanager
     def _get_cursor(self, database_name):
-        with self._check.connection.open_managed_default_connection(key_prefix=KEY_PREFIX):
-            with self._check.connection.get_managed_cursor(key_prefix=KEY_PREFIX) as cursor:
+        with self._check.connection.open_managed_default_connection(KEY_PREFIX):
+            with self._check.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                 switch_db_statement = construct_use_statement(database_name)
                 cursor.execute(switch_db_statement)
                 query = self._get_tables_query()
@@ -167,8 +167,8 @@ def _map_row(self, database: DatabaseInfo, cursor_row) -> DatabaseObject:
         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=KEY_PREFIX_PRE_2017):
-                with self._check.connection.get_managed_cursor(key_prefix=KEY_PREFIX_PRE_2017) as cursor:
+            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"))

sqlserver/datadog_checks/sqlserver/sqlserver.py

@@ -124,6 +124,8 @@
 
 set_default_driver_conf()
 
+KEY_PREFIX = "dbm-sqlserver-"
+
 
 class SQLServer(DatabaseCheck):
     __NAMESPACE__ = "sqlserver"
@@ -411,8 +413,8 @@ def load_static_information(self):
         }
         missing_keys = expected_keys - set(self.static_info_cache.keys())
         if missing_keys:
-            with self.connection.open_managed_default_connection():
-                with self.connection.get_managed_cursor() as cursor:
+            with self.connection.open_managed_default_connection(KEY_PREFIX):
+                with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                     if STATIC_INFO_VERSION not in self.static_info_cache:
                         cursor.execute("select @@version")
                         results = cursor.fetchall()
@@ -533,8 +535,8 @@ def make_metric_list_to_collect(self):
                         self.log.warning("Database %s does not exist. Disabling checks for this instance.", context)
                         return
             if self.instance.get("stored_procedure") is None:
-                with self.connection.open_managed_default_connection():
-                    with self.connection.get_managed_cursor() as cursor:
+                with self.connection.open_managed_default_connection(KEY_PREFIX):
+                    with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                         self.autodiscover_databases(cursor)
                     self._make_metric_list_to_collect(self._config.custom_metrics)
         except SQLConnectionError:
@@ -754,7 +756,7 @@ def get_sql_counter_type(self, counter_name):
         cached = self._sql_counter_types.get(counter_name)
         if cached:
             return cached
-        with self.connection.get_managed_cursor() as cursor:
+        with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
             cursor.execute(COUNTER_TYPE_QUERY, (counter_name,))
             (sql_counter_type,) = cursor.fetchone()
             if sql_counter_type == PERF_LARGE_RAW_BASE:
@@ -832,8 +834,8 @@ def _check_connections_by_connecting_to_db(self):
                     self.log.warning("failed service check for auto discovered database: %s", e)
 
     def _check_connections_by_use_db(self):
-        with self.connection.open_managed_default_connection():
-            with self.connection.get_managed_cursor() as cursor:
+        with self.connection.open_managed_default_connection(KEY_PREFIX):
+            with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                 for db in self.databases:
                     check_err_message = "Database {} connection service check failed: {}"
                     try:
@@ -1030,34 +1032,34 @@ def load_basic_metrics(self, cursor):
 
     def collect_metrics(self):
         """Fetch the metrics from all the associated database tables."""
-        with self.connection.open_managed_default_connection():
+        with self.connection.open_managed_default_connection(KEY_PREFIX):
             if not self._config.only_custom_queries:
-                with self.connection.get_managed_cursor() as cursor:
+                with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                     self.load_basic_metrics(cursor)
 
             # Neither pyodbc nor adodbapi are able to read results of a query if the number of rows affected
             # statement are returned as part of the result set, so we disable for the entire connection
             # this is important mostly for custom_queries or the stored_procedure feature
             # https://docs.microsoft.com/en-us/sql/t-sql/statements/set-nocount-transact-sql
-            with self.connection.get_managed_cursor() as cursor:
+            with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                 cursor.execute("SET NOCOUNT ON")
             try:
                 if not self._config.only_custom_queries:
                     # restore the current database after executing dynamic queries
                     # this is to ensure the current database context is not changed
-                    with self.connection.restore_current_database_context():
+                    with self.connection.restore_current_database_context(KEY_PREFIX):
                         if self.database_metrics:
                             for database_metric in self.database_metrics:
                                 database_metric.execute()
 
                 # reuse the connection for custom queries
                 self._query_manager.execute()
             finally:
-                with self.connection.get_managed_cursor() as cursor:
+                with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
                     cursor.execute("SET NOCOUNT OFF")
 
     def execute_query_raw(self, query, db=None):
-        with self.connection.get_managed_cursor() as cursor:
+        with self.connection.get_managed_cursor(KEY_PREFIX) as cursor:
             if db:
                 ctx = construct_use_statement(db)
                 self.log.debug("changing cursor context via use statement: %s", ctx)

sqlserver/datadog_checks/sqlserver/statements.py

@@ -531,8 +531,8 @@ def collect_statement_metrics_and_plans(self):
 
         # re-use the check's conn module, but set extra_key=dbm- to ensure we get our own
         # raw connection. adodbapi and pyodbc modules are thread safe, but connections are not.
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 rows = self._collect_metrics_rows(cursor)
                 if not rows:
                     return

sqlserver/datadog_checks/sqlserver/stored_procedures.py

@@ -169,8 +169,8 @@ def collect_procedure_metrics(self):
         """
         # re-use the check's conn module, but set extra_key=dbm- to ensure we get our own
         # raw connection. adodbapi and pyodbc modules are thread safe, but connections are not.
-        with self._check.connection.open_managed_default_connection(key_prefix=self._conn_key_prefix):
-            with self._check.connection.get_managed_cursor(key_prefix=self._conn_key_prefix) as cursor:
+        with self._check.connection.open_managed_default_connection(self._conn_key_prefix):
+            with self._check.connection.get_managed_cursor(self._conn_key_prefix) as cursor:
                 rows = self._collect_metrics_rows(cursor)
                 if not rows:
                     self.log.debug("collect_procedure_metrics: no rows returned")