Merge pull request #15 from AllenInstitute/feature/copy-changes-from-ocs

Copy over changes from OCS

Author: rpmcginty

src/aibs_informatics_aws_utils/data_sync/operations.py

@@ -34,6 +34,7 @@
 
 MAX_LOCK_WAIT_TIME_IN_SECS = 60 * 60 * 6  # 6 hours
 
+LOCK_ROOT_ENV_VAR = "DATA_SYNC_LOCK_ROOT"
 
 LocalPath = Union[Path, EFSPath]
 
@@ -100,7 +101,7 @@ def sync_s3_to_local(self, source_path: S3URI, destination_path: LocalPath):
             @retry(CannotAcquirePathLockError, tries=tries, delay=delay, backoff=1)
             @functools.wraps(sync_paths)
             def sync_paths_with_lock(*args, **kwargs):
-                with PathLock(destination_path) as lock:
+                with PathLock(destination_path, lock_root=os.getenv(LOCK_ROOT_ENV_VAR)) as lock:
                     response = sync_paths(*args, **kwargs)
                 return response
 

src/aibs_informatics_aws_utils/dynamodb/conditions.py

@@ -1,7 +1,7 @@
 import re
 from collections import defaultdict
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List, Mapping, Match, Union, cast
+from typing import Any, Dict, Iterable, List, Mapping, Match, Optional, Union, cast
 
 from aibs_informatics_core.collections import ValidatedStr
 from aibs_informatics_core.models.aws.dynamodb import (
@@ -24,6 +24,39 @@
 logger = get_logger(__name__)
 
 
+def condition_to_str(condition: Optional[ConditionBase]) -> Optional[str]:
+    """Converts a ConditionBase Boto3 object to its str representation
+
+    NOTE: Function should be removed if this PR is merged: https://github.com/boto/boto3/pull/3254
+
+    Examples:
+
+    >>> condition_to_str(Key("name").eq("new_name") & Attr("description").begins_with("new"))
+    (name = new_name AND begins_with(description, new))
+
+    >>> condition_to_str(Attr("description").contains("cool"))
+    contains(description, cool)
+
+    >>> condition_to_str(None)
+    None
+    """
+    if condition is None:
+        return None
+
+    builder = ConditionExpressionBuilder()
+    expression = builder.build_expression(condition)
+
+    condition_expression = expression.condition_expression
+
+    for name_placeholder, actual_name in expression.attribute_name_placeholders.items():
+        condition_expression = condition_expression.replace(name_placeholder, str(actual_name))
+
+    for value_placeholder, actual_value in expression.attribute_value_placeholders.items():
+        condition_expression = condition_expression.replace(value_placeholder, str(actual_value))
+
+    return condition_expression
+
+
 @dataclass
 class ExpressionComponentsBase:
     expression: str

src/aibs_informatics_aws_utils/dynamodb/table.py

@@ -5,6 +5,7 @@
     Dict,
     Generic,
     List,
+    Literal,
     Mapping,
     MutableMapping,
     Optional,
@@ -14,6 +15,7 @@
     TypeVar,
     Union,
     cast,
+    overload,
 )
 
 from aibs_informatics_core.env import EnvBase
@@ -42,6 +44,7 @@
 from botocore.exceptions import ClientError
 
 from aibs_informatics_aws_utils.core import get_client_error_code
+from aibs_informatics_aws_utils.dynamodb.conditions import condition_to_str
 from aibs_informatics_aws_utils.dynamodb.functions import (
     convert_floats_to_decimals,
     execute_partiql_statement,
@@ -72,8 +75,7 @@ def check_db_query_unique(
     key_condition_expression: Optional[ConditionBase] = None,
     filter_expression: Optional[ConditionBase] = None,
 ):
-    # TODO: this should be len(query_result) > 1
-    if len(query_result) != 1:
+    if len(query_result) > 1:
         readable_key_expression: Optional[BuiltConditionExpression] = None
         if key_condition_expression:
             expression_builder = ConditionExpressionBuilder()
@@ -245,8 +247,6 @@ def build_optimized_condition_expression_set(
 
 @dataclass
 class DynamoDBTable(LoggingMixin, Generic[DB_MODEL, DB_INDEX]):
-    # env_base: EnvBase = field(default_factory=EnvBase.from_env)
-
     def __post_init__(self):
         check_table_name_and_index_match(self.table_name, self.get_db_index_cls())
 
@@ -432,10 +432,9 @@ def query(
             expect_non_empty (bool, optional): Whether the resulting query should return at least
                 one result. An error will be raised if expect_non_empty=True and 0 results were
                 returned by the query.
-            expect_unique (bool, option): Whether the result of the query is expected to be
-                unique (i.e. returns 1 and ONLY 1 result). An error will be raised if 0 or
-                more than 1 results were returned for the query.
-
+            expect_unique (bool, option): Whether the result of the query is expected to
+                return AT MOST one result. An error will be raised if expect_unique=True and MORE
+                than 1 result was returned for the query.
         Returns:
             Sequence[Dict[str, Any]]: A sequence of dictionaries representing database rows
                 where partition_key/sort_key and filter conditions are satisfied.
@@ -455,7 +454,8 @@ def query(
 
         self.log.info(
             f"Calling query on {self.table_name} table (index: {index_name}, "
-            f"key condition: {key_condition_expression}, filters: {filter_expression})"
+            f"key condition: {condition_to_str(key_condition_expression)}, "
+            f"filters: {condition_to_str(filter_expression)})"
         )
 
         items = table_query(
@@ -507,9 +507,9 @@ def scan(
             expect_non_empty (bool, optional): Whether the resulting query should return at least
                 one result. An error will be raised if expect_non_empty=True and 0 results were
                 returned by the query.
-            expect_unique (bool, option): Whether the result of the query is expected to be
-                unique (i.e. returns 1 and ONLY 1 result). An error will be raised if 0 or
-                more than 1 results were returned for the query.
+            expect_unique (bool, option): Whether the result of the query is expected to
+                return AT MOST one result. An error will be raised if expect_unique=True and MORE
+                than 1 result was returned for the query.
 
         Returns:
             Sequence[Dict[str, Any]]: A sequence of dictionaries representing database rows
@@ -524,7 +524,7 @@ def scan(
 
         self.log.info(
             f"Calling scan on {self.table_name} table (index: {index_name},"
-            f" filters: {filter_expression})"
+            f" filters: {condition_to_str(filter_expression)})"
         )
 
         items = table_scan(
@@ -611,17 +611,9 @@ def put(
         condition_expression: Optional[ConditionBase] = None,
         **table_put_item_kwargs,
     ) -> DB_MODEL:
-        # Convert our ConditionBase class to a readable string (for logging/debugging purposes)
-        if condition_expression:
-            expression = ConditionExpressionBuilder().build_expression(condition_expression)
-            expression_str = (
-                f"({expression.condition_expression}, "
-                f"{expression.attribute_name_placeholders}, "
-                f"{expression.attribute_value_placeholders})"
-            )
-        else:
-            expression_str = "None"
-        put_summary = f"(entry: {entry}, condition_expression: {expression_str})"
+        put_summary = (
+            f"(entry: {entry}, condition_expression: {condition_to_str(condition_expression)})"
+        )
         self.log.debug(f"{self.table_name} - Putting new entry: {put_summary}")
 
         e_msg_intro = f"{self.table_name} - Error putting entry: {put_summary}."
@@ -659,22 +651,30 @@ def update(
 
         for k in key:
             new_attributes.pop(k, None)
+        # Add k:v pair from new_attributes if new != old value for a given key
+        new_clean_attrs: Dict[str, Any] = {}
+        if old_entry:
+            for k, new_v in new_attributes.items():
+                if getattr(old_entry, k) != new_v:
+                    new_clean_attrs[k] = new_v
+        else:
+            new_clean_attrs = new_attributes
 
-        if not new_attributes:
+        if not new_clean_attrs:
             self.log.debug(
                 f"{self.table_name} - No attr_updates to do! Skipping _update_entry call."
             )
             if not old_entry:
                 old_entry = self.get(key)
             return old_entry
 
-        update_summary = f"(old_entry: {old_entry}, new_attributes: {new_attributes})"
+        update_summary = f"(old_entry: {old_entry}, new_attributes: {new_clean_attrs})"
         self.log.debug(f"{self.table_name} - Updating entry: {update_summary}")
         try:
             updated_item = table_update_item(
                 table_name=self.table_name,
                 key=key,
-                attributes=new_attributes,
+                attributes=new_clean_attrs,
                 return_values="ALL_NEW",
                 **table_update_item_kwargs,
             )
@@ -694,6 +694,29 @@ def update(
         self.log.debug(f"{self.table_name} - Successfully updated entry: {updated_entry}")
         return updated_entry
 
+    @overload
+    def delete(
+        self,
+        key: Union[DynamoDBKey, DB_MODEL],
+        error_on_nonexistent: Literal[True],
+    ) -> DB_MODEL:
+        ...
+
+    @overload
+    def delete(
+        self,
+        key: Union[DynamoDBKey, DB_MODEL],
+        error_on_nonexistent: Literal[False],
+    ) -> Optional[DB_MODEL]:
+        ...
+
+    @overload
+    def delete(
+        self,
+        key: Union[DynamoDBKey, DB_MODEL],
+    ) -> Optional[DB_MODEL]:
+        ...
+
     def delete(
         self,
         key: Union[DynamoDBKey, DB_MODEL],

src/aibs_informatics_aws_utils/sqs.py

@@ -1,6 +1,8 @@
 import json
 import logging
+from typing import Optional, Type
 
+from aibs_informatics_core.utils.json import DecimalEncoder
 from botocore.exceptions import ClientError
 
 from aibs_informatics_aws_utils.core import AWSService, get_region
@@ -13,7 +15,7 @@
 get_sqs_resource = AWSService.SQS.get_resource
 
 
-def delete_from_queue(queue_name: str, receipt_handle: str, region: str = None):
+def delete_from_queue(queue_name: str, receipt_handle: str, region: Optional[str] = None):
     sqs = get_sqs_client()
     queue_url_response = sqs.get_queue_url(QueueName=queue_name)
     queue_url = queue_url_response["QueueUrl"]
@@ -40,3 +42,63 @@ def send_to_dispatch_queue(payload: dict, env_base: str):
     response = sqs.send_message(QueueUrl=queue_url, MessageBody=json.dumps(payload))
 
     return response["MD5OfMessageBody"]
+
+
+def send_sqs_message(
+    queue_name: str,
+    payload: dict,
+    message_deduplication_id: Optional[str] = None,
+    message_group_id: Optional[str] = None,
+    payload_json_encoder: Type[json.JSONEncoder] = DecimalEncoder,
+) -> str:
+    """Send a message to an SQS queue by providing a queue name
+
+    Args:
+        queue_name (str): The name of the queue that you want to send a message to.
+            (e.g. 'aps-sync-request-queue.fifo')
+        payload (dict): A dictionary representing the message payload you would like to send.
+        message_deduplication_id (Optional[str], optional): An ID that can be used by SQS
+            to remove messages that have the same deduplication_id. Do not set if your
+            SQS queue already uses content based deduplication. Defaults to None.
+        message_group_id (Optional[str], optional): Required for FIFO queues.
+            Messages sent with the same message_group_id will obey FIFO rules. Messages with
+            different message_group_ids may be interleaved. Defaults to None.
+        payload_json_encoder (Type[json.JSONEncoder], optional): The JSONEncoder
+            class that should be used to covert the input `payload` dictionary into
+            a json string. By default uses a DecimalEncoder which can handle decimal.Decimal types.
+
+    Raises:
+        AWSError: If the provided queue_name cannot be resolved to an SQS url.
+            HINT: Does the code calling this function have the correct SQS permissions?
+        RuntimeError: If the destination queue is a FIFO queue, then `message_group_id` MUST
+            be provided.
+
+    Returns:
+        str: Returns an MD5 digest of the send message body.
+    """
+    sqs = get_sqs_client(region=get_region())
+    try:
+        queue_url_response = sqs.get_queue_url(QueueName=queue_name)
+    except ClientError as e:
+        raise AWSError(
+            f"Could not find SQS queue with name: {queue_name}. "
+            "Does the code calling send_sqs_message() have sqs:GetQueueUrl permissions?"
+        )
+
+    send_sqs_message_args = {
+        "QueueUrl": queue_url_response["QueueUrl"],
+        "MessageBody": json.dumps(payload, cls=payload_json_encoder),
+    }
+
+    if message_group_id is not None:
+        send_sqs_message_args["MessageGroupId"] = message_group_id
+    else:
+        if queue_name.endswith(".fifo"):
+            raise RuntimeError("SQS messages for a FIFO queue *must* include a message_group_id!")
+
+    if message_deduplication_id is not None:
+        send_sqs_message_args["MessageDeduplicationId"] = message_deduplication_id
+
+    response = sqs.send_message(**send_sqs_message_args)  # type: ignore  # complains about valid kwargs
+
+    return response["MD5OfMessageBody"]