Reformatted with black

Author: fatimarahman

CHANGELOG.md

@@ -1,6 +1,6 @@
 # Changelog
 ## v1.1.6 6/26/24
-- Generalized Avro functions and separated encoding/decoding behavior.
+- Generalized Avro functions and separated encoding/decoding behavior
 
 ## v1.1.5 6/6/24
 - Use executemany instead of execute when appropriate in RedshiftClient.execute_transaction

src/nypl_py_utils/classes/avro_client.py

@@ -9,46 +9,53 @@
 from nypl_py_utils.functions.log_helper import create_log
 from requests.exceptions import JSONDecodeError, RequestException
 
+
 class AvroClient:
     """
     Base class for Avro schema interaction. Takes as input the
     Platform API endpoint from which to fetch the schema in JSON format.
     """
-    
+
     def __init__(self, platform_schema_url):
-        self.logger = create_log('avro_encoder')
-        self.schema = avro.schema.parse(
-            self.get_json_schema(platform_schema_url))
-    
+        self.logger = create_log("avro_encoder")
+        self.schema = avro.schema.parse(self.get_json_schema(platform_schema_url))
+
     def get_json_schema(self, platform_schema_url):
         """
         Fetches a JSON response from the input Platform API endpoint and
         interprets it as an Avro schema.
         """
-        self.logger.info('Fetching Avro schema from {}'.format(
-            platform_schema_url))
+        self.logger.info("Fetching Avro schema from {}".format(platform_schema_url))
         try:
             response = requests.get(platform_schema_url)
             response.raise_for_status()
         except RequestException as e:
             self.logger.error(
-                'Failed to retrieve schema from {url}: {error}'.format(
-                    url=platform_schema_url, error=e))
+                "Failed to retrieve schema from {url}: {error}".format(
+                    url=platform_schema_url, error=e
+                )
+            )
             raise AvroClientError(
-                'Failed to retrieve schema from {url}: {error}'.format(
-                    url=platform_schema_url, error=e)) from None
+                "Failed to retrieve schema from {url}: {error}".format(
+                    url=platform_schema_url, error=e
+                )
+            ) from None
 
         try:
             json_response = response.json()
-            return json_response['data']['schema']
+            return json_response["data"]["schema"]
         except (JSONDecodeError, KeyError) as e:
             self.logger.error(
-                'Retrieved schema is malformed: {errorType} {errorMessage}'
-                .format(errorType=type(e), errorMessage=e))
+                "Retrieved schema is malformed: {errorType} {errorMessage}".format(
+                    errorType=type(e), errorMessage=e
+                )
+            )
             raise AvroClientError(
-                'Retrieved schema is malformed: {errorType} {errorMessage}'
-                .format(errorType=type(e), errorMessage=e)) from None
-    
+                "Retrieved schema is malformed: {errorType} {errorMessage}".format(
+                    errorType=type(e), errorMessage=e
+                )
+            ) from None
+
 
 class AvroEncoder(AvroClient):
     """
@@ -63,18 +70,17 @@ def encode_record(self, record):
         Returns the encoded record as a byte string.
         """
         self.logger.debug(
-            'Encoding record using {schema} schema'.format(
-               schema=self.schema.name))
+            "Encoding record using {schema} schema".format(schema=self.schema.name)
+        )
         datum_writer = DatumWriter(self.schema)
         with BytesIO() as output_stream:
             encoder = BinaryEncoder(output_stream)
             try:
                 datum_writer.write(record, encoder)
                 return output_stream.getvalue()
             except AvroException as e:
-                self.logger.error('Failed to encode record: {}'.format(e))
-                raise AvroClientError(
-                    'Failed to encode record: {}'.format(e)) from None
+                self.logger.error("Failed to encode record: {}".format(e))
+                raise AvroClientError("Failed to encode record: {}".format(e)) from None
 
     def encode_batch(self, record_list):
         """
@@ -83,8 +89,10 @@ def encode_batch(self, record_list):
         Returns a list of byte strings where each string is an encoded record.
         """
         self.logger.info(
-            'Encoding ({num_rec}) records using {schema} schema'.format(
-                num_rec=len(record_list), schema=self.schema.name))
+            "Encoding ({num_rec}) records using {schema} schema".format(
+                num_rec=len(record_list), schema=self.schema.name
+            )
+        )
         encoded_records = []
         datum_writer = DatumWriter(self.schema)
         with BytesIO() as output_stream:
@@ -96,9 +104,10 @@ def encode_batch(self, record_list):
                     output_stream.seek(0)
                     output_stream.truncate(0)
                 except AvroException as e:
-                    self.logger.error('Failed to encode record: {}'.format(e))
+                    self.logger.error("Failed to encode record: {}".format(e))
                     raise AvroClientError(
-                        'Failed to encode record: {}'.format(e)) from None
+                        "Failed to encode record: {}".format(e)
+                    ) from None
         return encoded_records
 
 
@@ -110,34 +119,37 @@ class AvroDecoder(AvroClient):
 
     def decode_record(self, record, encoding="binary"):
         """
-        Decodes a single record represented either as a byte or 
+        Decodes a single record represented either as a byte or
         base64 string, using the given Avro schema.
 
         Returns a dictionary where each key is a field in the schema.
         """
-        self.logger.info('Decoding {rec} of type {type} using {schema} schema'.format(
-            rec=record, type=encoding, schema=self.schema.name))
-        
+        self.logger.info(
+            "Decoding {rec} of type {type} using {schema} schema".format(
+                rec=record, type=encoding, schema=self.schema.name
+            )
+        )
+
         if encoding == "base64":
             return self._decode_base64(record)
         elif encoding == "binary":
             return self._decode_binary(record)
         else:
-            self.logger.error('Failed to decode record due to encoding type: {}'.format(encoding))
-            raise AvroClientError(
-                'Invalid encoding type: {}'.format(encoding))
-        
+            self.logger.error(
+                "Failed to decode record due to encoding type: {}".format(encoding)
+            )
+            raise AvroClientError("Invalid encoding type: {}".format(encoding))
+
     def _decode_base64(self, record):
-        decoded_data = base64.b64decode(record).decode("utf-8")
+        decoded_data = base64.b64decode(record)
         try:
             return json.loads(decoded_data)
         except Exception as e:
             if isinstance(decoded_data, bytes):
-                self._decode_binary(decoded_data) 
+                return self._decode_binary(decoded_data)
             else:
-                self.logger.error('Failed to decode record: {}'.format(e))
-                raise AvroClientError(
-                    'Failed to decode record: {}'.format(e)) from None
+                self.logger.error("Failed to decode record: {}".format(e))
+                raise AvroClientError("Failed to decode record: {}".format(e)) from None
 
     def _decode_binary(self, record):
         datum_reader = DatumReader(self.schema)
@@ -146,11 +158,10 @@ def _decode_binary(self, record):
             try:
                 return datum_reader.read(decoder)
             except Exception as e:
-                self.logger.error('Failed to decode record: {}'.format(e))
-                raise AvroClientError(
-                    'Failed to decode record: {}'.format(e)) from None
+                self.logger.error("Failed to decode record: {}".format(e))
+                raise AvroClientError("Failed to decode record: {}".format(e)) from None
 
 
 class AvroClientError(Exception):
     def __init__(self, message=None):
-        self.message = message
+        self.message = message

src/nypl_py_utils/classes/kinesis_client.py

@@ -17,20 +17,19 @@ class KinesisClient:
     """
 
     def __init__(self, stream_arn, batch_size, max_retries=5):
-        self.logger = create_log('kinesis_client')
+        self.logger = create_log("kinesis_client")
         self.stream_arn = stream_arn
         self.batch_size = batch_size
         self.max_retries = max_retries
 
         try:
             self.kinesis_client = boto3.client(
-                'kinesis', region_name=os.environ.get('AWS_REGION',
-                                                      'us-east-1'))
+                "kinesis", region_name=os.environ.get("AWS_REGION", "us-east-1")
+            )
         except ClientError as e:
-            self.logger.error(
-                'Could not create Kinesis client: {err}'.format(err=e))
+            self.logger.error("Could not create Kinesis client: {err}".format(err=e))
             raise KinesisClientError(
-                'Could not create Kinesis client: {err}'.format(err=e)
+                "Could not create Kinesis client: {err}".format(err=e)
             ) from None
 
     def close(self):
@@ -45,10 +44,11 @@ def send_records(self, records):
         """
         records_sent_since_pause = 0
         for i in range(0, len(records), self.batch_size):
-            encoded_batch = records[i:i + self.batch_size]
-            kinesis_records = [{'Data': record, 'PartitionKey':
-                                str(int(time.time() * 1000000000))}
-                               for record in encoded_batch]
+            encoded_batch = records[i : i + self.batch_size]
+            kinesis_records = [
+                {"Data": record, "PartitionKey": str(int(time.time() * 1000000000))}
+                for record in encoded_batch
+            ]
 
             if records_sent_since_pause + len(encoded_batch) > 1000:
                 records_sent_since_pause = 0
