adding joint data loading flag (#79)

stanislav-bedunkevich · web-flow · commit 46f8588ae7fe · 2024-03-21T14:26:35.000+01:00
diff --git a/cli-e2e-test/rel/json_schema_mapping.rel b/cli-e2e-test/rel/json_schema_mapping.rel
@@ -9,12 +9,16 @@
 
 def JsonSourceConfigKey(k) { json_source_config(k, x...) from x... }
 
+@function
 def json_source:value[src in JsonSourceConfigKey] = source_catalog[src, _, :value]
+@function
 def json_source:child[src in JsonSourceConfigKey] = source_catalog[src, _, :child]
 def json_source:root[src in JsonSourceConfigKey] = source_catalog[src, _, :root]
 def json_source:array[src in JsonSourceConfigKey] = source_catalog[src, _, :array]
 
+@function
 def json_source:value[src in JsonSourceConfigKey] = simple_source_catalog[src, :value]
+@function
 def json_source:child[src in JsonSourceConfigKey] = simple_source_catalog[src, :child]
 def json_source:root[src in JsonSourceConfigKey] = simple_source_catalog[src, :root]
 def json_source:array[src in JsonSourceConfigKey] = simple_source_catalog[src, :array]
diff --git a/cli-e2e-test/test_e2e.py b/cli-e2e-test/test_e2e.py
@@ -40,6 +40,16 @@ def test_scenario1_model(self):
         self.assertNotEqual(rsp, 1)
         self.assert_output_dir_files(self.test_scenario1_model.__name__)
 
+    def test_scenario1_load_data_jointly(self):
+        # when
+        test_args = ["--batch-config", "./config/model/scenario1.json",
+                     "--end-date", "20220105",
+                     "--drop-db", "--load-data-jointly"]
+        rsp = call(self.cmd_with_common_arguments + test_args)
+        # then
+        self.assertNotEqual(rsp, 1)
+        self.assert_output_dir_files(self.test_scenario1_model.__name__)
+
     def test_scenario1_model_yaml(self):
         # when
         test_args = ["--batch-config", "./config/model/scenario1.yaml",
@@ -122,7 +132,7 @@ def test_scenario3_model_single_partition_change_for_date_partitioned(self):
         rsp_json = workflow.rai.execute_relation_json(self.logger, rai_config, self.env_config, RESOURCES_TO_DELETE_REL)
         self.assertEqual(rsp_json, [{'partition': 2023090800001, 'relation': 'city_data'}])
 
-    def test_scenario3_model_two_partitions_overriden_by_one_for_date_partitioned(self):
+    def test_scenario3_model_two_partitions_overridden_by_one_for_date_partitioned(self):
         # when
         test_args = ["--batch-config", "./config/model/scenario3.json",
                      "--start-date", "20230908",
diff --git a/cli/args.py b/cli/args.py
@@ -81,6 +81,12 @@ def parse() -> Namespace:
         action=BooleanOptionalAction,
         default=True
     )
+    parser.add_argument(
+        "--load-data-jointly",
+        help="When loading data, load all sources and partitions in one transaction",
+        action=BooleanOptionalAction,
+        default=False
+    )
     parser.add_argument(
         "--log-level",
         help="Set log level",
diff --git a/cli/runner.py b/cli/runner.py
@@ -51,7 +51,8 @@ def start(factories: dict[str, workflow.executor.WorkflowStepFactory] = MappingP
             workflow.constants.END_DATE: args.end_date,
             workflow.constants.FORCE_REIMPORT: args.force_reimport,
             workflow.constants.FORCE_REIMPORT_NOT_CHUNK_PARTITIONED: args.force_reimport_not_chunk_partitioned,
-            workflow.constants.COLLAPSE_PARTITIONS_ON_LOAD: args.collapse_partitions_on_load
+            workflow.constants.COLLAPSE_PARTITIONS_ON_LOAD: args.collapse_partitions_on_load,
+            workflow.constants.LOAD_DATA_JOINTLY: args.load_data_jointly
         }
         config = workflow.executor.WorkflowConfig(env_config, workflow.common.BatchConfig(args.batch_config_name,
                                                                                           batch_config_json),
diff --git a/test/test_cfg_src_step_factory.py b/test/test_cfg_src_step_factory.py
@@ -76,7 +76,8 @@ def _create_wf_cfg(env_config: EnvConfig, batch_config: BatchConfig) -> Workflow
         constants.END_DATE: "2021-01-01",
         constants.FORCE_REIMPORT: False,
         constants.FORCE_REIMPORT_NOT_CHUNK_PARTITIONED: False,
-        constants.COLLAPSE_PARTITIONS_ON_LOAD: False
+        constants.COLLAPSE_PARTITIONS_ON_LOAD: False,
+        constants.LOAD_DATA_JOINTLY: False
     }
     return WorkflowConfig(
         env=env_config,
diff --git a/workflow/__init__.py b/workflow/__init__.py
@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version_info__ = (0, 0, 44)
+__version_info__ = (0, 0, 45)
 __version__ = ".".join(map(str, __version_info__))
diff --git a/workflow/constants.py b/workflow/constants.py
@@ -80,6 +80,7 @@
 FORCE_REIMPORT = "force_reimport"
 FORCE_REIMPORT_NOT_CHUNK_PARTITIONED = "force_reimport_not_chunk_partitioned"
 COLLAPSE_PARTITIONS_ON_LOAD = "collapse_partitions_on_load"
+LOAD_DATA_JOINTLY = "load_data_jointly"
 
 # Snowflake constants
 
diff --git a/workflow/executor.py b/workflow/executor.py
@@ -1,10 +1,9 @@
+import asyncio
+import concurrent.futures
 import dataclasses
 import logging
 import subprocess
 import time
-import asyncio
-import concurrent.futures
-from workflow import snow
 from datetime import datetime
 from enum import Enum
 from itertools import groupby
@@ -14,9 +13,10 @@
 from more_itertools import peekable
 
 from workflow import query as q, paths, rai, constants
-from workflow.exception import StepTimeOutException, CommandExecutionException
+from workflow import snow
 from workflow.common import EnvConfig, RaiConfig, Source, BatchConfig, Export, FileType, ContainerType, Container, \
     FileMetadata
+from workflow.exception import StepTimeOutException, CommandExecutionException
 from workflow.manager import ResourceManager
 from workflow.utils import save_csv_output, format_duration, build_models, extract_date_range, build_relation_path, \
     get_common_model_relative_path, get_or_create_eventloop
@@ -327,10 +327,12 @@ def _parse_sources(step: dict, env_config: EnvConfig) -> List[Source]:
 
 class LoadDataWorkflowStep(WorkflowStep):
     collapse_partitions_on_load: bool
+    load_jointly: bool
 
-    def __init__(self, idt, name, type_value, state, timing, engine_size, collapse_partitions_on_load):
+    def __init__(self, idt, name, type_value, state, timing, engine_size, collapse_partitions_on_load, load_jointly):
         super().__init__(idt, name, type_value, state, timing, engine_size)
         self.collapse_partitions_on_load = collapse_partitions_on_load
+        self.load_jointly = load_jointly
 
     def _execute(self, logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig):
         rai.execute_query(logger, rai_config, env_config, q.DELETE_REFRESHED_SOURCES_DATA, readonly=False)
@@ -349,20 +351,22 @@ def _execute(self, logger: logging.Logger, env_config: EnvConfig, rai_config: Ra
             else:
                 simple_resources.append(src)
 
-        for src in simple_resources:
-            self._load_source(logger, env_config, rai_config, src)
+        self._load_simple_resources(logger, env_config, rai_config, simple_resources)
 
         if async_resources:
-            for src in async_resources:
-                self._load_source(logger, env_config, rai_config, src)
-            self.await_pending(env_config, logger, missed_resources)
+            self._load_async_resources(logger, env_config, rai_config, async_resources)
 
-    def await_pending(self, env_config, logger, missed_resources):
+    def _load_async_resources(self, logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig,
+                              async_resources) -> None:
+        for src in async_resources:
+            self._load_async_resource(logger, env_config, rai_config, src)
+        self._await_pending(env_config, logger, async_resources)
+
+    def _await_pending(self, env_config, logger, pending_resources):
         loop = get_or_create_eventloop()
         if loop.is_running():
             raise Exception('Waiting for resource would interrupt unexpected event loop - aborting to avoid confusion')
-        pending = [src for src in missed_resources if self._resource_is_async(src)]
-        pending_cos = [self._await_async_resource(logger, env_config, resource) for resource in pending]
+        pending_cos = [self._await_async_resource(logger, env_config, resource) for resource in pending_resources]
         loop.run_until_complete(asyncio.gather(*pending_cos))
 
     async def _await_async_resource(self, logger: logging.Logger, env_config: EnvConfig, src):
@@ -371,69 +375,103 @@ async def _await_async_resource(self, logger: logging.Logger, env_config: EnvCon
         if ContainerType.SNOWFLAKE == container.type:
             await snow.await_data_sync(logger, config, src["resources"])
 
-    def _load_source(self, logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig, src):
-        source_name = src["source"]
-        if 'is_date_partitioned' in src and src['is_date_partitioned'] == 'Y':
-            logger.info(f"Loading source '{source_name}' partitioned by date")
-            if self.collapse_partitions_on_load:
-                srcs = src["dates"]
-                first_date = srcs[0]["date"]
-                last_date = srcs[-1]["date"]
-
-                logger.info(
-                    f"Loading '{source_name}' all date partitions simultaneously, range {first_date} to {last_date}")
-
-                resources = []
-                for d in srcs:
-                    resources += d["resources"]
-                self._load_resource(logger, env_config, rai_config, resources, src)
-            else:
-                logger.info(f"Loading '{source_name}' one date partition at a time")
-                for d in src["dates"]:
-                    logger.info(f"Loading partition for date {d['date']}")
-
-                    for res in d["resources"]:
-                        self._load_resource(logger, env_config, rai_config, [res], src)
+    def _load_simple_resources(self, logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig,
+                               simple_resources) -> None:
+        # prepare queries for simple resources
+        query_batches = []
+        for src in simple_resources:
+            # now add all the items returned by _get_data_load_query` to the `query_batches` list
+            query_batches.extend(self._get_data_load_query(logger, env_config, src))
+
+        # execute queries for simple resources, if `load_jointly` is set to True then execute all queries in one txn
+        if self.load_jointly:
+            logger.info("Loading all CSV/JSON(L) sources jointly")
+            query = ""
+            inputs = {}
+            for query_with_input in query_batches:
+                query += query_with_input.query
+                inputs.update(query_with_input.inputs)
+            rai.execute_query(logger, rai_config, env_config, query, inputs, readonly=False)
         else:
-            logger.info(f"Loading source '{source_name}' not partitioned by date")
-            if self.collapse_partitions_on_load:
-                logger.info(f"Loading '{source_name}' all chunk partitions simultaneously")
-                self._load_resource(logger, env_config, rai_config, src["resources"], src)
-            else:
-                logger.info(f"Loading '{source_name}' one chunk partition at a time")
-                for res in src["resources"]:
-                    self._load_resource(logger, env_config, rai_config, [res], src)
-
-    @staticmethod
-    def _resource_is_async(src):
-        return True if ContainerType.SNOWFLAKE == ContainerType.from_source(src) else False
+            for query_with_input in query_batches:
+                rai.execute_query(logger, rai_config, env_config, query_with_input.query, query_with_input.inputs,
+                                  readonly=False)
 
-    @staticmethod
-    def _load_resource(logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig, resources, src) -> None:
+    def _get_data_load_query(self, logger: logging.Logger, env_config: EnvConfig, src) -> list:
         try:
             container = env_config.get_container(src["container"])
             config = EnvConfig.get_config(container)
-            if ContainerType.LOCAL == container.type or ContainerType.AZURE == container.type:
-                query_with_input = q.load_resources(logger, config, resources, src)
-                rai.execute_query(logger, rai_config, env_config, query_with_input.query, query_with_input.inputs,
-                                  readonly=False)
-            elif ContainerType.SNOWFLAKE == container.type:
-                snow.begin_data_sync(logger, config, rai_config, resources, src)
+            if 'is_date_partitioned' in src and src['is_date_partitioned'] == 'Y':
+                return self._get_date_part_load_query(logger, config, src)
+            else:
+                return self._get_simple_src_load_query(logger, config, src)
         except KeyError as e:
             logger.error(f"Unsupported file type: {src['file_type']}. Skip the source: {src}", e)
         except ValueError as e:
             logger.error(f"Unsupported source type. Skip the source: {src}", e)
+        return []  # return empty list if source is not supported
+
+    def _get_date_part_load_query(self, logger: logging.Logger, config, src):
+        source_name = src["source"]
+        logger.info(f"Loading source '{source_name}' partitioned by date")
+        if self.collapse_partitions_on_load:
+            srcs = src["dates"]
+            first_date = srcs[0]["date"]
+            last_date = srcs[-1]["date"]
+
+            logger.info(
+                f"Loading '{source_name}' all date partitions simultaneously, range {first_date} to {last_date}")
+
+            resources = []
+            for d in srcs:
+                resources += d["resources"]
+            return [q.load_resources(logger, config, resources, src)]
+        else:
+            logger.info(f"Loading '{source_name}' one date partition at a time")
+            batch = []
+            for d in src["dates"]:
+                logger.info(f"Loading partition for date {d['date']}")
+
+                for res in d["resources"]:
+                    batch.append(q.load_resources(logger, config, [res], src))
+            return batch
+
+    def _get_simple_src_load_query(self, logger: logging.Logger, config, src):
+        source_name = src["source"]
+        logger.info(f"Loading source '{source_name}' not partitioned by date")
+        if self.collapse_partitions_on_load:
+            logger.info(f"Loading '{source_name}' all chunk partitions simultaneously")
+            return [q.load_resources(logger, config, src["resources"], src)]
+        else:
+            logger.info(f"Loading '{source_name}' one chunk partition at a time")
+            batch = []
+            for res in src["resources"]:
+                batch.append(q.load_resources(logger, config, [res], src))
+            return batch
+
+    @staticmethod
+    def _resource_is_async(src):
+        return True if ContainerType.SNOWFLAKE == ContainerType.from_source(src) else False
+
+    @staticmethod
+    def _load_async_resource(logger: logging.Logger, env_config: EnvConfig, rai_config: RaiConfig, resources, src) ->\
+            None:
+        container = env_config.get_container(src["container"])
+        config = EnvConfig.get_config(container)
+        snow.begin_data_sync(logger, config, rai_config, resources, src)
 
 
 class LoadDataWorkflowStepFactory(WorkflowStepFactory):
 
     def _required_params(self, config: WorkflowConfig) -> List[str]:
-        return [constants.COLLAPSE_PARTITIONS_ON_LOAD]
+        return [constants.COLLAPSE_PARTITIONS_ON_LOAD, constants.LOAD_DATA_JOINTLY]
 
     def _get_step(self, logger: logging.Logger, config: WorkflowConfig, idt, name, type_value, state, timing,
                   engine_size, step: dict) -> WorkflowStep:
         collapse_partitions_on_load = config.step_params[constants.COLLAPSE_PARTITIONS_ON_LOAD]
-        return LoadDataWorkflowStep(idt, name, type_value, state, timing, engine_size, collapse_partitions_on_load)
+        load_jointly = config.step_params[constants.LOAD_DATA_JOINTLY]
+        return LoadDataWorkflowStep(idt, name, type_value, state, timing, engine_size, collapse_partitions_on_load,
+                                    load_jointly)
 
 
 class MaterializeWorkflowStep(WorkflowStep):
diff --git a/workflow/query.py b/workflow/query.py
@@ -307,10 +307,10 @@ def _local_load_multipart_query(rel_name: str, file_type: FileType, parts) -> Qu
     load_config = _multi_part_load_config_query(rel_name, file_type,
                                                 _local_multipart_config_integration(raw_data_rel_name))
 
-    query = f"{_part_index_relation(part_indexes)}\n" \
+    query = f"{_part_index_relation(rel_name, part_indexes)}\n" \
             f"{raw_text}\n" \
             f"{load_config}\n" \
-            f"{insert_text}"
+            f"{insert_text}\n"
 
     return QueryWithInputs(query, inputs)
 
@@ -330,10 +330,10 @@ def _azure_load_multipart_query(rel_name: str, file_type: FileType, parts, confi
     load_config = _multi_part_load_config_query(rel_name, file_type,
                                                 _azure_multipart_config_integration(path_rel_name, config))
 
-    return f"{_part_index_relation(part_indexes)}\n" \
+    return f"{_part_index_relation(rel_name, part_indexes)}\n" \
            f"{_path_rel_name_relation(path_rel_name, part_uri_map)}\n" \
            f"{load_config}\n" \
-           f"{insert_text}"
+           f"{insert_text}\n"
 
 
 def _multi_part_load_config_query(rel_name: str, file_type: FileType, config_integration: str) -> str:
@@ -345,7 +345,7 @@ def _multi_part_load_config_query(rel_name: str, file_type: FileType, config_int
     return f"""
 bound {IMPORT_CONFIG_REL}:{rel_name}:schema
 bound {IMPORT_CONFIG_REL}:{rel_name}:syntax:header
-module {_config_rel_name(rel_name)}[i in part_indexes]
+module {_config_rel_name(rel_name)}[i in part_indexes:{rel_name}]
     {schema}
     {config_integration}
 end
@@ -383,25 +383,25 @@ def _indexed_literal(raw_data_rel_name: str, index: int) -> str:
 def _config_rel_name(rel: str) -> str: return f"load_{rel}_config"
 
 
-def _part_index_relation(part_index: str) -> str:
-    return f"def part_index_config:schema:INDEX = \"int\"\n" \
-           f"def part_index_config:data = \"\"\"\n" \
+def _part_index_relation(rel_name: str, part_index: str) -> str:
+    return f"def part_index_config:{rel_name}:schema:INDEX = \"int\"\n" \
+           f"def part_index_config:{rel_name}:data = \"\"\"\n" \
            f"INDEX\n" \
            f"{part_index}\n" \
            f"\"\"\"" \
-           f"def part_indexes_csv = load_csv[part_index_config]\n" \
-           f"def part_indexes = part_indexes_csv:INDEX[_]"
+           f"def part_indexes_csv:{rel_name} = load_csv[part_index_config:{rel_name}]\n" \
+           f"def part_indexes:{rel_name} = part_indexes_csv:{rel_name}:INDEX[_]"
 
 
 def _path_rel_name_relation(path_rel_name: str, part_uri_map: str) -> str:
-    return f"def part_uri_map_config:schema:INDEX = \"int\"\n" \
-           f"def part_uri_map_config:schema:URI = \"string\"\n" \
-           f"def part_uri_map_config:data = \"\"\"\n" \
+    return f"def part_uri_map_config:{path_rel_name}:schema:INDEX = \"int\"\n" \
+           f"def part_uri_map_config:{path_rel_name}:schema:URI = \"string\"\n" \
+           f"def part_uri_map_config:{path_rel_name}:data = \"\"\"\n" \
            f"INDEX,URI\n" \
            f"{part_uri_map}\n" \
            f"\"\"\"" \
-           f"def part_uri_map_csv = load_csv[part_uri_map_config]\n" \
-           f"def {path_rel_name}(i, u) {{ part_uri_map_csv:INDEX(row, i) and part_uri_map_csv:URI(row, u) from row }}"
+           f"def part_uri_map_csv:{path_rel_name} = load_csv[part_uri_map_config:{path_rel_name}]\n" \
+           f"def {path_rel_name}(i, u) {{ part_uri_map_csv:{path_rel_name}:INDEX(row, i) and part_uri_map_csv:{path_rel_name}:URI(row, u) from row }}"
 
 
 def _export_relation_as_csv_local(rel_name) -> str:
