diff --git a/awsimple/__version__.py b/awsimple/__version__.py index 81dcf89..eb0f108 100644 --- a/awsimple/__version__.py +++ b/awsimple/__version__.py @@ -1,7 +1,7 @@ __application_name__ = "awsimple" __title__ = __application_name__ __author__ = "abel" -__version__ = "7.2.1" +__version__ = "7.3.0" __author_email__ = "j@abel.co" __url__ = "https://github.com/jamesabel/awsimple" __download_url__ = "https://github.com/jamesabel/awsimple" diff --git a/awsimple/aws.py b/awsimple/aws.py index 8fe8535..deb7665 100644 --- a/awsimple/aws.py +++ b/awsimple/aws.py @@ -1,4 +1,5 @@ import os +import threading from typing import Union, Any from logging import getLogger @@ -8,16 +9,51 @@ from botocore.credentials import Credentials from awsimple import __application_name__, is_mock, is_using_localstack +from .exceptions import AWSimpleException # noqa: F401 (re-exported for backwards compatibility) log = getLogger(__application_name__) +# Process-wide moto mock state. Multiple AWSAccess instances share one moto mock and one saved copy of the AWS +# environment variables, so instance creation/deletion order can't corrupt the environment or stop a mock another +# instance still needs. +_moto_lock = threading.Lock() +_moto_state = {"count": 0, "mock": None, "saved_env": {}} # type: dict +_mock_env_keys = ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SECURITY_TOKEN", "AWS_SESSION_TOKEN"] -class AWSimpleException(Exception): - pass + +def _moto_acquire(): + with _moto_lock: + if _moto_state["count"] == 0: + for aws_key in _mock_env_keys: + _moto_state["saved_env"][aws_key] = os.environ.get(aws_key) # will be None if not set + os.environ[aws_key] = "testing" + + from moto import mock_aws + + _moto_state["mock"] = mock_aws() + _moto_state["mock"].start() + _moto_state["count"] += 1 + + +def _moto_release(): + with _moto_lock: + if _moto_state["count"] > 0: + _moto_state["count"] -= 1 + if _moto_state["count"] == 0: + if _moto_state["mock"] is not None: + _moto_state["mock"].stop() + _moto_state["mock"] = None + for aws_key, value in _moto_state["saved_env"].items(): + if value is None: + os.environ.pop(aws_key, None) + else: + os.environ[aws_key] = value + _moto_state["saved_env"] = {} def boto_error_to_string(boto_error) -> Union[str, None]: - if (response := boto_error.response) is None: + # BotoCoreError subclasses (e.g. HTTPClientError) don't have a .response attribute - only ClientError does + if (response := getattr(boto_error, "response", None)) is None: most_recent_error = str(boto_error) else: if (response_error := response.get("Error")) is None: @@ -61,8 +97,7 @@ def __init__( # string representation of AWS most recent error code self.most_recent_error = None # type: Union[str, None] - self._moto_mock = None - self._aws_keys_save = {} + self._is_mocked = False # use keys in AWS config # https://docs.aws.amazon.com/cli/latest/userguide/cli-config-files.html @@ -75,14 +110,8 @@ def __init__( self.client = None # type: Any if is_mock(): # moto mock AWS - for aws_key in ["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", "AWS_SECURITY_TOKEN", "AWS_SESSION_TOKEN"]: - self._aws_keys_save[aws_key] = os.environ.get(aws_key) # will be None if not set - os.environ[aws_key] = "testing" - - from moto import mock_aws - - self._moto_mock = mock_aws() - self._moto_mock.start() + _moto_acquire() + self._is_mocked = True region = "us-east-1" if self.resource_name == "logs" or self.resource_name is None: # logs don't have a resource @@ -100,7 +129,11 @@ def __init__( self.aws_access_key_id = "test" self.aws_secret_access_key = "test" self.region_name = "us-west-2" - if self.resource_name is not None: + if self.resource_name is None: + # just the session, but not the client or resource + self.client = None + self.resource = None + else: if self.resource_name == "logs": # logs don't have resource self.resource = None @@ -161,14 +194,23 @@ def get_account_id(self): def test(self) -> bool: """ - Basic connection/capability test + Basic connection/credentials test. Calls STS GetCallerIdentity, which requires no special permissions but does require valid credentials. + Returns False if the credentials are invalid. Configuration errors (e.g. botocore's ProfileNotFound) still raise. - :return: True if connection OK + :return: True if connection OK, False if the credentials are invalid """ + from botocore.exceptions import ClientError, NoCredentialsError # import here to facilitate mocking - resources = self.session.get_available_resources() # boto3 will throw an error if there's an issue here - if self.resource_name is not None and self.resource_name not in resources: - raise PermissionError(self.resource_name) # we don't have permission to the specified resource + if is_using_localstack(): + sts_client = self.session.client("sts", endpoint_url=self._get_localstack_endpoint_url()) + else: + sts_client = self.session.client("sts") + try: + sts_client.get_caller_identity() # raises if credentials are invalid + except (ClientError, NoCredentialsError) as e: + log.info(f"{self.profile_name=} {e}") + self.most_recent_error = boto_error_to_string(e) + return False return True # if we got here, we were successful def is_mocked(self) -> bool: @@ -177,20 +219,15 @@ def is_mocked(self) -> bool: :return: True if mocked """ - return self._moto_mock is not None + return self._is_mocked def clear_most_recent_error(self): self.most_recent_error = None def __del__(self): - if self._moto_mock is not None: - # if mocking, put everything back - - for aws_key, value in self._aws_keys_save.items(): - if value is None: - del os.environ[aws_key] - else: - os.environ[aws_key] = value - - self._moto_mock.stop() - self._moto_mock = None # mock is "done" + if getattr(self, "_is_mocked", False): + self._is_mocked = False + try: + _moto_release() + except Exception: # noqa: S110 - interpreter may be shutting down + pass diff --git a/awsimple/cache.py b/awsimple/cache.py index f8080c9..aef1bb2 100644 --- a/awsimple/cache.py +++ b/awsimple/cache.py @@ -51,7 +51,7 @@ def lru_cache_write(new_data: Union[Path, bytes], cache_dir: Path, cache_file_na wrote_to_cache = False try: - max_free_absolute = max_free_portion * get_disk_free() if max_free_portion is not None else None + max_free_absolute = max_free_portion * get_disk_free(cache_dir) if max_free_portion is not None else None values = [v for v in [max_free_absolute, max_absolute_cache_size] if v is not None] max_cache_size = min(values) if len(values) > 0 else None log.info(f"{max_cache_size=}") @@ -81,6 +81,8 @@ def lru_cache_write(new_data: Union[Path, bytes], cache_dir: Path, cache_file_na least_recently_used_access_time = None least_recently_used_size = None for file_path in cache_dir.rglob("*"): + if not file_path.is_file(): + continue # only files are evicted (rglob can also return directories) access_time = os.path.getatime(file_path) if least_recently_used_path is None or least_recently_used_access_time is None or access_time < least_recently_used_access_time: least_recently_used_path = file_path @@ -91,9 +93,8 @@ def lru_cache_write(new_data: Union[Path, bytes], cache_dir: Path, cache_file_na log.debug(f"evicting {least_recently_used_path=} {least_recently_used_access_time=} {least_recently_used_size=}") least_recently_used_path.unlink() if least_recently_used_size is None: - AWSimpleException(f"{least_recently_used_size=}") - else: - overage -= least_recently_used_size + raise AWSimpleException(f"{least_recently_used_size=}") + overage -= least_recently_used_size if overage == starting_overage: # tried to free up space but were unsuccessful, so give up diff --git a/awsimple/dynamodb.py b/awsimple/dynamodb.py index ae82a78..35d98b7 100644 --- a/awsimple/dynamodb.py +++ b/awsimple/dynamodb.py @@ -8,6 +8,7 @@ from collections import OrderedDict, defaultdict, namedtuple import datetime from pathlib import Path +import os from os.path import getsize, getmtime from typing import List, Union, Any, Type, Dict, Callable from pprint import pformat @@ -16,12 +17,11 @@ from enum import Enum import decimal from decimal import Decimal -from functools import lru_cache from collections.abc import Iterable from logging import getLogger -from botocore.exceptions import EndpointConnectionError, ClientError +from botocore.exceptions import ClientError from boto3.dynamodb.conditions import Key from typeguard import typechecked from dictim import dictim # type: ignore @@ -161,9 +161,8 @@ def dict_to_dynamodb(input_value: Any, convert_images: bool = True, raise_except # converts tuple to list resp = [dict_to_dynamodb(v, convert_images, raise_exception) for v in input_value] elif type(input_value) is str: - # DynamoDB does not allow zero length strings - if len(input_value) > 0: - resp = input_value + # DynamoDB supports zero length strings for non-key attributes (since May 2020) + resp = input_value elif type(input_value) is bool or input_value is None or type(input_value) is decimal.Decimal: resp = input_value # native DynamoDB types elif type(input_value) is float or type(input_value) is int: @@ -239,6 +238,10 @@ def __init__(self, table_name: Union[str, None] = None, **kwargs): self.cache_hit = False self.secondary_index_postfix = "-index" + # per-instance caches (do not use functools.lru_cache on methods - it holds a global reference to self, so instances are never garbage collected) + self._primary_keys_cache = None # type: Union[Dict[KeyType, str], None] + self._secondary_indexes_cache = None # type: Union[List[Dict[KeyType, str]], None] + self.table_name = table_name # can be None (the default) if we're only doing things that don't require a table name such as get_table_names() # avoid recursion @@ -322,25 +325,19 @@ def scan_table(self) -> list: assert self.resource is not None table = self.resource.Table(self.table_name) + # connection errors are deliberately not caught here - returning partial results as if they were the complete table would silently corrupt caller data (and caches) more_to_evaluate = True exclusive_start_key = None while more_to_evaluate: - try: - if exclusive_start_key is None: - response = table.scan() - else: - response = table.scan(ExclusiveStartKey=exclusive_start_key) - except EndpointConnectionError as e: - log.warning(e) - response = None + if exclusive_start_key is None: + response = table.scan() + else: + response = table.scan(ExclusiveStartKey=exclusive_start_key) + items.extend(response["Items"]) + if "LastEvaluatedKey" not in response: more_to_evaluate = False - - if response is not None: - items.extend(response["Items"]) - if "LastEvaluatedKey" not in response: - more_to_evaluate = False - else: - exclusive_start_key = response["LastEvaluatedKey"] + else: + exclusive_start_key = response["LastEvaluatedKey"] log.info(f"read {len(items)} items from {self.table_name}") @@ -383,9 +380,9 @@ def scan_table_cached(self, invalidate_cache: bool = False) -> list: self.cache_hit = False now = time.time() try: - if now <= getmtime(str(cache_file_path)) + self.cache_life: + cache_file_mtime = getmtime(str(cache_file_path)) + if now <= cache_file_mtime + self.cache_life: # cache file exists and is current, see if it has expired - cache_file_mtime = getmtime(str(cache_file_path)) if self.metadata_table is None: table_mtime_f = None else: @@ -397,10 +394,12 @@ def scan_table_cached(self, invalidate_cache: bool = False) -> list: # determine if table has been updated since local cache file was written # (assumes the clock of the system that wrote the table is in sync with the clock of this system within the clock skew) self.cache_hit = table_mtime_f is not None and table_mtime_f + get_accommodated_clock_skew() <= cache_file_mtime - with open(cache_file_path, "rb") as f: - log.info(f"{self.table_name=},{cache_file_path=}") - table_data = pickle.load(f) - log.debug(f"done reading {cache_file_path=}") + if self.cache_hit: + # only unpickle if we're actually going to use the cached data + with open(cache_file_path, "rb") as f: + log.info(f"{self.table_name=},{cache_file_path=}") + table_data = pickle.load(f) + log.debug(f"done reading {cache_file_path=}") except FileNotFoundError: self.cache_hit = False # simple cache miss except (EOFError, OSError, pickle.PickleError) as e: @@ -416,6 +415,10 @@ def scan_table_cached(self, invalidate_cache: bool = False) -> list: # update local data cache with open(cache_file_path, "wb") as f: pickle.dump(table_data, f) + # stamp the cache file with the same clock used for the metadata table mtimes - on Windows the filesystem's own timestamp can lag + # time.time() by a clock tick (~15 mS), which can make a just-written cache file appear older than the table's metadata mtime + write_time = time.time() + os.utime(cache_file_path, (write_time, write_time)) except (DynamoDBTableNotFound, self.client.exceptions.ResourceNotFoundException) as e: log.debug(f"{self.table_name=},{e}") table_data = [] @@ -453,9 +456,9 @@ def create_table( partition_key: str, sort_key: Union[str, None] = None, secondary_index: Union[str, None] = None, - partition_key_type: Union[Type[str], Type[int], Type[bool]] = str, - sort_key_type: Union[Type[str], Type[int], Type[bool]] = str, - secondary_key_type: Union[Type[str], Type[int], Type[bool]] = str, + partition_key_type: Union[Type[str], Type[int], Type[bytes]] = str, + sort_key_type: Union[Type[str], Type[int], Type[bytes]] = str, + secondary_key_type: Union[Type[str], Type[int], Type[bytes]] = str, ) -> bool: """ Create a DynamoDB table. @@ -463,14 +466,14 @@ def create_table( :param partition_key: DynamoDB partition key (AKA hash key) :param sort_key: DynamoDB sort key :param secondary_index: secondary index key - :param partition_key_type: partition key type of str, int, bool (str default) - :param sort_key_type: sort key type of str, int, bool (str default) - :param secondary_key_type: secondary key type of str, int, bool (str default) + :param partition_key_type: partition key type of str, int, bytes (str default) + :param sort_key_type: sort key type of str, int, bytes (str default) + :param secondary_key_type: secondary key type of str, int, bytes (str default) :return: True if table created """ def add_key(k, t, kt): - assert t in ("S", "N", "B") # DynamoDB key types (string, number, bool) + assert t in ("S", "N", "B") # DynamoDB key types (string, number, binary) assert kt in ("HASH", "RANGE") definition = {"AttributeName": k, "AttributeType": t} schema = {"AttributeName": k, "KeyType": kt} @@ -481,8 +484,11 @@ def type_to_attribute_type(t): ts = "S" elif t is int: ts = "N" + elif t is bytes: + ts = "B" # binary elif t is bool: - ts = "B" + # DynamoDB "B" is *binary*, not boolean - booleans are not a valid DynamoDB key type + raise ValueError("DynamoDB does not support boolean key attributes - key types are str ('S'), int ('N'), and bytes ('B' binary)") else: raise ValueError(t) return ts @@ -538,45 +544,44 @@ def _get_keys_from_schema(self, table_schema: List) -> Dict[KeyType, str]: key_schema[aws_name_to_key_type[table_key_schema["KeyType"]]] = table_key_schema["AttributeName"] return key_schema - @lru_cache() def get_primary_keys_dict(self) -> Dict[KeyType, str]: """ - Get the table's primary keys. Raise TableNotFound if table does not exist. + Get the table's primary keys (cached per instance). Raise TableNotFound if table does not exist. :return: a dict with the primary key partition key and (optionally) sort key """ - assert self.resource is not None - try: - table = self.resource.Table(self.table_name) - key_schema = table.key_schema - except self.client.exceptions.ResourceNotFoundException: - raise DynamoDBTableNotFound(str(self.table_name)) - keys = self._get_keys_from_schema(key_schema) - return keys + if self._primary_keys_cache is None: + assert self.resource is not None + try: + table = self.resource.Table(self.table_name) + key_schema = table.key_schema + except self.client.exceptions.ResourceNotFoundException: + raise DynamoDBTableNotFound(str(self.table_name)) + self._primary_keys_cache = self._get_keys_from_schema(key_schema) + return self._primary_keys_cache - @lru_cache() def get_primary_partition_key(self) -> str: primary_keys = self.get_primary_keys_dict() return primary_keys[KeyType.partition] - @lru_cache() def get_primary_sort_key(self) -> Union[str, None]: primary_keys = self.get_primary_keys_dict() return primary_keys.get(KeyType.sort) - @lru_cache() def get_secondary_indexes(self) -> List[Dict[KeyType, str]]: """ - Get the secondary indexes as a list of dicts with the key as the KeyType. + Get the secondary indexes as a list of dicts with the key as the KeyType (cached per instance). :return: list of dicts with secondary keys """ - secondary_indexes = [] - assert self.resource is not None - for table_secondary_index in self.resource.Table(self.table_name).global_secondary_indexes: - secondary_indexes.append(self._get_keys_from_schema(table_secondary_index["KeySchema"])) - return secondary_indexes + if self._secondary_indexes_cache is None: + secondary_indexes = [] + assert self.resource is not None + for table_secondary_index in self.resource.Table(self.table_name).global_secondary_indexes: + secondary_indexes.append(self._get_keys_from_schema(table_secondary_index["KeySchema"])) + self._secondary_indexes_cache = secondary_indexes + return self._secondary_indexes_cache def _query(self, comp: str, *args) -> List[dict]: """ @@ -839,7 +844,8 @@ def delete_item(self, partition_key: Union[str, None] = None, partition_value: U if sort_key is not None: key[sort_key] = sort_value table.delete_item(Key=key) - self.metadata_table.update_table_mtime() + if self.metadata_table is not None: + self.metadata_table.update_table_mtime() # cant' do a @typechecked() since optional item requires a single type def upsert_item( @@ -866,22 +872,32 @@ def upsert_item( sort_key = self.get_primary_sort_key() if item is None: - AWSimpleException(f"{item=}") - else: - assert self.resource is not None - table = self.resource.Table(self.table_name) - key = {partition_key: partition_value} # type: dict[str, Any] - if sort_key is not None: - key[sort_key] = sort_value + raise AWSimpleException(f"no item given ({item=})") - # create the required boto3 strings and dict for the update - update_expression = "SET " - expression_attribute_values = {} - for k, v in item.items(): - update_expression += f"{k} = :{k} " - expression_attribute_values[f":{k}"] = v + assert self.resource is not None + table = self.resource.Table(self.table_name) + key = {partition_key: partition_value} # type: dict[str, Any] + if sort_key is not None: + key[sort_key] = sort_value - table.update_item(Key=key, UpdateExpression=update_expression, ExpressionAttributeValues=expression_attribute_values) + if len(item) == 0: + raise AWSimpleException(f"empty item ({item=})") + + # create the required boto3 strings and dicts for the update + # (use expression attribute name placeholders so DynamoDB reserved words like "name" or "status" work, and separate clauses with commas) + update_clauses = [] + expression_attribute_names = {} + expression_attribute_values = {} + for attribute_number, (k, v) in enumerate(item.items()): + name_placeholder = f"#a{attribute_number}" + value_placeholder = f":v{attribute_number}" + expression_attribute_names[name_placeholder] = k + expression_attribute_values[value_placeholder] = v + update_clauses.append(f"{name_placeholder} = {value_placeholder}") + update_expression = "SET " + ", ".join(update_clauses) + + table.update_item(Key=key, UpdateExpression=update_expression, ExpressionAttributeNames=expression_attribute_names, ExpressionAttributeValues=expression_attribute_values) + if self.metadata_table is not None: self.metadata_table.update_table_mtime() def delete_all_items(self) -> int: @@ -908,7 +924,8 @@ def delete_all_items(self) -> int: key[sort_key] = item[sort_key] table.delete_item(Key=key) count += 1 - self.metadata_table.update_table_mtime() + if self.metadata_table is not None: + self.metadata_table.update_table_mtime() return count @typechecked() @@ -983,10 +1000,14 @@ def update_table_mtime(self): self.mtime_human_string: datetime.datetime.fromtimestamp(m_time).astimezone().isoformat(), } ) - self.put_item(item=item) + try: + self.put_item(item=item) + except DynamoDBTableNotFound: + self.create_metadata_table() + self.put_item(item=item) @typechecked() - def get_table_mtime_f(self) -> Union[float | None]: + def get_table_mtime_f(self) -> Union[float, None]: """ Get a table's mtime from the metadata table. :return: table's mtime as a float or None if table hasn't been written to diff --git a/awsimple/dynamodb_miv.py b/awsimple/dynamodb_miv.py index 8d1cc5d..d3038e8 100644 --- a/awsimple/dynamodb_miv.py +++ b/awsimple/dynamodb_miv.py @@ -48,8 +48,8 @@ def create_table( # type: ignore self, partition_key: str, secondary_index: Union[str, None] = None, - partition_key_type: Union[Type[str], Type[int], Type[bool]] = str, - secondary_key_type: Union[Type[str], Type[int], Type[bool]] = str, + partition_key_type: Union[Type[str], Type[int], Type[bytes]] = str, + secondary_key_type: Union[Type[str], Type[int], Type[bytes]] = str, ) -> bool: return super().create_table(partition_key, miv_string, secondary_index, partition_key_type, int, secondary_key_type) @@ -64,30 +64,46 @@ def put_item(self, item: dict, time_us: Union[int, None] = None): assert self.resource is not None table = self.resource.Table(self.table_name) + new_item = deepcopy(item) + # Determine new miv. The miv is an int to avoid comparison or specification problems that can arise with floats. For example, when it comes time to delete an item. if time_us is None: - # get the miv for the existing entries partition_key = self.get_primary_partition_key() partition_value = item[partition_key] - try: - existing_most_senior_item = self.get_most_senior_item(partition_key, partition_value) - existing_miv_ui = existing_most_senior_item[miv_string] - except DBItemNotFound: - existing_miv_ui = None - - current_time_us = get_time_us() - if existing_miv_ui is None or current_time_us > existing_miv_ui: - new_miv_ui = current_time_us - else: - # the prior writer seems to be from the future (from our perspective), so just increment the existing miv by the smallest increment and go with that - new_miv_ui = existing_miv_ui + 1 + + # A conditional put makes the read-compute-write sequence safe against concurrent writers: if another writer takes our miv first, the condition + # fails and we recompute rather than silently overwriting their item. + retries_remaining = 10 + while True: + # get the miv for the existing entries + try: + existing_most_senior_item = self.get_most_senior_item(partition_key, partition_value) + existing_miv_ui = existing_most_senior_item[miv_string] + except DBItemNotFound: + existing_miv_ui = None + + current_time_us = get_time_us() + if existing_miv_ui is None or current_time_us > existing_miv_ui: + new_miv_ui = current_time_us + else: + # the prior writer seems to be from the future (from our perspective), so just increment the existing miv by the smallest increment and go with that + new_miv_ui = existing_miv_ui + 1 + + new_item[miv_string] = new_miv_ui + try: + table.put_item(Item=new_item, ConditionExpression="attribute_not_exists(#miv)", ExpressionAttributeNames={"#miv": miv_string}) + break + except self.client.exceptions.ConditionalCheckFailedException: + retries_remaining -= 1 + if retries_remaining <= 0: + raise else: - new_miv_ui = time_us + new_item[miv_string] = time_us + table.put_item(Item=new_item) - # make the new item with the new miv and put it into the DB table - new_item = deepcopy(item) - new_item[miv_string] = new_miv_ui - table.put_item(Item=new_item) + # keep the metadata table's mtime current so scan_table_cached() invalidates properly (DynamoDBAccess.put_item does this too, but we write directly to the table here) + if self.metadata_table is not None: + self.metadata_table.update_table_mtime() @typechecked() def get_most_senior_item(self, partition_key: str, partition_value: Union[str, int]) -> dict: diff --git a/awsimple/exceptions.py b/awsimple/exceptions.py index d44c188..0fc451d 100644 --- a/awsimple/exceptions.py +++ b/awsimple/exceptions.py @@ -1,5 +1,11 @@ -class AWSimpleExceptionBase(Exception): - """Base exception for AWSimple errors.""" +class AWSimpleException(Exception): + """Base exception for all AWSimple errors.""" + + pass + + +class AWSimpleExceptionBase(AWSimpleException): + """Deprecated intermediate base exception, retained for backwards compatibility. Use AWSimpleException.""" pass diff --git a/awsimple/logs.py b/awsimple/logs.py index 6b64df0..fa17bb4 100644 --- a/awsimple/logs.py +++ b/awsimple/logs.py @@ -3,6 +3,8 @@ from pathlib import Path from datetime import datetime +from botocore.exceptions import ClientError + from .aws import AWSAccess from .platform import get_user_name, get_computer_name @@ -42,6 +44,11 @@ def put(self, message: str): self.client.create_log_stream(logGroupName=self.log_group, logStreamName=self.get_stream_name()) self._put(message) + def _put_log_events(self, stream_name: str, log_events: list, sequence_token: Union[str, None]): + if sequence_token is None: + return self.client.put_log_events(logGroupName=self.log_group, logStreamName=stream_name, logEvents=log_events) + return self.client.put_log_events(logGroupName=self.log_group, logStreamName=stream_name, logEvents=log_events, sequenceToken=sequence_token) + def _put(self, message: str): """ Perform the put log event. Internal method to enable try/except in the regular .put() method. @@ -61,15 +68,16 @@ def _put(self, message: str): # timestamp defined by AWS to be mS since epoch log_events = [{"timestamp": int(round(time.time() * 1000)), "message": message}] try: - if self._upload_sequence_token is None: - put_response = self.client.put_log_events(logGroupName=self.log_group, logStreamName=stream_name, logEvents=log_events) - else: - put_response = self.client.put_log_events(logGroupName=self.log_group, logStreamName=stream_name, logEvents=log_events, sequenceToken=self._upload_sequence_token) + put_response = self._put_log_events(stream_name, log_events, self._upload_sequence_token) except self.client.exceptions.InvalidSequenceTokenException as e: - # something went terribly wrong in logging, so write what happened somewhere safe - with Path(Path.home(), "awsimple_exception.txt").open("w") as f: - f.write(f"{datetime.now().astimezone().isoformat()},{self.log_group=},{stream_name=},{self._upload_sequence_token=},{e}\n") - put_response = None + # our token is stale - retry once with the token AWS says it expects, so the message isn't lost + try: + put_response = self._put_log_events(stream_name, log_events, e.response.get("expectedSequenceToken")) + except ClientError as retry_exception: + # something went terribly wrong in logging, so write what happened somewhere safe (append so prior records aren't lost) + with Path(Path.home(), "awsimple_exception.txt").open("a") as f: + f.write(f"{datetime.now().astimezone().isoformat()},{self.log_group=},{stream_name=},{self._upload_sequence_token=},{e},{retry_exception}\n") + put_response = None if put_response is None: self._upload_sequence_token = None diff --git a/awsimple/mock.py b/awsimple/mock.py index 4fe6463..6690d75 100644 --- a/awsimple/mock.py +++ b/awsimple/mock.py @@ -1,5 +1,4 @@ import os -from functools import cache from tobool import to_bool_strict @@ -7,7 +6,6 @@ use_localstack_env_var = "AWSIMPLE_USE_LOCALSTACK" -@cache def is_mock() -> bool: """ Is using moto mock? @@ -16,7 +14,6 @@ def is_mock() -> bool: return to_bool_strict(os.environ.get(use_moto_mock_env_var, "0")) -@cache def is_using_localstack() -> bool: """ Is using localstack? diff --git a/awsimple/pubsub.py b/awsimple/pubsub.py index 7de3240..ae26acc 100644 --- a/awsimple/pubsub.py +++ b/awsimple/pubsub.py @@ -9,7 +9,7 @@ from threading import Thread, Event from queue import Queue from queue import Empty -from logging import Logger +from logging import getLogger import json from typeguard import typechecked @@ -18,11 +18,11 @@ from .sns import SNSAccess from .sqs import SQSPollAccess, get_all_sqs_queues -from .dynamodb import _DynamoDBMetadataTable +from .dynamodb import _DynamoDBMetadataTable, DBItemNotFound, DynamoDBTableNotFound from .platform import get_node_name from .__version__ import __application_name__ -log = Logger(__application_name__) +log = getLogger(__application_name__) # getLogger (not a raw Logger instance) so the application's logging configuration applies queue_timeout = timedelta(days=30).total_seconds() @@ -42,11 +42,14 @@ def remove_old_queues( if len(channel) < 2: # avoid deleting all queues log.warning(f"blank channel ({channel=}) - not deleting any queues") return removed - for sqs_queue_name in get_all_sqs_queues(channel): + for sqs_queue_name in get_all_sqs_queues(channel, profile_name=profile_name, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, region_name=region_name): sqs_metadata = _DynamoDBMetadataTable( SQS_NAME, sqs_queue_name, profile_name=profile_name, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, region_name=region_name ) - mtime = sqs_metadata.get_table_mtime_f() + try: + mtime = sqs_metadata.get_table_mtime_f() + except (DBItemNotFound, DynamoDBTableNotFound): + mtime = None # queue has no metadata entry (e.g. it was never used) - leave it alone if mtime is not None and time.time() - mtime > queue_timeout: sqs = SQSPollAccess(sqs_queue_name, profile_name=profile_name, aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key, region_name=region_name) try: @@ -80,22 +83,24 @@ def _connect_sns_to_sqs(sqs: SQSPollAccess, sns: SNSAccess) -> None: subscription = topic.subscribe(Protocol="sqs", Endpoint=sqs_arn) log.info(f"Subscribed {sqs.queue_name} to topic {topic_arn}. Subscription ARN: {subscription.arn}") - # Update queue policy to allow SNS -> SQS - policy = { - "Version": "2012-10-17", - "Id": "sns-sqs-subscription-policy", - "Statement": [ - { - "Sid": "Allow-SNS-SendMessage", - "Effect": "Allow", - "Principal": {"Service": "sns.amazonaws.com"}, - "Action": "sqs:SendMessage", - "Resource": sqs_arn, - "Condition": {"ArnEquals": {"aws:SourceArn": topic_arn}}, - } - ], + # Update queue policy to allow SNS -> SQS (merge into any existing policy - replacing it would revoke other topics' permissions) + statement = { + "Sid": f"AllowSNSSendMessage{topic_arn.split(':')[-1]}", # Sid must be unique within the policy, so include the topic name + "Effect": "Allow", + "Principal": {"Service": "sns.amazonaws.com"}, + "Action": "sqs:SendMessage", + "Resource": sqs_arn, + "Condition": {"ArnEquals": {"aws:SourceArn": topic_arn}}, } assert sqs.queue is not None + existing_policy_string = sqs.client.get_queue_attributes(QueueUrl=sqs.queue.url, AttributeNames=["Policy"]).get("Attributes", {}).get("Policy") + if existing_policy_string is None: + policy = {"Version": "2012-10-17", "Id": "sns-sqs-subscription-policy", "Statement": [statement]} + else: + policy = json.loads(existing_policy_string) + statements = policy.setdefault("Statement", []) + if statement not in statements: + statements.append(statement) sqs.queue.set_attributes(Attributes={"Policy": json.dumps(policy)}) log.debug(f"Queue {sqs.queue_name} policy updated to allow topic {topic_arn}.") @@ -115,11 +120,21 @@ def __init__(self, sqs: SQSPollAccess, new_event: Event) -> None: def run(self): while not self._exit_event.is_set(): - messages = self._sqs.receive_messages() # long poll + try: + messages = self._sqs.receive_messages() # long poll + except Exception as e: + # don't let a transient AWS error silently kill the polling thread + log.warning(f"{self._sqs.queue_name=} receive failed : {e}") + time.sleep(1.0) + continue for message in messages: - parsed = json.loads(message.message) - self.sub_queue.put(parsed["Message"]) - self._new_event.set() + try: + parsed = json.loads(message.message) + self.sub_queue.put(parsed["Message"]) + except (json.JSONDecodeError, KeyError, TypeError) as e: + log.warning(f"{self._sqs.queue_name=} malformed message : {e}") + else: + self._new_event.set() def request_exit(self): self._exit_event.set() @@ -134,7 +149,8 @@ def make_name_aws_safe(*args: str) -> str: :return: AWS safe name """ - base36 = strif.hash_string("".join(args)).base36.strip() + # join with a separator so e.g. ("a", "b") and ("ab",) hash differently (the separator can't appear in the hash output, avoiding cross-boundary collisions) + base36 = strif.hash_string("\x1f".join(args)).base36.strip() assert 30 <= len(base36) <= 31 return base36 @@ -164,7 +180,8 @@ def __init__( """ self.channel = AWS_RESOURCE_PREFIX + make_name_aws_safe(channel) # prefix with ps (pubsub) to avoid collisions with other uses of SNS topics and SQS queues self.node_name = get_node_name() if node_name is None else node_name - self.sqs_queue_name = AWS_RESOURCE_PREFIX + make_name_aws_safe(self.channel, self.node_name) + # queue name is the channel name plus a node hash, so all of a channel's queues share the channel name as a prefix (this is what lets remove_old_queues() find them) + self.sqs_queue_name = self.channel + make_name_aws_safe(self.node_name) self.sub_callback = sub_callback self.use_sub_queue = use_sub_queue @@ -183,6 +200,13 @@ def __init__( super().__init__(daemon=True) # make daemon so an instance of this thread exits when the main program exits def run(self): + try: + self._run() + except Exception: + # a daemon thread that dies silently leaves publish() queueing into the void, so at least make the failure visible + log.exception(f"pubsub thread failed,{self.channel=}") + + def _run(self): sns = SNSAccess( self.channel, @@ -218,7 +242,14 @@ def run(self): _connect_sns_to_sqs(sqs, sns) sqs_metadata.update_table_mtime() # update SQS use time (the existing infrastructure calls it a "table", but we're using it for the SQS queue) - remove_old_queues(self.channel) # clean up old queues + # clean up old queues (using the same credentials as this instance) + remove_old_queues( + self.channel, + profile_name=self.profile_name, + aws_access_key_id=self.aws_access_key_id, + aws_secret_access_key=self.aws_secret_access_key, + region_name=self.region_name, + ) if self.sub_callback is None and not self.use_sub_queue: # not being used as a subscriber @@ -229,30 +260,9 @@ def run(self): while not self._exit_event.is_set(): - # pub - try: - message = self._pub_queue.get(False) - message_string = json.dumps(message) - sns.publish(message_string) - except Empty: - pass - except RuntimeError as e: - log.info(f"SQS,{self.sqs_queue_name=},{e}") - - # sub + self._drain_pub_queue(sns) if sqs_thread is not None: - try: - message_string = sqs_thread.sub_queue.get(False) - if self.use_sub_queue: - self._sub_queue.put(message_string) - if self.sub_callback is not None: - message = json.loads(message_string) - self.sub_callback(message) - sqs_metadata.update_table_mtime() - except Empty: - pass # no message - except RuntimeError as e: - log.info(f"SQS,{self.sqs_queue_name=},{e}") + self._drain_sub_queue(sqs_thread, sqs_metadata) if self._new_event.wait(self._new_event_wait_time): # timeout in case the new event technique fails self._new_event.clear() @@ -263,6 +273,42 @@ def run(self): if sqs_thread.is_alive(): log.error("sqs_thread did not exit cleanly") + def _drain_pub_queue(self, sns: SNSAccess) -> None: + # drain all queued messages (the "new" event is binary, so handling only one message per wait cycle would throttle bursts to one message per timeout) + while True: + try: + message = self._pub_queue.get(False) + except Empty: + break + try: + message_string = json.dumps(message) + sns.publish(message_string) + except Exception as e: + log.warning(f"SNS publish failed,{self.channel=},{e}") + + def _drain_sub_queue(self, sqs_thread: _SubscriptionThread, sqs_metadata: _DynamoDBMetadataTable) -> None: + # drain all received messages + got_message = False + while True: + try: + message_string = sqs_thread.sub_queue.get(False) + except Empty: + break + got_message = True + try: + if self.use_sub_queue: + self._sub_queue.put(message_string) + if self.sub_callback is not None: + message = json.loads(message_string) + self.sub_callback(message) + except Exception as e: + log.warning(f"SQS,{self.sqs_queue_name=},{e}") + if got_message: + try: + sqs_metadata.update_table_mtime() + except Exception as e: + log.warning(f"SQS metadata update failed,{self.sqs_queue_name=},{e}") + @typechecked() def publish(self, message: dict) -> None: """ diff --git a/awsimple/s3.py b/awsimple/s3.py index 89b269d..6b3a6c1 100644 --- a/awsimple/s3.py +++ b/awsimple/s3.py @@ -13,7 +13,6 @@ import json from logging import getLogger -import boto3 from botocore.client import Config from botocore.exceptions import ClientError, EndpointConnectionError, ConnectionClosedError, SSLError from boto3.s3.transfer import TransferConfig @@ -23,7 +22,16 @@ from hashy import get_string_sha512, get_file_sha512, get_bytes_sha512, get_dls_sha512 from yasf import sf -from awsimple import CacheAccess, __application_name__, lru_cache_write, AWSimpleException, convert_serializable_special_cases, S3BucketAlreadyExistsNotOwnedByYou +from awsimple import ( + CacheAccess, + __application_name__, + lru_cache_write, + AWSimpleException, + convert_serializable_special_cases, + S3BucketAlreadyExistsNotOwnedByYou, + is_using_localstack, + boto_error_to_string, +) # Use this project's name as a prefix to avoid string collisions. Use dashes instead of underscore since that's AWS's convention. sha512_string = f"{__application_name__}-sha512" @@ -109,6 +117,7 @@ def __init__(self, bucket_name: Union[str, None] = None, **kwargs): self.retry_count = 10 self.public_readable = False self.download_status = S3DownloadStatus() + self._bucket_region = None # type: Union[str, None] # lazily determined and cached super().__init__(resource_name="s3", **kwargs) def get_s3_transfer_config(self) -> TransferConfig: @@ -194,7 +203,7 @@ def upload(self, file_path: Union[str, Path], s3_key: str, force: bool = False) :param file_path: path to file to upload :param s3_key: S3 key :param force: True to force the upload, even if the file hash matches the S3 contents - :return: True if uploaded + :return: True if uploaded, False if the S3 object was already up to date. Raises AWSimpleException if the upload fails after all retries. """ log.info(f'S3 upload : "{file_path}" to {self.bucket_name}/{s3_key}') @@ -210,9 +219,9 @@ def upload(self, file_path: Union[str, Path], s3_key: str, force: bool = False) if self.object_exists(s3_key): s3_object_metadata = self.get_s3_object_metadata(s3_key) log.info(f"{s3_object_metadata=}") - if s3_object_metadata.get_sha512() is not None and file_sha512 is not None: - # use the hash provided by awsimple, if it exists - upload_flag = file_sha512 != s3_object_metadata.get_sha512() + if s3_object_metadata.sha512 is not None: + # use the hash provided by awsimple, if it exists (note that .get_sha512() never returns None - it synthesizes a substitute hash - so check .sha512 itself) + upload_flag = file_sha512 != s3_object_metadata.sha512 else: # if not, use mtime upload_flag = not isclose(file_mtime, s3_object_metadata.mtime.timestamp(), abs_tol=self.mtime_abs_tol) @@ -242,6 +251,9 @@ def upload(self, file_path: Union[str, Path], s3_key: str, force: bool = False) transfer_retry_count += 1 + if not uploaded_flag: + raise AWSimpleException(f"couldn't upload {file_path} to {self.bucket_name}/{s3_key} after {self.retry_count} attempts") + else: log.info(f"file hash of {file_sha512} is the same as is already on S3 and force={force} - not uploading") @@ -265,9 +277,9 @@ def upload_object_as_json(self, json_serializable_object: Union[List, Dict], s3_ if not force and self.object_exists(s3_key): s3_object_metadata = self.get_s3_object_metadata(s3_key) log.info(f"{s3_object_metadata=}") - if s3_object_metadata.get_sha512() is not None and json_sha512 is not None: - # use the hash provided by awsimple, if it exists - upload_flag = json_sha512 != s3_object_metadata.get_sha512() + if s3_object_metadata.sha512 is not None: + # use the hash provided by awsimple, if it exists (note that .get_sha512() never returns None - it synthesizes a substitute hash - so check .sha512 itself) + upload_flag = json_sha512 != s3_object_metadata.sha512 uploaded_flag = False if upload_flag: @@ -290,6 +302,9 @@ def upload_object_as_json(self, json_serializable_object: Union[List, Dict], s3_ transfer_retry_count += 1 time.sleep(self.retry_sleep_time) + if not uploaded_flag: + raise AWSimpleException(f"couldn't upload JSON to {self.bucket_name}/{s3_key} after {self.retry_count} attempts") + else: log.info(f"file hash of {json_sha512} is the same as is already on S3 and force={force} - not uploading") @@ -302,13 +317,13 @@ def download(self, s3_key: str, dest_path: Union[str, Path]) -> bool: :param s3_key: S3 key :param dest_path: destination file or directory path. If the path is a directory, the file will be downloaded to that directory with the same name as the S3 key. - :return: True if downloaded successfully + :return: True if downloaded successfully. Raises AWSimpleException if the download fails after all retries. """ if isinstance(dest_path, str): log.info(f"{dest_path} is not Path object. Non-Path objects will be deprecated in the future") + dest_path = Path(dest_path) - assert isinstance(dest_path, Path) if dest_path.is_dir(): dest_path = Path(dest_path, s3_key) @@ -334,6 +349,8 @@ def download(self, s3_key: str, dest_path: Union[str, Path]) -> bool: time.sleep(self.retry_sleep_time) transfer_retry_count += 1 log.debug(sf(transfer_retry_count=transfer_retry_count, success=success, bucket_name=self.bucket_name, s3_key=s3_key, dest_path=dest_path)) + if not success: + raise AWSimpleException(f"couldn't download {self.bucket_name}/{s3_key} to {dest_path} after {self.retry_count} attempts") return success @typechecked() @@ -369,7 +386,7 @@ def download_cached(self, s3_key: str, dest_path: Path) -> S3DownloadStatus: if not self.download_status.cache_hit: log.info(f"{self.bucket_name=}/{s3_key=} cache miss : {dest_path=} ({dest_path.absolute()})") - self.download(s3_key, dest_path) + self.download(s3_key, dest_path) # raises AWSimpleException on failure, so we don't write a bad cache entry or falsely report success self.cache_dir.mkdir(parents=True, exist_ok=True) self.download_status.cache_write = lru_cache_write(dest_path, self.cache_dir, sha512, self.cache_max_absolute, self.cache_max_of_free) self.download_status.success = True @@ -436,36 +453,36 @@ def get_s3_object_url(self, s3_key: str) -> str: :param s3_key: S3 key :return: object URL """ - bucket_location = self.client.get_bucket_location(Bucket=self.bucket_name) - location = bucket_location["LocationConstraint"] - url = f"https://{self.bucket_name}.s3-{location}.amazonaws.com/{s3_key}" + if self._bucket_region is None: + bucket_location = self.client.get_bucket_location(Bucket=self.bucket_name) + self._bucket_region = bucket_location["LocationConstraint"] or "us-east-1" # LocationConstraint is None for us-east-1 + url = f"https://{self.bucket_name}.s3.{self._bucket_region}.amazonaws.com/{s3_key}" return url @typechecked() def get_s3_object_metadata(self, s3_key: str) -> S3ObjectMetadata: """ - Get S3 object metadata + Get S3 object metadata. Raises AWSimpleException if the object does not exist. :param s3_key: S3 key - :return: S3ObjectMetadata or None if object does not exist + :return: S3ObjectMetadata """ - assert self.resource is not None - bucket_resource = self.resource.Bucket(self.bucket_name) - if self.object_exists(s3_key): - bucket_object = bucket_resource.Object(s3_key) - assert isinstance(self.bucket_name, str) # mainly for mypy - s3_object_metadata = S3ObjectMetadata( - self.bucket_name, - s3_key, - bucket_object.content_length, - bucket_object.last_modified, - bucket_object.e_tag[1:-1].lower(), - bucket_object.metadata.get(sha512_string), - self.get_s3_object_url(s3_key), - ) - - else: - raise AWSimpleException(f"{self.bucket_name=} {s3_key=} does not exist") + try: + head = self.client.head_object(Bucket=self.bucket_name, Key=s3_key) + except ClientError as e: + if boto_error_to_string(e) in ("404", "NoSuchKey", "NotFound"): + raise AWSimpleException(f"{self.bucket_name=} {s3_key=} does not exist") from e + raise + assert isinstance(self.bucket_name, str) # mainly for mypy + s3_object_metadata = S3ObjectMetadata( + self.bucket_name, + s3_key, + head["ContentLength"], + head["LastModified"], + head["ETag"][1:-1].lower(), + head.get("Metadata", {}).get(sha512_string), + self.get_s3_object_url(s3_key), + ) log.debug(f"{s3_object_metadata=}") return s3_object_metadata @@ -493,8 +510,11 @@ def bucket_exists(self) -> bool: """ # use a "custom" config so that .head_bucket() doesn't take a really long time if the bucket does not exist - config = Config(connect_timeout=5, retries={"max_attempts": 3, "mode": "standard"}) - s3 = boto3.client("s3", config=config) + if self.is_mocked() or is_using_localstack(): + s3 = self.client # the existing client is already pointed at the mock or localstack endpoint + else: + config = Config(connect_timeout=5, retries={"max_attempts": 3, "mode": "standard"}) + s3 = self.session.client("s3", config=config) # use the session so the configured profile/keys/region are honored assert self.bucket_name is not None try: s3.head_bucket(Bucket=self.bucket_name) diff --git a/awsimple/sns.py b/awsimple/sns.py index c5539d5..4aa98df 100644 --- a/awsimple/sns.py +++ b/awsimple/sns.py @@ -3,7 +3,6 @@ """ from typing import Union, Dict, Any -from functools import cache from typeguard import typechecked @@ -22,26 +21,30 @@ def __init__(self, topic_name: str, auto_create: bool = False, **kwargs): super().__init__(resource_name="sns", **kwargs) self.topic_name = topic_name self.auto_create = auto_create + # per-instance cache (do not use functools.cache on methods - it holds a global reference to self, so instances are never garbage collected) + self._topic = None # type: Any + + def _find_topic(self) -> Any: + assert self.resource is not None + for t in self.resource.topics.all(): + if t.arn.split(":")[-1] == self.topic_name: + return t + return None - @cache def get_topic(self) -> Any: """ - gets the associated SNS Topic instance + gets the associated SNS Topic instance (cached per instance) :return: sns.Topic instance """ - topic = None - assert self.resource is not None - for t in self.resource.topics.all(): - if t.arn.split(":")[-1] == self.topic_name: - topic = t - if self.auto_create and topic is None: - self.create_topic() - self.auto_create = False # only do this once - topic = self.get_topic() - return topic - - @cache + if self._topic is None: + topic = self._find_topic() + if self.auto_create and topic is None: + self.create_topic() + topic = self._find_topic() + self._topic = topic + return self._topic + def get_arn(self) -> str: """ get topic ARN from topic name diff --git a/awsimple/sqs.py b/awsimple/sqs.py index 37bf86a..d892ef9 100644 --- a/awsimple/sqs.py +++ b/awsimple/sqs.py @@ -170,12 +170,16 @@ def exists(self) -> bool: """ return self._get_queue() is not None - def calculate_nominal_work_time(self) -> int: + def calculate_nominal_work_time(self) -> float: response_times = [] for begin, end in self.response_history.values(): if end is not None: response_times.append(end - begin) - nominal_work_time = max(statistics.median(response_times), self.minimum_nominal_work_time) # tolerate in case the measured work is very short + if len(response_times) == 0: + # no completed messages in the history yet, so be conservative + nominal_work_time = timedelta(hours=1).total_seconds() + else: + nominal_work_time = max(statistics.median(response_times), self.minimum_nominal_work_time) # tolerate in case the measured work is very short log.debug(f"{nominal_work_time=}") return nominal_work_time @@ -211,7 +215,8 @@ def _receive(self, max_number_of_messages_parameter: Union[int, None] = None) -> log.warning(f'JSONDecodeError : "{self._get_response_history_file_path()}" : {e}') if len(self.response_history) == 0: now = time.time() - self.response_history[None] = (now, now + timedelta(hours=1).total_seconds()) # we have no history, so the initial nominal run time is a long time + # we have no history, so the initial nominal run time is a long time (a string sentinel key is used since this dict is JSON round-tripped and None would become "null") + self.response_history["__initial__"] = (now, now + timedelta(hours=1).total_seconds()) # receive the message(s) messages = [] # type: List[Any] @@ -352,18 +357,25 @@ def add_permission(self, source_arn: str): """ - # a little brute-force, but this is the only way I could assign SQS policy to accept messages from SNS - policy = { - "Version": "2012-10-17", - "Statement": [{"Effect": "Allow", "Principal": "*", "Action": "SQS:SendMessage", "Resource": self.get_arn(), "Condition": {"StringEquals": {"aws:SourceArn": source_arn}}}], - } - - policy_string = json.dumps(policy) - log.info(f"{policy_string=}") if (queue := self._get_queue()) is None: log.warning(f"could not get queue {self.queue_name}") + return + + statement = {"Effect": "Allow", "Principal": "*", "Action": "SQS:SendMessage", "Resource": self.get_arn(), "Condition": {"StringEquals": {"aws:SourceArn": source_arn}}} + + # merge into any existing policy - replacing the whole policy would revoke permissions granted to other sources (e.g. other SNS topics) + existing_policy_string = self.client.get_queue_attributes(QueueUrl=queue.url, AttributeNames=["Policy"]).get("Attributes", {}).get("Policy") + if existing_policy_string is None: + policy = {"Version": "2012-10-17", "Statement": [statement]} else: - self.client.set_queue_attributes(QueueUrl=queue.url, Attributes={"Policy": policy_string}) + policy = json.loads(existing_policy_string) + statements = policy.setdefault("Statement", []) + if statement not in statements: + statements.append(statement) + + policy_string = json.dumps(policy) + log.info(f"{policy_string=}") + self.client.set_queue_attributes(QueueUrl=queue.url, Attributes={"Policy": policy_string}) def purge(self): """ @@ -400,17 +412,18 @@ def __init__(self, queue_name: str, **kwargs): @typechecked() -def get_all_sqs_queues(prefix: str = "") -> List[str]: +def get_all_sqs_queues(prefix: str = "", **kwargs) -> List[str]: """ get all SQS queues :param prefix: prefix to filter queue names (empty string for all queues) + :param kwargs: kwargs for AWSAccess (e.g. profile_name, region_name) so the listing uses the same credentials as the caller :return: list of queue names """ queue_names = [] - sqs = AWSAccess("sqs") + sqs = AWSAccess("sqs", **kwargs) for queue in list(sqs.resource.queues.all()): queue_name = queue.url.split("/")[-1] if queue_name.startswith(prefix): diff --git a/examples/aws_access_test.py b/examples/aws_access_test.py index 29be7b8..58c306a 100644 --- a/examples/aws_access_test.py +++ b/examples/aws_access_test.py @@ -1,4 +1,4 @@ from awsimple import AWSAccess # In this example we're using the default profile -print(AWSAccess().test()) # Should be 'True' +print(AWSAccess().test()) # 'True' if your default-profile AWS credentials are valid, 'False' if not diff --git a/pyproject.toml b/pyproject.toml index 7e93b20..6559af6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,2 +1,7 @@ [tool.black] line-length = 192 + +[tool.pytest.ini_options] +# only collect the test suite - the examples directory contains scripts (e.g. aws_access_test.py) that match +# pytest's *_test.py pattern and hit real AWS at import time +testpaths = ["test_awsimple"] diff --git a/test_awsimple/conftest.py b/test_awsimple/conftest.py index 48fb82d..1740a6d 100644 --- a/test_awsimple/conftest.py +++ b/test_awsimple/conftest.py @@ -5,7 +5,7 @@ from botocore.exceptions import EndpointConnectionError -from awsimple import is_mock, use_moto_mock_env_var, S3Access, is_using_localstack, dynamodb +from awsimple import is_mock, use_moto_mock_env_var, AWSAccess, S3Access, is_using_localstack, dynamodb from test_awsimple import test_awsimple_str, temp_dir, cache_dir @@ -42,6 +42,14 @@ def emit(self, record): assert False +@pytest.fixture(scope="session", autouse=True) +def moto_session(): + # Hold one AWSAccess instance for the whole session so the (reference counted) moto mock and its state persist across all tests. + # Without this, moto state would reset whenever no AWSAccess instances happen to be alive, and tests that build on prior tests' AWS state would fail. + _aws_access = AWSAccess() if is_mock() else None + yield _aws_access + + @pytest.fixture(scope="session", autouse=True) def session_fixture(): temp_dir.mkdir(parents=True, exist_ok=True) diff --git a/test_awsimple/test_dynamodb.py b/test_awsimple/test_dynamodb.py index a1a3cca..345bea8 100644 --- a/test_awsimple/test_dynamodb.py +++ b/test_awsimple/test_dynamodb.py @@ -103,7 +103,7 @@ def test_dynamodb(): assert dynamodb_dict["a_tuple"] == [1, 2, 3] assert dynamodb_dict["42"] == "my_key_is_an_int" # test conversion of an int key to a string assert dynamodb_dict["test_date_time"] == "2019-06-04T20:18:55+00:00" - assert dynamodb_dict["zero_len_string"] is None + assert dynamodb_dict["zero_len_string"] == "" # DynamoDB supports zero length strings for non-key attributes (since May 2020) assert dynamodb_dict["A"] == "i am A" # Enum key (conversion uses the Enum name) assert dynamodb_dict["Y"] == "why" # StrEnum key diff --git a/test_awsimple/test_pubsub/test_pubsub_list_queues.py b/test_awsimple/test_pubsub/test_pubsub_list_queues.py index 7ef152e..aa56847 100644 --- a/test_awsimple/test_pubsub/test_pubsub_list_queues.py +++ b/test_awsimple/test_pubsub/test_pubsub_list_queues.py @@ -1,10 +1,14 @@ -from awsimple import get_all_sqs_queues, is_mock +from awsimple import get_all_sqs_queues, SQSAccess from awsimple.pubsub import AWS_RESOURCE_PREFIX def test_pubsub_list_queues(): + # create a queue with the pubsub prefix so this test doesn't depend on queues left over from other tests + sqs_access = SQSAccess(f"{AWS_RESOURCE_PREFIX}testlistqueues", auto_create=True) + sqs_access.create_queue() + queues_names = get_all_sqs_queues() print(queues_names) diff --git a/test_awsimple/test_pubsub/test_pubsub_make_name_aws_safe.py b/test_awsimple/test_pubsub/test_pubsub_make_name_aws_safe.py index 6104f8f..eff38f8 100644 --- a/test_awsimple/test_pubsub/test_pubsub_make_name_aws_safe.py +++ b/test_awsimple/test_pubsub/test_pubsub_make_name_aws_safe.py @@ -13,4 +13,6 @@ def test_pubsub_make_name_aws_safe(): assert make_name_aws_safe("Invalid#Name$With%Special&Chars*") == "jbudqy4oq2aqhcxgs40b0nmahou3pgq" assert make_name_aws_safe("ab") == "phbce4exwcst2t3d0hqd8k81nc27kd8" - assert make_name_aws_safe("a", "b") == "phbce4exwcst2t3d0hqd8k81nc27kd8" + # multiple args must hash differently than their concatenation, otherwise e.g. channel "ab" collides with channel "a" + node "b" + assert make_name_aws_safe("a", "b") != make_name_aws_safe("ab") + assert 30 <= len(make_name_aws_safe("a", "b")) <= 31