Add wrapper for submitting AWS Batch jobs

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

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]'
 ```

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",

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

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")