Process query rows one at a time to reduce memory footprint (#15268)

* Process query rows one at a time to reduce memory footprint * Remove `fetchall` from mock
DataDog · Jul 17, 2023 · e97c04d · e97c04d
1 parent 31c4aa7
commit e97c04d
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git a/snowflake/datadog_checks/snowflake/check.py b/snowflake/datadog_checks/snowflake/check.py
@@ -155,8 +155,9 @@ def execute_query_raw(self, query):
 
             if cursor.rowcount is None or cursor.rowcount < 1:
                 self.log.debug("Failed to fetch records from query: `%s`", query)
-                return []
-            return cursor.fetchall()
+                return
+            # Iterating on the cursor provides one row at a time without loading all of them at once
+            yield from cursor
 
     def connect(self):
         self.log.debug(
@@ -209,8 +210,8 @@ def connect(self):
     @AgentCheck.metadata_entrypoint
     def _collect_version(self):
         try:
-            raw_version = self.execute_query_raw("select current_version();")
-            version = raw_version[0][0]
+            raw_version = next(self.execute_query_raw("select current_version();"))
+            version = raw_version[0]
         except Exception as e:
             self.log.error("Error collecting version for Snowflake: %s", e)
         else:

diff --git a/snowflake/tests/snowflake_connector_patch/_snowflake_connector_patch/connector.py b/snowflake/tests/snowflake_connector_patch/_snowflake_connector_patch/connector.py
@@ -4,6 +4,7 @@
 import re
 
 import requests
+from snowflake.connector.cursor import SnowflakeCursor
 
 from . import tables
 
@@ -49,14 +50,20 @@ def execute(self, query):
             if self.schema == 'ORGANIZATION_USAGE':
                 table_prefix = 'ORGANIZATION_'
             table_attr = "{}{}".format(table_prefix, table_name)
-            self.__data = getattr(tables, table_attr, [])
+            self.__data = list(getattr(tables, table_attr, []))
         elif query == 'select current_version();':
             self.__data = [('4.30.2',)]
         else:
             self.__data = []
 
-    def fetchall(self):
-        return self.__data
+    def fetchone(self):
+        try:
+            return self.__data.pop(0)
+        except IndexError:
+            return None
+
+    def __iter__(self):
+        return SnowflakeCursor.__iter__(self)
 
     def close(self):
-        pass
+        self.__data = []
diff --git a/snowflake/tests/test_snowflake.py b/snowflake/tests/test_snowflake.py
@@ -123,7 +123,7 @@ def test_version_metadata(dd_run_check, instance, datadog_agent):
         'version.raw': '4.30.2',
         'version.scheme': 'semver',
     }
-    with mock.patch('datadog_checks.snowflake.SnowflakeCheck.execute_query_raw', return_value=expected_version):
+    with mock.patch('datadog_checks.snowflake.SnowflakeCheck.execute_query_raw', return_value=iter(expected_version)):
         check = SnowflakeCheck(CHECK_NAME, {}, [instance])
         check.check_id = 'test:123'
         check._conn = mock.MagicMock()