Table

DynamoDB table wrapper and operations.

---

DynamoDBTable dataclass

DynamoDBTable()

Bases: LoggingMixin, Generic[DB_MODEL, DB_INDEX]

batch_get

batch_get(keys, partial=False, ignore_missing=False)

Batch get items from a DynamoDB table by providing a list of keys.

Parameters:

Name	Type	Description	Default
keys	Union[List[DynamoDBKey], List[DynamoDBItemValue], List[Tuple[DynamoDBItemValue, DynamoDBItemValue]]]	The partition key values that should be used to search the table. Each key can be one of: single partition key value a tuple of partition and sort key value a dictionary of attribute:value	required
partial	bool	Whether partial values are allowed. Defaults to False.	False
ignore_missing	bool	If true, suppress errors for keys that are not found. Defaults to False.	False

Raises:

Type	Description
DBReadException	If any of the items could not be found.

Returns:

Type	Description
List[DB_MODEL]	List of database entries that were found.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def batch_get(
    self,
    keys: Union[
        List[DynamoDBKey],
        List[DynamoDBItemValue],
        List[Tuple[DynamoDBItemValue, DynamoDBItemValue]],
    ],
    partial: bool = False,
    ignore_missing: bool = False,
) -> List[DB_MODEL]:
    """Batch get items from a DynamoDB table by providing a list of keys.

    Args:
        keys: The partition key values that should be used to search the table.
            Each key can be one of:

            - **single partition key value**
            - **a tuple of partition and sort key value**
            - **a dictionary of attribute:value**

        partial: Whether partial values are allowed. Defaults to False.
        ignore_missing: If true, suppress errors for keys that are not found.
            Defaults to False.

    Raises:
        DBReadException: If any of the items could not be found.

    Returns:
        List of database entries that were found.
    """
    if not keys:
        return []

    index = self.index_or_default()
    item_keys = [self.build_key(key, index=index) for key in keys]

    items = table_get_items(table_name=self.table_name, keys=item_keys)
    if len(items) != len(item_keys) and not ignore_missing:
        missing_keys = set(
            [(_[index.key_name], _.get(index.sort_key_name or "")) for _ in item_keys]
        ).difference((_[index.key_name], _.get(index.sort_key_name or "")) for _ in items)

        raise DBReadException(f"Could not find items for {missing_keys}")
    entries = [self.build_entry(_, partial=partial) for _ in items]
    return entries

delete

delete(
    key: Union[DynamoDBKey, DB_MODEL],
    error_on_nonexistent: Literal[True],
) -> DB_MODEL

delete(
    key: Union[DynamoDBKey, DB_MODEL],
    error_on_nonexistent: Literal[False],
) -> Optional[DB_MODEL]

delete(
    key: Union[DynamoDBKey, DB_MODEL]
) -> Optional[DB_MODEL]

delete(key, error_on_nonexistent=False)

Delete an item from the DynamoDB table.

Parameters:

Name	Type	Description	Default
key	Union[DynamoDBKey, DB_MODEL]	The primary key of the item to be deleted, or a DB_MODEL instance representing the item to be deleted.	required
error_on_nonexistent	bool	Whether to raise an error if the item does not exist. Defaults to False.	False

Raises:

Type	Description
DBWriteException	If there was an error deleting the entry.
DBWriteException	If error_on_nonexistent is True and the item does not exist.

Returns:

