From c9b1ddfd1514adc3f6c9a6ee34aac2be9c50bbf5 Mon Sep 17 00:00:00 2001
From: vitor-aignosi <vitor.santos@aignosi.com.br>
Date: Mon, 17 Nov 2025 10:38:44 -0300
Subject: [PATCH] SIENTIAPDE-1273  Refactor data export in Drift,
 SimpleMetrics, and FormatAndExportPrediction workflows to improve data
 handling. Changed the export format from a list of dictionaries to direct
 objects for enhanced clarity and compatibility with downstream systems.

---
 laborious/workflows/drift.py                                  | 1 +
 laborious/workflows/simple_metrics.py                         | 3 ++-
 .../workflows/sub_workflows/format_and_export_prediction.py   | 4 ++--
 3 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/laborious/workflows/drift.py b/laborious/workflows/drift.py
index 60bb274..d556dac 100644
--- a/laborious/workflows/drift.py
+++ b/laborious/workflows/drift.py
@@ -52,6 +52,7 @@ class Drift:
                 **metadata,
                 'query': gathering_query,
                 'datetime_columns': ['timestamp', 'created_at'],
+                'orient': 'records',
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=300),
diff --git a/laborious/workflows/simple_metrics.py b/laborious/workflows/simple_metrics.py
index f001962..ce950af 100644
--- a/laborious/workflows/simple_metrics.py
+++ b/laborious/workflows/simple_metrics.py
@@ -53,6 +53,7 @@ class SimpleMetrics:
                 **metadata,
                 'query': query,
                 'datetime_columns': ['timestamp'],
+                'orient': 'records',
             },
             retry_policy=retry_policy,
             start_to_close_timeout=timedelta(seconds=300),
@@ -82,7 +83,7 @@ class SimpleMetrics:
             Activities.export_data_to_postgres,
             {
                 **metadata,
-                'data': simple_metrics.to_dict(orient='records'),
+                'data': simple_metrics,
                 'schema': input_data['schema'],
                 'table_name': input_data['target_table_name'],
                 'timestamp_conversion': {
diff --git a/laborious/workflows/sub_workflows/format_and_export_prediction.py b/laborious/workflows/sub_workflows/format_and_export_prediction.py
index 1d59e7a..4280ccf 100644
--- a/laborious/workflows/sub_workflows/format_and_export_prediction.py
+++ b/laborious/workflows/sub_workflows/format_and_export_prediction.py
@@ -107,7 +107,7 @@ class FormatAndExportPrediction:
                         **metadata,
                         'schema': input_data['schema'],
                         'table_name': input_data['transform_table_name'],
-                        'data': transformed.to_dict(orient='records'),
+                        'data': transformed,
                         'timestamp_conversion': {
                             'column': 'timestamp',
                             'format': DATETIME_FORMAT_WITH_TZ,
@@ -156,7 +156,7 @@ class FormatAndExportPrediction:
                 **metadata,
                 'schema': input_data['schema'],
                 'table_name': input_data['table_name'],
-                'data': prediction.to_dict(orient='records'),
+                'data': prediction,
                 'timestamp_conversion': {'column': 'timestamp', 'format': DATETIME_FORMAT_WITH_TZ},
             },
             retry_policy=retry_policy,