datadog_checks_base: Fix obfuscate_sql_with_metadata wrapper not handling json.loads() edge case

datadog_checks_base/datadog_checks/base/utils/db/utils.py

@@ -182,19 +182,22 @@ def default_json_event_encoding(o):
 
 
 def obfuscate_sql_with_metadata(query, options=None):
-    try:
-        statement = datadog_agent.obfuscate_sql(query, options)
-        statement_with_metadata = json.loads(statement)
-    except ValueError:  # Backwards compatible exception between py38 and py27 for json.loads()
-        # Assume we're running against an older agent and return the obfuscated query without metadata.
-        return {'query': statement, 'metadata': {}}
-    except Exception as e:
-        raise e
-    metadata = statement_with_metadata.get('metadata', {})
-    tables = metadata.pop('tables_csv', None)
-    tables = [table.strip() for table in tables.split(',') if table != ''] if tables else None
-    statement_with_metadata['metadata']['tables'] = tables
-    return statement_with_metadata
+    def _load_metadata(statement):
+        try:
+            statement_with_metadata = json.loads(statement)
+            metadata = statement_with_metadata.get('metadata', {})
+            tables = metadata.pop('tables_csv', None)
+            tables = [table.strip() for table in tables.split(',') if table != ''] if tables else None
+            statement_with_metadata['metadata']['tables'] = tables
+            return statement_with_metadata
+        except ValueError:
+            # Assume we're running against an older agent and return the obfuscated query without metadata.
+            return {'query': statement, 'metadata': {}}
+
+    obfuscated_statement = datadog_agent.obfuscate_sql(query, options)
+    if options and json.loads(options).get('return_json_metadata', False):
+        return _load_metadata(obfuscated_statement)
+    return {'query': obfuscated_statement, 'metadata': {}}
 
 
 class DBMAsyncJob(object):

datadog_checks_base/tests/base/utils/db/test_util.py

@@ -78,7 +78,7 @@ class TestDBExcepption(BaseException):
 
 
 @pytest.mark.parametrize(
-    "obfusactor_return_value,expected_value",
+    "obfuscator_return_value,expected_value,return_json_metadata",
     [
         (
             json.dumps(
@@ -91,6 +91,7 @@ class TestDBExcepption(BaseException):
                 'query': 'SELECT * FROM datadog',
                 'metadata': {'commands': ['SELECT'], 'comments': None, 'tables': ['datadog']},
             },
+            True,
         ),
         (
             json.dumps(
@@ -111,6 +112,7 @@ class TestDBExcepption(BaseException):
                     'tables': ['datadog', 'datadog2'],
                 },
             },
+            True,
         ),
         (
             json.dumps(
@@ -123,23 +125,38 @@ class TestDBExcepption(BaseException):
                 'query': 'COMMIT',
                 'metadata': {'commands': ['COMMIT'], 'comments': None, 'tables': None},
             },
+            True,
         ),
         (
             'SELECT * FROM datadog',
             {
                 'query': 'SELECT * FROM datadog',
                 'metadata': {},
             },
+            # This test ensures that we handle the failed json decoder when requesting for metadata.
+            # This scenario could happen when newer integrations run against older agents which might not have the
+            # metadata API.
+            True,
+        ),
+        (
+            'SELECT * FROM datadog',
+            {
+                'query': 'SELECT * FROM datadog',
+                'metadata': {},
+            },
+            False,
         ),
     ],
 )
-def test_obfuscate_sql_with_metadata(obfusactor_return_value, expected_value):
+def test_obfuscate_sql_with_metadata(obfuscator_return_value, expected_value, return_json_metadata):
     def _mock_obfuscate_sql(query, options=None):
-        return obfusactor_return_value
+        return obfuscator_return_value
 
     with mock.patch.object(datadog_agent, 'obfuscate_sql', passthrough=True) as mock_agent:
         mock_agent.side_effect = _mock_obfuscate_sql
-        statement = obfuscate_sql_with_metadata('query here does not matter')
+        statement = obfuscate_sql_with_metadata(
+            'query here does not matter', json.dumps({'return_json_metadata': return_json_metadata})
+        )
         assert statement == expected_value