Type	Description
Optional[DB_MODEL]	The deleted database model entry, or None if the item did not exist and error_on_nonexistent is False.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def delete(
    self,
    key: Union[DynamoDBKey, DB_MODEL],
    error_on_nonexistent: bool = False,
) -> Optional[DB_MODEL]:
    """Delete an item from the DynamoDB table.

    Args:
        key (Union[DynamoDBKey, DB_MODEL]): The primary key of the item to be deleted,
            or a DB_MODEL instance representing the item to be deleted.
        error_on_nonexistent (bool, optional): Whether to raise an error if the item
            does not exist. Defaults to False.

    Raises:
        DBWriteException: If there was an error deleting the entry.
        DBWriteException: If error_on_nonexistent is True and the item does not exist.

    Returns:
        The deleted database model entry, or None if the item did not exist and
            error_on_nonexistent is False.
    """
    if isinstance(key, self.get_db_model_cls()):
        key = self.build_key_from_entry(key)
    delete_summary = f"(db_primary_key: {key})"
    self.log.debug(f"{self.table_name} - Deleting entry with: {delete_summary}")
    e_msg = f"{self.table_name} - Delete failed for the following primary key: {key}"
    try:
        deleted_attributes = table_delete_item(
            table_name=self.table_name,
            key=cast(DynamoDBKey, key),
            return_values="ALL_OLD",  # type: ignore[arg-type] # expected type more general than specified here
        )

        if not deleted_attributes:
            self.log.info(f"{self.table_name} - Nothing deleted for primary key: {key}")
            if error_on_nonexistent:
                raise DBWriteException(e_msg)
            return None
        else:
            deleted_entry = self.build_entry(deleted_attributes)
            self.log.info(f"{self.table_name} - Deleted entry: {deleted_attributes}")
            return deleted_entry
    except ClientError as e:
        detailed_e_msg = f"{e_msg}. Details: {get_client_error_code(e)}"
        raise DBWriteException(detailed_e_msg)

get

get(key, partial=False)

Get a single item from a DynamoDB table by providing a key.

Parameters:

Name	Type	Description	Default
key	Union[DynamoDBKey, DynamoDBItemValue, Tuple[DynamoDBItemValue, DynamoDBItemValue]]	The partition key value that should be used to search the table. Can be a single partition key value, a tuple of partition and sort key value, or a dictionary of attribute:value.	required
partial	bool	Whether partial values are allowed. Defaults to False.	False

Raises:

Type	Description
DBReadException	If the item could not be found.

Returns:

Type	Description
DB_MODEL	The database entry that was found.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def get(
    self,
    key: Union[DynamoDBKey, DynamoDBItemValue, Tuple[DynamoDBItemValue, DynamoDBItemValue]],
    partial: bool = False,
) -> DB_MODEL:
    """Get a single item from a DynamoDB table by providing a key.

    Args:
        key: The partition key value that should be used to search the table.
            Can be a single partition key value, a tuple of partition and sort key value,
            or a dictionary of attribute:value.
        partial: Whether partial values are allowed. Defaults to False.

    Raises:
        DBReadException: If the item could not be found.

    Returns:
        The database entry that was found.
    """
    item_key = self.build_key(key)

    item = table_get_item(table_name=self.table_name, key=item_key)
    if item is None:
        raise DBReadException(f"Could not resolve item from {item_key}")
    return self.build_entry(item, partial=partial)

put

put(
    entry,
    condition_expression=None,
    **table_put_item_kwargs
)

Put a new item into the DynamoDB table.

Parameters:

Name	Type	Description	Default
entry	DB_MODEL	The database model entry to be put into the table.	required
condition_expression	Optional[ConditionBase]	An optional condition expression that must be satisfied for the put operation to succeed. Defaults to None.	None
**table_put_item_kwargs		Additional keyword arguments to be passed to the table_put_item function.	{}

Raises:

Type	Description
DBWriteException	If there was an error putting the entry.
DBWriteException	If the HTTP response code indicates a failure.

Returns:

