databuild/databuild/test/app/dsl/generated/dsl_job_lookup.py

53 lines
2 KiB
Python
Executable file

#!/usr/bin/env python3
"""
Generated job lookup for DataBuild DSL graph.
Maps partition patterns to job targets.
"""
import sys
import re
import json
from collections import defaultdict
# Mapping from partition patterns to job targets
JOB_MAPPINGS = {
r"daily_color_votes/(?P<data_date>\d{4}-\d{2}-\d{2})/(?P<color>[^/]+)": "//databuild/test/app/dsl/generated:ingest_color_votes",
r"color_votes_1m/(?P<data_date>\d{4}-\d{2}-\d{2})/(?P<color>[^/]+)": "//databuild/test/app/dsl/generated:trailing_color_votes",
r"color_votes_1w/(?P<data_date>\d{4}-\d{2}-\d{2})/(?P<color>[^/]+)": "//databuild/test/app/dsl/generated:trailing_color_votes",
r"daily_votes/(?P<data_date>\d{4}-\d{2}-\d{2})": "//databuild/test/app/dsl/generated:aggregate_color_votes",
r"votes_1w/(?P<data_date>\d{4}-\d{2}-\d{2})": "//databuild/test/app/dsl/generated:aggregate_color_votes",
r"votes_1m/(?P<data_date>\d{4}-\d{2}-\d{2})": "//databuild/test/app/dsl/generated:aggregate_color_votes",
r"color_vote_report/(?P<data_date>\d{4}-\d{2}-\d{2})/(?P<color>[^/]+)": "//databuild/test/app/dsl/generated:color_vote_report_calc",
}
def lookup_job_for_partition(partition_ref: str) -> str:
"""Look up which job can build the given partition reference."""
for pattern, job_target in JOB_MAPPINGS.items():
if re.match(pattern, partition_ref):
return job_target
raise ValueError(f"No job found for partition: {partition_ref}")
def main():
if len(sys.argv) < 2:
print("Usage: job_lookup.py <partition_ref> [partition_ref...]", file=sys.stderr)
sys.exit(1)
results = defaultdict(list)
try:
for partition_ref in sys.argv[1:]:
job_target = lookup_job_for_partition(partition_ref)
results[job_target].append(partition_ref)
# Output the results as JSON (matching existing lookup format)
print(json.dumps(dict(results)))
except ValueError as e:
print(f"ERROR: {e}", file=sys.stderr)
sys.exit(1)
if __name__ == "__main__":
main()