feat: fix storage_org

weni-ai · May 24, 2024 · 6b9374e · 6b9374e
1 parent 3841190
commit 6b9374e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/run.py b/run.py
@@ -17,7 +17,7 @@
 
 
 from flowrun.storage.elasticsearch import FlowRunElasticSearch
-from flowrun.storage.postgresql import FlowRunPostgreSQL
+from flowrun.storage.postgresql import FlowRunPostgreSQL, OrgPostgreSQL
 from flowrun.transformer import (
     bulk_flowrun_sql_to_elasticsearch_transformer,
     flowrun_sql_to_elasticsearch_transformer,
@@ -27,14 +27,15 @@
 
 FlowRunPGtoES = ObjectETLProcessor(
     object_transformer=flowrun_sql_to_elasticsearch_transformer,
-    storage_from=FlowRunPostgreSQL(),
     storage_to=FlowRunElasticSearch(),
+    storage_from=FlowRunPostgreSQL(),
 )
 
 BulkFlowRunPGtoES = BulkObjectETLProcessor(
     object_transformer=bulk_flowrun_sql_to_elasticsearch_transformer,
-    storage_from=FlowRunPostgreSQL(),
     storage_to=FlowRunElasticSearch(),
+    storage_from=FlowRunPostgreSQL(),
+    storage_org=OrgPostgreSQL(),
 )
 
 

diff --git a/shared/processors.py b/shared/processors.py
@@ -36,12 +36,16 @@ def execute(self, object_identifier: str):
 
 class BulkObjectETLProcessor:
     def __init__(
-        self, object_transformer: Callable, storage_from: object, storage_to: object
+        self,
+        object_transformer: Callable,
+        storage_from: object,
+        storage_to: object,
+        storage_org: object,
     ) -> None:
         self.object_transformer = object_transformer
         self.storage_from = storage_from
         self.storage_to = storage_to
-        self.storage_org
+        self.storage_org = storage_org
 
     def execute(self):
         orgs = self.storage_org.list_active()