Type	Description
DB_MODEL	The database model entry that was put into the table.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def put(
    self,
    entry: DB_MODEL,
    condition_expression: Optional[ConditionBase] = None,
    **table_put_item_kwargs,
) -> DB_MODEL:
    """Put a new item into the DynamoDB table.

    Args:
        entry (DB_MODEL): The database model entry to be put into the table.
        condition_expression (Optional[ConditionBase], optional): An optional condition
            expression that must be satisfied for the put operation to succeed.
            Defaults to None.
        **table_put_item_kwargs: Additional keyword arguments to be passed to the
            table_put_item function.

    Raises:
        DBWriteException: If there was an error putting the entry.
        DBWriteException: If the HTTP response code indicates a failure.

    Returns:
        The database model entry that was put into the table.
    """
    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}."
    try:
        put_response = table_put_item(
            table_name=self.table_name,
            item=self.build_item(entry),
            condition_expression=condition_expression,
            **table_put_item_kwargs,
        )
    except ClientError as e:
        e_msg_client = f"{e_msg_intro} Details: {get_client_error_code(e)}"
        self.log.error(e_msg_client)
        raise DBWriteException(e_msg_client)
    else:
        if not (200 <= put_response["ResponseMetadata"]["HTTPStatusCode"] < 300):
            e_msg_http = f"{e_msg_intro} Table put_item response: {put_response}"
            raise DBWriteException(e_msg_http)
    self.log.debug(f"{self.table_name} - Put successful: {put_response}")
    return entry

query

query(
    index,
    partition_key,
    sort_key_condition_expression=None,
    filters=None,
    consistent_read=False,
    expect_non_empty=False,
    expect_unique=False,
    allow_partial=False,
)

Query a DynamoDB table by providing a DBIndex, partition_key, optional sort_key, and optional filter conditions.

Parameters:

Name	Type	Description	Default
index	DB_INDEX	Specifies the specific table index (e.g. main table, global secondary index, or local secondary index) that should be queried.	required
partition_key	Union[DynamoDBPrimaryKeyItemValue, ConditionBase]	The partition key value that should be used to search the table.	required
sort_key_condition_expression	Optional[ConditionBase]	The sort key condition expression to be used to query the table. Example: from boto3.dynamodb.conditions import Key Key('my_sort_key_name').begins_with('prefix_')	None
filters	Optional[List[ConditionBase]]	A list of ConditionBase expressions where query results must satisfy.	None
consistent_read	bool	Whether a strongly consistent read should be used for the query. By default False which returns eventually consistent reads.	False
expect_non_empty	bool	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.	False
expect_unique	bool	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.	False
allow_partial	bool	Whether to allow partial entries. Defaults to False.	False

Returns:

Type	Description
List[DB_MODEL]	A list of database model entries where partition_key/sort_key and filter conditions are satisfied.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def query(
    self,
    index: DB_INDEX,
    partition_key: Union[DynamoDBPrimaryKeyItemValue, ConditionBase],
    sort_key_condition_expression: Optional[ConditionBase] = None,
    filters: Optional[List[ConditionBase]] = None,
    consistent_read: bool = False,
    expect_non_empty: bool = False,
    expect_unique: bool = False,
    allow_partial: bool = False,
) -> List[DB_MODEL]:
    """Query a DynamoDB table by providing a DBIndex, partition_key, optional sort_key, and optional filter conditions.

    Args:
        index: Specifies the specific table index (e.g. main table, global secondary
            index, or local secondary index) that should be queried.
        partition_key: The partition key value that should be used to search the table.
        sort_key_condition_expression: The sort key condition expression to be used
            to query the table. Example:

            ```python
            from boto3.dynamodb.conditions import Key
            Key('my_sort_key_name').begins_with('prefix_')
            ```

        filters: A list of ConditionBase expressions where query results must satisfy.
        consistent_read: Whether a strongly consistent read should be used
            for the query. By default False which returns eventually consistent reads.
        expect_non_empty: 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: 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.
        allow_partial: Whether to allow partial entries. Defaults to False.

    Returns:
        A list of database model entries where partition_key/sort_key and
            filter conditions are satisfied.
    """  # noqa: E501

    if consistent_read:
        check_index_supports_strongly_consistent_read(index=index)

    index_name = self.get_index_name(index)

    key_condition_expression = self._build_key_condition_expression(
        index=index,
        partition_key=partition_key,
        sort_key_condition_expression=sort_key_condition_expression,
    )
    filter_expression = self._build_filter_condition_expression(filters=filters)

    self.log.info(
        f"Calling query on {self.table_name} table (index: {index_name}, "
        f"key condition: {condition_to_str(key_condition_expression)}, "
        f"filters: {condition_to_str(filter_expression)})"
    )

    items = table_query(
        table_name=self.table_name,
        index_name=index_name,
        key_condition_expression=key_condition_expression,
        filter_expression=filter_expression,
        consistent_read=consistent_read,
    )
    if expect_unique:
        check_db_query_unique(
            index=index,
            query_result=items,
            key_condition_expression=key_condition_expression,
            filter_expression=filter_expression,
        )
    if expect_non_empty:
        check_db_query_non_empty(
            index=index,
            query_result=items,
            key_condition_expression=key_condition_expression,
            filter_expression=filter_expression,
        )
    entries = [self.build_entry(_, partial=True) for _ in items]
    if not allow_partial:
        entries = self._fill_values(entries)
    return entries