@@ -63,32 +63,41 @@ def _send_kinesis_format_records(self, kinesis_records, call_count):
         """
         if call_count > self.max_retries:
             self.logger.error(
-                'Failed to send records to Kinesis {} times in a row'.format(
-                    call_count-1))
+                "Failed to send records to Kinesis {} times in a row".format(
+                    call_count - 1
+                )
+            )
             raise KinesisClientError(
-                'Failed to send records to Kinesis {} times in a row'.format(
-                    call_count-1)) from None
+                "Failed to send records to Kinesis {} times in a row".format(
+                    call_count - 1
+                )
+            ) from None
 
         try:
             self.logger.info(
-                'Sending ({count}) records to {arn} Kinesis stream'.format(
-                    count=len(kinesis_records), arn=self.stream_arn))
+                "Sending ({count}) records to {arn} Kinesis stream".format(
+                    count=len(kinesis_records), arn=self.stream_arn
+                )
+            )
             response = self.kinesis_client.put_records(
-                Records=kinesis_records, StreamARN=self.stream_arn)
-            if response['FailedRecordCount'] > 0:
+                Records=kinesis_records, StreamARN=self.stream_arn
+            )
+            if response["FailedRecordCount"] > 0:
                 self.logger.warning(
-                    'Failed to send {} records to Kinesis'.format(
-                        response['FailedRecordCount']))
+                    "Failed to send {} records to Kinesis".format(
+                        response["FailedRecordCount"]
+                    )
+                )
                 failed_records = []
-                for i in range(len(response['Records'])):
-                    if 'ErrorCode' in response['Records'][i]:
+                for i in range(len(response["Records"])):
+                    if "ErrorCode" in response["Records"][i]:
                         failed_records.append(kinesis_records[i])
-                self._send_kinesis_format_records(failed_records, call_count+1)
+                self._send_kinesis_format_records(failed_records, call_count + 1)
         except ClientError as e:
-            self.logger.error(
-                'Error sending records to Kinesis: {}'.format(e))
+            self.logger.error("Error sending records to Kinesis: {}".format(e))
             raise KinesisClientError(
-                'Error sending records to Kinesis: {}'.format(e)) from None
+                "Error sending records to Kinesis: {}".format(e)
+            ) from None
 
 
 class KinesisClientError(Exception):

src/nypl_py_utils/classes/kms_client.py

@@ -11,16 +11,17 @@ class KmsClient:
     """Client for interacting with a KMS client"""
 
     def __init__(self):
-        self.logger = create_log('kms_client')
+        self.logger = create_log("kms_client")
 
         try:
             self.kms_client = boto3.client(
-                'kms', region_name=os.environ.get('AWS_REGION', 'us-east-1'))
+                "kms", region_name=os.environ.get("AWS_REGION", "us-east-1")
+            )
         except ClientError as e:
-            self.logger.error(
-                'Could not create KMS client: {err}'.format(err=e))
+            self.logger.error("Could not create KMS client: {err}".format(err=e))
             raise KmsClientError(
-                'Could not create KMS client: {err}'.format(err=e)) from None
+                "Could not create KMS client: {err}".format(err=e)
+            ) from None
 
     def close(self):
         self.kms_client.close()
@@ -30,16 +31,19 @@ def decrypt(self, encrypted_text):
         This method takes a base 64 KMS-encoded string and uses the KMS client
         to decrypt it into a usable string.
         """
-        self.logger.debug('Decrypting \'{}\''.format(encrypted_text))
+        self.logger.debug("Decrypting '{}'".format(encrypted_text))
         try:
             decoded_text = b64decode(encrypted_text)
             return self.kms_client.decrypt(CiphertextBlob=decoded_text)[
-                'Plaintext'].decode('utf-8')
+                "Plaintext"
+            ].decode("utf-8")
         except (ClientError, base64Error, TypeError) as e:
-            self.logger.error('Could not decrypt \'{val}\': {err}'.format(
-                val=encrypted_text, err=e))
-            raise KmsClientError('Could not decrypt \'{val}\': {err}'.format(
-                val=encrypted_text, err=e)) from None
+            self.logger.error(
+                "Could not decrypt '{val}': {err}".format(val=encrypted_text, err=e)
+            )
+            raise KmsClientError(
+                "Could not decrypt '{val}': {err}".format(val=encrypted_text, err=e)
+            ) from None
 
 
 class KmsClientError(Exception):