|  | 
| 20 | 20 | class AsyncSQSConnection(AsyncAWSQueryConnection): | 
| 21 | 21 |     """Async SQS Connection.""" | 
| 22 | 22 | 
 | 
| 23 |  | -    def __init__(self, sqs_connection, debug=0, region=None, fetch_message_attributes=None, **kwargs): | 
|  | 23 | +    def __init__( | 
|  | 24 | +        self, | 
|  | 25 | +        sqs_connection, | 
|  | 26 | +        debug=0, | 
|  | 27 | +        region=None, | 
|  | 28 | +        message_system_attribute_names=None, | 
|  | 29 | +        message_attribute_names=None, | 
|  | 30 | +        **kwargs | 
|  | 31 | +    ): | 
| 24 | 32 |         if boto3 is None: | 
| 25 | 33 |             raise ImportError('boto3 is not installed') | 
| 26 | 34 |         super().__init__( | 
| 27 | 35 |             sqs_connection, | 
| 28 | 36 |             region_name=region, debug=debug, | 
| 29 | 37 |             **kwargs | 
| 30 | 38 |         ) | 
| 31 |  | -        self.fetch_message_attributes = ( | 
| 32 |  | -            fetch_message_attributes if fetch_message_attributes is not None | 
| 33 |  | -            else ["ApproximateReceiveCount"] | 
|  | 39 | +        self.message_system_attribute_names = ( | 
|  | 40 | +            message_system_attribute_names if message_system_attribute_names else ["ApproximateReceiveCount"] | 
|  | 41 | +        ) | 
|  | 42 | +        self.message_attribute_names = ( | 
|  | 43 | +            [message_attribute_names] if isinstance(message_attribute_names, str) | 
|  | 44 | +            else (message_attribute_names or []) | 
| 34 | 45 |         ) | 
| 35 | 46 | 
 | 
| 36 | 47 |     def _create_query_request(self, operation, params, queue_url, method): | 
| @@ -160,13 +171,17 @@ def receive_message( | 
| 160 | 171 |     ): | 
| 161 | 172 |         params = {'MaxNumberOfMessages': number_messages} | 
| 162 | 173 |         proto_params = {'query': {}, 'json': {}} | 
| 163 |  | -        attrs = attributes if attributes is not None else self.fetch_message_attributes | 
|  | 174 | +        attrs = attributes if attributes is not None else self.message_system_attribute_names | 
|  | 175 | +        msg_attr_names = self.message_attribute_names if self.message_attribute_names else None | 
| 164 | 176 | 
 | 
| 165 | 177 |         if visibility_timeout: | 
| 166 | 178 |             params['VisibilityTimeout'] = visibility_timeout | 
| 167 | 179 |         if attrs: | 
| 168 |  | -            proto_params['json'].update({'AttributeNames': list(attrs)}) | 
| 169 |  | -            proto_params['query'].update(_query_object_encode({'AttributeName': list(attrs)})) | 
|  | 180 | +            proto_params['json'].update({'MessageSystemAttributeNames': list(attrs)}) | 
|  | 181 | +            proto_params['query'].update(_query_object_encode({'MessageSystemAttributeName': list(attrs)})) | 
|  | 182 | +        if msg_attr_names: | 
|  | 183 | +            proto_params['json'].update({'MessageAttributeNames': list(msg_attr_names)}) | 
|  | 184 | +            proto_params['query'].update(_query_object_encode({'MessageAttributeNames': list(msg_attr_names)})) | 
| 170 | 185 |         if wait_time_seconds is not None: | 
| 171 | 186 |             params['WaitTimeSeconds'] = wait_time_seconds | 
| 172 | 187 | 
 | 
|  | 
0 commit comments