scan

scan(
    index=None,
    filters=None,
    consistent_read=False,
    expect_non_empty=False,
    expect_unique=False,
    allow_partial=False,
)

Scan a DynamoDB table by providing a DBIndex and optional filter conditions.

Parameters:

Name	Type	Description	Default
index	Optional[DB_INDEX]	Specifies the specific table index (e.g. main table, global secondary index, or local secondary index) that should be scanned.	None
filters	Optional[List[ConditionBase]]	A list of ConditionBase expressions where scan results must satisfy.	None
consistent_read	bool	Whether a strongly consistent read should be used for the scan. By default False which returns eventually consistent reads.	False
expect_non_empty	bool	Whether the resulting scan should return at least one result. An error will be raised if expect_non_empty=True and 0 results were returned by the scan.	False
expect_unique	bool	Whether the result of the scan 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 scan.	False
allow_partial	bool	Whether to allow partial entries. Defaults to False.	False

Returns:

Type	Description
List[DB_MODEL]	A list of database model entries where filter conditions are satisfied.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def scan(
    self,
    index: Optional[DB_INDEX] = None,
    filters: Optional[List[ConditionBase]] = None,
    consistent_read: bool = False,
    expect_non_empty: bool = False,
    expect_unique: bool = False,
    allow_partial: bool = False,
) -> List[DB_MODEL]:
    """Scan a DynamoDB table by providing a DBIndex and optional filter conditions.

    Args:
        index: Specifies the specific table index (e.g. main table, global secondary
            index, or local secondary index) that should be scanned.
        filters: A list of ConditionBase expressions where scan results must satisfy.
        consistent_read: Whether a strongly consistent read should be used
            for the scan. By default False which returns eventually consistent reads.
        expect_non_empty: Whether the resulting scan should return at least
            one result. An error will be raised if expect_non_empty=True and 0 results were
            returned by the scan.
        expect_unique: Whether the result of the scan 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 scan.
        allow_partial: Whether to allow partial entries. Defaults to False.

    Returns:
        A list of database model entries where filter conditions are satisfied.
    """
    index = self.index_or_default(index)
    if consistent_read:
        check_index_supports_strongly_consistent_read(index=index)

    index_name = self.get_index_name(index)
    filter_expression = self._build_filter_condition_expression(filters=filters)

    self.log.info(
        f"Calling scan on {self.table_name} table (index: {index_name},"
        f" filters: {condition_to_str(filter_expression)})"
    )

    items = table_scan(
        table_name=self.table_name,
        index_name=index_name,
        filter_expression=filter_expression,
        consistent_read=consistent_read,
    )
    if expect_unique:
        check_db_query_unique(
            index=index,
            query_result=items,
            filter_expression=filter_expression,
        )
    if expect_non_empty:
        check_db_query_non_empty(
            index=index,
            query_result=items,
            filter_expression=filter_expression,
        )
    entries = [self.build_entry(_, partial=True) for _ in items]
    if not allow_partial:
        entries = self._fill_values(entries)
    return entries

