AAP-62673: library: collect to dataframe (#327)

* library/collector: make copy_table return DataFrame

drop the output_dir logic from copy_table, no longer returns file list,
but a pandas DataFrame

* add dataframe_to_csv_files helper

so that a collector can still be made to save to csv and go through the csvfilesplitter

* load_anonymized_rollup_data: expect dataframe list instead of file list

but also accept one dataframe, and a mixed list of filenames and dataframes

this is probably going away completely once we move the logic to tasks

* copy_table: pd.read_sql needs sqlalchemy, do the same manually with the psycopg connection instead

* EE: handle both t/f bools and True/False bools

* address cursor review - do not create empty csv files

have dataframe_to_csv_file_list only create files for non-empty
dataframes or dataframes with columns

* tests: update collector tests to mock copy_table returning DataFrames

Mock copy_table to return pd.DataFrame instead of file paths, matching
the new interface. Remove pointless assertions that only verified the
mock return value length.

Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>

* test_gather_jobs..: handle different data formats between expected csv and returned dataframe

the postgres to dataframe conversion outputs timezones differently (+00 vs +00:00),
bools (t vs True), and nullable-int fields (as float)

DataFrame.convert_dtypes() in copy_table might be a solution, at about 40% perf cost, or knowing the types
IMO the collector is the place that knows the formats and might handle that best by itself

for now, just updating the test to expect the data it gets

---------

Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>

Author: himdel

metrics_utility/anonymized_rollups/anonymized_rollups.py

@@ -479,18 +479,23 @@ def compute_anonymized_rollup_from_raw_data(input_data, salt):
     return anonymized_rollup
 
 
-# loads data from tarballs located in base_path/data/year/month/day/*{collector_name}*.tar.gz
-# inside tarball is file named {collector_name}.csv
-# this goes to dataframe, then filter_function is applied to the dataframe
+# loads data from a list of dataframes
+# then filter_function is applied to the dataframe
 # all result dataframes are concatenated into one dataframe
-def load_anonymized_rollup_data(rollup_object: BaseAnonymizedRollup, file_list: []):
-    # file_list - list of csv files that needs to be read
+def load_anonymized_rollup_data(rollup_object: BaseAnonymizedRollup, dataframe_list):
+    # compat for one dataframe
+    if isinstance(dataframe_list, pd.DataFrame):
+        prepared_data = rollup_object.prepare(dataframe_list)
+        return rollup_object.merge(None, prepared_data)
 
     concat_data = None
 
-    for file in file_list:
-        df = pd.read_csv(file, encoding='utf-8')
-        prepared_data = rollup_object.prepare(df)
+    for dataframe in dataframe_list:
+        # compat for CSVs
+        if isinstance(dataframe, str):
+            dataframe = pd.read_csv(dataframe, encoding='utf-8')
+
+        prepared_data = rollup_object.prepare(dataframe)
         concat_data = rollup_object.merge(concat_data, prepared_data)
 
     return concat_data

metrics_utility/anonymized_rollups/execution_environments_anonymized_rollup.py

@@ -26,7 +26,7 @@ def base(self, dataframe):
             }
 
         execution_environments_total = int(len(dataframe))
-        dataframe['managed'] = dataframe['managed'].map({'t': True, 'f': False})
+        dataframe['managed'] = dataframe['managed'].map({'t': True, 'f': False, True: True, False: False})
         execution_environments_default_total = int(dataframe['managed'].sum())
         execution_environments_custom_total = execution_environments_total - execution_environments_default_total
 

metrics_utility/library/README.md

@@ -11,8 +11,8 @@ It provides an abstraction over collectors, packaging and storage, extraction, r
 
 Collector is python function which accepts params, gathers data, and returns it in one of the supported formats.
 
