import io import logging import shutil import tempfile import uuid from datetime import datetime, timedelta from pathlib import Path from typing import Any from uuid import UUID from aioboto3 import Session from aiobotocore.config import AioConfig from pydantic import UUID4 from sqlalchemy import delete, select import product_staging.models.bulk_session_ingestion_execution # noqa: F401 import product_staging.models.bulk_session_ingestion_product # noqa: F401 import product_staging.models.bulk_session_metadata_file # noqa: F401 from product_staging.api.datasources import DATA_SOURCES, S3_CLIENT_KEY from product_staging.api.schemas.bulk_session import AssetStatus, MetadataStatus from product_staging.logic import metadata_json as metadata_json_logic from product_staging.models.bulk_session import BulkSession from product_staging.models.bulk_session_asset import BulkSessionAsset from product_staging.models.bulk_session_asset_cloud_transfer_job import ( BulkSessionAssetCloudTransferJob, ) from product_staging.models.bulk_session_asset_file import BulkSessionAssetFile from product_staging.models.bulk_session_ingestion import BulkSessionIngestion from product_staging.models.bulk_session_ingestion_execution import ( BulkSessionIngestionExecution, ) from product_staging.models.bulk_session_ingestion_product import ( BulkSessionIngestionProduct, ) from product_staging.models.bulk_session_metadata_file import BulkSessionMetadataFile from tests.integration.config import db_session from tests.integration.spreadsheet import randomize_codes logger = logging.getLogger(__name__) async def create_session_with_metadata( vendor_uuid: str, identity_uuid: str, created_on: datetime | None = None ) -> UUID4: """ Creates a bulk session which has validated metadata. ``created_on`` optionally pins the row's created_on instead of relying on the server default (NOW()) - used by slug tests to avoid same-second collisions. """ filename = "Kitty Wizard Catalog.xlsx" file_path = Path() / "tests" / "integration" / "data" / filename bulk_session_id = None s3_client_cm = Session().client( "s3", config=AioConfig(s3={"use_accelerate_endpoint": True}) ) s3_client = await s3_client_cm.__aenter__() DATA_SOURCES[S3_CLIENT_KEY] = s3_client with tempfile.NamedTemporaryFile(delete=True, suffix=".xlsx") as temp_file: shutil.copyfile(file_path, temp_file.name) temp_file.flush() randomize_codes(Path(temp_file.name)) xlsx_stream = io.BytesIO(temp_file.read()) async with db_session() as session: bulk_session_id = uuid.uuid4() bulk_session_metadata_file_id = uuid.uuid4() bulk_session_kwargs: dict[str, Any] = dict( bulk_session_id=str(bulk_session_id), created_by=str(identity_uuid), updated_by=str(identity_uuid), vendor_uuid=str(vendor_uuid), metadata_status=MetadataStatus.valid, asset_status=AssetStatus.default, ) if created_on is not None: bulk_session_kwargs["created_on"] = created_on bulk_session = BulkSession(**bulk_session_kwargs) session.add(bulk_session) await session.flush() stmt = select(BulkSession).where( BulkSession.bulk_session_id == str(bulk_session_id) ) result = await session.execute(stmt) bulk_session = result.scalar_one_or_none() assert bulk_session success_file = BulkSessionMetadataFile( bulk_session_metadata_file_id=str(bulk_session_metadata_file_id), bulk_session_id=str(bulk_session_id), created_by=str(identity_uuid), s3_filename=f"metadata/{uuid.uuid4()}", original_filename=filename, file_status="success", ) session.add(success_file) await session.flush() await session.commit() result = await metadata_json_logic.finalize_metadata_file( bulk_session_id=bulk_session_id, bulk_session_metadata_file_id=bulk_session_metadata_file_id, is_classical=False, identity_uuid=UUID(identity_uuid), xlsx_stream=xlsx_stream, session=session, ) await session.flush() await session.commit() return bulk_session_id async def delete_bulk_sessions_by_vendor_and_created_on( vendor_uuid: str, created_on: datetime ) -> None: """Delete any bulk_sessions (and their asset/metadata-file rows) already in this exact (vendor, second) bucket. The QA integration DB is shared and not reset between CI runs, so a prior run's leftover row in the same bucket would collide with a freshly seeded one. Callers should use a dead-zone `created_on` (e.g. a year-2099 timestamp) that no other test or live traffic will ever produce, so this only ever removes this test's own leftovers. """ async with db_session() as session: stmt = select(BulkSession.bulk_session_id).where( BulkSession.vendor_uuid == str(vendor_uuid), BulkSession.created_on >= created_on, BulkSession.created_on < created_on + timedelta(seconds=1), ) result = await session.execute(stmt) stale_ids = [row[0] for row in result.all()] if not stale_ids: return await session.execute( delete(BulkSessionAsset).where( BulkSessionAsset.bulk_session_id.in_(stale_ids) ) ) await session.execute( delete(BulkSessionMetadataFile).where( BulkSessionMetadataFile.bulk_session_id.in_(stale_ids) ) ) await session.execute( delete(BulkSession).where(BulkSession.bulk_session_id.in_(stale_ids)) ) await session.commit() async def create_asset_transfer_job( bulk_session_id: UUID4, identity_uuid: str ) -> UUID4: """ Creates a transfer job for a given bulk session. Returns the bulk_session_asset_cloud_transfer_job_id. """ transfer_job_id = uuid.uuid4() async with db_session() as session: job = BulkSessionAssetCloudTransferJob( bulk_session_asset_cloud_transfer_job_id=str(transfer_job_id), bulk_session_id=str(bulk_session_id), job_status="in_progress", job_type="google", created_by=str(identity_uuid), ) session.add(job) await session.flush() await session.commit() return transfer_job_id async def create_bulk_session_metadata_file( vendor_uuid: str, identity_uuid: str ) -> tuple[UUID4, UUID4]: """Creates a bulk session with a metadata file record. Returns (bulk_session_id, bulk_session_metadata_file_id). """ async with db_session() as session: bulk_session_id = uuid.uuid4() metadata_file_id = uuid.uuid4() bulk_session = BulkSession( bulk_session_id=str(bulk_session_id), created_by=str(identity_uuid), updated_by=str(identity_uuid), vendor_uuid=str(vendor_uuid), metadata_status=MetadataStatus.valid, asset_status=AssetStatus.default, ) session.add(bulk_session) await session.flush() metadata_file = BulkSessionMetadataFile( bulk_session_metadata_file_id=str(metadata_file_id), bulk_session_id=str(bulk_session_id), s3_filename=f"metadata/{uuid.uuid4()}", original_filename="Kitty Wizard Catalog.xlsx", created_by=str(identity_uuid), file_status="success", ) session.add(metadata_file) await session.commit() return bulk_session_id, metadata_file_id async def create_bulk_session_asset_file( vendor_uuid: str, identity_uuid: str ) -> tuple[UUID4, UUID4]: """Creates a bulk session with an asset file record. Returns (bulk_session_id, bulk_session_asset_file_id). """ bulk_session_id = await create_session_with_metadata(vendor_uuid, identity_uuid) async with db_session() as session: asset_file_id = uuid.uuid4() asset_file = BulkSessionAssetFile( bulk_session_asset_file_id=str(asset_file_id), bulk_session_id=str(bulk_session_id), s3_filename=f"assets/{uuid.uuid4()}/test.jpg", original_filename="test.jpg", created_by=str(identity_uuid), file_status="uploading", file_type="image", ) session.add(asset_file) await session.commit() return bulk_session_id, asset_file_id async def create_bulk_session_with_asset_s3_filename( vendor_uuid: str, identity_uuid: str ) -> tuple[UUID4, str]: """Creates a bulk session with an asset file record. Returns (bulk_session_id, s3_filename). """ bulk_session_id = await create_session_with_metadata(vendor_uuid, identity_uuid) async with db_session() as session: asset_file_id = uuid.uuid4() s3_filename = f"assets/{uuid.uuid4()}/test.jpg" asset_file = BulkSessionAssetFile( bulk_session_asset_file_id=str(asset_file_id), bulk_session_id=str(bulk_session_id), s3_filename=s3_filename, original_filename="test.jpg", created_by=str(identity_uuid), file_status="uploading", file_type="image", ) session.add(asset_file) await session.commit() return bulk_session_id, s3_filename async def create_bulk_session_with_metadata_json_s3_filename( vendor_uuid: str, identity_uuid: str ) -> tuple[UUID4, str, str]: """Creates a bulk session with a finalized metadata JSON file in S3 and a matching bulk_session_ingestion_execution row. The /metadata-hydration endpoint calls log_for_execution, which asserts an execution row exists for the given execution_arn — so the test fixture has to seed both the JSON file and the execution. Returns (bulk_session_id, json_s3_filename, execution_arn) where json_s3_filename is the `metadata_json/{uuid}.json` key produced by finalize_metadata_file. """ bulk_session_id, bulk_session_ingestion_id = await create_bulk_session_ingestion( vendor_uuid, identity_uuid ) execution_arn = f"arn:aws:states:us-east-1:123456789012:execution:my-state-machine:{uuid.uuid4()}" async with db_session() as session: stmt = select(BulkSessionMetadataFile.json_file_name).where( BulkSessionMetadataFile.bulk_session_id == str(bulk_session_id), BulkSessionMetadataFile.file_status == "success", ) json_file_name = (await session.execute(stmt)).scalar_one() execution = BulkSessionIngestionExecution( bulk_session_ingestion_execution_id=str(uuid.uuid4()), bulk_session_ingestion_id=str(bulk_session_ingestion_id), execution_arn=execution_arn, ) session.add(execution) await session.commit() assert json_file_name return bulk_session_id, json_file_name, execution_arn async def create_bulk_session_ingestion( vendor_uuid: str, identity_uuid: str ) -> tuple[UUID4, UUID4]: """Creates a bulk session with a bulk_session_ingestion record. Returns (bulk_session_id, bulk_session_ingestion_id). """ bulk_session_id = await create_session_with_metadata(vendor_uuid, identity_uuid) async with db_session() as session: ingestion_id = uuid.uuid4() ingestion = BulkSessionIngestion( bulk_session_ingestion_id=str(ingestion_id), bulk_session_id=str(bulk_session_id), ingestion_status="in_progress", assets_required=False, submit_products=False, send_notifications=False, created_by=str(identity_uuid), updated_by=str(identity_uuid), ) session.add(ingestion) await session.commit() return bulk_session_id, ingestion_id async def create_bulk_session_ingestion_product( vendor_uuid: str, identity_uuid: str ) -> tuple[UUID4, UUID4, UUID4, str]: """Creates a bulk session with an ingestion record and an ingestion product. Returns (bulk_session_id, bulk_session_ingestion_id, bulk_session_ingestion_product_id, product_code). """ bulk_session_id, bulk_session_ingestion_id = await create_bulk_session_ingestion( vendor_uuid, identity_uuid ) async with db_session() as session: ingestion_product_id = uuid.uuid4() product_code = f"PC-{uuid.uuid4()}" ingestion_product = BulkSessionIngestionProduct( bulk_session_ingestion_product_id=str(ingestion_product_id), bulk_session_ingestion_id=str(bulk_session_ingestion_id), product_id=1, upc=str(uuid.uuid4().int)[:12], product_code=product_code, submit_status="default", ingestion_status="in_progress", ) session.add(ingestion_product) await session.commit() return ( bulk_session_id, bulk_session_ingestion_id, ingestion_product_id, product_code, ) async def get_session_product_asset_filenames(bulk_session_id: UUID4) -> list[str]: """ Retrieves the original filenames of all product assets associated with a bulk session. """ async with db_session() as session: stmt = select(BulkSessionAsset.original_filename).where( BulkSessionAsset.bulk_session_id == str(bulk_session_id) ) result = await session.execute(stmt) return list(result.scalars().all())