smart_query

smart_query(
    *filters,
    consistent_read=False,
    expect_non_empty=False,
    expect_unique=False,
    allow_partial=False,
    allow_scan=True,
    **kw_filters
)

Perform a smart query on the DynamoDB table by automatically determining whether to use a query or scan operation based on the provided filters.

Parameters:

Name	Type	Description	Default
*filters	Union[DynamoDBKey, ConditionBase]	Varargs of DynamoDBKey or ConditionBase to filter the query/scan	()
consistent_read	bool	Whether a strongly consistent read should be used for the query/scan. Defaults to False which returns eventually consistent reads.	False
expect_non_empty	bool	Whether the resulting query/scan should return at least one result. An error will be raised if expect_non_empty=True and 0 results were returned.	False
expect_unique	bool	Whether the result of the query/scan is expected to return AT MOST one result. An error will be raised if expect_unique=True and MORE than 1 result was returned.	False
allow_partial	bool	Whether to allow partial entries. Defaults to False.	False
allow_scan	bool	Whether to allow a scan operation if no partition key is found in the filters. Defaults to True.	True

Raises:

Type	Description
DBQueryException	If no partition key is found and allow_scan is False.

Returns:

Type	Description
List[DB_MODEL]	A list of database model entries where partition_key/sort_key and filter conditions are satisfied.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def smart_query(
    self,
    *filters: Union[DynamoDBKey, ConditionBase],
    consistent_read: bool = False,
    expect_non_empty: bool = False,
    expect_unique: bool = False,
    allow_partial: bool = False,
    allow_scan: bool = True,
    **kw_filters: Any,
) -> List[DB_MODEL]:
    """
    Perform a smart query on the DynamoDB table by automatically determining whether to
    use a query or scan operation based on the provided filters.

    Args:
        *filters (Union[DynamoDBKey, ConditionBase]):
            Varargs of DynamoDBKey or ConditionBase to filter the query/scan
        consistent_read (bool):
            Whether a strongly consistent read should be used for the query/scan.
            Defaults to False which returns eventually consistent reads.
        expect_non_empty (bool):
            Whether the resulting query/scan should return at least one result.
            An error will be raised if expect_non_empty=True and 0 results were returned.
        expect_unique (bool):
            Whether the result of the query/scan is expected to return AT MOST one result.
            An error will be raised if expect_unique=True and MORE than 1 result was returned.
        allow_partial (bool):
            Whether to allow partial entries. Defaults to False.
        allow_scan (bool):
            Whether to allow a scan operation if no partition key is found in the filters.
            Defaults to True.

    Raises:
        DBQueryException:
            If no partition key is found and allow_scan is False.

    Returns:
        A list of database model entries where partition_key/sort_key and filter
            conditions are satisfied.
    """  # noqa: E501
    (
        index,
        partition_key,
        sort_key_condition_expression,
        filter_expressions,
    ) = build_optimized_condition_expression_set(
        self.get_db_index_cls(),
        *filters,
        **kw_filters,
    )
    if partition_key:
        return self.query(
            index=self.index_or_default(index),
            partition_key=partition_key,
            sort_key_condition_expression=sort_key_condition_expression,
            filters=filter_expressions,
            consistent_read=consistent_read,
            expect_non_empty=expect_non_empty,
            expect_unique=expect_unique,
            allow_partial=allow_partial,
        )
    else:
        if not allow_scan:
            raise DBQueryException(
                f"Could not find a partition key for {self.table_name} table! "
                "Please provide a partition key or allow_scan=True."
            )
        if sort_key_condition_expression:
            filter_expressions.append(sort_key_condition_expression)
        return self.scan(
            index=index,
            filters=filter_expressions,
            consistent_read=consistent_read,
            expect_non_empty=expect_non_empty,
            expect_unique=expect_unique,
            allow_partial=allow_partial,
        )

update

update(
    key,
    new_entry,
    old_entry=None,
    **table_update_item_kwargs
)

