-
Notifications
You must be signed in to change notification settings - Fork 1.2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Add remote invoke implementation for Kinesis stream service for…
… put_record API (#6063) * Add mypy-boto3 stubs for kinesis * Add remote invoke implementation for kinesis service * unit tests * fix typo in comment Co-authored-by: Wing Fung Lau <[email protected]> --------- Co-authored-by: Wing Fung Lau <[email protected]>
- Loading branch information
Showing
5 changed files
with
364 additions
and
1 deletion.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,132 @@ | ||
| """ | ||
| Remote invoke executor implementation for Kinesis streams | ||
| """ | ||
| import logging | ||
| import uuid | ||
| from dataclasses import asdict, dataclass | ||
| from typing import cast | ||
|
|
||
| from botocore.exceptions import ClientError, ParamValidationError | ||
| from mypy_boto3_kinesis import KinesisClient | ||
|
|
||
| from samcli.lib.remote_invoke.exceptions import ( | ||
| ErrorBotoApiCallException, | ||
| InvalidResourceBotoParameterException, | ||
| ) | ||
| from samcli.lib.remote_invoke.remote_invoke_executors import ( | ||
| BotoActionExecutor, | ||
| RemoteInvokeIterableResponseType, | ||
| RemoteInvokeOutputFormat, | ||
| RemoteInvokeResponse, | ||
| ) | ||
|
|
||
| LOG = logging.getLogger(__name__) | ||
| STREAM_NAME = "StreamName" | ||
| DATA = "Data" | ||
| PARTITION_KEY = "PartitionKey" | ||
|
|
||
|
|
||
| @dataclass | ||
| class KinesisStreamPutRecordTextOutput: | ||
| """ | ||
| Dataclass that stores put_record boto3 API fields used to create | ||
| text output. | ||
| """ | ||
|
|
||
| ShardId: str | ||
| SequenceNumber: str | ||
|
|
||
| def get_output_response_dict(self) -> dict: | ||
| """ | ||
| Returns a dict of existing dataclass fields. | ||
| Returns | ||
| ------- | ||
| dict | ||
| Returns the dict of the fields that will be used as the output response for | ||
| text format output. | ||
| """ | ||
| return asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if v is not None}) | ||
|
|
||
|
|
||
| class KinesisPutDataExecutor(BotoActionExecutor): | ||
| """ | ||
| Calls "put_record" method of "Kinesis stream" service with given input. | ||
| If a file location provided, the file handle will be passed as input object. | ||
| """ | ||
|
|
||
| _kinesis_client: KinesisClient | ||
| _stream_name: str | ||
| _remote_output_format: RemoteInvokeOutputFormat | ||
| request_parameters: dict | ||
|
|
||
| def __init__(self, kinesis_client: KinesisClient, physical_id: str, remote_output_format: RemoteInvokeOutputFormat): | ||
| self._kinesis_client = kinesis_client | ||
| self._remote_output_format = remote_output_format | ||
| self._stream_name = physical_id | ||
| self.request_parameters = {} | ||
|
|
||
| def validate_action_parameters(self, parameters: dict) -> None: | ||
| """ | ||
| Validates the input boto parameters and prepares the parameters for calling the API. | ||
| Parameters | ||
| ---------- | ||
| parameters: dict | ||
| Boto parameters provided as input | ||
| """ | ||
| for parameter_key, parameter_value in parameters.items(): | ||
| if parameter_key == STREAM_NAME: | ||
| LOG.warning("StreamName is defined using the value provided for resource_id argument.") | ||
| elif parameter_key == DATA: | ||
| LOG.warning("Data is defined using the value provided for either --event or --event-file options.") | ||
| else: | ||
| self.request_parameters[parameter_key] = parameter_value | ||
|
|
||
| if PARTITION_KEY not in self.request_parameters: | ||
| self.request_parameters[PARTITION_KEY] = str(uuid.uuid4()) | ||
|
|
||
| def _execute_action(self, payload: str) -> RemoteInvokeIterableResponseType: | ||
| """ | ||
| Calls "put_record" method to write single data record to Kinesis data stream. | ||
| Parameters | ||
| ---------- | ||
| payload: str | ||
| The Data record which will be sent to the Kinesis stream | ||
| Yields | ||
| ------ | ||
| RemoteInvokeIterableResponseType | ||
| Response that is consumed by remote invoke consumers after execution | ||
| """ | ||
| if payload: | ||
| self.request_parameters[DATA] = payload | ||
| else: | ||
| self.request_parameters[DATA] = "{}" | ||
| LOG.debug("Input event not found, putting a record with Data {}") | ||
| self.request_parameters[STREAM_NAME] = self._stream_name | ||
| LOG.debug( | ||
| "Calling kinesis_client.put_record with StreamName:%s, Data:%s", | ||
| self.request_parameters[STREAM_NAME], | ||
| payload, | ||
| ) | ||
| try: | ||
| put_record_response = cast(dict, self._kinesis_client.put_record(**self.request_parameters)) | ||
|
|
||
| if self._remote_output_format == RemoteInvokeOutputFormat.JSON: | ||
| yield RemoteInvokeResponse(put_record_response) | ||
| if self._remote_output_format == RemoteInvokeOutputFormat.TEXT: | ||
| put_record_text_output = KinesisStreamPutRecordTextOutput( | ||
| ShardId=put_record_response["ShardId"], | ||
| SequenceNumber=put_record_response["SequenceNumber"], | ||
| ) | ||
| output_data = put_record_text_output.get_output_response_dict() | ||
| yield RemoteInvokeResponse(output_data) | ||
| except ParamValidationError as param_val_ex: | ||
| raise InvalidResourceBotoParameterException( | ||
| f"Invalid parameter key provided." | ||
| f" {str(param_val_ex).replace(f'{STREAM_NAME}, ', '').replace(f'{DATA}, ', '')}" | ||
| ) | ||
| except ClientError as client_ex: | ||
| raise ErrorBotoApiCallException(client_ex) from client_ex |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
142 changes: 142 additions & 0 deletions
142
tests/unit/lib/remote_invoke/test_kinesis_invoke_executors.py
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,142 @@ | ||
| from unittest import TestCase | ||
| from unittest.mock import patch, Mock | ||
|
|
||
| from parameterized import parameterized, parameterized_class | ||
| from samcli.lib.remote_invoke.kinesis_invoke_executors import ( | ||
| RemoteInvokeOutputFormat, | ||
| KinesisPutDataExecutor, | ||
| ParamValidationError, | ||
| InvalidResourceBotoParameterException, | ||
| ErrorBotoApiCallException, | ||
| ClientError, | ||
| KinesisStreamPutRecordTextOutput, | ||
| ) | ||
| from samcli.lib.remote_invoke.remote_invoke_executors import RemoteInvokeResponse | ||
|
|
||
|
|
||
| class TestKinesisStreamPutRecordTextOutput(TestCase): | ||
| @parameterized.expand( | ||
| [ | ||
| ("mock-shard-id", "mock-sequence-number"), | ||
| ] | ||
| ) | ||
| def test_kinesis_put_record_text_output(self, shard_id, sequence_number): | ||
| text_output = KinesisStreamPutRecordTextOutput(ShardId=shard_id, SequenceNumber=sequence_number) | ||
| self.assertEqual(text_output.ShardId, shard_id) | ||
| self.assertEqual(text_output.SequenceNumber, sequence_number) | ||
|
|
||
| @parameterized.expand( | ||
| [ | ||
| ( | ||
| "mock-shard-id", | ||
| "mock-sequence-number", | ||
| { | ||
| "ShardId": "mock-shard-id", | ||
| "SequenceNumber": "mock-sequence-number", | ||
| }, | ||
| ), | ||
| ] | ||
| ) | ||
| def test_get_output_response_dict(self, shard_id, sequence_number, expected_output): | ||
| text_output = KinesisStreamPutRecordTextOutput(ShardId=shard_id, SequenceNumber=sequence_number) | ||
| output_response_dict = text_output.get_output_response_dict() | ||
| self.assertEqual(output_response_dict, expected_output) | ||
|
|
||
|
|
||
| @parameterized_class( | ||
| "output", | ||
| [[RemoteInvokeOutputFormat.TEXT], [RemoteInvokeOutputFormat.JSON]], | ||
| ) | ||
| class TestKinesisPutDataExecutor(TestCase): | ||
| output: RemoteInvokeOutputFormat | ||
|
|
||
| def setUp(self) -> None: | ||
| self.kinesis_client = Mock() | ||
| self.stream_name = "mock-kinesis-stream" | ||
| self.kinesis_put_data_executor = KinesisPutDataExecutor(self.kinesis_client, self.stream_name, self.output) | ||
|
|
||
| @patch("samcli.lib.remote_invoke.kinesis_invoke_executors.uuid") | ||
| def test_execute_action_successful(self, patched_uuid): | ||
| mock_uuid_value = "patched-uuid-value" | ||
| patched_uuid.uuid4.return_value = mock_uuid_value | ||
| given_input_data = "hello world" | ||
| mock_shard_id = "shardId-000000000000" | ||
| mock_sequence_number = "2941492a-5847-4ebb-a8a3-58c07ce9f198" | ||
| mock_text_response = { | ||
| "ShardId": mock_shard_id, | ||
| "SequenceNumber": mock_sequence_number, | ||
| } | ||
|
|
||
| mock_json_response = { | ||
| "ShardId": mock_shard_id, | ||
| "SequenceNumber": mock_sequence_number, | ||
| "ResponseMetadata": {}, | ||
| } | ||
| self.kinesis_client.put_record.return_value = { | ||
| "ShardId": mock_shard_id, | ||
| "SequenceNumber": mock_sequence_number, | ||
| "ResponseMetadata": {}, | ||
| } | ||
| self.kinesis_put_data_executor.validate_action_parameters({}) | ||
| result = self.kinesis_put_data_executor._execute_action(given_input_data) | ||
| if self.output == RemoteInvokeOutputFormat.JSON: | ||
| self.assertEqual(list(result), [RemoteInvokeResponse(mock_json_response)]) | ||
| else: | ||
| self.assertEqual(list(result), [RemoteInvokeResponse(mock_text_response)]) | ||
|
|
||
| self.kinesis_client.put_record.assert_called_with( | ||
| Data=given_input_data, StreamName=self.stream_name, PartitionKey=mock_uuid_value | ||
| ) | ||
|
|
||
| @parameterized.expand( | ||
| [ | ||
| ({}, {"PartitionKey": "mock-uuid-value"}), | ||
| ( | ||
| {"ExplicitHashKey": "mock-explicit-hash-key", "SequenceNumberForOrdering": "1"}, | ||
| { | ||
| "PartitionKey": "mock-uuid-value", | ||
| "ExplicitHashKey": "mock-explicit-hash-key", | ||
| "SequenceNumberForOrdering": "1", | ||
| }, | ||
| ), | ||
| ( | ||
| { | ||
| "PartitionKey": "override-partition-key", | ||
| }, | ||
| { | ||
| "PartitionKey": "override-partition-key", | ||
| }, | ||
| ), | ||
| ( | ||
| {"StreamName": "mock-stream-name", "Data": "mock-data"}, | ||
| {"PartitionKey": "mock-uuid-value"}, | ||
| ), | ||
| ( | ||
| {"invalidParameterKey": "invalidParameterValue"}, | ||
| {"invalidParameterKey": "invalidParameterValue", "PartitionKey": "mock-uuid-value"}, | ||
| ), | ||
| ] | ||
| ) | ||
| @patch("samcli.lib.remote_invoke.kinesis_invoke_executors.uuid") | ||
| def test_validate_action_parameters(self, parameters, expected_boto_parameters, patched_uuid): | ||
| mock_uuid_value = "mock-uuid-value" | ||
| patched_uuid.uuid4.return_value = mock_uuid_value | ||
| self.kinesis_put_data_executor.validate_action_parameters(parameters) | ||
| self.assertEqual(self.kinesis_put_data_executor.request_parameters, expected_boto_parameters) | ||
|
|
||
| @parameterized.expand( | ||
| [ | ||
| (ParamValidationError(report="Invalid parameters"), InvalidResourceBotoParameterException), | ||
| ( | ||
| ClientError(error_response={"Error": {"Code": "MockException"}}, operation_name="send_message"), | ||
| ErrorBotoApiCallException, | ||
| ), | ||
| ] | ||
| ) | ||
| def test_execute_action_put_record_throws_boto_errors(self, boto_error, expected_error_thrown): | ||
| given_input_message = "hello world" | ||
| self.kinesis_client.put_record.side_effect = boto_error | ||
| with self.assertRaises(expected_error_thrown): | ||
| self.kinesis_put_data_executor.validate_action_parameters({}) | ||
| for _ in self.kinesis_put_data_executor._execute_action(given_input_message): | ||
| pass |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters