jam-cloud/lambda/jamtrack-importer/batch/run_batch_job.py

87 lines
2.9 KiB
Python

import os
import json
import subprocess
import sys
def load_env_vars(json_file):
"""
Load environment variables from a JSON file and set them in the current process.
"""
try:
with open(json_file, 'r') as f:
env_vars = json.load(f)
if not isinstance(env_vars, dict):
raise ValueError("JSON file must contain key-value pairs.")
# Set each key-value pair as an environment variable
for key, value in env_vars.items():
os.environ[key] = str(value)
print(f"Set environment variable: {key}={value}")
except FileNotFoundError:
print(f"Error: File not found - {json_file}")
sys.exit(1)
except json.JSONDecodeError:
print(f"Error: Failed to parse JSON file - {json_file}")
sys.exit(1)
except Exception as e:
print(f"Error: {e}")
sys.exit(1)
def run_docker(image, env_vars, docker_args, jam_track_id):
"""
Run a Docker container with the specified image, environment variables, and additional arguments.
"""
try:
# Build the Docker run command with environment variables
docker_cmd = ["docker", "run", "--rm"]
for key, value in env_vars.items():
docker_cmd.extend(["-e", f"{key}={value}"])
docker_cmd.extend(["-e", f"JAM_TRACK_ID={jam_track_id}"])
# Add the Docker image and additional arguments
docker_cmd.append(image)
docker_cmd.extend(docker_args)
# Execute the Docker run command
print(f"Running Docker command: {' '.join(docker_cmd)}")
process = subprocess.Popen(docker_cmd, stdout = sys.stdout, stderr = sys.stderr, text=True)
#result = subprocess.run(docker_cmd, capture_output=False, text=True)
process.wait()
# Print the output or handle errors
if process.returncode == 0:
print("Docker runner succeeded")
else:
print(f"Docker runner failed #{process.returncode}:")
except Exception as e:
print(f"Error running Docker: {e}")
sys.exit(1)
if __name__ == "__main__":
if len(sys.argv) < 3:
print("Usage: python run_batch_job.py <docker_image> <json_file> [additional_args...]")
sys.exit(1)
docker_image = sys.argv[1]
json_file = sys.argv[2]
jam_track_id = sys.argv[3]
docker_args = sys.argv[4:] # All remaining arguments are passed to the Docker command
# Load environment variables from JSON
load_env_vars(json_file)
# Extract current environment variables (after setting from JSON)
current_env = {key: value for key, value in os.environ.items() if key in json.load(open(json_file))}
current_env["AWS_ACCESS_KEY_ID"] = os.environ["AWS_KEY"]
current_env["AWS_SECRET_ACCESS_KEY"] = os.environ["AWS_SECRET"]
# Run the Docker container
run_docker(docker_image, current_env, docker_args, jam_track_id)