Update an existing item in the DynamoDB table.

Parameters:

Name	Type	Description	Default
key	DynamoDBKey	The primary key of the item to be updated.	required
new_entry	Union[Mapping[str, Any], DB_MODEL]	The new values for the item. Can be a dictionary of attribute:value pairs or a DB_MODEL instance.	required
old_entry	Optional[DB_MODEL]	The existing entry before the update. If provided, only attributes that differ from the old_entry will be updated. Defaults to None.	None
**table_update_item_kwargs		Additional keyword arguments to be passed to the table_update_item function.	{}

Raises:

Type	Description
DBWriteException	If there was an error updating the entry.
DBWriteException	If no attributes need to be updated.

Returns:

Type	Description
DB_MODEL	The updated database model entry.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def update(
    self,
    key: DynamoDBKey,
    new_entry: Union[Mapping[str, Any], DB_MODEL],
    old_entry: Optional[DB_MODEL] = None,
    **table_update_item_kwargs,
) -> DB_MODEL:
    """Update an existing item in the DynamoDB table.

    Args:
        key (DynamoDBKey): The primary key of the item to be updated.
        new_entry (Union[Mapping[str, Any], DB_MODEL]): The new values for the item.
            Can be a dictionary of attribute:value pairs or a DB_MODEL instance.
        old_entry (Optional[DB_MODEL], optional): The existing entry before the update.
            If provided, only attributes that differ from the old_entry will be updated.
            Defaults to None.
        **table_update_item_kwargs: Additional keyword arguments to be passed to the
            table_update_item function.

    Raises:
        DBWriteException: If there was an error updating the entry.
        DBWriteException: If no attributes need to be updated.

    Returns:
        The updated database model entry.
    """
    new_attributes: Dict[str, Any] = {}
    if isinstance(new_entry, self.get_db_model_cls()):
        new_attributes = self.build_item(new_entry, partial=True)
    elif new_entry:
        # we still need to convert floats to decimals if new_entry is a dict
        new_attributes = convert_floats_to_decimals(new_entry)  # type: ignore[arg-type]

    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_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_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_clean_attrs,
            return_values="ALL_NEW",
            **table_update_item_kwargs,
        )
        # table_update_item will always return a dict if ReturnValues != "NONE"
        if updated_item is None:
            raise DBWriteException(
                f"{self.table_name} - Error updating entry: {update_summary}"
            )
        updated_entry = self.build_entry(updated_item)
    except ClientError as e:
        e_msg = (
            f"{self.table_name} - Error updating entry: {update_summary}. "
            f"Details: {get_client_error_code(e)}"
        )
        self.log.error(e_msg)
        raise DBWriteException(e_msg)
    self.log.debug(f"{self.table_name} - Successfully updated entry: {updated_entry}")
    return updated_entry

build_optimized_condition_expression_set

build_optimized_condition_expression_set(
    candidate_indexes, *args, **kwargs
)

Builds an optimized set of conditions for a query or scan

Parameters:

Name	Type	Description	Default
candidate_indexes	Type[DB_INDEX] | Sequence[DB_INDEX]	index class or subset of indexes the order of the indexes matters! The first index that matches the provided conditions will be used.	required
*args	Union[DynamoDBKey, ConditionBase]	varargs of DynamoDBKey or ConditionBase	()
**kwargs	Any	kwargs of DynamoDBKey or ConditionBase	{}

Returns:

Type	Description
Tuple[Optional[DB_INDEX], Optional[ConditionBase], Optional[ConditionBase], List[ConditionBase]]	A tuple containing: the target index, partition key condition, sort key condition, and filter expressions.

Source code in src/aibs_informatics_aws_utils/dynamodb/table.py

