diff --git a/CHANGELOG.md b/CHANGELOG.md index 8dc0569..a444daa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,4 +1,7 @@ # Changelog +## v1.13.0 10/15/26 +- Add Batch client + ## v1.12.3 9/18/26 - Enable S3 client to download arbitrary files as BytesIO streams - Enable SFTP client to list available files diff --git a/README.md b/README.md index 87f5139..62605ab 100644 --- a/README.md +++ b/README.md @@ -8,6 +8,7 @@ This package contains common Python utility classes and functions. * Decrypting values with KMS * Encoding and decoding records using a given Avro schema * Retrieving secrets from AWS Secrets Manager +* Submitting jobs to an AWS Batch queue * Downloading files from a remote SSH SFTP server * Connecting to and querying an Azure SQL database * Connecting to and querying a MySQL database @@ -72,7 +73,7 @@ pip install \ --platform manylinux2014_x86_64 \ --target=./package \ --implementation cp \ - --python 3.9 \ + --python 3.13 \ --only-binary=:all: --upgrade \ 'psycopg[binary]' ``` diff --git a/pyproject.toml b/pyproject.toml index 82590df..2ce7180 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "nypl_py_utils" -version = "1.12.3" +version = "1.13.0" authors = [ { name="Aaron Friedman", email="aaronfriedman@nypl.org" }, ] @@ -33,6 +33,11 @@ azure-client = [ "mssql-python==1.15.0", "pandas==3.0.6" ] +batch-client = [ + "nypl_py_utils[log-helper]", + "boto3==1.43.97", + "botocore==1.43.97" +] cloudlibrary-client = [ "nypl_py_utils[log-helper]", "requests==2.34.2" @@ -102,7 +107,7 @@ research-catalog-identifier-helper = [ "requests==2.34.2" ] development = [ - "nypl_py_utils[avro-client,azure-client,cloudlibrary-client,kinesis-client,kms-client,mysql-client,oauth2-api-client,postgresql-client,redshift-client,s3-client,secrets-manager-client,sftp-client,snowflake-client,config-helper,log-helper,obfuscation-helper,patron-data-helper,research-catalog-identifier-helper]", + "nypl_py_utils[avro-client,azure-client,batch-client,cloudlibrary-client,kinesis-client,kms-client,mysql-client,oauth2-api-client,postgresql-client,redshift-client,s3-client,secrets-manager-client,sftp-client,snowflake-client,config-helper,log-helper,obfuscation-helper,patron-data-helper,research-catalog-identifier-helper]", "flake8==7.3.0", "freezegun==1.5.5", "mock==5.2.0", diff --git a/src/nypl_py_utils/classes/batch_client.py b/src/nypl_py_utils/classes/batch_client.py new file mode 100644 index 0000000..8773662 --- /dev/null +++ b/src/nypl_py_utils/classes/batch_client.py @@ -0,0 +1,76 @@ +import boto3 +import os + +from botocore.exceptions import ClientError +from nypl_py_utils.functions.log_helper import create_log + + +class BatchClient: + """Client for interacting with AWS Batch""" + + def __init__(self): + self.logger = create_log("batch_client") + + try: + self.batch_client = boto3.client( + "batch", region_name=os.environ.get("AWS_REGION", "us-east-1") + ) + except ClientError as e: + self.logger.error( + "Could not create Batch client: {err}".format(err=e) + ) + raise BatchClientError( + "Could not create Batch client: {err}".format(err=e) + ) from None + + def close(self): + self.batch_client.close() + + def submit_job( + self, job_name, job_definition, queue_name=None, env_vars={} + ): + """ + Submits a job to be run by AWS Batch. + + Parameters + ---------- + job_name: str + The name of the job. This is the developer's choice and is NOT the + same as the job_definition, which must correspond to a known value. + job_definition: str + Name of the existing, registered job definition in AWS to run + queue_name: str, optional + Name of the job queue to use. If None, defaults to the + job_definition string. + env_vars: dict, optional + Environment variables to use when running the job + """ + formatted_env = [{"name": k, "value": v} for k, v in env_vars.items()] + if queue_name is None: + queue_name = job_definition + + self.logger.info(f"Submitting {job_definition} as {job_name}") + try: + if formatted_env: + self.batch_client.submit_job( + jobName=job_name, + jobQueue=queue_name, + jobDefinition=job_definition, + containerOverrides={"environment": formatted_env}, + ) + else: + self.batch_client.submit_job( + jobName=job_name, + jobQueue=queue_name, + jobDefinition=job_definition, + ) + except ClientError as e: + self.logger.error(f"Could not run {job_definition} job: {str(e)}") + raise BatchClientError( + f"Could not run {job_definition} job: {str(e)}" + ) from None + + +class BatchClientError(Exception): + def __init__(self, message=None): + self.message = message diff --git a/tests/test_batch_client.py b/tests/test_batch_client.py new file mode 100644 index 0000000..65a77c1 --- /dev/null +++ b/tests/test_batch_client.py @@ -0,0 +1,42 @@ +import pytest + +from botocore.exceptions import ClientError +from nypl_py_utils.classes.batch_client import BatchClient, BatchClientError + + +class TestBatchClient: + + @pytest.fixture + def test_instance(self, mocker): + mocker.patch('boto3.client') + return BatchClient() + + def test_submit_job_defaults(self, test_instance): + test_instance.submit_job("job_name", "job_def") + test_instance.batch_client.submit_job.assert_called_once_with( + jobName="job_name", + jobQueue="job_def", + jobDefinition="job_def" + ) + + def test_submit_job_overrides(self, test_instance): + test_instance.submit_job( + "job_name", "job_def", + queue_name="queue_name", + env_vars={"var1": "val1", "var2": 2} + ) + test_instance.batch_client.submit_job.assert_called_once_with( + jobName="job_name", + jobQueue="queue_name", + jobDefinition="job_def", + containerOverrides={"environment": [ + {"name": "var1", "value": "val1"}, + {"name": "var2", "value": 2} + ]} + ) + + def test_submit_job_error(self, test_instance): + test_instance.batch_client.submit_job.side_effect = \ + ClientError({}, 'SubmitJob') + with pytest.raises(BatchClientError): + test_instance.submit_job("job_name", "job_def")