o-nikolas commented on code in PR #48513:
URL: https://github.com/apache/airflow/pull/48513#discussion_r2029458311
##########
task-sdk/src/airflow/sdk/execution_time/execute_workload.py:
##########
@@ -74,16 +71,44 @@ def main():
parser = argparse.ArgumentParser(
description="Execute a workload in a Containerised executor using the
task SDK."
)
- parser.add_argument(
- "input_file", help="Path to the input JSON file containing the
execution workload payload."
+
+ # Create a mutually exclusive group to ensure that only one of the flags
is set
+ group = parser.add_mutually_exclusive_group(required=True)
+ group.add_argument(
+ "--json-path",
+ help="Path to the input JSON file containing the execution workload
payload.",
+ type=str,
+ )
+ group.add_argument(
+ "--json-string",
+ help="The JSON string itself containing the execution workload
payload.",
+ type=str,
)
args = parser.parse_args()
- with open(args.input_file) as file:
- input_data = file.read()
+ from pydantic import TypeAdapter
- execute_workload(input_data)
+ from airflow.executors import workloads
+
+ decoder = TypeAdapter[workloads.All](workloads.All)
+ if args.json_path:
+ try:
+ with open(args.json_path) as file:
+ input_data = file.read()
+ workload = decoder.validate_json(input_data)
+ except Exception as e:
+ log.error("Failed to read file", error=str(e))
+ sys.exit(1)
+
+ elif args.json_string:
+ try:
+ workload = decoder.validate_json(args.json_string)
+ except Exception as e:
+ log.error("Failed to parse input JSON string", error=str(e))
+ sys.exit(1)
Review Comment:
Argparse will do that for us, the mutually exclusive group is set to be
required (see `required=True` on line 76), so if one isn't provided argparse
will throw an exception with a message saying so.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
