override parser messy but working

rialto/runner/config_loader.py

@@ -21,6 +21,7 @@
 from pydantic import BaseModel
 
 from rialto.common.utils import load_yaml
+from rialto.runner.config_overrides import override_config
 
 
 class IntervalConfig(BaseModel):
@@ -91,6 +92,11 @@ class PipelinesConfig(BaseModel):
     pipelines: list[PipelineConfig]
 
 
-def get_pipelines_config(path) -> PipelinesConfig:
+def get_pipelines_config(path: str, overrides: Dict) -> PipelinesConfig:
     """Load and parse yaml config"""
-    return PipelinesConfig(**load_yaml(path))
+    raw_config = load_yaml(path)
+    if overrides:
+        cfg = override_config(raw_config, overrides)
+        return PipelinesConfig(**cfg)
+    else:
+        return PipelinesConfig(**raw_config)

rialto/runner/config_overrides.py

@@ -0,0 +1,49 @@
+from typing import Dict
+
+from loguru import logger
+
+
+def _override(config, path, value) -> Dict:
+    key = path[0]
+    if "[" in key:
+        name = key.split("[")[0]
+        index = key.split("[")[1].replace("]", "")
+        if "=" in index:
+            index_key, index_value = index.split("=")
+            position = next(i for i, x in enumerate(config[name]) if x.get(index_key) == index_value)
+            if len(path) == 1:
+                config[name][position] = value
+            else:
+                config[name][position] = _override(config[name][position], path[1:], value)
+        else:
+            index = int(index)
+            if index >= 0:
+                if len(path) == 1:
+                    config[name][index] = value
+                else:
+                    config[name][index] = _override(config[name][index], path[1:], value)
+            else:
+                if len(path) == 1:
+                    config[name].append(value)
+                else:
+                    raise ValueError(f"Invalid index {index} for key {name} in path {path}")
+    else:
+        if len(path) == 1:
+            config[key] = value
+        else:
+            config[key] = _override(config[key], path[1:], value)
+    return config
+
+
+def override_config(config: Dict, overrides: Dict) -> Dict:
+    """Override config with user input
+
+    :param config: Config dictionary
+    :param overrides: Dictionary of overrides
+    :return: Overridden config
+    """
+    for path, value in overrides.items():
+        logger.info("Applying override: ", path, value)
+        config = _override(config, path.split("."), value)
+
+    return config

rialto/runner/runner.py

@@ -16,7 +16,7 @@
 
 import datetime
 from datetime import date
-from typing import List, Tuple
+from typing import Dict, List, Tuple
 
 import pyspark.sql.functions as F
 from loguru import logger
@@ -44,9 +44,10 @@ def __init__(
         rerun: bool = False,
         op: str = None,
         skip_dependencies: bool = False,
+        overrides: Dict = None,
     ):
         self.spark = spark
-        self.config = get_pipelines_config(config_path)
+        self.config = get_pipelines_config(config_path, overrides)
         self.reader = TableReader(spark)
 
         self.date_from = date_from

tests/runner/overrider.yaml

@@ -0,0 +1,96 @@
+#  Copyright 2022 ABSA Group Limited
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+runner:
+  watched_period_units: "months"
+  watched_period_value: 2
+  mail:
+    sender: [email protected]
+    smtp: server.test
+    to:
+      - [email protected]
+      - [email protected]
+    subject: test report
+pipelines:
+  - name: SimpleGroup
+    module:
+      python_module: tests.runner.transformations
+      python_class: SimpleGroup
+    schedule:
+      frequency: weekly
+      day: 7
+      info_date_shift:
+        - value: 3
+          units: days
+        - value: 2
+          units: weeks
+    dependencies:
+      - table: source.schema.dep1
+        interval:
+          units: "days"
+          value: 1
+        date_col: "DATE"
+      - table: source.schema.dep2
+        interval:
+          units: "months"
+          value: 3
+        date_col: "DATE"
+    target:
+      target_schema: catalog.schema
+      target_partition_column: "INFORMATION_DATE"
+    loader:
+      config_path: path/to/config.yaml
+      feature_schema: catalog.feature_tables
+      metadata_schema: catalog.metadata
+    metadata_manager:
+      metadata_schema: catalog.metadata
+  - name: GroupNoDeps
+    module:
+      python_module: tests.runner.transformations
+      python_class: SimpleGroup
+    schedule:
+      frequency: weekly
+      day: 7
+      info_date_shift:
+        value: 3
+        units: days
+  - name: NamedDeps
+    module:
+     python_module: tests.runner.transformations
+     python_class: SimpleGroup
+    schedule:
+     frequency: weekly
+     day: 7
+     info_date_shift:
+       value: 3
+       units: days
+    dependencies:
+     - table: source.schema.dep1
+       name: source1
+       interval:
+         units: "days"
+         value: 1
+       date_col: "DATE"
+     - table: source.schema.dep2
+       name: source2
+       interval:
+         units: "months"
+         value: 3
+       date_col: "batch"
+    target:
+      target_schema: catalog.schema
+      target_partition_column: "INFORMATION_DATE"
+    extras:
+      some_value: 3
+      some_other_value: cat

tests/runner/test_overrides.py

@@ -0,0 +1,41 @@
+from rialto.runner import Runner
+
+
+def test_overrides_simple(spark):
+    runner = Runner(
+        spark,
+        config_path="tests/runner/transformations/config.yaml",
+        run_date="2023-03-31",
+        overrides={"runner.mail.to": ["[email protected]", "[email protected]", "[email protected]"]},
+    )
+    assert runner.config.runner.mail.to == ["[email protected]", "[email protected]", "[email protected]"]
+
+
+def test_overrides_array_index(spark):
+    runner = Runner(
+        spark,
+        config_path="tests/runner/transformations/config.yaml",
+        run_date="2023-03-31",
+        overrides={"runner.mail.to[1]": "[email protected]"},
+    )
+    assert runner.config.runner.mail.to == ["[email protected]", "[email protected]"]
+
+
+def test_overrides_array_append(spark):
+    runner = Runner(
+        spark,
+        config_path="tests/runner/transformations/config.yaml",
+        run_date="2023-03-31",
+        overrides={"runner.mail.to[-1]": "test"},
+    )
+    assert runner.config.runner.mail.to == ["[email protected]", "[email protected]", "test"]
+
+
+def test_overrides_array_lookup(spark):
+    runner = Runner(
+        spark,
+        config_path="tests/runner/transformations/config.yaml",
+        run_date="2023-03-31",
+        overrides={"pipelines[name=SimpleGroup].target.target_schema": "new_schema"},
+    )
+    assert runner.config.pipelines[0].target.target_schema == "new_schema"