-It either returns a python dict, which gets serialized into JSON,
-or a list of filenames of temporary files it created.
+It either returns a python dict (for snapshot collectors like config),
+or a pandas DataFrame (for SQL-based collectors).
 
 It's exported decorated to wrap calls into BaseCollector subclass instances, so that param passing can happen separately from .gather().
 The wrapper ensures that any calls to `my_collector(db=connection).gather()` do the same thing as an undecorated `my_collector(db=connection)` - this is so that initialization can happen before db locks are acquired.
@@ -28,18 +28,30 @@ Currently supported:
 
 Controller collectors (in `metrics_utility.library.collectors.controller`):
 * `config(db, billing_provider_params).gather() -> Dict`
-* `execution_environments(db, [output_dir]).gather() -> [filenames]`
-* `job_host_summary(db, since, until, [output_dir]).gather() -> [filenames]`
-* `job_host_summary_service(db, since, until, [output_dir]).gather() -> [filenames]`
-* `main_host(db, [output_dir]).gather() -> [filenames]`
-* `main_indirectmanagednodeaudit(db, since, until, [output_dir]).gather() -> [filenames]`
-* `main_jobevent(db, since, until, [output_dir]).gather() -> [filenames]`
-* `main_jobevent_service(db, since, until, [output_dir]).gather() -> [filenames]`
-* `unified_jobs(db, since, until, [output_dir]).gather() -> [filenames]`
+* `execution_environments(db).gather() -> DataFrame`
+* `job_host_summary(db, since, until).gather() -> DataFrame`
+* `job_host_summary_service(db, since, until).gather() -> DataFrame`
+* `main_host(db).gather() -> DataFrame`
+* `main_host_daily(db, since, until).gather() -> DataFrame`
+* `main_indirectmanagednodeaudit(db, since, until).gather() -> DataFrame`
+* `main_jobevent(db, since, until).gather() -> DataFrame`
+* `main_jobevent_service(db, since, until).gather() -> DataFrame`
+* `unified_jobs(db, since, until).gather() -> DataFrame`
 
 Other collectors (in `metrics_utility.library.collectors.others`):
 * `total_workers_vcpu(cluster_name, metering_enabled, prometheus_url, ca_cert_path, token) -> Dict`
 
+For CLI usage or when CSV files are needed, use the `dataframe_to_csv_files()` helper from `metrics_utility.library.csv_utils`:
+
+```python
+from metrics_utility.library.csv_utils import dataframe_to_csv_files
+
+df = execution_environments(db=db).gather()
+csv_files = dataframe_to_csv_files(df, 'main_executionenvironment', '/tmp/output')
+# Returns: ['/tmp/output/main_executionenvironment_table.csv']
+# or ['.._split0.csv', '.._split1.csv', ...] for large datasets
+```
+
 
 #### Package
 

metrics_utility/library/collectors/controller/config.py

@@ -96,10 +96,10 @@ def _get_install_type():
     return 'traditional'
 
 
+# FIXME: psycopg.sql
 def _get_controller_settings(db, keys):
     settings = {}
     with db.cursor() as cursor:
-        # FIXME: psycopg.sql ?
         in_sql = "'" + "', '".join(keys) + "'"
         cursor.execute(f'SELECT key, value FROM conf_setting WHERE key IN ({in_sql})')
         for key, value in cursor.fetchall():

metrics_utility/library/collectors/controller/credentials_service.py

@@ -2,7 +2,7 @@
 
 
 @collector
-def credentials_service(*, db=None, since=None, until=None, output_dir=None):
+def credentials_service(*, db=None, since=None, until=None):
     query = f"""
         SELECT
             main_credentialtype.name as credential_type,
@@ -19,4 +19,4 @@ def credentials_service(*, db=None, since=None, until=None, output_dir=None):
         ORDER BY main_unifiedjob.id ASC, main_credentialtype.name ASC
     """
 
-    return copy_table(db=db, table='credentials', query=query, output_dir=output_dir)
+    return copy_table(db=db, query=query)

metrics_utility/library/collectors/controller/execution_environments.py

@@ -2,7 +2,7 @@
 
 
 @collector
-def execution_environments(*, db=None, output_dir=None):
+def execution_environments(*, db=None):
     query = """
         SELECT
             id,
@@ -20,4 +20,4 @@ def execution_environments(*, db=None, output_dir=None):
         FROM main_executionenvironment
     """
 
-    return copy_table(db=db, table='main_executionenvironment', query=query, output_dir=output_dir)
+    return copy_table(db=db, query=query)

metrics_utility/library/collectors/controller/job_host_summary.py

@@ -2,7 +2,7 @@
 
 
 @collector
-def job_host_summary(*, db=None, since=None, until=None, output_dir=None):
+def job_host_summary(*, db=None, since=None, until=None):
     where = ' AND '.join(
         [
             f"main_jobhostsummary.modified >= '{since.isoformat()}'",
@@ -78,4 +78,4 @@ def job_host_summary(*, db=None, since=None, until=None, output_dir=None):
         ORDER BY main_jobhostsummary.modified ASC
     """
 
-    return copy_table(db=db, table='main_jobhostsummary', query=query, prepend_query=True, output_dir=output_dir)
+    return copy_table(db=db, query=query, prepend_query=True)

metrics_utility/library/collectors/controller/job_host_summary_service.py

@@ -2,7 +2,7 @@
 
 
 @collector
-def job_host_summary_service(*, db=None, since=None, until=None, output_dir=None):
+def job_host_summary_service(*, db=None, since=None, until=None):
     where = ' AND '.join(
         [
             f"mu.finished >= '{since.isoformat()}'",
@@ -85,4 +85,4 @@ def job_host_summary_service(*, db=None, since=None, until=None, output_dir=None
         ORDER BY mu.finished ASC
     """
 
-    return copy_table(db=db, table='main_jobhostsummary', query=query, prepend_query=True, output_dir=output_dir)
+    return copy_table(db=db, query=query, prepend_query=True)

metrics_utility/library/collectors/controller/main_host.py

@@ -87,13 +87,13 @@ def _main_host_query(where):
 
 
 @collector
-def main_host(*, db=None, output_dir=None):
+def main_host(*, db=None):
     query = _main_host_query("enabled='t'")
-    return copy_table(db=db, table='main_host', query=query, prepend_query=True, output_dir=output_dir)
+    return copy_table(db=db, query=query, prepend_query=True)
 
 
 @collector
-def main_host_daily(*, db=None, since=None, until=None, output_dir=None):
+def main_host_daily(*, db=None, since=None, until=None):
     # prefer running with until=False, to not skip hosts that keep being modified
 
     where = f"""
@@ -102,4 +102,4 @@ def main_host_daily(*, db=None, since=None, until=None, output_dir=None):
         OR {date_where('main_host.modified', since, until)})
     """
     query = _main_host_query(where)
-    return copy_table(db=db, table='main_host_daily', query=query, prepend_query=True, output_dir=output_dir)
+    return copy_table(db=db, query=query, prepend_query=True)

metrics_utility/library/collectors/controller/main_indirectmanagednodeaudit.py

@@ -2,7 +2,7 @@
 
 
 @collector
-def main_indirectmanagednodeaudit(*, db=None, since=None, until=None, output_dir=None):
+def main_indirectmanagednodeaudit(*, db=None, since=None, until=None):
     where = ' AND '.join(
         [
             f"main_indirectmanagednodeaudit.created >= '{since.isoformat()}'",
@@ -40,4 +40,4 @@ def main_indirectmanagednodeaudit(*, db=None, since=None, until=None, output_dir
         ORDER BY main_indirectmanagednodeaudit.created ASC
     """
 
-    return copy_table(db=db, table='main_indirectmanagednodeaudit', query=query, output_dir=output_dir)
+    return copy_table(db=db, query=query)