def build_optimized_condition_expression_set(
    candidate_indexes: Union[Type[DB_INDEX], Sequence[DB_INDEX]],
    *args: Union[DynamoDBKey, ConditionBase],
    **kwargs: Any,
) -> Tuple[
    Optional[DB_INDEX], Optional[ConditionBase], Optional[ConditionBase], List[ConditionBase]
]:
    """Builds an optimized set of conditions for a query or scan


    Args:
        candidate_indexes (Type[DB_INDEX]|Sequence[DB_INDEX]): index class or subset of indexes
            the order of the indexes matters! The first index that matches the provided
            conditions will be used.
        *args (Union[DynamoDBKey, ConditionBase]): varargs of DynamoDBKey or ConditionBase
        **kwargs (Any): kwargs of DynamoDBKey or ConditionBase

    Returns:
        A tuple containing:
            the target index,
            partition key condition,
            sort key condition,
            and filter expressions.

    """
    target_index: Optional[DB_INDEX] = None
    partition_key: Optional[ConditionBase] = None
    sort_key_condition_expression: Optional[ConditionBase] = None
    filter_expressions: List[ConditionBase] = []

    if not args and not kwargs:
        return target_index, partition_key, sort_key_condition_expression, filter_expressions

    if not isinstance(candidate_indexes, Sequence):
        candidate_indexes = [ci for ci in candidate_indexes]
    index_all_key_names = set(
        {
            *{_.key_name for _ in candidate_indexes},
            *{_.sort_key_name for _ in candidate_indexes if _.sort_key_name},
        }
    )

    SupportedKeyComparisonTypes = (
        Equals,
        GreaterThan,
        GreaterThanEquals,
        LessThan,
        LessThanEquals,
        BeginsWith,
        Between,
    )

    candidate_conditions: Dict[
        str,
        Union[
            Equals, GreaterThan, GreaterThanEquals, LessThan, LessThanEquals, BeginsWith, Between
        ],
    ] = {}
    non_candidate_conditions: List[ConditionBase] = []

    for _ in (kwargs,) + args:
        if not isinstance(_, ConditionBase):
            for k, v in _.items():
                if k not in index_all_key_names:
                    non_candidate_conditions.append(Attr(k).eq(v))
                    continue
                new_condition = Key(k).eq(v)
                if (
                    k in candidate_conditions
                    and candidate_conditions[k]._values[1:] != new_condition._values[1:]  # type: ignore[attr-defined,union-attr]
                ):
                    raise DBQueryException(f"Multiple values provided for attribute {k}!")
                candidate_conditions[k] = Key(k).eq(v)
        elif len(_._values) and isinstance(_._values[0], (Key, Attr)):  # type: ignore[attr-defined,union-attr]
            attr_name = cast(str, _._values[0].name)  # type: ignore[attr-defined,union-attr]
            if attr_name not in index_all_key_names or not isinstance(
                _, SupportedKeyComparisonTypes
            ):
                non_candidate_conditions.append(_)
                continue
            if (
                attr_name in candidate_conditions
                and candidate_conditions[attr_name]._values[1:] != _._values[1:]  # type: ignore[union-attr]
            ):
                raise DBQueryException(f"Multiple values provided for attribute {attr_name}!")
            candidate_conditions[attr_name] = _
        else:
            non_candidate_conditions.append(_)

    for index in candidate_indexes:
        if index.key_name in candidate_conditions and isinstance(
            candidate_conditions[index.key_name], Equals
        ):
            target_index = index
            partition_key = candidate_conditions.pop(index.key_name)
            partition_key._values = (Key(index.key_name), *partition_key._values[1:])  # type: ignore[union-attr]
            if index.sort_key_name is not None and index.sort_key_name in candidate_conditions:
                sort_key_condition_expression = candidate_conditions.pop(index.sort_key_name)
                sort_key_condition_expression._values = (  # type: ignore[union-attr]
                    Key(index.sort_key_name),
                    *sort_key_condition_expression._values[1:],  # type: ignore[union-attr]
                )
            break

    # convert all remaining to filters
    filter_expressions = list(candidate_conditions.values()) + non_candidate_conditions

    return target_index, partition_key, sort_key_condition_expression, filter_expressions