From 7fe8c6320f7193b506f84997524f31d4fc55f841 Mon Sep 17 00:00:00 2001 From: fpin Date: Sun, 26 Mar 2023 10:26:31 +0200 Subject: [PATCH 1/3] refactor: Simplify code for loading optional methods in action files --- karadoc/common/job_core/has_before_after.py | 29 +++--- karadoc/common/job_core/load.py | 88 ++++--------------- karadoc/common/job_core/package.py | 51 +++++++++++ karadoc/spark/job_core/has_external_inputs.py | 60 ++++--------- .../spark/job_core/has_external_outputs.py | 55 ++++-------- .../job_core/has_stream_external_inputs.py | 52 ++++------- .../job_core/has_stream_external_output.py | 30 ++++--- .../spark/job_core/test_has_batch_output.py | 7 -- .../has_inputs/QUALITY_CHECK.py | 2 +- 9 files changed, 147 insertions(+), 227 deletions(-) create mode 100644 karadoc/common/job_core/package.py diff --git a/karadoc/common/job_core/has_before_after.py b/karadoc/common/job_core/has_before_after.py index 831b6e7..e755d2d 100644 --- a/karadoc/common/job_core/has_before_after.py +++ b/karadoc/common/job_core/has_before_after.py @@ -1,39 +1,40 @@ +from karadoc.common.job_core.package import OptionalMethod + + class HasBeforeAfter: - def __init__(self) -> None: - # Attributes that the user is not supposed to change - self.__before_each = None - self.__after_each = None - self.__before_all = None - self.__after_all = None + def __init__(self): + def default() -> None: + pass + + self.before_each = OptionalMethod(default, "before_each") + self.after_each = OptionalMethod(default, "after_each") + self.before_all = OptionalMethod(default, "before_all") + self.after_all = OptionalMethod(default, "after_all") def before_each(self) -> None: """This method will be called before executing each action method defined in the job's definition file. This is similar to unittest's `setup()` method. """ - if self.__before_each is not None: - return self.__before_each() + pass def after_each(self) -> None: """This method will be called after executing each action method defined in the job's definition file. This is similar to unittest's `teardown()` method. """ - if self.__after_each is not None: - return self.__after_each() + pass def before_all(self) -> None: """This method will be called once before executing all the action methods defined in the job's definition file. This is similar to unittest's `setupClass()` method. """ - if self.__before_all is not None: - return self.__before_all() + pass def after_all(self) -> None: """This method will be called once after executing all the action methods defined in the job's definition file. This is similar to unittest's `teardownClass()` method. """ - if self.__after_all is not None: - return self.__after_all() + pass diff --git a/karadoc/common/job_core/load.py b/karadoc/common/job_core/load.py index 6b84683..41b9872 100644 --- a/karadoc/common/job_core/load.py +++ b/karadoc/common/job_core/load.py @@ -4,36 +4,17 @@ from os.path import isfile, splitext from pathlib import Path from types import FunctionType, ModuleType -from typing import Any, Callable, Dict, Optional, Type, TypeVar, Union, cast +from typing import Dict, Optional, Type, TypeVar, Union, cast from karadoc.common import conf from karadoc.common.exceptions import ActionFileLoadingError, ForbiddenActionError -from karadoc.common.job_core.has_before_after import HasBeforeAfter from karadoc.common.job_core.has_vars import HasVars from karadoc.common.job_core.job_base import JobBase +from karadoc.common.job_core.package import OptionalMethod from karadoc.common.model import file_index from karadoc.common.table_utils import parse_table_name from karadoc.common.utils.assert_utils import assert_true -from karadoc.spark.job_core.has_external_inputs import ( - HasExternalInputs, - _read_external_input_signature_check, - _read_external_inputs_signature_check, -) -from karadoc.spark.job_core.has_external_outputs import ( - HasExternalOutputs, - _write_external_output_signature_check, - _write_external_outputs_signature_check, -) from karadoc.spark.job_core.has_spark import HasSpark -from karadoc.spark.job_core.has_stream_external_inputs import ( - HasStreamExternalInputs, - _read_stream_external_input_signature_check, - _read_stream_external_inputs_signature_check, -) -from karadoc.spark.job_core.has_stream_external_output import ( - HasStreamExternalOutput, - _write_stream_external_output_signature_check, -) from karadoc.spark.quality.checks import Alert, Metric Job = TypeVar("Job", JobBase, JobBase) @@ -175,15 +156,16 @@ def __load_module_file(module_name: str, module_path: Union[str, Path]) -> Modul return mod -def __load_optional_method( - mod: ModuleType, tpe: type, method_name: str, signature_method: Optional[Callable[..., Any]] = None -) -> None: - if hasattr(mod, method_name): - private_method_name = "_" + tpe.__name__ + "__" + method_name - if signature_method is not None: - signature_method = cast(FunctionType, signature_method) - check_method_signatures(method_name, getattr(mod, method_name), signature_method) - setattr(mod.job, private_method_name, getattr(mod, method_name)) +def __load_optional_method(mod: ModuleType, method: OptionalMethod) -> None: + if hasattr(mod, method.name): + method_defined_in_file = getattr(mod, method.name) + method.set_method(method_defined_in_file) + + +def __load_optional_methods(mod: ModuleType) -> None: + optional_methods = [method for name, method in inspect.getmembers(mod.job) if isinstance(method, OptionalMethod)] + for method in optional_methods: + __load_optional_method(mod, method) def _set_spark_batch_job(mod: ModuleType) -> None: @@ -211,10 +193,6 @@ def _set_quality_check_job(mod: ModuleType) -> None: from karadoc.spark.quality.quality_check_job import QualityCheckJob if isinstance(mod.job, QualityCheckJob): - __load_optional_method(mod, HasBeforeAfter, "before_all") - __load_optional_method(mod, HasBeforeAfter, "after_all") - __load_optional_method(mod, HasBeforeAfter, "before_each") - __load_optional_method(mod, HasBeforeAfter, "after_each") found_alerts = [obj for name, obj in inspect.getmembers(mod) if isinstance(obj, Alert)] for alert in found_alerts: mod.job.add_alert(alert) @@ -223,48 +201,12 @@ def _set_quality_check_job(mod: ModuleType) -> None: mod.job.add_metric(metric) -def _set_has_external_inputs(mod: ModuleType) -> None: - if isinstance(mod.job, HasExternalInputs): - __load_optional_method(mod, HasExternalInputs, "read_external_input", _read_external_input_signature_check()) - __load_optional_method(mod, HasExternalInputs, "read_external_inputs", _read_external_inputs_signature_check()) - - -def _set_has_stream_external_inputs(mod: ModuleType) -> None: - if isinstance(mod.job, HasStreamExternalInputs): - __load_optional_method( - mod, HasStreamExternalInputs, "read_external_input", _read_stream_external_input_signature_check() - ) - __load_optional_method( - mod, HasStreamExternalInputs, "read_external_inputs", _read_stream_external_inputs_signature_check() - ) - - -def _set_has_external_outputs(mod: ModuleType) -> None: - if isinstance(mod.job, HasExternalOutputs): - __load_optional_method( - mod, HasExternalOutputs, "write_external_output", _write_external_output_signature_check() - ) - __load_optional_method( - mod, HasExternalOutputs, "write_external_outputs", _write_external_outputs_signature_check() - ) - - -def _set_has_stream_external_output(mod: ModuleType) -> None: - if isinstance(mod.job, HasStreamExternalOutput): - __load_optional_method( - mod, HasStreamExternalOutput, "write_external_output", _write_stream_external_output_signature_check() - ) - - def __set_job_from_module(mod: ModuleType) -> JobBase: + __load_optional_methods(mod) _set_spark_batch_job(mod) _set_spark_stream_job(mod) _set_analyze_job(mod) _set_quality_check_job(mod) - _set_has_external_inputs(mod) - _set_has_stream_external_inputs(mod) - _set_has_external_outputs(mod) - _set_has_stream_external_output(mod) return mod.job @@ -281,8 +223,8 @@ def __load_file( file_path: str, file_type: str, passed_vars: Optional[Dict[str, str]], - job_type: type, -) -> JobBase: + job_type: Type[Job], +) -> Job: full_table_name = schema_name + "." + table_name if file_path is None: raise ActionFileLoadingError(f"Could not find a {file_type} for table {full_table_name}") diff --git a/karadoc/common/job_core/package.py b/karadoc/common/job_core/package.py new file mode 100644 index 0000000..8bafc34 --- /dev/null +++ b/karadoc/common/job_core/package.py @@ -0,0 +1,51 @@ +import inspect +from types import FunctionType +from typing import Callable, Optional + + +def method_signature_str(method_name: str, full_arg_spec: inspect.FullArgSpec) -> str: + annotations = full_arg_spec.annotations + + def arg_str(arg_name: str) -> str: + if arg_name in annotations: + return f"{arg_name}: {annotations[arg_name]}" + else: + return arg_name + + args = [arg_str(arg_name) for arg_name in full_arg_spec.args] + return_str = f" -> {annotations['return']}" if "return" in annotations else "" + return f"def {method_name}({', '.join(args)}){return_str}:" + + +def check_method_signatures(method_name: str, actual: Callable, expected: Callable) -> None: + """Ensures that the signature of a given method matches the expected signature taken from another method. + + :param method_name: name of the method to check + :param actual: function to check + :param expected: control function with the expected signature + :return: + """ + if not isinstance(actual, FunctionType): + raise TypeError("%s is not a function" % method_name) + actual_args = inspect.getfullargspec(actual) + expected_args = inspect.getfullargspec(expected) + if actual_args != expected_args: + raise TypeError( + f"The method {method_name} should have the following signature\n" + + method_signature_str(method_name, expected_args) + ) + + +class OptionalMethod: + def __init__(self, default_func: Callable, name: Optional[str] = None): + self.func = default_func + if name is None: + name = default_func.__name__ + self.name = name + + def __call__(self, *args, **kwargs): + return self.func(*args, **kwargs) + + def set_method(self, func: Callable): + check_method_signatures(self.name, func, self.func) + self.func = func diff --git a/karadoc/spark/job_core/has_external_inputs.py b/karadoc/spark/job_core/has_external_inputs.py index 993645e..a76ca96 100644 --- a/karadoc/spark/job_core/has_external_inputs.py +++ b/karadoc/spark/job_core/has_external_inputs.py @@ -1,6 +1,7 @@ -from typing import TYPE_CHECKING, Callable, Dict, Optional, Union +from typing import TYPE_CHECKING, Dict, Optional, Union from karadoc.common.conf import CONNECTION_GROUP +from karadoc.common.job_core.package import OptionalMethod from karadoc.spark.job_core.has_spark import HasSpark from karadoc.spark.spark_connector import SparkConnector @@ -8,34 +9,22 @@ from pyspark.sql import DataFrame -def _read_external_input_signature_check(): - from pyspark.sql import DataFrame - - def read_external_input(source: Dict) -> DataFrame: - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - pass - - return read_external_input - - -def _read_external_inputs_signature_check(): - from pyspark.sql import DataFrame - - def read_external_inputs() -> Dict[str, DataFrame]: - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - pass - - return read_external_inputs - - class HasExternalInputs(HasSpark): def __init__(self) -> None: super().__init__() + # Private attributes - self.__read_external_input: Optional[Callable] = None - """Stores the optional override of the read_external_input method by the user""" - self.__read_external_inputs: Optional[Callable] = None - """Stores the optional override of the read_external_inputs method by the user""" + from pyspark.sql import DataFrame + + def read_external_input(source: Dict) -> DataFrame: + connector = self.get_input_connector(source) + return connector.read(source) + + def read_external_inputs() -> Dict[str, DataFrame]: + return {source_alias: self.read_external_input(source_alias) for source_alias in self.external_inputs} + + self.__read_external_input = OptionalMethod(read_external_input) + self.__read_external_inputs = OptionalMethod(read_external_inputs) self._limit_external_inputs: Optional[int] = None """When set, limits the size of every external input DataFrame to this number of row @@ -58,26 +47,14 @@ def __init__(self) -> None: } """ - def _read_external_input_default(self, source: Dict) -> "DataFrame": - connector = self.get_input_connector(source) - return connector.read(source) - - def _read_external_inputs_default(self) -> Dict[str, "DataFrame"]: - return {source_alias: self.read_external_input(source_alias) for source_alias in self.external_inputs} - def read_external_input(self, source_alias: str) -> "DataFrame": """Reads a given external input and returns it as a Spark DataFrame - :param source_alias: the alias of the source in job.external_inputs + :param source_alias: alias of a source declared in job.external_inputs :return: a DataFrame """ source = self.external_inputs[source_alias] - external_input: "DataFrame" - if self.__read_external_input is None: - external_input = self._read_external_input_default(source) - else: - external_input = self.__read_external_input(source) - + external_input = self.__read_external_input(source) if self._limit_external_inputs is not None: external_input = external_input.limit(self._limit_external_inputs) return external_input @@ -88,10 +65,7 @@ def read_external_inputs(self) -> Dict[str, "DataFrame"]: :return: a Dict[alias, DataFrame] """ external_inputs: Dict[str, "DataFrame"] - if self.__read_external_inputs is None: - external_inputs = self._read_external_inputs_default() - else: - external_inputs = self.__read_external_inputs() + external_inputs = self.__read_external_inputs() if self._limit_external_inputs is not None: for alias, df in external_inputs.items(): diff --git a/karadoc/spark/job_core/has_external_outputs.py b/karadoc/spark/job_core/has_external_outputs.py index 7990ede..9c50fe8 100644 --- a/karadoc/spark/job_core/has_external_outputs.py +++ b/karadoc/spark/job_core/has_external_outputs.py @@ -1,51 +1,34 @@ from typing import TYPE_CHECKING, Dict, Union from karadoc.common.conf import CONNECTION_GROUP +from karadoc.common.job_core.package import OptionalMethod from karadoc.spark.job_core.has_spark import HasSpark if TYPE_CHECKING: from pyspark.sql import DataFrame -def _write_external_output_signature_check(): - from pyspark.sql import DataFrame - - def write_external_output_signature(df: DataFrame, dest: Dict) -> None: - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - pass - - return write_external_output_signature - +class HasExternalOutputs(HasSpark): + def __init__(self) -> None: + super().__init__() -def _write_external_outputs_signature_check(): - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - from pyspark.sql import DataFrame + # Private attributes + from pyspark.sql import DataFrame - def write_external_outputs(df: DataFrame) -> None: - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - pass + def write_external_output(df: DataFrame, dest: Dict) -> None: + connector = self.get_output_connector(dest) + connector.write(df, dest) - return write_external_outputs + def write_external_outputs(df: DataFrame) -> None: + for _, dest in self.external_outputs.items(): + self.write_external_output(df, dest) + self.__write_external_output = OptionalMethod(write_external_output) + self.__write_external_outputs = OptionalMethod(write_external_outputs) -class HasExternalOutputs(HasSpark): - def __init__(self) -> None: - super().__init__() # Attributes that the user may change self.external_outputs: Dict[str, dict] = {} - # Attributes that the user is not supposed to change - self.__write_external_output = None - self.__write_external_outputs = None - - def _write_external_output_default(self, df: "DataFrame", dest: Dict) -> None: - connector = self.get_output_connector(dest) - connector.write(df, dest) - - def _write_external_outputs_default(self, df: "DataFrame") -> None: - for _, dest in self.external_outputs.items(): - self.write_external_output(df, dest) - def write_external_output(self, df: "DataFrame", dest: Dict) -> None: """Writes a given DataFrame to a given external output @@ -53,10 +36,7 @@ def write_external_output(self, df: "DataFrame", dest: Dict) -> None: :param dest: The alias of the external output to write to :return: nothing """ - if self.__write_external_output is None: - return self._write_external_output_default(df, dest) - else: - return self.__write_external_output(df, dest) + return self.__write_external_output(df, dest) def write_external_outputs(self, df: "DataFrame") -> None: """Writes a given DataFrame to all the declared external outputs @@ -64,10 +44,7 @@ def write_external_outputs(self, df: "DataFrame") -> None: :param df: The DataFrame to write :return: nothing """ - if self.__write_external_outputs is None: - return self._write_external_outputs_default(df) - else: - return self.__write_external_outputs(df) + return self.__write_external_outputs(df) def get_output_connector(self, dest: Union[str, Dict]): if type(dest) == str: diff --git a/karadoc/spark/job_core/has_stream_external_inputs.py b/karadoc/spark/job_core/has_stream_external_inputs.py index d829a84..3e15960 100644 --- a/karadoc/spark/job_core/has_stream_external_inputs.py +++ b/karadoc/spark/job_core/has_stream_external_inputs.py @@ -1,38 +1,29 @@ from typing import TYPE_CHECKING, Dict, Union from karadoc.common.conf import CONNECTION_GROUP +from karadoc.common.job_core.package import OptionalMethod from karadoc.spark.job_core.has_spark import HasSpark if TYPE_CHECKING: from pyspark.sql import DataFrame -def _read_stream_external_input_signature_check(): - from pyspark.sql import DataFrame - - def read_stream_external_input(source: Dict) -> DataFrame: - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - pass - - return read_stream_external_input - - -def _read_stream_external_inputs_signature_check(): - from pyspark.sql import DataFrame - - def read_stream_external_inputs() -> Dict[str, DataFrame]: - """Empty method used to check the signature of the equivalent method defined in the POPULATE files""" - pass - - return read_stream_external_inputs - - class HasStreamExternalInputs(HasSpark): def __init__(self) -> None: super().__init__() + # Private attributes - self.__read_external_input = None - self.__read_external_inputs = None + from pyspark.sql import DataFrame + + def read_external_input(source: Dict) -> DataFrame: + connector = self.get_input_connector(source) + return connector.read_stream(source) + + def read_external_inputs() -> Dict[str, DataFrame]: + return {source_alias: self.read_external_input(source_alias) for source_alias in self.external_inputs} + + self.__read_external_input = OptionalMethod(read_external_input) + self.__read_external_inputs = OptionalMethod(read_external_inputs) # Attributes that the user may change in action files self.external_inputs: Dict[str, dict] = {} @@ -51,13 +42,6 @@ def __init__(self) -> None: } """ - def _read_external_input_default(self, source: Dict) -> "DataFrame": - connector = self.get_input_connector(source) - return connector.read_stream(source) - - def _read_external_inputs_default(self) -> Dict[str, "DataFrame"]: - return {source_alias: self.read_external_input(source_alias) for source_alias in self.external_inputs} - def read_external_input(self, source_alias: str) -> "DataFrame": """Reads a given external input and returns it as a Spark DataFrame @@ -65,20 +49,14 @@ def read_external_input(self, source_alias: str) -> "DataFrame": :return: a DataFrame """ source = self.external_inputs[source_alias] - if self.__read_external_input is None: - return self._read_external_input_default(source) - else: - return self.__read_external_input(source) + return self.__read_external_input(source) def read_external_inputs(self) -> Dict[str, "DataFrame"]: """Reads all declared external inputs and returns them as Spark DataFrames :return: a Dict[alias, DataFrame] """ - if self.__read_external_inputs is None: - return self._read_external_inputs_default() - else: - return self.__read_external_inputs() + return self.__read_external_inputs() def load_external_inputs_as_views(self, cache=False): for alias in self.external_inputs: diff --git a/karadoc/spark/job_core/has_stream_external_output.py b/karadoc/spark/job_core/has_stream_external_output.py index 30ee847..9d593c1 100644 --- a/karadoc/spark/job_core/has_stream_external_output.py +++ b/karadoc/spark/job_core/has_stream_external_output.py @@ -1,6 +1,8 @@ from typing import TYPE_CHECKING, Dict, Optional, Union from karadoc.common.conf import CONNECTION_GROUP +from karadoc.common.job_core.package import OptionalMethod +from karadoc.spark.job_core.has_spark import HasSpark if TYPE_CHECKING: from pyspark.sql import DataFrame @@ -18,29 +20,31 @@ def write_stream_external_output_signature_check(df: DataFrame, dest: Dict) -> D return write_stream_external_output_signature_check -class HasStreamExternalOutput: +class HasStreamExternalOutput(HasSpark): def __init__(self) -> None: - # Attributes that the user may change - self.external_output: Optional[dict] = None + super().__init__() + + # Private attributes + from pyspark.sql import DataFrame + from pyspark.sql.streaming import DataStreamWriter - # Attributes that the user is not supposed to change - self.__write_external_output = None + def write_external_output(df: DataFrame, dest: Dict) -> DataStreamWriter: + connector = self.get_output_connector(dest) + return connector.write_stream(df, dest) - def _write_external_output_default(self, df: "DataFrame", dest: Dict) -> "DataStreamWriter": - connector = self.get_output_connector(dest) - return connector.write_stream(df, dest) + self.__write_external_output = OptionalMethod(write_external_output) + + # Attributes that the user may change + self.external_output: Optional[dict] = None def write_external_output(self, df: "DataFrame", dest: Dict) -> "DataStreamWriter": """Writes a given DataFrame to a given external output :param df: The DataFrame to write :param dest: The alias of the external output to write to - :return: nothing + :return: a DataStreamWriter """ - if self.__write_external_output is None: - return self._write_external_output_default(df, dest) - else: - return self.__write_external_output(df, dest) + return self.__write_external_output(df, dest) def get_output_connector(self, dest: Union[str, Dict]): if type(dest) == str: diff --git a/tests/karadoc/spark/job_core/test_has_batch_output.py b/tests/karadoc/spark/job_core/test_has_batch_output.py index 98ae6f4..3fe95f1 100644 --- a/tests/karadoc/spark/job_core/test_has_batch_output.py +++ b/tests/karadoc/spark/job_core/test_has_batch_output.py @@ -18,13 +18,6 @@ def test_output_partition_to_dynamic_partitions_list(self): actual = has_output._output_partition_to_dynamic_partitions(["k0", ("k1", "v1"), "k2"]) self.assertEqual(actual, ["k0", "k2"]) - def test_output_partitioning_type(self): - job = HasBatchOutput() - job.output_partition = [("day", "2020-10-10")] - self.assertEqual(job.output_partitioning_type, "static") - job.output_partition = ["day"] - self.assertEqual(job.output_partitioning_type, "dynamic") - def test_output_partition_names(self): job = HasBatchOutput() job.output_partition = [("part1_name", "part1_value"), ("part2_name", "part2_value")] diff --git a/tests/resources/karadoc/cli/commands/test_quality_check/models/test_schema.db/has_inputs/QUALITY_CHECK.py b/tests/resources/karadoc/cli/commands/test_quality_check/models/test_schema.db/has_inputs/QUALITY_CHECK.py index f8641b4..8df9fe8 100644 --- a/tests/resources/karadoc/cli/commands/test_quality_check/models/test_schema.db/has_inputs/QUALITY_CHECK.py +++ b/tests/resources/karadoc/cli/commands/test_quality_check/models/test_schema.db/has_inputs/QUALITY_CHECK.py @@ -5,7 +5,7 @@ job.inputs = {"input_table": "test_schema.input_table"} -def before_all(): +def before_all() -> None: job.load_inputs_as_views() From aacb3355e3d3d0ab611c9bf1c7c02c366337a856 Mon Sep 17 00:00:00 2001 From: fpin Date: Sun, 26 Mar 2023 10:28:34 +0200 Subject: [PATCH 2/3] fix: fix bug when using mixed partitioning --- karadoc/spark/job_core/has_output.py | 27 ++------------------------- 1 file changed, 2 insertions(+), 25 deletions(-) diff --git a/karadoc/spark/job_core/has_output.py b/karadoc/spark/job_core/has_output.py index b3ead4c..478181c 100644 --- a/karadoc/spark/job_core/has_output.py +++ b/karadoc/spark/job_core/has_output.py @@ -20,6 +20,7 @@ def _output_partition_to_dynamic_partitions(partition) -> List[str]: class HasOutput(HasSpark, ABC): def __init__(self) -> None: + super().__init__() # Attributes that the user may change self.output_partition: List[Union[Tuple[str, ...], str]] = [] self.output_mode = "OVERWRITE" @@ -29,24 +30,6 @@ def __init__(self) -> None: self.output: Optional[str] = None self.output_warehouse_dir = conf.get_warehouse_folder_location() - @property - def output_partitioning_type(self) -> Optional[str]: - """Indicates the type of partitioning of the output. - Possible return values are: - - None: no partitioning - - "static": dynamic partitioning - - "dynamic": static partitioning - - :return: either None, "static", or "dynamic" - """ - if self.output_partition: - if type(self.output_partition[0]) == str: - return "dynamic" - else: - return "static" - else: - return None - @property def output_partition_names(self) -> List[str]: """Returns the names of the output partitions. @@ -54,13 +37,7 @@ def output_partition_names(self) -> List[str]: :return: a list of string """ - output_partitioning_type = self.output_partitioning_type - if output_partitioning_type == "dynamic": - return [p_name for p_name in self.output_partition] - elif output_partitioning_type == "static": - return [p_name for p_name, p_value in self.output_partition] - else: - return [] + return [p[0] if isinstance(p, Tuple) else p for p in self.output_partition] @abstractmethod def write_table(self, df: "DataFrame"): From 19a0e374573090de52d92a5c44572d5658c0bc72 Mon Sep 17 00:00:00 2001 From: fpin Date: Sun, 26 Mar 2023 11:13:42 +0200 Subject: [PATCH 3/3] refactor: Simplify code for loading required methods in action files --- karadoc/common/commands/command.py | 2 +- karadoc/common/job_core/has_before_after.py | 1 + karadoc/common/job_core/job_base.py | 3 - karadoc/common/job_core/load.py | 59 +++++++------------ karadoc/common/job_core/package.py | 33 ++++++++++- karadoc/spark/analyze/analyze_job.py | 5 +- karadoc/spark/batch/spark_batch_job.py | 4 +- karadoc/spark/stream/spark_stream_job.py | 5 +- tests/karadoc/spark/batch/test_batch_exec.py | 2 +- .../karadoc/spark/stream/test_stream_exec.py | 2 +- 10 files changed, 63 insertions(+), 53 deletions(-) diff --git a/karadoc/common/commands/command.py b/karadoc/common/commands/command.py index b080ed9..408fc59 100644 --- a/karadoc/common/commands/command.py +++ b/karadoc/common/commands/command.py @@ -48,7 +48,7 @@ def add_arguments(parser: ArgumentParser) -> None: @abstractmethod def do_command(args: Namespace) -> Optional[ReturnCode]: """Entrypoint of the command. - This method can return nothing if it's only way of failing is raising an Exception. + This method can return nothing if its only way of failing is raising an Exception. Commands that can fail without raising exceptions (e.g. a validation command) should return a ReturnCode. """ pass diff --git a/karadoc/common/job_core/has_before_after.py b/karadoc/common/job_core/has_before_after.py index e755d2d..642bb6a 100644 --- a/karadoc/common/job_core/has_before_after.py +++ b/karadoc/common/job_core/has_before_after.py @@ -4,6 +4,7 @@ class HasBeforeAfter: def __init__(self): def default() -> None: + # We do nothing by default pass self.before_each = OptionalMethod(default, "before_each") diff --git a/karadoc/common/job_core/job_base.py b/karadoc/common/job_core/job_base.py index 405d15d..5e38355 100644 --- a/karadoc/common/job_core/job_base.py +++ b/karadoc/common/job_core/job_base.py @@ -1,11 +1,8 @@ -from typing import Optional - from karadoc.common.utils.assert_utils import assert_true class JobBase: _action_file_name_conf_key: str - _run_method_name: Optional[str] = None output: str @classmethod diff --git a/karadoc/common/job_core/load.py b/karadoc/common/job_core/load.py index 41b9872..25927f6 100644 --- a/karadoc/common/job_core/load.py +++ b/karadoc/common/job_core/load.py @@ -10,7 +10,7 @@ from karadoc.common.exceptions import ActionFileLoadingError, ForbiddenActionError from karadoc.common.job_core.has_vars import HasVars from karadoc.common.job_core.job_base import JobBase -from karadoc.common.job_core.package import OptionalMethod +from karadoc.common.job_core.package import ActionFileMethod from karadoc.common.model import file_index from karadoc.common.table_utils import parse_table_name from karadoc.common.utils.assert_utils import assert_true @@ -39,15 +39,7 @@ def load_non_runnable_action_file(full_table_name: str, job_type: Type[Job]) -> passed_vars = None job = __load_action_file(job_type, full_table_name, passed_vars) - if job._run_method_name is not None: - - def empty_run() -> None: - raise ForbiddenActionError( - f"The {job._run_method_name} method of a job returned by the `load_non_runnable_action_file` " - "method cannot be called. use load_runnable_action_file instead." - ) - - job.__setattr__(job._run_method_name, empty_run) + __unset_all_action_methods(job) return job @@ -156,37 +148,29 @@ def __load_module_file(module_name: str, module_path: Union[str, Path]) -> Modul return mod -def __load_optional_method(mod: ModuleType, method: OptionalMethod) -> None: - if hasattr(mod, method.name): - method_defined_in_file = getattr(mod, method.name) - method.set_method(method_defined_in_file) - - -def __load_optional_methods(mod: ModuleType) -> None: - optional_methods = [method for name, method in inspect.getmembers(mod.job) if isinstance(method, OptionalMethod)] - for method in optional_methods: - __load_optional_method(mod, method) +def __load_action_methods(mod: ModuleType) -> None: + action_methods = [method for name, method in inspect.getmembers(mod.job) if isinstance(method, ActionFileMethod)] + for method in action_methods: + if hasattr(mod, method.name): + method_defined_in_file = getattr(mod, method.name) + method.set_method(method_defined_in_file) -def _set_spark_batch_job(mod: ModuleType) -> None: - from karadoc.spark.batch.spark_batch_job import SparkBatchJob - - if isinstance(mod.job, SparkBatchJob): - mod.job.run = mod.run - - -def _set_spark_stream_job(mod: ModuleType) -> None: - from karadoc.spark.stream.spark_stream_job import SparkStreamJob +def __unset_action_method(job: Job, method: ActionFileMethod): + def empty_run() -> None: + raise ForbiddenActionError( + f"The {method.name} method of a job returned by the `load_non_runnable_action_file` " + "method cannot be called. Use load_runnable_action_file instead." + ) - if isinstance(mod.job, SparkStreamJob): - mod.job.stream = mod.stream + job.__setattr__(method.name, empty_run) -def _set_analyze_job(mod: ModuleType) -> None: - from karadoc.spark.analyze.analyze_job import AnalyzeJob +def __unset_all_action_methods(job: Job) -> None: + action_methods = [method for name, method in inspect.getmembers(job) if isinstance(method, ActionFileMethod)] - if isinstance(mod.job, AnalyzeJob): - mod.job.analyze = mod.analyze + for method in action_methods: + __unset_action_method(job, method) def _set_quality_check_job(mod: ModuleType) -> None: @@ -202,10 +186,7 @@ def _set_quality_check_job(mod: ModuleType) -> None: def __set_job_from_module(mod: ModuleType) -> JobBase: - __load_optional_methods(mod) - _set_spark_batch_job(mod) - _set_spark_stream_job(mod) - _set_analyze_job(mod) + __load_action_methods(mod) _set_quality_check_job(mod) return mod.job diff --git a/karadoc/common/job_core/package.py b/karadoc/common/job_core/package.py index 8bafc34..e4dd313 100644 --- a/karadoc/common/job_core/package.py +++ b/karadoc/common/job_core/package.py @@ -1,4 +1,5 @@ import inspect +from abc import ABC, abstractmethod from types import FunctionType from typing import Callable, Optional @@ -36,7 +37,20 @@ def check_method_signatures(method_name: str, actual: Callable, expected: Callab ) -class OptionalMethod: +class ActionFileMethod(ABC): + name: str + + @abstractmethod + def __call__(self, *args, **kwargs): + pass + + @abstractmethod + def set_method(self, func: Callable) -> None: + """Description of the command which will be displayed in the help""" + pass + + +class OptionalMethod(ActionFileMethod): def __init__(self, default_func: Callable, name: Optional[str] = None): self.func = default_func if name is None: @@ -46,6 +60,21 @@ def __init__(self, default_func: Callable, name: Optional[str] = None): def __call__(self, *args, **kwargs): return self.func(*args, **kwargs) - def set_method(self, func: Callable): + def set_method(self, func: Callable) -> None: check_method_signatures(self.name, func, self.func) self.func = func + + +class RequiredMethod(ActionFileMethod): + def __init__(self, name: str, signature_func: Optional[Callable] = None): + self.name = name + self.signature_func = signature_func + self.func = None + + def __call__(self, *args, **kwargs) -> None: + return self.func(*args, **kwargs) + + def set_method(self, func: Callable) -> None: + if self.signature_func is not None: + check_method_signatures(self.name, func, self.signature_func) + self.func = func diff --git a/karadoc/spark/analyze/analyze_job.py b/karadoc/spark/analyze/analyze_job.py index 28cd347..3e8ee52 100644 --- a/karadoc/spark/analyze/analyze_job.py +++ b/karadoc/spark/analyze/analyze_job.py @@ -1,11 +1,11 @@ from karadoc.common.job_core.has_vars import HasVars +from karadoc.common.job_core.package import RequiredMethod from karadoc.spark.job_core.has_batch_inputs import HasBatchInputs from karadoc.spark.job_core.has_spark import HasSpark class AnalyzeJob(HasBatchInputs, HasVars, HasSpark): _action_file_name_conf_key = "spark.analyze_timeline" - _run_method_name = "analyze" def __init__(self) -> None: HasSpark.__init__(self) @@ -16,4 +16,5 @@ def __init__(self) -> None: self.reference_time_col = "application_date" self.cohorts = ["cohort"] self.nb_buckets = 5 - self.analyze = None + + self.analyze = RequiredMethod("analyze") diff --git a/karadoc/spark/batch/spark_batch_job.py b/karadoc/spark/batch/spark_batch_job.py index 59cde6e..4831ec6 100644 --- a/karadoc/spark/batch/spark_batch_job.py +++ b/karadoc/spark/batch/spark_batch_job.py @@ -1,6 +1,7 @@ from karadoc.common.job_core.has_disable import HasDisable from karadoc.common.job_core.has_keys import HasKeys from karadoc.common.job_core.has_vars import HasVars +from karadoc.common.job_core.package import RequiredMethod from karadoc.spark.job_core.has_batch_inputs import HasBatchInputs from karadoc.spark.job_core.has_batch_output import HasBatchOutput from karadoc.spark.job_core.has_external_inputs import HasExternalInputs @@ -19,7 +20,6 @@ class SparkBatchJob( HasSpark, ): _action_file_name_conf_key = "spark.batch" - _run_method_name = "run" def __init__(self) -> None: HasSpark.__init__(self) @@ -31,4 +31,4 @@ def __init__(self) -> None: HasKeys.__init__(self) HasDisable.__init__(self) - self.run = None + self.run = RequiredMethod("run") diff --git a/karadoc/spark/stream/spark_stream_job.py b/karadoc/spark/stream/spark_stream_job.py index 26ec529..3524f01 100644 --- a/karadoc/spark/stream/spark_stream_job.py +++ b/karadoc/spark/stream/spark_stream_job.py @@ -1,5 +1,6 @@ from karadoc.common.job_core.has_disable import HasDisable from karadoc.common.job_core.has_vars import HasVars +from karadoc.common.job_core.package import RequiredMethod from karadoc.spark.job_core.has_spark import HasSpark from karadoc.spark.job_core.has_stream_external_inputs import HasStreamExternalInputs from karadoc.spark.job_core.has_stream_external_output import HasStreamExternalOutput @@ -17,7 +18,6 @@ class SparkStreamJob( HasSpark, ): _action_file_name_conf_key = "spark.stream" - _run_method_name = "stream" def __init__(self) -> None: HasSpark.__init__(self) @@ -27,4 +27,5 @@ def __init__(self) -> None: HasStreamInputs.__init__(self) HasStreamOutput.__init__(self) HasDisable.__init__(self) - self.stream = None + + self.stream = RequiredMethod("stream") diff --git a/tests/karadoc/spark/batch/test_batch_exec.py b/tests/karadoc/spark/batch/test_batch_exec.py index 69db446..c1077b7 100644 --- a/tests/karadoc/spark/batch/test_batch_exec.py +++ b/tests/karadoc/spark/batch/test_batch_exec.py @@ -24,7 +24,7 @@ def test_load_populate(self): job.init() with self.assertRaises(Exception) as cm: job.run() - self.assertIn("use load_runnable_action_file instead", str(cm.exception)) + self.assertIn("Use load_runnable_action_file instead", str(cm.exception)) def test_load_runnable_populate_with_relative_imports(self): job = load_runnable_populate("test_schema.relative_import", {}) diff --git a/tests/karadoc/spark/stream/test_stream_exec.py b/tests/karadoc/spark/stream/test_stream_exec.py index 6eb17be..be9c624 100644 --- a/tests/karadoc/spark/stream/test_stream_exec.py +++ b/tests/karadoc/spark/stream/test_stream_exec.py @@ -40,4 +40,4 @@ def test_load_stream_file(self): job.init() with self.assertRaises(Exception) as cm: job.stream() - self.assertIn("use load_runnable_action_file instead", str(cm.exception)) + self.assertIn("Use load_runnable_action_file instead", str(cm.exception))