"""Marshmallow schemas for search.""" import datetime import zoneinfo from typing import Any import marshmallow from marshmallow import decorators, exceptions, fields, validate from vectororder import config from vectororder.constants import ( delivery_statuses, encoders, encoding_statuses, error, search, ) from vectororder.models.schemas import JobStatus, Priority MYSQL_INT_SIGNED_MAX_VALUE = 2147483647 MYSQL_INT_UNSIGNED_MAX_VALUE = 4294967295 MYSQL_SMALLINT_UNSIGNED_MAX_VALUE = 65535 DISPLAY_UPC_MIN_LENGTH = 1 DISPLAY_UPC_MAX_LENGTH = 15 DATE_RANGE_FIELDS = ( # 'from' date, 'to' date, destination key (("created_at_from", "created_at_to"), "created_at"), (("encoding_start_from", "encoding_start_to"), "encoding_started"), (("encoding_end_from", "encoding_end_to"), "encoding_ended"), (("delivery_start_from", "delivery_start_to"), "delivery_started"), (("delivery_end_from", "delivery_end_to"), "delivery_ended"), ) # Input dates are US/Eastern; ES stores UTC. All date → datetime conversions # use this timezone so DST transitions are handled automatically. _EASTERN = zoneinfo.ZoneInfo("US/Eastern") class SearchPOSTSchema(marshmallow.Schema): """Vector order search POST Schema. The following operations are performed: - validate individual input fields values - convert incoming field names to corresponding search index fields - post-process certain field values according to search index schema """ types = fields.List( fields.String(validate=validate.OneOf(search.VO_TYPES)), attribute="meta_update", ) priorities = fields.List( fields.Int(validate=validate.OneOf([priority.value for priority in Priority])), attribute="priority", ) display_upcs = fields.List( fields.String( validate=[ validate.Length( min=DISPLAY_UPC_MIN_LENGTH, error=error.ERROR_MESSAGE_DISPLAY_UPC_SHORT.format( DISPLAY_UPC_MIN_LENGTH ), ), validate.Length( max=DISPLAY_UPC_MAX_LENGTH, error=error.ERROR_MESSAGE_DISPLAY_UPC_LONG.format( DISPLAY_UPC_MAX_LENGTH ), ), validate.Regexp( r"[0-9]+", error=error.ERROR_MESSAGE_DISPLAY_UPC_INVALID ), ], ), attribute="display_upc", ) product_ids = fields.List( fields.Int( validate=validate.Range( min=0, max=MYSQL_INT_UNSIGNED_MAX_VALUE, error=error.ERROR_MESSAGE_PRODUCT_ID_OUT_OF_RANGE.format( MYSQL_INT_UNSIGNED_MAX_VALUE ), ), error_messages={"invalid": error.ERROR_MESSAGE_PRODUCT_ID_INVALID}, ), attribute="product_id", ) order_ids = fields.List( fields.Int(validate=validate.Range(min=0, max=MYSQL_INT_SIGNED_MAX_VALUE)), attribute="order_id", ) statuses = fields.List( fields.String(validate=validate.Regexp(r"\w+")), attribute="status", ) user_ids = fields.List( fields.Int(validate=validate.Range(min=0, max=MYSQL_INT_UNSIGNED_MAX_VALUE)), attribute="user_id", ) encoder_ids = fields.List( fields.Int(validate=validate.OneOf(encoders.ENCODERS.keys())), attribute="encoder_id", ) encoding = fields.List( fields.String( validate=validate.OneOf(encoding_statuses.ENCODING_STATUSES.keys()) ) ) encoding_queue_ids = fields.List( fields.Int(validate=validate.Range(min=0, max=MYSQL_INT_SIGNED_MAX_VALUE)), attribute="encoding_queue_id", ) delivery = fields.List( fields.String( validate=validate.OneOf(delivery_statuses.DELIVERY_STATUSES.keys()) ) ) store_ids = fields.List( fields.Int( validate=validate.Range(min=0, max=MYSQL_SMALLINT_UNSIGNED_MAX_VALUE) ), attribute="store_id", ) error_log = fields.String() # Date/time fields accept ISO 8601. User input is Eastern dates (YYYY-MM-DD); # _to_date_ranges converts them to UTC datetimes for ES range queries. created_at_from = fields.DateTime() created_at_to = fields.DateTime() encoding_start_from = fields.DateTime() encoding_start_to = fields.DateTime() encoding_end_from = fields.DateTime() encoding_end_to = fields.DateTime() delivery_start_from = fields.DateTime() delivery_start_to = fields.DateTime() delivery_end_from = fields.DateTime() delivery_end_to = fields.DateTime() @decorators.post_load def post_load( self, data: dict[str, Any], partial: bool = False, many: bool = False ) -> dict[str, Any]: """Perform processing after the load step. Modifies passed data inplace. Args: data (dict): Loaded data. """ # We have to guarantee the order of these calls, so can't use # separate @post_load decorated methods. # This should come before date range processing, since it can override # encoding_started and delivery_started dates. data = self._process_statuses(data) data = self._to_date_ranges(data) # We don't have the return value of the latest call. return self._order_types_to_meta_update(data) @decorators.validates_schema def validate_dates( self, data: dict[str, Any], partial: bool = False, many: bool = False ) -> None: """Validate input date pairs.""" for (date_from_field, date_to_field), _ in DATE_RANGE_FIELDS: date_from = data.get(date_from_field) date_to = data.get(date_to_field) # If any of the two fields is not set - the range is valid. if not all((date_to, date_from)): continue if ( isinstance(date_from, datetime.datetime) and isinstance(date_to, datetime.datetime) and date_from > date_to ): raise exceptions.ValidationError( error.ERROR_MESSAGE_INVALID_DATE_RANGE.format( date_from_field, date_to_field ) ) @decorators.validates_schema(pass_original=True) def validate_unknown_fields( self, _: Any, original_data: dict[str, Any], partial: bool = False, many: bool = False, ) -> None: """Validate data to have no unknown fields present.""" unknown = set(original_data) - set(self.fields) if unknown: raise exceptions.ValidationError( error.ERROR_MESSAGE_UNKNOWN_FIELD, str(unknown) ) @staticmethod def _to_date_ranges(data: dict[str, Any]) -> dict[str, Any]: """Convert input date pairs into date ranges under different keys. Modifies passed data inplace. Args: data (dict): Loaded data. Returns (dict): Processed data. """ for (date_from_field, date_to_field), dest_field in DATE_RANGE_FIELDS: date_from = data.pop(date_from_field, None) date_to = data.pop(date_to_field, None) # If both values are not set - do not insert any date range. if not any((date_from, date_to)): continue # Inputs are Eastern; convert to UTC for ES. Naive datetimes # produced by date-only strings (e.g. "2026-04-20") are treated # as Eastern by attaching the timezone without shifting the clock. if isinstance(date_from, datetime.datetime): date_from = date_from.replace(tzinfo=_EASTERN).astimezone(datetime.UTC) if isinstance(date_to, datetime.datetime): if not any( (date_to.hour, date_to.minute, date_to.second, date_to.microsecond) ): # No time provided — cover the whole Eastern day. date_to = datetime.datetime( date_to.year, date_to.month, date_to.day, 23, 59, 59, 999999, tzinfo=_EASTERN, ).astimezone(datetime.UTC) else: date_to = date_to.replace(tzinfo=_EASTERN).astimezone(datetime.UTC) data[dest_field] = (date_from, date_to) return data @staticmethod def _process_statuses(data: dict[str, Any]) -> dict[str, Any]: """Convert input *_stuck statuses into date ranges and real statuses. Modifies passed data inplace. After processing regular values remain on the top level. Logical disjunction inner queries are represented by tuples. Result example: { 'encoder_id': [10, 17, 18, 19, 23], '_statuses': ( {'status': ['encoding'], 'encoding_started': (None, datetime.datetime(2018, 5, 13, 19, 3, 48, 277749))}, {'status': ['cancelled']}, ), } Args: data (dict): Loaded data. Returns: dict: Processed data. """ stuck_datetime = datetime.datetime.now(datetime.UTC) - datetime.timedelta( minutes=config.STATUS_STUCK_MINUTES ) status_list = data.pop("status", []) encoding_stuck_fields: dict[str, Any] = {} delivering_stuck_fields: dict[str, Any] = {} if JobStatus.ENCODING_STUCK.value in status_list: status_list.remove(JobStatus.ENCODING_STUCK.value) encoding_stuck_fields["status"] = [JobStatus.ENCODING.value] encoding_stuck_fields["encoding_started"] = (None, stuck_datetime) if JobStatus.DELIVERING_STUCK.value in status_list: status_list.remove(JobStatus.DELIVERING_STUCK.value) delivering_stuck_fields["status"] = [JobStatus.DELIVERING.value] delivering_stuck_fields["delivery_started"] = (None, stuck_datetime) status_fields = {"status": status_list} if status_list else {} # This inner local disjunction holds a list of simple statuses and/or # a combination of status and its start date(s) dicts. all_status_fields = ( encoding_stuck_fields, delivering_stuck_fields, status_fields, ) non_empty_status_fields = tuple(f for f in all_status_fields if f) if non_empty_status_fields: data[search.VO_STATUSES] = non_empty_status_fields return data @staticmethod def _order_types_to_meta_update(data: dict[str, Any]) -> dict[str, Any]: """Convert incoming order types to a single meta_update field. Modifies passed data inplace. Args: data (dict): Loaded data. Returns (dict): Processed data. """ field_name = "meta_update" # matches the 'types' field attribute. vo_types = data.get(field_name) if not vo_types: return data vo_types = set(vo_types) # If both possible values are present - do not include this field. if vo_types == set(search.VO_TYPES): data.pop(field_name) # Convert to: {meta_update: True} elif vo_types == {search.VO_TYPE_META_UPDATE}: data[field_name] = True # Convert to: {meta_update: False} elif vo_types == {search.VO_TYPE_DELIVERY}: data[field_name] = False # Note: invalid input is checked by the schema field validators. return data