diff --git a/application/single_app/app.py b/application/single_app/app.py index 05fa0c8de..faf66d428 100644 --- a/application/single_app/app.py +++ b/application/single_app/app.py @@ -286,7 +286,7 @@ def start_background_tasks(): print("Background tasks disabled for this web process.") _background_tasks_started = True return - start_background_task_threads() + start_background_task_threads(app=app) _background_tasks_started = True diff --git a/application/single_app/background_tasks.py b/application/single_app/background_tasks.py index a10edd908..3dea3258f 100644 --- a/application/single_app/background_tasks.py +++ b/application/single_app/background_tasks.py @@ -700,14 +700,14 @@ def run_tabular_generated_output_scheduler_loop(): time.sleep(30) -def run_data_management_scheduler_loop(): - """Queue due Data Management backup jobs across scaled-out workers.""" +def run_data_management_scheduler_loop(app=None): + """Queue due backup and recoverable migration jobs across scaled-out workers.""" while True: lock_document = None try: lock_document = acquire_distributed_task_lock('data_management_scheduler_scan', lease_seconds=300) if lock_document: - check_due_data_management_jobs_once() + check_due_data_management_jobs_once(app=app) except Exception as exc: print(f"Error in Data Management scheduler check: {exc}") log_event(f"[DataManagement] Error in scheduler check: {exc}", level=logging.ERROR) @@ -752,7 +752,7 @@ def run_app_maintenance_loop(): time.sleep(max(int(sleep_seconds or 3600), 15)) -def start_background_task_threads(): +def start_background_task_threads(app=None): """Start all background task loops for the current process.""" task_specs = [ ('Logging timer background task started.', run_logging_timer_loop), @@ -763,7 +763,7 @@ def start_background_task_threads(): ('Workflow scheduler background task started.', run_workflow_scheduler_loop), ('File Sync scheduler background task started.', run_file_sync_scheduler_loop), ('Tabular generated-output scheduler background task started.', run_tabular_generated_output_scheduler_loop), - ('Data Management scheduler background task started.', run_data_management_scheduler_loop), + ('Data Management scheduler background task started.', lambda: run_data_management_scheduler_loop(app=app)), ('App maintenance background task started.', run_app_maintenance_loop), ] diff --git a/application/single_app/config.py b/application/single_app/config.py index a3f37abd0..a69d7c20c 100644 --- a/application/single_app/config.py +++ b/application/single_app/config.py @@ -97,7 +97,7 @@ EXECUTOR_TYPE = 'thread' EXECUTOR_MAX_WORKERS = 30 SESSION_TYPE = 'filesystem' -VERSION = "0.250.070" +VERSION = "0.250.071" IS_DEVELOPMENT = is_development_env_enabled() SESSION_COOKIE_SAMESITE = os.getenv('SESSION_COOKIE_SAMESITE', 'Lax') diff --git a/application/single_app/functions_data_management.py b/application/single_app/functions_data_management.py index 376b2f5dd..10049e4ef 100644 --- a/application/single_app/functions_data_management.py +++ b/application/single_app/functions_data_management.py @@ -2,6 +2,7 @@ """Data Management settings, schedules, and durable job records.""" import copy +import base64 import hashlib import json import logging @@ -11,17 +12,21 @@ import time import tempfile import uuid +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait from datetime import datetime, timedelta, timezone +from threading import Event, Lock, Thread +from urllib.parse import urlparse from azure.core import MatchConditions from azure.core.credentials import AzureKeyCredential +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError import azure.cosmos as azure_cosmos from azure.cosmos import PartitionKey from azure.cosmos.exceptions import CosmosResourceNotFoundError from azure.identity import DefaultAzureCredential from azure.search.documents import SearchClient from azure.search.documents.indexes import SearchIndexClient -from azure.search.documents.indexes.models import SearchIndex +from azure.search.documents.indexes.models import SearchField, SearchFieldDataType, SearchIndex from azure.storage.blob import BlobServiceClient, ContentSettings from cryptography.fernet import Fernet @@ -34,12 +39,59 @@ cosmos_settings_container, ) from functions_appinsights import log_event +from functions_cosmos_throughput import ( + CosmosThroughputError, + get_container_throughput, + get_database_throughput, + set_database_throughput, +) +from functions_data_management_migration_state import ( + MIGRATION_RESOURCE_STATUS_COMPLETED, + MIGRATION_RESOURCE_STATUS_FAILED, + build_transfer_metrics, + complete_migration_resource, + fail_migration_resource, + initialize_migration_state, + is_migration_resource_completed, + start_migration_resource, + update_migration_resource, +) +from functions_data_management_search_write_fence import ( + acquire_data_management_search_write_fence, + acquire_data_management_target_migration_coordinator, + release_data_management_search_write_fence, + release_data_management_target_migration_coordinator, + renew_data_management_search_write_fence, + renew_data_management_target_migration_coordinator, +) +from functions_migration_provenance import ( + BLOB_MIGRATION_STATUS_METADATA_KEY, + COSMOS_MIGRATION_PROVENANCE_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + add_cosmos_migration_provenance, + add_search_migration_provenance, + create_migration_provenance_context, + get_blob_migration_provenance, + get_cosmos_migration_provenance, + get_search_migration_provenance, + is_successful_migration_record, + merge_blob_migration_metadata, + migration_record_matches_source, + should_skip_migration_record, +) DATA_MANAGEMENT_SETTINGS_ID = "backup_settings" DATA_MANAGEMENT_SETTINGS_TYPE = "data_management_settings" DATA_MANAGEMENT_JOB_TYPE = "data_management_job" DATA_MANAGEMENT_JOB_ITEM_TYPE = "data_management_job_item" +DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE = "data_management_migration_manifest_batch" +DATA_MANAGEMENT_MIRROR_DELETION_BATCH_TYPE = "data_management_mirror_deletion_batch" +DATA_MANAGEMENT_MIGRATION_LOCK_TYPE = "data_management_migration_lock" DATA_MANAGEMENT_OPERATION_BACKUP = "backup" DATA_MANAGEMENT_OPERATION_RESTORE = "restore" @@ -83,8 +135,40 @@ DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS = 900 DATA_MANAGEMENT_DEFAULT_STALE_SECONDS = 1200 DATA_MANAGEMENT_DEFAULT_JOB_LIMIT = 25 +DATA_MANAGEMENT_DEFAULT_RECOVERY_JOB_LIMIT = 25 +DATA_MANAGEMENT_RECOVERY_QUEUE_DELAY_SECONDS = 60 +DATA_MANAGEMENT_RECOVERY_RESUBMIT_DELAY_SECONDS = 120 +DATA_MANAGEMENT_MIGRATION_LOCK_PREFIX = "data_management_migration_lock" DATA_MANAGEMENT_MIGRATION_CATALOG_LIMIT = 50 DATA_MANAGEMENT_MIGRATION_BATCH_SIZE = 500 +DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE = 100 +DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE = 1000 +DATA_MANAGEMENT_SEARCH_SCOPE_FILTER_BATCH_SIZE = 100 +DATA_MANAGEMENT_BLOB_SERVICE_TIMEOUT_SECONDS = 120 +DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS = 30 +DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS = 2.0 +DATA_MANAGEMENT_MIGRATION_HEARTBEAT_POLL_SECONDS = 1.0 +DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS = 120 +DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY = "new_only" +DATA_MANAGEMENT_MIGRATION_MODE_DELTA_UPSERT = "delta_upsert" +DATA_MANAGEMENT_MIGRATION_MODE_MIRROR = "mirror_with_deletions" +DATA_MANAGEMENT_MIGRATION_MODES = { + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + DATA_MANAGEMENT_MIGRATION_MODE_DELTA_UPSERT, + DATA_MANAGEMENT_MIGRATION_MODE_MIRROR, +} +DATA_MANAGEMENT_MIRROR_CONFIRMATION = "MIRROR WITH DELETIONS" +DATA_MANAGEMENT_SEARCH_WRITE_FREEZE_CONFIRMATION_ERROR = ( + "Confirm that external destination AI Search writers are frozen before migrating AI Search documents." +) +DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS = 8 +DATA_MANAGEMENT_MIGRATION_MAX_PARALLEL_OPERATIONS = 32 +DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT = 5 +DATA_MANAGEMENT_MIGRATION_MAX_RETRY_COUNT = 10 +DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS = 0 +DATA_MANAGEMENT_MIGRATION_MAX_SKIP_WITHIN_HOURS = 8760 +DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU = 10000 +DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU = 10000 DATA_MANAGEMENT_KEY_VAULT_SCOPE_VALUE = "data-management" DATA_MANAGEMENT_ENCRYPTION_SECRET_NAME = "backup-encryption-key" DATA_MANAGEMENT_REDACTED_VALUE = "***REDACTED***" @@ -92,6 +176,8 @@ "We suggest not running backups, restores, or migrations during your operational business hours. " "These jobs run inside the App Service environment and can affect application performance." ) +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME = {} +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK = Lock() DATA_MANAGEMENT_DEFAULT_SETTINGS = { "enabled": False, @@ -104,6 +190,8 @@ "target_cosmos_endpoint": "", "target_cosmos_key": "", "target_cosmos_database_name": DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, + "target_cosmos_subscription_id": "", + "target_cosmos_resource_group": "", "target_ai_search_authentication_type": "managed_identity", "target_ai_search_endpoint": "", "target_ai_search_key": "", @@ -122,6 +210,11 @@ "include_source_blobs": True, "low_impact_mode": True, "max_parallel_operations": 1, + "migration_max_parallel_operations": DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS, + "migration_retry_count": DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT, + "migration_skip_recent_within_hours": DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS, + "migration_temporary_destination_ru_enabled": False, + "migration_temporary_destination_ru": DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, "next_full_backup_run_at": None, "next_partial_backup_run_at": None, "last_full_backup_completed_at": None, @@ -137,6 +230,14 @@ class DataManagementSettingsValidationError(ValueError): class DataManagementCosmosEditorError(ValueError): """Raised when Cosmos editor input is unsafe or incomplete.""" + +class DataManagementMigrationLeaseLostError(RuntimeError): + """Raised when a stale migration worker no longer owns its job lease.""" + + +class DataManagementMigrationCanceledError(RuntimeError): + """Raised when an administrator requests cooperative migration cancellation.""" + DATA_MANAGEMENT_FRONTEND_SECRET_FIELDS = { "backup_storage_connection_string", "encryption_key_reference", @@ -146,6 +247,7 @@ class DataManagementCosmosEditorError(ValueError): } DATA_MANAGEMENT_MIGRATION_TARGET_TYPES = {"users", "groups", "public_workspaces"} +DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER = ("users", "groups", "public_workspaces") DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { "users": [ @@ -498,6 +600,8 @@ def normalize_data_management_settings(payload=None, existing_settings=None, cur source["target_cosmos_authentication_type"] = "managed_identity" source["target_cosmos_endpoint"] = _safe_text(source.get("target_cosmos_endpoint")) source["target_cosmos_key"] = _safe_text(source.get("target_cosmos_key")) + source["target_cosmos_subscription_id"] = _safe_text(source.get("target_cosmos_subscription_id")) + source["target_cosmos_resource_group"] = _safe_text(source.get("target_cosmos_resource_group")) if source["target_cosmos_authentication_type"] == "managed_identity": source["target_cosmos_key"] = "" source["target_cosmos_database_name"] = DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME @@ -534,6 +638,34 @@ def normalize_data_management_settings(payload=None, existing_settings=None, cur source["include_source_blobs"] = False source["low_impact_mode"] = _safe_bool(source.get("low_impact_mode"), True) source["max_parallel_operations"] = _safe_int(source.get("max_parallel_operations"), default=1, minimum=1, maximum=5) + source["migration_max_parallel_operations"] = _safe_int( + source.get("migration_max_parallel_operations"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_PARALLEL_OPERATIONS, + ) + source["migration_retry_count"] = _safe_int( + source.get("migration_retry_count"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_RETRY_COUNT, + ) + source["migration_skip_recent_within_hours"] = _safe_int( + source.get("migration_skip_recent_within_hours"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS, + minimum=0, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_SKIP_WITHIN_HOURS, + ) + source["migration_temporary_destination_ru_enabled"] = _safe_bool( + source.get("migration_temporary_destination_ru_enabled"), + False, + ) + source["migration_temporary_destination_ru"] = _safe_int( + source.get("migration_temporary_destination_ru"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, + minimum=1000, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU, + ) for key in ("next_full_backup_run_at", "next_partial_backup_run_at", "last_full_backup_completed_at", "last_partial_backup_completed_at"): parsed = _parse_iso_datetime(source.get(key)) @@ -698,16 +830,47 @@ def _normalize_data_management_settings_from_payload(settings=None): return normalize_data_management_settings(existing_settings=existing_settings, application_settings=application_settings) -def test_target_cosmos_connection(settings=None): +def test_target_cosmos_connection(settings=None, migration_plan=None): normalized_settings = _normalize_data_management_settings_from_payload(settings) target_database = _get_target_cosmos_database(normalized_settings) properties = target_database.read() - return { + normalized_plan = normalize_data_management_migration_plan({ + "migration_plan": migration_plan if isinstance(migration_plan, dict) else {}, + }) + selected_total = sum( + len((normalized_plan.get(target_type) or {}).get("ids") or []) + if (normalized_plan.get(target_type) or {}).get("mode") == "selected" + else 1 if (normalized_plan.get(target_type) or {}).get("mode") == "all" + else 0 + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER + ) + migration_access = None + if selected_total: + migration_access = _preflight_target_cosmos_migration_access( + normalized_settings, + normalized_plan, + ) + capacity = None + if normalized_settings.get("migration_temporary_destination_ru_enabled"): + inspected_capacity = _inspect_target_cosmos_migration_capacity( + normalized_settings, + normalized_plan, + ) + capacity = { + "target_ru": inspected_capacity.get("target_ru"), + "database_mode": inspected_capacity.get("database_mode"), + "database_current_ru": inspected_capacity.get("database_current_ru"), + "targets": inspected_capacity.get("targets"), + } + result = { "success": True, "target": "cosmos", "database_name": properties.get("id") or DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, "authentication_type": normalized_settings.get("target_cosmos_authentication_type"), + "migration_access": migration_access, + "capacity": capacity, } + return result def test_target_search_connection(settings=None): @@ -715,7 +878,13 @@ def test_target_search_connection(settings=None): endpoint = _safe_text(normalized_settings.get("target_ai_search_endpoint")) if not endpoint: raise ValueError("Target Search endpoint is required.") - index_client = SearchIndexClient(endpoint=endpoint, credential=_get_target_ai_search_credential(normalized_settings)) + index_client = SearchIndexClient( + endpoint=endpoint, + credential=_get_target_ai_search_credential(normalized_settings), + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) existing_indexes = set(index_client.list_index_names()) expected_indexes = [artifact["index_name"] for artifact in DATA_MANAGEMENT_SEARCH_ARTIFACTS] return { @@ -784,7 +953,7 @@ def _normalize_migration_selection(selection=None): } -def _get_target_cosmos_database(settings): +def _get_target_cosmos_client(settings): endpoint = _safe_text((settings or {}).get("target_cosmos_endpoint")) if not endpoint: raise ValueError("Target Cosmos endpoint is required before running migration.") @@ -792,304 +961,5474 @@ def _get_target_cosmos_database(settings): key = _safe_text((settings or {}).get("target_cosmos_key")) if not key: raise ValueError("Target Cosmos account key is required when account key authentication is selected.") - client = azure_cosmos.CosmosClient(endpoint, credential=key, consistency_level="Session") - else: - client = azure_cosmos.CosmosClient(endpoint, credential=DefaultAzureCredential(), consistency_level="Session") - return client.create_database_if_not_exists(DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME) + return azure_cosmos.CosmosClient( + endpoint, + credential=key, + consistency_level="Session", + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + return azure_cosmos.CosmosClient( + endpoint, + credential=DefaultAzureCredential(), + consistency_level="Session", + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + + +def _get_target_cosmos_database(settings): + return _get_target_cosmos_client(settings).create_database_if_not_exists( + DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME + ) + + +def _get_existing_target_cosmos_database(settings): + """Open the destination database without creating resources during preview.""" + database = _get_target_cosmos_client(settings).get_database_client( + DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME + ) + database.read() + return database def _get_target_cosmos_container(target_database, container_name, partition_key_path): - return target_database.create_container_if_not_exists( + target_container = target_database.create_container_if_not_exists( id=container_name, partition_key=PartitionKey(path=partition_key_path), ) + _validate_target_cosmos_container_partition_key( + target_container, + container_name, + partition_key_path, + ) + return target_container + + +def _get_target_data_management_search_write_gate_container(target_database): + """Open the target SimpleChat job container used to coordinate Search writers.""" + return _get_target_cosmos_container( + target_database, + getattr( + app_config, + "cosmos_data_management_jobs_container_name", + "data_management_jobs", + ), + "/id", + ) -def _iter_selected_cosmos_records(container_definition, selection): - source_container = getattr(app_config, container_definition["container_attr"], None) - if not source_container: - return - if container_definition.get("documents") and not selection.get("include_documents"): - return - mode = selection.get("mode") - if mode == "all": - yield from _iter_cosmos_container_items(source_container) - return - if mode != "selected": - return - - ids = selection.get("ids") or [] - if container_definition.get("id_field") == "id": - for item_id in ids: - try: - yield _strip_cosmos_system_fields(source_container.read_item(item=item_id, partition_key=item_id)) - except CosmosResourceNotFoundError: - continue - return - - filter_field = container_definition.get("filter_field") - if not filter_field: - return - for selected_id in ids: - query = f"SELECT * FROM c WHERE c.{filter_field} = @selected_id" - for item in source_container.query_items( - query=query, - parameters=[{"name": "@selected_id", "value": selected_id}], - enable_cross_partition_query=True, - ): - yield _strip_cosmos_system_fields(item) +def _get_target_search_write_fence_lease_seconds(settings): + """Keep uncertain target Search writes quarantined beyond their bounded request lifetime.""" + return min( + _get_migration_lock_lease_seconds(settings), + DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + + DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS, + ) -def _copy_cosmos_records_to_target(target_database, target_type, selection): - copied = [] - for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: - target_container_name = getattr(app_config, container_definition["container_name_attr"], container_definition["name"]) - target_container = _get_target_cosmos_container( - target_database, - target_container_name, - container_definition["partition_key_path"], - ) - count = 0 - for item in _iter_selected_cosmos_records(container_definition, selection) or []: - target_container.upsert_item(item) - count += 1 - copied.append({ - "name": container_definition["name"], - "type": "cosmos_container", - "target_type": target_type, - "container_name": target_container_name, - "item_count": count, - "partition_key_path": container_definition["partition_key_path"], - }) - return copied +def _set_target_migration_coordinator_runtime(job_id, container, coordinator): + """Keep non-serializable target SDK state out of durable migration job records.""" + with DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK: + DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME[_safe_text(job_id)] = { + "container": container, + "coordinator": coordinator, + } -def _escape_search_filter_value(value): - return _safe_text(value).replace("'", "''") +def _get_target_migration_coordinator_runtime(job_id): + with DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK: + return DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME.get(_safe_text(job_id)) -def _build_search_filter(field_name, selection): - if selection.get("mode") == "all": - return None - if selection.get("mode") != "selected": - return "id eq '__no_migration_selection__'" - conditions = [f"{field_name} eq '{_escape_search_filter_value(item_id)}'" for item_id in selection.get("ids") or []] - if not conditions: - return "id eq '__no_migration_selection__'" - return " or ".join(conditions) +def _clear_target_migration_coordinator_runtime(job_id): + with DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME_LOCK: + return DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_RUNTIME.pop( + _safe_text(job_id), + None, + ) -def _get_target_ai_search_credential(settings): - if (settings or {}).get("target_ai_search_authentication_type") == "key": - key = _safe_text((settings or {}).get("target_ai_search_key")) - if not key: - raise ValueError("Target AI Search key is required when key authentication is selected.") - return AzureKeyCredential(key) - return DefaultAzureCredential() +def _acquire_target_migration_coordinator(job, target_database, settings): + """Acquire a destination-wide lease before this source performs any migration work.""" + target_container = _get_target_data_management_search_write_gate_container( + target_database + ) + coordinator = acquire_data_management_target_migration_coordinator( + target_container, + job.get("id"), + _get_migration_lock_lease_seconds(settings), + existing_lock=( + job.get("target_migration_coordinator") + if isinstance(job.get("target_migration_coordinator"), dict) + else None + ), + ) + job["target_migration_coordinator"] = copy.deepcopy(coordinator) + _set_target_migration_coordinator_runtime( + job.get("id"), + target_container, + coordinator, + ) + return coordinator -def _get_target_search_client(settings, index_name): - endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) - if not endpoint: - raise ValueError("Target AI Search endpoint is required before running AI Search migration.") - return SearchClient(endpoint=endpoint, index_name=index_name, credential=_get_target_ai_search_credential(settings)) +def _renew_target_migration_coordinator(job, settings): + """Renew the destination-wide coordinator that protects independent source deployments.""" + runtime = _get_target_migration_coordinator_runtime(job.get("id")) + if not isinstance(runtime, dict): + return + coordinator = runtime.get("coordinator") + target_container = runtime.get("container") + if not isinstance(coordinator, dict) or target_container is None: + return + try: + renewed_coordinator = renew_data_management_target_migration_coordinator( + target_container, + coordinator, + _get_migration_lock_lease_seconds({ + "data_management_job_lease_seconds": ( + (settings or {}).get("data_management_job_lease_seconds") or + coordinator.get("lease_seconds") + ), + }), + ) + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "The target migration coordinator was lost or could not be renewed." + ) from exc + job["target_migration_coordinator"] = copy.deepcopy(renewed_coordinator) -def _ensure_target_search_index(settings, index_name, schema_file): - endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) - if not endpoint: - raise ValueError("Target AI Search endpoint is required before running AI Search migration.") - index_client = SearchIndexClient(endpoint=endpoint, credential=_get_target_ai_search_credential(settings)) +def _release_target_migration_coordinator(job): + """Release a cleanly completed or drained-canceled destination coordinator lease.""" + runtime = _clear_target_migration_coordinator_runtime(job.get("id")) + if not isinstance(runtime, dict): + return False try: - index_client.get_index(index_name) - return "exists" + return release_data_management_target_migration_coordinator( + runtime.get("container"), + runtime.get("coordinator"), + ) except Exception: - schema = _get_search_schema(schema_file) - schema = {key: value for key, value in schema.items() if not key.startswith("@odata.")} - try: - index_client.create_or_update_index(SearchIndex.from_dict(schema)) - return "created" - except Exception as exc: - raise ValueError(f"Target AI Search index {index_name} is missing and could not be created: {exc}") from exc + return False -def _upload_search_documents_in_batches(search_client, documents): - uploaded = 0 - batch = [] - for document in documents: - batch.append(document) - if len(batch) >= DATA_MANAGEMENT_MIGRATION_BATCH_SIZE: - search_client.upload_documents(documents=batch) - uploaded += len(batch) - batch = [] - if batch: - search_client.upload_documents(documents=batch) - uploaded += len(batch) - return uploaded +def _validate_target_cosmos_container_partition_key(target_container, container_name, expected_path): + """Reject existing target containers whose partition key differs from SimpleChat's contract.""" + container_read = getattr(target_container, "read", None) + if not callable(container_read): + return + properties = container_read() + partition_key = properties.get("partitionKey") if isinstance(properties, dict) else None + actual_paths = partition_key.get("paths") if isinstance(partition_key, dict) else None + if not isinstance(actual_paths, list) or actual_paths != [expected_path]: + raise DataManagementSettingsValidationError( + f"Destination Cosmos container '{container_name}' has partition key paths " + f"{actual_paths!r}; expected [{expected_path!r}]." + ) -def _copy_ai_search_to_target(settings, migration_plan): - copied = [] - if not migration_plan.get("include_ai_search"): - return copied - search_mappings = [ - ("users", "user_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), - ("groups", "group_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), - ("public_workspaces", "public_workspace_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), - ] - for target_type, field_name, artifact in search_mappings: - selection = migration_plan.get(target_type) or {} - if selection.get("mode") == "none" or not selection.get("include_documents"): - continue - source_client = CLIENTS.get(artifact["client_key"]) - if not source_client: - copied.append({"name": artifact["name"], "type": "ai_search_documents", "status": "skipped", "warning": "Source AI Search client is not initialized."}) +def _iter_migration_cosmos_container_definitions(migration_plan): + """Yield only destination containers that the selected migration can write.""" + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = (migration_plan or {}).get(target_type) or {} + if selection.get("mode") == "none": continue - index_status = _ensure_target_search_index(settings, artifact["index_name"], artifact["schema_file"]) - target_client = _get_target_search_client(settings, artifact["index_name"]) - search_filter = _build_search_filter(field_name, selection) - documents = _iter_search_documents_for_filter(source_client, search_filter) - uploaded = _upload_search_documents_in_batches(target_client, documents) - copied.append({ - "name": artifact["name"], - "type": "ai_search_documents", - "target_type": target_type, - "index_name": artifact["index_name"], - "item_count": uploaded, - "partial_filter": search_filter, - "index_status": index_status, - }) - return copied - - -def _iter_search_documents_for_filter(search_client, search_filter=None): - for result in search_client.search(search_text="*", filter=search_filter, include_total_count=True): - document = dict(result) - yield {key: value for key, value in document.items() if not key.startswith("@search.")} + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + yield target_type, container_definition -def _get_target_enhanced_citations_blob_client(settings): - auth_type = _safe_text((settings or {}).get("target_enhanced_citations_storage_authentication_type"), "managed_identity") - if auth_type == "connection_string": - connection_string = _safe_text((settings or {}).get("target_enhanced_citations_storage_connection_string")) - if not connection_string: - raise ValueError("Destination Enhanced Citations storage connection string is required when migrating source document blobs.") - return BlobServiceClient.from_connection_string(connection_string) - blob_endpoint = _safe_text((settings or {}).get("target_enhanced_citations_storage_blob_endpoint")) - if not blob_endpoint: - raise ValueError("Destination Enhanced Citations Blob endpoint is required when migrating source document blobs.") - return BlobServiceClient(account_url=blob_endpoint, credential=DefaultAzureCredential()) +def _target_cosmos_container_name(container_definition): + return getattr( + app_config, + container_definition["container_name_attr"], + container_definition["name"], + ) -def _document_blob_reference(document_item): - if not isinstance(document_item, dict): - return None - container_name = document_item.get("blob_container") or _blob_container_for_document(document_item) - blob_path = document_item.get("blob_path") or document_item.get("archived_blob_path") - if not blob_path and document_item.get("file_name"): - scope_id = document_item.get("public_workspace_id") or document_item.get("group_id") or document_item.get("user_id") - blob_path = f"{scope_id}/{document_item.get('file_name')}" if scope_id else "" - if not container_name or not blob_path: - return None - return container_name, blob_path +def _get_migration_parallel_operations(settings): + return _safe_int( + (settings or {}).get("migration_max_parallel_operations"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_PARALLEL_OPERATIONS, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_PARALLEL_OPERATIONS, + ) -def _blob_container_for_document(document_item): - if document_item.get("public_workspace_id"): - return app_config.storage_account_public_documents_container_name - if document_item.get("group_id"): - return app_config.storage_account_group_documents_container_name - return app_config.storage_account_user_documents_container_name +def _get_migration_retry_count(settings): + return _safe_int( + (settings or {}).get("migration_retry_count"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_RETRY_COUNT, + minimum=1, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_RETRY_COUNT, + ) -def _iter_selected_document_records_for_blob_migration(migration_plan): - container_map = [ - ("users", app_config.cosmos_user_documents_container, "user_id"), - ("groups", app_config.cosmos_group_documents_container, "group_id"), - ("public_workspaces", app_config.cosmos_public_documents_container, "public_workspace_id"), +def _get_target_cosmos_account_name(endpoint): + normalized_endpoint = _safe_text(endpoint) + if not normalized_endpoint: + return "" + parsed_endpoint = urlparse( + normalized_endpoint if "://" in normalized_endpoint else f"https://{normalized_endpoint}" + ) + hostname = parsed_endpoint.hostname or "" + return hostname.split(".")[0].strip() + + +def _get_target_cosmos_management_settings(settings): + account_name = _get_target_cosmos_account_name((settings or {}).get("target_cosmos_endpoint")) + management_settings = { + "cosmos_throughput_subscription_id": _safe_text( + (settings or {}).get("target_cosmos_subscription_id") + ), + "cosmos_throughput_resource_group": _safe_text( + (settings or {}).get("target_cosmos_resource_group") + ), + "cosmos_throughput_account_name": account_name, + "cosmos_throughput_database_name": DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, + } + missing_values = [ + field_name + for field_name, value in management_settings.items() + if not value ] - for target_type, container, filter_field in container_map: - selection = migration_plan.get(target_type) or {} - if selection.get("mode") == "none" or not selection.get("include_documents"): - continue - if selection.get("mode") == "all": - yield from _iter_cosmos_container_items(container) - continue - for selected_id in selection.get("ids") or []: - query = f"SELECT * FROM c WHERE c.{filter_field} = @selected_id" - for item in container.query_items( - query=query, - parameters=[{"name": "@selected_id", "value": selected_id}], - enable_cross_partition_query=True, - ): - yield _strip_cosmos_system_fields(item) - - -def _copy_source_blobs_to_target(settings, migration_plan): - if not migration_plan.get("include_source_blobs"): - return [] - source_client = _get_source_blob_service_client() - if not source_client: - return [{"name": "source_blobs", "type": "source_blobs", "status": "skipped", "warning": "Source Enhanced Citations storage is not configured."}] - target_client = _get_target_enhanced_citations_blob_client(settings) - copied = {} - warnings = [] - for document in _iter_selected_document_records_for_blob_migration(migration_plan): - reference = _document_blob_reference(document) - if not reference: - continue - container_name, blob_path = reference - artifact = copied.setdefault(container_name, {"name": container_name, "type": "source_blob_container", "container_name": container_name, "blob_count": 0, "bytes": 0}) - try: - target_container = target_client.create_container(container_name) - except Exception: - target_container = target_client.get_container_client(container_name) - try: - source_blob = source_client.get_blob_client(container=container_name, blob=blob_path) - blob_bytes = source_blob.download_blob().readall() - target_container.upload_blob(name=blob_path, data=blob_bytes, overwrite=True, content_settings=ContentSettings(content_type="application/octet-stream")) - artifact["blob_count"] += 1 - artifact["bytes"] += len(blob_bytes) - except Exception as exc: - warnings.append(f"{container_name}/{blob_path}: {exc}") - artifacts = list(copied.values()) - if warnings: - artifacts.append({"name": "source_blob_warnings", "type": "source_blobs", "status": "warning", "warning": "; ".join(warnings[:10])}) - return artifacts + if missing_values: + raise DataManagementSettingsValidationError( + "Destination Cosmos capacity management requires " + f"{', '.join(missing_values)}." + ) + return management_settings -def _migration_selection_summary(migration_plan): +def _build_migration_configuration_snapshot(settings, migration_plan): + """Build a durable state fingerprint without copying credential material.""" return { - target_type: { - "mode": (migration_plan.get(target_type) or {}).get("mode"), - "selected_count": len((migration_plan.get(target_type) or {}).get("ids") or []), - "include_documents": bool((migration_plan.get(target_type) or {}).get("include_documents")), - } - for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES + "migration_plan": copy.deepcopy(migration_plan), + "parallel_operations": _get_migration_parallel_operations(settings), + "retry_count": _get_migration_retry_count(settings), + "skip_recent_within_hours": _safe_int( + (settings or {}).get("migration_skip_recent_within_hours"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_SKIP_WITHIN_HOURS, + minimum=0, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_SKIP_WITHIN_HOURS, + ), + "temporary_destination_ru_enabled": _safe_bool( + (settings or {}).get("migration_temporary_destination_ru_enabled"), + False, + ), + "temporary_destination_ru": _safe_int( + (settings or {}).get("migration_temporary_destination_ru"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, + minimum=1000, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU, + ), + "target_cosmos_endpoint": _safe_text((settings or {}).get("target_cosmos_endpoint")), + "target_cosmos_authentication_type": _safe_text( + (settings or {}).get("target_cosmos_authentication_type") + ), + "target_cosmos_subscription_id": _safe_text( + (settings or {}).get("target_cosmos_subscription_id") + ), + "target_cosmos_resource_group": _safe_text( + (settings or {}).get("target_cosmos_resource_group") + ), + "target_ai_search_endpoint": _safe_text((settings or {}).get("target_ai_search_endpoint")), + "target_ai_search_authentication_type": _safe_text( + (settings or {}).get("target_ai_search_authentication_type") + ), + "target_enhanced_citations_storage_endpoint": _safe_text( + (settings or {}).get("target_enhanced_citations_storage_blob_endpoint") + ), + "target_enhanced_citations_storage_authentication_type": _safe_text( + (settings or {}).get("target_enhanced_citations_storage_authentication_type") + ), } -def _job_lease_holder_id(): - return f"{socket.gethostname()}:{uuid.uuid4().hex}" - - -def _is_stale_job(job, stale_seconds=DATA_MANAGEMENT_DEFAULT_STALE_SECONDS): - last_heartbeat = _parse_iso_datetime((job or {}).get("last_heartbeat_at") or (job or {}).get("updated_at")) - if not last_heartbeat: - return True - return last_heartbeat <= _now_utc() - timedelta(seconds=stale_seconds) +def _migration_plan_scope_signature(migration_plan): + """Return the source scope fields that must match an incremental baseline.""" + signature = { + "include_ai_search": bool((migration_plan or {}).get("include_ai_search")), + "include_source_blobs": bool((migration_plan or {}).get("include_source_blobs")), + } + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = (migration_plan or {}).get(target_type) or {} + signature[target_type] = { + "mode": _safe_text(selection.get("mode"), "none"), + "ids": sorted(_dedupe_limited_strings(selection.get("ids"), limit=2000)), + "include_documents": bool(selection.get("include_documents")), + } + return signature -def _read_job(job_id): - return cosmos_data_management_jobs_container.read_item(item=job_id, partition_key=job_id) +def _migration_destination_signature(configuration, migration_plan): + """Return normalized destination identities without credential material.""" + return { + "cosmos": _safe_text((configuration or {}).get("target_cosmos_endpoint")).lower(), + "ai_search": ( + _safe_text((configuration or {}).get("target_ai_search_endpoint")).lower() + if (migration_plan or {}).get("include_ai_search") else "" + ), + "source_blobs": ( + _safe_text( + (configuration or {}).get("target_enhanced_citations_storage_endpoint") + ).lower() + if (migration_plan or {}).get("include_source_blobs") else "" + ), + } -def _replace_job(job): - return cosmos_data_management_jobs_container.replace_item( - item=job.get("id"), - body=job, - etag=job.get("_etag"), - match_condition=MatchConditions.IfNotModified, +def _validate_incremental_baseline_job(candidate_job, current_job, current_configuration, migration_plan): + """Return a compatible completed cutoff or reject an unsafe baseline.""" + candidate_state = candidate_job.get("migration_state") if isinstance(candidate_job, dict) else None + candidate_configuration = ( + candidate_state.get("configuration") + if isinstance(candidate_state, dict) and isinstance(candidate_state.get("configuration"), dict) + else {} + ) + candidate_plan = candidate_configuration.get("migration_plan") + candidate_plan = candidate_plan if isinstance(candidate_plan, dict) else {} + candidate_result = candidate_job.get("result") if isinstance(candidate_job, dict) else {} + reconciliation_resource = ( + (candidate_state.get("resources") or {}).get("reconciliation:cutover") + if isinstance(candidate_state, dict) else None + ) + reconciliation_result = ( + reconciliation_resource.get("result") + if isinstance(reconciliation_resource, dict) and + isinstance(reconciliation_resource.get("result"), dict) + else {} + ) + if ( + not isinstance(candidate_job, dict) or + candidate_job.get("id") == (current_job or {}).get("id") or + candidate_job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION or + candidate_job.get("status") not in { + DATA_MANAGEMENT_STATUS_COMPLETED, + DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS, + } or + not isinstance(candidate_state, dict) or + candidate_state.get("status") != MIGRATION_RESOURCE_STATUS_COMPLETED or + (isinstance(candidate_result, dict) and candidate_result.get("dry_run")) or + reconciliation_result.get("readiness") not in {"ready", "ready_with_warnings"} + ): + raise DataManagementSettingsValidationError( + "Incremental migration baseline must be a completed, reconciled, non-dry-run migration." + ) + if _migration_plan_scope_signature(candidate_plan) != _migration_plan_scope_signature(migration_plan): + raise DataManagementSettingsValidationError( + "Incremental migration baseline uses a different source scope." + ) + if ( + _migration_destination_signature(candidate_configuration, candidate_plan) != + _migration_destination_signature(current_configuration, migration_plan) + ): + raise DataManagementSettingsValidationError( + "Incremental migration baseline targets a different destination." + ) + source_cutoff_at = _safe_text(candidate_state.get("source_cutoff_at")) + if _parse_iso_datetime(source_cutoff_at) is None: + raise DataManagementSettingsValidationError( + "Incremental migration baseline does not contain a valid source watermark." + ) + return source_cutoff_at + + +def _get_completed_migration_baseline_candidates(limit=100): + query = ( + "SELECT * FROM c WHERE c.type = @type AND c.operation = @operation " + "AND (c.status = @completed OR c.status = @completed_with_warnings) " + "ORDER BY c.completed_at DESC" + ) + parameters = [ + {"name": "@type", "value": DATA_MANAGEMENT_JOB_TYPE}, + {"name": "@operation", "value": DATA_MANAGEMENT_OPERATION_MIGRATION}, + {"name": "@completed", "value": DATA_MANAGEMENT_STATUS_COMPLETED}, + { + "name": "@completed_with_warnings", + "value": DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS, + }, + ] + return list(cosmos_data_management_jobs_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + max_item_count=limit, + ))[:limit] + + +def _resolve_data_management_migration_baseline(job, settings, migration_plan): + """Resolve one immutable prior cutoff for delta and mirror migration modes.""" + resolved_plan = copy.deepcopy(migration_plan) + migration_mode = resolved_plan.get("migration_mode", DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY) + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resolved_plan["baseline_job_id"] = "" + resolved_plan["baseline_source_cutoff_at"] = "" + return resolved_plan + + existing_state = job.get("migration_state") if isinstance(job, dict) else None + existing_configuration = ( + existing_state.get("configuration") + if isinstance(existing_state, dict) and isinstance(existing_state.get("configuration"), dict) + else {} + ) + existing_plan = existing_configuration.get("migration_plan") + if isinstance(existing_plan, dict) and existing_plan.get("baseline_source_cutoff_at"): + resolved_plan["baseline_job_id"] = _safe_text(existing_plan.get("baseline_job_id")) + resolved_plan["baseline_source_cutoff_at"] = _safe_text( + existing_plan.get("baseline_source_cutoff_at") + ) + return resolved_plan + + current_configuration = _build_migration_configuration_snapshot(settings, resolved_plan) + requested_baseline_job_id = _safe_text(resolved_plan.get("baseline_job_id")) + if requested_baseline_job_id: + candidate_jobs = [_read_job(requested_baseline_job_id)] + else: + candidate_jobs = _get_completed_migration_baseline_candidates() + + incompatibility_messages = [] + for candidate_job in candidate_jobs: + try: + source_cutoff_at = _validate_incremental_baseline_job( + candidate_job, + job, + current_configuration, + resolved_plan, + ) + except DataManagementSettingsValidationError as exc: + incompatibility_messages.append(str(exc)) + if requested_baseline_job_id: + raise + continue + resolved_plan["baseline_job_id"] = _safe_text(candidate_job.get("id")) + resolved_plan["baseline_source_cutoff_at"] = source_cutoff_at + return resolved_plan + + detail = incompatibility_messages[0] if incompatibility_messages else "No completed migration was found." + raise DataManagementSettingsValidationError( + f"Incremental migration requires a compatible completed baseline. {detail}" + ) + + +def _initialize_data_management_migration_state(job, settings, migration_plan): + """Persist the job's stable migration identity before target writes begin.""" + configuration = _build_migration_configuration_snapshot(settings, migration_plan) + state = initialize_migration_state( + job.get("migration_state"), + str(job.get("id")), + configuration, + ) + job["migration_state"] = state + context = create_migration_provenance_context( + migration_id=state["migration_id"], + migrated_at_utc=state["source_cutoff_at"], + skip_within_hours=configuration["skip_recent_within_hours"], + ) + migration_mode = migration_plan.get( + "migration_mode", + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + baseline_source_cutoff_at = _safe_text(migration_plan.get("baseline_source_cutoff_at")) + state["incremental"] = { + "mode": migration_mode, + "baseline_job_id": _safe_text(migration_plan.get("baseline_job_id")), + "baseline_source_cutoff_at": baseline_source_cutoff_at, + "source_cutoff_at": state["source_cutoff_at"], + "deletion_policy": ( + "confirmed_migration_owned_only" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR else + "disabled" + ), + } + state["watermarks"] = { + "baseline": baseline_source_cutoff_at, + "current_cutoff": state["source_cutoff_at"], + "cosmos": {"strategy": "_ts_content_hash_with_identity", "through": state["source_cutoff_at"]}, + "ai_search": {"strategy": "id_keyset_with_observed_source_hash", "through": None}, + "source_blobs": {"strategy": "stable_etag_with_path", "through": None}, + } + context.update({ + "migration_mode": migration_mode, + "baseline_job_id": state["incremental"]["baseline_job_id"], + "baseline_source_cutoff_at": baseline_source_cutoff_at, + }) + return state, context + + +def _set_document_path_value(document, partition_key_path, value): + """Set a simple or nested Cosmos partition-key value in a probe document.""" + path_parts = [part for part in _safe_text(partition_key_path).strip("/").split("/") if part] + if not path_parts: + raise DataManagementSettingsValidationError("Destination Cosmos container has an invalid partition key path.") + current = document + for path_part in path_parts[:-1]: + nested = current.get(path_part) + if not isinstance(nested, dict): + nested = {} + current[path_part] = nested + current = nested + current[path_parts[-1]] = value + return value + + +def _verify_target_cosmos_container_access(target_container, partition_key_path): + """Prove the destination has create, read, and delete data-plane access.""" + probe_id = f"simplechat-migration-preflight-{uuid.uuid4().hex}" + probe_document = { + "id": probe_id, + "type": "simplechat_migration_preflight", + "created_at": _now_iso(), + } + partition_key_value = _set_document_path_value(probe_document, partition_key_path, probe_id) + target_container.create_item(body=probe_document) + try: + target_container.read_item(item=probe_id, partition_key=partition_key_value) + finally: + target_container.delete_item(item=probe_id, partition_key=partition_key_value) + + +def _preflight_target_cosmos_migration_access(settings, migration_plan): + """Create planned containers and verify their destination data-plane permissions.""" + target_database = _get_target_cosmos_database(settings) + container_results = [] + for target_type, container_definition in _iter_migration_cosmos_container_definitions(migration_plan): + source_container = getattr(app_config, container_definition["container_attr"], None) + if source_container is None: + raise DataManagementSettingsValidationError( + f"Source Cosmos container '{container_definition['name']}' is not initialized." + ) + source_probe = source_container.query_items( + query="SELECT TOP 1 * FROM c", + enable_cross_partition_query=True, + ) + next(iter(source_probe), None) + target_container_name = _target_cosmos_container_name(container_definition) + target_container = _get_target_cosmos_container( + target_database, + target_container_name, + container_definition["partition_key_path"], + ) + _verify_target_cosmos_container_access( + target_container, + container_definition["partition_key_path"], + ) + container_results.append({ + "target_type": target_type, + "container_name": target_container_name, + "partition_key_path": container_definition["partition_key_path"], + "source_read_verified": True, + "destination_write_verified": True, + }) + return { + "database_name": DATA_MANAGEMENT_TARGET_COSMOS_DATABASE_NAME, + "container_count": len(container_results), + "containers": container_results, + } + + +def _inspect_target_cosmos_migration_capacity(settings, migration_plan): + """Read destination capacity targets before an optional temporary RU boost.""" + management_settings = _get_target_cosmos_management_settings(settings) + database_throughput = get_database_throughput(management_settings) + targets = [] + if database_throughput.get("is_scalable"): + targets.append({ + "scope": "database", + "container_name": "", + "mode": database_throughput.get("mode"), + "current_ru": database_throughput.get("current_ru"), + }) + else: + seen_container_names = set() + for _, container_definition in _iter_migration_cosmos_container_definitions(migration_plan): + container_name = _target_cosmos_container_name(container_definition) + if container_name in seen_container_names: + continue + seen_container_names.add(container_name) + container_throughput = get_container_throughput(management_settings, container_name) + if container_throughput.get("is_scalable"): + targets.append({ + "scope": "container", + "container_name": container_name, + "mode": container_throughput.get("mode"), + "current_ru": container_throughput.get("current_ru"), + }) + if not targets: + raise CosmosThroughputError( + "Destination Cosmos throughput is not scalable at the database or selected container level." + ) + return { + "target_ru": _safe_int( + (settings or {}).get("migration_temporary_destination_ru"), + default=DATA_MANAGEMENT_MIGRATION_DEFAULT_DESTINATION_RU, + minimum=1000, + maximum=DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU, + ), + "management_settings": management_settings, + "database_mode": database_throughput.get("mode"), + "database_current_ru": database_throughput.get("current_ru"), + "targets": targets, + } + + +def _get_data_management_job_lease_seconds(settings): + return _safe_int( + (settings or {}).get("data_management_job_lease_seconds"), + default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, + minimum=60, + maximum=7200, + ) + + +def _migration_resource_metrics(resource): + if not isinstance(resource, dict): + return {} + result = resource.get("result") + if isinstance(result, dict) and result: + return result + progress = resource.get("progress") + return progress if isinstance(progress, dict) else {} + + +def _update_migration_state_totals(state): + totals = { + "copied_count": 0, + "skipped_count": 0, + "failed_count": 0, + "collision_count": 0, + "processed_count": 0, + "bytes": 0, + "request_units": 0.0, + } + for resource in (state.get("resources") or {}).values(): + metrics = _migration_resource_metrics(resource) + for field_name in ( + "copied_count", + "skipped_count", + "failed_count", + "collision_count", + "processed_count", + "bytes", + ): + totals[field_name] += _safe_int(metrics.get(field_name), default=0, minimum=0) + try: + totals["request_units"] += max(0.0, float(metrics.get("request_units") or 0.0)) + except (TypeError, ValueError): + continue + totals["request_units"] = round(totals["request_units"], 3) + state["totals"] = totals + return totals + + +def _migration_progress_marker(progress): + marker_fields = ( + "copied_count", + "created_count", + "updated_count", + "unchanged_count", + "skipped_count", + "failed_count", + "collision_count", + "missing_count", + "not_applicable_count", + "processed_count", + "bytes", + "observed_bytes", + "request_units", + "source_read_count", + "destination_accepted_count", + "destination_failed_count", + "completed_count", + "services_completed", + ) + return tuple((field_name, (progress or {}).get(field_name)) for field_name in marker_fields) + + +def _persist_migration_checkpoint( + job, + state, + settings, + resource_name, + progress, + message, + allow_cancel_requested=False, +): + """Save a batch checkpoint and renew the worker lease without timeline noise.""" + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + previous_resource = (state.get("resources") or {}).get(resource_name) or {} + previous_progress = ( + previous_resource.get("progress") + if isinstance(previous_resource.get("progress"), dict) + else {} + ) + now = _now_utc() + if _migration_progress_marker(progress) != _migration_progress_marker(previous_progress): + progress["last_progress_at"] = now.isoformat() + state["last_progress_at"] = now.isoformat() + job["last_progress_at"] = now.isoformat() + else: + progress["last_progress_at"] = ( + previous_progress.get("last_progress_at") or + state.get("last_progress_at") or + job.get("last_progress_at") + ) + update_migration_resource(state, resource_name, progress) + _update_migration_state_totals(state) + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _persist_migration_heartbeat(job, settings, message): + """Renew worker and coordinator leases without advancing resource progress.""" + _assert_migration_job_lease(job) + now = _now_utc() + job.update({ + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + return _save_data_management_job(job) + + +def _complete_migration_resource_checkpoint(job, state, settings, resource_name, result, message): + """Complete a resource after its final durable checkpoint is written.""" + _assert_migration_job_lease(job) + complete_migration_resource(state, resource_name, result=result) + _update_migration_state_totals(state) + now = _now_utc() + result["last_progress_at"] = now.isoformat() + state["last_progress_at"] = now.isoformat() + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "last_progress_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _fail_migration_resource_checkpoint(job, state, settings, resource_name, error_message): + """Persist a failed resource before allowing the job to be retried.""" + _assert_migration_job_lease(job) + fail_migration_resource(state, resource_name, error_message) + _update_migration_state_totals(state) + now = _now_utc() + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": f"Migration resource failed: {resource_name}", + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _persist_migration_state( + job, + state, + settings, + message, + allow_cancel_requested=False, +): + """Persist non-resource migration state changes and renew the worker lease.""" + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + now = _now_utc() + job.update({ + "migration_state": state, + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_expires_at": ( + now + timedelta(seconds=_get_data_management_job_lease_seconds(settings)) + ).isoformat(), + "last_message": message, + }) + _save_data_management_job(job) + return job.get("migration_state") if isinstance(job.get("migration_state"), dict) else state + + +def _preflight_target_ai_search_migration_access(settings, migration_plan): + """Ensure target indexes are writable and have provenance fields before copying.""" + _validate_target_ai_search_migration_write_safety(settings, migration_plan) + target_database = None + target_search_write_gate_verified = False + results = [] + search_mappings = [ + ("users", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), + ("groups", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), + ("public_workspaces", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), + ] + for target_type, artifact in search_mappings: + selection = (migration_plan or {}).get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + if target_database is None: + target_database = _get_target_cosmos_database(settings) + target_search_write_gate_container = _get_target_data_management_search_write_gate_container( + target_database + ) + _verify_target_cosmos_container_access( + target_search_write_gate_container, + "/id", + ) + target_search_write_gate_verified = True + source_client = CLIENTS.get(artifact["client_key"]) + if not source_client: + raise DataManagementSettingsValidationError( + f"Source AI Search client '{artifact['name']}' is not initialized." + ) + next(iter(source_client.search( + search_text="*", + top=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )), None) + index_status = _ensure_target_search_index( + settings, + artifact["index_name"], + artifact["schema_file"], + ) + target_client = _get_target_search_client(settings, artifact["index_name"]) + probe_id = f"simplechat-migration-preflight-{uuid.uuid4().hex}" + probe_document = { + "id": probe_id, + SEARCH_MIGRATION_ID_FIELD: probe_id, + SEARCH_MIGRATED_AT_FIELD: _now_iso(), + SEARCH_MIGRATION_STATUS_FIELD: "preflight", + } + probe_written = False + probe_succeeded = False + try: + upload_results = list(target_client.upload_documents( + documents=[probe_document], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if not upload_results or not _get_search_result_succeeded(upload_results[0]): + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{artifact['index_name']}' rejected a write probe." + ) + probe_written = True + read_results = list(target_client.search( + search_text="*", + filter=f"id eq '{_escape_search_filter_value(probe_id)}'", + select=["id"], + top=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if not any(_safe_text(dict(item).get("id")) == probe_id for item in read_results): + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{artifact['index_name']}' did not return its write probe." + ) + probe_succeeded = True + finally: + if probe_written: + try: + delete_results = list(target_client.delete_documents( + documents=[{"id": probe_id}], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if delete_results and not _get_search_result_succeeded(delete_results[0]): + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{artifact['index_name']}' rejected write-probe cleanup." + ) + except Exception: + if probe_succeeded: + raise + results.append({ + "target_type": target_type, + "index_name": artifact["index_name"], + "source_read_verified": True, + "destination_write_verified": True, + "index_status": index_status, + }) + return { + "index_count": len(results), + "indexes": results, + "target_search_write_gate_verified": target_search_write_gate_verified, + } + + +def _preflight_target_blob_migration_access(settings, migration_plan): + """Verify source and target Blob clients before the long-running copy phase.""" + if not (migration_plan or {}).get("include_source_blobs"): + return {"enabled": False, "container_count": 0, "containers": []} + source_client = _get_source_blob_service_client() + if not source_client: + raise DataManagementSettingsValidationError( + "Source Enhanced Citations storage is not configured for blob migration." + ) + target_client = _get_target_enhanced_citations_blob_client(settings) + container_names = [] + for target_type, selection in (migration_plan or {}).items(): + if target_type not in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: + continue + if not isinstance(selection, dict) or selection.get("mode") == "none" or not selection.get("include_documents"): + continue + if target_type == "users": + container_names.append(app_config.storage_account_user_documents_container_name) + elif target_type == "groups": + container_names.append(app_config.storage_account_group_documents_container_name) + elif target_type == "public_workspaces": + container_names.append(app_config.storage_account_public_documents_container_name) + results = [] + for container_name in sorted(set(container_names)): + source_container = source_client.get_container_client(container_name) + next(iter(source_container.list_blobs()), None) + try: + target_client.create_container(container_name) + except ResourceExistsError: + target_client.get_container_client(container_name).get_container_properties() + probe_blob_name = f"simplechat-migration-preflight/{uuid.uuid4().hex}" + target_blob = target_client.get_blob_client( + container=container_name, + blob=probe_blob_name, + ) + probe_written = False + probe_succeeded = False + try: + target_blob.upload_blob( + data=b"", + overwrite=False, + metadata={"simplechatMigrationPreflight": "true"}, + ) + probe_written = True + target_blob.get_blob_properties() + probe_succeeded = True + finally: + if probe_written: + try: + target_blob.delete_blob() + except Exception: + if probe_succeeded: + raise + results.append({ + "container_name": container_name, + "source_read_verified": True, + "destination_write_verified": True, + }) + return {"enabled": True, "container_count": len(results), "containers": results} + + +def _run_data_management_migration_preflight(job, migration_state, settings, migration_plan): + """Persist preflight evidence before any selected user data is written.""" + migration_state["preflight"] = { + "status": "in_progress", + "started_at": _now_iso(), + } + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Validating source and destination migration access", + ) + _assert_migration_job_lease(job) + cosmos_result = _preflight_target_cosmos_migration_access(settings, migration_plan) + _assert_migration_job_lease(job) + search_result = _preflight_target_ai_search_migration_access(settings, migration_plan) + _assert_migration_job_lease(job) + blob_result = _preflight_target_blob_migration_access(settings, migration_plan) + capacity_result = None + if _safe_bool((settings or {}).get("migration_temporary_destination_ru_enabled"), False): + capacity_result = _inspect_target_cosmos_migration_capacity(settings, migration_plan) + migration_state["preflight"] = { + "status": "completed", + "completed_at": _now_iso(), + "cosmos": cosmos_result, + "ai_search": search_result, + "source_blobs": blob_result, + "destination_capacity": capacity_result, + } + return _persist_migration_state( + job, + migration_state, + settings, + "Migration source and destination preflight completed", + ) + + +def _get_destination_capacity_current(management_settings, capacity_target): + if capacity_target.get("scope") == "container": + return get_container_throughput( + management_settings, + capacity_target.get("container_name"), + ) + return get_database_throughput(management_settings) + + +def _apply_temporary_destination_capacity(job, migration_state, settings, migration_plan): + """Temporarily increase only migration-targeted destination throughput to <=10,000 RU/s.""" + _assert_migration_job_lease(job) + if not _safe_bool((settings or {}).get("migration_temporary_destination_ru_enabled"), False): + migration_state["capacity"] = {"status": "not_requested", "restore_pending": False} + return _persist_migration_state( + job, + migration_state, + settings, + "Temporary destination Cosmos capacity boost is disabled", + ) + + capacity = migration_state.get("capacity") if isinstance(migration_state.get("capacity"), dict) else {} + if capacity.get("restore_pending") and capacity.get("targets"): + if capacity.get("status") in {"applying", "restore_pending"}: + _restore_temporary_destination_capacity(job, migration_state, settings) + capacity = migration_state.get("capacity") if isinstance(migration_state.get("capacity"), dict) else {} + if capacity.get("restore_pending"): + raise DataManagementSettingsValidationError( + "Destination Cosmos capacity restoration is still pending. Resolve the recorded restore error before retrying migration." + ) + return _persist_migration_state( + job, + migration_state, + settings, + "Recovered pending destination Cosmos capacity without reapplying a boost", + ) + else: + return migration_state + + inspection = _inspect_target_cosmos_migration_capacity(settings, migration_plan) + target_ru = inspection["target_ru"] + capacity = { + "status": "applying", + "restore_pending": True, + "target_ru": target_ru, + "management_settings": inspection["management_settings"], + "targets": [], + "started_at": _now_iso(), + } + migration_state["capacity"] = capacity + migration_state = _persist_migration_state( + job, + migration_state, + settings, + f"Preparing temporary destination Cosmos capacity up to {target_ru} RU/s", + ) + capacity = migration_state.setdefault("capacity", capacity) + + for target in inspection["targets"]: + _assert_migration_job_lease(job) + original_ru = _safe_int(target.get("current_ru"), default=0, minimum=0) + target_snapshot = { + "scope": target.get("scope"), + "container_name": target.get("container_name") or "", + "mode": target.get("mode"), + "original_ru": original_ru, + "target_ru": target_ru, + "changed": False, + "boost_attempted": False, + "restore_status": "not_required", + } + capacity["targets"].append(target_snapshot) + if original_ru and original_ru < target_ru: + target_snapshot["boost_attempted"] = True + target_snapshot["restore_status"] = "pending" + migration_state = _persist_migration_state( + job, + migration_state, + settings, + f"Recorded destination Cosmos capacity recovery snapshot for {target_snapshot['scope']}", + ) + capacity = migration_state.setdefault("capacity", capacity) + target_snapshot = capacity["targets"][-1] + _assert_migration_job_lease(job) + try: + scale_result = set_database_throughput( + inspection["management_settings"], + target_ru, + initiated_by=f"data_management_migration:{job.get('id')}", + reason="temporary_migration_capacity_boost", + decision={ + "scope": target_snapshot["scope"], + "container_name": target_snapshot["container_name"], + "target_mode": target_snapshot["mode"], + }, + ) + except Exception as exc: + target_snapshot["boost_error"] = str(exc)[:300] + migration_state["capacity"] = capacity + _persist_migration_state( + job, + migration_state, + settings, + f"Destination Cosmos capacity boost failed for {target_snapshot['scope']}", + ) + raise + target_snapshot["changed"] = True + target_snapshot["boosted_to_ru"] = scale_result.get("to_ru", target_ru) + target_snapshot["restore_status"] = "pending" + else: + target_snapshot["restore_status"] = "not_required" + migration_state = _persist_migration_state( + job, + migration_state, + settings, + f"Updated temporary destination Cosmos capacity for {target_snapshot['scope']}", + ) + capacity = migration_state.setdefault("capacity", capacity) + + capacity["status"] = "boosted" + capacity["completed_at"] = _now_iso() + return _persist_migration_state( + job, + migration_state, + settings, + f"Temporary destination Cosmos capacity is ready up to {target_ru} RU/s", + ) + + +def _restore_temporary_destination_capacity( + job, + migration_state, + settings, + allow_cancel_requested=False, +): + """Restore only the capacity values this migration changed, without clobbering outside updates.""" + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + capacity = migration_state.get("capacity") if isinstance(migration_state.get("capacity"), dict) else {} + if not capacity.get("restore_pending"): + return [], migration_state + management_settings = capacity.get("management_settings") or {} + restore_warnings = [] + for target in reversed(capacity.get("targets") or []): + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) + if not (target.get("changed") or target.get("boost_attempted")): + continue + try: + current = _get_destination_capacity_current(management_settings, target) + current_ru = _safe_int(current.get("current_ru"), default=0, minimum=0) + boosted_to_ru = _safe_int(target.get("boosted_to_ru"), default=target.get("target_ru"), minimum=0) + original_ru = _safe_int(target.get("original_ru"), default=0, minimum=0) + if current_ru == original_ru: + target["restore_status"] = "not_changed" + continue + if current_ru != boosted_to_ru: + target["restore_status"] = "skipped_external_change" + restore_warnings.append( + f"Did not restore destination {target.get('scope')} capacity because it changed after the migration boost." + ) + continue + scale_result = set_database_throughput( + management_settings, + original_ru, + initiated_by=f"data_management_migration:{job.get('id')}", + reason="restore_temporary_migration_capacity", + decision={ + "scope": target.get("scope"), + "container_name": target.get("container_name") or "", + "target_mode": target.get("mode"), + }, + ) + target["restore_status"] = "restored" + target["restored_to_ru"] = scale_result.get("to_ru", original_ru) + except Exception as exc: + target["restore_status"] = "restore_failed" + restore_warnings.append( + f"Failed to restore destination {target.get('scope')} capacity: {str(exc)[:300]}" + ) + + unresolved_restore_statuses = {"pending", "restore_failed"} + capacity["restore_pending"] = any( + target.get("restore_status") in unresolved_restore_statuses + for target in capacity.get("targets") or [] + ) + capacity["restored_at"] = _now_iso() + if capacity["restore_pending"]: + capacity["status"] = "restore_pending" + else: + capacity["status"] = "restored_with_warnings" if restore_warnings else "restored" + capacity["restore_warnings"] = restore_warnings + migration_state["capacity"] = capacity + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Restored temporary destination Cosmos capacity", + allow_cancel_requested=allow_cancel_requested, + ) + return restore_warnings, migration_state + + +def _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=None, + source_start_epoch=None, + include_source_version=False, +): + source_container = getattr(app_config, container_definition["container_attr"], None) + if not source_container: + return + if container_definition.get("documents") and not selection.get("include_documents"): + return + mode = selection.get("mode") + + def prepare_item(item): + cleaned_item = _strip_cosmos_system_fields(item) + if include_source_version: + return cleaned_item, _safe_text(item.get("_ts")) + return cleaned_item + + if mode == "all": + query = "SELECT * FROM c" + parameters = [] + conditions = [] + if source_start_epoch is not None: + conditions.append("c._ts >= @source_start_epoch") + parameters.append({"name": "@source_start_epoch", "value": source_start_epoch}) + if source_cutoff_epoch is not None: + conditions.append("c._ts <= @source_cutoff_epoch") + parameters = [{"name": "@source_cutoff_epoch", "value": source_cutoff_epoch}] + if source_start_epoch is not None: + parameters.insert(0, {"name": "@source_start_epoch", "value": source_start_epoch}) + if conditions: + query = f"SELECT * FROM c WHERE {' AND '.join(conditions)}" + for item in source_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + ): + yield prepare_item(item) + return + if mode != "selected": + return + + ids = selection.get("ids") or [] + if container_definition.get("id_field") == "id": + for item_id in ids: + try: + item = source_container.read_item(item=item_id, partition_key=item_id) + except CosmosResourceNotFoundError: + continue + if ( + source_cutoff_epoch is not None and + item.get("_ts") is not None and + _safe_int(item.get("_ts"), default=0, minimum=0) > source_cutoff_epoch + ): + continue + if ( + source_start_epoch is not None and + item.get("_ts") is not None and + _safe_int(item.get("_ts"), default=0, minimum=0) < source_start_epoch + ): + continue + yield prepare_item(item) + return + + filter_field = container_definition.get("filter_field") + if not filter_field: + return + for selected_id in ids: + query = f"SELECT * FROM c WHERE c.{filter_field} = @selected_id" + parameters = [{"name": "@selected_id", "value": selected_id}] + if source_start_epoch is not None: + query += " AND c._ts >= @source_start_epoch" + parameters.append({"name": "@source_start_epoch", "value": source_start_epoch}) + if source_cutoff_epoch is not None: + query += " AND c._ts <= @source_cutoff_epoch" + parameters.append({"name": "@source_cutoff_epoch", "value": source_cutoff_epoch}) + for item in source_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + ): + yield prepare_item(item) + + +def _build_cosmos_document_identity(document_id, partition_key_value): + if document_id is None or partition_key_value is None: + return "" + return json.dumps( + {"id": str(document_id), "partition_key": partition_key_value}, + ensure_ascii=True, + sort_keys=True, + default=str, + separators=(",", ":"), + ) + + +def _get_cosmos_document_identity(document, partition_key_path): + if not isinstance(document, dict): + return "" + return _build_cosmos_document_identity( + document.get("id"), + _get_document_path_value(document, partition_key_path), + ) + + +def _build_cosmos_source_hash(document): + """Hash source content without service fields or migration provenance.""" + canonical_document = { + key: value + for key, value in (document or {}).items() + if key not in { + COSMOS_MIGRATION_PROVENANCE_FIELD, + "_etag", + "_rid", + "_self", + "_attachments", + "_ts", + } + } + encoded = json.dumps( + canonical_document, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=_json_default, + ).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _get_cosmos_partition_key_select_expression(partition_key_path): + path_parts = [part for part in _safe_text(partition_key_path).strip("/").split("/") if part] + if not path_parts or any(not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", part) for part in path_parts): + raise DataManagementSettingsValidationError("Destination Cosmos container has an unsupported partition key path.") + return "c." + ".".join(path_parts) + + +def _get_target_cosmos_document(target_container, document, partition_key_path): + """Read one exact destination item without allocating an account-wide key set.""" + document_id = _safe_text((document or {}).get("id")) + partition_key_value = _get_document_path_value(document, partition_key_path) + if not document_id or partition_key_value is None: + raise DataManagementSettingsValidationError( + "Migration source record is missing a required ID or partition key value." + ) + try: + return target_container.read_item( + item=document_id, + partition_key=partition_key_value, + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return None + + +def _classify_target_cosmos_document( + target_document, + provenance_context, + container_name, + migration_mode=DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + source_hash="", + source_version="", +): + """Allow only empty or migration-owned target identities to be written.""" + if target_document is None: + return "create" + target_provenance = get_cosmos_migration_provenance(target_document) + if is_successful_migration_record(target_provenance): + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) + ): + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "resume_verified" + return "update" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + return "unchanged" + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "unchanged" + return "update" + raise DataManagementSettingsValidationError( + f"Destination Cosmos container '{container_name}' contains an unowned record that conflicts with this migration." + ) + + +def _get_cosmos_response_request_units(headers): + for header_name, header_value in (headers or {}).items(): + if str(header_name).lower() != "x-ms-request-charge": + continue + try: + return max(0.0, float(header_value)) + except (TypeError, ValueError): + return 0.0 + return 0.0 + + +def _drain_migration_futures_with_heartbeat( + pending_futures, + consume_result, + persist_heartbeat, + cancel_event, + persist_stopped_results=None, +): + pending_by_future = ( + dict(pending_futures) + if isinstance(pending_futures, (dict, list, tuple)) + else {future: None for future in pending_futures} + ) + last_heartbeat = time.monotonic() + stop_error = None + while pending_by_future: + completed_futures, _pending_futures = wait( + set(pending_by_future), + timeout=DATA_MANAGEMENT_MIGRATION_HEARTBEAT_POLL_SECONDS, + return_when=FIRST_COMPLETED, + ) + for future in completed_futures: + context = pending_by_future.pop(future) + if future.cancelled(): + continue + try: + consume_result(future.result(), context) + except Exception as exc: + if stop_error is None: + stop_error = exc + cancel_event.set() + for pending_future in pending_by_future: + pending_future.cancel() + now = time.monotonic() + if ( + pending_by_future and + stop_error is None and + now - last_heartbeat >= DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + try: + persist_heartbeat() + except Exception as exc: + stop_error = exc + cancel_event.set() + for future in pending_by_future: + future.cancel() + last_heartbeat = now + if stop_error is not None: + if callable(persist_stopped_results): + try: + persist_stopped_results() + except Exception: + pass + raise stop_error + + +def _write_cosmos_migration_record( + target_container, + document, + partition_key_path, + provenance_context, + retry_count, + disposition="create", + target_document=None, + source_hash="", + source_version="", + cancel_event=None, +): + """Write one provenance-tagged Cosmos record with bounded transient retries.""" + writable_document = copy.deepcopy(document) + add_cosmos_migration_provenance( + writable_document, + provenance_context, + source_hash=source_hash, + source_version=source_version, + ) + payload_bytes = len(json.dumps(writable_document, default=_json_default).encode("utf-8")) + request_units = 0.0 + request_unit_lock = Lock() + + def response_hook(headers, _response): + nonlocal request_units + with request_unit_lock: + request_units += _get_cosmos_response_request_units(headers) + + started_at = time.perf_counter() + last_error = None + for attempt in range(1, retry_count + 1): + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Cosmos migration stopped after cancellation or lease loss." + ) + try: + if disposition == "update": + target_container.replace_item( + item=target_document, + body=writable_document, + etag=(target_document or {}).get("_etag"), + match_condition=MatchConditions.IfNotModified, + response_hook=response_hook, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + else: + target_container.create_item( + writable_document, + response_hook=response_hook, + retry_write=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + ) + return { + "copied": True, + "created": disposition == "create", + "updated": disposition == "update", + "bytes": payload_bytes, + "request_units": request_units, + "attempt": attempt, + "elapsed_seconds": time.perf_counter() - started_at, + } + except Exception as exc: + last_error = exc + status_code = getattr(exc, "status_code", None) + if status_code in {409, 412}: + target_document = _get_target_cosmos_document( + target_container, + document, + partition_key_path, + ) + try: + disposition = _classify_target_cosmos_document( + target_document, + provenance_context, + "migration target", + migration_mode=(provenance_context or {}).get( + "migration_mode", + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ), + source_hash=source_hash, + source_version=source_version, + ) + except DataManagementSettingsValidationError as collision_error: + return { + "copied": False, + "skipped": False, + "collision": True, + "bytes": 0, + "request_units": request_units, + "attempt": attempt, + "elapsed_seconds": time.perf_counter() - started_at, + "error": str(collision_error), + } + if disposition in {"unchanged", "resume_verified"}: + return { + "copied": False, + "skipped": True, + "resume_verified": disposition == "resume_verified", + "bytes": 0, + "request_units": request_units, + "attempt": attempt, + "elapsed_seconds": time.perf_counter() - started_at, + } + if disposition == "update": + continue + retryable = status_code in {408, 429, 449, 500, 503} or status_code is None + if not retryable or attempt >= retry_count: + break + retry_delay = min(30, 2 ** (attempt - 1)) + if cancel_event is not None: + if cancel_event.wait(retry_delay): + raise DataManagementMigrationCanceledError( + "Cosmos migration stopped during retry backoff." + ) + else: + time.sleep(retry_delay) + + return { + "copied": False, + "skipped": False, + "collision": False, + "bytes": 0, + "request_units": request_units, + "attempt": retry_count, + "elapsed_seconds": time.perf_counter() - started_at, + "error": str(last_error or "Cosmos document write failed.")[:500], + } + + +def _copy_cosmos_records_to_target( + target_database, + target_type, + selection, + job, + migration_state, + provenance_context, + settings, +): + copied = [] + source_cutoff = _parse_iso_datetime(migration_state.get("source_cutoff_at")) + source_cutoff_epoch = int(source_cutoff.timestamp()) if source_cutoff else None + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + baseline_cutoff = _parse_iso_datetime( + (provenance_context or {}).get("baseline_source_cutoff_at") + ) + source_start_epoch = ( + int(baseline_cutoff.timestamp()) + if baseline_cutoff and migration_mode != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY + else None + ) + parallel_operations = _get_migration_parallel_operations(settings) + retry_count = _get_migration_retry_count(settings) + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + target_container_name = _target_cosmos_container_name(container_definition) + resource_name = f"cosmos:{target_type}:{container_definition['name']}" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + previous_result = completed_resource.get("result") if isinstance(completed_resource, dict) else {} + copied.append({ + "name": container_definition["name"], + "type": "cosmos_container", + "target_type": target_type, + "container_name": target_container_name, + "partition_key_path": container_definition["partition_key_path"], + "status": "checkpoint_completed", + **(previous_result if isinstance(previous_result, dict) else {}), + }) + continue + + target_container = _get_target_cosmos_container( + target_database, + target_container_name, + container_definition["partition_key_path"], + ) + resource = start_migration_resource(migration_state, resource_name) + resource_started_at = resource.get("attempt_started_at") or resource.get("started_at") + previous_progress = resource.get("progress") if isinstance(resource.get("progress"), dict) else {} + copied_count = _safe_int(previous_progress.get("copied_count"), default=0, minimum=0) + skipped_count = _safe_int(previous_progress.get("skipped_count"), default=0, minimum=0) + prior_failed_count = _safe_int(previous_progress.get("failed_count"), default=0, minimum=0) + failed_count = 0 + collision_count = _safe_int(previous_progress.get("collision_count"), default=0, minimum=0) + destination_provenance_skip_count = _safe_int( + previous_progress.get("destination_provenance_skip_count"), + default=0, + minimum=0, + ) + created_count = _safe_int(previous_progress.get("created_count"), default=0, minimum=0) + updated_count = _safe_int(previous_progress.get("updated_count"), default=0, minimum=0) + unchanged_count = _safe_int(previous_progress.get("unchanged_count"), default=0, minimum=0) + source_read_count = _safe_int(previous_progress.get("source_read_count"), default=0, minimum=0) + retry_attempt_count = _safe_int(previous_progress.get("retry_attempt_count"), default=0, minimum=0) + byte_count = _safe_int(previous_progress.get("bytes"), default=0, minimum=0) + request_units = float(previous_progress.get("request_units") or 0.0) + errors = [] + append_manifest, flush_manifest = _create_migration_manifest_writer( + job.get("id"), + resource_name, + ) + transfer_cancel_event = Event() + + def persist_progress(allow_cancel_requested=False): + flush_manifest() + progress = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + request_units=request_units, + ) + progress.update({ + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "source_read_count": source_read_count, + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + "prior_failed_count": prior_failed_count, + }) + return _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + f"Migrating Cosmos container {target_container_name}", + allow_cancel_requested=allow_cancel_requested, + ) + + def persist_heartbeat(): + return _persist_migration_heartbeat( + job, + settings, + f"Waiting for Cosmos writes to {target_container_name}", + ) + + def persist_stopped_results(): + return persist_progress(allow_cancel_requested=True) + + def consume_write_result(write_result, manifest_context): + nonlocal copied_count, skipped_count, failed_count + nonlocal collision_count, destination_provenance_skip_count + nonlocal created_count, updated_count, unchanged_count + nonlocal byte_count, request_units, retry_attempt_count + request_units += write_result.get("request_units", 0.0) + retry_attempt_count += max(0, _safe_int(write_result.get("attempt"), default=1) - 1) + if write_result.get("copied"): + copied_count += 1 + if write_result.get("updated"): + updated_count += 1 + else: + created_count += 1 + byte_count += write_result.get("bytes", 0) + elif write_result.get("skipped"): + if not write_result.get("resume_verified"): + skipped_count += 1 + unchanged_count += 1 + destination_provenance_skip_count += 1 + else: + failed_count += 1 + if write_result.get("collision"): + collision_count += 1 + errors.append(write_result.get("error", "Cosmos write failed.")) + status = ( + "updated" + if write_result.get("updated") else + "created" + if write_result.get("copied") else + "resume_verified" + if write_result.get("resume_verified") else + "unchanged" + if write_result.get("skipped") else + "collision" + if write_result.get("collision") else + "failed" + ) + append_manifest({ + **manifest_context, + "status": status, + "attempt": write_result.get("attempt"), + "bytes": write_result.get("bytes", 0), + "error": write_result.get("error"), + }) + + try: + pending_writes = [] + with ThreadPoolExecutor(max_workers=parallel_operations) as executor: + for item, source_version in _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=source_cutoff_epoch, + source_start_epoch=source_start_epoch, + include_source_version=True, + ) or []: + source_read_count += 1 + source_hash = _build_cosmos_source_hash(item) + manifest_context = { + "service": "cosmos", + "resource_name": resource_name, + "target_type": target_type, + "source_identity": _hash_migration_manifest_identity( + target_type, + container_definition["name"], + item.get("id"), + _get_document_path_value( + item, + container_definition["partition_key_path"], + ), + ), + "destination_identity": _hash_migration_manifest_identity( + target_container_name, + item.get("id"), + _get_document_path_value( + item, + container_definition["partition_key_path"], + ), + ), + "source_version": source_version, + "source_hash": source_hash, + "_locator": { + "service": "cosmos", + "resource_name": resource_name, + "target_type": target_type, + "document_id": item.get("id"), + "partition_key": _get_document_path_value( + item, + container_definition["partition_key_path"], + ), + "container_name": target_container_name, + }, + } + target_document = _get_target_cosmos_document( + target_container, + item, + container_definition["partition_key_path"], + ) + disposition = _classify_target_cosmos_document( + target_document, + provenance_context, + target_container_name, + migration_mode=migration_mode, + source_hash=source_hash, + source_version=source_version, + ) + if disposition in {"unchanged", "resume_verified"}: + if disposition == "unchanged": + skipped_count += 1 + unchanged_count += 1 + destination_provenance_skip_count += 1 + append_manifest({**manifest_context, "status": disposition}) + if (copied_count + skipped_count + failed_count) % max(100, parallel_operations * 4) == 0: + migration_state = persist_progress() + continue + + if not pending_writes: + _assert_migration_job_lease(job) + pending_writes.append(( + executor.submit( + _write_cosmos_migration_record, + target_container, + item, + container_definition["partition_key_path"], + provenance_context, + retry_count, + disposition=disposition, + target_document=target_document, + source_hash=source_hash, + source_version=source_version, + cancel_event=transfer_cancel_event, + ), + manifest_context, + )) + if len(pending_writes) < parallel_operations: + continue + + _drain_migration_futures_with_heartbeat( + pending_writes, + consume_write_result, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + pending_writes = [] + migration_state = persist_progress() + + _drain_migration_futures_with_heartbeat( + pending_writes, + consume_write_result, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + + flush_manifest() + + result = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + request_units=request_units, + ) + result.update({ + "name": container_definition["name"], + "type": "cosmos_container", + "target_type": target_type, + "container_name": target_container_name, + "partition_key_path": container_definition["partition_key_path"], + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "source_read_count": source_read_count, + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + "prior_failed_count": prior_failed_count, + }) + if failed_count: + error_summary = "; ".join(errors[:3]) + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + error_summary, + ) + raise RuntimeError( + f"Cosmos migration failed for container '{target_container_name}' after {failed_count} item write failure(s)." + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + f"Completed Cosmos container {target_container_name}", + ) + copied.append(result) + except (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError): + flush_manifest() + raise + except Exception as exc: + flush_manifest() + if not failed_count: + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(exc), + ) + raise + return copied + + +def _escape_search_filter_value(value): + return _safe_text(value).replace("'", "''") + + +def _build_search_filter(field_name, selection): + if selection.get("mode") == "all": + return None + if selection.get("mode") != "selected": + return "id eq '__no_migration_selection__'" + conditions = [f"{field_name} eq '{_escape_search_filter_value(item_id)}'" for item_id in selection.get("ids") or []] + if not conditions: + return "id eq '__no_migration_selection__'" + return " or ".join(conditions) + + +def _combine_search_filters(*search_filters): + normalized_filters = [ + _safe_text(search_filter) + for search_filter in search_filters + if _safe_text(search_filter) + ] + if not normalized_filters: + return None + return " and ".join(f"({search_filter})" for search_filter in normalized_filters) + + +def _build_search_scope_filter_batches(field_name, selection): + """Bound selected-scope filters so large migrations stay within Search limits.""" + if selection.get("mode") == "all": + return [None] + selected_ids = sorted({ + _safe_text(item_id) + for item_id in (selection.get("ids") or []) + if _safe_text(item_id) + }) + if selection.get("mode") != "selected" or not selected_ids: + return ["id eq '__no_migration_selection__'"] + return [ + _build_search_filter(field_name, { + "mode": "selected", + "ids": selected_ids[offset:offset + DATA_MANAGEMENT_SEARCH_SCOPE_FILTER_BATCH_SIZE], + }) + for offset in range(0, len(selected_ids), DATA_MANAGEMENT_SEARCH_SCOPE_FILTER_BATCH_SIZE) + ] + + +def _iter_search_document_pages( + search_client, + search_filter=None, + page_size=None, + last_id="", + select=None, +): + """Yield stable keyset pages without allowing the SDK to fall back to deep skip.""" + safe_page_size = _safe_int( + page_size, + default=DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE, + minimum=1, + maximum=DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE, + ) + cursor_id = _safe_text(last_id) + while True: + cursor_filter = ( + f"id gt '{_escape_search_filter_value(cursor_id)}'" + if cursor_id + else None + ) + search_kwargs = { + "search_text": "*", + "filter": _combine_search_filters(search_filter, cursor_filter), + "order_by": ["id asc"], + "top": safe_page_size, + "include_total_count": False, + } + if select: + search_kwargs["select"] = select + results = search_client.search( + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + **search_kwargs, + ) + page = [] + previous_id = cursor_id + for result in results: + document = { + key: value + for key, value in dict(result).items() + if not key.startswith("@search.") + } + document_id = _safe_text(document.get("id")) + if not document_id: + raise DataManagementSettingsValidationError( + "Source AI Search returned a document without an ID during keyset pagination." + ) + if previous_id and document_id <= previous_id: + raise DataManagementSettingsValidationError( + "Source AI Search did not honor the required strictly increasing ID order." + ) + previous_id = document_id + page.append(document) + if not page: + break + cursor_id = _safe_text(page[-1].get("id")) + yield page, cursor_id + if len(page) < safe_page_size: + break + + +def _get_target_ai_search_credential(settings): + if (settings or {}).get("target_ai_search_authentication_type") == "key": + key = _safe_text((settings or {}).get("target_ai_search_key")) + if not key: + raise ValueError("Target AI Search key is required when key authentication is selected.") + return AzureKeyCredential(key) + return DefaultAzureCredential() + + +def _get_target_search_client(settings, index_name): + endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) + if not endpoint: + raise ValueError("Target AI Search endpoint is required before running AI Search migration.") + return SearchClient( + endpoint=endpoint, + index_name=index_name, + credential=_get_target_ai_search_credential(settings), + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + + +def _normalize_search_service_endpoint(endpoint): + """Compare Search service endpoints without credentials or request paths.""" + normalized_endpoint = _safe_text(endpoint).strip().rstrip("/").lower() + parsed_endpoint = urlparse( + normalized_endpoint if "://" in normalized_endpoint else f"https://{normalized_endpoint}" + ) + return _safe_text(parsed_endpoint.netloc or normalized_endpoint).rstrip("/").lower() + + +def _get_search_client_service_endpoint(search_client): + """Read the SDK endpoint when the source client exposes one.""" + return _normalize_search_service_endpoint( + getattr(search_client, "_endpoint", None) or + getattr(search_client, "endpoint", None) + ) + + +def _migration_plan_includes_ai_search_documents(migration_plan): + """Return whether this plan will write selected document chunks to AI Search.""" + if not (migration_plan or {}).get("include_ai_search"): + return False + return any( + (migration_plan.get(target_type) or {}).get("mode") != "none" and + bool((migration_plan.get(target_type) or {}).get("include_documents")) + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES + ) + + +def _validate_target_ai_search_migration_write_safety(settings, migration_plan): + """Require a frozen, distinct target before non-conditional Search document writes.""" + if not _migration_plan_includes_ai_search_documents(migration_plan): + return + if not bool((migration_plan or {}).get("target_ai_search_writes_frozen")): + raise DataManagementSettingsValidationError( + DATA_MANAGEMENT_SEARCH_WRITE_FREEZE_CONFIRMATION_ERROR + ) + target_endpoint = _normalize_search_service_endpoint( + (settings or {}).get("target_ai_search_endpoint") + ) + if not target_endpoint: + return + for artifact in DATA_MANAGEMENT_SEARCH_ARTIFACTS: + source_client = CLIENTS.get(artifact["client_key"]) + source_endpoint = _get_search_client_service_endpoint(source_client) + if source_endpoint and source_endpoint == target_endpoint: + raise DataManagementSettingsValidationError( + "Target AI Search endpoint must differ from the selected source Search service." + ) + + +def _get_migration_search_provenance_fields(): + return [ + SearchField( + name=SEARCH_MIGRATION_ID_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=True, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATED_AT_FIELD, + type=SearchFieldDataType.DateTimeOffset, + filterable=True, + sortable=True, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATION_STATUS_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=True, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATION_SOURCE_HASH_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=False, + facetable=False, + retrievable=True, + ), + SearchField( + name=SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + type=SearchFieldDataType.String, + searchable=False, + filterable=True, + sortable=False, + facetable=False, + retrievable=True, + ), + ] + + +def _ensure_target_search_migration_provenance_fields(index): + """Add compatible queryable provenance fields without rebuilding an index.""" + existing_fields = { + field.name: field + for field in (getattr(index, "fields", None) or []) + if getattr(field, "name", None) + } + fields = list(getattr(index, "fields", None) or []) + added_fields = 0 + for required_field in _get_migration_search_provenance_fields(): + existing_field = existing_fields.get(required_field.name) + if existing_field is None: + fields.append(required_field) + added_fields += 1 + continue + if ( + str(getattr(existing_field, "type", "")) != str(required_field.type) or + getattr(existing_field, "filterable", False) is not True or + getattr(existing_field, "retrievable", False) is not True + ): + raise DataManagementSettingsValidationError( + f"Target AI Search index '{getattr(index, 'name', '')}' has incompatible " + f"migration provenance field '{required_field.name}'." + ) + if added_fields: + index.fields = fields + return added_fields + + +def _ensure_target_search_index(settings, index_name, schema_file): + endpoint = _safe_text((settings or {}).get("target_ai_search_endpoint")) + if not endpoint: + raise ValueError("Target AI Search endpoint is required before running AI Search migration.") + index_client = SearchIndexClient( + endpoint=endpoint, + credential=_get_target_ai_search_credential(settings), + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + try: + index = index_client.get_index(index_name) + added_fields = _ensure_target_search_migration_provenance_fields(index) + if added_fields: + index_client.create_or_update_index(index) + return "updated_with_migration_provenance" + return "exists_with_migration_provenance" + except ResourceNotFoundError: + schema = _get_search_schema(schema_file) + schema = {key: value for key, value in schema.items() if not key.startswith("@odata.")} + try: + index = SearchIndex.from_dict(schema) + _ensure_target_search_migration_provenance_fields(index) + index_client.create_or_update_index(index) + return "created_with_migration_provenance" + except Exception as exc: + raise ValueError(f"Target AI Search index {index_name} is missing and could not be created: {exc}") from exc + + +def _get_target_search_documents_by_ids(search_client, document_ids): + """Read only the current bounded upload batch from the target index.""" + normalized_ids = [ + _safe_text(document_id) + for document_id in document_ids + if _safe_text(document_id) + ] + if not normalized_ids: + return {} + conditions = [ + f"id eq '{_escape_search_filter_value(document_id)}'" + for document_id in normalized_ids + ] + results = search_client.search( + search_text="*", + filter=" or ".join(conditions), + select=[ + "id", + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + ], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + documents_by_id = {} + for result in results: + document = dict(result) + document_id = _safe_text(document.get("id")) + if document_id: + documents_by_id[document_id] = document + return documents_by_id + + +def _build_search_source_hash(document): + """Hash canonical source fields while excluding destination migration metadata.""" + excluded_fields = { + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + } + canonical_document = { + key: value + for key, value in (document or {}).items() + if key not in excluded_fields and not key.startswith("@search.") + } + encoded = json.dumps( + canonical_document, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=_json_default, + ).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _classify_target_search_document( + target_document, + provenance_context, + index_name, + migration_mode=DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + source_hash="", +): + """Allow only empty or migration-owned target keys to be uploaded.""" + if target_document is None: + return "create" + target_provenance = get_search_migration_provenance(target_document) + if is_successful_migration_record(target_provenance): + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) + ): + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + ): + return "resume_verified" + return "update" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + return "unchanged" + if migration_record_matches_source(target_provenance, source_hash=source_hash): + return "unchanged" + return "update" + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{index_name}' contains an unowned document that conflicts with this migration." + ) + + +def _get_search_result_succeeded(result): + if isinstance(result, dict): + return bool(result.get("succeeded", result.get("status", False))) + return bool(getattr(result, "succeeded", getattr(result, "status", False))) + + +def _reclassify_pending_search_migration_entries( + search_client, + pending_entries, + provenance_context, + index_name, + migration_mode, +): + """Allow only verified response-lost writes to leave an ambiguous Search retry.""" + document_ids = [ + _safe_text((entry.get("document") or {}).get("id")) + for entry in pending_entries + ] + target_documents = _get_target_search_documents_by_ids(search_client, document_ids) + verified_entries = [] + retry_entries = [] + for entry in pending_entries: + document = entry.get("document") if isinstance(entry, dict) else {} + document_id = _safe_text((document or {}).get("id")) + target_document = target_documents.get(document_id) + if target_document is None: + retry_entries.append(entry) + continue + source_hash = _safe_text((document or {}).get(SEARCH_MIGRATION_SOURCE_HASH_FIELD)) + disposition = _classify_target_search_document( + target_document, + provenance_context, + index_name, + migration_mode=migration_mode, + source_hash=source_hash, + ) + if disposition != "resume_verified": + raise DataManagementSettingsValidationError( + f"Destination AI Search index '{index_name}' changed after an ambiguous " + f"migration write for document '{document_id}'. The migration stopped before retrying it." + ) + verified_entries.append(entry) + return verified_entries, retry_entries + + +def _upload_search_migration_batch( + search_client, + documents, + retry_count, + dispositions=None, + manifest_contexts=None, + cancel_event=None, + provenance_context=None, + index_name="", + migration_mode=DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, +): + """Upload one idempotent batch and surface partial indexing failures.""" + pending_entries = [ + { + "document": document, + "disposition": ( + dispositions[index] + if isinstance(dispositions, list) and index < len(dispositions) + else "create" + ), + "bytes": len(json.dumps(document, default=_json_default).encode("utf-8")), + "manifest_context": ( + manifest_contexts[index] + if isinstance(manifest_contexts, list) and index < len(manifest_contexts) + else {} + ), + } + for index, document in enumerate(documents) + ] + started_at = time.perf_counter() + last_error = "" + copied_count = 0 + created_count = 0 + updated_count = 0 + accepted_bytes = 0 + manifest_entries = [] + + def record_accepted_entry(entry, attempt): + nonlocal copied_count, created_count, updated_count, accepted_bytes + copied_count += 1 + accepted_bytes += entry["bytes"] + if entry["disposition"] == "update": + updated_count += 1 + else: + created_count += 1 + manifest_entries.append({ + **entry["manifest_context"], + "status": entry["disposition"] + "d", + "attempt": attempt, + "bytes": entry["bytes"], + }) + + for attempt in range(1, retry_count + 1): + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "AI Search migration stopped after cancellation or lease loss." + ) + try: + results = list(search_client.upload_documents( + documents=[entry["document"] for entry in pending_entries], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + if len(results) != len(pending_entries): + last_error = "AI Search returned an incomplete document result set." + else: + failed_entries = [] + for entry, result in zip(pending_entries, results): + if not _get_search_result_succeeded(result): + failed_entries.append(entry) + continue + record_accepted_entry(entry, attempt) + pending_entries = failed_entries + if not pending_entries: + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": 0, + "bytes": accepted_bytes, + "attempt": attempt, + "manifest_entries": manifest_entries, + "elapsed_seconds": time.perf_counter() - started_at, + } + if not last_error: + last_error = f"AI Search returned {len(pending_entries)} failed document result(s)." + except DataManagementSettingsValidationError: + raise + except Exception as exc: + last_error = str(exc)[:500] + + if attempt < retry_count: + if provenance_context and index_name: + try: + verified_entries, pending_entries = _reclassify_pending_search_migration_entries( + search_client, + pending_entries, + provenance_context, + index_name, + migration_mode, + ) + except DataManagementSettingsValidationError: + raise + except Exception as exc: + last_error = ( + "AI Search migration could not verify unresolved documents after an " + f"ambiguous write: {str(exc)[:400]}" + ) + break + for entry in verified_entries: + record_accepted_entry(entry, attempt) + if not pending_entries: + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": 0, + "bytes": accepted_bytes, + "attempt": attempt, + "manifest_entries": manifest_entries, + "elapsed_seconds": time.perf_counter() - started_at, + } + retry_delay = min(30, 2 ** (attempt - 1)) + if cancel_event is not None: + if cancel_event.wait(retry_delay): + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": 0, + "bytes": accepted_bytes, + "attempt": attempt, + "manifest_entries": manifest_entries, + "interrupted": True, + "elapsed_seconds": time.perf_counter() - started_at, + } + else: + time.sleep(retry_delay) + + manifest_entries.extend({ + **entry["manifest_context"], + "status": "failed", + "attempt": retry_count, + "bytes": 0, + "error": last_error or "AI Search document indexing failed.", + } for entry in pending_entries) + return { + "copied": copied_count, + "created": created_count, + "updated": updated_count, + "failed": len(pending_entries), + "bytes": accepted_bytes, + "attempt": retry_count, + "manifest_entries": manifest_entries, + "elapsed_seconds": time.perf_counter() - started_at, + "error": last_error or "AI Search document indexing failed.", + } + + +def _copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_search_write_fence=None, +): + copied = [] + if not migration_plan.get("include_ai_search"): + return copied + search_mappings = [ + ("users", "user_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), + ("groups", "group_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), + ("public_workspaces", "public_workspace_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), + ] + parallel_operations = _get_migration_parallel_operations(settings) + retry_count = _get_migration_retry_count(settings) + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + batch_size = min(50, DATA_MANAGEMENT_MIGRATION_BATCH_SIZE) + search_page_size = min( + DATA_MANAGEMENT_SEARCH_KEYSET_PAGE_SIZE, + max(batch_size, batch_size * parallel_operations), + ) + gate_container = None + gate_fence = None + if isinstance(target_search_write_fence, tuple) and len(target_search_write_fence) == 2: + gate_container, gate_fence = target_search_write_fence + + def renew_target_search_write_fence(): + if gate_container is None or gate_fence is None: + return + renew_data_management_search_write_fence( + gate_container, + gate_fence, + _get_target_search_write_fence_lease_seconds(settings), + ) + + for target_type, field_name, artifact in search_mappings: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + source_client = CLIENTS.get(artifact["client_key"]) + if not source_client: + copied.append({"name": artifact["name"], "type": "ai_search_documents", "status": "skipped", "warning": "Source AI Search client is not initialized."}) + continue + resource_name = f"ai_search:{target_type}:{artifact['index_name']}" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + previous_result = completed_resource.get("result") if isinstance(completed_resource, dict) else {} + copied.append({ + "name": artifact["name"], + "type": "ai_search_documents", + "target_type": target_type, + "index_name": artifact["index_name"], + "status": "checkpoint_completed", + **(previous_result if isinstance(previous_result, dict) else {}), + }) + continue + + index_status = _ensure_target_search_index(settings, artifact["index_name"], artifact["schema_file"]) + target_client = _get_target_search_client(settings, artifact["index_name"]) + search_filters = _build_search_scope_filter_batches(field_name, selection) + resource = start_migration_resource(migration_state, resource_name) + resource_started_at = resource.get("attempt_started_at") or resource.get("started_at") + previous_progress = resource.get("progress") if isinstance(resource.get("progress"), dict) else {} + previous_cursor = previous_progress.get("keyset_cursor") + previous_cursor = previous_cursor if isinstance(previous_cursor, dict) else {} + copied_count = _safe_int(previous_progress.get("copied_count"), default=0, minimum=0) + skipped_count = _safe_int(previous_progress.get("skipped_count"), default=0, minimum=0) + prior_failed_count = _safe_int(previous_progress.get("failed_count"), default=0, minimum=0) + failed_count = 0 + collision_count = _safe_int(previous_progress.get("collision_count"), default=0, minimum=0) + destination_provenance_skip_count = _safe_int( + previous_progress.get("destination_provenance_skip_count"), + default=0, + minimum=0, + ) + byte_count = _safe_int(previous_progress.get("bytes"), default=0, minimum=0) + source_read_count = _safe_int(previous_progress.get("source_read_count"), default=0, minimum=0) + created_count = _safe_int(previous_progress.get("created_count"), default=0, minimum=0) + updated_count = _safe_int(previous_progress.get("updated_count"), default=0, minimum=0) + unchanged_count = _safe_int(previous_progress.get("unchanged_count"), default=0, minimum=0) + retry_attempt_count = _safe_int(previous_progress.get("retry_attempt_count"), default=0, minimum=0) + checkpoint_count = _safe_int(previous_progress.get("checkpoint_count"), default=0, minimum=0) + cursor_scope_batch_index = _safe_int( + previous_cursor.get("scope_batch_index"), + default=0, + minimum=0, + maximum=len(search_filters), + ) + cursor_last_id = _safe_text(previous_cursor.get("last_id")) + errors = [] + append_manifest, flush_manifest = _create_migration_manifest_writer( + job.get("id"), + resource_name, + ) + transfer_cancel_event = Event() + + def persist_progress(allow_cancel_requested=False): + renew_target_search_write_fence() + flush_manifest() + progress = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + progress.update({ + "parallel_operations": parallel_operations, + "batch_size": batch_size, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "source_read_count": source_read_count, + "destination_accepted_count": copied_count, + "destination_failed_count": failed_count, + "prior_failed_count": prior_failed_count, + "retry_attempt_count": retry_attempt_count, + "checkpoint_count": checkpoint_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "keyset_cursor": { + "scope_batch_index": cursor_scope_batch_index, + "scope_batch_count": len(search_filters), + "last_id": cursor_last_id, + }, + }) + return _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + f"Migrating AI Search index {artifact['index_name']}", + allow_cancel_requested=allow_cancel_requested, + ) + + def persist_heartbeat(): + result = _persist_migration_heartbeat( + job, + settings, + f"Waiting for AI Search uploads to {artifact['index_name']}", + ) + renew_target_search_write_fence() + return result + + def persist_stopped_results(): + return persist_progress(allow_cancel_requested=True) + + def prepare_upload_batch(source_batch): + nonlocal collision_count + source_ids = [_safe_text(document.get("id")) for document in source_batch] + if any(not document_id for document_id in source_ids): + raise DataManagementSettingsValidationError( + f"Source AI Search index '{artifact['index_name']}' contains a document without an ID." + ) + _assert_migration_job_lease(job) + renew_target_search_write_fence() + target_documents = _get_target_search_documents_by_ids( + target_client, + source_ids, + ) + upload_documents = [] + upload_dispositions = [] + upload_manifest_contexts = [] + deferred_outcomes = { + "skipped_count": 0, + "unchanged_count": 0, + "destination_provenance_skip_count": 0, + "manifest_entries": [], + } + for source_document, source_document_id in zip(source_batch, source_ids): + source_hash = _build_search_source_hash(source_document) + manifest_context = { + "service": "ai_search", + "resource_name": resource_name, + "target_type": target_type, + "source_identity": _hash_migration_manifest_identity( + artifact["index_name"], + source_document_id, + ), + "destination_identity": _hash_migration_manifest_identity( + "destination", + artifact["index_name"], + source_document_id, + ), + "source_hash": source_hash, + "_locator": { + "service": "ai_search", + "resource_name": resource_name, + "target_type": target_type, + "document_id": source_document_id, + "index_name": artifact["index_name"], + }, + } + try: + disposition = _classify_target_search_document( + target_documents.get(source_document_id), + provenance_context, + artifact["index_name"], + migration_mode=migration_mode, + source_hash=source_hash, + ) + except DataManagementSettingsValidationError: + collision_count += 1 + raise + if disposition in {"unchanged", "resume_verified"}: + if disposition == "unchanged": + deferred_outcomes["skipped_count"] += 1 + deferred_outcomes["unchanged_count"] += 1 + deferred_outcomes["destination_provenance_skip_count"] += 1 + deferred_outcomes["manifest_entries"].append({ + **manifest_context, + "status": disposition, + }) + continue + upload_documents.append(add_search_migration_provenance( + copy.deepcopy(source_document), + provenance_context, + source_hash=source_hash, + )) + upload_dispositions.append(disposition) + upload_manifest_contexts.append(manifest_context) + return ( + upload_documents, + upload_dispositions, + upload_manifest_contexts, + deferred_outcomes, + ) + + def consume_batch_result(batch_result): + nonlocal copied_count, failed_count, byte_count, retry_attempt_count + nonlocal created_count, updated_count + copied_count += batch_result.get("copied", 0) + created_count += batch_result.get("created", 0) + updated_count += batch_result.get("updated", 0) + failed_count += batch_result.get("failed", 0) + byte_count += batch_result.get("bytes", 0) + retry_attempt_count += max(0, _safe_int(batch_result.get("attempt"), default=1) - 1) + if batch_result.get("error"): + errors.append(batch_result["error"]) + for manifest_entry in batch_result.get("manifest_entries") or []: + append_manifest(manifest_entry) + + def consume_batch_future(batch_result, _context): + consume_batch_result(batch_result) + + try: + with ThreadPoolExecutor(max_workers=parallel_operations) as executor: + for scope_batch_index, search_filter in enumerate(search_filters): + if scope_batch_index < cursor_scope_batch_index: + continue + last_id = cursor_last_id if scope_batch_index == cursor_scope_batch_index else "" + for source_page, page_last_id in _iter_search_document_pages( + source_client, + search_filter=search_filter, + page_size=search_page_size, + last_id=last_id, + ): + page_failed_count = failed_count + pending_batches = [] + page_outcomes = { + "skipped_count": 0, + "unchanged_count": 0, + "destination_provenance_skip_count": 0, + "manifest_entries": [], + } + for offset in range(0, len(source_page), batch_size): + source_batch = [] + source_batch.extend(source_page[offset:offset + batch_size]) + if not source_batch: + continue + ( + upload_documents, + upload_dispositions, + upload_manifest_contexts, + deferred_outcomes, + ) = prepare_upload_batch(source_batch) + for field_name in ( + "skipped_count", + "unchanged_count", + "destination_provenance_skip_count", + ): + page_outcomes[field_name] += deferred_outcomes[field_name] + page_outcomes["manifest_entries"].extend( + deferred_outcomes["manifest_entries"] + ) + if upload_documents: + pending_batches.append(executor.submit( + _upload_search_migration_batch, + target_client, + upload_documents, + retry_count, + upload_dispositions, + upload_manifest_contexts, + transfer_cancel_event, + provenance_context=provenance_context, + index_name=artifact["index_name"], + migration_mode=migration_mode, + )) + if len(pending_batches) < parallel_operations: + continue + _drain_migration_futures_with_heartbeat( + {future: None for future in pending_batches}, + consume_batch_future, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + pending_batches = [] + _drain_migration_futures_with_heartbeat( + {future: None for future in pending_batches}, + consume_batch_future, + persist_heartbeat, + transfer_cancel_event, + persist_stopped_results=persist_stopped_results, + ) + + if failed_count > page_failed_count: + migration_state = persist_progress() + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + "; ".join(errors[:3]) or "AI Search document indexing failed.", + ) + raise RuntimeError( + f"AI Search migration failed for index '{artifact['index_name']}' " + "before its keyset cursor could advance." + ) + + source_read_count += len(source_page) + skipped_count += page_outcomes["skipped_count"] + unchanged_count += page_outcomes["unchanged_count"] + destination_provenance_skip_count += page_outcomes[ + "destination_provenance_skip_count" + ] + for manifest_entry in page_outcomes["manifest_entries"]: + append_manifest(manifest_entry) + cursor_scope_batch_index = scope_batch_index + cursor_last_id = page_last_id + checkpoint_count += 1 + migration_state = persist_progress() + + cursor_scope_batch_index = scope_batch_index + 1 + cursor_last_id = "" + checkpoint_count += 1 + migration_state = persist_progress() + + result = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + flush_manifest() + result.update({ + "name": artifact["name"], + "type": "ai_search_documents", + "target_type": target_type, + "index_name": artifact["index_name"], + "scope_filter_batch_count": len(search_filters), + "index_status": index_status, + "parallel_operations": parallel_operations, + "batch_size": batch_size, + "search_page_size": search_page_size, + "retry_count": retry_count, + "destination_provenance_skip_count": destination_provenance_skip_count, + "collision_count": collision_count, + "source_read_count": source_read_count, + "destination_accepted_count": copied_count, + "destination_failed_count": failed_count, + "prior_failed_count": prior_failed_count, + "retry_attempt_count": retry_attempt_count, + "checkpoint_count": checkpoint_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "keyset_cursor": { + "scope_batch_index": len(search_filters), + "scope_batch_count": len(search_filters), + "last_id": "", + "completed": True, + }, + }) + if failed_count: + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + "; ".join(errors[:3]) or "AI Search document indexing failed.", + ) + raise RuntimeError( + f"AI Search migration failed for index '{artifact['index_name']}' after {failed_count} document failure(s)." + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + f"Completed AI Search index {artifact['index_name']}", + ) + copied.append(result) + except (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError): + flush_manifest() + raise + except Exception as exc: + flush_manifest() + if not failed_count: + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(exc), + ) + raise + return copied + + +def _get_target_enhanced_citations_blob_client(settings): + auth_type = _safe_text((settings or {}).get("target_enhanced_citations_storage_authentication_type"), "managed_identity") + if auth_type == "connection_string": + connection_string = _safe_text((settings or {}).get("target_enhanced_citations_storage_connection_string")) + if not connection_string: + raise ValueError("Destination Enhanced Citations storage connection string is required when migrating source document blobs.") + return BlobServiceClient.from_connection_string(connection_string) + blob_endpoint = _safe_text((settings or {}).get("target_enhanced_citations_storage_blob_endpoint")) + if not blob_endpoint: + raise ValueError("Destination Enhanced Citations Blob endpoint is required when migrating source document blobs.") + return BlobServiceClient(account_url=blob_endpoint, credential=DefaultAzureCredential()) + + +def _document_blob_reference(document_item): + if not isinstance(document_item, dict): + return None + container_name = document_item.get("blob_container") or _blob_container_for_document(document_item) + blob_path = document_item.get("blob_path") or document_item.get("archived_blob_path") + if not container_name or not blob_path: + return None + return container_name, blob_path + + +def _document_migration_scope_hash(document_item): + if not isinstance(document_item, dict): + return "" + for scope_type, field_name in ( + ("public_workspaces", "public_workspace_id"), + ("groups", "group_id"), + ("users", "user_id"), + ): + scope_id = _safe_text(document_item.get(field_name)) + if scope_id: + encoded = f"{scope_type}:{scope_id}".encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + return "" + + +def _blob_container_for_document(document_item): + if document_item.get("public_workspace_id"): + return app_config.storage_account_public_documents_container_name + if document_item.get("group_id"): + return app_config.storage_account_group_documents_container_name + return app_config.storage_account_user_documents_container_name + + +def _iter_selected_document_records_for_blob_migration(migration_plan, source_cutoff_epoch=None): + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + container_definition = next( + ( + definition + for definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type] + if definition.get("documents") + ), + None, + ) + if not container_definition: + continue + yield from _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=source_cutoff_epoch, + ) or [] + + +def _get_blob_properties_or_none(blob_client): + try: + return blob_client.get_blob_properties() + except ResourceNotFoundError: + return None + + +def _get_blob_content_md5(blob_properties): + content_settings = getattr(blob_properties, "content_settings", None) + content_md5 = getattr(content_settings, "content_md5", None) + if isinstance(content_md5, (bytes, bytearray)): + return base64.b64encode(bytes(content_md5)).decode("ascii") + return _safe_text(content_md5) + + +def _build_blob_source_fingerprint(blob_properties): + """Build stable source content and version markers from service properties.""" + content_md5 = _get_blob_content_md5(blob_properties) + last_modified = getattr(blob_properties, "last_modified", None) + if isinstance(last_modified, datetime): + last_modified = last_modified.astimezone(timezone.utc).isoformat() + version_payload = { + "etag": _safe_text(getattr(blob_properties, "etag", None)), + "last_modified": _safe_text(last_modified), + "size": _safe_int(getattr(blob_properties, "size", 0), default=0, minimum=0), + "blob_type": _safe_text(getattr(blob_properties, "blob_type", None)), + } + encoded = json.dumps( + version_payload, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + source_hash = f"md5:{content_md5}" if content_md5 else "" + source_version = f"sha256:{hashlib.sha256(encoded).hexdigest()}" + return source_hash, source_version + + +def _get_source_blob_tags(source_blob_client, source_properties): + source_tags = getattr(source_properties, "tags", None) + if isinstance(source_tags, dict): + return source_tags + if _safe_int(getattr(source_properties, "tag_count", 0), default=0, minimum=0) <= 0: + return None + get_blob_tags = getattr(source_blob_client, "get_blob_tags", None) + if not callable(get_blob_tags): + return None + tags = get_blob_tags() + return tags if isinstance(tags, dict) else None + + +def _classify_target_blob( + target_properties, + provenance_context, + migration_mode, + source_hash, + source_version, +): + if target_properties is None: + return "create" + target_provenance = get_blob_migration_provenance( + getattr(target_properties, "metadata", None) + ) + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) and + _safe_text(target_provenance.get("status")).lower() == "pending" + ): + return "update" + if is_successful_migration_record(target_provenance): + if ( + _safe_text(target_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) + ): + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "resume_verified" + return "update" + if migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + return "unchanged" + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + return "unchanged" + return "update" + return "collision" + + +def _build_pending_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash, + source_version, + scope_hash, +): + metadata = merge_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash=source_hash, + source_version=source_version, + scope_hash=scope_hash, + ) + metadata[BLOB_MIGRATION_STATUS_METADATA_KEY] = "pending" + return metadata + + +def _iter_migration_blob_chunks( + source_download, + progress_callback=None, + cancel_event=None, +): + for chunk in source_download.chunks(): + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Blob transfer stopped after migration cancellation or lease loss." + ) + if progress_callback: + progress_callback(len(chunk)) + yield chunk + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Blob transfer stopped after migration cancellation or lease loss." + ) + + +def _copy_source_blob_migration_record( + source_blob_client, + target_blob_client, + provenance_context, + retry_count, + scope_hash="", + progress_callback=None, + cancel_event=None, +): + """Stream one blob to the target and stamp metadata only on success.""" + try: + source_properties = source_blob_client.get_blob_properties() + except ResourceNotFoundError: + return {"status": "missing", "bytes": 0} + + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + source_hash, source_version = _build_blob_source_fingerprint(source_properties) + existing_target_properties = _get_blob_properties_or_none(target_blob_client) + disposition = _classify_target_blob( + existing_target_properties, + provenance_context, + migration_mode, + source_hash, + source_version, + ) + if disposition in {"unchanged", "resume_verified"}: + return { + "status": disposition, + "bytes": 0, + "source_hash": source_hash, + "source_version": source_version, + } + if disposition == "collision": + return { + "status": "collision", + "bytes": 0, + "error": "Destination blob exists without successful migration provenance.", + } + + source_metadata = getattr(source_properties, "metadata", None) or {} + pending_metadata = _build_pending_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash, + source_version, + scope_hash, + ) + destination_metadata = merge_blob_migration_metadata( + source_metadata, + provenance_context, + source_hash=source_hash, + source_version=source_version, + scope_hash=scope_hash, + ) + source_size = _safe_int(getattr(source_properties, "size", 0), default=0, minimum=0) + content_settings = getattr(source_properties, "content_settings", None) + source_tags = _get_source_blob_tags(source_blob_client, source_properties) + source_blob_tier = getattr(source_properties, "blob_tier", None) + source_blob_type = getattr(source_properties, "blob_type", None) + source_etag = _safe_text(getattr(source_properties, "etag", None)) + started_at = time.perf_counter() + last_error = "" + + for attempt in range(1, retry_count + 1): + try: + if progress_callback: + progress_callback(0, reset=True) + if cancel_event is not None and cancel_event.is_set(): + raise DataManagementMigrationCanceledError( + "Blob transfer stopped after migration cancellation or lease loss." + ) + download_kwargs = { + "max_concurrency": 1, + "timeout": DATA_MANAGEMENT_BLOB_SERVICE_TIMEOUT_SECONDS, + } + if source_etag: + download_kwargs.update({ + "etag": source_etag, + "match_condition": MatchConditions.IfNotModified, + }) + source_download = source_blob_client.download_blob(**download_kwargs) + upload_kwargs = { + "data": _iter_migration_blob_chunks( + source_download, + progress_callback=progress_callback, + cancel_event=cancel_event, + ), + "length": source_size, + "overwrite": disposition == "update", + "metadata": pending_metadata, + "content_settings": content_settings, + "max_concurrency": 1, + "timeout": DATA_MANAGEMENT_BLOB_SERVICE_TIMEOUT_SECONDS, + } + if source_tags is not None: + upload_kwargs["tags"] = source_tags + if source_blob_tier is not None: + upload_kwargs["standard_blob_tier"] = source_blob_tier + if source_blob_type is not None: + upload_kwargs["blob_type"] = source_blob_type + if disposition == "update" and existing_target_properties is not None: + target_etag = getattr(existing_target_properties, "etag", None) + if target_etag: + upload_kwargs["etag"] = target_etag + upload_kwargs["match_condition"] = MatchConditions.IfNotModified + target_blob_client.upload_blob(**upload_kwargs) + destination_properties = _get_blob_properties_or_none(target_blob_client) + destination_size = _safe_int( + getattr(destination_properties, "size", source_size), + default=source_size, + minimum=0, + ) + destination_md5 = _get_blob_content_md5(destination_properties) + if destination_size != source_size or ( + source_hash and destination_md5 and source_hash != f"md5:{destination_md5}" + ): + raise RuntimeError( + "Destination blob verification did not match the source content properties." + ) + current_source_properties = source_blob_client.get_blob_properties() + current_source_etag = _safe_text(getattr(current_source_properties, "etag", None)) + if source_etag and current_source_etag != source_etag: + raise RuntimeError("Source blob changed while it was being transferred.") + destination_etag = _safe_text(getattr(destination_properties, "etag", None)) + set_metadata_kwargs = {"metadata": destination_metadata} + if destination_etag: + set_metadata_kwargs.update({ + "etag": destination_etag, + "match_condition": MatchConditions.IfNotModified, + }) + target_blob_client.set_blob_metadata(**set_metadata_kwargs) + return { + "status": "updated" if disposition == "update" else "copied", + "bytes": source_size, + "attempt": attempt, + "source_size": source_size, + "destination_size": destination_size, + "source_hash": source_hash, + "source_version": source_version, + "elapsed_seconds": time.perf_counter() - started_at, + } + except Exception as exc: + if isinstance(exc, (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError)): + raise + last_error = str(exc)[:500] + status_code = getattr(exc, "status_code", None) + if isinstance(exc, ResourceExistsError) or status_code in {409, 412}: + concurrent_target_properties = _get_blob_properties_or_none(target_blob_client) + concurrent_disposition = _classify_target_blob( + concurrent_target_properties, + provenance_context, + migration_mode, + source_hash, + source_version, + ) + if concurrent_disposition in {"unchanged", "resume_verified"}: + return {"status": concurrent_disposition, "bytes": 0} + if concurrent_disposition == "update": + disposition = "update" + existing_target_properties = concurrent_target_properties + continue + if concurrent_disposition == "collision": + return { + "status": "collision", + "bytes": 0, + "error": "Destination blob changed without migration ownership during this migration.", + } + retryable = status_code in {408, 429, 500, 503} or status_code is None + if not retryable or attempt >= retry_count: + break + time.sleep(min(30, 2 ** (attempt - 1))) + + if disposition == "create": + failed_target_properties = _get_blob_properties_or_none(target_blob_client) + failed_provenance = get_blob_migration_provenance( + getattr(failed_target_properties, "metadata", None) + ) if failed_target_properties else {} + if ( + _safe_text(failed_provenance.get("migrationId")) == + _safe_text((provenance_context or {}).get("migration_id")) and + _safe_text(failed_provenance.get("status")).lower() == "pending" + ): + delete_kwargs = {} + failed_etag = _safe_text(getattr(failed_target_properties, "etag", None)) + if failed_etag: + delete_kwargs.update({ + "etag": failed_etag, + "match_condition": MatchConditions.IfNotModified, + }) + try: + target_blob_client.delete_blob(**delete_kwargs) + except Exception: + pass + + return { + "status": "failed", + "bytes": 0, + "attempt": retry_count, + "elapsed_seconds": time.perf_counter() - started_at, + "error": last_error or "Source blob copy failed.", + } + + +def _copy_source_blobs_to_target(settings, migration_plan, job, migration_state, provenance_context): + if not migration_plan.get("include_source_blobs"): + return [] + source_client = _get_source_blob_service_client() + if not source_client: + return [{"name": "source_blobs", "type": "source_blobs", "status": "skipped", "warning": "Source Enhanced Citations storage is not configured."}] + + resource_name = "source_blobs:selected_documents" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + previous_result = completed_resource.get("result") if isinstance(completed_resource, dict) else {} + return [{ + "name": "source_blobs", + "type": "source_blobs", + "status": "checkpoint_completed", + **(previous_result if isinstance(previous_result, dict) else {}), + }] + + target_client = _get_target_enhanced_citations_blob_client(settings) + resource = start_migration_resource(migration_state, resource_name) + resource_started_at = resource.get("attempt_started_at") or resource.get("started_at") + source_cutoff = _parse_iso_datetime(migration_state.get("source_cutoff_at")) + source_cutoff_epoch = int(source_cutoff.timestamp()) if source_cutoff else None + parallel_operations = _get_migration_parallel_operations(settings) + retry_count = _get_migration_retry_count(settings) + migration_mode = _safe_text( + (provenance_context or {}).get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ) + previous_progress = resource.get("progress") if isinstance(resource.get("progress"), dict) else {} + copied_count = _safe_int(previous_progress.get("copied_count"), default=0, minimum=0) + skipped_count = _safe_int(previous_progress.get("skipped_count"), default=0, minimum=0) + prior_failed_count = _safe_int(previous_progress.get("failed_count"), default=0, minimum=0) + prior_collision_count = _safe_int(previous_progress.get("collision_count"), default=0, minimum=0) + prior_missing_count = _safe_int(previous_progress.get("missing_count"), default=0, minimum=0) + failed_count = 0 + collision_count = 0 + missing_count = 0 + not_applicable_count = _safe_int(previous_progress.get("not_applicable_count"), default=0, minimum=0) + created_count = _safe_int(previous_progress.get("created_count"), default=0, minimum=0) + updated_count = _safe_int(previous_progress.get("updated_count"), default=0, minimum=0) + unchanged_count = _safe_int(previous_progress.get("unchanged_count"), default=0, minimum=0) + retry_attempt_count = _safe_int(previous_progress.get("retry_attempt_count"), default=0, minimum=0) + byte_count = _safe_int(previous_progress.get("bytes"), default=0, minimum=0) + container_stats = {} + errors = [] + append_manifest, flush_manifest = _create_migration_manifest_writer( + job.get("id"), + resource_name, + ) + transfer_progress = {} + transfer_progress_lock = Lock() + transfer_cancel_event = Event() + + def persist_progress(message="Migrating selected source document blobs"): + nonlocal migration_state + flush_manifest() + with transfer_progress_lock: + in_flight_bytes = sum(transfer_progress.values()) + progress = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + progress.update({ + "missing_count": missing_count, + "not_applicable_count": not_applicable_count, + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "collision_count": collision_count, + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + "prior_failed_count": prior_failed_count, + "prior_collision_count": prior_collision_count, + "prior_missing_count": prior_missing_count, + "in_flight_bytes": in_flight_bytes, + "observed_bytes": byte_count + in_flight_bytes, + "observed_bytes_per_second": round( + (byte_count + in_flight_bytes) / + max(0.001, float(progress.get("elapsed_seconds") or 0.001)), + 3, + ), + "active_transfer_count": len(transfer_progress), + }) + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + message, + ) + return migration_state + + def make_transfer_progress_callback(reference_key): + def report_progress(byte_delta, reset=False): + with transfer_progress_lock: + if reset: + transfer_progress[reference_key] = 0 + else: + transfer_progress[reference_key] = ( + transfer_progress.get(reference_key, 0) + + _safe_int(byte_delta, default=0, minimum=0) + ) + return report_progress + + def consume_copy_future(result_container_name, reference_key, future): + nonlocal copied_count, skipped_count, failed_count, collision_count, missing_count, byte_count + nonlocal created_count, updated_count, unchanged_count, retry_attempt_count + with transfer_progress_lock: + transfer_progress.pop(reference_key, None) + copy_result = future.result() + artifact = container_stats[result_container_name] + artifact["blob_count"] += 1 + status = copy_result.get("status") + retry_attempt_count += max(0, _safe_int(copy_result.get("attempt"), default=1) - 1) + if status in {"copied", "updated"}: + copied_count += 1 + if status == "updated": + updated_count += 1 + artifact["updated_count"] += 1 + else: + created_count += 1 + artifact["created_count"] += 1 + byte_count += copy_result.get("bytes", 0) + artifact["copied_count"] += 1 + artifact["bytes"] += copy_result.get("bytes", 0) + elif status in {"unchanged", "resume_verified"}: + if status == "unchanged": + skipped_count += 1 + unchanged_count += 1 + artifact["skipped_count"] += 1 + if status == "unchanged": + artifact["unchanged_count"] += 1 + elif status == "missing": + missing_count += 1 + artifact["missing_count"] += 1 + else: + failed_count += 1 + if status == "collision": + collision_count += 1 + artifact["collision_count"] += 1 + artifact["failed_count"] += 1 + errors.append(copy_result.get("error", "Source blob copy failed.")) + append_manifest({ + "service": "source_blobs", + "resource_name": resource_name, + "source_identity": _hash_migration_manifest_identity( + "source", + result_container_name, + reference_key.split("\n", 1)[-1], + ), + "destination_identity": _hash_migration_manifest_identity( + "destination", + result_container_name, + reference_key.split("\n", 1)[-1], + ), + "status": status, + "attempt": copy_result.get("attempt"), + "bytes": copy_result.get("bytes", 0), + "source_version": copy_result.get("source_version"), + "source_hash": copy_result.get("source_hash"), + "error": copy_result.get("error"), + "_locator": { + "service": "source_blobs", + "resource_name": resource_name, + "container_name": result_container_name, + "blob_name": reference_key.split("\n", 1)[-1], + }, + }) + + def process_pending_copies(pending_copies): + pending_by_future = { + future: (result_container_name, reference_key) + for result_container_name, reference_key, future in pending_copies + } + last_heartbeat = 0.0 + while pending_by_future: + completed_futures, _pending_futures = wait( + set(pending_by_future), + timeout=1.0, + return_when=FIRST_COMPLETED, + ) + for future in completed_futures: + result_container_name, reference_key = pending_by_future.pop(future) + consume_copy_future(result_container_name, reference_key, future) + now = time.monotonic() + if pending_by_future and now - last_heartbeat >= 2.0: + try: + persist_progress("Migrating source blobs; worker heartbeat renewed") + except (DataManagementMigrationCanceledError, DataManagementMigrationLeaseLostError): + transfer_cancel_event.set() + for future in pending_by_future: + future.cancel() + raise + last_heartbeat = now + + try: + pending_copies = [] + pending_reference_keys = set() + with ThreadPoolExecutor(max_workers=parallel_operations) as executor: + for document in _iter_selected_document_records_for_blob_migration( + migration_plan, + source_cutoff_epoch=source_cutoff_epoch, + ): + reference = _document_blob_reference(document) + if not reference: + not_applicable_count += 1 + append_manifest({ + "service": "source_blobs", + "resource_name": resource_name, + "source_identity": _hash_migration_manifest_identity( + "no-source-blob", + document.get("id"), + ), + "status": "not_applicable", + "_locator": { + "service": "source_blobs", + "resource_name": resource_name, + "document_id": document.get("id"), + }, + }) + continue + container_name, blob_path = reference + reference_key = f"{container_name}\n{blob_path}" + if reference_key in pending_reference_keys: + continue + artifact = container_stats.setdefault(container_name, { + "name": container_name, + "type": "source_blob_container", + "container_name": container_name, + "blob_count": 0, + "copied_count": 0, + "created_count": 0, + "updated_count": 0, + "unchanged_count": 0, + "skipped_count": 0, + "missing_count": 0, + "collision_count": 0, + "failed_count": 0, + "bytes": 0, + }) + try: + target_client.create_container(container_name) + except ResourceExistsError: + pass + source_blob_client = source_client.get_blob_client( + container=container_name, + blob=blob_path, + ) + target_blob_client = target_client.get_blob_client( + container=container_name, + blob=blob_path, + ) + if not pending_copies: + _assert_migration_job_lease(job) + pending_reference_keys.add(reference_key) + progress_callback = make_transfer_progress_callback(reference_key) + pending_copies.append((container_name, reference_key, executor.submit( + _copy_source_blob_migration_record, + source_blob_client, + target_blob_client, + provenance_context, + retry_count, + _document_migration_scope_hash(document), + progress_callback, + transfer_cancel_event, + ))) + if len(pending_copies) < parallel_operations: + continue + process_pending_copies(pending_copies) + pending_copies = [] + pending_reference_keys.clear() + persist_progress() + + process_pending_copies(pending_copies) + + flush_manifest() + + result = build_transfer_metrics( + resource_started_at, + copied_count=copied_count, + skipped_count=skipped_count, + failed_count=failed_count, + byte_count=byte_count, + ) + result.update({ + "name": "source_blobs", + "type": "source_blobs", + "missing_count": missing_count, + "not_applicable_count": not_applicable_count, + "parallel_operations": parallel_operations, + "retry_count": retry_count, + "collision_count": collision_count, + "container_count": len(container_stats), + "migration_mode": migration_mode, + "baseline_source_cutoff_at": _safe_text( + (provenance_context or {}).get("baseline_source_cutoff_at") + ), + "created_count": created_count, + "updated_count": updated_count, + "unchanged_count": unchanged_count, + "retry_attempt_count": retry_attempt_count, + }) + if failed_count or missing_count: + failure_message = ( + f"{missing_count} persisted source blob reference(s) could not be resolved." + if missing_count else + "; ".join(errors[:3]) or "Source blob copy failed." + ) + persist_progress("Source blob migration requires remediation before retry") + migration_state = _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + failure_message, + ) + raise RuntimeError( + "Source blob migration requires remediation before retry: " + f"{failed_count} failed and {missing_count} unexpectedly missing blob(s)." + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + "Completed selected source document blob migration", + ) + artifacts = list(container_stats.values()) + artifacts.append(result) + return artifacts + except Exception as exc: + flush_manifest() + if not failed_count: + _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(exc), + ) + raise + + +def _iter_target_cosmos_records(target_container, container_definition, selection): + """Enumerate the exact destination scope used by the migration plan.""" + mode = selection.get("mode") + if mode == "all": + yield from target_container.query_items( + query="SELECT * FROM c", + enable_cross_partition_query=True, + ) + return + if mode != "selected": + return + if container_definition.get("id_field") == "id": + for item_id in selection.get("ids") or []: + try: + yield target_container.read_item(item=item_id, partition_key=item_id) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + continue + return + filter_field = container_definition.get("filter_field") + if not filter_field: + return + for selected_id in selection.get("ids") or []: + yield from target_container.query_items( + query=f"SELECT * FROM c WHERE c.{filter_field} = @selected_id", + parameters=[{"name": "@selected_id", "value": selected_id}], + enable_cross_partition_query=True, + ) + + +def _delete_target_cosmos_record(target_container, document, partition_key_path): + delete_kwargs = { + "item": document.get("id"), + "partition_key": _get_document_path_value(document, partition_key_path), + } + if document.get("_etag"): + delete_kwargs.update({ + "etag": document.get("_etag"), + "match_condition": MatchConditions.IfNotModified, + }) + target_container.delete_item(**delete_kwargs) + + +def _get_reconciliation_target_cosmos_container( + target_database, + container_name, + partition_key_path, + read_only, +): + if not read_only: + return _get_target_cosmos_container( + target_database, + container_name, + partition_key_path, + ) + get_container_client = getattr(target_database, "get_container_client", None) + if not callable(get_container_client): + return _get_target_cosmos_container( + target_database, + container_name, + partition_key_path, + ) + target_container = get_container_client(container_name) + try: + _validate_target_cosmos_container_partition_key( + target_container, + container_name, + partition_key_path, + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return None + return target_container + + +def _migration_blob_container_target_type(container_name): + container_mappings = { + getattr(app_config, "storage_account_user_documents_container_name", ""): "users", + getattr(app_config, "storage_account_group_documents_container_name", ""): "groups", + getattr(app_config, "storage_account_public_documents_container_name", ""): "public_workspaces", + } + return container_mappings.get(container_name) + + +def _selected_scope_hashes(target_type, selection): + if selection.get("mode") == "all": + return None + return { + hashlib.sha256(f"{target_type}:{scope_id}".encode("utf-8")).hexdigest() + for scope_id in _dedupe_limited_strings(selection.get("ids"), limit=2000) + } + + +def _new_reconciliation_report(service, include_blob_fields=False): + report = { + "service": service, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + "resources": [], + } + if include_blob_fields: + report.update({ + "unresolved_scope_count": 0, + "not_applicable_count": 0, + "source_missing_count": 0, + }) + return report + + +def _start_reconciliation_heartbeat_thread( + job, + settings, + state_holder, + stop_event, + failure_holder, + persistence_lock, +): + def run_heartbeat(): + while not stop_event.wait(2.0): + try: + with persistence_lock: + state = state_holder["state"] + state_holder["state"] = _persist_migration_state( + job, + state, + settings, + "Reconciliation worker heartbeat renewed", + ) + except Exception as exc: + failure_holder["error"] = exc + stop_event.set() + return + + heartbeat_thread = Thread(target=run_heartbeat, daemon=True) + heartbeat_thread.start() + return heartbeat_thread + + +def _raise_reconciliation_heartbeat_failure(stop_event, failure_holder): + if not stop_event.is_set(): + return + error = failure_holder.get("error") + if error: + raise error + raise DataManagementMigrationCanceledError("Reconciliation was canceled.") + + +def _append_reconciliation_resource(report, resource_report): + report["resources"].append(resource_report) + for field_name in report: + if field_name in {"service", "resources"} or str(field_name).startswith("_"): + continue + report[field_name] += _safe_int(resource_report.get(field_name), default=0, minimum=0) + + +def _emit_reconciliation_candidate(report, candidate, deletion_candidate_callback): + if callable(deletion_candidate_callback): + deletion_candidate_callback(_sanitize_mirror_deletion_candidate(candidate)) + + +def _reconcile_cosmos_migration( + target_database, + migration_plan, + migration_state, + provenance_context, + apply_deletions=False, + heartbeat_callback=None, + deletion_candidate_callback=None, + stop_event=None, + failure_holder=None, +): + """Reconcile Cosmos with bounded point reads and one destination scan.""" + report = _new_reconciliation_report("cosmos") + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none": + continue + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + container_name = _target_cosmos_container_name(container_definition) + target_container = _get_reconciliation_target_cosmos_container( + target_database, + container_name, + container_definition["partition_key_path"], + read_only=True, + ) + source_container = getattr(app_config, container_definition["container_attr"], None) + resource_report = { + "target_type": target_type, + "container_name": container_name, + "source_count": 0, + "target_count": 0, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + } + for item, source_version in _iter_selected_cosmos_records( + container_definition, + selection, + include_source_version=True, + ) or []: + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + resource_report["source_count"] += 1 + if heartbeat_callback and resource_report["source_count"] % 250 == 0: + heartbeat_callback( + f"Reconciling source Cosmos container {container_name}", + resource_report["source_count"], + ) + target_document = ( + _get_target_cosmos_document( + target_container, + item, + container_definition["partition_key_path"], + ) + if target_container is not None else None + ) + if target_document is None: + resource_report["missing_count"] += 1 + resource_report["create_count"] += 1 + continue + target_provenance = get_cosmos_migration_provenance(target_document) + if not is_successful_migration_record(target_provenance): + resource_report["conflict_count"] += 1 + continue + resource_report["matched_count"] += 1 + source_hash = _build_cosmos_source_hash(item) + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + resource_report["unchanged_count"] += 1 + else: + resource_report["stale_count"] += 1 + if migration_plan.get("migration_mode") != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resource_report["update_count"] += 1 + + target_records = ( + _iter_target_cosmos_records(target_container, container_definition, selection) + if target_container is not None else [] + ) + for target_document in target_records or []: + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + resource_report["target_count"] += 1 + if heartbeat_callback and resource_report["target_count"] % 250 == 0: + heartbeat_callback( + f"Reconciling destination Cosmos container {container_name}", + resource_report["target_count"], + ) + target_provenance = get_cosmos_migration_provenance(target_document) + if _source_cosmos_document_exists( + source_container, + target_document, + container_definition["partition_key_path"], + ): + continue + if is_successful_migration_record(target_provenance): + resource_report["destination_only_owned_count"] += 1 + resource_report["remaining_destination_only_owned_count"] += 1 + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR: + resource_report["delete_candidate_count"] += 1 + _emit_reconciliation_candidate(report, { + "service": "cosmos", + "target_type": target_type, + "container_name": container_name, + "document_id": target_document.get("id"), + "partition_key": _get_document_path_value( + target_document, + container_definition["partition_key_path"], + ), + "target_etag": target_document.get("_etag"), + }, deletion_candidate_callback) + else: + resource_report["destination_only_unowned_count"] += 1 + _append_reconciliation_resource(report, resource_report) + return report + + +def _source_cosmos_document_exists(source_container, target_document, partition_key_path): + if source_container is None: + return True + try: + source_container.read_item( + item=target_document.get("id"), + partition_key=_get_document_path_value(target_document, partition_key_path), + ) + return True + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return False + + +def _iter_search_documents_keyset(search_client, search_filter=None, select=None): + for page, _last_id in _iter_search_document_pages( + search_client, + search_filter=search_filter, + select=select, + ): + yield from page + + +def _next_or_none(iterator): + try: + return next(iterator) + except StopIteration: + return None + + +def _reconcile_ai_search_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=None, + deletion_candidate_callback=None, + stop_event=None, + failure_holder=None, +): + """Reconcile Search by merging two ordered keyset streams.""" + report = _new_reconciliation_report("ai_search") + if not migration_plan.get("include_ai_search"): + return report + search_mappings = [ + ("users", "user_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[0]), + ("groups", "group_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[1]), + ("public_workspaces", "public_workspace_id", DATA_MANAGEMENT_SEARCH_ARTIFACTS[2]), + ] + provenance_select = [ + "id", + SEARCH_MIGRATION_ID_FIELD, + SEARCH_MIGRATED_AT_FIELD, + SEARCH_MIGRATION_STATUS_FIELD, + SEARCH_MIGRATION_SOURCE_HASH_FIELD, + SEARCH_MIGRATION_SOURCE_VERSION_FIELD, + ] + for target_type, field_name, artifact in search_mappings: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none" or not selection.get("include_documents"): + continue + source_client = CLIENTS.get(artifact["client_key"]) + if not source_client: + continue + target_client = _get_target_search_client(settings, artifact["index_name"]) + resource_report = { + "target_type": target_type, + "index_name": artifact["index_name"], + "source_count": 0, + "target_count": 0, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + } + for search_filter in _build_search_scope_filter_batches(field_name, selection): + source_iterator = iter(_iter_search_documents_keyset( + source_client, + search_filter=search_filter, + )) + try: + target_iterator = iter(_iter_search_documents_keyset( + target_client, + search_filter=search_filter, + select=provenance_select, + )) + source_document = _next_or_none(source_iterator) + target_document = _next_or_none(target_iterator) + except ResourceNotFoundError: + source_document = _next_or_none(source_iterator) + target_document = None + target_iterator = iter(()) + while source_document is not None or target_document is not None: + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + source_id = _safe_text((source_document or {}).get("id")) + target_id = _safe_text((target_document or {}).get("id")) + if target_document is None or (source_document is not None and source_id < target_id): + resource_report["source_count"] += 1 + resource_report["missing_count"] += 1 + resource_report["create_count"] += 1 + source_document = _next_or_none(source_iterator) + elif source_document is None or target_id < source_id: + resource_report["target_count"] += 1 + target_provenance = get_search_migration_provenance(target_document) + if is_successful_migration_record(target_provenance): + resource_report["destination_only_owned_count"] += 1 + resource_report["remaining_destination_only_owned_count"] += 1 + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR: + resource_report["delete_candidate_count"] += 1 + _emit_reconciliation_candidate(report, { + "service": "ai_search", + "target_type": target_type, + "index_name": artifact["index_name"], + "document_id": target_id, + }, deletion_candidate_callback) + else: + resource_report["destination_only_unowned_count"] += 1 + target_document = _next_or_none(target_iterator) + else: + resource_report["source_count"] += 1 + resource_report["target_count"] += 1 + target_provenance = get_search_migration_provenance(target_document) + if not is_successful_migration_record(target_provenance): + resource_report["conflict_count"] += 1 + else: + resource_report["matched_count"] += 1 + if migration_record_matches_source( + target_provenance, + source_hash=_build_search_source_hash(source_document), + ): + resource_report["unchanged_count"] += 1 + else: + resource_report["stale_count"] += 1 + if migration_plan.get("migration_mode") != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resource_report["update_count"] += 1 + source_document = _next_or_none(source_iterator) + target_document = _next_or_none(target_iterator) + processed_count = resource_report["source_count"] + resource_report["target_count"] + if heartbeat_callback and processed_count % 1000 == 0: + heartbeat_callback( + f"Reconciling AI Search index {artifact['index_name']}", + processed_count, + ) + _append_reconciliation_resource(report, resource_report) + return report + + +def _source_blob_reference_exists_in_plan( + migration_plan, + container_name, + blob_name, + scope_hash, +): + target_type = _migration_blob_container_target_type(container_name) + selection = migration_plan.get(target_type) or {} + allowed_scope_hashes = _selected_scope_hashes(target_type, selection) + if allowed_scope_hashes is not None and scope_hash not in allowed_scope_hashes: + return False + container_definition = next(( + definition + for definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS.get(target_type, []) + if definition.get("documents") + ), None) + source_container = ( + getattr(app_config, container_definition["container_attr"], None) + if container_definition else None + ) + if source_container is None: + return True + query = ( + "SELECT * FROM c WHERE c.blob_path = @blob_path " + "OR c.archived_blob_path = @blob_path" + ) + for document in source_container.query_items( + query=query, + parameters=[{"name": "@blob_path", "value": blob_name}], + enable_cross_partition_query=True, + ): + if allowed_scope_hashes is None or _document_migration_scope_hash(document) == scope_hash: + return True + return False + + +def _reconcile_blob_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=None, + deletion_candidate_callback=None, + stop_event=None, + failure_holder=None, +): + """Reconcile Blob paths with bounded point reads and a destination scan.""" + report = _new_reconciliation_report("source_blobs", include_blob_fields=True) + if not migration_plan.get("include_source_blobs"): + return report + target_service = _get_target_enhanced_citations_blob_client(settings) + source_service = _get_source_blob_service_client() + resource_reports = {} + for document_count, document in enumerate( + _iter_selected_document_records_for_blob_migration(migration_plan), + start=1, + ): + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + reference = _document_blob_reference(document) + if not reference: + report["not_applicable_count"] += 1 + continue + container_name, blob_name = reference + target_type = _migration_blob_container_target_type(container_name) + resource_report = resource_reports.setdefault(container_name, { + "target_type": target_type, + "container_name": container_name, + "source_count": 0, + "target_count": 0, + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 0, + "unresolved_scope_count": 0, + "not_applicable_count": 0, + "source_missing_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "stale_count": 0, + "delete_candidate_count": 0, + }) + resource_report["source_count"] += 1 + source_blob = source_service.get_blob_client(container=container_name, blob=blob_name) + try: + source_properties = source_blob.get_blob_properties() + except ResourceNotFoundError: + resource_report["source_missing_count"] += 1 + continue + target_blob = target_service.get_blob_client(container=container_name, blob=blob_name) + target_properties = _get_blob_properties_or_none(target_blob) + if target_properties is None: + resource_report["missing_count"] += 1 + resource_report["create_count"] += 1 + continue + target_provenance = get_blob_migration_provenance( + getattr(target_properties, "metadata", None) + ) + if not is_successful_migration_record(target_provenance): + resource_report["conflict_count"] += 1 + continue + resource_report["matched_count"] += 1 + source_hash, source_version = _build_blob_source_fingerprint(source_properties) + if migration_record_matches_source( + target_provenance, + source_hash=source_hash, + source_version=source_version, + ): + resource_report["unchanged_count"] += 1 + else: + resource_report["stale_count"] += 1 + if migration_plan.get("migration_mode") != DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY: + resource_report["update_count"] += 1 + if heartbeat_callback and document_count % 250 == 0: + heartbeat_callback("Reconciling source Blob references", document_count) + + for container_name, resource_report in resource_reports.items(): + target_type = resource_report["target_type"] + selection = migration_plan.get(target_type) or {} + allowed_scope_hashes = _selected_scope_hashes(target_type, selection) + container_client = target_service.get_container_client(container_name) + try: + target_blob_items = container_client.list_blobs(include=["metadata"]) + except ResourceNotFoundError: + target_blob_items = [] + for target_count, blob_item in enumerate(target_blob_items, start=1): + if stop_event is not None: + _raise_reconciliation_heartbeat_failure(stop_event, failure_holder or {}) + metadata = ( + blob_item.get("metadata") + if isinstance(blob_item, dict) else getattr(blob_item, "metadata", None) + ) or {} + provenance = get_blob_migration_provenance(metadata) + blob_name = _safe_text( + blob_item.get("name") if isinstance(blob_item, dict) else getattr(blob_item, "name", None) + ) + if not blob_name: + continue + scope_hash = _safe_text(provenance.get("scopeHash")) + source_reference_exists = _source_blob_reference_exists_in_plan( + migration_plan, + container_name, + blob_name, + scope_hash, + ) + if is_successful_migration_record(provenance): + if allowed_scope_hashes is not None and not scope_hash: + resource_report["unresolved_scope_count"] += 1 + continue + if allowed_scope_hashes is not None and scope_hash not in allowed_scope_hashes: + continue + elif allowed_scope_hashes is not None and not source_reference_exists: + continue + resource_report["target_count"] += 1 + if source_reference_exists: + continue + if is_successful_migration_record(provenance): + resource_report["destination_only_owned_count"] += 1 + resource_report["remaining_destination_only_owned_count"] += 1 + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR: + resource_report["delete_candidate_count"] += 1 + _emit_reconciliation_candidate(report, { + "service": "source_blobs", + "target_type": target_type, + "container_name": container_name, + "blob_name": blob_name, + "target_etag": ( + blob_item.get("etag") + if isinstance(blob_item, dict) else getattr(blob_item, "etag", None) + ), + }, deletion_candidate_callback) + else: + resource_report["destination_only_unowned_count"] += 1 + if heartbeat_callback and target_count % 250 == 0: + heartbeat_callback(f"Reconciling Blob container {container_name}", target_count) + _append_reconciliation_resource(report, resource_report) + return report + + +def _source_cosmos_identity_exists(candidate): + container_definition = _get_migration_cosmos_container_definition( + candidate.get("target_type"), + candidate.get("container_name"), + ) + source_container = ( + getattr(app_config, container_definition["container_attr"], None) + if container_definition else None + ) + if source_container is None: + return True + try: + source_container.read_item( + item=candidate.get("document_id"), + partition_key=candidate.get("partition_key"), + ) + return True + except (CosmosResourceNotFoundError, ResourceNotFoundError): + return False + + +def _source_search_identity_exists(candidate): + artifact = _get_migration_search_artifact(candidate.get("target_type")) + source_client = CLIENTS.get(artifact["client_key"]) if artifact else None + document_id = _safe_text(candidate.get("document_id")) + if source_client is None or not document_id: + return True + results = source_client.search( + search_text="*", + filter=f"id eq '{_escape_search_filter_value(document_id)}'", + select=["id"], + top=1, + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + ) + return any(_safe_text(dict(result).get("id")) == document_id for result in results) + + +def _source_blob_identity_exists(candidate): + source_service = _get_source_blob_service_client() + if source_service is None: + return True + source_blob = source_service.get_blob_client( + container=candidate.get("container_name"), + blob=candidate.get("blob_name"), + ) + try: + source_blob.get_blob_properties() + return True + except ResourceNotFoundError: + return False + + +def _get_migration_cosmos_container_definition(target_type, container_name): + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS.get( + _safe_text(target_type), + [], + ): + if _target_cosmos_container_name(container_definition) == _safe_text(container_name): + return container_definition + return None + + +def _get_migration_search_artifact(target_type): + artifacts_by_type = { + "users": DATA_MANAGEMENT_SEARCH_ARTIFACTS[0], + "groups": DATA_MANAGEMENT_SEARCH_ARTIFACTS[1], + "public_workspaces": DATA_MANAGEMENT_SEARCH_ARTIFACTS[2], + } + return artifacts_by_type.get(_safe_text(target_type)) + + +def _validate_mirror_deletion_candidate(candidate, settings=None, target_database=None): + """Revalidate current source absence and destination ownership before any delete.""" + service = candidate.get("service") + if service == "cosmos": + if _source_cosmos_identity_exists(candidate): + return False, "A Cosmos deletion candidate currently exists in the source." + container_definition = _get_migration_cosmos_container_definition( + candidate.get("target_type"), + candidate.get("container_name"), + ) + if not container_definition or target_database is None: + return False, "A Cosmos deletion candidate could not resolve its resource definition." + target_container = _get_reconciliation_target_cosmos_container( + target_database, + candidate.get("container_name"), + container_definition["partition_key_path"], + read_only=True, + ) + try: + current_target = target_container.read_item( + item=candidate.get("document_id"), + partition_key=candidate.get("partition_key"), + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + current_target = None + if not current_target or not is_successful_migration_record( + get_cosmos_migration_provenance(current_target) + ): + return False, "A Cosmos deletion candidate lost migration ownership." + if not candidate.get("target_etag") or current_target.get("_etag") != candidate.get("target_etag"): + return False, "A Cosmos deletion candidate changed after reconciliation." + candidate["target_container"] = target_container + candidate["container_definition"] = container_definition + candidate["target_document"] = current_target + return True, "" + if service == "ai_search": + if _source_search_identity_exists(candidate): + return False, "An AI Search deletion candidate currently exists in the source." + artifact = _get_migration_search_artifact(candidate.get("target_type")) + if not artifact: + return False, "An AI Search deletion candidate could not resolve its index definition." + target_client = _get_target_search_client(settings, artifact["index_name"]) + target_documents = _get_target_search_documents_by_ids( + target_client, + [candidate.get("document_id")], + ) + current_target = target_documents.get(_safe_text(candidate.get("document_id"))) + if not current_target or not is_successful_migration_record( + get_search_migration_provenance(current_target) + ): + return False, "An AI Search deletion candidate lost migration ownership." + candidate["target_client"] = target_client + candidate["index_name"] = artifact["index_name"] + return True, "" + if service == "source_blobs": + if _source_blob_identity_exists(candidate): + return False, "A Blob deletion candidate currently exists in the source." + target_service = _get_target_enhanced_citations_blob_client(settings) + target_blob = target_service.get_blob_client( + container=candidate.get("container_name"), + blob=candidate.get("blob_name"), + ) + target_properties = _get_blob_properties_or_none(target_blob) + if not target_properties or not is_successful_migration_record( + get_blob_migration_provenance(getattr(target_properties, "metadata", None)) + ): + return False, "A Blob deletion candidate lost migration ownership." + current_etag = _safe_text(getattr(target_properties, "etag", None)) + if not candidate.get("target_etag") or current_etag != _safe_text(candidate.get("target_etag")): + return False, "A Blob deletion candidate changed after reconciliation." + candidate["target_blob"] = target_blob + candidate["target_etag"] = current_etag + return True, "" + return False, "An unsupported mirror deletion candidate was produced." + + +def _validate_mirror_deletion_candidates( + candidates, + heartbeat_callback=None, + settings=None, + target_database=None, +): + validated_count = 0 + blockers = set() + last_heartbeat_at = 0.0 + for index, candidate in enumerate(candidates, start=1): + now = time.monotonic() + if heartbeat_callback and ( + index == 1 or + index % 100 == 0 or + now - last_heartbeat_at >= DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + heartbeat_callback("Revalidating mirror deletion candidates", index) + last_heartbeat_at = now + is_valid, blocker = _validate_mirror_deletion_candidate( + candidate, + settings=settings, + target_database=target_database, + ) + if is_valid: + validated_count += 1 + elif len(blockers) < 20: + blockers.add(blocker) + return validated_count, sorted(blockers) + + +def _apply_mirror_deletion_candidates( + candidates, + heartbeat_callback=None, + settings=None, + target_database=None, + manifest_append=None, + manifest_flush=None, +): + deleted_counts = {"cosmos": 0, "ai_search": 0, "source_blobs": 0} + search_batch = [] + search_batch_key = None + + def record_deleted(candidate): + if not callable(manifest_append): + return + manifest_append({ + "service": candidate.get("service"), + "resource_name": "reconciliation:cutover", + "target_type": candidate.get("target_type"), + "source_identity": _hash_migration_manifest_identity( + "deleted-source", + candidate.get("service"), + candidate.get("container_name") or candidate.get("index_name"), + candidate.get("document_id") or candidate.get("blob_name"), + ), + "destination_identity": _hash_migration_manifest_identity( + "deleted-destination", + candidate.get("service"), + candidate.get("container_name") or candidate.get("index_name"), + candidate.get("document_id") or candidate.get("blob_name"), + ), + "status": "deleted", + "_locator": { + "service": candidate.get("service"), + "resource_name": "reconciliation:cutover", + "target_type": candidate.get("target_type"), + "document_id": candidate.get("document_id"), + "partition_key": candidate.get("partition_key"), + "container_name": candidate.get("container_name"), + "blob_name": candidate.get("blob_name"), + "index_name": candidate.get("index_name"), + }, + }) + if callable(manifest_flush): + manifest_flush() + + def flush_search_batch(): + nonlocal search_batch, search_batch_key + if not search_batch: + return + target_client = search_batch[0].get("target_client") + last_heartbeat_at = 0.0 + for index, candidate in enumerate(search_batch, start=1): + now = time.monotonic() + if heartbeat_callback and ( + index == 1 or + index % 100 == 0 or + now - last_heartbeat_at >= + DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + heartbeat_callback( + "Revalidating mirror AI Search deletions", + sum(deleted_counts.values()), + ) + last_heartbeat_at = now + is_valid, blocker = _validate_mirror_deletion_candidate( + candidate, + settings=settings, + target_database=target_database, + ) + if not is_valid: + raise DataManagementSettingsValidationError(blocker) + if heartbeat_callback: + heartbeat_callback( + "Applying mirror AI Search deletions", + sum(deleted_counts.values()), + ) + results = list(target_client.delete_documents( + documents=[{"id": candidate.get("document_id")} for candidate in search_batch], + connection_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + read_timeout=DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS, + retry_total=0, + )) + successful_count = 0 + failure_count = 0 + for candidate, delete_result in zip(search_batch, results): + if _get_search_result_succeeded(delete_result): + successful_count += 1 + record_deleted(candidate) + else: + failure_count += 1 + failure_count += max(0, len(search_batch) - len(results)) + deleted_counts["ai_search"] += successful_count + if failure_count: + raise RuntimeError( + f"AI Search mirror reconciliation failed to delete {failure_count} document(s)." + ) + search_batch = [] + search_batch_key = None + + for candidate in candidates: + is_valid, blocker = _validate_mirror_deletion_candidate( + candidate, + settings=settings, + target_database=target_database, + ) + if not is_valid: + raise DataManagementSettingsValidationError(blocker) + if heartbeat_callback: + heartbeat_callback( + "Applying mirror deletions", + sum(deleted_counts.values()), + ) + if candidate.get("service") == "ai_search": + candidate_key = candidate.get("index_name") + if search_batch and ( + candidate_key != search_batch_key or + len(search_batch) >= DATA_MANAGEMENT_MIGRATION_BATCH_SIZE + ): + flush_search_batch() + search_batch_key = candidate_key + search_batch.append(candidate) + continue + flush_search_batch() + if candidate.get("service") == "cosmos": + _delete_target_cosmos_record( + candidate.get("target_container"), + candidate.get("target_document"), + (candidate.get("container_definition") or {}).get("partition_key_path"), + ) + deleted_counts["cosmos"] += 1 + record_deleted(candidate) + elif candidate.get("service") == "source_blobs": + candidate.get("target_blob").delete_blob( + etag=candidate.get("target_etag"), + match_condition=MatchConditions.IfNotModified, + ) + deleted_counts["source_blobs"] += 1 + record_deleted(candidate) + flush_search_batch() + return deleted_counts + + +def _public_reconciliation_report(report): + return { + key: value + for key, value in (report or {}).items() + if not str(key).startswith("_") + } + + +def _run_data_management_migration_reconciliation( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_database, + preview_snapshot=None, + migration_artifacts=None, +): + """Persist final source/destination reconciliation and mirror deletion outcomes.""" + resource_name = "reconciliation:cutover" + completed_resource = migration_state.get("resources", {}).get(resource_name) + if is_migration_resource_completed(migration_state, resource_name): + return completed_resource.get("result") if isinstance(completed_resource, dict) else {} + start_migration_resource(migration_state, resource_name) + deletion_plan_id = str(uuid.uuid4()) + append_deletion_candidate, flush_deletion_candidates = ( + _create_mirror_deletion_candidate_writer(job.get("id"), deletion_plan_id) + ) + append_reconciliation_manifest, flush_reconciliation_manifest = ( + _create_migration_manifest_writer(job.get("id"), resource_name) + ) + reconciliation_stop_event = Event() + reconciliation_failure_holder = {} + reconciliation_state_holder = {"state": migration_state} + reconciliation_persistence_lock = Lock() + reconciliation_heartbeat_thread = _start_reconciliation_heartbeat_thread( + job, + settings, + reconciliation_state_holder, + reconciliation_stop_event, + reconciliation_failure_holder, + reconciliation_persistence_lock, + ) + last_heartbeat_at = 0.0 + heartbeat_count = 0 + aggregate_fields = ( + "matched_count", + "missing_count", + "destination_only_owned_count", + "remaining_destination_only_owned_count", + "destination_only_unowned_count", + "conflict_count", + "deleted_count", + "unresolved_scope_count", + "create_count", + "update_count", + "unchanged_count", + "stale_count", + "delete_candidate_count", + "not_applicable_count", + "source_missing_count", + ) + service_reports = [] + + def stop_reconciliation_heartbeat(): + reconciliation_stop_event.set() + reconciliation_heartbeat_thread.join(timeout=5.0) + + def build_partial_result(error=None): + partial_result = { + field_name: sum( + _safe_int(report.get(field_name), default=0, minimum=0) + for report in service_reports + ) + for field_name in aggregate_fields + } + partial_result.update({ + "name": "migration_reconciliation", + "type": "migration_reconciliation", + "migration_mode": migration_plan.get("migration_mode"), + "readiness": "not_ready", + "services_completed": len(service_reports), + "services_total": 3, + "services": [ + _public_reconciliation_report(report) + for report in service_reports + ], + }) + if error: + partial_result["error"] = str(error)[:1000] + return partial_result + + def fail_reconciliation(error, failure_result=None): + nonlocal migration_state + stop_reconciliation_heartbeat() + migration_state = reconciliation_state_holder["state"] + persisted_result = copy.deepcopy( + failure_result if isinstance(failure_result, dict) else build_partial_result(error) + ) + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = persisted_result + fail_migration_resource( + migration_state, + resource_name, + str(error)[:1000], + ) + job["migration_state"] = migration_state + cleanup_errors = [] + for cleanup_name, cleanup in ( + ("deletion_plan", flush_deletion_candidates), + ("manifest", flush_reconciliation_manifest), + ): + try: + cleanup() + except Exception as cleanup_exc: + cleanup_errors.append(f"{cleanup_name}: {str(cleanup_exc)[:300]}") + if cleanup_errors: + persisted_result["persistence_warnings"] = cleanup_errors + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = persisted_result + try: + _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + str(error)[:1000], + ) + except Exception as checkpoint_exc: + persisted_result.setdefault("persistence_warnings", []).append( + f"checkpoint: {str(checkpoint_exc)[:300]}" + ) + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = persisted_result + resource["status"] = MIGRATION_RESOURCE_STATUS_FAILED + resource["last_error"] = str(error)[:1000] + job["migration_state"] = migration_state + + def persist_reconciliation_checkpoint(progress, message): + nonlocal migration_state + with reconciliation_persistence_lock: + migration_state = reconciliation_state_holder["state"] + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + progress, + message, + ) + reconciliation_state_holder["state"] = migration_state + return migration_state + + def heartbeat(message, completed_count=0, force=False): + nonlocal migration_state, last_heartbeat_at, heartbeat_count + _assert_migration_job_lease(job) + heartbeat_count += 1 + now = time.monotonic() + if not force and now - last_heartbeat_at < 2.0: + return + _raise_reconciliation_heartbeat_failure( + reconciliation_stop_event, + reconciliation_failure_holder, + ) + with reconciliation_persistence_lock: + migration_state = reconciliation_state_holder["state"] + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + resource_name, + { + "state": "reconciling", + "heartbeat_count": heartbeat_count, + "completed_count": completed_count, + }, + message, + ) + reconciliation_state_holder["state"] = migration_state + last_heartbeat_at = now + + try: + service_reports.append(_reconcile_cosmos_migration( + target_database, + migration_plan, + migration_state, + provenance_context, + apply_deletions=False, + heartbeat_callback=heartbeat, + deletion_candidate_callback=append_deletion_candidate, + stop_event=reconciliation_stop_event, + failure_holder=reconciliation_failure_holder, + )) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + flush_deletion_candidates() + persist_reconciliation_checkpoint( + {"services_completed": 1, "services_total": 3}, + "Reconciling migrated Cosmos records", + ) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + service_reports.append(_reconcile_ai_search_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat, + deletion_candidate_callback=append_deletion_candidate, + stop_event=reconciliation_stop_event, + failure_holder=reconciliation_failure_holder, + )) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + flush_deletion_candidates() + persist_reconciliation_checkpoint( + {"services_completed": 2, "services_total": 3}, + "Reconciling migrated AI Search documents", + ) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + service_reports.append(_reconcile_blob_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat, + deletion_candidate_callback=append_deletion_candidate, + stop_event=reconciliation_stop_event, + failure_holder=reconciliation_failure_holder, + )) + except Exception as exc: + fail_reconciliation(exc) + raise + try: + flush_deletion_candidates() + heartbeat("Completed read-only migration reconciliation", 3, force=True) + except Exception as exc: + fail_reconciliation(exc) + raise + result = { + field_name: sum(_safe_int(report.get(field_name), default=0, minimum=0) for report in service_reports) + for field_name in aggregate_fields + } + result.update({ + "name": "migration_reconciliation", + "type": "migration_reconciliation", + "migration_mode": migration_plan.get("migration_mode"), + "deletion_policy": ( + "migration_owned_only" + if migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR else + "retain_destination_only" + ), + "services": [_public_reconciliation_report(report) for report in service_reports], + }) + actual_outcomes = { + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "delete_count": result.get("delete_candidate_count", 0), + "not_applicable_count": 0, + "missing_count": 0, + "conflict_count": 0, + "failed_count": 0, + } + for artifact in migration_artifacts or []: + if artifact.get("type") not in { + "cosmos_container", + "ai_search_documents", + "source_blobs", + }: + continue + actual_outcomes["create_count"] += _safe_int( + artifact.get("created_count"), + default=0, + minimum=0, + ) + actual_outcomes["update_count"] += _safe_int( + artifact.get("updated_count"), + default=0, + minimum=0, + ) + actual_outcomes["unchanged_count"] += _safe_int( + artifact.get("unchanged_count"), + default=0, + minimum=0, + ) + actual_outcomes["not_applicable_count"] += _safe_int( + artifact.get("not_applicable_count"), + default=0, + minimum=0, + ) + actual_outcomes["missing_count"] += _safe_int( + artifact.get("missing_count"), + default=0, + minimum=0, + ) + actual_outcomes["conflict_count"] += _safe_int( + artifact.get("collision_count"), + default=0, + minimum=0, + ) + actual_outcomes["failed_count"] += _safe_int( + artifact.get("failed_count"), + default=0, + minimum=0, + ) + preview_counts = ( + preview_snapshot.get("estimated_outcomes") + if isinstance(preview_snapshot, dict) and + isinstance(preview_snapshot.get("estimated_outcomes"), dict) + else {} + ) + result["preview_snapshot"] = _sanitize_activity_value( + preview_snapshot if isinstance(preview_snapshot, dict) else {} + ) + result["actual_outcomes"] = actual_outcomes + result["preview_actual_divergence"] = { + field_name: actual_outcomes.get(field_name, 0) - _safe_int( + preview_counts.get(field_name), + default=0, + ) + for field_name in actual_outcomes + } + critical_mismatch_count = sum(( + result["missing_count"], + result["conflict_count"], + result["stale_count"], + result["source_missing_count"], + result["unresolved_scope_count"], + actual_outcomes["failed_count"], + )) + if critical_mismatch_count: + result["readiness"] = "not_ready" + elif ( + result["remaining_destination_only_owned_count"] or + result["destination_only_unowned_count"] or + result["unresolved_scope_count"] + ): + result["readiness"] = "ready_with_warnings" + else: + result["readiness"] = "ready" + + is_mirror = ( + migration_plan.get("migration_mode") == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR and + migration_plan.get("mirror_deletions_confirmed") is True + ) + if is_mirror: + preview_divergence = result["preview_actual_divergence"] + deletion_blockers = [] + if not preview_counts: + deletion_blockers.append("Mirror deletion requires a server-owned inventory preview.") + if any(value != 0 for value in preview_divergence.values()): + deletion_blockers.append("Migration outcomes diverged from the queued preview.") + if critical_mismatch_count: + deletion_blockers.append("Reconciliation found missing, stale, failed, or conflicting items.") + if result["destination_only_unowned_count"]: + deletion_blockers.append("Unowned destination-only items prevent an exact mirror.") + candidate_iterator = _iter_mirror_deletion_candidates( + job.get("id"), + deletion_plan_id, + ) + try: + validated_candidate_count, candidate_blockers = _validate_mirror_deletion_candidates( + candidate_iterator, + heartbeat_callback=heartbeat, + settings=settings, + target_database=target_database, + ) + except Exception as exc: + failure_result = copy.deepcopy(result) + failure_result.update({ + "readiness": "not_ready", + "error": str(exc)[:1000], + }) + fail_reconciliation(exc, failure_result=failure_result) + raise + deletion_blockers.extend(candidate_blockers) + if validated_candidate_count != result["delete_candidate_count"]: + deletion_blockers.append("Mirror deletion candidate count changed during validation.") + if deletion_blockers: + result["readiness"] = "not_ready" + result["deletion_status"] = "blocked" + result["deletion_blockers"] = sorted(set(deletion_blockers))[:20] + else: + target_search_write_gate_container = None + target_search_write_fence = None + requires_search_write_fence = any( + report.get("service") == "ai_search" and + _safe_int(report.get("delete_candidate_count"), default=0) > 0 + for report in service_reports + ) + + last_target_search_fence_renewal = 0.0 + + def mirror_deletion_heartbeat(message, completed_count=0): + nonlocal last_target_search_fence_renewal + heartbeat(message, completed_count) + if ( + target_search_write_gate_container is not None and + target_search_write_fence is not None and + time.monotonic() - last_target_search_fence_renewal >= + DATA_MANAGEMENT_MIGRATION_HEARTBEAT_INTERVAL_SECONDS + ): + renew_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + _get_target_search_write_fence_lease_seconds(settings), + ) + last_target_search_fence_renewal = time.monotonic() + + try: + if requires_search_write_fence: + target_search_write_gate_container = _get_target_data_management_search_write_gate_container( + target_database + ) + target_search_write_fence = acquire_data_management_search_write_fence( + target_search_write_gate_container, + job.get("id"), + _get_target_search_write_fence_lease_seconds(settings), + heartbeat_callback=lambda: heartbeat( + "Waiting for target AI Search writers to drain before mirror deletion", + result.get("deleted_count", 0), + ), + ) + deletion_counts = _apply_mirror_deletion_candidates( + _iter_mirror_deletion_candidates(job.get("id"), deletion_plan_id), + heartbeat_callback=mirror_deletion_heartbeat, + settings=settings, + target_database=target_database, + manifest_append=append_reconciliation_manifest, + manifest_flush=flush_reconciliation_manifest, + ) + if target_search_write_gate_container is not None: + release_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + ) + except DataManagementMigrationCanceledError: + if target_search_write_gate_container is not None: + release_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + ) + raise + except Exception as exc: + failure_result = copy.deepcopy(result) + failure_result.update({ + "readiness": "not_ready", + "error": str(exc)[:1000], + }) + fail_reconciliation(exc, failure_result=failure_result) + raise + flush_reconciliation_manifest() + result["deleted_count"] = sum(deletion_counts.values()) + result["remaining_destination_only_owned_count"] = max( + 0, + result["destination_only_owned_count"] - result["deleted_count"], + ) + result["actual_outcomes"]["delete_count"] = result["deleted_count"] + result["preview_actual_divergence"]["delete_count"] = ( + result["deleted_count"] - _safe_int(preview_counts.get("delete_count"), default=0) + ) + result["deletion_status"] = "completed" + result["deletion_counts"] = deletion_counts + result["readiness"] = ( + "ready" + if not result["preview_actual_divergence"]["delete_count"] else + "not_ready" + ) + heartbeat("Completed guarded mirror deletions", result["deleted_count"], force=True) + else: + result["actual_outcomes"]["delete_count"] = 0 + result["preview_actual_divergence"]["delete_count"] = ( + 0 - _safe_int(preview_counts.get("delete_count"), default=0) + ) + if result["readiness"] == "not_ready": + stop_reconciliation_heartbeat() + flush_reconciliation_manifest() + migration_state = reconciliation_state_holder["state"] + resource = (migration_state.get("resources") or {}).get(resource_name) + if isinstance(resource, dict): + resource["result"] = copy.deepcopy(result) + _fail_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + "Migration reconciliation is not ready for cutover.", + ) + raise RuntimeError("Migration reconciliation is not ready for cutover.") + stop_reconciliation_heartbeat() + migration_state = reconciliation_state_holder["state"] + _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + resource_name, + result, + "Completed migration source/destination reconciliation", + ) + return result + + +def _migration_selection_summary(migration_plan): + return { + target_type: { + "mode": (migration_plan.get(target_type) or {}).get("mode"), + "selected_count": len((migration_plan.get(target_type) or {}).get("ids") or []), + "include_documents": bool((migration_plan.get(target_type) or {}).get("include_documents")), + } + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES + } + + +def _job_lease_holder_id(): + return f"{socket.gethostname()}:{uuid.uuid4().hex}" + + +def _is_stale_job(job, stale_seconds=DATA_MANAGEMENT_DEFAULT_STALE_SECONDS): + last_heartbeat = _parse_iso_datetime((job or {}).get("last_heartbeat_at") or (job or {}).get("updated_at")) + if not last_heartbeat: + return True + return last_heartbeat <= _now_utc() - timedelta(seconds=stale_seconds) + + +def _read_job(job_id): + return cosmos_data_management_jobs_container.read_item(item=job_id, partition_key=job_id) + + +def _get_migration_destination_lock_id(settings, migration_plan): + """Use one global coordinator so partially overlapping migrations cannot race.""" + return f"{DATA_MANAGEMENT_MIGRATION_LOCK_PREFIX}_global" + + +def _get_migration_lock_lease_seconds(settings): + configured_lease_seconds = _safe_int( + (settings or {}).get("data_management_job_lease_seconds"), + default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, + minimum=60, + maximum=7200, + ) + minimum_quarantine_seconds = ( + DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + + DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS + ) + maximum_recovery_safe_seconds = ( + DATA_MANAGEMENT_DEFAULT_STALE_SECONDS - + DATA_MANAGEMENT_RECOVERY_RESUBMIT_DELAY_SECONDS + ) + return min( + max(configured_lease_seconds, minimum_quarantine_seconds), + maximum_recovery_safe_seconds, + ) + + +def _has_active_migration_destination_lock(job): + """Return whether a coordinator lock remains active and clear only conclusively stale metadata.""" + lock = job.get("migration_coordinator_lock") if isinstance(job, dict) else None + if not isinstance(lock, dict) or not lock.get("id") or not lock.get("lock_token"): + if isinstance(job, dict): + job.pop("migration_coordinator_lock", None) + return False + try: + current_lock = _read_job(lock["id"]) + except (CosmosResourceNotFoundError, ResourceNotFoundError, KeyError): + job.pop("migration_coordinator_lock", None) + return False + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease could not be read while claiming a retry." + ) from exc + expires_at = _parse_iso_datetime(current_lock.get("expires_at")) + lock_is_active = bool( + current_lock.get("type") == DATA_MANAGEMENT_MIGRATION_LOCK_TYPE and + expires_at and + expires_at > _now_utc() + ) + if not lock_is_active: + job.pop("migration_coordinator_lock", None) + return lock_is_active + + +def _acquire_migration_destination_lock(job, settings, migration_plan): + """Acquire a destination-scoped lease so independent migration jobs cannot overlap.""" + if not isinstance(job, dict) or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + return None + lock_id = _get_migration_destination_lock_id(settings, migration_plan) + now = _now_utc() + lease_seconds = _get_migration_lock_lease_seconds(settings) + lock_document = { + "id": lock_id, + "type": DATA_MANAGEMENT_MIGRATION_LOCK_TYPE, + "migration_job_id": job.get("id"), + "lock_token": str(uuid.uuid4()), + "acquired_at": now.isoformat(), + "updated_at": now.isoformat(), + "lease_seconds": lease_seconds, + "expires_at": (now + timedelta(seconds=lease_seconds)).isoformat(), + } + try: + cosmos_data_management_jobs_container.create_item(body=lock_document) + except Exception as exc: + if getattr(exc, "status_code", None) != 409: + raise + existing_lock = _read_job(lock_id) + expires_at = _parse_iso_datetime(existing_lock.get("expires_at")) + if expires_at and expires_at > now: + raise DataManagementSettingsValidationError( + "Another Data Management migration is active for this destination. Wait for it to complete or expire before retrying." + ) + replacement_lock = dict(existing_lock) + replacement_lock.update(lock_document) + try: + cosmos_data_management_jobs_container.replace_item( + item=lock_id, + body=replacement_lock, + etag=existing_lock.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as replace_exc: + if getattr(replace_exc, "status_code", None) in {409, 412}: + raise DataManagementSettingsValidationError( + "Another Data Management migration acquired this destination before the retry could start." + ) from replace_exc + raise + + job["migration_coordinator_lock"] = { + "id": lock_id, + "lock_token": lock_document["lock_token"], + "lease_seconds": lease_seconds, + "expires_at": lock_document["expires_at"], + } + return job["migration_coordinator_lock"] + + +def _renew_migration_destination_lock(job, settings=None): + """Renew and verify the exclusive destination lease held by this worker.""" + lock = job.get("migration_coordinator_lock") if isinstance(job, dict) else None + if not isinstance(lock, dict) or not lock.get("id"): + return + now = _now_utc() + try: + current_lock = _read_job(lock["id"]) + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease could not be read." + ) from exc + if ( + current_lock.get("type") != DATA_MANAGEMENT_MIGRATION_LOCK_TYPE or + current_lock.get("migration_job_id") != job.get("id") or + current_lock.get("lock_token") != lock.get("lock_token") + ): + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease was superseded." + ) + expires_at = _parse_iso_datetime(current_lock.get("expires_at")) + if expires_at is None or expires_at <= now: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease expired." + ) + lease_seconds = _get_migration_lock_lease_seconds({ + "data_management_job_lease_seconds": ( + (settings or {}).get("data_management_job_lease_seconds") or + lock.get("lease_seconds") + ), + }) + current_lock["updated_at"] = now.isoformat() + current_lock["expires_at"] = ( + now + timedelta(seconds=lease_seconds) + ).isoformat() + try: + renewed_lock = cosmos_data_management_jobs_container.replace_item( + item=current_lock.get("id"), + body=current_lock, + etag=current_lock.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if getattr(exc, "status_code", None) in {409, 412}: + raise DataManagementMigrationLeaseLostError( + "Migration destination coordinator lease changed during renewal." + ) from exc + raise + lock["expires_at"] = renewed_lock.get("expires_at", current_lock["expires_at"]) + lock["lease_seconds"] = lease_seconds + + +def _release_migration_destination_lock(job): + """Release only the destination coordinator lease token owned by this job.""" + lock = job.get("migration_coordinator_lock") if isinstance(job, dict) else None + if not isinstance(lock, dict) or not lock.get("id"): + return + try: + current_lock = _read_job(lock["id"]) + except Exception: + return + if ( + current_lock.get("type") != DATA_MANAGEMENT_MIGRATION_LOCK_TYPE or + current_lock.get("migration_job_id") != job.get("id") or + current_lock.get("lock_token") != lock.get("lock_token") + ): + return + try: + cosmos_data_management_jobs_container.delete_item( + item=current_lock.get("id"), + partition_key=current_lock.get("id"), + etag=current_lock.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception: + return + + +def _assert_migration_job_lease(job, allow_cancel_requested=False): + """Stop a stale migration worker before it schedules another remote action.""" + if not isinstance(job, dict) or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + return + lease_holder_id = _safe_text(job.get("lease_holder_id")) + if not lease_holder_id: + return + try: + persisted_job = _read_job(job.get("id")) + except Exception as exc: + raise DataManagementMigrationLeaseLostError( + "Migration worker could not verify its durable job lease." + ) from exc + + persisted_expiry = _parse_iso_datetime(persisted_job.get("lease_expires_at")) + if ( + persisted_job.get("status") != DATA_MANAGEMENT_STATUS_RUNNING or + _safe_text(persisted_job.get("lease_holder_id")) != lease_holder_id or + persisted_expiry is None or + persisted_expiry <= _now_utc() + ): + raise DataManagementMigrationLeaseLostError( + "Migration worker lease was superseded or expired." + ) + if persisted_job.get("cancel_requested_at") and not allow_cancel_requested: + raise DataManagementMigrationCanceledError( + "Migration cancellation was requested by an administrator." + ) + if persisted_job.get("_etag"): + job["_etag"] = persisted_job.get("_etag") + _renew_migration_destination_lock(job) + _renew_target_migration_coordinator(job, None) + + +def _replace_job(job): + return cosmos_data_management_jobs_container.replace_item( + item=job.get("id"), + body=job, + etag=job.get("_etag"), + match_condition=MatchConditions.IfNotModified, ) @@ -1196,6 +6535,10 @@ def queue_data_management_job(operation, backup_type=None, requested_by=None, re "last_heartbeat_at": None, "last_message": "Queued data management job", "last_error": None, + "cancel_requested_at": None, + "cancel_requested_by": None, + "cancel_requested_by_email": None, + "cancel_reason": None, "requested_by": _safe_text(requested_by), "requested_by_email": _safe_text(requested_by_email), "scheduled": bool(scheduled), @@ -1241,9 +6584,164 @@ def get_data_management_jobs(limit=DATA_MANAGEMENT_DEFAULT_JOB_LIMIT): ))[:safe_limit] +def get_recoverable_data_management_migration_jobs( + current_time=None, + limit=DATA_MANAGEMENT_DEFAULT_RECOVERY_JOB_LIMIT, +): + """Return queued or stale migration jobs that can safely be resubmitted.""" + safe_limit = _safe_int( + limit, + default=DATA_MANAGEMENT_DEFAULT_RECOVERY_JOB_LIMIT, + minimum=1, + maximum=100, + ) + now = current_time if isinstance(current_time, datetime) else _now_utc() + query = ( + "SELECT * FROM c WHERE c.type = @type AND c.operation = @operation " + "AND (c.status = @queued OR c.status = @running) ORDER BY c.updated_at ASC" + ) + parameters = [ + {"name": "@type", "value": DATA_MANAGEMENT_JOB_TYPE}, + {"name": "@operation", "value": DATA_MANAGEMENT_OPERATION_MIGRATION}, + {"name": "@queued", "value": DATA_MANAGEMENT_STATUS_QUEUED}, + {"name": "@running", "value": DATA_MANAGEMENT_STATUS_RUNNING}, + ] + candidates = list(cosmos_data_management_jobs_container.query_items( + query=query, + parameters=parameters, + enable_cross_partition_query=True, + max_item_count=safe_limit, + ))[:safe_limit] + recoverable_jobs = [] + for job in candidates: + status = job.get("status") + queued_since = _parse_iso_datetime(job.get("updated_at") or job.get("created_at")) + last_recovery = _parse_iso_datetime(job.get("last_recovery_submitted_at")) + resubmit_due = ( + last_recovery is None or + last_recovery <= now - timedelta(seconds=DATA_MANAGEMENT_RECOVERY_RESUBMIT_DELAY_SECONDS) + ) + if status == DATA_MANAGEMENT_STATUS_QUEUED: + is_queued_long_enough = ( + queued_since is None or + queued_since <= now - timedelta(seconds=DATA_MANAGEMENT_RECOVERY_QUEUE_DELAY_SECONDS) + ) + if is_queued_long_enough and resubmit_due: + recoverable_jobs.append({"job": job, "reason": "queued_recovery"}) + elif status == DATA_MANAGEMENT_STATUS_RUNNING and _is_stale_job(job) and resubmit_due: + recoverable_jobs.append({"job": job, "reason": "stale_recovery"}) + return recoverable_jobs + + +def recover_data_management_migration_jobs(app=None, settings=None, current_time=None): + """Resubmit recoverable migration jobs through the executor, never inline.""" + now = current_time if isinstance(current_time, datetime) else _now_utc() + recovery_results = [] + for recovery in get_recoverable_data_management_migration_jobs(current_time=now): + job = recovery["job"] + reason = recovery["reason"] + job.update({ + "last_recovery_submitted_at": now.isoformat(), + "recovery_attempt_count": _safe_int(job.get("recovery_attempt_count"), default=0, minimum=0) + 1, + "updated_at": now.isoformat(), + "last_message": ( + "Detected a stale migration worker; scheduling durable recovery." + if reason == "stale_recovery" else + "Detected a delayed queued migration; scheduling executor recovery." + ), + }) + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + "migration-stall-detected" if reason == "stale_recovery" else "migration-queued-recovery", + saved_job, + status=saved_job.get("status"), + message=saved_job.get("last_message"), + details={ + "recovery_attempt_count": saved_job.get("recovery_attempt_count"), + "cancel_requested": bool(saved_job.get("cancel_requested_at")), + }, + ) + submitted = submit_data_management_job(app, saved_job.get("id")) + if submitted: + _record_data_management_job_event( + saved_job.get("id"), + "migration-recovery-submitted", + saved_job, + status=saved_job.get("status"), + message="Submitted migration recovery to the executor.", + details={"reason": reason}, + ) + else: + log_event( + "[DataManagement] Migration recovery could not submit to the executor.", + {"job_id": saved_job.get("id"), "reason": reason}, + level=logging.WARNING, + ) + recovery_results.append({ + "job_id": saved_job.get("id"), + "reason": reason, + "submitted": submitted, + }) + return recovery_results + + +def _sanitize_data_management_migration_state_for_admin(state): + if not isinstance(state, dict): + return {} + resources = {} + for resource_name, resource in (state.get("resources") or {}).items(): + if not isinstance(resource, dict): + continue + resources[_safe_text(resource_name)] = { + "status": resource.get("status"), + "attempts": resource.get("attempts"), + "started_at": resource.get("started_at"), + "attempt_started_at": resource.get("attempt_started_at"), + "updated_at": resource.get("updated_at"), + "completed_at": resource.get("completed_at"), + "last_error": _safe_text(resource.get("last_error"))[:500], + "progress": _sanitize_activity_value(resource.get("progress") if isinstance(resource.get("progress"), dict) else {}), + "result": _sanitize_activity_value(resource.get("result") if isinstance(resource.get("result"), dict) else {}), + } + capacity = state.get("capacity") if isinstance(state.get("capacity"), dict) else {} + return { + "schema_version": state.get("schema_version"), + "migration_id": state.get("migration_id"), + "status": state.get("status"), + "created_at": state.get("created_at"), + "updated_at": state.get("updated_at"), + "completed_at": state.get("completed_at"), + "source_cutoff_at": state.get("source_cutoff_at"), + "last_progress_at": state.get("last_progress_at"), + "resume_count": state.get("resume_count"), + "last_error": _safe_text(state.get("last_error"))[:500], + "totals": _sanitize_activity_value(state.get("totals") if isinstance(state.get("totals"), dict) else {}), + "preflight": _sanitize_activity_value(state.get("preflight") if isinstance(state.get("preflight"), dict) else {}), + "capacity": { + "status": capacity.get("status"), + "restore_pending": bool(capacity.get("restore_pending")), + "target_ru": capacity.get("target_ru"), + "started_at": capacity.get("started_at"), + "completed_at": capacity.get("completed_at"), + "restored_at": capacity.get("restored_at"), + "restore_warnings": _sanitize_activity_value( + capacity.get("restore_warnings") if isinstance(capacity.get("restore_warnings"), list) else [] + ), + "targets": _sanitize_activity_value(capacity.get("targets") if isinstance(capacity.get("targets"), list) else []), + }, + "resources": resources, + } + + def sanitize_data_management_job_for_admin(job): if not isinstance(job, dict): return None + result = _sanitize_activity_value(job.get("result") if isinstance(job.get("result"), dict) else {}) + if isinstance(result, dict) and "migration_state" in result: + result["migration_state"] = _sanitize_data_management_migration_state_for_admin( + job.get("migration_state") + ) return { "id": job.get("id"), "operation": job.get("operation"), @@ -1254,13 +6752,28 @@ def sanitize_data_management_job_for_admin(job): "started_at": job.get("started_at"), "completed_at": job.get("completed_at"), "last_heartbeat_at": job.get("last_heartbeat_at"), + "last_progress_at": job.get("last_progress_at"), "last_message": job.get("last_message"), "last_error": job.get("last_error"), + "cancel_requested_at": job.get("cancel_requested_at"), "requested_by_email": job.get("requested_by_email"), "scheduled": bool(job.get("scheduled")), "progress": job.get("progress") if isinstance(job.get("progress"), dict) else {}, "warnings": job.get("warnings") if isinstance(job.get("warnings"), list) else [], - "result": job.get("result") if isinstance(job.get("result"), dict) else {}, + "result": result, + "migration_state": _sanitize_data_management_migration_state_for_admin(job.get("migration_state")), + "can_retry": bool( + job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION and + ( + job.get("status") in {DATA_MANAGEMENT_STATUS_FAILED, DATA_MANAGEMENT_STATUS_CANCELED} or + (job.get("status") == DATA_MANAGEMENT_STATUS_RUNNING and _is_stale_job(job)) + ) + ), + "can_cancel": bool( + job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION and + job.get("status") in {DATA_MANAGEMENT_STATUS_QUEUED, DATA_MANAGEMENT_STATUS_RUNNING} and + not job.get("cancel_requested_at") + ), } @@ -1284,9 +6797,173 @@ def get_data_management_job(job_id): if not safe_job_id: return None try: - return _read_job(safe_job_id) + job = _read_job(safe_job_id) except CosmosResourceNotFoundError: return None + return job if job.get("type") == DATA_MANAGEMENT_JOB_TYPE else None + + +def request_data_management_migration_cancellation( + job_id, + requested_by=None, + requested_by_email=None, + reason="", +): + """Request cooperative cancellation without invalidating durable checkpoints.""" + job = get_data_management_job(job_id) + if not job: + raise DataManagementSettingsValidationError("Data Management job was not found.") + if job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Only migration jobs can be canceled from migration checkpoints.") + if job.get("status") == DATA_MANAGEMENT_STATUS_CANCELED: + return job + if job.get("status") in DATA_MANAGEMENT_TERMINAL_STATUSES: + raise DataManagementSettingsValidationError("Only queued or running migration jobs can be canceled.") + + now = _now_iso() + safe_reason = _safe_text(reason)[:500] + migration_state = job.get("migration_state") + if isinstance(migration_state, dict): + migration_state.update({ + "cancel_requested_at": now, + "cancel_requested_by": _safe_text(requested_by), + "cancel_reason": safe_reason, + "updated_at": now, + }) + + job.update({ + "cancel_requested_at": now, + "cancel_requested_by": _safe_text(requested_by), + "cancel_requested_by_email": _safe_text(requested_by_email), + "cancel_reason": safe_reason, + "updated_at": now, + "last_message": "Migration cancellation requested. The worker will stop at its next durable checkpoint.", + "migration_state": migration_state, + }) + if job.get("status") == DATA_MANAGEMENT_STATUS_QUEUED: + job.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "completed_at": now, + "last_heartbeat_at": now, + "last_message": "Queued migration canceled before execution started.", + "lease_holder_id": None, + "lease_expires_at": None, + }) + if isinstance(migration_state, dict): + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "completed_at": now, + "updated_at": now, + }) + event_name = "migration-canceled" + event_message = "Queued migration canceled before execution started." + else: + event_name = "migration-cancel-requested" + event_message = "Migration cancellation requested; worker will stop at the next durable checkpoint." + + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + event_name, + saved_job, + status=saved_job.get("status"), + message=event_message, + details={ + "cancel_requested_at": now, + "has_reason": bool(safe_reason), + }, + ) + return saved_job + + +def retry_data_management_migration_job(job_id): + """Requeue a failed migration without changing its provenance identity or checkpoints.""" + job = get_data_management_job(job_id) + if not job: + raise DataManagementSettingsValidationError("Data Management job was not found.") + if job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Only migration jobs can be retried from migration checkpoints.") + retryable_statuses = {DATA_MANAGEMENT_STATUS_FAILED, DATA_MANAGEMENT_STATUS_CANCELED} + if job.get("status") == DATA_MANAGEMENT_STATUS_RUNNING and _is_stale_job(job): + retryable_statuses.add(DATA_MANAGEMENT_STATUS_RUNNING) + if job.get("status") not in retryable_statuses: + raise DataManagementSettingsValidationError( + "Only failed, canceled, or stale migration jobs can be retried." + ) + migration_state = job.get("migration_state") + canceled_before_start = ( + job.get("status") == DATA_MANAGEMENT_STATUS_CANCELED and + not isinstance(migration_state, dict) + ) + if not canceled_before_start and ( + not isinstance(migration_state, dict) or not migration_state.get("migration_id") + ): + raise DataManagementSettingsValidationError("This migration does not have durable checkpoint state to retry.") + if ( + not canceled_before_start and + _safe_text(migration_state.get("migration_id")) != _safe_text(job.get("id")) + ): + raise DataManagementSettingsValidationError("Migration checkpoint identity does not match the job identity.") + + now = _now_iso() + if isinstance(migration_state, dict): + migration_state["last_cancellation"] = { + "requested_at": migration_state.get("cancel_requested_at"), + "requested_by": migration_state.get("cancel_requested_by"), + "reason": migration_state.get("cancel_reason"), + } + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_QUEUED, + "updated_at": now, + "last_error": None, + "retry_requested_at": now, + "cancel_requested_at": None, + "cancel_requested_by": None, + "cancel_reason": None, + }) + job.update({ + "status": DATA_MANAGEMENT_STATUS_QUEUED, + "updated_at": now, + "completed_at": None, + "last_heartbeat_at": None, + "last_message": "Migration retry queued from durable checkpoints", + "last_error": None, + "last_cancellation": { + "requested_at": job.get("cancel_requested_at"), + "requested_by": job.get("cancel_requested_by"), + "requested_by_email": job.get("cancel_requested_by_email"), + "reason": job.get("cancel_reason"), + }, + "cancel_requested_at": None, + "cancel_requested_by": None, + "cancel_requested_by_email": None, + "cancel_reason": None, + "lease_holder_id": None, + "lease_expires_at": None, + "migration_state": migration_state, + "result": {}, + }) + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + "migration-retry-queued", + saved_job, + status=DATA_MANAGEMENT_STATUS_QUEUED, + message="Migration retry queued from durable checkpoints", + details={ + "migration_id": ( + migration_state.get("migration_id") + if isinstance(migration_state, dict) else + job.get("id") + ), + "resume_count": ( + migration_state.get("resume_count") + if isinstance(migration_state, dict) else + 0 + ), + }, + ) + return saved_job def get_data_management_job_items(job_id, limit=200): @@ -1311,8 +6988,42 @@ def get_data_management_job_detail(job_id): job = get_data_management_job(job_id) if not job: return None + sanitized_job = sanitize_data_management_job_for_admin(job) + reconciliation_resource = ( + ((job.get("migration_state") or {}).get("resources") or {}).get( + "reconciliation:cutover" + ) + if isinstance(job.get("migration_state"), dict) else None + ) + reconciliation_result = ( + reconciliation_resource.get("result") + if isinstance(reconciliation_resource, dict) and + isinstance(reconciliation_resource.get("result"), dict) + else None + ) + if reconciliation_result and isinstance(sanitized_job, dict): + sanitized_result = ( + sanitized_job.get("result") + if isinstance(sanitized_job.get("result"), dict) + else {} + ) + artifacts = ( + list(sanitized_result.get("artifacts")) + if isinstance(sanitized_result.get("artifacts"), list) + else [] + ) + if not any( + isinstance(artifact, dict) and artifact.get("type") == "migration_reconciliation" + for artifact in artifacts + ): + reconciliation_artifact = _summarize_backup_artifact(reconciliation_result) + if reconciliation_artifact: + artifacts.append(reconciliation_artifact) + sanitized_result["artifacts"] = artifacts + sanitized_result["artifact_count"] = len(artifacts) + sanitized_job["result"] = sanitized_result return { - "job": sanitize_data_management_job_for_admin(job), + "job": sanitized_job, "items": [ sanitized_item for sanitized_item in ( @@ -1324,6 +7035,27 @@ def get_data_management_job_detail(job_id): } +def get_data_management_job_progress(job_id): + job = get_data_management_job(job_id) + if not job: + return None + public_job = sanitize_data_management_job_for_admin(job) + return { + "id": public_job.get("id"), + "operation": public_job.get("operation"), + "status": public_job.get("status"), + "updated_at": public_job.get("updated_at"), + "last_heartbeat_at": public_job.get("last_heartbeat_at"), + "last_progress_at": public_job.get("last_progress_at"), + "last_message": public_job.get("last_message"), + "last_error": public_job.get("last_error"), + "progress": public_job.get("progress"), + "migration_state": public_job.get("migration_state"), + "can_retry": public_job.get("can_retry"), + "can_cancel": public_job.get("can_cancel"), + } + + def _summarize_backup_artifact(artifact): if not isinstance(artifact, dict): return None @@ -1335,6 +7067,31 @@ def _summarize_backup_artifact(artifact): "path", "bytes", "item_count", + "copied_count", + "created_count", + "updated_count", + "unchanged_count", + "skipped_count", + "failed_count", + "collision_count", + "missing_count", + "not_applicable_count", + "processed_count", + "request_units", + "elapsed_seconds", + "items_per_second", + "bytes_per_second", + "request_units_per_second", + "parallel_operations", + "batch_size", + "retry_count", + "retry_attempt_count", + "prior_failed_count", + "checkpoint_count", + "source_read_count", + "destination_accepted_count", + "destination_failed_count", + "destination_provenance_skip_count", "blob_count", "encrypted", "container_name", @@ -1343,6 +7100,21 @@ def _summarize_backup_artifact(artifact): "partial_since_epoch", "partial_filter", "prefix", + "migration_mode", + "readiness", + "deletion_policy", + "deletion_status", + "deletion_blockers", + "deleted_count", + "delete_candidate_count", + "remaining_destination_only_owned_count", + "destination_only_unowned_count", + "unresolved_scope_count", + "stale_count", + "keyset_cursor", + "actual_outcomes", + "preview_actual_divergence", + "services", "warning", ] return { @@ -1372,7 +7144,11 @@ def _backup_artifact_totals(artifacts): continue totals["artifact_count"] += 1 totals["bytes"] += _safe_int(artifact.get("bytes"), default=0, minimum=0) - totals["record_count"] += _safe_int(artifact.get("item_count"), default=0, minimum=0) + totals["record_count"] += _safe_int( + artifact.get("item_count", artifact.get("copied_count")), + default=0, + minimum=0, + ) totals["blob_count"] += _safe_int(artifact.get("blob_count"), default=0, minimum=0) if artifact.get("warning") or artifact.get("status") == "warning": totals["warning_count"] += 1 @@ -1564,12 +7340,46 @@ def normalize_data_management_migration_plan(options): raw_plan = options.get("migration_plan") if isinstance(options, dict) else {} if not isinstance(raw_plan, dict): raw_plan = {} + migration_mode = _safe_text( + raw_plan.get("migration_mode"), + DATA_MANAGEMENT_MIGRATION_MODE_NEW_ONLY, + ).lower() + if migration_mode not in DATA_MANAGEMENT_MIGRATION_MODES: + raise DataManagementSettingsValidationError( + "Migration mode must be new_only, delta_upsert, or mirror_with_deletions." + ) + baseline_job_id = _safe_text(raw_plan.get("baseline_job_id")) + if baseline_job_id: + try: + baseline_job_id = str(uuid.UUID(baseline_job_id)) + except (AttributeError, TypeError, ValueError) as exc: + raise DataManagementSettingsValidationError( + "Migration baseline job ID must be a valid GUID." + ) from exc + mirror_confirmation = _safe_text(raw_plan.get("mirror_confirmation")) + mirror_deletions_confirmed = ( + migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR and + mirror_confirmation == DATA_MANAGEMENT_MIRROR_CONFIRMATION + ) + if ( + migration_mode == DATA_MANAGEMENT_MIGRATION_MODE_MIRROR and + not mirror_deletions_confirmed + ): + raise DataManagementSettingsValidationError( + f"Mirror migration requires the exact confirmation phrase: {DATA_MANAGEMENT_MIRROR_CONFIRMATION}" + ) plan = { "users": _normalize_migration_selection(raw_plan.get("users")), "groups": _normalize_migration_selection(raw_plan.get("groups")), "public_workspaces": _normalize_migration_selection(raw_plan.get("public_workspaces")), "include_ai_search": raw_plan.get("include_ai_search") is not False, "include_source_blobs": bool(raw_plan.get("include_source_blobs")), + "target_ai_search_writes_frozen": ( + raw_plan.get("target_ai_search_writes_frozen") is True + ), + "migration_mode": migration_mode, + "baseline_job_id": baseline_job_id, + "mirror_deletions_confirmed": mirror_deletions_confirmed, } for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: if plan[target_type].get("mode") == "none": @@ -1577,83 +7387,516 @@ def normalize_data_management_migration_plan(options): return plan -def _resolve_plan_scope_ids(target_type, plan_entry): - if plan_entry.get("mode") == "all": - catalog = get_data_management_migration_catalog(target_type, limit=1000) - return [item.get("id") for item in catalog.get("items", []) if item.get("id")] - return _dedupe_limited_strings(plan_entry.get("ids")) +def _resolve_plan_scope_ids(target_type, plan_entry): + if plan_entry.get("mode") == "all": + catalog = get_data_management_migration_catalog(target_type, limit=1000) + return [item.get("id") for item in catalog.get("items", []) if item.get("id")] + return _dedupe_limited_strings(plan_entry.get("ids")) + + +def summarize_data_management_migration_plan(options): + plan = normalize_data_management_migration_plan(options or {}) + summary = {} + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: + ids = _resolve_plan_scope_ids(target_type, plan[target_type]) + summary[target_type] = { + "mode": plan[target_type].get("mode"), + "count": len(ids), + "include_documents": bool(plan[target_type].get("include_documents")), + "ids": ids[:50], + } + summary["include_ai_search"] = bool(plan.get("include_ai_search")) + summary["include_source_blobs"] = bool(plan.get("include_source_blobs")) + summary["target_ai_search_writes_frozen"] = bool( + plan.get("target_ai_search_writes_frozen") + ) + summary["migration_mode"] = plan.get("migration_mode") + summary["baseline_job_id"] = plan.get("baseline_job_id") + summary["mirror_deletions_confirmed"] = bool(plan.get("mirror_deletions_confirmed")) + return summary + + +def _count_preview_source_cosmos_records(migration_plan, source_cutoff_at): + source_cutoff = _parse_iso_datetime(source_cutoff_at) + source_cutoff_epoch = int(source_cutoff.timestamp()) if source_cutoff else None + record_count = 0 + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none": + continue + for container_definition in DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS[target_type]: + if container_definition.get("documents") and not selection.get("include_documents"): + continue + record_count += sum( + 1 + for _item in _iter_selected_cosmos_records( + container_definition, + selection, + source_cutoff_epoch=source_cutoff_epoch, + ) or [] + ) + return record_count + + +def preview_data_management_migration_plan( + settings, + options, + resolved_migration_plan=None, + source_cutoff_at=None, + migration_id=None, + heartbeat_callback=None, +): + """Build a read-only server-owned inventory preview for one migration plan.""" + if isinstance(resolved_migration_plan, dict): + migration_plan = copy.deepcopy(resolved_migration_plan) + else: + migration_plan = normalize_data_management_migration_plan(options or {}) + preview_job = {"id": str(uuid.uuid4())} + migration_plan = _resolve_data_management_migration_baseline( + preview_job, + settings, + migration_plan, + ) + source_cutoff_at = _safe_text(source_cutoff_at) or _now_iso() + preview_migration_id = _safe_text(migration_id) or str(uuid.uuid4()) + migration_state = {"source_cutoff_at": source_cutoff_at} + provenance_context = create_migration_provenance_context( + migration_id=preview_migration_id, + migrated_at_utc=source_cutoff_at, + ) + provenance_context.update({ + "migration_mode": migration_plan.get("migration_mode"), + "baseline_job_id": migration_plan.get("baseline_job_id"), + "baseline_source_cutoff_at": migration_plan.get("baseline_source_cutoff_at"), + }) + try: + target_database = _get_existing_target_cosmos_database(settings) + cosmos_report = _reconcile_cosmos_migration( + target_database, + migration_plan, + migration_state, + provenance_context, + apply_deletions=False, + heartbeat_callback=heartbeat_callback, + ) + except (CosmosResourceNotFoundError, ResourceNotFoundError): + source_record_count = _count_preview_source_cosmos_records( + migration_plan, + source_cutoff_at, + ) + cosmos_report = { + "service": "cosmos", + "create_count": source_record_count, + "update_count": 0, + "unchanged_count": 0, + "delete_candidate_count": 0, + "conflict_count": 0, + "missing_count": 0, + "warning": "Destination Cosmos database does not exist; selected source records are new.", + } + search_report = _reconcile_ai_search_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat_callback, + ) + blob_report = _reconcile_blob_migration( + settings, + migration_plan, + apply_deletions=False, + heartbeat_callback=heartbeat_callback, + ) + service_reports = [cosmos_report, search_report, blob_report] + estimated_outcomes = { + "create_count": sum(_safe_int(report.get("create_count"), default=0) for report in service_reports), + "update_count": sum(_safe_int(report.get("update_count"), default=0) for report in service_reports), + "unchanged_count": sum(_safe_int(report.get("unchanged_count"), default=0) for report in service_reports), + "delete_count": sum(_safe_int(report.get("delete_candidate_count"), default=0) for report in service_reports), + "not_applicable_count": sum(_safe_int(report.get("not_applicable_count"), default=0) for report in service_reports), + "missing_count": sum( + _safe_int(report.get("source_missing_count"), default=0) + for report in service_reports + ), + "conflict_count": sum(_safe_int(report.get("conflict_count"), default=0) for report in service_reports), + "failed_count": 0, + } + normalized_preview_plan = copy.deepcopy(migration_plan) + normalized_preview_plan.pop("mirror_confirmation", None) + return { + "captured_at": source_cutoff_at, + "migration_mode": migration_plan.get("migration_mode"), + "baseline_job_id": migration_plan.get("baseline_job_id"), + "baseline_source_cutoff_at": migration_plan.get("baseline_source_cutoff_at"), + "plan_fingerprint": hashlib.sha256(json.dumps( + normalized_preview_plan, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=str, + ).encode("utf-8")).hexdigest(), + "estimated_outcomes": estimated_outcomes, + "services": service_reports, + } + + +def _try_claim_data_management_job(job_id, settings=None): + try: + job = _read_job(job_id) + except CosmosResourceNotFoundError: + return None + + status = _safe_text(job.get("status")) + if status in DATA_MANAGEMENT_TERMINAL_STATUSES: + return None + if status == DATA_MANAGEMENT_STATUS_RUNNING and not _is_stale_job(job): + return None + if job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION: + try: + if _has_active_migration_destination_lock(job): + return None + except DataManagementMigrationLeaseLostError as exc: + log_event( + "[DataManagement] Migration retry claim deferred because its coordinator lock could not be verified.", + {"job_id": job_id, "error": str(exc)}, + level=logging.WARNING, + ) + return None + + lease_seconds = _safe_int( + (settings or {}).get("data_management_job_lease_seconds"), + default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, + minimum=60, + maximum=7200, + ) + now = _now_utc() + job.update({ + "status": DATA_MANAGEMENT_STATUS_RUNNING, + "started_at": job.get("started_at") or now.isoformat(), + "updated_at": now.isoformat(), + "last_heartbeat_at": now.isoformat(), + "lease_holder_id": _job_lease_holder_id(), + "lease_expires_at": (now + timedelta(seconds=lease_seconds)).isoformat(), + "last_message": "Data management job claimed by a worker", + }) + if job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION: + job["migration_attempt_id"] = str(uuid.uuid4()) + try: + claimed_job = _replace_job(job) + _record_data_management_job_event( + claimed_job.get("id"), + "claimed", + claimed_job, + status=DATA_MANAGEMENT_STATUS_RUNNING, + message="Data management job claimed by a worker", + details={"lease_expires_at": claimed_job.get("lease_expires_at")}, + ) + return claimed_job + except Exception as exc: + status_code = getattr(exc, "status_code", None) + if status_code not in (409, 412): + log_event( + "[DataManagement] Job claim failed.", + {"job_id": job_id, "status_code": status_code, "error": str(exc)}, + level=logging.WARNING, + ) + return None + + +def _safe_job_item_id_part(value): + normalized = _safe_text(value, "event") + safe_value = "".join(character if character.isalnum() or character in {"-", "_"} else "-" for character in normalized) + return (safe_value.strip("-") or "event")[:80] + + +def _sanitize_migration_manifest_entry(entry): + allowed_fields = { + "service", + "resource_name", + "target_type", + "source_identity", + "destination_identity", + "item_ref", + "status", + "attempt", + "bytes", + "source_version", + "source_hash", + "error", + "recorded_at", + } + sanitized = { + field_name: _sanitize_activity_value((entry or {}).get(field_name)) + for field_name in allowed_fields + if (entry or {}).get(field_name) is not None + } + sanitized["error"] = _safe_text(sanitized.get("error"))[:500] + sanitized["recorded_at"] = _safe_text(sanitized.get("recorded_at")) or _now_iso() + return sanitized + + +def _hash_migration_manifest_identity(*identity_parts): + encoded = json.dumps( + [str(part) for part in identity_parts], + ensure_ascii=True, + separators=(",", ":"), + default=str, + ).encode("utf-8") + return f"sha256:{hashlib.sha256(encoded).hexdigest()}" + + +def _build_migration_manifest_item_ref( + job_id, + service="", + resource_name="", + target_type="", + document_id="", + partition_key=None, + container_name="", + blob_name="", + index_name="", +): + return str(uuid.uuid4()) + + +def resolve_data_management_migration_manifest_item(job_id, item_ref): + safe_job_id = _safe_text(job_id) + safe_item_ref = _safe_text(item_ref) + if not safe_job_id or not safe_item_ref: + raise DataManagementSettingsValidationError("Migration manifest item reference is required.") + job = get_data_management_job(safe_job_id) + if not job or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Data Management migration job was not found.") + query = ( + "SELECT * FROM c WHERE c.job_id = @job_id AND c.type = @type " + "ORDER BY c.created_at ASC" + ) + parameters = [ + {"name": "@job_id", "value": safe_job_id}, + {"name": "@type", "value": DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE}, + ] + batches = cosmos_data_management_job_items_container.query_items( + query=query, + parameters=parameters, + partition_key=safe_job_id, + max_item_count=100, + ) + for batch in batches: + for locator in batch.get("private_locators") or []: + if _safe_text(locator.get("item_ref")) != safe_item_ref: + continue + return { + "service": _safe_text(locator.get("service")), + "resource_name": _safe_text(locator.get("resource_name")), + "target_type": _safe_text(locator.get("target_type")), + "document_id": _safe_text(locator.get("document_id")), + "partition_key": locator.get("partition_key"), + "container_name": _safe_text(locator.get("container_name")), + "blob_name": _safe_text(locator.get("blob_name")), + "index_name": _safe_text(locator.get("index_name")), + } + raise DataManagementSettingsValidationError( + "Migration manifest item reference was not found." + ) + + +def _write_migration_manifest_batch(job_id, resource_name, entries): + normalized_entries = [] + private_locators = [] + for entry in entries or []: + if not isinstance(entry, dict): + continue + locator = entry.get("_locator") if isinstance(entry.get("_locator"), dict) else None + item_ref = _safe_text(entry.get("item_ref")) + if locator and not item_ref: + item_ref = _build_migration_manifest_item_ref(job_id, **locator) + public_entry = dict(entry) + public_entry.pop("_locator", None) + if item_ref: + public_entry["item_ref"] = item_ref + normalized_entries.append(_sanitize_migration_manifest_entry(public_entry)) + if locator and item_ref: + private_locators.append({ + "item_ref": item_ref, + "service": _safe_text(locator.get("service")), + "resource_name": _safe_text(locator.get("resource_name")), + "target_type": _safe_text(locator.get("target_type")), + "document_id": _safe_text(locator.get("document_id")), + "partition_key": copy.deepcopy(locator.get("partition_key")), + "container_name": _safe_text(locator.get("container_name")), + "blob_name": _safe_text(locator.get("blob_name")), + "index_name": _safe_text(locator.get("index_name")), + }) + if not normalized_entries: + return None + create_item = getattr(cosmos_data_management_job_items_container, "create_item", None) + if not callable(create_item): + raise RuntimeError("Migration manifest storage is not available.") + now = _now_iso() + body = { + "id": ( + f"{_safe_job_item_id_part(job_id)}:manifest:" + f"{_safe_job_item_id_part(resource_name)}:{uuid.uuid4().hex}" + ), + "job_id": _safe_text(job_id), + "type": DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE, + "resource_name": _safe_text(resource_name), + "created_at": now, + "updated_at": now, + "entry_count": len(normalized_entries), + "entries": normalized_entries, + "private_locators": private_locators, + } + return create_item(body=body) + + +def _create_migration_manifest_writer(job_id, resource_name): + buffer = [] + + def append(entry): + buffer.append(entry) + if len(buffer) >= DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE: + flush() + + def flush(): + if not buffer: + return + pending_entries = list(buffer) + _write_migration_manifest_batch(job_id, resource_name, pending_entries) + del buffer[:len(pending_entries)] + + return append, flush + + +def iter_data_management_migration_manifest_entries(job_id, statuses=None): + safe_job_id = _safe_text(job_id) + if not safe_job_id: + return + normalized_statuses = { + _safe_text(status).lower() + for status in (statuses or []) + if _safe_text(status) + } + query = ( + "SELECT * FROM c WHERE c.job_id = @job_id AND c.type = @type " + "ORDER BY c.created_at ASC" + ) + parameters = [ + {"name": "@job_id", "value": safe_job_id}, + {"name": "@type", "value": DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_TYPE}, + ] + batches = cosmos_data_management_job_items_container.query_items( + query=query, + parameters=parameters, + partition_key=safe_job_id, + max_item_count=100, + ) + for batch in batches: + for entry in batch.get("entries") or []: + sanitized = _sanitize_migration_manifest_entry(entry) + if normalized_statuses and _safe_text(sanitized.get("status")).lower() not in normalized_statuses: + continue + yield sanitized + + +def export_data_management_migration_manifest(job_id, statuses=None): + job = get_data_management_job(job_id) + if not job or job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION: + raise DataManagementSettingsValidationError("Data Management migration job was not found.") + entries = list(iter_data_management_migration_manifest_entries(job_id, statuses=statuses)) + lines = [ + json.dumps(entry, ensure_ascii=True, separators=(",", ":"), default=str) + for entry in entries + ] + return { + "content": "\n".join(lines) + ("\n" if lines else ""), + "entry_count": len(entries), + "statuses": sorted({ + _safe_text(status).lower() + for status in (statuses or []) + if _safe_text(status) + }), + } + + +def _sanitize_mirror_deletion_candidate(candidate): + allowed_fields = { + "service", + "target_type", + "container_name", + "document_id", + "partition_key", + "target_etag", + "index_name", + "blob_name", + } + return { + field_name: copy.deepcopy((candidate or {}).get(field_name)) + for field_name in allowed_fields + if (candidate or {}).get(field_name) is not None + } -def summarize_data_management_migration_plan(options): - plan = normalize_data_management_migration_plan(options or {}) - summary = {} - for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: - ids = _resolve_plan_scope_ids(target_type, plan[target_type]) - summary[target_type] = { - "mode": plan[target_type].get("mode"), - "count": len(ids), - "include_documents": bool(plan[target_type].get("include_documents")), - "ids": ids[:50], - } - summary["include_ai_search"] = bool(plan.get("include_ai_search")) - summary["include_source_blobs"] = bool(plan.get("include_source_blobs")) - return summary +def _write_mirror_deletion_candidate_batch(job_id, plan_id, candidates): + normalized_candidates = [ + _sanitize_mirror_deletion_candidate(candidate) + for candidate in (candidates or []) + if isinstance(candidate, dict) + ] + if not normalized_candidates: + return None + now = _now_iso() + return cosmos_data_management_job_items_container.create_item(body={ + "id": ( + f"{_safe_job_item_id_part(job_id)}:mirror-plan:" + f"{_safe_job_item_id_part(plan_id)}:{uuid.uuid4().hex}" + ), + "job_id": _safe_text(job_id), + "type": DATA_MANAGEMENT_MIRROR_DELETION_BATCH_TYPE, + "plan_id": _safe_text(plan_id), + "created_at": now, + "updated_at": now, + "candidate_count": len(normalized_candidates), + "candidates": normalized_candidates, + }) -def _try_claim_data_management_job(job_id, settings=None): - try: - job = _read_job(job_id) - except CosmosResourceNotFoundError: - return None +def _create_mirror_deletion_candidate_writer(job_id, plan_id): + buffer = [] - status = _safe_text(job.get("status")) - if status in DATA_MANAGEMENT_TERMINAL_STATUSES: - return None - if status == DATA_MANAGEMENT_STATUS_RUNNING and not _is_stale_job(job): - return None + def append(candidate): + buffer.append(candidate) + if len(buffer) >= DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE: + flush() - lease_seconds = _safe_int( - (settings or {}).get("data_management_job_lease_seconds"), - default=DATA_MANAGEMENT_DEFAULT_LEASE_SECONDS, - minimum=60, - maximum=7200, - ) - now = _now_utc() - job.update({ - "status": DATA_MANAGEMENT_STATUS_RUNNING, - "started_at": job.get("started_at") or now.isoformat(), - "updated_at": now.isoformat(), - "last_heartbeat_at": now.isoformat(), - "lease_holder_id": _job_lease_holder_id(), - "lease_expires_at": (now + timedelta(seconds=lease_seconds)).isoformat(), - "last_message": "Data management job claimed by a worker", - }) - try: - claimed_job = _replace_job(job) - _record_data_management_job_event( - claimed_job.get("id"), - "claimed", - claimed_job, - status=DATA_MANAGEMENT_STATUS_RUNNING, - message="Data management job claimed by a worker", - details={"lease_expires_at": claimed_job.get("lease_expires_at")}, - ) - return claimed_job - except Exception as exc: - status_code = getattr(exc, "status_code", None) - if status_code not in (409, 412): - log_event( - "[DataManagement] Job claim failed.", - {"job_id": job_id, "status_code": status_code, "error": str(exc)}, - level=logging.WARNING, - ) - return None + def flush(): + if not buffer: + return + pending_candidates = list(buffer) + _write_mirror_deletion_candidate_batch(job_id, plan_id, pending_candidates) + del buffer[:len(pending_candidates)] + return append, flush -def _safe_job_item_id_part(value): - normalized = _safe_text(value, "event") - safe_value = "".join(character if character.isalnum() or character in {"-", "_"} else "-" for character in normalized) - return (safe_value.strip("-") or "event")[:80] + +def _iter_mirror_deletion_candidates(job_id, plan_id): + query = ( + "SELECT * FROM c WHERE c.job_id = @job_id AND c.type = @type " + "AND c.plan_id = @plan_id ORDER BY c.created_at ASC" + ) + parameters = [ + {"name": "@job_id", "value": _safe_text(job_id)}, + {"name": "@type", "value": DATA_MANAGEMENT_MIRROR_DELETION_BATCH_TYPE}, + {"name": "@plan_id", "value": _safe_text(plan_id)}, + ] + batches = cosmos_data_management_job_items_container.query_items( + query=query, + parameters=parameters, + partition_key=_safe_text(job_id), + max_item_count=100, + ) + for batch in batches: + for candidate in batch.get("candidates") or []: + yield _sanitize_mirror_deletion_candidate(candidate) def create_data_management_job_item(job_id, step_name, status=DATA_MANAGEMENT_STATUS_QUEUED, message=None, details=None): @@ -2116,20 +8359,50 @@ def save_data_management_cosmos_editor_document(container_name, document_id, par def _save_data_management_job(job): body = _strip_cosmos_system_fields(job) - saved = cosmos_data_management_jobs_container.upsert_item(body) + job_etag = job.get("_etag") if isinstance(job, dict) else None + if job_etag and hasattr(cosmos_data_management_jobs_container, "replace_item"): + try: + saved = cosmos_data_management_jobs_container.replace_item( + item=job.get("id"), + body=body, + etag=job_etag, + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if getattr(exc, "status_code", None) in {409, 412}: + raise DataManagementMigrationLeaseLostError( + "Migration job state was changed by another worker." + ) from exc + raise + else: + saved = cosmos_data_management_jobs_container.upsert_item(body) job.clear() job.update(saved) return job -def _set_job_progress(job, message, completed_steps, total_steps, current_step=None, status=DATA_MANAGEMENT_STATUS_RUNNING): +def _set_job_progress( + job, + message, + completed_steps, + total_steps, + current_step=None, + status=DATA_MANAGEMENT_STATUS_RUNNING, + allow_cancel_requested=False, +): total_steps = max(1, total_steps) + _assert_migration_job_lease( + job, + allow_cancel_requested=allow_cancel_requested, + ) completed_steps = max(0, min(completed_steps, total_steps)) percent_complete = int((completed_steps / total_steps) * 100) + now = _now_iso() job.update({ "status": status, - "updated_at": _now_iso(), - "last_heartbeat_at": _now_iso(), + "updated_at": now, + "last_heartbeat_at": now, + "last_progress_at": now, "last_message": message, "progress": { "total_steps": total_steps, @@ -2573,73 +8846,396 @@ def execute_backup_job(job, settings): def execute_migration_job(job, settings): options = job.get("options") if isinstance(job.get("options"), dict) else {} migration_plan = normalize_data_management_migration_plan(options) + total_steps = 10 + _set_job_progress( + job, + "Resolving migration plan and baseline", + 0, + total_steps, + current_step="plan", + allow_cancel_requested=True, + ) + migration_plan = _resolve_data_management_migration_baseline( + job, + settings, + migration_plan, + ) dry_run = bool(options.get("dry_run")) warnings = [] artifacts = [] - total_steps = 5 plan_summary = summarize_data_management_migration_plan({"migration_plan": migration_plan}) + plan_summary["baseline_source_cutoff_at"] = migration_plan.get("baseline_source_cutoff_at") + selected_total = sum(plan_summary[target_type]["count"] for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES) + if selected_total == 0: + raise DataManagementSettingsValidationError("Choose at least one user, group, or public workspace before running migration.") + + migration_state, provenance_context = _initialize_data_management_migration_state( + job, + settings, + migration_plan, + ) + try: + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Initialized durable migration provenance and checkpoint state", + ) + except DataManagementMigrationCanceledError as exc: + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "canceled_at": _now_iso(), + "last_error": None, + "cancel_message": str(exc), + }) + _persist_migration_state( + job, + migration_state, + settings, + "Migration cancellation acknowledged before execution started", + allow_cancel_requested=True, + ) + raise _set_job_progress(job, "Validated migration selection plan", 1, total_steps, current_step="plan") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state _record_data_management_job_event( job.get("id"), "migration-plan", job, status=DATA_MANAGEMENT_STATUS_RUNNING, message="Migration selection plan validated", - details={"migration_plan": plan_summary, "dry_run": dry_run}, + details={ + "migration_plan": plan_summary, + "dry_run": dry_run, + "migration_id": provenance_context["migration_id"], + "parallel_operations": _get_migration_parallel_operations(settings), + }, ) - selected_total = sum(plan_summary[target_type]["count"] for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES) - if selected_total == 0: - raise DataManagementSettingsValidationError("Choose at least one user, group, or public workspace before running migration.") - if dry_run: warning = "Migration dry run completed. No destination records were written." warnings.append(warning) + migration_state.update({ + "status": MIGRATION_RESOURCE_STATUS_COMPLETED, + "completed_at": _now_iso(), + }) + migration_state = _persist_migration_state(job, migration_state, settings, warning) _set_job_progress(job, warning, total_steps, total_steps, current_step="dry_run", status=DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS) return { "migration_plan": plan_summary, "dry_run": True, + "migration_id": provenance_context["migration_id"], + "migration_state": migration_state, "artifacts": [], "artifact_totals": _backup_artifact_totals([]), "warnings": warnings, } + _validate_target_ai_search_migration_write_safety(settings, migration_plan) + _acquire_migration_destination_lock(job, settings, migration_plan) + try: + _save_data_management_job(job) + except Exception: + _release_migration_destination_lock(job) + raise + + # The empty destination database must exist before it can host the target-wide coordinator. target_database = _get_target_cosmos_database(settings) - _set_job_progress(job, "Connected to target Cosmos database", 2, total_steps, current_step="target_cosmos") + try: + _acquire_target_migration_coordinator(job, target_database, settings) + _save_data_management_job(job) + except Exception: + _release_target_migration_coordinator(job) + _release_migration_destination_lock(job) + raise - for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPES: - selection = migration_plan.get(target_type) or {} - if selection.get("mode") == "none": - continue - copied = _copy_cosmos_records_to_target(target_database, target_type, selection) - artifacts.extend(copied) + try: + _set_job_progress( + job, + "Building server-owned migration inventory", + 1, + total_steps, + current_step="inventory", + ) + preview_resource_name = "migration_preview:inventory" + preview_resource = migration_state.get("resources", {}).get(preview_resource_name) + migration_preview = ( + preview_resource.get("result") + if is_migration_resource_completed(migration_state, preview_resource_name) and + isinstance(preview_resource, dict) and + isinstance(preview_resource.get("result"), dict) + else None + ) + if not isinstance(migration_preview, dict): + start_migration_resource(migration_state, preview_resource_name) + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + preview_resource_name, + {"state": "inventory", "heartbeat_count": 0}, + "Building server-owned migration inventory preview", + ) + preview_heartbeat_count = 0 + preview_last_heartbeat_at = 0.0 + + def preview_heartbeat(message, completed_count=0): + nonlocal migration_state, preview_heartbeat_count, preview_last_heartbeat_at + preview_heartbeat_count += 1 + now = time.monotonic() + if now - preview_last_heartbeat_at < 2.0: + return + migration_state = _persist_migration_checkpoint( + job, + migration_state, + settings, + preview_resource_name, + { + "state": "inventory", + "heartbeat_count": preview_heartbeat_count, + "completed_count": completed_count, + }, + message, + ) + preview_last_heartbeat_at = now + + migration_preview = preview_data_management_migration_plan( + settings, + options, + resolved_migration_plan=migration_plan, + source_cutoff_at=migration_state.get("source_cutoff_at"), + migration_id=job.get("id"), + heartbeat_callback=preview_heartbeat, + ) + migration_state = _complete_migration_resource_checkpoint( + job, + migration_state, + settings, + preview_resource_name, + migration_preview, + "Completed server-owned migration inventory preview", + ) + options["migration_preview"] = migration_preview + job["options"] = options + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Pinned server-owned migration inventory preview", + ) + + _set_job_progress(job, "Migration inventory completed", 2, total_steps, current_step="inventory") + _set_job_progress(job, "Validating migration destinations", 2, total_steps, current_step="preflight") + migration_state = _run_data_management_migration_preflight( + job, + migration_state, + settings, + migration_plan, + ) + _set_job_progress(job, "Destination migration preflight completed", 3, total_steps, current_step="preflight") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state _record_data_management_job_event( job.get("id"), - f"migration-cosmos-{target_type}", + "migration-preflight", job, status=DATA_MANAGEMENT_STATUS_RUNNING, - message=f"Migrated {target_type.replace('_', ' ')} Cosmos records", - details={"target_type": target_type, "artifacts": copied}, + message="Verified source and destination migration access", + details=migration_state.get("preflight") if isinstance(migration_state.get("preflight"), dict) else {}, ) - _set_job_progress(job, "Cosmos migration completed", 3, total_steps, current_step="cosmos") - try: - search_artifacts = _copy_ai_search_to_target(settings, migration_plan) + _set_job_progress(job, "Preparing destination Cosmos capacity", 3, total_steps, current_step="capacity") + migration_state = _apply_temporary_destination_capacity( + job, + migration_state, + settings, + migration_plan, + ) + _set_job_progress(job, "Destination Cosmos capacity prepared", 4, total_steps, current_step="capacity") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + + _set_job_progress(job, "Migrating Cosmos records", 4, total_steps, current_step="cosmos") + for target_type in DATA_MANAGEMENT_MIGRATION_TARGET_TYPE_ORDER: + selection = migration_plan.get(target_type) or {} + if selection.get("mode") == "none": + continue + copied = _copy_cosmos_records_to_target( + target_database, + target_type, + selection, + job, + migration_state, + provenance_context, + settings, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + artifacts.extend(copied) + _record_data_management_job_event( + job.get("id"), + f"migration-cosmos-{target_type}", + job, + status=DATA_MANAGEMENT_STATUS_RUNNING, + message=f"Migrated {target_type.replace('_', ' ')} Cosmos records", + details={"target_type": target_type, "artifacts": copied}, + ) + _set_job_progress(job, "Cosmos migration completed", 5, total_steps, current_step="cosmos") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + + search_artifacts = [] + if _migration_plan_includes_ai_search_documents(migration_plan): + _set_job_progress(job, "Freezing target AI Search writers", 5, total_steps, current_step="ai_search") + target_search_write_gate_container = _get_target_data_management_search_write_gate_container( + target_database + ) + target_search_write_fence = acquire_data_management_search_write_fence( + target_search_write_gate_container, + job.get("id"), + _get_target_search_write_fence_lease_seconds(settings), + heartbeat_callback=lambda: _persist_migration_heartbeat( + job, + settings, + "Waiting for target AI Search writers to drain", + ), + ) + release_target_search_write_fence = False + try: + _set_job_progress(job, "Migrating AI Search documents", 5, total_steps, current_step="ai_search") + search_artifacts = _copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_search_write_fence=( + target_search_write_gate_container, + target_search_write_fence, + ), + ) + release_target_search_write_fence = True + except DataManagementMigrationCanceledError: + release_target_search_write_fence = True + raise + finally: + if release_target_search_write_fence: + release_data_management_search_write_fence( + target_search_write_gate_container, + target_search_write_fence, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state artifacts.extend(search_artifacts) - except Exception as exc: - warnings.append(f"AI Search migration warning: {exc}") - _set_job_progress(job, "AI Search migration step completed", 4, total_steps, current_step="ai_search") + _set_job_progress(job, "AI Search migration completed", 6, total_steps, current_step="ai_search") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state - try: - source_blob_artifacts = _copy_source_blobs_to_target(settings, migration_plan) + _set_job_progress(job, "Migrating source document blobs", 6, total_steps, current_step="source_blobs") + source_blob_artifacts = _copy_source_blobs_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state artifacts.extend(source_blob_artifacts) + _set_job_progress(job, "Source blob migration completed", 7, total_steps, current_step="source_blobs") + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + + _set_job_progress(job, "Reconciling source and destination", 7, total_steps, current_step="reconciliation") + reconciliation_artifact = _run_data_management_migration_reconciliation( + settings, + migration_plan, + job, + migration_state, + provenance_context, + target_database, + preview_snapshot=migration_preview, + migration_artifacts=artifacts, + ) + migration_state = job.get("migration_state") if isinstance(job.get("migration_state"), dict) else migration_state + artifacts.append(reconciliation_artifact) + _set_job_progress(job, "Migration reconciliation completed", 8, total_steps, current_step="reconciliation") + _record_data_management_job_event( + job.get("id"), + "migration-reconciliation", + job, + status=DATA_MANAGEMENT_STATUS_RUNNING, + message="Reconciled migration source and destination identities", + details=reconciliation_artifact, + ) + except DataManagementMigrationCanceledError as exc: + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "canceled_at": _now_iso(), + "last_error": None, + "cancel_message": str(exc), + }) + _persist_migration_state( + job, + migration_state, + settings, + "Migration cancellation acknowledged at a durable checkpoint", + allow_cancel_requested=True, + ) + raise except Exception as exc: - warnings.append(f"Enhanced Citations source blob migration warning: {exc}") + migration_state["status"] = MIGRATION_RESOURCE_STATUS_FAILED + migration_state["last_error"] = str(exc)[:1000] + _persist_migration_state(job, migration_state, settings, "Migration execution failed") + raise + finally: + try: + _set_job_progress( + job, + "Restoring destination Cosmos capacity", + 8, + total_steps, + current_step="capacity_restore", + allow_cancel_requested=True, + ) + restore_warnings, migration_state = _restore_temporary_destination_capacity( + job, + migration_state, + settings, + allow_cancel_requested=True, + ) + warnings.extend(restore_warnings) + except Exception as restore_exc: + warnings.append(f"Failed to restore temporary destination Cosmos capacity: {str(restore_exc)[:300]}") + + capacity = ( + migration_state.get("capacity") + if isinstance(migration_state.get("capacity"), dict) + else {} + ) + if capacity.get("restore_pending"): + migration_state.update({ + "status": MIGRATION_RESOURCE_STATUS_FAILED, + "last_error": "Destination Cosmos capacity restoration remains pending.", + }) + _persist_migration_state( + job, + migration_state, + settings, + "Destination Cosmos capacity restoration remains pending; migration can be retried", + ) + raise RuntimeError( + "Destination Cosmos capacity restoration remains pending. Retry the migration to restore the recorded capacity snapshot." + ) + for artifact in artifacts: if artifact.get("warning"): warnings.append(artifact.get("warning")) - _set_job_progress(job, "Source blob migration step completed", 5, total_steps, current_step="source_blobs") + migration_state.update({ + "status": MIGRATION_RESOURCE_STATUS_COMPLETED, + "completed_at": _now_iso(), + }) + _update_migration_state_totals(migration_state) + migration_state = _persist_migration_state( + job, + migration_state, + settings, + "Migration execution completed", + ) + _set_job_progress(job, "Migration execution completed", 10, total_steps, current_step="complete") artifact_summaries = summarize_backup_artifacts(artifacts) artifact_totals = _backup_artifact_totals(artifact_summaries) @@ -2651,6 +9247,12 @@ def execute_migration_job(job, settings): message="Migration execution completed", details={ "migration_plan": plan_summary, + "migration_id": provenance_context["migration_id"], + "migration_state": { + "resume_count": migration_state.get("resume_count"), + "totals": migration_state.get("totals"), + "capacity": migration_state.get("capacity"), + }, "artifact_count": len(artifacts), "artifact_totals": artifact_totals, "artifacts": artifact_summaries, @@ -2660,6 +9262,8 @@ def execute_migration_job(job, settings): return { "migration_plan": plan_summary, "dry_run": False, + "migration_id": provenance_context["migration_id"], + "migration_state": migration_state, "artifact_count": len(artifacts), "artifact_totals": artifact_totals, "artifacts": artifact_summaries, @@ -2672,6 +9276,7 @@ def process_data_management_job(job_id): job = _try_claim_data_management_job(job_id, settings=settings) if not job: return None + release_migration_lock = job.get("operation") != DATA_MANAGEMENT_OPERATION_MIGRATION try: if job.get("operation") == DATA_MANAGEMENT_OPERATION_BACKUP: @@ -2681,6 +9286,7 @@ def process_data_management_job(job_id): message = "Backup completed with warnings" if warnings else "Backup completed successfully" elif job.get("operation") == DATA_MANAGEMENT_OPERATION_MIGRATION: result = execute_migration_job(job, settings) + release_migration_lock = True warnings = list(job.get("warnings") or []) + result.get("warnings", []) status = DATA_MANAGEMENT_STATUS_COMPLETED_WITH_WARNINGS if warnings else DATA_MANAGEMENT_STATUS_COMPLETED message = "Migration completed with warnings" if warnings else "Migration completed successfully" @@ -2702,6 +9308,24 @@ def process_data_management_job(job_id): message = "Data Management job foundation completed with implementation warnings" now = _now_iso() + completed_progress = ( + copy.deepcopy(job.get("progress")) + if isinstance(job.get("progress"), dict) + else {} + ) + if not completed_progress: + completed_progress = { + "total_steps": 1, + "completed_steps": 1, + "current_step": "complete", + "percent_complete": 100, + } + else: + completed_progress.update({ + "completed_steps": completed_progress.get("total_steps", 1), + "current_step": "complete", + "percent_complete": 100, + }) job.update({ "status": status, "updated_at": now, @@ -2713,12 +9337,7 @@ def process_data_management_job(job_id): "lease_expires_at": None, "warnings": warnings, "result": result, - "progress": { - "total_steps": 1, - "completed_steps": 1, - "current_step": "complete", - "percent_complete": 100, - }, + "progress": completed_progress, }) saved_job = _save_data_management_job(job) _record_data_management_job_event( @@ -2730,6 +9349,54 @@ def process_data_management_job(job_id): details={"warnings": warnings, "result": result}, ) return saved_job + except DataManagementMigrationCanceledError as exc: + release_migration_lock = True + now = _now_iso() + migration_state = job.get("migration_state") + if isinstance(migration_state, dict): + migration_state.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "canceled_at": now, + "updated_at": now, + "last_error": None, + }) + job.update({ + "status": DATA_MANAGEMENT_STATUS_CANCELED, + "updated_at": now, + "completed_at": now, + "last_heartbeat_at": now, + "last_message": "Migration canceled at a durable checkpoint.", + "last_error": None, + "lease_holder_id": None, + "lease_expires_at": None, + "migration_state": migration_state, + "progress": { + "total_steps": job.get("progress", {}).get("total_steps", 0), + "completed_steps": job.get("progress", {}).get("completed_steps", 0), + "current_step": "canceled", + "percent_complete": job.get("progress", {}).get("percent_complete", 0), + }, + }) + saved_job = _save_data_management_job(job) + _record_data_management_job_event( + saved_job.get("id"), + "migration-canceled", + saved_job, + status=DATA_MANAGEMENT_STATUS_CANCELED, + message="Migration canceled at a durable checkpoint.", + details={ + "cancel_requested_at": saved_job.get("cancel_requested_at"), + "reason_present": bool(saved_job.get("cancel_reason")), + }, + ) + return saved_job + except DataManagementMigrationLeaseLostError as exc: + log_event( + "[DataManagement] Migration worker stopped after losing its job lease.", + {"job_id": job_id, "operation": job.get("operation"), "error": str(exc)}, + level=logging.WARNING, + ) + return None except Exception as exc: now = _now_iso() job.update({ @@ -2758,12 +9425,14 @@ def process_data_management_job(job_id): details={"error": str(exc)}, ) return saved_job + finally: + if release_migration_lock: + _release_target_migration_coordinator(job) + _release_migration_destination_lock(job) def submit_data_management_job(app, job_id): - if not app: - return False - executor = app.extensions.get("executor") + executor = app.extensions.get("executor") if app else None if executor and hasattr(executor, "submit_stored"): executor.submit_stored( f"data_management_{job_id}", @@ -2774,7 +9443,14 @@ def submit_data_management_job(app, job_id): if executor and hasattr(executor, "submit"): executor.submit(process_data_management_job, job_id) return True - return False + worker_thread = Thread( + target=process_data_management_job, + kwargs={"job_id": job_id}, + daemon=True, + name=f"data-management-{job_id[:12]}", + ) + worker_thread.start() + return True def build_scheduled_occurrence_id(backup_type, run_at): @@ -2785,10 +9461,15 @@ def build_scheduled_occurrence_id(backup_type, run_at): def check_due_data_management_jobs_once(app=None): settings = get_data_management_settings() + current_time = _now_utc() + recovery_results = recover_data_management_migration_jobs( + app=app, + settings=settings, + current_time=current_time, + ) if not settings.get("enabled"): - return [] + return recovery_results - current_time = _now_utc() queued_jobs = [] for backup_type, next_key in ( (DATA_MANAGEMENT_BACKUP_FULL, "next_full_backup_run_at"), @@ -2823,4 +9504,4 @@ def check_due_data_management_jobs_once(app=None): if not submitted: process_data_management_job(job.get("id")) - return queued_jobs \ No newline at end of file + return queued_jobs + recovery_results \ No newline at end of file diff --git a/application/single_app/functions_data_management_migration_state.py b/application/single_app/functions_data_management_migration_state.py new file mode 100644 index 000000000..fd26a1ceb --- /dev/null +++ b/application/single_app/functions_data_management_migration_state.py @@ -0,0 +1,203 @@ +# functions_data_management_migration_state.py +"""Durable, secret-free checkpoint state for Data Management migrations.""" + +import copy +import hashlib +import json +from datetime import datetime, timezone + + +MIGRATION_STATE_SCHEMA_VERSION = 1 +MIGRATION_RESOURCE_STATUS_PENDING = "pending" +MIGRATION_RESOURCE_STATUS_IN_PROGRESS = "in_progress" +MIGRATION_RESOURCE_STATUS_COMPLETED = "completed" +MIGRATION_RESOURCE_STATUS_FAILED = "failed" + + +def _utc_now_iso(current_time=None): + """Return a normalized UTC timestamp without accepting local wall-clock time.""" + timestamp = current_time if isinstance(current_time, datetime) else datetime.now(timezone.utc) + if timestamp.tzinfo is None: + timestamp = timestamp.replace(tzinfo=timezone.utc) + return timestamp.astimezone(timezone.utc).isoformat() + + +def _parse_utc_datetime(value): + """Parse a durable checkpoint timestamp when it is valid.""" + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except (TypeError, ValueError): + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def build_migration_configuration_fingerprint(configuration): + """Hash a caller-supplied, secret-free configuration snapshot.""" + encoded = json.dumps( + configuration if isinstance(configuration, dict) else {}, + ensure_ascii=True, + sort_keys=True, + separators=(",", ":"), + default=str, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def initialize_migration_state(existing_state, migration_id, configuration, current_time=None): + """Create or resume an in-job migration state with fingerprint protection.""" + fingerprint = build_migration_configuration_fingerprint(configuration) + timestamp = _utc_now_iso(current_time) + state = copy.deepcopy(existing_state) if isinstance(existing_state, dict) else {} + + if state: + if state.get("schema_version") != MIGRATION_STATE_SCHEMA_VERSION: + raise ValueError("Migration state uses an unsupported schema version.") + if state.get("migration_id") != migration_id: + raise ValueError("Migration state belongs to a different migration ID.") + if state.get("configuration_fingerprint") != fingerprint: + raise ValueError("Migration settings changed after the job was checkpointed. Queue a new migration job.") + state["status"] = MIGRATION_RESOURCE_STATUS_IN_PROGRESS + state["updated_at"] = timestamp + state["resume_count"] = int(state.get("resume_count") or 0) + 1 + state.setdefault("resources", {}) + state.setdefault("preflight", {}) + state.setdefault("capacity", {}) + state.setdefault("totals", {}) + return state + + return { + "schema_version": MIGRATION_STATE_SCHEMA_VERSION, + "migration_id": migration_id, + "configuration": copy.deepcopy(configuration if isinstance(configuration, dict) else {}), + "configuration_fingerprint": fingerprint, + "status": MIGRATION_RESOURCE_STATUS_IN_PROGRESS, + "created_at": timestamp, + "updated_at": timestamp, + "completed_at": None, + "source_cutoff_at": timestamp, + "resume_count": 0, + "resources": {}, + "preflight": {}, + "capacity": {}, + "totals": {}, + } + + +def get_migration_resource(state, resource_name): + """Return a resource checkpoint without allocating a missing record.""" + resources = state.get("resources") if isinstance(state, dict) else None + return resources.get(resource_name) if isinstance(resources, dict) else None + + +def is_migration_resource_completed(state, resource_name): + """Return whether a resource completed successfully in this migration.""" + resource = get_migration_resource(state, resource_name) + return isinstance(resource, dict) and resource.get("status") == MIGRATION_RESOURCE_STATUS_COMPLETED + + +def start_migration_resource(state, resource_name, current_time=None): + """Start or resume one resource while preserving its last durable counters.""" + if not isinstance(state, dict): + raise ValueError("Migration state must be a dictionary.") + timestamp = _utc_now_iso(current_time) + resources = state.setdefault("resources", {}) + existing = resources.get(resource_name) if isinstance(resources, dict) else None + if isinstance(existing, dict) and existing.get("status") == MIGRATION_RESOURCE_STATUS_COMPLETED: + return existing + + progress = copy.deepcopy(existing.get("progress") if isinstance(existing, dict) else {}) + attempts = int(existing.get("attempts") or 0) + 1 if isinstance(existing, dict) else 1 + resource = { + "status": MIGRATION_RESOURCE_STATUS_IN_PROGRESS, + "attempts": attempts, + "started_at": existing.get("started_at") if isinstance(existing, dict) and existing.get("started_at") else timestamp, + "attempt_started_at": timestamp, + "updated_at": timestamp, + "completed_at": None, + "last_error": None, + "progress": progress, + "result": {}, + } + resources[resource_name] = resource + state["updated_at"] = timestamp + return resource + + +def update_migration_resource(state, resource_name, progress, current_time=None): + """Persist counters for an active resource checkpoint.""" + resource = get_migration_resource(state, resource_name) + if not isinstance(resource, dict): + raise ValueError(f"Migration resource '{resource_name}' was not started.") + if resource.get("status") == MIGRATION_RESOURCE_STATUS_COMPLETED: + raise ValueError(f"Migration resource '{resource_name}' is already completed.") + timestamp = _utc_now_iso(current_time) + resource["progress"] = copy.deepcopy(progress if isinstance(progress, dict) else {}) + resource["updated_at"] = timestamp + state["updated_at"] = timestamp + return resource + + +def complete_migration_resource(state, resource_name, result=None, current_time=None): + """Mark a resource completed only after its final target writes are known.""" + resource = get_migration_resource(state, resource_name) + if not isinstance(resource, dict): + raise ValueError(f"Migration resource '{resource_name}' was not started.") + timestamp = _utc_now_iso(current_time) + resource.update({ + "status": MIGRATION_RESOURCE_STATUS_COMPLETED, + "updated_at": timestamp, + "completed_at": timestamp, + "last_error": None, + "result": copy.deepcopy(result if isinstance(result, dict) else {}), + }) + state["updated_at"] = timestamp + return resource + + +def fail_migration_resource(state, resource_name, error_message, current_time=None): + """Persist an interrupted resource so the same job can resume it later.""" + resource = get_migration_resource(state, resource_name) + if not isinstance(resource, dict): + resource = start_migration_resource(state, resource_name, current_time=current_time) + timestamp = _utc_now_iso(current_time) + resource.update({ + "status": MIGRATION_RESOURCE_STATUS_FAILED, + "updated_at": timestamp, + "last_error": str(error_message or "Migration resource failed."), + }) + state["updated_at"] = timestamp + return resource + + +def build_transfer_metrics( + started_at, + copied_count=0, + skipped_count=0, + failed_count=0, + byte_count=0, + request_units=0.0, + current_time=None, +): + """Calculate bounded, JSON-safe counters and transfer rates for job details.""" + started = _parse_utc_datetime(started_at) + now = _parse_utc_datetime(current_time) if current_time is not None else datetime.now(timezone.utc) + elapsed_seconds = max(0.001, (now - started).total_seconds()) if started else 0.001 + copied = max(0, int(copied_count or 0)) + skipped = max(0, int(skipped_count or 0)) + failed = max(0, int(failed_count or 0)) + bytes_transferred = max(0, int(byte_count or 0)) + consumed_request_units = max(0.0, float(request_units or 0.0)) + return { + "copied_count": copied, + "skipped_count": skipped, + "failed_count": failed, + "processed_count": copied + skipped + failed, + "bytes": bytes_transferred, + "request_units": round(consumed_request_units, 3), + "elapsed_seconds": round(elapsed_seconds, 3), + "items_per_second": round((copied + skipped + failed) / elapsed_seconds, 3), + "bytes_per_second": round(bytes_transferred / elapsed_seconds, 3), + "request_units_per_second": round(consumed_request_units / elapsed_seconds, 3), + } \ No newline at end of file diff --git a/application/single_app/functions_data_management_search_write_fence.py b/application/single_app/functions_data_management_search_write_fence.py new file mode 100644 index 000000000..cb246230e --- /dev/null +++ b/application/single_app/functions_data_management_search_write_fence.py @@ -0,0 +1,644 @@ +# functions_data_management_search_write_fence.py +"""Cross-app Azure AI Search write fencing for Data Management migrations.""" + +import copy +import time +import uuid +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone + +from azure.core import MatchConditions +from azure.core.exceptions import ResourceNotFoundError + + +DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID = "data_management_search_write_gate_global" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE = "data_management_search_write_gate" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_OPEN = "open" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING = "closing" +DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_FROZEN = "frozen" +DATA_MANAGEMENT_SEARCH_WRITE_REQUEST_TIMEOUT_SECONDS = 30 +DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS = 150 +DATA_MANAGEMENT_SEARCH_WRITE_GATE_POLL_SECONDS = 0.1 +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID = ( + "data_management_target_migration_coordinator_global" +) +DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE = ( + "data_management_target_migration_coordinator" +) + + +class DataManagementSearchWriteGateError(RuntimeError): + """Raised when the durable Search write gate cannot be coordinated safely.""" + + +class DataManagementSearchWritesFrozenError(DataManagementSearchWriteGateError): + """Raised when a target migration has frozen SimpleChat Search writes.""" + + +class DataManagementSearchWriteFenceLostError(DataManagementSearchWriteGateError): + """Raised when a migration no longer owns its target Search write fence.""" + + +class DataManagementTargetMigrationCoordinatorError(DataManagementSearchWriteGateError): + """Raised when a target migration coordinator cannot be acquired safely.""" + + +class DataManagementTargetMigrationCoordinatorLostError( + DataManagementTargetMigrationCoordinatorError +): + """Raised when a migration no longer owns its target coordinator lease.""" + + +def _now_utc(): + return datetime.now(timezone.utc) + + +def _parse_utc_datetime(value): + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except (TypeError, ValueError): + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _safe_int(value, default=0, minimum=0): + try: + normalized = int(value) + except (TypeError, ValueError): + return default + return max(minimum, normalized) + + +def _is_not_found_error(exc): + return ( + getattr(exc, "status_code", None) == 404 or + isinstance(exc, (KeyError, ResourceNotFoundError)) + ) + + +def _is_conflict_error(exc): + return getattr(exc, "status_code", None) in {409, 412} + + +def _read_gate(container): + try: + gate = container.read_item( + item=DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + partition_key=DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + ) + except Exception as exc: + if _is_not_found_error(exc): + return None + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate could not be read." + ) from exc + if not isinstance(gate, dict): + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate has an invalid record." + ) + return gate + + +def _new_open_gate(now=None): + timestamp = now or _now_utc() + return { + "id": DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + "type": DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE, + "state": DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_OPEN, + "writer_leases": [], + "active_writer_count": 0, + "updated_at": timestamp.isoformat(), + } + + +def _create_or_read_gate(container): + gate = _read_gate(container) + if gate is not None: + return gate + try: + return container.create_item(body=_new_open_gate()) + except Exception as exc: + if _is_conflict_error(exc): + gate = _read_gate(container) + if gate is not None: + return gate + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate could not be created." + ) from exc + + +def _replace_gate(container, existing_gate, replacement_gate): + try: + return container.replace_item( + item=DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + body=replacement_gate, + etag=existing_gate.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if _is_conflict_error(exc): + return None + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate could not be updated." + ) from exc + + +def _active_writer_leases(gate, now=None): + timestamp = now or _now_utc() + active_leases = [] + for lease in (gate or {}).get("writer_leases") or []: + if not isinstance(lease, dict) or not lease.get("token"): + continue + expires_at = _parse_utc_datetime(lease.get("expires_at")) + if expires_at and expires_at > timestamp: + active_leases.append({ + "token": str(lease.get("token")), + "expires_at": expires_at.isoformat(), + }) + return active_leases + + +def _is_active_migration_fence(gate, now=None): + timestamp = now or _now_utc() + expires_at = _parse_utc_datetime((gate or {}).get("expires_at")) + return bool( + (gate or {}).get("type") == DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE and + (gate or {}).get("state") in { + DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING, + DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_FROZEN, + } and + expires_at and + expires_at > timestamp + ) + + +def _open_expired_gate(container, gate, now=None): + timestamp = now or _now_utc() + if _is_active_migration_fence(gate, timestamp): + return gate + replacement = _new_open_gate(timestamp) + replacement["writer_leases"] = _active_writer_leases(gate, timestamp) + replacement["active_writer_count"] = len(replacement["writer_leases"]) + return _replace_gate(container, gate, replacement) + + +def acquire_data_management_search_write_slot(container): + """Reserve one bounded target Search write before issuing the data-plane request.""" + for _attempt in range(12): + now = _now_utc() + gate = _create_or_read_gate(container) + if gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE: + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate has an unexpected record type." + ) + if _is_active_migration_fence(gate, now): + raise DataManagementSearchWritesFrozenError( + "AI Search writes are temporarily frozen while a Data Management migration is running." + ) + if gate.get("state") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_OPEN: + if _open_expired_gate(container, gate, now) is None: + continue + continue + lease_token = uuid.uuid4().hex + replacement = copy.deepcopy(gate) + active_leases = _active_writer_leases(gate, now) + active_leases.append({ + "token": lease_token, + "expires_at": ( + now + timedelta(seconds=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS) + ).isoformat(), + }) + replacement.update({ + "writer_leases": active_leases, + "active_writer_count": len(active_leases), + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is not None: + return lease_token + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate changed too often to reserve a write slot." + ) + + +def release_data_management_search_write_slot(container, lease_token): + """Release one target Search write slot without masking the caller's result.""" + for _attempt in range(12): + try: + gate = _read_gate(container) + except DataManagementSearchWriteGateError: + return False + if gate is None or gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE: + return False + now = _now_utc() + replacement = copy.deepcopy(gate) + active_leases = [ + lease + for lease in _active_writer_leases(gate, now) + if lease.get("token") != lease_token + ] + replacement.update({ + "writer_leases": active_leases, + "active_writer_count": len(active_leases), + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is not None: + return True + return False + + +@contextmanager +def hold_data_management_search_write_slot(container): + """Hold a target Search write slot until a response is known or its ambiguity lease expires.""" + lease_token = acquire_data_management_search_write_slot(container) + response_confirmed = False + try: + yield + response_confirmed = True + finally: + if response_confirmed: + release_data_management_search_write_slot(container, lease_token) + + +def acquire_data_management_search_write_fence( + container, + migration_id, + lease_seconds, + heartbeat_callback=None, +): + """Close target SimpleChat Search writes, drain active slots, and freeze the gate.""" + normalized_migration_id = str(migration_id or "").strip() + if not normalized_migration_id: + raise DataManagementSearchWriteGateError( + "A migration ID is required to acquire the target Search write fence." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + fence_token = uuid.uuid4().hex + last_heartbeat = 0.0 + last_fence_renewal = 0.0 + while True: + now = _now_utc() + gate = _create_or_read_gate(container) + if gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE: + raise DataManagementSearchWriteGateError( + "The Data Management Search write gate has an unexpected record type." + ) + state = gate.get("state") + same_fence = ( + gate.get("migration_id") == normalized_migration_id and + gate.get("fence_token") == fence_token + ) + if _is_active_migration_fence(gate, now) and not same_fence: + raise DataManagementSearchWritesFrozenError( + "Another migration is already freezing target AI Search writes." + ) + if state != DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING or not same_fence: + replacement = copy.deepcopy(gate) + active_leases = _active_writer_leases(gate, now) + replacement.update({ + "state": DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING, + "migration_id": normalized_migration_id, + "fence_token": fence_token, + "expires_at": ( + now + timedelta(seconds=normalized_lease_seconds) + ).isoformat(), + "writer_leases": active_leases, + "active_writer_count": len(active_leases), + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is None: + continue + gate = replacement + if ( + state == DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_CLOSING and + same_fence and + time.monotonic() - last_fence_renewal >= 2.0 + ): + replacement = copy.deepcopy(gate) + replacement.update({ + "expires_at": ( + now + timedelta(seconds=normalized_lease_seconds) + ).isoformat(), + "updated_at": now.isoformat(), + }) + persisted = _replace_gate(container, gate, replacement) + if persisted is None: + continue + gate = persisted + last_fence_renewal = time.monotonic() + active_leases = _active_writer_leases(gate, now) + if not active_leases: + replacement = copy.deepcopy(gate) + replacement.update({ + "state": DATA_MANAGEMENT_SEARCH_WRITE_GATE_STATE_FROZEN, + "writer_leases": [], + "active_writer_count": 0, + "updated_at": now.isoformat(), + }) + if _replace_gate(container, gate, replacement) is not None: + return { + "id": DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID, + "migration_id": normalized_migration_id, + "fence_token": fence_token, + "lease_seconds": normalized_lease_seconds, + "expires_at": replacement.get("expires_at"), + } + continue + if callable(heartbeat_callback) and time.monotonic() - last_heartbeat >= 2.0: + heartbeat_callback() + last_heartbeat = time.monotonic() + time.sleep(DATA_MANAGEMENT_SEARCH_WRITE_GATE_POLL_SECONDS) + + +def renew_data_management_search_write_fence(container, fence, lease_seconds): + """Renew a migration-owned frozen target Search gate.""" + if not isinstance(fence, dict) or not fence.get("fence_token"): + return None + now = _now_utc() + gate = _read_gate(container) + if ( + not isinstance(gate, dict) or + gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE or + gate.get("migration_id") != fence.get("migration_id") or + gate.get("fence_token") != fence.get("fence_token") or + not _is_active_migration_fence(gate, now) + ): + raise DataManagementSearchWriteFenceLostError( + "The target AI Search write fence was lost or expired." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=fence.get("lease_seconds"), + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + replacement = copy.deepcopy(gate) + replacement.update({ + "expires_at": (now + timedelta(seconds=normalized_lease_seconds)).isoformat(), + "updated_at": now.isoformat(), + }) + renewed = _replace_gate(container, gate, replacement) + if renewed is None: + raise DataManagementSearchWriteFenceLostError( + "The target AI Search write fence changed during renewal." + ) + fence["expires_at"] = renewed.get("expires_at", replacement.get("expires_at")) + fence["lease_seconds"] = normalized_lease_seconds + return fence + + +def release_data_management_search_write_fence(container, fence): + """Open only the exact target Search gate owned by the completed migration.""" + if not isinstance(fence, dict) or not fence.get("fence_token"): + return False + for _attempt in range(12): + gate = _read_gate(container) + if ( + not isinstance(gate, dict) or + gate.get("type") != DATA_MANAGEMENT_SEARCH_WRITE_GATE_TYPE or + gate.get("migration_id") != fence.get("migration_id") or + gate.get("fence_token") != fence.get("fence_token") + ): + return False + now = _now_utc() + replacement = _new_open_gate(now) + replacement["writer_leases"] = _active_writer_leases(gate, now) + replacement["active_writer_count"] = len(replacement["writer_leases"]) + if _replace_gate(container, gate, replacement) is not None: + return True + return False + + +def _read_target_migration_coordinator(container): + try: + coordinator = container.read_item( + item=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + partition_key=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + ) + except Exception as exc: + if _is_not_found_error(exc): + return None + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator could not be read." + ) from exc + if not isinstance(coordinator, dict): + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator has an invalid record." + ) + return coordinator + + +def _replace_target_migration_coordinator(container, existing, replacement): + try: + return container.replace_item( + item=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + body=replacement, + etag=existing.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if _is_conflict_error(exc): + return None + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator could not be updated." + ) from exc + + +def _target_migration_coordinator_is_active(coordinator, now=None): + timestamp = now or _now_utc() + expires_at = _parse_utc_datetime((coordinator or {}).get("expires_at")) + return bool( + (coordinator or {}).get("type") == + DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE and + expires_at and + expires_at > timestamp + ) + + +def _new_target_migration_coordinator(migration_id, lock_token, lease_seconds, now=None): + timestamp = now or _now_utc() + return { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "type": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE, + "migration_id": str(migration_id), + "lock_token": str(lock_token), + "acquired_at": timestamp.isoformat(), + "updated_at": timestamp.isoformat(), + "lease_seconds": int(lease_seconds), + "expires_at": (timestamp + timedelta(seconds=int(lease_seconds))).isoformat(), + } + + +def acquire_data_management_target_migration_coordinator( + container, + migration_id, + lease_seconds, + existing_lock=None, +): + """Acquire one destination-wide coordinator before any migration mutation begins.""" + normalized_migration_id = str(migration_id or "").strip() + if not normalized_migration_id: + raise DataManagementTargetMigrationCoordinatorError( + "A migration ID is required to acquire the target migration coordinator." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + existing_lock = existing_lock if isinstance(existing_lock, dict) else {} + existing_token = str(existing_lock.get("lock_token") or "").strip() + for _attempt in range(12): + now = _now_utc() + coordinator = _read_target_migration_coordinator(container) + if coordinator is None: + lock_token = existing_token or uuid.uuid4().hex + replacement = _new_target_migration_coordinator( + normalized_migration_id, + lock_token, + normalized_lease_seconds, + now, + ) + try: + persisted = container.create_item(body=replacement) + except Exception as exc: + if _is_conflict_error(exc): + continue + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator could not be created." + ) from exc + return { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "migration_id": normalized_migration_id, + "lock_token": lock_token, + "lease_seconds": normalized_lease_seconds, + "expires_at": persisted.get("expires_at", replacement["expires_at"]), + } + + if coordinator.get("type") != DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE: + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator has an unexpected record type." + ) + if _target_migration_coordinator_is_active(coordinator, now): + if ( + coordinator.get("migration_id") == normalized_migration_id and + existing_token and + coordinator.get("lock_token") == existing_token + ): + reused_lock = { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "migration_id": normalized_migration_id, + "lock_token": existing_token, + "lease_seconds": normalized_lease_seconds, + "expires_at": coordinator.get("expires_at"), + } + return renew_data_management_target_migration_coordinator( + container, + reused_lock, + normalized_lease_seconds, + ) + raise DataManagementTargetMigrationCoordinatorError( + "Another SimpleChat source is actively migrating to this destination." + ) + + replacement = _new_target_migration_coordinator( + normalized_migration_id, + existing_token or uuid.uuid4().hex, + normalized_lease_seconds, + now, + ) + persisted = _replace_target_migration_coordinator( + container, + coordinator, + replacement, + ) + if persisted is None: + continue + return { + "id": DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + "migration_id": normalized_migration_id, + "lock_token": replacement["lock_token"], + "lease_seconds": normalized_lease_seconds, + "expires_at": persisted.get("expires_at", replacement["expires_at"]), + } + raise DataManagementTargetMigrationCoordinatorError( + "The target migration coordinator changed too often to acquire safely." + ) + + +def renew_data_management_target_migration_coordinator(container, lock, lease_seconds): + """Renew an exact target-wide coordinator lease held by this migration.""" + if not isinstance(lock, dict) or not lock.get("lock_token"): + raise DataManagementTargetMigrationCoordinatorLostError( + "The target migration coordinator handle is missing." + ) + coordinator = _read_target_migration_coordinator(container) + now = _now_utc() + if ( + not isinstance(coordinator, dict) or + coordinator.get("type") != DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE or + coordinator.get("migration_id") != lock.get("migration_id") or + coordinator.get("lock_token") != lock.get("lock_token") or + not _target_migration_coordinator_is_active(coordinator, now) + ): + raise DataManagementTargetMigrationCoordinatorLostError( + "The target migration coordinator was lost, superseded, or expired." + ) + normalized_lease_seconds = _safe_int( + lease_seconds, + default=lock.get("lease_seconds"), + minimum=DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS, + ) + replacement = copy.deepcopy(coordinator) + replacement.update({ + "lease_seconds": normalized_lease_seconds, + "updated_at": now.isoformat(), + "expires_at": ( + now + timedelta(seconds=normalized_lease_seconds) + ).isoformat(), + }) + persisted = _replace_target_migration_coordinator( + container, + coordinator, + replacement, + ) + if persisted is None: + raise DataManagementTargetMigrationCoordinatorLostError( + "The target migration coordinator changed during renewal." + ) + lock["lease_seconds"] = normalized_lease_seconds + lock["expires_at"] = persisted.get("expires_at", replacement["expires_at"]) + return lock + + +def release_data_management_target_migration_coordinator(container, lock): + """Release only the exact target coordinator held by a finished migration.""" + if not isinstance(lock, dict) or not lock.get("lock_token"): + return False + try: + coordinator = _read_target_migration_coordinator(container) + except DataManagementTargetMigrationCoordinatorError: + return False + if ( + not isinstance(coordinator, dict) or + coordinator.get("type") != DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_TYPE or + coordinator.get("migration_id") != lock.get("migration_id") or + coordinator.get("lock_token") != lock.get("lock_token") + ): + return False + try: + container.delete_item( + item=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + partition_key=DATA_MANAGEMENT_TARGET_MIGRATION_COORDINATOR_ID, + etag=coordinator.get("_etag"), + match_condition=MatchConditions.IfNotModified, + ) + except Exception as exc: + if _is_conflict_error(exc): + return False + return False + return True \ No newline at end of file diff --git a/application/single_app/functions_documents.py b/application/single_app/functions_documents.py index df5d1da88..b9bdb48f3 100644 --- a/application/single_app/functions_documents.py +++ b/application/single_app/functions_documents.py @@ -19,6 +19,10 @@ sync_document_access_index_for_document_fail_open, validate_document_access_index_shadow, ) +from functions_data_management_search_write_fence import ( + DataManagementSearchWritesFrozenError, + hold_data_management_search_write_slot, +) from functions_visio import build_visio_page_markdown, parse_vsdx_pages from functions_content import * from functions_settings import * @@ -32,6 +36,39 @@ _AUDIO_RUNTIME_CAPABILITIES_CACHE = None + +class DocumentSearchAclProjectionDeferredError(RuntimeError): + """Raised when an authorization-reducing Search ACL update must be retried safely.""" + + +def _search_indexing_results_succeeded(results): + """Return whether Azure AI Search acknowledged every requested document mutation.""" + normalized_results = list(results or []) + return bool(normalized_results) and all( + bool( + result.get("succeeded", result.get("status", False)) + if isinstance(result, dict) else + getattr(result, "succeeded", getattr(result, "status", False)) + ) + for result in normalized_results + ) + + +def _execute_document_search_write(search_client, operation_name, *args, **kwargs): + """Serialize bounded target Search writes with an active Data Management migration fence.""" + kwargs.update({ + "connection_timeout": 30, + "read_timeout": 30, + "retry_total": 0, + }) + with hold_data_management_search_write_slot(cosmos_data_management_jobs_container): + results = getattr(search_client, operation_name)(*args, **kwargs) + if not _search_indexing_results_succeeded(results): + raise RuntimeError( + f"Azure AI Search did not acknowledge every {operation_name} document mutation." + ) + return results + def allowed_file(filename, allowed_extensions=None): if not allowed_extensions: allowed_extensions = ALLOWED_EXTENSIONS @@ -786,7 +823,11 @@ def set_document_chunk_visibility(document_item, active=True): documents_to_update.append(chunk_item) - search_client.upload_documents(documents=documents_to_update) + _execute_document_search_write( + search_client, + "upload_documents", + documents=documents_to_update, + ) return len(documents_to_update) @@ -1333,7 +1374,11 @@ def save_video_chunk( # 3) upload to search index try: debug_print(f"[VIDEO CHUNK] Uploading chunk to search index") - client.upload_documents(documents=[chunk]) + _execute_document_search_write( + client, + "upload_documents", + documents=[chunk], + ) debug_print(f"[VIDEO CHUNK] Upload successful for chunk: {chunk_id}") print(f"[VideoChunk] UPLOAD OK for {chunk_id}", flush=True) except Exception as e: @@ -2675,7 +2720,11 @@ def save_chunks(page_text_content, page_number, file_name, user_id, document_id, else: search_client = CLIENTS["search_client_user"] # Upload as a single-document list - search_client.upload_documents(documents=[chunk_document]) + _execute_document_search_write( + search_client, + "upload_documents", + documents=[chunk_document], + ) except Exception as e: print(f"Error uploading chunk document for document {document_id}: {e}") @@ -2861,7 +2910,11 @@ def save_chunks_batch(chunks_data, user_id, document_id, group_id=None, public_w upload_batch_size = 32 for i in range(0, len(chunk_documents), upload_batch_size): sub_batch = chunk_documents[i:i + upload_batch_size] - search_client.upload_documents(documents=sub_batch) + _execute_document_search_write( + search_client, + "upload_documents", + documents=sub_batch, + ) except Exception as e: log_event(f"[save_chunks_batch] Error uploading batch to AI Search for document {document_id}: {e}", level=logging.ERROR) @@ -3030,7 +3083,11 @@ def update_chunk_metadata(chunk_id, user_id, group_id=None, public_workspace_id= else: chunk_item[field] = kwargs[field] - search_client.upload_documents(documents=[chunk_item]) + _execute_document_search_write( + search_client, + "upload_documents", + documents=[chunk_item], + ) except Exception as e: print(f"Error updating chunk metadata for chunk {chunk_id}: {e}") @@ -3944,7 +4001,11 @@ def delete_document_chunks(document_id, group_id=None, public_workspace_id=None) documents_to_delete = [{"id": doc_id} for doc_id in ids_to_delete] batch = IndexDocumentsBatch() batch.add_delete_actions(documents_to_delete) - result = search_client.index_documents(batch) + result = _execute_document_search_write( + search_client, + "index_documents", + batch, + ) except Exception as e: raise @@ -3955,7 +4016,9 @@ def delete_document_version_chunks(document_id, version, group_id=None, public_w search_client = CLIENTS["search_client_public"] if is_public_workspace else CLIENTS["search_client_group"] if is_group else CLIENTS["search_client_user"] - search_client.delete_documents( + _execute_document_search_write( + search_client, + "delete_documents", actions=[ {"@search.action": "delete", "id": chunk['id']} for chunk in search_client.search( @@ -9110,41 +9173,41 @@ def unshare_document_from_user(document_id, owner_user_id, target_user_id): # Remove all entries for the target user (by oid prefix) new_shared_user_ids = [entry for entry in shared_user_ids if not entry.startswith(f"{target_user_id},")] if len(new_shared_user_ids) != len(shared_user_ids): + # A revoked user must never retain search access after the API reports success. + # Update every chunk projection before committing the authoritative Cosmos ACL. + chunks = get_all_chunks(document_id, actual_owner_id) + try: + for chunk in chunks: + chunk_id = chunk.get('id') + if not chunk_id: + continue + update_chunk_metadata( + chunk_id=chunk_id, + user_id=actual_owner_id, + group_id=None, + public_workspace_id=None, + document_id=document_id, + shared_user_ids=new_shared_user_ids + ) + except DataManagementSearchWritesFrozenError as exc: + raise DocumentSearchAclProjectionDeferredError( + "Document access was not changed because target Search writes are temporarily frozen. Retry after the migration finishes." + ) from exc + document_item['shared_user_ids'] = new_shared_user_ids document_item['last_updated'] = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ') - # Update the document _upsert_document_and_sync_access_index( cosmos_user_documents_container, document_item, operation='document_unshared_from_user', ) - # Update all chunks with the new shared_user_ids - try: - chunks = get_all_chunks(document_id, actual_owner_id) - for chunk in chunks: - chunk_id = chunk.get('id') - if chunk_id: - try: - update_chunk_metadata( - chunk_id=chunk_id, - user_id=actual_owner_id, - group_id=None, - public_workspace_id=None, - document_id=document_id, - shared_user_ids=new_shared_user_ids - ) - except Exception as chunk_e: - print(f"Warning: Failed to update chunk {chunk_id}: {chunk_e}") - # Continue with other chunks - except Exception as e: - print(f"Warning: Failed to update chunks for document {document_id}: {e}") - # Don't fail the whole operation if chunk update fails - return True except CosmosResourceNotFoundError: return False + except DocumentSearchAclProjectionDeferredError: + raise except Exception as e: print(f"Error unsharing document {document_id}: {e}") return False diff --git a/application/single_app/functions_migration_provenance.py b/application/single_app/functions_migration_provenance.py new file mode 100644 index 000000000..92134623d --- /dev/null +++ b/application/single_app/functions_migration_provenance.py @@ -0,0 +1,251 @@ +# functions_migration_provenance.py +"""Durable provenance metadata shared by in-app migration destinations.""" + +from datetime import datetime, timedelta, timezone +import uuid + + +MIGRATION_PROVENANCE_STATUS_SUCCEEDED = "succeeded" +COSMOS_MIGRATION_PROVENANCE_FIELD = "simplechatMigration" +SEARCH_MIGRATION_ID_FIELD = "simplechatMigrationId" +SEARCH_MIGRATED_AT_FIELD = "simplechatMigratedAtUtc" +SEARCH_MIGRATION_STATUS_FIELD = "simplechatMigrationStatus" +SEARCH_MIGRATION_SOURCE_HASH_FIELD = "simplechatMigrationSourceHash" +SEARCH_MIGRATION_SOURCE_VERSION_FIELD = "simplechatMigrationSourceVersion" +BLOB_MIGRATION_ID_METADATA_KEY = "simplechatMigrationId" +BLOB_MIGRATED_AT_METADATA_KEY = "simplechatMigratedAtUtc" +BLOB_MIGRATION_STATUS_METADATA_KEY = "simplechatMigrationStatus" +BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY = "simplechatMigrationSourceHash" +BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY = "simplechatMigrationSourceVersion" +BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY = "simplechatMigrationScopeHash" + + +def _parse_utc_datetime(value): + """Return an aware UTC timestamp when the supplied value is valid.""" + if isinstance(value, datetime): + parsed = value + else: + try: + parsed = datetime.fromisoformat(str(value or "").replace("Z", "+00:00")) + except (TypeError, ValueError): + return None + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=timezone.utc) + return parsed.astimezone(timezone.utc) + + +def _utc_iso(value): + """Format a UTC datetime using the durable migration metadata format.""" + parsed = _parse_utc_datetime(value) + if parsed is None: + raise ValueError("Migration timestamp must be a valid UTC datetime.") + return parsed.isoformat() + + +def _record_value(record, *names): + """Read the first available provenance value without changing stored data.""" + if not isinstance(record, dict): + return None + for name in names: + if name in record: + return record.get(name) + return None + + +def create_migration_provenance_context(migration_id=None, migrated_at_utc=None, skip_within_hours=0): + """Create a validated provenance context for one durable migration job.""" + try: + normalized_migration_id = str(uuid.UUID(str(migration_id or uuid.uuid4()))) + except (AttributeError, TypeError, ValueError) as exc: + raise ValueError("Migration ID must be a valid GUID.") from exc + + try: + normalized_skip_hours = int(skip_within_hours) + except (TypeError, ValueError) as exc: + raise ValueError("Migration provenance skip window must be an integer number of hours.") from exc + if normalized_skip_hours < 0 or normalized_skip_hours > 8760: + raise ValueError("Migration provenance skip window must be between 0 and 8760 hours.") + + return { + "migration_id": normalized_migration_id, + "migrated_at_utc": _utc_iso(migrated_at_utc or datetime.now(timezone.utc)), + "skip_within_hours": normalized_skip_hours, + } + + +def build_migration_provenance_record( + context, + source_hash="", + source_version="", + scope_hash="", +): + """Return the neutral, destination-independent success record.""" + record = { + "migrationId": str((context or {}).get("migration_id") or ""), + "migratedAtUtc": str((context or {}).get("migrated_at_utc") or ""), + "status": MIGRATION_PROVENANCE_STATUS_SUCCEEDED, + } + if source_hash: + record["sourceHash"] = str(source_hash) + if source_version: + record["sourceVersion"] = str(source_version) + if scope_hash: + record["scopeHash"] = str(scope_hash) + return record + + +def is_successful_migration_record(record): + """Return whether a destination item has durable successful migration ownership.""" + if not isinstance(record, dict): + return False + status = str(_record_value(record, "status", "migrationStatus") or "").strip().lower() + migration_id = str(_record_value(record, "migrationId", "migration_id") or "").strip() + return status == MIGRATION_PROVENANCE_STATUS_SUCCEEDED and bool(migration_id) + + +def migration_record_matches_source(record, source_hash="", source_version=""): + """Compare an owned destination marker with the current source fingerprint.""" + if not is_successful_migration_record(record): + return False + normalized_hash = str(source_hash or "").strip() + normalized_version = str(source_version or "").strip() + if not normalized_hash and not normalized_version: + return False + if normalized_hash and ( + str(_record_value(record, "sourceHash", "source_hash") or "").strip() != normalized_hash + ): + return False + if normalized_version and ( + str(_record_value(record, "sourceVersion", "source_version") or "").strip() != + normalized_version + ): + return False + return True + + +def should_skip_migration_record(record, context, current_time=None): + """Return whether a successful current or recent migration owns an item.""" + status = str(_record_value(record, "status", "migrationStatus") or "").strip().lower() + if status != MIGRATION_PROVENANCE_STATUS_SUCCEEDED: + return False + + migration_id = str(_record_value(record, "migrationId", "migration_id") or "").strip() + context_migration_id = str((context or {}).get("migration_id") or "").strip() + if migration_id and migration_id.lower() == context_migration_id.lower(): + return True + + try: + skip_within_hours = int((context or {}).get("skip_within_hours") or 0) + except (TypeError, ValueError): + return False + if skip_within_hours <= 0: + return False + + migrated_at = _parse_utc_datetime(_record_value(record, "migratedAtUtc", "migrated_at_utc")) + now = _parse_utc_datetime(current_time or datetime.now(timezone.utc)) + if migrated_at is None or now is None: + return False + return migrated_at >= now - timedelta(hours=skip_within_hours) + + +def add_cosmos_migration_provenance(document, context, source_hash="", source_version=""): + """Add migration provenance to a writable Cosmos document.""" + if not isinstance(document, dict): + raise ValueError("Cosmos migration documents must be dictionaries.") + document[COSMOS_MIGRATION_PROVENANCE_FIELD] = build_migration_provenance_record( + context, + source_hash=source_hash, + source_version=source_version, + ) + return document + + +def get_cosmos_migration_provenance(document): + """Get a Cosmos document's migration record when present.""" + if not isinstance(document, dict): + return None + value = document.get(COSMOS_MIGRATION_PROVENANCE_FIELD) + return value if isinstance(value, dict) else None + + +def add_search_migration_provenance(document, context, source_hash="", source_version=""): + """Add migration provenance to a writable Azure AI Search document.""" + if not isinstance(document, dict): + raise ValueError("AI Search migration documents must be dictionaries.") + record = build_migration_provenance_record( + context, + source_hash=source_hash, + source_version=source_version, + ) + document[SEARCH_MIGRATION_ID_FIELD] = record["migrationId"] + document[SEARCH_MIGRATED_AT_FIELD] = record["migratedAtUtc"] + document[SEARCH_MIGRATION_STATUS_FIELD] = record["status"] + if record.get("sourceHash"): + document[SEARCH_MIGRATION_SOURCE_HASH_FIELD] = record["sourceHash"] + if record.get("sourceVersion"): + document[SEARCH_MIGRATION_SOURCE_VERSION_FIELD] = record["sourceVersion"] + return document + + +def get_search_migration_provenance(document): + """Translate Azure AI Search provenance fields into the neutral record.""" + if not isinstance(document, dict): + return None + record = { + "migrationId": document.get(SEARCH_MIGRATION_ID_FIELD), + "migratedAtUtc": document.get(SEARCH_MIGRATED_AT_FIELD), + "status": document.get(SEARCH_MIGRATION_STATUS_FIELD), + } + if document.get(SEARCH_MIGRATION_SOURCE_HASH_FIELD): + record["sourceHash"] = document.get(SEARCH_MIGRATION_SOURCE_HASH_FIELD) + if document.get(SEARCH_MIGRATION_SOURCE_VERSION_FIELD): + record["sourceVersion"] = document.get(SEARCH_MIGRATION_SOURCE_VERSION_FIELD) + return record + + +def merge_blob_migration_metadata( + metadata, + context, + source_hash="", + source_version="", + scope_hash="", +): + """Preserve existing Blob metadata while appending migration provenance.""" + merged = { + str(key): str(value) + for key, value in (metadata or {}).items() + if key is not None and value is not None + } + record = build_migration_provenance_record( + context, + source_hash=source_hash, + source_version=source_version, + scope_hash=scope_hash, + ) + merged[BLOB_MIGRATION_ID_METADATA_KEY] = record["migrationId"] + merged[BLOB_MIGRATED_AT_METADATA_KEY] = record["migratedAtUtc"] + merged[BLOB_MIGRATION_STATUS_METADATA_KEY] = record["status"] + if record.get("sourceHash"): + merged[BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY] = record["sourceHash"] + if record.get("sourceVersion"): + merged[BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY] = record["sourceVersion"] + if record.get("scopeHash"): + merged[BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY] = record["scopeHash"] + return merged + + +def get_blob_migration_provenance(metadata): + """Translate Blob metadata provenance fields into the neutral record.""" + values = metadata or {} + record = { + "migrationId": values.get(BLOB_MIGRATION_ID_METADATA_KEY), + "migratedAtUtc": values.get(BLOB_MIGRATED_AT_METADATA_KEY), + "status": values.get(BLOB_MIGRATION_STATUS_METADATA_KEY), + } + if values.get(BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY): + record["sourceHash"] = values.get(BLOB_MIGRATION_SOURCE_HASH_METADATA_KEY) + if values.get(BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY): + record["sourceVersion"] = values.get(BLOB_MIGRATION_SOURCE_VERSION_METADATA_KEY) + if values.get(BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY): + record["scopeHash"] = values.get(BLOB_MIGRATION_SCOPE_HASH_METADATA_KEY) + return record \ No newline at end of file diff --git a/application/single_app/route_backend_data_management.py b/application/single_app/route_backend_data_management.py index e1e1de7b0..c1992ee4a 100644 --- a/application/single_app/route_backend_data_management.py +++ b/application/single_app/route_backend_data_management.py @@ -3,7 +3,7 @@ import logging -from flask import current_app, jsonify, request, session +from flask import Response, current_app, jsonify, request, session from functions_activity_logging import log_general_admin_action from functions_appinsights import log_event @@ -15,17 +15,23 @@ DATA_MANAGEMENT_OPERATION_RESTORE, DataManagementCosmosEditorError, DataManagementSettingsValidationError, + export_data_management_migration_manifest, generate_data_management_encryption_key, get_data_management_cosmos_editor_containers, get_data_management_cosmos_editor_document, get_data_management_backup_summary, get_data_management_job_detail, + get_data_management_job_progress, get_data_management_jobs, get_data_management_migration_catalog, get_data_management_settings, log_data_management_cosmos_editor_activity, + preview_data_management_migration_plan, queue_data_management_job, query_data_management_cosmos_editor_documents, + request_data_management_migration_cancellation, + resolve_data_management_migration_manifest_item, + retry_data_management_migration_job, sanitize_data_management_job_for_admin, sanitize_data_management_settings_for_admin, save_data_management_cosmos_editor_document, @@ -182,8 +188,12 @@ def test_admin_data_management_storage(): def test_admin_data_management_target_cosmos(): payload = request.get_json(silent=True) or {} settings_payload = payload.get("settings") if isinstance(payload.get("settings"), dict) else None + migration_plan = payload.get("migration_plan") if isinstance(payload.get("migration_plan"), dict) else None try: - result = test_target_cosmos_connection(settings=settings_payload) + result = test_target_cosmos_connection( + settings=settings_payload, + migration_plan=migration_plan, + ) except Exception as exc: log_event( "[DataManagement] Target Cosmos connection test failed.", @@ -406,6 +416,116 @@ def get_admin_data_management_job_detail(job_id): return jsonify({"success": False, "error": "Data Management job was not found."}), 404 return jsonify({"success": True, **detail}), 200 + @bp.route("/api/admin/data-management/jobs//progress", methods=["GET"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def get_admin_data_management_job_progress(job_id): + progress = get_data_management_job_progress(job_id) + if not progress: + return jsonify({"success": False, "error": "Data Management job was not found."}), 404 + return jsonify({"success": True, "job": progress}), 200 + + @bp.route("/api/admin/data-management/jobs//migration-manifest", methods=["GET"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def download_admin_data_management_migration_manifest(job_id): + statuses = { + status.strip().lower() + for status in str(request.args.get("statuses") or "").split(",") + if status.strip() + } + try: + export = export_data_management_migration_manifest(job_id, statuses=statuses) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 404 + response = Response( + export["content"], + status=200, + content_type="application/x-ndjson; charset=utf-8", + ) + suffix = "-failures" if statuses else "" + response.headers["Content-Disposition"] = ( + f'attachment; filename="migration-{job_id}{suffix}.jsonl"' + ) + response.headers["X-Migration-Manifest-Entries"] = str(export["entry_count"]) + return response + + @bp.route( + "/api/admin/data-management/jobs//migration-manifest/items/", + methods=["GET"], + ) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def resolve_admin_data_management_migration_manifest_item(job_id, item_ref): + try: + item = resolve_data_management_migration_manifest_item(job_id, item_ref) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 404 + return jsonify({"success": True, "item": item}), 200 + + @bp.route("/api/admin/data-management/jobs//retry", methods=["POST"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def retry_admin_data_management_migration_job(job_id): + try: + job = retry_data_management_migration_job(job_id) + submitted = submit_data_management_job(current_app._get_current_object(), job.get("id")) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 400 + except Exception as exc: + log_event( + "[DataManagement] Failed to retry migration job.", + {"job_id": job_id, "error": str(exc)}, + level=logging.ERROR, + exceptionTraceback=True, + ) + return jsonify({"success": False, "error": "Migration retry could not be queued."}), 400 + + _log_data_management_admin_action( + "data_management_migration_retry_queued", + "Queued a Data Management migration retry from durable checkpoints.", + {"job_id": job.get("id"), "submitted": submitted}, + ) + public_job = sanitize_data_management_job_for_admin(job) + public_job["submitted_to_executor"] = submitted + return jsonify({"success": True, "job": public_job}), 202 + + @bp.route("/api/admin/data-management/jobs//cancel", methods=["POST"]) + @swagger_route(security=get_auth_security()) + @login_required + @admin_required + def cancel_admin_data_management_migration_job(job_id): + payload = request.get_json(silent=True) or {} + admin_user_id, admin_email = _get_admin_context() + try: + job = request_data_management_migration_cancellation( + job_id, + requested_by=admin_user_id, + requested_by_email=admin_email, + reason=payload.get("reason"), + ) + except DataManagementSettingsValidationError as exc: + return jsonify({"success": False, "error": str(exc)}), 400 + except Exception as exc: + log_event( + "[DataManagement] Failed to request migration cancellation.", + {"job_id": job_id, "error": str(exc)}, + level=logging.ERROR, + exceptionTraceback=True, + ) + return jsonify({"success": False, "error": "Migration cancellation could not be requested."}), 400 + + _log_data_management_admin_action( + "data_management_migration_cancel_requested", + "Requested cancellation of a Data Management migration.", + {"job_id": job.get("id"), "status": job.get("status")}, + ) + return jsonify({"success": True, "job": sanitize_data_management_job_for_admin(job)}), 202 + @bp.route("/api/admin/data-management/backups", methods=["GET"]) @swagger_route(security=get_auth_security()) @login_required @@ -436,9 +556,26 @@ def summarize_admin_data_management_migration(): payload = request.get_json(silent=True) or {} try: summary = summarize_data_management_migration_plan(payload) + preview = None + if payload.get("include_inventory") is True: + preview = preview_data_management_migration_plan( + get_data_management_settings(), + payload, + ) except DataManagementSettingsValidationError as exc: return jsonify({"success": False, "error": str(exc)}), 400 - return jsonify({"success": True, "summary": summary}), 200 + except Exception as exc: + log_event( + "[DataManagement] Migration inventory preview failed.", + {"error": str(exc)}, + level=logging.ERROR, + exceptionTraceback=True, + ) + return jsonify({ + "success": False, + "error": "Migration inventory preview could not be completed.", + }), 400 + return jsonify({"success": True, "summary": summary, "preview": preview}), 200 @bp.route("/api/admin/data-management/jobs", methods=["POST"]) @swagger_route(security=get_auth_security()) @@ -458,12 +595,17 @@ def create_admin_data_management_job(): admin_user_id, admin_email = _get_admin_context() try: + job_options = ( + dict(payload.get("options")) + if isinstance(payload.get("options"), dict) + else {} + ) job = queue_data_management_job( operation, backup_type=backup_type, requested_by=admin_user_id, requested_by_email=admin_email, - options=payload.get("options") if isinstance(payload.get("options"), dict) else {}, + options=job_options, ) submitted = submit_data_management_job(current_app._get_current_object(), job.get("id")) except Exception as exc: diff --git a/application/single_app/route_backend_documents.py b/application/single_app/route_backend_documents.py index f7d0a6b77..18ac783d2 100644 --- a/application/single_app/route_backend_documents.py +++ b/application/single_app/route_backend_documents.py @@ -2052,6 +2052,10 @@ def api_unshare_document(document_id): return jsonify({'message': 'Document unshared successfully'}), 200 else: return jsonify({'error': 'Failed to unshare document'}), 500 + except DocumentSearchAclProjectionDeferredError as exc: + response = jsonify({'error': str(exc)}) + response.headers['Retry-After'] = '150' + return response, 503 except exceptions.CosmosResourceNotFoundError: return jsonify({'error': 'Document not found or access denied'}), 404 except Exception as e: diff --git a/application/single_app/static/js/admin/admin_data_management.js b/application/single_app/static/js/admin/admin_data_management.js index 6476ec031..5c79699f9 100644 --- a/application/single_app/static/js/admin/admin_data_management.js +++ b/application/single_app/static/js/admin/admin_data_management.js @@ -7,6 +7,7 @@ const backupStorageAuthManagedIdentity = "managed_identity"; const backupStorageAuthConnectionString = "connection_string"; const targetCosmosDatabaseName = "SimpleChat"; const cosmosEditorConfirmationPhrase = "I understand this can damage system data"; +const migrationMirrorConfirmationPhrase = "MIRROR WITH DELETIONS"; const elements = {}; let dataManagementModified = false; let storedBackupConnectionStringAvailable = false; @@ -22,8 +23,9 @@ let cosmosEditorContinuationToken = null; let cosmosEditorResultCount = 0; let cosmosEditorSelectedDocument = null; let cosmosEditorPendingDocument = null; +let pendingMigrationCancellationJobId = null; -const jobDetailRefreshIntervalMs = 4000; +const jobDetailRefreshIntervalMs = 2000; const activeJobStatuses = new Set(["queued", "running"]); const migrationTargetTypes = ["users", "groups", "public_workspaces"]; const migrationSelections = { @@ -83,6 +85,8 @@ function bindElements() { "data_management_target_cosmos_database", "data-management-target-cosmos-key-field", "data_management_target_cosmos_key", + "data_management_target_cosmos_subscription_id", + "data_management_target_cosmos_resource_group", "data-management-test-target-cosmos-btn", "data-management-target-ai-search-section", "data_management_target_ai_search_auth", @@ -98,6 +102,15 @@ function bindElements() { "data_management_target_ec_connection_string", "data-management-test-target-ec-storage-btn", "data-management-refresh-migration-summary-btn", + "data-management-migration-mode-section", + "data_management_migration_mode_new_only", + "data_management_migration_mode_delta_upsert", + "data_management_migration_mode_mirror_with_deletions", + "data-management-migration-mode-description", + "data-management-migration-baseline-field", + "data_management_migration_baseline_job_id", + "data-management-migration-mirror-confirmation", + "data_management_migration_mirror_confirmation_phrase", "data_management_migration_users_mode", "data_management_migration_users_search", "data-management-search-users-btn", @@ -117,7 +130,16 @@ function bindElements() { "data-management-migration-public-workspaces-selected", "data_management_migration_public_workspaces_documents", "data_management_migration_include_ai_search", + "data-management-migration-search-write-freeze", + "data_management_migration_target_search_writes_frozen", "data_management_migration_include_source_blobs", + "data-management-test-migration-access-btn", + "data_management_migration_max_parallel_operations", + "data_management_migration_retry_count", + "data_management_migration_skip_recent_within_hours", + "data_management_migration_temporary_destination_ru_enabled", + "data-management-migration-temporary-ru-field", + "data_management_migration_temporary_destination_ru", "data-management-migration-summary", "data-management-migration-preview-btn", "data-management-execute-migration-btn", @@ -139,11 +161,15 @@ function bindElements() { "data-management-job-detail-subtitle", "data-management-job-detail-refresh-state", "data-management-job-detail-summary", + "data-management-job-detail-actions", "data-management-job-detail-progress", "data-management-job-items-tbody", "data-management-job-artifacts-tbody", "data-management-job-manifest-detail", "data-management-job-warnings", + "data-management-migration-cancel-modal", + "data-management-migration-cancel-message", + "data-management-confirm-migration-cancel-btn", "data-management-cosmos-editor-section", "data-management-cosmos-editor-open-danger-btn", "data-management-cosmos-editor-locked-message", @@ -205,6 +231,7 @@ function bindEvents() { elements.dataManagementGenerateKeyBtn?.addEventListener("click", generateEncryptionKey); elements.dataManagementTestStorageBtn?.addEventListener("click", testBackupStorage); elements.dataManagementTestTargetCosmosBtn?.addEventListener("click", testTargetCosmos); + elements.dataManagementTestMigrationAccessBtn?.addEventListener("click", testMigrationAccess); elements.dataManagementTestTargetSearchBtn?.addEventListener("click", testTargetSearch); elements.dataManagementTestTargetEcStorageBtn?.addEventListener("click", testTargetEnhancedCitationStorage); elements.dataManagementRunFullBackupBtn?.addEventListener("click", () => queueBackup("full")); @@ -219,6 +246,10 @@ function bindEvents() { elements.dataManagementViewAllBackupsBtn?.addEventListener("click", () => setBackupFilter("all")); elements.dataManagementRefreshJobsBtn?.addEventListener("click", loadDataManagementJobs); elements.dataManagementJobDetailModal?.addEventListener("hidden.bs.modal", () => stopJobDetailAutoRefresh({ clearJob: true })); + elements.dataManagementMigrationCancelModal?.addEventListener("hidden.bs.modal", () => { + pendingMigrationCancellationJobId = null; + }); + elements.dataManagementConfirmMigrationCancelBtn?.addEventListener("click", requestMigrationCancellation); elements.dataManagementKeyVaultLink?.addEventListener("click", openKeyVaultSettings); elements.dataManagementCosmosEditorOpenDangerBtn?.addEventListener("click", showCosmosEditorDangerModal); elements.datamanagementcosmoseditordangeraccept?.addEventListener("change", updateCosmosEditorDangerAcceptState); @@ -231,9 +262,25 @@ function bindEvents() { elements.dataManagementCosmosEditorSaveBtn?.addEventListener("click", openCosmosEditorSaveModal); elements.datamanagementcosmoseditorconfirmationphrase?.addEventListener("input", updateCosmosEditorConfirmSaveState); elements.dataManagementCosmosEditorConfirmSaveBtn?.addEventListener("click", saveCosmosEditorDocument); + elements.datamanagementmigrationtemporarydestinationruenabled?.addEventListener("change", updateMigrationCapacityVisibility); + [ + elements.datamanagementmigrationmodenewonly, + elements.datamanagementmigrationmodedeltaupsert, + elements.datamanagementmigrationmodemirrorwithdeletions, + ].forEach((modeElement) => modeElement?.addEventListener("change", () => { + updateMigrationModeVisibility(); + loadMigrationSummary(); + })); + elements.datamanagementmigrationbaselinejobid?.addEventListener("change", loadMigrationSummary); + elements.datamanagementmigrationmirrorconfirmationphrase?.addEventListener("input", () => { + elements.datamanagementmigrationmirrorconfirmationphrase?.classList.remove("is-invalid"); + }); bindDataManagementChangeTracking(); setStorageAuthVisibility(); setMigrationTargetVisibility(); + updateMigrationCapacityVisibility(); + updateMigrationModeVisibility(); + updateMigrationSearchWriteFreezeVisibility(); updateConnectionStringStatus(); updateDataManagementSaveButtonState(); } @@ -253,7 +300,10 @@ function bindMigrationPickerEvents() { } }); }); - elements.datamanagementmigrationincludeaisearch?.addEventListener("change", loadMigrationSummary); + elements.datamanagementmigrationincludeaisearch?.addEventListener("change", () => { + updateMigrationSearchWriteFreezeVisibility(); + loadMigrationSummary(); + }); elements.datamanagementmigrationincludesourceblobs?.addEventListener("change", loadMigrationSummary); } @@ -419,12 +469,19 @@ function populateSettings(settings) { setValue(elements.datamanagementtargetcosmosendpoint, settings.target_cosmos_endpoint || ""); setValue(elements.datamanagementtargetcosmosdatabase, targetCosmosDatabaseName); setValue(elements.datamanagementtargetcosmoskey, settings.target_cosmos_key || ""); + setValue(elements.datamanagementtargetcosmossubscriptionid, settings.target_cosmos_subscription_id || ""); + setValue(elements.datamanagementtargetcosmosresourcegroup, settings.target_cosmos_resource_group || ""); setValue(elements.datamanagementtargetaisearchauth, settings.target_ai_search_authentication_type || "managed_identity"); setValue(elements.datamanagementtargetaisearchendpoint, settings.target_ai_search_endpoint || ""); setValue(elements.datamanagementtargetaisearchkey, settings.target_ai_search_key || ""); setValue(elements.datamanagementtargetecstorageauth, settings.target_enhanced_citations_storage_authentication_type || "managed_identity"); setValue(elements.datamanagementtargetecblobendpoint, settings.target_enhanced_citations_storage_blob_endpoint || ""); setValue(elements.datamanagementtargetecconnectionstring, settings.target_enhanced_citations_storage_connection_string || ""); + setValue(elements.datamanagementmigrationmaxparalleloperations, settings.migration_max_parallel_operations || 8); + setValue(elements.datamanagementmigrationretrycount, settings.migration_retry_count || 5); + setValue(elements.datamanagementmigrationskiprecentwithinhours, settings.migration_skip_recent_within_hours || 0); + setChecked(elements.datamanagementmigrationtemporarydestinationruenabled, Boolean(settings.migration_temporary_destination_ru_enabled)); + setValue(elements.datamanagementmigrationtemporarydestinationru, settings.migration_temporary_destination_ru || 10000); if (elements.dataManagementKeyStorage) { elements.dataManagementKeyStorage.textContent = formatKeyStorage(settings.encryption_key_storage); @@ -436,6 +493,7 @@ function populateSettings(settings) { updateKeyStorageExperience(settings); setStorageAuthVisibility(); setMigrationTargetVisibility(); + updateMigrationCapacityVisibility(); migrationTargetTypes.forEach(updateMigrationPickerVisibility); updateConnectionStringStatus(); resetDataManagementModified(); @@ -529,6 +587,14 @@ function setMigrationTargetVisibility() { } } +function updateMigrationCapacityVisibility() { + const enabled = Boolean(elements.datamanagementmigrationtemporarydestinationruenabled?.checked); + setElementVisible(elements.dataManagementMigrationTemporaryRuField, enabled); + if (elements.datamanagementmigrationtemporarydestinationru) { + elements.datamanagementmigrationtemporarydestinationru.disabled = !enabled; + } +} + function setKeyStorageAlert(variant, iconClass, title, message, linkText) { const alertElement = elements.dataManagementKeyStorageAlert; const iconElement = elements.dataManagementKeyStorageAlertIcon; @@ -581,12 +647,19 @@ function collectSettings() { target_cosmos_endpoint: getValue(elements.datamanagementtargetcosmosendpoint), target_cosmos_database_name: targetCosmosDatabaseName, target_cosmos_key: getValue(elements.datamanagementtargetcosmoskey), + target_cosmos_subscription_id: getValue(elements.datamanagementtargetcosmossubscriptionid), + target_cosmos_resource_group: getValue(elements.datamanagementtargetcosmosresourcegroup), target_ai_search_authentication_type: getValue(elements.datamanagementtargetaisearchauth) || "managed_identity", target_ai_search_endpoint: getValue(elements.datamanagementtargetaisearchendpoint), target_ai_search_key: getValue(elements.datamanagementtargetaisearchkey), target_enhanced_citations_storage_authentication_type: getValue(elements.datamanagementtargetecstorageauth) || backupStorageAuthManagedIdentity, target_enhanced_citations_storage_blob_endpoint: getValue(elements.datamanagementtargetecstorageauth) === backupStorageAuthManagedIdentity ? getValue(elements.datamanagementtargetecblobendpoint) : "", target_enhanced_citations_storage_connection_string: getValue(elements.datamanagementtargetecstorageauth) === backupStorageAuthConnectionString ? getValue(elements.datamanagementtargetecconnectionstring) : "", + migration_max_parallel_operations: getNumberValue(elements.datamanagementmigrationmaxparalleloperations, 8), + migration_retry_count: getNumberValue(elements.datamanagementmigrationretrycount, 5), + migration_skip_recent_within_hours: getNumberValue(elements.datamanagementmigrationskiprecentwithinhours, 0), + migration_temporary_destination_ru_enabled: Boolean(elements.datamanagementmigrationtemporarydestinationruenabled?.checked), + migration_temporary_destination_ru: getNumberValue(elements.datamanagementmigrationtemporarydestinationru, 10000), }; } @@ -1165,9 +1238,12 @@ async function testTargetCosmos() { try { const data = await requestJson("/api/admin/data-management/target/cosmos/test", { method: "POST", - body: JSON.stringify({ settings: collectSettings() }), + body: JSON.stringify({ settings: collectSettings(), migration_plan: buildMigrationPlan() }), }); - setStatus(`Target Cosmos connection succeeded. Database: ${data.database_name || targetCosmosDatabaseName}.`, "success"); + const verifiedCount = Number(data.migration_access?.container_count || 0); + const accessText = verifiedCount ? ` Verified ${formatNumber(verifiedCount)} planned container(s).` : ""; + const capacityText = data.capacity?.target_ru ? ` Temporary capacity can reach ${formatNumber(data.capacity.target_ru)} RU/s.` : ""; + setStatus(`Target Cosmos connection succeeded. Database: ${data.database_name || targetCosmosDatabaseName}.${accessText}${capacityText}`, "success"); showToast("Target Cosmos connection succeeded.", "success"); } catch (error) { setStatus(error.message || "Target Cosmos connection test failed.", "danger"); @@ -1177,6 +1253,25 @@ async function testTargetCosmos() { } } +async function testMigrationAccess() { + setBusy(elements.dataManagementTestMigrationAccessBtn, true, "Validating..."); + try { + const data = await requestJson("/api/admin/data-management/target/cosmos/test", { + method: "POST", + body: JSON.stringify({ settings: collectSettings(), migration_plan: buildMigrationPlan() }), + }); + const verifiedCount = Number(data.migration_access?.container_count || 0); + const capacityText = data.capacity?.target_ru ? ` Destination capacity management is ready up to ${formatNumber(data.capacity.target_ru)} RU/s.` : ""; + setStatus(`Cosmos migration access validation succeeded. ${formatNumber(verifiedCount)} planned Cosmos container(s) can be read and written.${capacityText}`, "success"); + showToast("Cosmos migration access validation succeeded.", "success"); + } catch (error) { + setStatus(error.message || "Cosmos migration access validation failed.", "danger"); + showToast(error.message || "Cosmos migration access validation failed.", "danger"); + } finally { + setBusy(elements.dataManagementTestMigrationAccessBtn, false); + } +} + async function testTargetSearch() { setBusy(elements.dataManagementTestTargetSearchBtn, true, "Testing..."); try { @@ -1247,9 +1342,31 @@ async function queueOperation(operation, backupType = null, options = {}, trigge function queueMigration(dryRun) { const button = elements.dataManagementExecuteMigrationBtn; + const migrationPlan = buildMigrationPlan(); + if ( + migrationPlan.include_ai_search && + !migrationPlan.target_ai_search_writes_frozen + ) { + const message = "Confirm that external destination AI Search writers are frozen before migrating Search documents."; + elements.datamanagementmigrationtargetsearchwritesfrozen?.focus(); + setStatus(message, "danger"); + showToast(message, "danger"); + return Promise.resolve(); + } + if ( + migrationPlan.migration_mode === "mirror_with_deletions" && + migrationPlan.mirror_confirmation !== migrationMirrorConfirmationPhrase + ) { + const message = `Type ${migrationMirrorConfirmationPhrase} before running a mirror migration.`; + elements.datamanagementmigrationmirrorconfirmationphrase?.classList.add("is-invalid"); + elements.datamanagementmigrationmirrorconfirmationphrase?.focus(); + setStatus(message, "danger"); + showToast(message, "danger"); + return Promise.resolve(); + } return queueOperation("migration", null, { dry_run: Boolean(dryRun), - migration_plan: buildMigrationPlan(), + migration_plan: migrationPlan, }, button); } @@ -1403,13 +1520,46 @@ function createSmallMutedElement(text) { } function buildMigrationPlan() { + const migrationMode = getMigrationTransferMode(); return { users: buildMigrationSelection("users"), groups: buildMigrationSelection("groups"), public_workspaces: buildMigrationSelection("public_workspaces"), include_ai_search: Boolean(elements.datamanagementmigrationincludeaisearch?.checked), + target_ai_search_writes_frozen: Boolean(elements.datamanagementmigrationtargetsearchwritesfrozen?.checked), include_source_blobs: Boolean(elements.datamanagementmigrationincludesourceblobs?.checked), + migration_mode: migrationMode, + baseline_job_id: migrationMode === "new_only" ? "" : getValue(elements.datamanagementmigrationbaselinejobid), + mirror_confirmation: migrationMode === "mirror_with_deletions" ? getValue(elements.datamanagementmigrationmirrorconfirmationphrase) : "", + }; +} + +function updateMigrationSearchWriteFreezeVisibility() { + const includesAiSearch = Boolean(elements.datamanagementmigrationincludeaisearch?.checked); + setElementVisible(elements.dataManagementMigrationSearchWriteFreeze, includesAiSearch); + if (!includesAiSearch && elements.datamanagementmigrationtargetsearchwritesfrozen) { + elements.datamanagementmigrationtargetsearchwritesfrozen.checked = false; + } +} + +function getMigrationTransferMode() { + return document.querySelector("input[name='data_management_migration_mode']:checked")?.value || "new_only"; +} + +function updateMigrationModeVisibility() { + const migrationMode = getMigrationTransferMode(); + const usesBaseline = migrationMode !== "new_only"; + const isMirror = migrationMode === "mirror_with_deletions"; + elements.dataManagementMigrationBaselineField?.classList.toggle("d-none", !usesBaseline); + elements.dataManagementMigrationMirrorConfirmation?.classList.toggle("d-none", !isMirror); + const descriptions = { + new_only: "Copies source items that are absent from the destination. Existing destination data is never updated or deleted.", + delta_upsert: "Copies new items and updates changed migration-owned items since the prior successful watermark. Destination-only data is retained.", + mirror_with_deletions: "Runs delta/upsert, then deletes destination-only items that carry successful SimpleChat migration ownership. Unowned data is retained.", }; + if (elements.dataManagementMigrationModeDescription) { + elements.dataManagementMigrationModeDescription.textContent = descriptions[migrationMode] || descriptions.new_only; + } } function buildMigrationSelection(targetType) { @@ -1430,9 +1580,12 @@ async function loadMigrationSummary(triggerButton = null, showSuccess = false) { try { const data = await requestJson("/api/admin/data-management/migration/summary", { method: "POST", - body: JSON.stringify({ migration_plan: buildMigrationPlan() }), + body: JSON.stringify({ + migration_plan: buildMigrationPlan(), + include_inventory: Boolean(showSuccess), + }), }); - renderMigrationSummary(data.summary || {}); + renderMigrationSummary(data.summary || {}, data.preview || null); if (showSuccess) { setStatus("Migration preview refreshed.", "success"); showToast("Migration preview refreshed.", "success"); @@ -1448,7 +1601,7 @@ async function loadMigrationSummary(triggerButton = null, showSuccess = false) { } } -function renderMigrationSummary(summary) { +function renderMigrationSummary(summary, preview = null) { const summaryElement = elements.dataManagementMigrationSummary; if (!summaryElement) { return; @@ -1465,11 +1618,42 @@ function renderMigrationSummary(summary) { options.className = "d-flex flex-wrap gap-2"; options.appendChild(createBadge(summary.include_ai_search ? "AI Search included" : "AI Search skipped", summary.include_ai_search ? "bg-info text-dark" : "bg-secondary")); options.appendChild(createBadge(summary.include_source_blobs ? "Source blobs included" : "Source blobs skipped", summary.include_source_blobs ? "bg-info text-dark" : "bg-secondary")); + options.appendChild(createBadge(`Mode: ${formatActivityLabel(summary.migration_mode || "new_only")}`, summary.migration_mode === "mirror_with_deletions" ? "bg-danger" : "bg-primary")); optionsColumn.appendChild(options); container.appendChild(optionsColumn); + if (preview?.estimated_outcomes) { + container.appendChild(createMigrationPreviewOutcomes(preview)); + } summaryElement.replaceChildren(container); } +function createMigrationPreviewOutcomes(preview) { + const column = document.createElement("div"); + column.className = "col-12 border-top pt-3"; + const heading = document.createElement("div"); + heading.className = "fw-semibold mb-2"; + heading.textContent = "Estimated destination changes"; + const outcomes = document.createElement("div"); + outcomes.className = "d-flex flex-wrap gap-2"; + [ + ["Create", "create_count", "bg-success"], + ["Update", "update_count", "bg-primary"], + ["Unchanged", "unchanged_count", "bg-secondary"], + ["Delete", "delete_count", "bg-danger"], + ["Not applicable", "not_applicable_count", "bg-light text-dark border"], + ["Missing", "missing_count", "bg-warning text-dark"], + ["Conflicts", "conflict_count", "bg-warning text-dark"], + ].forEach(([label, fieldName, badgeClass]) => { + outcomes.appendChild(createBadge(`${label}: ${formatNumber(preview.estimated_outcomes[fieldName] || 0)}`, badgeClass)); + }); + const metadata = document.createElement("div"); + metadata.className = "small text-muted mt-2"; + const baselineText = preview.baseline_job_id ? ` Baseline ${preview.baseline_job_id}.` : ""; + metadata.textContent = `Captured ${formatDate(preview.captured_at)}.${baselineText} Actual completion counts will report any divergence from this snapshot.`; + column.append(heading, outcomes, metadata); + return column; +} + function createMigrationSummaryCard(targetType, targetSummary) { const column = document.createElement("div"); column.className = "col-md-4"; @@ -1707,7 +1891,7 @@ function createJobRow(job) { row.appendChild(createStatusCell(job.status || "unknown")); row.appendChild(createCell(formatProgress(job.progress))); row.appendChild(createCell(job.last_message || job.last_error || "")); - row.appendChild(createJobActionCell(job.id)); + row.appendChild(createJobActionCell(job.id, job)); return row; } @@ -1721,7 +1905,7 @@ function createMessageRow(message, variant, columnSpan) { return row; } -function createJobActionCell(jobId) { +function createJobActionCell(jobId, job = {}) { const cell = document.createElement("td"); const button = document.createElement("button"); button.type = "button"; @@ -1733,9 +1917,98 @@ function createJobActionCell(jobId) { button.appendChild(document.createTextNode("View Log")); button.addEventListener("click", () => loadDataManagementJobDetail(jobId)); cell.appendChild(button); + if (job.operation === "migration" && job.can_retry === true) { + const retryButton = document.createElement("button"); + retryButton.type = "button"; + retryButton.className = "btn btn-outline-warning btn-sm ms-1"; + retryButton.disabled = !jobId; + const retryIcon = document.createElement("i"); + retryIcon.className = "bi bi-arrow-repeat me-1"; + retryButton.append(retryIcon, document.createTextNode(job.status === "running" ? "Resume" : "Retry")); + retryButton.addEventListener("click", () => retryMigrationJob(jobId, retryButton)); + cell.appendChild(retryButton); + } + if (job.operation === "migration" && job.can_cancel === true) { + const cancelButton = document.createElement("button"); + cancelButton.type = "button"; + cancelButton.className = "btn btn-outline-danger btn-sm ms-1"; + cancelButton.disabled = !jobId; + const cancelIcon = document.createElement("i"); + cancelIcon.className = "bi bi-stop-circle me-1"; + cancelButton.append(cancelIcon, document.createTextNode("Cancel")); + cancelButton.addEventListener("click", () => openMigrationCancellationModal(job)); + cell.appendChild(cancelButton); + } return cell; } +function openMigrationCancellationModal(job) { + if (!job?.id || !elements.dataManagementMigrationCancelModal || !window.bootstrap?.Modal) { + return; + } + pendingMigrationCancellationJobId = job.id; + setText( + elements.dataManagementMigrationCancelMessage, + `Request cancellation for migration ${job.id}?` + ); + window.bootstrap.Modal.getOrCreateInstance(elements.dataManagementMigrationCancelModal).show(); +} + +async function requestMigrationCancellation() { + const jobId = pendingMigrationCancellationJobId; + if (!jobId) { + return; + } + setBusy(elements.dataManagementConfirmMigrationCancelBtn, true, "Requesting..."); + try { + const data = await requestJson( + `/api/admin/data-management/jobs/${encodeURIComponent(jobId)}/cancel`, + { method: "POST", body: JSON.stringify({}) } + ); + const canceledImmediately = data.job?.status === "canceled"; + setStatus( + canceledImmediately + ? "Queued migration canceled before execution started." + : "Migration cancellation requested. The worker will stop at its next durable checkpoint.", + "success" + ); + showToast( + canceledImmediately ? "Migration canceled." : "Migration cancellation requested.", + "success" + ); + window.bootstrap.Modal.getOrCreateInstance(elements.dataManagementMigrationCancelModal).hide(); + renderJobs([data.job]); + loadDataManagementJobs(); + if (currentJobDetailId === jobId) { + loadDataManagementJobDetail(jobId, { showModal: false, liveRefresh: true }); + } + } catch (error) { + setStatus(error.message || "Migration cancellation could not be requested.", "danger"); + showToast(error.message || "Migration cancellation could not be requested.", "danger"); + } finally { + setBusy(elements.dataManagementConfirmMigrationCancelBtn, false); + } +} + +async function retryMigrationJob(jobId, button) { + setBusy(button, true, "Retrying..."); + try { + const data = await requestJson(`/api/admin/data-management/jobs/${encodeURIComponent(jobId)}/retry`, { + method: "POST", + }); + setStatus("Migration retry queued from durable checkpoints.", "success"); + showToast("Migration retry queued.", "success"); + renderJobs([data.job]); + loadDataManagementJobs(); + loadDataManagementJobDetail(jobId, { showModal: false, liveRefresh: true }); + } catch (error) { + setStatus(error.message || "Migration retry could not be queued.", "danger"); + showToast(error.message || "Migration retry could not be queued.", "danger"); + } finally { + setBusy(button, false); + } +} + function createCell(text) { const cell = document.createElement("td"); cell.textContent = text ?? ""; @@ -1777,6 +2050,7 @@ function renderJobDetailModal(job, items) { setText(elements.dataManagementJobDetailTitle, formatOperation(job.operation || "job", job.backup_type || "")); setText(elements.dataManagementJobDetailSubtitle, job.id ? `Job ID: ${job.id}` : "Job ID not available"); renderJobDetailSummary(job); + renderJobDetailActions(job); renderJobDetailProgress(job); renderJobItems(items); const artifacts = getJobArtifacts(job, items); @@ -1785,12 +2059,82 @@ function renderJobDetailModal(job, items) { renderJobWarnings(job, artifacts); } +function renderJobDetailActions(job) { + const container = elements.dataManagementJobDetailActions; + if (!container) { + return; + } + container.replaceChildren(); + if (job.operation !== "migration") { + container.classList.add("d-none"); + return; + } + const fullManifestLink = createMigrationManifestDownloadLink( + job.id, + "Download manifest", + "", + "btn-outline-primary" + ); + const failuresLink = createMigrationManifestDownloadLink( + job.id, + "Download failures", + "failed,missing,collision", + "btn-outline-danger" + ); + container.append(fullManifestLink, failuresLink); + if (job.can_retry === true) { + const retryButton = document.createElement("button"); + retryButton.type = "button"; + retryButton.className = "btn btn-outline-warning btn-sm"; + const retryIcon = document.createElement("i"); + retryIcon.className = "bi bi-arrow-repeat me-1"; + retryButton.append( + retryIcon, + document.createTextNode(job.status === "running" ? "Resume" : "Retry") + ); + retryButton.addEventListener("click", () => retryMigrationJob(job.id, retryButton)); + container.appendChild(retryButton); + } + if (job.can_cancel === true) { + const cancelButton = document.createElement("button"); + cancelButton.type = "button"; + cancelButton.className = "btn btn-outline-danger btn-sm"; + const cancelIcon = document.createElement("i"); + cancelIcon.className = "bi bi-stop-circle me-1"; + cancelButton.append(cancelIcon, document.createTextNode("Cancel")); + cancelButton.addEventListener("click", () => openMigrationCancellationModal(job)); + container.appendChild(cancelButton); + } + container.classList.toggle("d-none", container.childElementCount === 0); +} + +function createMigrationManifestDownloadLink(jobId, label, statuses, buttonClass) { + const link = document.createElement("a"); + link.className = `btn ${buttonClass} btn-sm`; + const encodedJobId = encodeURIComponent(jobId || ""); + const query = statuses ? `?statuses=${encodeURIComponent(statuses)}` : ""; + link.href = `/api/admin/data-management/jobs/${encodedJobId}/migration-manifest${query}`; + link.download = ""; + const icon = document.createElement("i"); + icon.className = "bi bi-download me-1"; + link.append(icon, document.createTextNode(label)); + return link; +} + function renderJobDetailSummary(job) { const container = elements.dataManagementJobDetailSummary; if (!container) { return; } const result = job.result && typeof job.result === "object" ? job.result : {}; + const migrationState = job.migration_state && typeof job.migration_state === "object" ? job.migration_state : {}; + const migrationTotals = migrationState.totals && typeof migrationState.totals === "object" ? migrationState.totals : {}; + const migrationTiles = job.operation === "migration" ? [ + createSummaryTile("Migration ID", migrationState.migration_id || job.id || "N/A"), + createSummaryTile("Processed", formatNumber(migrationTotals.processed_count || 0)), + createSummaryTile("Transferred", formatBytes(migrationTotals.bytes || 0)), + createSummaryTile("Request units", formatNumber(migrationTotals.request_units || 0)), + ] : []; container.replaceChildren( createSummaryTile("Status", createStatusBadge(job.status || "unknown")), createSummaryTile("Progress", formatProgress(job.progress)), @@ -1800,7 +2144,8 @@ function renderJobDetailSummary(job) { createSummaryTile("Created", formatDate(job.created_at)), createSummaryTile("Started", formatDate(job.started_at)), createSummaryTile("Completed", formatDate(job.completed_at)), - createSummaryTile("Artifacts", formatNumber(result.artifact_count || getArtifactTotals(result.artifacts).artifactCount || 0)) + createSummaryTile("Artifacts", formatNumber(result.artifact_count || getArtifactTotals(result.artifacts).artifactCount || 0)), + ...migrationTiles ); } @@ -1830,6 +2175,11 @@ function renderJobDetailProgress(job) { const percent = getProgressPercent(progress); const completedSteps = Number(progress.completed_steps || 0); const totalSteps = Number(progress.total_steps || 0); + const isMigration = job.operation === "migration"; + const isMigrationActive = isMigration && activeJobStatuses.has(job.status); + const displayedStage = isMigrationActive + ? Math.min(totalSteps, completedSteps + 1) + : completedSteps; const card = document.createElement("div"); card.className = "card border-info-subtle"; const body = document.createElement("div"); @@ -1842,29 +2192,124 @@ function renderJobDetailProgress(job) { title.textContent = progress.current_step ? formatActivityLabel(progress.current_step) : "Job progress"; const stepText = document.createElement("small"); stepText.className = "text-muted"; - stepText.textContent = totalSteps > 0 ? `${formatNumber(completedSteps)} of ${formatNumber(totalSteps)} steps complete` : "Waiting for step details"; + stepText.textContent = isMigration + ? totalSteps > 0 ? `Stage ${formatNumber(displayedStage)} of ${formatNumber(totalSteps)}` : "Waiting for stage details" + : totalSteps > 0 ? `${formatNumber(completedSteps)} of ${formatNumber(totalSteps)} steps complete` : "Waiting for step details"; header.append(title, stepText); const progressWrapper = document.createElement("div"); progressWrapper.className = "progress"; progressWrapper.setAttribute("role", "progressbar"); - progressWrapper.setAttribute("aria-valuenow", String(percent)); progressWrapper.setAttribute("aria-valuemin", "0"); progressWrapper.setAttribute("aria-valuemax", "100"); const progressBar = document.createElement("div"); - progressBar.className = `progress-bar ${job.status === "failed" ? "bg-danger" : "bg-primary"}`; - progressBar.style.width = `${percent}%`; - progressBar.textContent = `${percent}%`; + if (isMigrationActive) { + progressWrapper.setAttribute("aria-valuetext", "Migration stage is active; measured throughput is shown below."); + progressBar.className = "progress-bar progress-bar-striped progress-bar-animated bg-primary"; + progressBar.style.width = "100%"; + progressBar.textContent = "Active stage"; + } else { + progressWrapper.setAttribute("aria-valuenow", String(percent)); + progressBar.className = `progress-bar ${job.status === "failed" ? "bg-danger" : "bg-primary"}`; + progressBar.style.width = `${percent}%`; + progressBar.textContent = `${percent}%`; + } progressWrapper.appendChild(progressBar); const message = document.createElement("p"); message.className = "text-muted small mt-2 mb-0"; message.textContent = job.last_message || job.last_error || "No job message has been recorded yet."; body.append(header, progressWrapper, message); + if (isMigration) { + const migrationMetrics = getMigrationLiveMetrics(job); + if (migrationMetrics.length) { + const metricGrid = document.createElement("div"); + metricGrid.className = "row g-2 small mt-3"; + migrationMetrics.forEach((metric) => { + metricGrid.appendChild(createDetailMetric(metric.label, metric.value)); + }); + body.appendChild(metricGrid); + } + } card.appendChild(body); container.replaceChildren(card); } +function getMigrationLiveMetrics(job) { + const migrationState = job?.migration_state; + if (!migrationState || typeof migrationState !== "object") { + return []; + } + const totals = migrationState.totals && typeof migrationState.totals === "object" ? migrationState.totals : {}; + const resources = migrationState.resources && typeof migrationState.resources === "object" ? Object.values(migrationState.resources) : []; + const activeResource = resources.find((resource) => resource?.status === "in_progress") || {}; + const activeMetrics = activeResource.progress && typeof activeResource.progress === "object" ? activeResource.progress : activeResource.result || {}; + const capacity = migrationState.capacity && typeof migrationState.capacity === "object" ? migrationState.capacity : {}; + const metrics = [ + { label: "Processed", value: formatNumber(totals.processed_count || 0) }, + { label: "Transferred", value: formatBytes(totals.bytes || 0) }, + { label: "Copied / skipped", value: `${formatNumber(totals.copied_count || 0)} / ${formatNumber(totals.skipped_count || 0)}` }, + { label: "Failures", value: formatNumber(totals.failed_count || 0) }, + { label: "Collisions", value: formatNumber(totals.collision_count || 0) }, + ]; + if (activeMetrics.observed_bytes !== undefined) { + metrics.push({ label: "Observed transferred", value: formatBytes(activeMetrics.observed_bytes || 0) }); + } + if (activeMetrics.items_per_second !== undefined) { + metrics.push({ label: "Items per second", value: formatNumber(activeMetrics.items_per_second) }); + } + const observedRate = activeMetrics.observed_bytes_per_second ?? activeMetrics.bytes_per_second; + if (observedRate !== undefined) { + metrics.push({ label: "Transfer rate", value: `${formatBytes(observedRate || 0)}/s` }); + } + if (activeMetrics.request_units_per_second !== undefined) { + metrics.push({ label: "RU per second", value: formatNumber(activeMetrics.request_units_per_second) }); + } + if (capacity.status) { + metrics.push({ label: "Destination capacity", value: formatActivityLabel(capacity.status) }); + } + const heartbeatAge = formatTimestampAge(job?.last_heartbeat_at); + const progressAge = formatTimestampAge(job?.last_progress_at || migrationState.last_progress_at); + if (heartbeatAge) { + metrics.push({ label: "Last heartbeat", value: heartbeatAge }); + } + if (progressAge) { + metrics.push({ label: "Last progress", value: progressAge }); + } + if (heartbeatAge) { + const hasRecentProgress = timestampAgeSeconds( + job?.last_progress_at || migrationState.last_progress_at + ) <= 10; + metrics.push({ + label: "Liveness", + value: hasRecentProgress ? "Running - progress active" : "Running - alive, no recent progress", + }); + } + return metrics; +} + +function timestampAgeSeconds(value) { + const timestamp = Date.parse(value || ""); + if (!Number.isFinite(timestamp)) { + return Number.POSITIVE_INFINITY; + } + return Math.max(0, Math.floor((Date.now() - timestamp) / 1000)); +} + +function formatTimestampAge(value) { + const ageSeconds = timestampAgeSeconds(value); + if (!Number.isFinite(ageSeconds)) { + return ""; + } + if (ageSeconds < 60) { + return `${ageSeconds}s ago`; + } + if (ageSeconds < 3600) { + return `${Math.floor(ageSeconds / 60)}m ago`; + } + return `${Math.floor(ageSeconds / 3600)}h ago`; +} + function renderJobItems(items) { const container = elements.dataManagementJobItemsTbody; if (!container) { @@ -1956,15 +2401,41 @@ function createArtifactCard(artifact) { const metrics = document.createElement("div"); metrics.className = "row g-2 small mb-2"; + const copiedCount = artifact.copied_count ?? artifact.item_count ?? 0; metrics.append( - createDetailMetric("Records", formatNumber(artifact.item_count || 0)), + createDetailMetric("Copied", formatNumber(copiedCount)), + createDetailMetric("Created", formatNumber(artifact.created_count || 0)), + createDetailMetric("Updated", formatNumber(artifact.updated_count || 0)), + createDetailMetric("Unchanged", formatNumber(artifact.unchanged_count || 0)), + createDetailMetric("Skipped", formatNumber(artifact.skipped_count || 0)), + createDetailMetric("Failed", formatNumber(artifact.failed_count || 0)), + createDetailMetric("Collisions", formatNumber(artifact.collision_count || 0)), createDetailMetric("Blobs", formatNumber(artifact.blob_count || 0)), createDetailMetric("Size", formatBytes(artifact.bytes || 0)) ); + if (artifact.type === "migration_reconciliation") { + metrics.append( + createDetailMetric("Readiness", formatActivityLabel(artifact.readiness || "unknown")), + createDetailMetric("Deleted", formatNumber(artifact.deleted_count || 0)), + createDetailMetric("Delete candidates", formatNumber(artifact.delete_candidate_count || 0)) + ); + } + if (artifact.items_per_second !== undefined || artifact.bytes_per_second !== undefined) { + metrics.append( + createDetailMetric("Items / second", formatNumber(artifact.items_per_second || 0)), + createDetailMetric("Transfer rate", `${formatBytes(artifact.bytes_per_second || 0)}/s`) + ); + } const location = createDetailBlock("Location", artifact.path || artifact.prefix || "Not recorded"); const notes = createArtifactNoteGroup(artifact); body.append(header, metrics, location, notes); + if (artifact.preview_actual_divergence && typeof artifact.preview_actual_divergence === "object") { + const divergenceTitle = document.createElement("div"); + divergenceTitle.className = "text-muted small mt-2"; + divergenceTitle.textContent = "Preview versus actual"; + body.append(divergenceTitle, createDetailChipGroup(artifact.preview_actual_divergence)); + } card.appendChild(body); return card; } @@ -2094,6 +2565,24 @@ function createArtifactNoteGroup(artifact) { { label: "Index", value: artifact.index_name }, { label: "Partial since", value: artifact.partial_since_epoch }, { label: "Filter", value: artifact.partial_filter }, + { label: "Request units", value: artifact.request_units }, + { label: "RU per second", value: artifact.request_units_per_second }, + { label: "Workers", value: artifact.parallel_operations }, + { label: "Retries", value: artifact.retry_count }, + { label: "Retry delays", value: artifact.retry_attempt_count }, + { label: "Prior failures", value: artifact.prior_failed_count }, + { label: "Checkpoints", value: artifact.checkpoint_count }, + { label: "Readiness", value: artifact.readiness, className: artifact.readiness === "ready" ? "bg-success" : "bg-warning text-dark" }, + { label: "Deletion", value: artifact.deletion_status }, + { label: "Remaining owned extras", value: artifact.remaining_destination_only_owned_count }, + { label: "Unowned extras", value: artifact.destination_only_unowned_count }, + { label: "Unresolved scope", value: artifact.unresolved_scope_count }, + { label: "Stale", value: artifact.stale_count }, + { label: "Deletion blockers", value: Array.isArray(artifact.deletion_blockers) ? artifact.deletion_blockers.join("; ") : artifact.deletion_blockers, className: "bg-warning text-dark" }, + { label: "Prior migration skips", value: artifact.destination_provenance_skip_count }, + { label: "Missing blobs", value: artifact.missing_count }, + { label: "No source blob", value: artifact.not_applicable_count }, + { label: "Collisions", value: artifact.collision_count }, { label: "Warning", value: artifact.warning, className: "bg-warning text-dark" }, ].filter((note) => isPresent(note.value)); @@ -2184,7 +2673,7 @@ function getArtifactTotals(artifacts) { } totals.artifactCount += 1; totals.bytes += Number(artifact.bytes || 0); - totals.recordCount += Number(artifact.item_count || 0); + totals.recordCount += Number(artifact.item_count ?? artifact.copied_count ?? 0); totals.blobCount += Number(artifact.blob_count || 0); }); return totals; @@ -2261,7 +2750,23 @@ async function refreshOpenJobDetail() { } jobDetailRefreshInFlight = true; try { - await loadDataManagementJobDetail(currentJobDetailId, { showModal: false, liveRefresh: true }); + const data = await requestJson( + `/api/admin/data-management/jobs/${encodeURIComponent(currentJobDetailId)}/progress`, + { method: "GET" } + ); + const job = data.job || {}; + renderJobDetailSummary(job); + renderJobDetailActions(job); + renderJobDetailProgress(job); + updateJobDetailAutoRefresh(job); + if (!activeJobStatuses.has(String(job.status || "unknown"))) { + await loadDataManagementJobDetail( + currentJobDetailId, + { showModal: false, liveRefresh: true } + ); + } + } catch (error) { + stopJobDetailAutoRefresh({ message: "Live refresh paused after a request failure." }); } finally { jobDetailRefreshInFlight = false; } diff --git a/application/single_app/templates/admin_settings.html b/application/single_app/templates/admin_settings.html index fc7af65e8..b53d44cd7 100644 --- a/application/single_app/templates/admin_settings.html +++ b/application/single_app/templates/admin_settings.html @@ -5071,6 +5071,16 @@
Target Cosmos Database
+
+
+ + +
+
+ + +
+
@@ -5141,6 +5151,32 @@
Migration Workflow
Refresh Summary
+
+ Synchronization mode +
+ + + + + + +
+

+ Copies source items that are absent from the destination. Existing destination data is never updated or deleted. +

+
+
+ + +
The baseline fixes the prior service watermarks used for this catch-up run.
+
+
+ +
@@ -5215,6 +5251,11 @@
Public Workspaces
+
+ + +
Azure AI Search does not support conditional document writes. SimpleChat target indexing pauses automatically; freeze all external writers before copying Search documents.
+
@@ -5223,6 +5264,44 @@
Public Workspaces
+
+
+
+
Performance And Resume
+

Migration uses durable per-resource checkpoints and the same migration ID after a retry.

+
+ +
+
+
+ + +
+
+ + +
+
+ + +
+
+
+ + +
+
+
+ + +
+
+
The migration records the prior database or dedicated-container capacity, raises only eligible targets up to 10,000 RU/s, and restores it after completion or failure. Capacity changes require destination ARM permission.
+
+
+

Choose migration targets to see the summary.

@@ -5521,6 +5600,7 @@
+
@@ -5568,6 +5648,25 @@
Warnings
+ +
diff --git a/docs/explanation/features/DATA_MANAGEMENT_BACKUP_MIGRATION.md b/docs/explanation/features/DATA_MANAGEMENT_BACKUP_MIGRATION.md index fe321b31d..5f6f4f8c4 100644 --- a/docs/explanation/features/DATA_MANAGEMENT_BACKUP_MIGRATION.md +++ b/docs/explanation/features/DATA_MANAGEMENT_BACKUP_MIGRATION.md @@ -1,7 +1,7 @@ # Data Management Backup and Migration Implemented in version: **0.241.211** -Updated in version: **0.241.221** +Updated in version: **0.250.071** ## Overview @@ -48,10 +48,14 @@ The Migration card supports a guided migration workflow: - Select whether to migrate no groups, all groups, or selected groups, with optional group document migration. - Select whether to migrate no public workspaces, all public workspaces, or selected public workspaces, with optional public workspace document migration. - Preview the migration plan to refresh counts and selected IDs before execution. +- Choose New only, Delta / upsert, or explicitly confirmed Mirror with deletions. Delta and mirror pin a compatible completed migration as their prior watermark. +- Run a read-only live inventory preview for estimated create, update, unchanged, delete, missing, not-applicable, and conflict counts. Execution also captures its own durable inventory after the worker acquires the global migration lock. - Execute Migration queues a durable Data Management migration job. The job history modal shows live progress, per-step timeline events, migrated artifact counts, and warnings. Migration execution currently copies selected SimpleChat Cosmos records, matching AI Search documents for selected document scopes, and source document blobs when Enhanced Citations source and destination storage are configured. +For durable provenance, destination access probes, collision protection, checkpoint retries, transfer telemetry, and temporary destination Cosmos capacity controls, see [Data Management Migration Resilience](DATA_MANAGEMENT_MIGRATION_RESILIENCE.md). + ### Security - All Data Management routes are admin-only. @@ -92,10 +96,10 @@ Data Management settings save through their own API and are excluded from the re 6. Queue a full or partial backup, or use the restore/migration dry-run buttons to create durable orchestration records. 7. Configure and test Target Cosmos, Target Search, and Target Enhanced Citation Storage before running an actual migration. 8. Use the Migration Workflow to choose users, groups, and public workspaces, then decide whether documents, AI Search entries, and source blobs should be included. -9. Preview the migration plan, then Execute Migration to queue the job. +9. Choose the synchronization mode, preview live destination changes, and then Execute Migration to queue the job. Mirror mode requires the exact destructive confirmation phrase. 10. Open Advanced backup scope only when you need to alter the default Cosmos DB, AI Search index, or source blob backup surfaces. 11. Use Backup Inventory to see available backups first, filter to full or partial backups, and open View Log for structured backup details. -12. Use Job History to inspect live progress, completed steps, warnings, storage/manifest details, and artifact contents for any Data Management job. +12. Use Job History to inspect live progress, completed steps, reconciliation readiness, preview divergence, and artifact contents. Migration detail also provides Retry/Cancel plus full or failure-only JSONL manifest downloads. Migration is used when moving SimpleChat data into another SimpleChat environment, rehearsing a cutover, or preparing a controlled environment transfer. The target Cosmos account, authentication type, and optional account key are configurable. The target database name is fixed to `SimpleChat` so future migration apply jobs use the standard SimpleChat container layout. @@ -110,9 +114,9 @@ For managed identity target Cosmos migration, assign this App Service identity C ## Limitations - Backup artifact export is implemented for Cosmos DB and AI Search, with optional source blob copying. -- Restore and migration apply logic currently create durable admin job records and warnings; the target import/apply engine is the next implementation layer. +- Restore execution is documented separately. Migration apply supports selected SimpleChat Cosmos, Search, and Enhanced Citations Blob surfaces rather than arbitrary Azure resources. ## Version References -- Application version updated in `application/single_app/config.py` to `0.241.221`. +- Application version updated in `application/single_app/config.py` to `0.250.071`. - Functional and UI tests include the same implementation version. \ No newline at end of file diff --git a/docs/explanation/features/DATA_MANAGEMENT_MIGRATION_RESILIENCE.md b/docs/explanation/features/DATA_MANAGEMENT_MIGRATION_RESILIENCE.md new file mode 100644 index 000000000..68757740c --- /dev/null +++ b/docs/explanation/features/DATA_MANAGEMENT_MIGRATION_RESILIENCE.md @@ -0,0 +1,151 @@ +# Data Management Migration Resilience + +Implemented in version: **0.250.071** + +GitHub issue: [#1043](https://github.com/microsoft/simplechat/issues/1043) + +## Overview + +Data Management migrations use durable, job-scoped provenance and checkpoints across selected Cosmos DB records, Azure AI Search documents, and Enhanced Citations source blobs. A migration retry retains the original migration GUID and resumes from verified service checkpoints. Operators can run create-only, delta/upsert, or explicitly confirmed mirror migrations, preview estimated destination changes, and review final source/destination reconciliation plus preview divergence. + +## Dependencies + +- Data Management administrator access. +- Source and destination data-plane permissions for the selected Cosmos containers, Azure AI Search indexes, and source blob containers. +- Destination Cosmos DB Data Contributor for target record writes. +- For temporary Cosmos capacity management, destination ARM access that can read and update database or dedicated-container throughput, plus the target subscription ID and resource group. + +## Technical Specifications + +### Durable Provenance + +Each migration job uses its job ID as a GUID-stable migration identity and records a single UTC start time. Successful copies add destination-specific metadata: + +- Cosmos documents: `simplechatMigration` with migration identity, status, and source `_ts` version. +- Azure AI Search documents: queryable migration fields plus a canonical SHA-256 source-document hash. Existing target indexes are updated with compatible provenance fields when necessary. +- Source blobs: migration identity, ETag/last-modified version fingerprint, available content MD5, and a non-reversible scope hash in Blob metadata. + +Only a `succeeded` marker is eligible for a skip. The current migration GUID always skips its completed work on resume. Cross-run age-based skipping is explicit and defaults to `0` hours. + +The migration does not overwrite an unowned destination collision. A Cosmos record, AI Search key, or Blob path that already exists without eligible provenance causes the affected resource to stop and record a collision count. This avoids replacing destination content that may have been created independently. Target Cosmos containers must also have the exact expected partition-key path before probing or writing. + +### Checkpoints And Throughput + +The migration state is stored on the existing Data Management job record. It contains a secret-free configuration fingerprint, resource status, attempts, copied/skipped/failed counts, bytes, Cosmos request units, rate calculations, preflight results, and capacity lifecycle state. Privacy-safe per-item outcomes are stored separately in bounded batches in the existing job-items container and use SHA-256 logical identities instead of raw document IDs or Blob paths. + +- Cosmos records use bounded parallel creates or ETag-conditional replacements with transient retry handling and response-charge capture. Updates are allowed only for successfully migration-owned records. +- AI Search uses deterministic `id asc` keyset pages capped at 1,000 records, a persisted per-index cursor, bounded selected-scope filters, concurrent upload batches, and per-document indexing outcomes. Migrations beyond 100,000 documents do not use deep `$skip` pagination. +- Before an AI Search transfer, the target SimpleChat `data_management_jobs` container closes and drains a durable Search write gate. Normal target SimpleChat document indexing takes a short bounded gate slot, so it cannot race migration uploads or final mirror deletes. The gate is renewed with migration heartbeats, released after a clean transfer or drained cancellation, and retained for the full uncertainty window after a lost response. +- Azure AI Search does not provide document-level ETag or create-only indexing actions. The migration therefore requires a checked external destination-writer freeze acknowledgement, rejects a destination that resolves to the source Search service, and re-reads unresolved keys after a timeout or partial response before retrying. External target writers remain an operational cutover precondition. +- Blob migration streams chunked source content instead of loading each blob into App Service memory. It renews the job lease while long futures are still active, reports in-flight bytes, observes cancellation between chunks, preserves content settings, metadata, tags, access tier, and blob type where supported, and uses source/target ETag conditions. +- Blob uploads use `pending` migration metadata during transfer. The worker verifies destination size/MD5 and source ETag stability before conditionally stamping `succeeded`; failed creates are removed while still owned by the pending attempt. +- Destination marker and Blob deduplication lookups are bounded to the active item or batch rather than retaining a full migration key set in App Service memory. +- ETag-based job replacement and worker-lease checks prevent a stale worker from silently overwriting a newer checkpoint. Each retry gets a fresh attempt timestamp so displayed rates exclude downtime. +- A target-side conditional coordinator lease in the destination `data_management_jobs` container prevents independent SimpleChat source deployments from overlapping before inventory, Cosmos, Search, or Blob writes. The source job persists only the target lock token and keeps the target SDK client process-local; clean completion and drained cancellation release the lease, while uncertain failures retain its expiry quarantine. + +Failed, canceled, or stale migration jobs can be retried from the existing job detail. Completed resource checkpoints and the migration GUID remain unchanged. + +### Incremental Modes And Watermarks + +- **New only** is the default. It creates missing destination identities and never updates or deletes existing data. +- **Delta / upsert** requires a compatible completed migration baseline. The backend uses an explicitly selected baseline job or pins the newest compatible completed job when the field is blank. It creates new items, updates changed migration-owned items, and retains destination-only data. +- **Mirror with deletions** runs delta/upsert and then removes destination-only items only when they carry successful migration ownership. The exact phrase `MIRROR WITH DELETIONS` is required. Unowned destination data is retained and reported. + +Each run captures a UTC cutoff for Cosmos and intentionally includes the baseline second again to avoid timestamp-resolution gaps. Cosmos combines `_ts` with a canonical content hash, so different revisions in the same second are still detected. AI Search has no shared modified-time field, so it uses the source version observed during resumable keyset enumeration plus a canonical source hash. Blob Storage uses the ETag observed immediately before transfer and verifies it again after upload. Search or Blob changes observed after inventory can therefore cause preview divergence or a retryable `not_ready` reconciliation; they are never represented as having a global timestamp cutoff. + +### Preview, Reconciliation, And Divergence + +The explicit Preview action performs a read-only source/destination inventory and estimates creates, updates, unchanged items, deletes, not-applicable blobs, missing sources, and conflicts. Execute Migration queues immediately. After the worker owns the global migration lock, it captures a fresh server-owned inventory as a durable, heartbeat-backed job stage and pins the resolved baseline before transfer begins. + +After copy stages complete, memory-bounded reconciliation compares selected source and destination identities across Cosmos, Search, and Blob Storage. Cosmos and Blob use bounded point reads; Search merges ordered source/destination keyset streams without retaining full key maps. Reconciliation records matched, missing, stale, destination-only owned, destination-only unowned, unresolved-scope, conflict, and deleted counts. A `not_ready` report fails retryably and cannot become an incremental baseline. + +Mirror deletion is two phase. The worker first persists a private bounded candidate plan and a read-only report. It blocks all deletion when transfer outcomes diverge from the queued preview, copied items are stale/missing/failed/conflicting, unowned extras remain, or any candidate changes. Only then does it stream the private plan again, revalidate current source absence and destination ownership/ETag immediately before each write, and apply guarded deletes. + +Job detail exposes readiness, actual outcomes, preview divergence, keyset cursor/checkpoint counters, and Retry/Cancel controls. Admins can download a sanitized JSONL manifest or a filtered failure list (`failed`, `missing`, and `collision`) for remediation; retry resumes the incomplete resource while verified successful items remain idempotent. + +### Cancellation And Automatic Recovery + +- An administrator can request cancellation from the migration job list. Queued migrations become terminally canceled immediately. Running migrations retain their lease until the worker reaches the next guarded checkpoint, then stop without starting additional migration work. +- Authorization-reducing personal document unshare operations fail retryably while the target Search write gate is frozen. The existing share grant remains authoritative until its Search ACL projection succeeds, so a revoked user is never reported as removed while stale Search chunks could still authorize retrieval. +- Temporary destination Cosmos capacity restoration is allowed during a cancellation request so a canceled migration does not leave the destination boosted. +- Canceled jobs retain their migration ID and completed resource checkpoints. The job detail exposes **Retry** or **Resume** once the canceled worker reaches its terminal state. +- The Data Management scheduler scans delayed queued and stale running migration jobs independently of scheduled backup enablement. It records recovery activity and submits the work asynchronously through the Flask executor or standalone scheduler worker thread. + +### Destination Preflight + +**Validate Cosmos Migration Access** checks the current selected migration plan for source Cosmos reads plus destination Cosmos create, read, and delete access. The full preflight runs when the migration worker starts: + +- Source Cosmos read access for each selected resource. +- Destination Cosmos create, read, and delete access using a short-lived probe item in every planned container. +- Source AI Search read access, destination index/provenance-field readiness, and a temporary Search upload/read/delete probe. +- Source and target blob container access when source-blob migration is selected, plus a temporary Blob upload/read/delete probe. + +Preflight results are saved in the migration state and shown through existing job details and activity history without credentials or source content. + +### Temporary Cosmos Capacity + +The optional **Temporarily increase destination Cosmos capacity during this migration** control raises eligible destination database or dedicated-container throughput to the configured target, capped at **10,000 RU/s**. + +- The original mode and RU/s value are persisted before every ARM capacity change. +- Only targets below the requested value are changed. +- Capacity is restored after successful or failed migration execution. +- Restoration is skipped if the target changed independently after the migration boost, preventing the migration from overwriting an external administrator's update. +- If capacity restoration fails or its outcome is uncertain after an interruption, the migration fails retryably with durable `restore_pending` state. Retry restores the recorded snapshot before any other work and does not reapply the boost. + +## Usage + +1. Open **Admin Settings > Migration** and configure the destination services. +2. Select the scopes and document surfaces to migrate. +3. Choose **New only**, **Delta / upsert**, or **Mirror with deletions**. Delta and mirror can name a completed baseline job or automatically use the latest compatible baseline. +4. Select **Preview Migration** and review estimated creates, updates, unchanged items, deletes, missing sources, and conflicts. For mirror, type the exact destructive confirmation phrase. +5. Select **Validate Cosmos Migration Access** to confirm the selected source and destination permissions. The worker performs complete Cosmos, Search, and Blob preflight before copying data. +6. Optionally enable the temporary destination Cosmos capacity boost, supply the target subscription and resource group, and leave the target at or below 10,000 RU/s. +7. Execute the migration and open the job log to observe durable inventory, stage state, transfer rates, in-flight Blob bytes, keyset checkpoints, outcome counts, reconciliation readiness, and preview divergence. Download the manifest or failure list when item-level remediation is required. +8. Use **Cancel** to request a cooperative stop. Use **Retry** or **Resume** on a failed, canceled, or stale job. The original migration GUID and completed destination markers are retained. + +### Recommended Cutover + +1. Run **New only** for the initial transfer. +2. Keep the source active and run one or more **Delta / upsert** catch-up migrations. +3. Resolve preview conflicts and any reconciliation state that is `not_ready`. +4. Freeze source writes and external destination AI Search writers for the final cutover window. Confirm the destination Search writer freeze in the migration workflow. +5. Run a final **Delta / upsert**, or a confirmed **Mirror with deletions** only when the destination must exactly follow source deletions. +6. Switch traffic only after final reconciliation is `ready` or its warnings are explicitly accepted. + +## Testing And Validation + +- `functional_tests/test_data_management_migration_provenance.py` +- `functional_tests/test_data_management_migration_state.py` +- `functional_tests/test_data_management_cosmos_migration_resilience.py` +- `functional_tests/test_data_management_ai_search_migration_resilience.py` +- `functional_tests/test_data_management_blob_migration_resilience.py` +- `functional_tests/test_data_management_destination_cosmos_capacity.py` +- `functional_tests/test_data_management_migration_orchestration.py` +- `functional_tests/test_data_management_migration_retry.py` +- `functional_tests/test_data_management_migration_coordinator.py` +- `functional_tests/test_data_management_migration_preflight.py` +- `functional_tests/test_data_management_migration_cancellation.py` +- `functional_tests/test_data_management_migration_recovery.py` +- `functional_tests/test_data_management_incremental_migration_modes.py` +- `functional_tests/test_data_management_migration_reconciliation.py` +- `functional_tests/test_data_management_migration_manifest.py` +- `functional_tests/test_data_management_security_patterns.py` +- `ui_tests/test_admin_data_management_settings_ui.py` + +The focused coverage verifies provenance and source fingerprints, unowned collision rejection, equal-second Cosmos updates, Search transfer and reconciliation beyond 100,000 documents, failed-page cursor retention, Blob mid-transfer heartbeats and pending-to-succeeded verification, mode and baseline validation, two-phase migration-owned mirror deletion, post-cutoff source protection, preview divergence, durable hashed manifests, global coordination, cancellation/retry recovery, route protection, and safe browser rendering. + +## Limitations + +- The application migrates selected SimpleChat scopes, not every arbitrary resource in the source account. +- Azure AI Search request-unit charges are not available through the data-plane response; job telemetry reports item and byte rates for Search. +- Throughput above 10,000 RU/s remains Azure portal managed and is not changed by SimpleChat. +- A migration that encounters unowned destination records, Search keys, or Blob paths must be reconciled by an administrator before it can continue. The job detail records the collision rather than overwriting destination data. +- A global destination coordinator prevents partially overlapping SimpleChat Data Management migrations from running concurrently. The target Search write gate also coordinates normal SimpleChat target indexing, but external writers remain outside that gate and must be frozen for final cutover. +- The browser workflow test requires a configured authenticated UI environment; the static UI contract remains runnable locally. +- AI Search has no shared source modified-time field, so delta mode keyset-scans selected documents and avoids writes through source-hash comparison. +- Selected-scope Blob mirror deletion requires the destination blob's scope hash. Older migration-owned blobs without that marker are retained and reported as unresolved rather than deleted. + +## Version References + +- Application version updated in `application/single_app/config.py` to `0.250.071`. +- This documentation and the related functional tests use the same implementation version. \ No newline at end of file diff --git a/docs/explanation/features/MIGRATION_PROVENANCE.md b/docs/explanation/features/MIGRATION_PROVENANCE.md new file mode 100644 index 000000000..dea63bcb0 --- /dev/null +++ b/docs/explanation/features/MIGRATION_PROVENANCE.md @@ -0,0 +1,88 @@ +# Migration Provenance + +Implemented in version: **0.250.071** + +## Overview + +Cosmos DB, Azure AI Search, and Azure Storage Account migration scripts now attach durable provenance to successfully migrated items. The provenance makes it possible to recognize a migration run without retaining local copies of the migrated content, and it prevents reprocessing items from the same run or a recent migration window. + +## Dependencies + +- PowerShell 7. +- Azure Cosmos DB SQL API access for Cosmos migrations. +- Azure AI Search admin keys or Azure resource permissions for Search migrations. +- Az.Accounts and Az.Storage access for Storage migrations. Storage metadata updates use the destination blob metadata object exposed by `Get-AzStorageBlob`. + +## Technical Specifications + +### Shared Migration Identity + +`scripts/Migration-State.ps1` now persists one migration GUID and its start timestamp in each state file. The scripts accept these common parameters: + +- `-MigrationId `: optional. When omitted, the script generates a GUID. Pass the same GUID to the Cosmos, AI Search, and Storage scripts to associate the complete environment migration. +- `-SkipMigratedWithinHours <0-8760>`: defaults to `24`. A value of `0` disables age-based bypassing but continues to bypass successful items associated with the current migration GUID. + +Resuming with the same state file preserves its GUID. Supplying a different GUID with an existing state file fails safely unless `-ResetState` is used. + +### Metadata Format + +All successful migrations record the following values: + +| Value | Purpose | +| --- | --- | +| `migrationId` / `simplechatMigrationId` | GUID for the migration run. | +| `migratedAtUtc` / `simplechatMigratedAtUtc` | UTC ISO 8601 migration timestamp. | +| `status` / `simplechatMigrationStatus` | `succeeded` after the target write completes. | + +The physical representation is destination-specific: + +- Cosmos documents receive a `simplechatMigration` object. +- AI Search indexes receive `simplechatMigrationId`, `simplechatMigratedAtUtc`, and `simplechatMigrationStatus` fields. The fields are filterable and retrievable so the migration can query them safely. +- Storage blobs receive the equivalent values as blob metadata. Existing blob metadata is read and merged before the migration values are written. + +The in-app Data Management migration extends this shared identity with source comparison fields in version `0.250.071`: Cosmos source `_ts` plus canonical content SHA-256, AI Search canonical source SHA-256, Blob ETag/last-modified/size plus MD5 when available, and a non-reversible Blob scope hash used to constrain selected-scope mirror deletion. Blob transfers use a `pending` status until source/destination verification succeeds. + +### Replay Avoidance + +- Cosmos differential migrations retain their existing conflict-based skip behavior. Full migrations query only destination records with successful matching or recent provenance, including their partition-key values, and skip the matching source identities. +- AI Search full migrations query only destination keys with successful matching or recent provenance. Differential migrations retain their existing destination-key skip behavior. +- Storage migrations evaluate destination metadata for every source blob. When all source blobs in a container are marked by the current or recent migration, the entire container bypasses AzCopy. Otherwise, the existing differential or full AzCopy behavior runs and successful destination blobs are stamped afterward. + +Failed Cosmos document writes are retained as explicit failure records in the migration state file. A destination item is never stamped `succeeded` unless its target write completed; an item with no durable destination copy remains eligible for a later retry. + +## Usage + +Generate one GUID before a coordinated migration: + +```powershell +$migrationId = [Guid]::NewGuid().ToString("D") +``` + +Pass it to each migration script: + +```powershell +.\scripts\Migration-Cosmos.ps1 -MigrationId $migrationId +.\scripts\Migration-AISearch.ps1 -MigrationId $migrationId +.\scripts\Migration-StorageAccount.ps1 -MigrationId $migrationId +``` + +To choose a six-hour replay window, include `-SkipMigratedWithinHours 6`. Use `-SkipMigratedWithinHours 0` when only the current migration GUID should bypass already marked items. + +## Testing And Validation + +- `functional_tests/test_migration_provenance.py` +- `functional_tests/test_cosmos_migration_document_skip_reporting.py` +- `functional_tests/test_ai_search_migration_provenance.py` +- `functional_tests/test_storage_migration_provenance.py` + +The coverage verifies GUID persistence, standardized metadata, preservation of existing blob metadata, full-mode Cosmos and AI Search destination bypassing, Storage replay avoidance, and existing error/retry behavior. + +## Limitations + +- Provenance applies to items successfully written to the destination. A failed write has no target item to tag, so its failure is recorded in migration state rather than as destination metadata. +- Storage container bypassing requires every current source blob to have eligible destination provenance. A container with new or unmarked blobs continues through the existing AzCopy flow, which remains non-destructive in differential mode. + +## Version References + +- Application implementation extended in `application/single_app/config.py` version `0.250.071`. +- Functional tests and this document use the same implementation version. \ No newline at end of file diff --git a/docs/explanation/features/index.md b/docs/explanation/features/index.md index 909bd5884..e0171d13b 100644 --- a/docs/explanation/features/index.md +++ b/docs/explanation/features/index.md @@ -37,6 +37,8 @@ category: Version History ## Versioned Features +- [Data Management Migration Resilience](DATA_MANAGEMENT_MIGRATION_RESILIENCE.md) +- [Migration Provenance](MIGRATION_PROVENANCE.md) - [Microsoft Teams App SSO](v0.242.072/TEAMS_APP_SSO.md) - [Tabular SK Large Result Pagination](v0.242.067/TABULAR_SK_LARGE_RESULT_PAGINATION.md) - [Model Endpoint Model Icon Picker](v0.242.060/MODEL_ENDPOINT_MODEL_ICON_PICKER.md) diff --git a/docs/explanation/release_notes.md b/docs/explanation/release_notes.md index 5cc2e5fea..7b5d298ae 100644 --- a/docs/explanation/release_notes.md +++ b/docs/explanation/release_notes.md @@ -2,6 +2,23 @@ For feature-focused and fix-focused drill-downs by version, see [Features by Version](/explanation/features/) and [Fixes by Version](/explanation/fixes/). +### **(v0.250.071)** + +#### New Features + +* **Resilient Data Management Migrations** + * Added durable migration provenance, per-resource checkpoints, bounded concurrent transfers, retry/resume controls, and long-running Cosmos, Search, Blob, inventory, and reconciliation heartbeats. + * Added New only, Delta / upsert, and explicitly confirmed Mirror with deletions modes with durable baseline watermarks, migration-owned update/delete protection, and server-owned live previews. + * Added post-copy Cosmos, AI Search, and Blob reconciliation with cutover readiness, actual outcome totals, and preview-versus-actual divergence reporting. + * Replaced deep AI Search pagination with persisted `id` keyset cursors, including selected-scope filter batching and coverage beyond 100,000 documents. + * Added memory-bounded reconciliation, two-phase mirror deletion with current-source/ETag revalidation, target-side coordination across independent SimpleChat sources, and retryable `not_ready` cutover failures. + * Preserved Blob content settings, metadata, tags, tier, type, and available content verification while adding mid-stream lease heartbeats, cooperative cancellation, source ETag stability, and pending-to-succeeded provenance. + * Added durable privacy-safe per-item outcome batches plus admin JSONL manifest and failure-list downloads. + * Made unresolved temporary Cosmos capacity restoration a retryable terminal state; retry restores the saved snapshot without reapplying the boost. + * Added a durable target AI Search write fence that drains normal SimpleChat indexing before transfer or mirror deletion, bounds Search requests, rejects self-targeting migrations, and requires explicit external-writer freeze acknowledgement before Search writes begin. + * Added target-side coordinator fencing across independent SimpleChat source deployments and retained uncertain target Search write slots through their full quarantine window. Authorization-reducing unshare requests now defer safely while target Search ACL writes are frozen. + * (Ref: microsoft/simplechat#1043, `functions_data_management.py`, `functions_migration_provenance.py`, `admin_data_management.js`, `DATA_MANAGEMENT_MIGRATION_RESILIENCE.md`) + ### **(v0.250.070)** #### Bug Fixes @@ -30,7 +47,6 @@ For feature-focused and fix-focused drill-downs by version, see [Features by Ver * Added Azure Blob Storage as an admin-controlled File Sync source for personal, group, and public workspaces, with account, container, prefix, selected-path, filter, tag, schedule, and remote-delete controls. * Added managed identity, Key Vault-backed service principal and connection string authentication, connection testing, virtual-folder browsing, ETag change detection, and streamed ingestion through the existing document pipeline. * (Ref: microsoft/simplechat#1027, `functions_file_sync.py`, `workspace-file-sync.js`, `AZURE_BLOB_STORAGE_FILE_SYNC.md`) - ### **(v0.250.066)** #### Bug Fixes diff --git a/functional_tests/test_ai_search_migration_provenance.py b/functional_tests/test_ai_search_migration_provenance.py new file mode 100644 index 000000000..5448403f8 --- /dev/null +++ b/functional_tests/test_ai_search_migration_provenance.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +""" +Functional test for AI Search migration provenance. +Version: 0.250.067 +Implemented in: 0.250.074 + +This test ensures AI Search migrations add provenance fields, tag copied +documents, and skip destination documents migrated within the configured window. +""" + +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile + + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCRIPT_PATH = REPO_ROOT / "scripts" / "Migration-AISearch.ps1" + + +def powershell_single_quote(value: Path) -> str: + """Escape a path for a PowerShell single-quoted string.""" + return str(value).replace("'", "''") + + +def test_ai_search_migration_provenance() -> None: + """Exercise schema tagging and recent destination document bypass behavior.""" + powershell = shutil.which("pwsh") + assert powershell, "PowerShell 7 is required to test the AI Search migration script." + + with tempfile.TemporaryDirectory(prefix="simplechat-ai-search-provenance-") as temp_dir: + state_path = Path(temp_dir) / "migration-state.json" + harness = r''' +$ErrorActionPreference = "Stop" +$scriptPath = '__SCRIPT_PATH__' +$statePath = '__STATE_PATH__' +$global:updatedIndex = $null +$global:writtenDocuments = [System.Collections.Generic.List[object]]::new() +$global:provenanceQueryCount = 0 + +function Convert-MockRequestBody { + param( + [AllowNull()] + [object]$Body + ) + + if ($Body -is [string]) { + return $Body | ConvertFrom-Json -Depth 100 + } + return $Body +} + +function Invoke-RestMethod { + [CmdletBinding()] + param( + [string]$Method, + [string]$Uri, + [hashtable]$Headers, + [AllowNull()] + [object]$Body, + [string]$ContentType, + [int]$TimeoutSec + ) + + $isSource = $Uri.StartsWith("https://source-search.") + $isDestination = $Uri.StartsWith("https://destination-search.") + if (-not $isSource -and -not $isDestination) { + throw "Unexpected endpoint: $Uri" + } + + if ($Method -eq "GET" -and $Uri -match "/synonymmaps\?") { + return [pscustomobject]@{ value = @() } + } + + if ($Method -eq "GET" -and $Uri -match "/indexes\?") { + $index = [pscustomobject]@{ + name = "documents" + fields = @( + [pscustomobject]@{ + name = "id" + type = "Edm.String" + key = $true + filterable = $true + sortable = $true + retrievable = $true + }, + [pscustomobject]@{ + name = "content" + type = "Edm.String" + searchable = $true + retrievable = $true + } + ) + } + return [pscustomobject]@{ value = @($index) } + } + + if ($Method -eq "GET" -and $Uri -match "/indexes/documents/docs/\$count\?") { + if ($isSource) { + return 2 + } + return 1 + } + + if ($Method -eq "PUT" -and $isDestination -and $Uri -match "/indexes/documents\?") { + $global:updatedIndex = Convert-MockRequestBody -Body $Body + return [pscustomobject]@{} + } + + if ($Method -eq "POST" -and $isDestination -and $Uri -match "/indexes/documents/docs/search\?") { + $request = Convert-MockRequestBody -Body $Body + if ([string]$request.filter -notmatch "simplechatMigrationStatus") { + throw "Destination provenance lookup did not include the migration filter." + } + $global:provenanceQueryCount++ + return [pscustomobject]@{ + value = @([pscustomobject]@{ id = "recent-migrated-doc" }) + } + } + + if ($Method -eq "POST" -and $isSource -and $Uri -match "/indexes/documents/docs/search\?") { + return [pscustomobject]@{ + value = @( + [pscustomobject]@{ id = "recent-migrated-doc"; content = "skip" }, + [pscustomobject]@{ id = "copy-doc"; content = "copy" } + ) + } + } + + if ($Method -eq "POST" -and $isDestination -and $Uri -match "/indexes/documents/docs/index\?") { + $request = Convert-MockRequestBody -Body $Body + foreach ($document in @($request.value)) { + $global:writtenDocuments.Add($document) + } + return [pscustomobject]@{ + value = @($request.value | ForEach-Object { + [pscustomobject]@{ + key = $_.id + status = $true + statusCode = 200 + errorMessage = "" + } + }) + } + } + + throw "Unexpected mock request: $Method $Uri" +} + +$parameters = @{ + SourceSearchService = "source-search" + SourceResourceGroup = "source-rg" + SourceSubscriptionId = "00000000-0000-0000-0000-000000000001" + SourceAdminKey = "source-admin-key" + DestinationSearchService = "destination-search" + DestinationResourceGroup = "destination-rg" + DestinationSubscriptionId = "00000000-0000-0000-0000-000000000002" + DestinationAdminKey = "destination-admin-key" + DifferentialMigration = $false + MigrationId = "11111111-1111-1111-1111-111111111111" + SkipMigratedWithinHours = 24 + ShowProgress = $false + BatchSize = 100 + MaxConcurrentBatches = 1 + PageSize = 100 + StateFilePath = $statePath +} + +& $scriptPath @parameters | Out-Null +$state = [IO.File]::ReadAllText($statePath) | ConvertFrom-Json -AsHashtable -Depth 100 +$fieldNames = @($global:updatedIndex.fields | ForEach-Object { $_.name }) +$writtenDocuments = @($global:writtenDocuments) + +if ( + $state.status -ne "completed" -or + $state.migrationId -ne $parameters.MigrationId -or + $state.summary.CopiedCount -ne 1 -or + $state.summary.SkippedCount -ne 1 +) { + throw "Unexpected migration state: $($state | ConvertTo-Json -Compress -Depth 100)" +} +if ( + $global:provenanceQueryCount -ne 1 -or + @("simplechatMigrationId", "simplechatMigratedAtUtc", "simplechatMigrationStatus" | Where-Object { + $fieldNames -notcontains $_ + }).Count -ne 0 +) { + throw "The destination index was not updated with migration provenance fields: $($global:updatedIndex | ConvertTo-Json -Compress -Depth 100)" +} +if ($writtenDocuments.Count -ne 1 -or $writtenDocuments[0].id -ne "copy-doc") { + throw "Recently migrated documents were not bypassed: $($writtenDocuments | ConvertTo-Json -Compress -Depth 100)" +} +$writtenDocument = $writtenDocuments[0] +if ( + $writtenDocument.simplechatMigrationId -ne $parameters.MigrationId -or + $writtenDocument.simplechatMigrationStatus -ne "succeeded" -or + [string]::IsNullOrWhiteSpace([string]$writtenDocument.simplechatMigratedAtUtc) +) { + throw "Copied document did not include migration provenance: $($writtenDocument | ConvertTo-Json -Compress -Depth 100)" +} +''' + harness = harness.replace( + "__SCRIPT_PATH__", powershell_single_quote(SCRIPT_PATH) + ).replace("__STATE_PATH__", powershell_single_quote(state_path)) + + result = subprocess.run( + [ + powershell, + "-NoLogo", + "-NoProfile", + "-ExecutionPolicy", + "Bypass", + "-Command", + harness, + ], + cwd=REPO_ROOT, + text=True, + capture_output=True, + check=False, + ) + + if result.returncode != 0: + print(result.stdout) + print(result.stderr, file=sys.stderr) + raise AssertionError( + f"AI Search migration provenance harness failed with exit code {result.returncode}." + ) + + +if __name__ == "__main__": + test_ai_search_migration_provenance() + print("AI Search migration provenance test passed.") \ No newline at end of file diff --git a/functional_tests/test_cosmos_migration_document_skip_reporting.py b/functional_tests/test_cosmos_migration_document_skip_reporting.py new file mode 100644 index 000000000..fa1a5562d --- /dev/null +++ b/functional_tests/test_cosmos_migration_document_skip_reporting.py @@ -0,0 +1,465 @@ +# test_cosmos_migration_document_skip_reporting.py +#!/usr/bin/env python3 +""" +Functional test for Cosmos migration JSON property and skip reporting. +Version: 0.250.067 +Implemented in: 0.250.064 + +This test ensures documents with empty or case-distinct JSON property names +are copied correctly and document-scoped write failures are recorded without +stopping later writes. +""" + +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile + + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCRIPT_PATH = REPO_ROOT / "scripts" / "Migration-Cosmos.ps1" + + +def _powershell_single_quote(value: Path) -> str: + """Escape a path for a PowerShell single-quoted string.""" + return str(value).replace("'", "''") + + +def test_cosmos_migration_document_skip_reporting() -> None: + """Exercise empty-name JSON fields and nonfatal write rejection reporting.""" + powershell = shutil.which("pwsh") + if not powershell: + raise AssertionError("PowerShell 7 is required to test the migration script.") + + with tempfile.TemporaryDirectory(prefix="simplechat-cosmos-skip-") as temp_dir: + state_path = Path(temp_dir) / "migration-state.json" + harness = r''' +$ErrorActionPreference = "Stop" +$scriptPath = '__SCRIPT_PATH__' +$statePath = '__STATE_PATH__' +$global:mockFatalWrites = $false +$global:mockThrottleResponsesRemaining = 0 +$global:mockSleepRecords = [System.Collections.Generic.List[object]]::new() +$global:mockProvenanceWriteCount = 0 +$global:mockDestinationProvenanceQueryCount = 0 + +function Start-Sleep { + [CmdletBinding()] + param( + [int]$Milliseconds, + [int]$Seconds + ) + + $global:mockSleepRecords.Add([pscustomobject]@{ + Milliseconds = $Milliseconds + Seconds = $Seconds + }) +} + +function Invoke-WebRequest { + [CmdletBinding()] + param( + [string]$Method, + [string]$Uri, + [hashtable]$Headers, + [AllowNull()] + [string]$Body, + [string]$ContentType, + [switch]$SkipHttpErrorCheck + ) + + $isSource = $Uri.StartsWith("https://source-cosmos.") + $isDestination = $Uri.StartsWith("https://destination-cosmos.") + if (-not $isSource -and -not $isDestination) { + throw "Unexpected endpoint: $Uri" + } + + if ($Method -eq "GET" -and $Uri -match "/dbs/SimpleChat$") { + return [pscustomobject]@{ + StatusCode = 200 + Headers = @{} + Content = '{"id":"SimpleChat"}' + } + } + + if ($Method -eq "GET" -and $Uri -match "/dbs/SimpleChat/colls$") { + return [pscustomobject]@{ + StatusCode = 200 + Headers = @{} + Content = '{"DocumentCollections":[{"id":"documents","partitionKey":{"paths":["/tenant"],"kind":"Hash","version":2},"statistics":[{"documentCount":5}]}]}' + } + } + + if ( + $Method -eq "GET" -and + $isSource -and + $Uri -match "/dbs/SimpleChat/colls/documents/docs$" + ) { + $recentMigrationTimestamp = [DateTime]::UtcNow.ToString("o") + return [pscustomobject]@{ + StatusCode = 200 + Headers = @{} + Content = ('{"Documents":[{"id":"empty-key-doc","tenant":"a","":"root-empty","Name":"upper","name":"lower","nested":{"":"nested-empty"},"_rid":"remove"},{"id":"bad-doc","tenant":"b","content":"reject"},{"tenant":"missing-id","content":"cannot write"},{"id":"after-doc","tenant":"c","content":"continues"},{"id":"recent-migrated-doc","tenant":"d","simplechatMigration":{"migrationId":"22222222-2222-2222-2222-222222222222","migratedAtUtc":"' + $recentMigrationTimestamp + '","status":"succeeded"}}],"_count":5}') + } + } + + if ( + $Method -eq "PUT" -and + $isDestination -and + $Uri -match "/dbs/SimpleChat/colls/documents$" + ) { + return [pscustomobject]@{ + StatusCode = 200 + Headers = @{} + Content = '{}' + } + } + + if ( + $Method -eq "POST" -and + $isDestination -and + $Uri -match "/dbs/SimpleChat/colls/documents/docs$" -and + $Headers.ContainsKey("x-ms-documentdb-isquery") + ) { + $global:mockDestinationProvenanceQueryCount++ + $recentDestinationTimestamp = [DateTime]::UtcNow.ToString("o") + return [pscustomobject]@{ + StatusCode = 200 + Headers = @{} + Content = ('{"Documents":[{"id":"after-doc","tenant":"c","simplechatMigration":{"migrationId":"11111111-1111-1111-1111-111111111111","migratedAtUtc":"' + $recentDestinationTimestamp + '","status":"succeeded"}}],"_count":1}') + } + } + + if ( + $Method -eq "POST" -and + $isDestination -and + $Uri -match "/dbs/SimpleChat/colls/documents/docs$" + ) { + $document = $Body | ConvertFrom-Json -AsHashtable -Depth 100 + $migrationMetadata = $document["simplechatMigration"] + if ( + $null -eq $migrationMetadata -or + $migrationMetadata["migrationId"] -ne "11111111-1111-1111-1111-111111111111" -or + $migrationMetadata["status"] -ne "succeeded" -or + [string]::IsNullOrWhiteSpace([string]$migrationMetadata["migratedAtUtc"]) + ) { + throw "Migration provenance metadata was missing or incomplete: $Body" + } + if ($document.id -eq "recent-migrated-doc") { + $global:mockProvenanceWriteCount++ + throw "A recently migrated document should not have reached the destination writer." + } + if ($global:mockFatalWrites -and $document.id -eq "after-doc") { + return [pscustomobject]@{ + StatusCode = 503 + Headers = @{} + Content = '{"code":"ServiceUnavailable","message":"mock outage"}' + } + } + if ($document.id -eq "empty-key-doc") { + if ( + $document[""] -ne "root-empty" -or + $document.Name -ne "upper" -or + $document.name -ne "lower" -or + $document.nested[""] -ne "nested-empty" -or + $document.Contains("_rid") + ) { + throw "The empty-name document was not preserved correctly: $Body" + } + } + if ($document.id -eq "bad-doc") { + return [pscustomobject]@{ + StatusCode = 400 + Headers = @{} + Content = '{"code":"BadRequest","message":"mock rejected document"}' + } + } + if ( + $document.id -eq "after-doc" -and + $global:mockThrottleResponsesRemaining -gt 0 + ) { + $global:mockThrottleResponsesRemaining-- + return [pscustomobject]@{ + StatusCode = 429 + Headers = @{} + Content = '{"code":"TooManyRequests","message":"mock throttling"}' + } + } + return [pscustomobject]@{ + StatusCode = 201 + Headers = @{} + Content = $Body + } + } + + throw "Unexpected mock request: $Method $Uri" +} + +$parameters = @{ + SourceCosmosAccount = "source-cosmos" + SourceResourceGroup = "source-rg" + SourceSubscriptionId = "00000000-0000-0000-0000-000000000001" + SourceDatabaseName = "SimpleChat" + SourcePrimaryKey = "c291cmNlLWtleQ==" + DestinationCosmosAccount = "destination-cosmos" + DestinationResourceGroup = "destination-rg" + DestinationSubscriptionId = "00000000-0000-0000-0000-000000000002" + DestinationDatabaseName = "SimpleChat" + DestinationPrimaryKey = "ZGVzdGluYXRpb24ta2V5" + Containers = @("documents") + DifferentialMigration = $true + MigrationId = "11111111-1111-1111-1111-111111111111" + ShowProgress = $false + PageSize = 100 + MaxConcurrentDocuments = 2 + MaxRetryCount = 2 + StateFilePath = $statePath +} + +$migrationOutput = @(& $scriptPath @parameters 3>&1 6>&1) +$outputText = $migrationOutput | Out-String +$state = [IO.File]::ReadAllText($statePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +$result = $state.resources.documents.result +$badSkippedDocument = @($result.SkippedDocuments | Where-Object { + $_.DocumentId -eq "bad-doc" +})[0] +$missingIdSkippedDocument = @($result.SkippedDocuments | Where-Object { + $_.DocumentId -eq "" +})[0] + +if ($state.status -ne "completed") { + throw "Migration state was not completed: $($state.status)" +} +if ( + $state.migrationId -ne $parameters.MigrationId -or + [string]::IsNullOrWhiteSpace([string]$state.migrationStartedUtc) +) { + throw "Migration state did not preserve provenance: $($state | ConvertTo-Json -Compress)" +} +if ( + $result.ProcessedCount -ne 5 -or + $result.CopiedCount -ne 2 -or + $result.SkippedCount -ne 3 -or + $result.ErrorSkippedCount -ne 2 -or + $global:mockProvenanceWriteCount -ne 0 +) { + throw "Unexpected result counts: $($result | ConvertTo-Json -Compress)" +} +if ( + $badSkippedDocument.ContainerName -ne "documents" -or + $badSkippedDocument.Stage -ne "Write" -or + $badSkippedDocument.Attempt -ne 1 -or + $badSkippedDocument.StatusCode -ne 400 -or + $badSkippedDocument.Reason -notmatch "mock rejected document" -or + [string]::IsNullOrWhiteSpace($badSkippedDocument.RecordedUtc) -or + $missingIdSkippedDocument.ContainerName -ne "documents" -or + $missingIdSkippedDocument.Stage -ne "Preparation" -or + $null -ne $missingIdSkippedDocument.StatusCode -or + $missingIdSkippedDocument.Reason -notmatch "without a valid id" -or + [string]::IsNullOrWhiteSpace($missingIdSkippedDocument.RecordedUtc) +) { + throw "Skipped document details were incomplete: $($result.SkippedDocuments | ConvertTo-Json -Compress)" +} +if ($state.summary.ErrorSkippedCount -ne 2) { + throw "Migration summary did not include the error-skipped count." +} +if ( + $outputText -notmatch "Admin review required: 2 document" -or + $outputText -notmatch "result\.SkippedDocuments" -or + $outputText -notmatch "entries in '[^']+migration-state\.json'" +) { + throw "The final admin warning was missing or incomplete: $outputText" +} + +$resumeOutput = @(& $scriptPath @parameters 3>&1 6>&1) | Out-String +$resumedState = [IO.File]::ReadAllText($statePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +if ( + $resumedState.status -ne "completed" -or + $resumedState.resumeCount -ne 1 -or + $resumedState.summary.ErrorSkippedCount -ne 2 -or + $resumeOutput -notmatch "Skipping completed container from migration state: documents" -or + $resumeOutput -notmatch "Admin review required: 2 document" +) { + throw "Resumed migration did not preserve skipped-document reporting: $resumeOutput" +} + +$sequentialStatePath = "$statePath.sequential.json" +$parameters.StateFilePath = $sequentialStatePath +$parameters.MaxConcurrentDocuments = 1 +$sequentialOutput = @(& $scriptPath @parameters 3>&1 6>&1) | Out-String +$sequentialState = [IO.File]::ReadAllText($sequentialStatePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +$sequentialResult = $sequentialState.resources.documents.result +if ( + $sequentialState.status -ne "completed" -or + $sequentialResult.ProcessedCount -ne 5 -or + $sequentialResult.CopiedCount -ne 2 -or + $sequentialResult.SkippedCount -ne 3 -or + $sequentialResult.ErrorSkippedCount -ne 2 -or + @($sequentialResult.SkippedDocuments | Where-Object { + $_.DocumentId -eq "bad-doc" + }).Count -ne 1 -or + @($sequentialResult.SkippedDocuments | Where-Object { + $_.DocumentId -eq "" + }).Count -ne 1 +) { + throw "Sequential migration did not copy and record skips correctly: $($sequentialResult | ConvertTo-Json -Compress)" +} +if ($sequentialOutput -notmatch "Admin review required: 2 document") { + throw "Sequential migration did not warn the admin about its skipped document." +} + +$recoveryStatePath = "$statePath.recovery.json" +$parameters.StateFilePath = $recoveryStatePath +$parameters.MaxConcurrentDocuments = 1 +$parameters.MaxRetryCount = 2 +$parameters.MaxThrottleRecoveryPauses = 2 +$parameters.ThrottleRecoveryPauseSeconds = 1 +$global:mockSleepRecords.Clear() +$global:mockThrottleResponsesRemaining = 2 +$recoveryOutput = @(& $scriptPath @parameters 3>&1 6>&1) | Out-String +$recoveryState = [IO.File]::ReadAllText($recoveryStatePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +$recoveryResult = $recoveryState.resources.documents.result +if ( + $recoveryState.status -ne "completed" -or + $recoveryResult.CopiedCount -ne 2 -or + $recoveryResult.ThrottleRecoveryPauseCount -ne 1 -or + $global:mockThrottleResponsesRemaining -ne 0 -or + @($global:mockSleepRecords | Where-Object { $_.Seconds -eq 1 }).Count -ne 1 +) { + throw "Exhausted 429 retries did not recover correctly: $($recoveryResult | ConvertTo-Json -Compress)" +} +if ($recoveryOutput -notmatch "Recovery pause 1/2 begins now") { + throw "Throttling recovery warning was missing: $recoveryOutput" +} + +$fullProvenanceStatePath = "$statePath.full-provenance.json" +$parameters.StateFilePath = $fullProvenanceStatePath +$parameters.DifferentialMigration = $false +$parameters.MaxConcurrentDocuments = 1 +$global:mockDestinationProvenanceQueryCount = 0 +$fullProvenanceOutput = @(& $scriptPath @parameters 3>&1 6>&1) | Out-String +$fullProvenanceState = [IO.File]::ReadAllText($fullProvenanceStatePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +$fullProvenanceResult = $fullProvenanceState.resources.documents.result +if ( + $fullProvenanceState.status -ne "completed" -or + $fullProvenanceResult.ProcessedCount -ne 5 -or + $fullProvenanceResult.CopiedCount -ne 1 -or + $fullProvenanceResult.SkippedCount -ne 4 -or + $fullProvenanceResult.ErrorSkippedCount -ne 2 -or + $global:mockDestinationProvenanceQueryCount -ne 1 -or + $fullProvenanceOutput -notmatch "Skipping 1 recently migrated document" +) { + throw "Full migration did not skip the destination-marked document: $($fullProvenanceResult | ConvertTo-Json -Compress)" +} +$parameters.DifferentialMigration = $true + +$exhaustedRecoveryStatePath = "$statePath.recovery-exhausted.json" +$parameters.StateFilePath = $exhaustedRecoveryStatePath +$parameters.MaxThrottleRecoveryPauses = 1 +$global:mockSleepRecords.Clear() +$global:mockThrottleResponsesRemaining = 4 +$exhaustedRecoveryError = $null +try { + & $scriptPath @parameters -ErrorAction Stop 3>&1 6>&1 | Out-Null +} +catch { + $exhaustedRecoveryError = $_ +} +if ( + $null -eq $exhaustedRecoveryError -or + $exhaustedRecoveryError.Exception.Message -notmatch "throttling recovery exhausted" -or + $exhaustedRecoveryError.Exception.Message -notmatch "1 recovery pause" -or + $exhaustedRecoveryError.Exception.Message -notmatch "Reduce -MaxConcurrentDocuments" -or + @($global:mockSleepRecords | Where-Object { $_.Seconds -eq 1 }).Count -ne 1 +) { + throw "Throttling recovery exhaustion was not reported cleanly: $exhaustedRecoveryError" +} +$exhaustedRecoveryState = [IO.File]::ReadAllText($exhaustedRecoveryStatePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +if ( + $exhaustedRecoveryState.status -ne "failed" -or + $exhaustedRecoveryState.lastError -notmatch "throttling recovery exhausted" +) { + throw "Throttling recovery exhaustion did not persist failed state." +} + +$fatalStatePath = "$statePath.fatal.json" +$parameters.StateFilePath = $fatalStatePath +$parameters.MaxRetryCount = 1 +$global:mockFatalWrites = $true +$fatalError = $null +$fatalOutput = [System.Collections.Generic.List[string]]::new() +try { + & $scriptPath @parameters -ErrorAction Stop 3>&1 6>&1 | + ForEach-Object { $fatalOutput.Add([string]$_) } +} +catch { + $fatalError = $_ +} +finally { + $global:mockFatalWrites = $false +} +if ($null -eq $fatalError -or $fatalError.Exception.Message -notmatch "HTTP 503") { + throw "A systemic destination failure was incorrectly skipped: $fatalError" +} +$fatalOutputText = $fatalOutput -join "`n" +if ( + $fatalOutputText -notmatch "Admin review required: 2 document" -or + $fatalOutputText -notmatch "progress\.SkippedDocuments" +) { + throw "The failed migration did not warn the admin about earlier document skips: $fatalOutputText" +} +$fatalState = [IO.File]::ReadAllText($fatalStatePath) | + ConvertFrom-Json -AsHashtable -Depth 100 +if ( + $fatalState.status -ne "failed" -or + $fatalState.resources.documents.status -ne "failed" -or + $fatalState.lastError -notmatch "HTTP 503" -or + $fatalState.resources.documents.progress.ErrorSkippedCount -ne 2 -or + @($fatalState.resources.documents.progress.SkippedDocuments | Where-Object { + $_.DocumentId -eq "bad-doc" -and $_.StatusCode -eq 400 + }).Count -ne 1 -or + @($fatalState.resources.documents.progress.SkippedDocuments | Where-Object { + $_.DocumentId -eq "" -and $_.Stage -eq "Preparation" + }).Count -ne 1 +) { + throw "Systemic destination failure was not persisted as failed state: $($fatalState | ConvertTo-Json -Compress -Depth 100)" +} +''' + harness = harness.replace( + "__SCRIPT_PATH__", _powershell_single_quote(SCRIPT_PATH) + ).replace("__STATE_PATH__", _powershell_single_quote(state_path)) + + result = subprocess.run( + [ + powershell, + "-NoLogo", + "-NoProfile", + "-ExecutionPolicy", + "Bypass", + "-Command", + harness, + ], + cwd=REPO_ROOT, + text=True, + capture_output=True, + check=False, + ) + + if result.returncode != 0: + print(result.stdout) + print(result.stderr, file=sys.stderr) + raise AssertionError( + f"Cosmos migration skip-reporting harness failed with exit code {result.returncode}." + ) + + +if __name__ == "__main__": + test_cosmos_migration_document_skip_reporting() + print("Cosmos migration document skip-reporting test passed.") \ No newline at end of file diff --git a/functional_tests/test_data_management_ai_search_migration_resilience.py b/functional_tests/test_data_management_ai_search_migration_resilience.py new file mode 100644 index 000000000..72376b889 --- /dev/null +++ b/functional_tests/test_data_management_ai_search_migration_resilience.py @@ -0,0 +1,838 @@ +# test_data_management_ai_search_migration_resilience.py +""" +Functional test for resilient Data Management AI Search migration. +Version: 0.250.071 +Implemented in: 0.250.075 +Updated in: 0.250.071 + +This test ensures the in-app migration skips matching destination provenance, +uploads bounded batches concurrently, and tags every copied Search document. +""" + +import copy +import importlib.util +from pathlib import Path +import re +import sys +import threading +import time +import types + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import initialize_migration_state +from functions_migration_provenance import create_migration_provenance_context + + +class FakeJobContainer: + """Persist deep copies like the production job container.""" + + def __init__(self): + self.manifest_batches = [] + + def upsert_item(self, body): + return copy.deepcopy(body) + + def create_item(self, body): + self.manifest_batches.append(copy.deepcopy(body)) + return copy.deepcopy(body) + + +class FakeSearchClient: + """Simulate Search reads, target provenance lookup, and parallel uploads.""" + + def __init__(self, documents=None, migration_id="", migrated_source_hash=""): + self.documents = documents or [] + self.migration_id = migration_id + self.migrated_source_hash = migrated_source_hash + self.source_queries = [] + self.uploaded = [] + self.active_uploads = 0 + self.maximum_active_uploads = 0 + self.lock = threading.Lock() + + def search(self, **kwargs): + selected_fields = kwargs.get("select") or [] + if "simplechatMigrationId" in selected_fields: + if "already-migrated" not in str(kwargs.get("filter") or ""): + return iter([]) + return iter([{ + "id": "already-migrated", + "simplechatMigrationId": self.migration_id, + "simplechatMigratedAtUtc": "2026-07-24T12:00:00+00:00", + "simplechatMigrationStatus": "succeeded", + "simplechatMigrationSourceHash": self.migrated_source_hash, + }]) + self.source_queries.append(copy.deepcopy(kwargs)) + assert "skip" not in kwargs + assert kwargs.get("order_by") == ["id asc"] + documents = sorted(copy.deepcopy(self.documents), key=lambda item: item["id"]) + cursor_match = re.search( + r"id gt '([^']+)'", + str(kwargs.get("filter") or ""), + ) + if cursor_match: + documents = [item for item in documents if item["id"] > cursor_match.group(1)] + return iter(documents[:int(kwargs.get("top") or len(documents))]) + + def upload_documents(self, documents, **_kwargs): + with self.lock: + self.active_uploads += 1 + self.maximum_active_uploads = max(self.maximum_active_uploads, self.active_uploads) + try: + time.sleep(0.03) + self.uploaded.extend(copy.deepcopy(documents)) + return [{"succeeded": True} for _ in documents] + finally: + with self.lock: + self.active_uploads -= 1 + + +class GeneratedSearchClient: + """Generate deterministic keyset pages without materializing a large index.""" + + def __init__(self, document_count): + self.document_count = document_count + self.source_queries = [] + + def search(self, **kwargs): + assert "skip" not in kwargs + assert kwargs.get("order_by") == ["id asc"] + page_size = int(kwargs.get("top") or 0) + assert 0 < page_size <= 1000 + self.source_queries.append(copy.deepcopy(kwargs)) + cursor_match = re.search( + r"id gt 'document-(\d+)'", + str(kwargs.get("filter") or ""), + ) + start_index = int(cursor_match.group(1)) + 1 if cursor_match else 0 + end_index = min(self.document_count, start_index + page_size) + return iter( + {"id": f"document-{index:06d}", "user_id": "user-1"} + for index in range(start_index, end_index) + ) + + +class CountingTargetSearchClient: + """Count uploaded Search documents while returning per-document outcomes.""" + + def __init__(self): + self.uploaded_count = 0 + self.lock = threading.Lock() + + def search(self, **_kwargs): + return iter([]) + + def upload_documents(self, documents, **_kwargs): + with self.lock: + self.uploaded_count += len(documents) + return [{"succeeded": True} for _ in documents] + + +class DeltaTargetSearchClient: + """Expose owned target hashes and capture create/update Search uploads.""" + + def __init__(self, documents): + self.documents = copy.deepcopy(documents) + self.uploaded = [] + + def search(self, **kwargs): + filter_text = str(kwargs.get("filter") or "") + requested_ids = re.findall(r"id eq '([^']+)'", filter_text) + return iter( + copy.deepcopy(self.documents[document_id]) + for document_id in requested_ids + if document_id in self.documents + ) + + def upload_documents(self, documents, **_kwargs): + self.uploaded.extend(copy.deepcopy(documents)) + for document in documents: + self.documents[document["id"]] = copy.deepcopy(document) + return [{"succeeded": True} for _ in documents] + + +class FailOnceTargetSearchClient(CountingTargetSearchClient): + """Fail one document result once, then accept the same key on retry.""" + + def __init__(self, failed_document_id): + super().__init__() + self.failed_document_id = failed_document_id + self.failure_returned = False + self.attempted_ids = [] + self.documents = {} + + def search(self, **kwargs): + requested_ids = re.findall( + r"id eq '([^']+)'", + str(kwargs.get("filter") or ""), + ) + return iter( + copy.deepcopy(self.documents[document_id]) + for document_id in requested_ids + if document_id in self.documents + ) + + def upload_documents(self, documents, **_kwargs): + self.attempted_ids.extend(document["id"] for document in documents) + results = [] + for document in documents: + should_fail = ( + document["id"] == self.failed_document_id and + not self.failure_returned + ) + if should_fail: + self.failure_returned = True + results.append({"succeeded": False}) + continue + with self.lock: + self.uploaded_count += 1 + self.documents[document["id"]] = copy.deepcopy(document) + results.append({"succeeded": True}) + return results + + +class AmbiguousTargetSearchClient: + """Simulate a response-lost Search write whose target state changes before retry.""" + + def __init__(self, response_lost_document): + self.documents = {} + self.response_lost_document = response_lost_document + self.upload_attempts = 0 + + def search(self, **kwargs): + requested_ids = re.findall( + r"id eq '([^']+)'", + str(kwargs.get("filter") or ""), + ) + return iter( + copy.deepcopy(self.documents[document_id]) + for document_id in requested_ids + if document_id in self.documents + ) + + def upload_documents(self, documents, **_kwargs): + self.upload_attempts += 1 + if self.upload_attempts > 1: + raise AssertionError("An ambiguous AI Search write must be reclassified before retry.") + document = documents[0] + self.documents[document["id"]] = copy.deepcopy(self.response_lost_document(document)) + raise TimeoutError("Search response was lost after the request reached the service.") + + +def load_data_management_module(monkeypatch, source_client, job_container): + """Load the production module without loading the full Flask application config.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {"search_client_user": source_client} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_ai_search_resilience_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_data_management_ai_search_migration_parallel_provenance(monkeypatch): + """Validate target marker skips and concurrent tagged uploads through the app path.""" + migration_id = "11111111-1111-1111-1111-111111111111" + source_client = FakeSearchClient([ + {"id": "already-migrated", "user_id": "user-1", "chunk_text": "skip"}, + {"id": "copy-one", "user_id": "user-1", "chunk_text": "one"}, + {"id": "copy-two", "user_id": "user-1", "chunk_text": "two"}, + {"id": "copy-three", "user_id": "user-1", "chunk_text": "three"}, + {"id": "copy-four", "user_id": "user-1", "chunk_text": "four"}, + ]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_client, job_container) + module.DATA_MANAGEMENT_MIGRATION_BATCH_SIZE = 2 + target_client = FakeSearchClient( + migration_id=migration_id, + migrated_source_hash=module._build_search_source_hash(source_client.documents[0]), + ) + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "updated_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + + settings = { + "migration_max_parallel_operations": 2, + "migration_retry_count": 2, + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "search"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-24T12:00:00+00:00", + ) + job = {"id": migration_id, "migration_state": migration_state} + + artifacts = module._copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + + assert len(target_client.uploaded) == 4 + assert target_client.maximum_active_uploads >= 2 + assert all( + document["simplechatMigrationId"] == migration_id + and document["simplechatMigrationStatus"] == "succeeded" + and document["simplechatMigratedAtUtc"] + for document in target_client.uploaded + ) + assert artifacts[0]["copied_count"] == 4 + assert artifacts[0]["skipped_count"] == 0 + assert artifacts[0]["destination_provenance_skip_count"] == 1 + assert artifacts[0]["batch_size"] == 2 + assert artifacts[0]["source_read_count"] == 5 + assert artifacts[0]["keyset_cursor"]["completed"] is True + assert artifacts[0]["checkpoint_count"] >= 2 + assert len(source_client.source_queries) >= 2 + assert any("id gt" in str(query.get("filter") or "") for query in source_client.source_queries) + resource = job["migration_state"]["resources"]["ai_search:users:simplechat-user-index"] + assert resource["status"] == "completed" + + try: + module._classify_target_search_document( + {"id": "unowned"}, + provenance_context, + "simplechat-user-index", + ) + except module.DataManagementSettingsValidationError as exc: + assert "unowned" in str(exc) + else: + raise AssertionError("An unowned destination AI Search document was accepted for overwrite.") + + +def test_target_search_schema_client_uses_migration_request_bounds(monkeypatch): + """Keep schema reads and writes inside the same bounded migration request budget.""" + source_client = FakeSearchClient() + module = load_data_management_module(monkeypatch, source_client, FakeJobContainer()) + captured_kwargs = {} + + class FakeIndexClient: + def __init__(self, **kwargs): + captured_kwargs.update(kwargs) + + def get_index(self, _index_name): + return types.SimpleNamespace(name="simplechat-user-index", fields=[]) + + def create_or_update_index(self, index): + return index + + monkeypatch.setattr(module, "SearchIndexClient", FakeIndexClient) + + assert module._ensure_target_search_index( + {"target_ai_search_endpoint": "https://target.search.windows.net"}, + "simplechat-user-index", + "unused.json", + ) == "updated_with_migration_provenance" + assert captured_kwargs["connection_timeout"] == ( + module.DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + ) + assert captured_kwargs["read_timeout"] == ( + module.DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + ) + assert captured_kwargs["retry_total"] == 0 + + +def test_data_management_ai_search_keyset_migrates_beyond_skip_limit(monkeypatch): + """Migrate more than 100,000 source keys without SDK continuation or deep skip.""" + document_count = 100_001 + migration_id = "22222222-2222-2222-2222-222222222222" + source_client = GeneratedSearchClient(document_count) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_client, job_container) + target_client = CountingTargetSearchClient() + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "exists_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + + settings = { + "migration_max_parallel_operations": 8, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "search-scale"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + job = {"id": migration_id, "migration_state": migration_state} + + artifacts = module._copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + + assert target_client.uploaded_count == document_count + assert artifacts[0]["copied_count"] == document_count + assert artifacts[0]["source_read_count"] == document_count + assert artifacts[0]["keyset_cursor"]["completed"] is True + assert len(source_client.source_queries) > 100 + assert all("skip" not in query for query in source_client.source_queries) + + +def test_data_management_ai_search_resumes_from_persisted_keyset_cursor(monkeypatch): + """Resume after a durable Search page without rescanning completed source keys.""" + migration_id = "33333333-3333-3333-3333-333333333333" + source_client = FakeSearchClient([ + {"id": f"document-{index:03d}", "user_id": "user-1"} + for index in range(10) + ]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_client, job_container) + module.DATA_MANAGEMENT_MIGRATION_BATCH_SIZE = 2 + target_client = CountingTargetSearchClient() + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "exists_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + + settings = { + "migration_max_parallel_operations": 2, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "search-resume"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + job = {"id": migration_id, "migration_state": migration_state} + persist_checkpoint = module._persist_migration_checkpoint + interrupted = False + + def interrupt_after_first_checkpoint(*args, **kwargs): + nonlocal interrupted + state = persist_checkpoint(*args, **kwargs) + if not interrupted: + interrupted = True + raise RuntimeError("simulated worker interruption") + return state + + monkeypatch.setattr(module, "_persist_migration_checkpoint", interrupt_after_first_checkpoint) + with pytest.raises(RuntimeError, match="simulated worker interruption"): + module._copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + + persisted_progress = job["migration_state"]["resources"][ + "ai_search:users:simplechat-user-index" + ]["progress"] + assert persisted_progress["keyset_cursor"]["last_id"] == "document-003" + first_attempt_query_count = len(source_client.source_queries) + monkeypatch.setattr(module, "_persist_migration_checkpoint", persist_checkpoint) + + artifacts = module._copy_ai_search_to_target( + settings, + migration_plan, + job, + job["migration_state"], + provenance_context, + ) + + resumed_queries = source_client.source_queries[first_attempt_query_count:] + assert resumed_queries + assert "id gt 'document-003'" in str(resumed_queries[0].get("filter") or "") + assert target_client.uploaded_count == 10 + assert artifacts[0]["copied_count"] == 10 + + +def test_data_management_ai_search_batches_large_selected_scope_filters(monkeypatch): + """Keep a 2,000-user selection below Search filter complexity limits.""" + source_client = FakeSearchClient() + module = load_data_management_module(monkeypatch, source_client, FakeJobContainer()) + selected_ids = [f"user-{index:04d}" for index in range(2000)] + + search_filters = module._build_search_scope_filter_batches( + "user_id", + {"mode": "selected", "ids": selected_ids}, + ) + + assert len(search_filters) == 20 + assert all(search_filter.count("user_id eq") <= 100 for search_filter in search_filters) + assert sum(search_filter.count("user_id eq") for search_filter in search_filters) == 2000 + + +def test_data_management_ai_search_delta_writes_only_changed_owned_documents(monkeypatch): + """Use source hashes to create new and update changed migration-owned Search keys.""" + migration_id = "66666666-6666-6666-6666-666666666666" + baseline_migration_id = "55555555-5555-5555-5555-555555555555" + source_documents = [ + {"id": "unchanged", "user_id": "user-1", "chunk_text": "same"}, + {"id": "changed", "user_id": "user-1", "chunk_text": "new"}, + {"id": "created", "user_id": "user-1", "chunk_text": "new"}, + ] + source_client = FakeSearchClient(source_documents) + module = load_data_management_module(monkeypatch, source_client, FakeJobContainer()) + baseline_context = create_migration_provenance_context( + migration_id=baseline_migration_id, + migrated_at_utc="2026-07-28T12:00:00+00:00", + ) + target_client = DeltaTargetSearchClient({ + "unchanged": module.add_search_migration_provenance( + {"id": "unchanged"}, + baseline_context, + source_hash=module._build_search_source_hash(source_documents[0]), + ), + "changed": module.add_search_migration_provenance( + {"id": "changed"}, + baseline_context, + source_hash="sha256:old", + ), + }) + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "exists_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + migration_state = initialize_migration_state(None, migration_id, {"test": "search-delta"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + provenance_context.update({ + "migration_mode": "delta_upsert", + "baseline_source_cutoff_at": "2026-07-28T12:00:00+00:00", + }) + job = {"id": migration_id, "migration_state": migration_state} + + artifacts = module._copy_ai_search_to_target( + { + "migration_max_parallel_operations": 2, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + }, + job, + migration_state, + provenance_context, + ) + + assert artifacts[0]["created_count"] == 1 + assert artifacts[0]["updated_count"] == 1 + assert artifacts[0]["unchanged_count"] == 1 + assert {document["id"] for document in target_client.uploaded} == {"created", "changed"} + assert all( + document["simplechatMigrationSourceHash"].startswith("sha256:") + for document in target_client.uploaded + ) + + +def test_data_management_ai_search_partial_failure_retries_same_keyset_page(monkeypatch): + """Keep the cursor before a partially failed page and clear attempt failures on retry.""" + migration_id = "77777777-7777-7777-7777-777777777777" + source_client = FakeSearchClient([ + {"id": f"document-{index:03d}", "user_id": "user-1"} + for index in range(4) + ]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_client, job_container) + module.DATA_MANAGEMENT_MIGRATION_BATCH_SIZE = 2 + target_client = FailOnceTargetSearchClient("document-001") + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "exists_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + settings = { + "migration_max_parallel_operations": 1, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "search-partial"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + job = {"id": migration_id, "migration_state": migration_state} + + with pytest.raises(RuntimeError, match="before its keyset cursor could advance"): + module._copy_ai_search_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + + failed_progress = job["migration_state"]["resources"][ + "ai_search:users:simplechat-user-index" + ]["progress"] + assert failed_progress["failed_count"] == 1 + assert failed_progress["keyset_cursor"]["last_id"] == "" + first_query_count = len(source_client.source_queries) + + artifacts = module._copy_ai_search_to_target( + settings, + migration_plan, + job, + job["migration_state"], + provenance_context, + ) + + resumed_queries = source_client.source_queries[first_query_count:] + assert resumed_queries + assert "id gt" not in str(resumed_queries[0].get("filter") or "") + assert target_client.attempted_ids.count("document-001") == 2 + assert artifacts[0]["failed_count"] == 0 + assert artifacts[0]["prior_failed_count"] == 1 + assert ( + artifacts[0]["created_count"] + + artifacts[0]["updated_count"] + + artifacts[0]["unchanged_count"] + ) == 4 + manifest_entries = [ + entry + for batch in job_container.manifest_batches + for entry in batch["entries"] + ] + failed_entries = [entry for entry in manifest_entries if entry["status"] == "failed"] + assert len(failed_entries) == 1 + assert failed_entries[0]["source_identity"].startswith("sha256:") + assert "document-001" not in failed_entries[0]["source_identity"] + + +def test_ai_search_ambiguous_write_reclassifies_before_retry(monkeypatch): + """Fail closed on a new unowned target key while recovering a response-lost own write.""" + migration_id = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa" + module = load_data_management_module(monkeypatch, FakeSearchClient(), FakeJobContainer()) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + source_document = {"id": "race", "user_id": "user-1", "chunk_text": "content"} + migration_document = module.add_search_migration_provenance( + copy.deepcopy(source_document), + provenance_context, + source_hash=module._build_search_source_hash(source_document), + ) + + unowned_target = AmbiguousTargetSearchClient( + lambda document: {"id": document["id"], "chunk_text": "external"} + ) + with pytest.raises(module.DataManagementSettingsValidationError, match="unowned document"): + module._upload_search_migration_batch( + unowned_target, + [migration_document], + retry_count=2, + dispositions=["create"], + provenance_context=provenance_context, + index_name="simplechat-user-index", + ) + assert unowned_target.upload_attempts == 1 + + response_lost_own_write = AmbiguousTargetSearchClient(lambda document: document) + result = module._upload_search_migration_batch( + response_lost_own_write, + [migration_document], + retry_count=2, + dispositions=["create"], + provenance_context=provenance_context, + index_name="simplechat-user-index", + ) + assert response_lost_own_write.upload_attempts == 1 + assert result["copied"] == 1 + assert result["created"] == 1 + assert result["failed"] == 0 + + +def test_ai_search_upload_renews_heartbeat_while_request_is_in_flight(monkeypatch): + """Renew the migration lease before a slow Search upload completes.""" + migration_id = "88888888-8888-8888-8888-888888888888" + source_client = FakeSearchClient([{"id": "slow", "user_id": "user-1"}]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_client, job_container) + release_upload = threading.Event() + upload_started = threading.Event() + upload_completed = threading.Event() + + class SlowTargetClient(CountingTargetSearchClient): + def upload_documents(self, documents, **kwargs): + upload_started.set() + assert release_upload.wait(timeout=5.0) + result = super().upload_documents(documents, **kwargs) + upload_completed.set() + return result + + target_client = SlowTargetClient() + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "exists_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + heartbeats = [] + original_heartbeat = module._persist_migration_heartbeat + + def capture_heartbeat(*args, **kwargs): + assert upload_started.is_set() + assert not upload_completed.is_set() + heartbeats.append(time.monotonic()) + release_upload.set() + return original_heartbeat(*args, **kwargs) + + monkeypatch.setattr(module, "_persist_migration_heartbeat", capture_heartbeat) + migration_state = initialize_migration_state(None, migration_id, {"test": "slow-search"}) + + module._copy_ai_search_to_target( + { + "migration_max_parallel_operations": 1, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + }, + {"id": migration_id, "migration_state": migration_state}, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + ) + + assert heartbeats + assert release_upload.is_set() + assert upload_completed.is_set() + + +def test_ai_search_cancellation_records_in_flight_success_before_stopping(monkeypatch): + """Persist an accepted upload without advancing its keyset cursor on cancellation.""" + migration_id = "99999999-9999-9999-9999-999999999999" + source_client = FakeSearchClient([{"id": "slow", "user_id": "user-1"}]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_client, job_container) + release_upload = threading.Event() + upload_started = threading.Event() + + class SlowTargetClient(CountingTargetSearchClient): + def upload_documents(self, documents, **kwargs): + upload_started.set() + assert release_upload.wait(timeout=5.0) + return super().upload_documents(documents, **kwargs) + + target_client = SlowTargetClient() + monkeypatch.setattr( + module, + "_ensure_target_search_index", + lambda *_args, **_kwargs: "exists_with_migration_provenance", + ) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + + def cancel_during_upload(*_args, **_kwargs): + assert upload_started.is_set() + release_upload.set() + raise module.DataManagementMigrationCanceledError("requested") + + monkeypatch.setattr(module, "_persist_migration_heartbeat", cancel_during_upload) + migration_state = initialize_migration_state(None, migration_id, {"test": "cancel-search"}) + job = {"id": migration_id, "migration_state": migration_state} + + with pytest.raises(module.DataManagementMigrationCanceledError, match="requested"): + module._copy_ai_search_to_target( + { + "migration_max_parallel_operations": 1, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + }, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + ) + + progress = job["migration_state"]["resources"][ + "ai_search:users:simplechat-user-index" + ]["progress"] + assert progress["copied_count"] == 1 + assert progress["created_count"] == 1 + assert progress["keyset_cursor"]["last_id"] == "" + manifest_entries = [ + entry + for batch in job_container.manifest_batches + for entry in batch["entries"] + ] + assert [entry["status"] for entry in manifest_entries] == ["created"] \ No newline at end of file diff --git a/functional_tests/test_data_management_blob_migration_resilience.py b/functional_tests/test_data_management_blob_migration_resilience.py new file mode 100644 index 000000000..361eac104 --- /dev/null +++ b/functional_tests/test_data_management_blob_migration_resilience.py @@ -0,0 +1,589 @@ +# test_data_management_blob_migration_resilience.py +""" +Functional test for resilient Data Management blob migration. +Version: 0.250.071 +Implemented in: 0.250.075 + +This test ensures blob migration streams source content, preserves metadata, +uses provenance skips, and exposes concurrent copy progress in job state. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import threading +import time +import types + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import initialize_migration_state +from functions_migration_provenance import create_migration_provenance_context + + +class FakeJobContainer: + """Persist deep copies like the production jobs container.""" + + def __init__(self): + self.manifest_batches = [] + + def upsert_item(self, body): + return copy.deepcopy(body) + + def create_item(self, body): + self.manifest_batches.append(copy.deepcopy(body)) + return copy.deepcopy(body) + + +class FakeSourceDocumentsContainer: + """Expose selected document records through a Cosmos-like iterator.""" + + def __init__(self, documents): + self.documents = documents + + def query_items(self, **_kwargs): + return iter(copy.deepcopy(self.documents)) + + +class FakeDownload: + """Expose only chunk streaming; readall intentionally does not exist.""" + + def __init__(self, payload): + self.payload = payload + + def chunks(self): + yield self.payload[:2] + yield self.payload[2:] + + +class FakeSourceBlob: + """Return source properties and a bounded streaming download.""" + + def __init__(self, payload): + self.payload = payload + + def get_blob_properties(self): + return types.SimpleNamespace( + metadata={"source": "preserved", "destination": "source-value"}, + size=len(self.payload), + content_settings="source-content-settings", + etag=f"etag-{self.payload.decode('utf-8')}", + last_modified=None, + blob_tier="Cool", + blob_type="BlockBlob", + tags={"classification": "internal"}, + ) + + def download_blob(self, **_kwargs): + if _kwargs.get("etag"): + assert _kwargs.get("etag") == f"etag-{self.payload.decode('utf-8')}" + assert _kwargs.get("match_condition") is not None + assert _kwargs.get("timeout") == 120 + return FakeDownload(self.payload) + + +class FakeSourceBlobService: + """Create fake source clients for every persisted blob path.""" + + def get_blob_client(self, **kwargs): + return FakeSourceBlob(f"payload:{kwargs['blob']}".encode("utf-8")) + + +class MissingSourceBlob(FakeSourceBlob): + """Represent a persisted source reference whose Blob no longer exists.""" + + def get_blob_properties(self): + from azure.core.exceptions import ResourceNotFoundError + + raise ResourceNotFoundError("not found") + + +class FakeTargetBlob: + """Store metadata and track concurrent streaming uploads.""" + + def __init__(self, existing_metadata, tracker): + self.existing_metadata = existing_metadata + self.tracker = tracker + self.uploads = [] + self.size = 0 + self.etag = "target-etag" + self.content_settings = None + + def get_blob_properties(self): + if self.existing_metadata is None: + return None + return types.SimpleNamespace( + metadata=copy.deepcopy(self.existing_metadata), + size=self.size, + etag=self.etag, + content_settings=self.content_settings, + ) + + def upload_blob(self, data, **kwargs): + with self.tracker["lock"]: + self.tracker["active"] += 1 + self.tracker["maximum"] = max(self.tracker["maximum"], self.tracker["active"]) + try: + time.sleep(0.03) + self.uploads.append({ + "data": b"".join(data), + "metadata": copy.deepcopy(kwargs["metadata"]), + "content_settings": kwargs["content_settings"], + "length": kwargs["length"], + "overwrite": kwargs["overwrite"], + "tags": copy.deepcopy(kwargs.get("tags")), + "standard_blob_tier": kwargs.get("standard_blob_tier"), + "blob_type": kwargs.get("blob_type"), + "etag": kwargs.get("etag"), + "match_condition": kwargs.get("match_condition"), + "timeout": kwargs.get("timeout"), + }) + self.existing_metadata = copy.deepcopy(kwargs["metadata"]) + self.size = kwargs["length"] + self.content_settings = kwargs["content_settings"] + finally: + with self.tracker["lock"]: + self.tracker["active"] -= 1 + + def set_blob_metadata(self, metadata, **kwargs): + assert kwargs.get("etag") == self.etag + assert kwargs.get("match_condition") is not None + self.existing_metadata = copy.deepcopy(metadata) + + def delete_blob(self, **_kwargs): + self.existing_metadata = None + self.size = 0 + + +class FakeTargetBlobService: + """Expose per-path target blobs and successful container creation.""" + + def __init__(self, migration_id): + self.tracker = {"active": 0, "maximum": 0, "lock": threading.Lock()} + self.blobs = { + "already-migrated.txt": FakeTargetBlob({ + "simplechatMigrationId": migration_id, + "simplechatMigratedAtUtc": "2026-07-24T12:00:00+00:00", + "simplechatMigrationStatus": "succeeded", + }, self.tracker), + "copy-one.txt": FakeTargetBlob(None, self.tracker), + "copy-two.txt": FakeTargetBlob(None, self.tracker), + } + + def create_container(self, _container_name): + return object() + + def get_blob_client(self, **kwargs): + return self.blobs[kwargs["blob"]] + + +def load_data_management_module(monkeypatch, source_documents, source_blob_service, job_container): + """Load the production module with migration dependencies replaced by fakes.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {"storage_account_office_docs_client": source_blob_service} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + config_module.cosmos_user_documents_container = source_documents + config_module.storage_account_user_documents_container_name = "user-documents" + config_module.storage_account_group_documents_container_name = "group-documents" + config_module.storage_account_public_documents_container_name = "public-documents" + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_blob_resilience_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_data_management_blob_migration_streams_and_tags_concurrently(monkeypatch): + """Validate target marker bypass, metadata preservation, and bounded streaming uploads.""" + migration_id = "11111111-1111-1111-1111-111111111111" + source_documents = FakeSourceDocumentsContainer([ + {"id": "one", "user_id": "user-1", "blob_path": "copy-one.txt", "_ts": 1}, + {"id": "two", "user_id": "user-1", "blob_path": "copy-two.txt", "_ts": 1}, + {"id": "already", "user_id": "user-1", "blob_path": "already-migrated.txt", "_ts": 1}, + {"id": "legacy", "user_id": "user-1", "file_name": "legacy.txt", "_ts": 1}, + ]) + source_blob_service = FakeSourceBlobService() + target_blob_service = FakeTargetBlobService(migration_id) + job_container = FakeJobContainer() + module = load_data_management_module( + monkeypatch, + source_documents, + source_blob_service, + job_container, + ) + monkeypatch.setattr( + module, + "_get_target_enhanced_citations_blob_client", + lambda *_args, **_kwargs: target_blob_service, + ) + already_source = source_blob_service.get_blob_client( + container="user-documents", + blob="already-migrated.txt", + ) + already_hash, already_version = module._build_blob_source_fingerprint( + already_source.get_blob_properties() + ) + target_blob_service.blobs["already-migrated.txt"].existing_metadata.update({ + "simplechatMigrationSourceHash": already_hash, + "simplechatMigrationSourceVersion": already_version, + }) + + settings = { + "migration_max_parallel_operations": 2, + "migration_retry_count": 2, + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_source_blobs": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "blobs"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-24T12:00:00+00:00", + ) + job = {"id": migration_id, "migration_state": migration_state} + + artifacts = module._copy_source_blobs_to_target( + settings, + migration_plan, + job, + migration_state, + provenance_context, + ) + + copy_one = target_blob_service.blobs["copy-one.txt"].uploads[0] + copy_two = target_blob_service.blobs["copy-two.txt"].uploads[0] + assert not target_blob_service.blobs["already-migrated.txt"].uploads + assert target_blob_service.tracker["maximum"] >= 2 + assert copy_one["metadata"]["destination"] == "source-value" + assert copy_one["metadata"]["source"] == "preserved" + assert copy_one["metadata"]["simplechatMigrationId"] == migration_id + assert copy_one["metadata"]["simplechatMigrationStatus"] == "pending" + assert target_blob_service.blobs["copy-one.txt"].existing_metadata[ + "simplechatMigrationStatus" + ] == "succeeded" + assert copy_one["metadata"]["simplechatMigrationScopeHash"] == module._document_migration_scope_hash({ + "user_id": "user-1", + }) + assert copy_two["content_settings"] == "source-content-settings" + assert copy_two["length"] == len(copy_two["data"]) + assert copy_one["tags"] == {"classification": "internal"} + assert copy_one["standard_blob_tier"] == "Cool" + assert copy_one["blob_type"] == "BlockBlob" + assert copy_one["overwrite"] is False + assert copy_two["overwrite"] is False + summary = next(artifact for artifact in artifacts if artifact.get("type") == "source_blobs") + assert summary["copied_count"] == 2 + assert summary["skipped_count"] == 0 + assert summary["not_applicable_count"] == 1 + resource = job["migration_state"]["resources"]["source_blobs:selected_documents"] + assert resource["status"] == "completed" + + unowned_target = FakeTargetBlob({"existing": "unowned"}, target_blob_service.tracker) + collision = module._copy_source_blob_migration_record( + FakeSourceBlob(b"collision"), + unowned_target, + provenance_context, + retry_count=1, + ) + assert collision["status"] == "collision" + assert not unowned_target.uploads + + +def test_data_management_blob_delta_updates_only_changed_owned_blobs(monkeypatch): + """Use source versions to create new and conditionally update changed owned blobs.""" + migration_id = "22222222-2222-2222-2222-222222222222" + baseline_migration_id = "11111111-1111-1111-1111-111111111111" + source_documents = FakeSourceDocumentsContainer([ + {"id": "unchanged", "user_id": "user-1", "blob_path": "unchanged.txt", "_ts": 1}, + {"id": "changed", "user_id": "user-1", "blob_path": "changed.txt", "_ts": 1}, + {"id": "created", "user_id": "user-1", "blob_path": "created.txt", "_ts": 1}, + ]) + source_blob_service = FakeSourceBlobService() + job_container = FakeJobContainer() + module = load_data_management_module( + monkeypatch, + source_documents, + source_blob_service, + job_container, + ) + baseline_context = create_migration_provenance_context( + migration_id=baseline_migration_id, + migrated_at_utc="2026-07-28T12:00:00+00:00", + ) + tracker = {"active": 0, "maximum": 0, "lock": threading.Lock()} + unchanged_source = source_blob_service.get_blob_client( + container="user-documents", + blob="unchanged.txt", + ) + unchanged_hash, unchanged_version = module._build_blob_source_fingerprint( + unchanged_source.get_blob_properties() + ) + target_blob_service = FakeTargetBlobService(migration_id) + target_blob_service.blobs = { + "unchanged.txt": FakeTargetBlob( + module.merge_blob_migration_metadata( + {}, + baseline_context, + source_hash=unchanged_hash, + source_version=unchanged_version, + ), + tracker, + ), + "changed.txt": FakeTargetBlob( + module.merge_blob_migration_metadata( + {}, + baseline_context, + source_version="sha256:old", + ), + tracker, + ), + "created.txt": FakeTargetBlob(None, tracker), + } + monkeypatch.setattr( + module, + "_get_target_enhanced_citations_blob_client", + lambda *_args, **_kwargs: target_blob_service, + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "blob-delta"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + provenance_context.update({ + "migration_mode": "delta_upsert", + "baseline_source_cutoff_at": "2026-07-28T12:00:00+00:00", + }) + job = {"id": migration_id, "migration_state": migration_state} + + artifacts = module._copy_source_blobs_to_target( + { + "migration_max_parallel_operations": 2, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_source_blobs": True, + }, + job, + migration_state, + provenance_context, + ) + + summary = next(artifact for artifact in artifacts if artifact.get("type") == "source_blobs") + assert summary["created_count"] == 1 + assert summary["updated_count"] == 1 + assert summary["unchanged_count"] == 1 + assert not target_blob_service.blobs["unchanged.txt"].uploads + changed_upload = target_blob_service.blobs["changed.txt"].uploads[0] + created_upload = target_blob_service.blobs["created.txt"].uploads[0] + assert changed_upload["overwrite"] is True + assert changed_upload["etag"] == "target-etag" + assert changed_upload["match_condition"] is not None + assert created_upload["overwrite"] is False + + +def test_blob_transfer_renews_heartbeat_before_long_copy_completes(monkeypatch): + """Renew the durable lease while a source Blob future is still streaming.""" + migration_id = "33333333-3333-3333-3333-333333333333" + + class SlowDownload(FakeDownload): + def chunks(self): + yield self.payload[:2] + time.sleep(2.2) + yield self.payload[2:] + + class SlowSourceBlob(FakeSourceBlob): + def download_blob(self, **kwargs): + super().download_blob(**kwargs) + return SlowDownload(self.payload) + + class SlowSourceBlobService(FakeSourceBlobService): + def get_blob_client(self, **kwargs): + return SlowSourceBlob(f"payload:{kwargs['blob']}".encode("utf-8")) + + source_documents = FakeSourceDocumentsContainer([ + {"id": "slow", "user_id": "user-1", "blob_path": "slow.txt", "_ts": 1}, + ]) + source_blob_service = SlowSourceBlobService() + job_container = FakeJobContainer() + module = load_data_management_module( + monkeypatch, + source_documents, + source_blob_service, + job_container, + ) + target_blob_service = FakeTargetBlobService(migration_id) + target_blob_service.blobs = { + "slow.txt": FakeTargetBlob(None, target_blob_service.tracker), + } + monkeypatch.setattr( + module, + "_get_target_enhanced_citations_blob_client", + lambda *_args, **_kwargs: target_blob_service, + ) + heartbeat_messages = [] + persist_checkpoint = module._persist_migration_checkpoint + + def capture_checkpoint(*args, **kwargs): + heartbeat_messages.append(args[5]) + return persist_checkpoint(*args, **kwargs) + + monkeypatch.setattr(module, "_persist_migration_checkpoint", capture_checkpoint) + migration_state = initialize_migration_state(None, migration_id, {"test": "blob-heartbeat"}) + job = {"id": migration_id, "migration_state": migration_state} + + module._copy_source_blobs_to_target( + { + "migration_max_parallel_operations": 1, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_source_blobs": True, + }, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + ) + + assert "Migrating source blobs; worker heartbeat renewed" in heartbeat_messages + + +def test_blob_verification_failure_never_leaves_succeeded_provenance(monkeypatch): + """Keep a failed create retryable instead of stamping a false success marker.""" + migration_id = "44444444-4444-4444-4444-444444444444" + source_blob = FakeSourceBlob(b"expected") + tracker = {"active": 0, "maximum": 0, "lock": threading.Lock()} + + class CorruptTargetBlob(FakeTargetBlob): + def get_blob_properties(self): + properties = super().get_blob_properties() + if properties is not None and self.uploads: + properties.size += 1 + return properties + + target_blob = CorruptTargetBlob(None, tracker) + module = load_data_management_module( + monkeypatch, + FakeSourceDocumentsContainer([]), + FakeSourceBlobService(), + FakeJobContainer(), + ) + + result = module._copy_source_blob_migration_record( + source_blob, + target_blob, + create_migration_provenance_context(migration_id=migration_id), + retry_count=1, + ) + + assert result["status"] == "failed" + assert target_blob.existing_metadata is None + + +def test_missing_persisted_blob_fails_while_legacy_no_blob_is_not_applicable(monkeypatch): + """Keep missing references retryable without warning on legacy standard citations.""" + migration_id = "55555555-5555-5555-5555-555555555555" + + class MissingSourceBlobService(FakeSourceBlobService): + def get_blob_client(self, **kwargs): + if kwargs["blob"] == "missing.txt": + return MissingSourceBlob(b"") + return super().get_blob_client(**kwargs) + + source_documents = FakeSourceDocumentsContainer([ + {"id": "missing", "user_id": "user-1", "blob_path": "missing.txt", "_ts": 1}, + {"id": "legacy", "user_id": "user-1", "file_name": "legacy.txt", "_ts": 1}, + ]) + source_blob_service = MissingSourceBlobService() + job_container = FakeJobContainer() + module = load_data_management_module( + monkeypatch, + source_documents, + source_blob_service, + job_container, + ) + target_blob_service = FakeTargetBlobService(migration_id) + target_blob_service.blobs = { + "missing.txt": FakeTargetBlob(None, target_blob_service.tracker), + } + monkeypatch.setattr( + module, + "_get_target_enhanced_citations_blob_client", + lambda *_args, **_kwargs: target_blob_service, + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "missing-blob"}) + job = {"id": migration_id, "migration_state": migration_state} + + with pytest.raises(RuntimeError, match="remediation before retry"): + module._copy_source_blobs_to_target( + { + "migration_max_parallel_operations": 1, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_source_blobs": True, + }, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + ) + + resource = job["migration_state"]["resources"]["source_blobs:selected_documents"] + assert resource["status"] == "failed" + assert resource["progress"]["missing_count"] == 1 + assert resource["progress"]["not_applicable_count"] == 1 + manifest_entries = [ + entry + for batch in job_container.manifest_batches + for entry in batch["entries"] + ] + assert {entry["status"] for entry in manifest_entries} == { + "missing", + "not_applicable", + } + assert all(entry["source_identity"].startswith("sha256:") for entry in manifest_entries) \ No newline at end of file diff --git a/functional_tests/test_data_management_cosmos_migration_resilience.py b/functional_tests/test_data_management_cosmos_migration_resilience.py new file mode 100644 index 000000000..d43d36d87 --- /dev/null +++ b/functional_tests/test_data_management_cosmos_migration_resilience.py @@ -0,0 +1,439 @@ +# test_data_management_cosmos_migration_resilience.py +""" +Functional test for resilient Data Management Cosmos migration. +Version: 0.250.071 +Implemented in: 0.250.075 + +This test ensures the in-app migration uses concurrent destination writes, +remote provenance skips, durable checkpoints, and request-unit throughput. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import threading +import time +import types + +from azure.core.exceptions import ResourceNotFoundError + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import initialize_migration_state +from functions_migration_provenance import create_migration_provenance_context + + +class FakeJobContainer: + """Persist a deep copy to emulate Cosmos serialization between checkpoints.""" + + def __init__(self): + self.manifest_batches = [] + + def upsert_item(self, body): + return copy.deepcopy(body) + + def create_item(self, body): + self.manifest_batches.append(copy.deepcopy(body)) + return copy.deepcopy(body) + + +class FakeSourceContainer: + """Return a stable source set below the migration cutoff.""" + + def __init__(self, documents): + self.documents = documents + self.queries = [] + + def query_items(self, **kwargs): + self.queries.append(copy.deepcopy(kwargs)) + parameters = { + parameter["name"]: parameter["value"] + for parameter in (kwargs.get("parameters") or []) + } + documents = copy.deepcopy(self.documents) + if "@source_start_epoch" in parameters: + documents = [ + document for document in documents + if document.get("_ts", 0) >= parameters["@source_start_epoch"] + ] + if "@source_cutoff_epoch" in parameters: + documents = [ + document for document in documents + if document.get("_ts", 0) <= parameters["@source_cutoff_epoch"] + ] + return iter(documents) + + +class FakeTargetContainer: + """Record concurrent writes and return a current-migration destination marker.""" + + def __init__(self, migration_id): + self.migration_id = migration_id + self.written = [] + self.updated = [] + self.active_writes = 0 + self.maximum_active_writes = 0 + self.lock = threading.Lock() + self.write_barrier = threading.Barrier(2) + self.documents = { + ("already-migrated", "already-migrated"): { + "id": "already-migrated", + "simplechatMigration": { + "migrationId": migration_id, + "migratedAtUtc": "2026-07-24T12:00:00+00:00", + "status": "succeeded", + "sourceHash": "", + "sourceVersion": "1", + }, + }, + } + + def read_item(self, item, partition_key, **_kwargs): + document = self.documents.get((item, partition_key)) + if document is None: + raise ResourceNotFoundError("not found") + return copy.deepcopy(document) + + def create_item(self, document, response_hook=None, retry_write=None, **_kwargs): + assert retry_write == 1 + key = (document["id"], document["id"]) + if key in self.documents: + conflict = RuntimeError("conflict") + conflict.status_code = 409 + raise conflict + with self.lock: + self.active_writes += 1 + self.maximum_active_writes = max(self.maximum_active_writes, self.active_writes) + try: + try: + self.write_barrier.wait(timeout=1.0) + except threading.BrokenBarrierError: + pass + time.sleep(0.03) + if response_hook: + response_hook({"x-ms-request-charge": "3.5"}, document) + self.written.append(copy.deepcopy(document)) + self.documents[key] = copy.deepcopy(document) + return document + finally: + with self.lock: + self.active_writes -= 1 + + def replace_item(self, item, body, etag=None, match_condition=None, response_hook=None, **_kwargs): + assert match_condition is not None + key = (body["id"], body["id"]) + assert self.documents.get(key) == item + if response_hook: + response_hook({"x-ms-request-charge": "4.0"}, body) + self.updated.append(copy.deepcopy(body)) + self.documents[key] = copy.deepcopy(body) + return body + + +class FakeTargetDatabase: + """Return the test target container for every requested migration resource.""" + + def __init__(self, target_container): + self.target_container = target_container + + def create_container_if_not_exists(self, **_kwargs): + return self.target_container + + +def load_data_management_module(monkeypatch, source_container, job_container): + """Load the production module with lightweight service dependencies.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + config_module.source_container = source_container + config_module.target_container_name = "target-container" + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_resilience_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_data_management_cosmos_migration_runs_concurrently_with_provenance(monkeypatch): + """Validate target markers skip and unmarked records checkpoint as parallel writes.""" + migration_id = "11111111-1111-1111-1111-111111111111" + source_container = FakeSourceContainer([ + {"id": "already-migrated", "_ts": 1}, + {"id": "copy-one", "_ts": 1}, + {"id": "copy-two", "_ts": 1}, + ]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_container, job_container) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "user_settings", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "id_field": "id", + }], + "groups": [], + "public_workspaces": [], + } + + settings = { + "migration_max_parallel_operations": 2, + "migration_retry_count": 2, + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": False}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + } + migration_state = initialize_migration_state( + None, + migration_id, + {"test": "cosmos"}, + ) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-24T12:00:00+00:00", + ) + job = {"id": migration_id, "migration_state": migration_state} + target_container = FakeTargetContainer(migration_id) + target_container.documents[("already-migrated", "already-migrated")][ + "simplechatMigration" + ]["sourceHash"] = module._build_cosmos_source_hash({"id": "already-migrated"}) + + artifacts = module._copy_cosmos_records_to_target( + FakeTargetDatabase(target_container), + "users", + migration_plan["users"], + job, + migration_state, + provenance_context, + settings, + ) + + assert len(target_container.written) == 2 + assert target_container.maximum_active_writes >= 2 + assert all( + document["simplechatMigration"]["migrationId"] == migration_id + and document["simplechatMigration"]["status"] == "succeeded" + for document in target_container.written + ) + assert artifacts[0]["copied_count"] == 2 + assert artifacts[0]["skipped_count"] == 0 + assert artifacts[0]["destination_provenance_skip_count"] == 1 + assert artifacts[0]["request_units"] == 7.0 + resource = job["migration_state"]["resources"]["cosmos:users:user_settings"] + assert resource["status"] == "completed" + assert resource["result"]["items_per_second"] > 0 + manifest_entries = [ + entry + for batch in job_container.manifest_batches + for entry in batch["entries"] + ] + assert len(manifest_entries) == 3 + assert all(entry["source_identity"].startswith("sha256:") for entry in manifest_entries) + + try: + module._classify_target_cosmos_document( + {"id": "unowned"}, + provenance_context, + "target-container", + ) + except module.DataManagementSettingsValidationError as exc: + assert "unowned" in str(exc) + else: + raise AssertionError("An unowned destination Cosmos record was accepted for overwrite.") + + +def test_data_management_cosmos_delta_upserts_only_changed_owned_records(monkeypatch): + """Use the baseline _ts watermark and update only migration-owned changed items.""" + migration_id = "22222222-2222-2222-2222-222222222222" + baseline_migration_id = "11111111-1111-1111-1111-111111111111" + baseline_epoch = 1_700_000_000 + source_documents = [ + {"id": "unchanged", "value": "same", "_ts": baseline_epoch}, + {"id": "changed", "value": "new", "_ts": baseline_epoch}, + {"id": "created", "value": "new", "_ts": baseline_epoch + 2}, + ] + source_container = FakeSourceContainer(source_documents) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_container, job_container) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "user_settings", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "id_field": "id", + }], + "groups": [], + "public_workspaces": [], + } + target_container = FakeTargetContainer(migration_id) + target_container.documents.update({ + ("unchanged", "unchanged"): { + "id": "unchanged", + "value": "same", + "_etag": "unchanged-etag", + "simplechatMigration": { + "migrationId": baseline_migration_id, + "migratedAtUtc": "2026-07-28T12:00:00+00:00", + "status": "succeeded", + "sourceHash": module._build_cosmos_source_hash(source_documents[0]), + "sourceVersion": str(baseline_epoch), + }, + }, + ("changed", "changed"): { + "id": "changed", + "value": "old", + "_etag": "changed-etag", + "simplechatMigration": { + "migrationId": baseline_migration_id, + "migratedAtUtc": "2026-07-28T12:00:00+00:00", + "status": "succeeded", + "sourceHash": module._build_cosmos_source_hash({ + "id": "changed", + "value": "old", + }), + "sourceVersion": str(baseline_epoch), + }, + }, + }) + migration_state = initialize_migration_state(None, migration_id, {"test": "cosmos-delta"}) + provenance_context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + provenance_context.update({ + "migration_mode": "delta_upsert", + "baseline_source_cutoff_at": "2023-11-14T22:13:20+00:00", + }) + job = {"id": migration_id, "migration_state": migration_state} + + artifacts = module._copy_cosmos_records_to_target( + FakeTargetDatabase(target_container), + "users", + {"mode": "all", "ids": [], "include_documents": False}, + job, + migration_state, + provenance_context, + { + "migration_max_parallel_operations": 2, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + ) + + assert artifacts[0]["created_count"] == 1 + assert artifacts[0]["updated_count"] == 1 + assert artifacts[0]["unchanged_count"] == 1 + assert artifacts[0]["source_read_count"] == 3 + assert len(target_container.written) == 1 + assert len(target_container.updated) == 1 + assert target_container.updated[0]["id"] == "changed" + assert target_container.updated[0]["simplechatMigration"]["sourceVersion"] == str( + baseline_epoch + ) + assert target_container.updated[0]["simplechatMigration"]["sourceHash"] == ( + module._build_cosmos_source_hash(source_documents[1]) + ) + assert any( + "c._ts >= @source_start_epoch" in query["query"] + for query in source_container.queries + ) + + +def test_cosmos_write_renews_heartbeat_while_request_is_in_flight(monkeypatch): + """Renew the migration lease before a slow Cosmos write completes.""" + migration_id = "33333333-3333-3333-3333-333333333333" + source_container = FakeSourceContainer([{"id": "slow", "_ts": 1}]) + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, source_container, job_container) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "user_settings", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "id_field": "id", + }], + "groups": [], + "public_workspaces": [], + } + release_write = threading.Event() + write_started = threading.Event() + write_completed = threading.Event() + + class SlowTargetContainer(FakeTargetContainer): + def create_item(self, document, response_hook=None, retry_write=None, **kwargs): + write_started.set() + assert release_write.wait(timeout=5.0) + result = super().create_item( + document, + response_hook=response_hook, + retry_write=retry_write, + **kwargs, + ) + write_completed.set() + return result + + target_container = SlowTargetContainer(migration_id) + target_container.write_barrier = threading.Barrier(1) + heartbeats = [] + original_heartbeat = module._persist_migration_heartbeat + + def capture_heartbeat(*args, **kwargs): + assert write_started.is_set() + assert not write_completed.is_set() + heartbeats.append(time.monotonic()) + release_write.set() + return original_heartbeat(*args, **kwargs) + + monkeypatch.setattr(module, "_persist_migration_heartbeat", capture_heartbeat) + migration_state = initialize_migration_state(None, migration_id, {"test": "slow-cosmos"}) + + module._copy_cosmos_records_to_target( + FakeTargetDatabase(target_container), + "users", + {"mode": "all", "ids": [], "include_documents": False}, + {"id": migration_id, "migration_state": migration_state}, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + { + "migration_max_parallel_operations": 1, + "migration_retry_count": 1, + "data_management_job_lease_seconds": 900, + }, + ) + + assert heartbeats + assert release_write.is_set() + assert write_completed.is_set() \ No newline at end of file diff --git a/functional_tests/test_data_management_destination_cosmos_capacity.py b/functional_tests/test_data_management_destination_cosmos_capacity.py new file mode 100644 index 000000000..5d1851214 --- /dev/null +++ b/functional_tests/test_data_management_destination_cosmos_capacity.py @@ -0,0 +1,316 @@ +# test_data_management_destination_cosmos_capacity.py +""" +Functional test for Data Management destination Cosmos migration controls. +Version: 0.250.071 +Implemented in: 0.250.075 + +This test ensures preflight proves destination create/read/delete access and +an opt-in 10,000 RU migration boost restores the original capacity afterward. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import initialize_migration_state + + +class FakeJobContainer: + """Persist deep-copy job state for checkpoint assertions.""" + + def upsert_item(self, body): + return copy.deepcopy(body) + + +class FakeSourceContainer: + """Return one source item to prove source data-plane access.""" + + def query_items(self, **_kwargs): + return iter([{"id": "source-item"}]) + + +class FakeTargetContainer: + """Track the temporary preflight item lifecycle.""" + + def __init__(self): + self.created = [] + self.read = [] + self.deleted = [] + + def create_item(self, body): + self.created.append(copy.deepcopy(body)) + return body + + def read_item(self, item, partition_key): + self.read.append((item, partition_key)) + return {"id": item} + + def delete_item(self, item, partition_key): + self.deleted.append((item, partition_key)) + + +class FakeTargetDatabase: + """Return a target container for each requested destination resource.""" + + def __init__(self, target_container): + self.target_container = target_container + + def create_container_if_not_exists(self, **_kwargs): + return self.target_container + + +def load_data_management_module(monkeypatch, source_container, job_container): + """Load Data Management with only the fake dependencies this test needs.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + config_module.source_container = source_container + config_module.target_container_name = "target-container" + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_destination_capacity_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_destination_cosmos_preflight_and_temporary_capacity_restore(monkeypatch): + """Validate destination access probe and safe 10,000 RU boost restoration.""" + migration_id = "11111111-1111-1111-1111-111111111111" + job_container = FakeJobContainer() + target_container = FakeTargetContainer() + module = load_data_management_module(monkeypatch, FakeSourceContainer(), job_container) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "user_settings", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "id_field": "id", + }], + "groups": [], + "public_workspaces": [], + } + monkeypatch.setattr(module, "_get_target_cosmos_database", lambda _settings: FakeTargetDatabase(target_container)) + + migration_plan = { + "users": {"mode": "all", "ids": [], "include_documents": False}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + } + preflight = module._preflight_target_cosmos_migration_access({}, migration_plan) + assert preflight["container_count"] == 1 + assert len(target_container.created) == len(target_container.read) == len(target_container.deleted) == 1 + assert target_container.created[0]["type"] == "simplechat_migration_preflight" + + capacity_calls = [] + monkeypatch.setattr( + module, + "_inspect_target_cosmos_migration_capacity", + lambda _settings, _plan: { + "target_ru": 10000, + "management_settings": {"target": "destination"}, + "targets": [{ + "scope": "database", + "container_name": "", + "mode": "autoscale", + "current_ru": 4000, + }, { + "scope": "container", + "container_name": "target-container", + "mode": "autoscale", + "current_ru": 5000, + }], + }, + ) + monkeypatch.setattr( + module, + "get_database_throughput", + lambda _settings: {"current_ru": 10000, "is_scalable": True}, + ) + monkeypatch.setattr( + module, + "get_container_throughput", + lambda _settings, _container_name: {"current_ru": 10000, "is_scalable": True}, + ) + + def set_capacity(_settings, target_ru, **kwargs): + target_scope = kwargs["decision"]["scope"] + target_container_name = kwargs["decision"]["container_name"] + if target_ru == 10000: + snapshots = job["migration_state"]["capacity"]["targets"] + assert any( + snapshot["scope"] == target_scope and + snapshot["container_name"] == target_container_name and + snapshot["boost_attempted"] is True and + snapshot["original_ru"] in {4000, 5000} + for snapshot in snapshots + ) + capacity_calls.append((target_ru, kwargs["decision"]["scope"], kwargs["decision"]["target_mode"])) + return {"to_ru": target_ru} + + monkeypatch.setattr(module, "set_database_throughput", set_capacity) + settings = { + "migration_temporary_destination_ru_enabled": True, + "migration_temporary_destination_ru": 10000, + "data_management_job_lease_seconds": 900, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "capacity"}) + job = {"id": migration_id, "migration_state": migration_state} + + state_after_boost = module._apply_temporary_destination_capacity( + job, + migration_state, + settings, + migration_plan, + ) + warnings, restored_state = module._restore_temporary_destination_capacity( + job, + state_after_boost, + settings, + ) + + assert capacity_calls == [ + (10000, "database", "autoscale"), + (10000, "container", "autoscale"), + (5000, "container", "autoscale"), + (4000, "database", "autoscale"), + ] + assert not warnings + assert restored_state["capacity"]["status"] == "restored" + assert restored_state["capacity"]["targets"][0]["restore_status"] == "restored" + assert restored_state["capacity"]["targets"][1]["restore_status"] == "restored" + + +def test_failed_capacity_restore_remains_pending_for_later_recovery(monkeypatch): + """Validate a failed restore leaves its durable snapshot available for a retry.""" + migration_id = "55555555-5555-5555-5555-555555555555" + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, FakeSourceContainer(), job_container) + monkeypatch.setattr( + module, + "get_database_throughput", + lambda _settings: {"current_ru": 10000, "is_scalable": True}, + ) + monkeypatch.setattr( + module, + "set_database_throughput", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("restore failed")), + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "restore-pending"}) + migration_state["capacity"] = { + "status": "boosted", + "restore_pending": True, + "management_settings": {"target": "destination"}, + "targets": [{ + "scope": "database", + "container_name": "", + "mode": "autoscale", + "original_ru": 4000, + "target_ru": 10000, + "boosted_to_ru": 10000, + "boost_attempted": True, + "changed": True, + "restore_status": "pending", + }], + } + job = {"id": migration_id, "migration_state": migration_state} + + warnings, restored_state = module._restore_temporary_destination_capacity( + job, + migration_state, + {"data_management_job_lease_seconds": 900}, + ) + + assert warnings + assert restored_state["capacity"]["restore_pending"] is True + assert restored_state["capacity"]["status"] == "restore_pending" + assert restored_state["capacity"]["targets"][0]["restore_status"] == "restore_failed" + + +def test_retry_recovers_pending_capacity_without_reboost(monkeypatch): + """Restore a durable pending snapshot before any new capacity inspection or boost.""" + migration_id = "66666666-6666-6666-6666-666666666666" + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, FakeSourceContainer(), job_container) + monkeypatch.setattr( + module, + "get_database_throughput", + lambda _settings: {"current_ru": 10000, "is_scalable": True}, + ) + capacity_calls = [] + + def restore_capacity(_settings, target_ru, **_kwargs): + capacity_calls.append(target_ru) + return {"to_ru": target_ru} + + monkeypatch.setattr(module, "set_database_throughput", restore_capacity) + monkeypatch.setattr( + module, + "_inspect_target_cosmos_migration_capacity", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("A recovery retry must not inspect or reapply the capacity boost.") + ), + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "restore-retry"}) + migration_state["capacity"] = { + "status": "restore_pending", + "restore_pending": True, + "management_settings": {"target": "destination"}, + "targets": [{ + "scope": "database", + "container_name": "", + "mode": "autoscale", + "original_ru": 4000, + "target_ru": 10000, + "boosted_to_ru": 10000, + "boost_attempted": True, + "changed": True, + "restore_status": "restore_failed", + }], + } + job = {"id": migration_id, "migration_state": migration_state} + + recovered_state = module._apply_temporary_destination_capacity( + job, + migration_state, + { + "migration_temporary_destination_ru_enabled": True, + "data_management_job_lease_seconds": 900, + }, + {"users": {"mode": "all", "include_documents": False}}, + ) + + assert capacity_calls == [4000] + assert recovered_state["capacity"]["restore_pending"] is False + assert recovered_state["capacity"]["status"] == "restored" \ No newline at end of file diff --git a/functional_tests/test_data_management_incremental_migration_modes.py b/functional_tests/test_data_management_incremental_migration_modes.py new file mode 100644 index 000000000..c2653ae46 --- /dev/null +++ b/functional_tests/test_data_management_incremental_migration_modes.py @@ -0,0 +1,366 @@ +# test_data_management_incremental_migration_modes.py +""" +Functional test for explicit Data Management incremental migration modes. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures mode defaults, baseline lineage, and destructive mirror +confirmation are normalized into the immutable migration plan. +""" + +import importlib.util +from pathlib import Path +import sys +import types + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + + +class FakeJobContainer: + """Provide explicit reads and ordered completed migration candidates.""" + + def __init__(self, jobs=None): + self.jobs = {job["id"]: job for job in (jobs or [])} + + def read_item(self, item, partition_key): + assert item == partition_key + return self.jobs[item] + + def query_items(self, **_kwargs): + return iter(sorted( + self.jobs.values(), + key=lambda job: job.get("completed_at") or "", + reverse=True, + )) + + +def load_data_management_module(monkeypatch, job_container=None): + """Load the production module with lightweight configuration dependencies.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.077" + config_module.cosmos_data_management_jobs_container = job_container or FakeJobContainer() + config_module.cosmos_data_management_job_items_container = object() + config_module.cosmos_settings_container = object() + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + throughput_module.CosmosThroughputError = RuntimeError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_incremental_modes_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def base_plan(): + """Return the smallest valid all-user migration plan.""" + return { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + "include_source_blobs": True, + } + + +def test_incremental_migration_mode_contract(monkeypatch): + """Validate the safe default, baseline GUID, and mirror confirmation guard.""" + module = load_data_management_module(monkeypatch) + + normalized_default = module.normalize_data_management_migration_plan({ + "migration_plan": base_plan(), + }) + assert normalized_default["migration_mode"] == "new_only" + assert normalized_default["baseline_job_id"] == "" + assert normalized_default["mirror_deletions_confirmed"] is False + assert normalized_default["target_ai_search_writes_frozen"] is False + + frozen_plan = base_plan() + frozen_plan["target_ai_search_writes_frozen"] = True + normalized_frozen_plan = module.normalize_data_management_migration_plan({ + "migration_plan": frozen_plan, + }) + assert normalized_frozen_plan["target_ai_search_writes_frozen"] is True + + baseline_job_id = "11111111-1111-1111-1111-111111111111" + delta_plan = base_plan() + delta_plan.update({ + "migration_mode": "delta_upsert", + "baseline_job_id": baseline_job_id, + }) + normalized_delta = module.normalize_data_management_migration_plan({ + "migration_plan": delta_plan, + }) + assert normalized_delta["migration_mode"] == "delta_upsert" + assert normalized_delta["baseline_job_id"] == baseline_job_id + + mirror_plan = base_plan() + mirror_plan["migration_mode"] = "mirror_with_deletions" + with pytest.raises(module.DataManagementSettingsValidationError, match="exact confirmation phrase"): + module.normalize_data_management_migration_plan({"migration_plan": mirror_plan}) + + mirror_plan["mirror_confirmation"] = "MIRROR WITH DELETIONS" + normalized_mirror = module.normalize_data_management_migration_plan({ + "migration_plan": mirror_plan, + }) + assert normalized_mirror["mirror_deletions_confirmed"] is True + + invalid_plan = base_plan() + invalid_plan["migration_mode"] = "overwrite_everything" + with pytest.raises(module.DataManagementSettingsValidationError, match="Migration mode"): + module.normalize_data_management_migration_plan({"migration_plan": invalid_plan}) + + +def completed_baseline_job( + module, + job_id, + migration_plan, + endpoint="https://target.documents.azure.com", + readiness="ready", +): + """Build a completed migration record with a durable cutoff and configuration.""" + return { + "id": job_id, + "type": module.DATA_MANAGEMENT_JOB_TYPE, + "operation": module.DATA_MANAGEMENT_OPERATION_MIGRATION, + "status": module.DATA_MANAGEMENT_STATUS_COMPLETED, + "completed_at": "2026-07-28T12:30:00+00:00", + "result": {"dry_run": False}, + "migration_state": { + "status": "completed", + "source_cutoff_at": "2026-07-28T12:00:00+00:00", + "resources": { + "reconciliation:cutover": { + "status": "completed", + "result": {"readiness": readiness}, + }, + }, + "configuration": { + "migration_plan": migration_plan, + "target_cosmos_endpoint": endpoint, + "target_ai_search_endpoint": "https://target.search.windows.net", + "target_enhanced_citations_storage_endpoint": "https://target.blob.core.windows.net", + }, + }, + } + + +def test_incremental_baseline_resolution_and_resume(monkeypatch): + """Resolve the latest compatible cutoff and reuse it on a retried job.""" + baseline_job_id = "22222222-2222-2222-2222-222222222222" + initial_module = load_data_management_module(monkeypatch) + baseline_plan = initial_module.normalize_data_management_migration_plan({ + "migration_plan": base_plan(), + }) + baseline_job = completed_baseline_job(initial_module, baseline_job_id, baseline_plan) + job_container = FakeJobContainer([baseline_job]) + module = load_data_management_module(monkeypatch, job_container=job_container) + delta_plan = base_plan() + delta_plan["migration_mode"] = "delta_upsert" + normalized_delta = module.normalize_data_management_migration_plan({ + "migration_plan": delta_plan, + }) + settings = { + "target_cosmos_endpoint": "https://target.documents.azure.com", + "target_ai_search_endpoint": "https://target.search.windows.net", + "target_enhanced_citations_storage_blob_endpoint": "https://target.blob.core.windows.net", + } + current_job = {"id": "33333333-3333-3333-3333-333333333333"} + + resolved = module._resolve_data_management_migration_baseline( + current_job, + settings, + normalized_delta, + ) + + assert resolved["baseline_job_id"] == baseline_job_id + assert resolved["baseline_source_cutoff_at"] == "2026-07-28T12:00:00+00:00" + + current_job["migration_state"] = { + "configuration": {"migration_plan": resolved}, + } + job_container.jobs.clear() + resumed = module._resolve_data_management_migration_baseline( + current_job, + settings, + normalized_delta, + ) + assert resumed["baseline_job_id"] == baseline_job_id + assert resumed["baseline_source_cutoff_at"] == resolved["baseline_source_cutoff_at"] + + +def test_incremental_baseline_rejects_incompatible_destination(monkeypatch): + """Reject an explicit baseline that targeted another destination account.""" + baseline_job_id = "44444444-4444-4444-4444-444444444444" + initial_module = load_data_management_module(monkeypatch) + baseline_plan = initial_module.normalize_data_management_migration_plan({ + "migration_plan": base_plan(), + }) + baseline_job = completed_baseline_job( + initial_module, + baseline_job_id, + baseline_plan, + endpoint="https://other.documents.azure.com", + ) + module = load_data_management_module( + monkeypatch, + job_container=FakeJobContainer([baseline_job]), + ) + delta_plan = base_plan() + delta_plan.update({ + "migration_mode": "delta_upsert", + "baseline_job_id": baseline_job_id, + }) + normalized_delta = module.normalize_data_management_migration_plan({ + "migration_plan": delta_plan, + }) + + with pytest.raises(module.DataManagementSettingsValidationError, match="different destination"): + module._resolve_data_management_migration_baseline( + {"id": "55555555-5555-5555-5555-555555555555"}, + { + "target_cosmos_endpoint": "https://target.documents.azure.com", + "target_ai_search_endpoint": "https://target.search.windows.net", + "target_enhanced_citations_storage_blob_endpoint": "https://target.blob.core.windows.net", + }, + normalized_delta, + ) + + +def test_migration_preview_aggregates_services_without_deletions(monkeypatch): + """Build server-owned estimates while keeping every reconciliation read-only.""" + module = load_data_management_module(monkeypatch) + deletion_flags = [] + monkeypatch.setattr(module, "_get_existing_target_cosmos_database", lambda _settings: object()) + monkeypatch.setattr( + module, + "_reconcile_cosmos_migration", + lambda *_args, **kwargs: deletion_flags.append(kwargs.get("apply_deletions")) or { + "service": "cosmos", + "create_count": 2, + "update_count": 1, + "unchanged_count": 3, + "delete_candidate_count": 0, + "conflict_count": 1, + "missing_count": 0, + }, + ) + monkeypatch.setattr( + module, + "_reconcile_ai_search_migration", + lambda *_args, **kwargs: deletion_flags.append(kwargs.get("apply_deletions")) or { + "service": "ai_search", + "create_count": 4, + "update_count": 2, + "unchanged_count": 6, + "delete_candidate_count": 0, + "conflict_count": 0, + "missing_count": 0, + }, + ) + monkeypatch.setattr( + module, + "_reconcile_blob_migration", + lambda *_args, **kwargs: deletion_flags.append(kwargs.get("apply_deletions")) or { + "service": "source_blobs", + "create_count": 1, + "update_count": 1, + "unchanged_count": 1, + "delete_candidate_count": 0, + "not_applicable_count": 2, + "source_missing_count": 1, + "conflict_count": 0, + "missing_count": 0, + }, + ) + + preview = module.preview_data_management_migration_plan( + {}, + {"migration_plan": base_plan()}, + ) + + assert deletion_flags == [False, False, False] + assert preview["estimated_outcomes"] == { + "create_count": 7, + "update_count": 4, + "unchanged_count": 10, + "delete_count": 0, + "not_applicable_count": 2, + "missing_count": 1, + "conflict_count": 1, + "failed_count": 0, + } + assert len(preview["plan_fingerprint"]) == 64 + + +def test_incremental_scope_signature_uses_all_two_thousand_ids(monkeypatch): + """Keep every supported selected ID in baseline compatibility decisions.""" + module = load_data_management_module(monkeypatch) + selected_ids = [f"user-{index:04d}" for index in range(2000)] + migration_plan = module.normalize_data_management_migration_plan({ + "migration_plan": { + **base_plan(), + "users": { + "mode": "selected", + "ids": selected_ids, + "include_documents": True, + }, + }, + }) + + signature = module._migration_plan_scope_signature(migration_plan) + + assert len(signature["users"]["ids"]) == 2000 + assert signature["users"]["ids"][-1] == "user-1999" + + +def test_non_destructive_warning_baseline_can_advance_watermark(monkeypatch): + """Allow retained destination-only data to remain the newest delta baseline.""" + baseline_job_id = "77777777-7777-7777-7777-777777777777" + initial_module = load_data_management_module(monkeypatch) + baseline_plan = initial_module.normalize_data_management_migration_plan({ + "migration_plan": base_plan(), + }) + baseline_job = completed_baseline_job( + initial_module, + baseline_job_id, + baseline_plan, + readiness="ready_with_warnings", + ) + module = load_data_management_module( + monkeypatch, + job_container=FakeJobContainer([baseline_job]), + ) + delta_plan = base_plan() + delta_plan["migration_mode"] = "delta_upsert" + + resolved = module._resolve_data_management_migration_baseline( + {"id": "88888888-8888-8888-8888-888888888888"}, + { + "target_cosmos_endpoint": "https://target.documents.azure.com", + "target_ai_search_endpoint": "https://target.search.windows.net", + "target_enhanced_citations_storage_blob_endpoint": "https://target.blob.core.windows.net", + }, + module.normalize_data_management_migration_plan({"migration_plan": delta_plan}), + ) + + assert resolved["baseline_job_id"] == baseline_job_id \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_cancellation.py b/functional_tests/test_data_management_migration_cancellation.py new file mode 100644 index 000000000..ed63bd660 --- /dev/null +++ b/functional_tests/test_data_management_migration_cancellation.py @@ -0,0 +1,194 @@ +# test_data_management_migration_cancellation.py +""" +Functional test for Data Management migration cancellation. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures queued migrations cancel immediately, running migrations stop +at a durable checkpoint, and a cooperative cancellation stays canceled rather +than becoming a failed job. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import initialize_migration_state + + +class FakeJobContainer: + """Store one migration job through Cosmos-like read and upsert operations.""" + + def __init__(self, job): + self.job = copy.deepcopy(job) + + def read_item(self, item, partition_key): + assert item == partition_key == self.job["id"] + return copy.deepcopy(self.job) + + def upsert_item(self, body): + self.job = copy.deepcopy(body) + return copy.deepcopy(self.job) + + +def load_data_management_module(monkeypatch, job_container): + """Load production cancellation helpers with lightweight infrastructure fakes.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.076" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_cancellation_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def build_migration_job(job_id, status, lease_holder_id=None): + """Build a minimal durable migration job fixture.""" + migration_state = initialize_migration_state(None, job_id, {"test": "cancellation"}) + return { + "id": job_id, + "type": "data_management_job", + "operation": "migration", + "status": status, + "lease_holder_id": lease_holder_id, + "lease_expires_at": "2099-01-01T00:00:00+00:00" if lease_holder_id else None, + "migration_state": migration_state, + "progress": {}, + "result": {}, + } + + +def test_queued_migration_cancels_immediately(monkeypatch): + """A queued job should become terminally canceled without a worker claim.""" + job_id = "11111111-1111-1111-1111-111111111111" + job_container = FakeJobContainer(build_migration_job(job_id, "queued")) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + + canceled_job = module.request_data_management_migration_cancellation( + job_id, + requested_by="admin-id", + requested_by_email="admin@example.com", + reason="Cutover paused", + ) + + assert canceled_job["status"] == "canceled" + assert canceled_job["cancel_requested_at"] + assert canceled_job["migration_state"]["status"] == "canceled" + public_job = module.sanitize_data_management_job_for_admin(canceled_job) + assert public_job["can_cancel"] is False + assert public_job["can_retry"] is True + + +def test_running_migration_cancellation_fences_next_checkpoint(monkeypatch): + """A running worker must observe a durable request before new migration work.""" + job_id = "22222222-2222-2222-2222-222222222222" + job_container = FakeJobContainer(build_migration_job(job_id, "running", "worker-1")) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + + requested_job = module.request_data_management_migration_cancellation(job_id) + assert requested_job["status"] == "running" + assert requested_job["cancel_requested_at"] + + try: + module._assert_migration_job_lease(requested_job) + except module.DataManagementMigrationCanceledError as exc: + assert "cancellation" in str(exc).lower() + else: + raise AssertionError("A cancellation request did not fence the running worker.") + + +def test_cancellation_exception_finalizes_job_without_failure(monkeypatch): + """The processor should record canceled rather than failed after safe stop.""" + job_id = "33333333-3333-3333-3333-333333333333" + job = build_migration_job(job_id, "running", "worker-1") + job["cancel_requested_at"] = "2026-07-24T12:00:00+00:00" + job_container = FakeJobContainer(job) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "get_data_management_settings", lambda: {}) + monkeypatch.setattr(module, "_try_claim_data_management_job", lambda *_args, **_kwargs: job) + monkeypatch.setattr( + module, + "execute_migration_job", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + module.DataManagementMigrationCanceledError("requested") + ), + ) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "_release_migration_destination_lock", lambda *_args, **_kwargs: None) + + canceled_job = module.process_data_management_job(job_id) + + assert canceled_job["status"] == "canceled" + assert canceled_job["last_error"] is None + assert canceled_job["progress"]["current_step"] == "canceled" + assert canceled_job["migration_state"]["status"] == "canceled" + + +def test_cancellation_before_execution_prevents_preflight_and_copy(monkeypatch): + """A persisted cancellation request must stop the executor before destination work.""" + job_id = "44444444-4444-4444-4444-444444444444" + job = build_migration_job(job_id, "running", "worker-1") + job.update({ + "cancel_requested_at": "2026-07-24T12:00:00+00:00", + "migration_state": None, + "options": { + "migration_plan": { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": False}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + }, + }, + }) + job_container = FakeJobContainer(job) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + module, + "_run_data_management_migration_preflight", + lambda *_args, **_kwargs: (_ for _ in ()).throw( + AssertionError("Cancellation should stop before migration preflight.") + ), + ) + + try: + module.execute_migration_job(job, {"data_management_job_lease_seconds": 900}) + except module.DataManagementMigrationCanceledError: + pass + else: + raise AssertionError("A pre-requested cancellation did not stop migration execution.") + + assert job["migration_state"]["status"] == "canceled" \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_coordinator.py b/functional_tests/test_data_management_migration_coordinator.py new file mode 100644 index 000000000..d4aa6ed69 --- /dev/null +++ b/functional_tests/test_data_management_migration_coordinator.py @@ -0,0 +1,310 @@ +# test_data_management_migration_coordinator.py +""" +Functional test for Data Management migration destination coordination. +Version: 0.250.071 +Implemented in: 0.250.075 +Updated in: 0.250.071 + +This test ensures only one migration can hold a destination coordinator lease +at a time and that release permits a later migration to acquire it safely. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + + +class ConflictError(Exception): + """Minimal Cosmos conflict used by the lock fake.""" + + status_code = 409 + + +class FakeCoordinatorContainer: + """Store job and coordinator documents keyed by their ID.""" + + def __init__(self): + self.items = {} + self.etag_counter = 0 + + def _save(self, body): + self.etag_counter += 1 + saved = copy.deepcopy(body) + saved["_etag"] = f"etag-{self.etag_counter}" + self.items[saved["id"]] = saved + return copy.deepcopy(saved) + + def create_item(self, body): + if body["id"] in self.items: + raise ConflictError("already exists") + return self._save(body) + + def read_item(self, item, partition_key): + assert item == partition_key + return copy.deepcopy(self.items[item]) + + def replace_item(self, item, body, **_kwargs): + assert item in self.items + return self._save(body) + + def delete_item(self, item, partition_key, **_kwargs): + assert item == partition_key + self.items.pop(item, None) + + def upsert_item(self, body): + return self._save(body) + + +class FakeTargetCoordinatorDatabase: + """Expose one shared target jobs container to otherwise independent source deployments.""" + + def __init__(self, container): + self.container = container + + def create_container_if_not_exists(self, **_kwargs): + return self.container + + +def load_data_management_module(monkeypatch, job_container): + """Load the production module with the coordinator storage fake.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_coordinator_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_destination_coordinator_blocks_overlapping_migrations(monkeypatch): + """Validate migrations to the same destination cannot overlap.""" + container = FakeCoordinatorContainer() + module = load_data_management_module(monkeypatch, container) + settings = { + "target_cosmos_endpoint": "https://target.documents.azure.com", + "target_ai_search_endpoint": "https://target.search.windows.net", + "target_enhanced_citations_storage_blob_endpoint": "https://target.blob.core.windows.net", + "data_management_job_lease_seconds": 900, + } + migration_plan = { + "include_ai_search": True, + "include_source_blobs": True, + } + first_job = {"id": "11111111-1111-1111-1111-111111111111", "operation": "migration"} + second_job = {"id": "22222222-2222-2222-2222-222222222222", "operation": "migration"} + + first_lock = module._acquire_migration_destination_lock(first_job, settings, migration_plan) + assert first_lock["id"] in container.items + assert first_lock["lease_seconds"] >= ( + module.DATA_MANAGEMENT_MIGRATION_REMOTE_REQUEST_TIMEOUT_SECONDS + + module.DATA_MANAGEMENT_MIGRATION_LOCK_RECOVERY_GRACE_SECONDS + ) + assert first_lock["lease_seconds"] < module.DATA_MANAGEMENT_DEFAULT_STALE_SECONDS + + try: + module._acquire_migration_destination_lock(second_job, settings, migration_plan) + except module.DataManagementSettingsValidationError as exc: + assert "active" in str(exc).lower() + else: + raise AssertionError("A second migration acquired the same destination coordinator lease.") + + module._release_migration_destination_lock(first_job) + second_lock = module._acquire_migration_destination_lock(second_job, settings, migration_plan) + assert second_lock["id"] in container.items + assert second_lock["lock_token"] != first_lock["lock_token"] + + +def test_coordinator_blocks_partially_overlapping_destination_plans(monkeypatch): + """Block concurrent jobs even when optional Search or Blob selections differ.""" + container = FakeCoordinatorContainer() + module = load_data_management_module(monkeypatch, container) + settings = { + "target_cosmos_endpoint": "https://target.documents.azure.com", + "target_ai_search_endpoint": "https://target.search.windows.net", + "target_enhanced_citations_storage_blob_endpoint": "https://target.blob.core.windows.net", + "data_management_job_lease_seconds": 900, + } + cosmos_only_job = { + "id": "33333333-3333-3333-3333-333333333333", + "operation": "migration", + } + full_job = { + "id": "44444444-4444-4444-4444-444444444444", + "operation": "migration", + } + + first_lock = module._acquire_migration_destination_lock( + cosmos_only_job, + settings, + {"include_ai_search": False, "include_source_blobs": False}, + ) + + try: + module._acquire_migration_destination_lock( + full_job, + settings, + {"include_ai_search": True, "include_source_blobs": True}, + ) + except module.DataManagementSettingsValidationError as exc: + assert "active" in str(exc).lower() + else: + raise AssertionError("A partially overlapping migration bypassed the global coordinator.") + + assert first_lock["id"].endswith("_global") + + +def test_retry_claim_waits_for_retained_coordinator_lock_expiry(monkeypatch): + """Do not claim a retry while an uncertain prior request remains quarantined.""" + container = FakeCoordinatorContainer() + module = load_data_management_module(monkeypatch, container) + settings = {"data_management_job_lease_seconds": 900} + job = { + "id": "55555555-5555-5555-5555-555555555555", + "type": "data_management_job", + "operation": "migration", + "status": "queued", + } + module._acquire_migration_destination_lock(job, settings, {}) + container._save(job) + + assert module._try_claim_data_management_job(job["id"], settings=settings) is None + + lock_id = job["migration_coordinator_lock"]["id"] + container.items[lock_id]["expires_at"] = "2000-01-01T00:00:00+00:00" + claimed_job = module._try_claim_data_management_job(job["id"], settings=settings) + + assert claimed_job["status"] == "running" + assert claimed_job["migration_attempt_id"] + + +def test_retry_replaces_released_or_expired_coordinator_lock_before_first_lease_check(monkeypatch): + """A retry must discard stale lock metadata before it verifies its new worker lease.""" + container = FakeCoordinatorContainer() + module = load_data_management_module(monkeypatch, container) + settings = {"data_management_job_lease_seconds": 900} + + released_job = { + "id": "66666666-6666-6666-6666-666666666666", + "type": "data_management_job", + "operation": "migration", + "status": "queued", + } + module._acquire_migration_destination_lock(released_job, settings, {}) + module._release_migration_destination_lock(released_job) + container._save(released_job) + + claimed_released_job = module._try_claim_data_management_job( + released_job["id"], + settings=settings, + ) + assert "migration_coordinator_lock" not in claimed_released_job + module._assert_migration_job_lease(claimed_released_job) + replacement_lock = module._acquire_migration_destination_lock( + claimed_released_job, + settings, + {}, + ) + assert container.items[replacement_lock["id"]]["migration_job_id"] == claimed_released_job["id"] + + second_container = FakeCoordinatorContainer() + second_module = load_data_management_module(monkeypatch, second_container) + expired_job = { + "id": "77777777-7777-7777-7777-777777777777", + "type": "data_management_job", + "operation": "migration", + "status": "queued", + } + expired_lock = second_module._acquire_migration_destination_lock(expired_job, settings, {}) + second_container.items[expired_lock["id"]]["expires_at"] = "2000-01-01T00:00:00+00:00" + second_container._save(expired_job) + + claimed_expired_job = second_module._try_claim_data_management_job( + expired_job["id"], + settings=settings, + ) + assert "migration_coordinator_lock" not in claimed_expired_job + second_module._assert_migration_job_lease(claimed_expired_job) + replacement_lock = second_module._acquire_migration_destination_lock( + claimed_expired_job, + settings, + {}, + ) + assert second_container.items[replacement_lock["id"]]["migration_job_id"] == claimed_expired_job["id"] + + +def test_target_coordinator_blocks_independent_source_job_stores(monkeypatch): + """Two source deployments must share the destination lock before either writes Cosmos data.""" + first_source_jobs = FakeCoordinatorContainer() + second_source_jobs = FakeCoordinatorContainer() + target_jobs = FakeCoordinatorContainer() + first_module = load_data_management_module(monkeypatch, first_source_jobs) + second_module = load_data_management_module(monkeypatch, second_source_jobs) + target_database = FakeTargetCoordinatorDatabase(target_jobs) + settings = {"data_management_job_lease_seconds": 900} + first_job = { + "id": "88888888-8888-8888-8888-888888888888", + "type": "data_management_job", + "operation": "migration", + } + second_job = { + "id": "99999999-9999-9999-9999-999999999999", + "type": "data_management_job", + "operation": "migration", + } + + first_module._acquire_target_migration_coordinator( + first_job, + target_database, + settings, + ) + assert "target_migration_coordinator" in first_job + assert "target_migration_coordinator_container" not in first_job + + with pytest.raises(RuntimeError, match="Another SimpleChat source"): + second_module._acquire_target_migration_coordinator( + second_job, + target_database, + settings, + ) + + assert first_module._release_target_migration_coordinator(first_job) is True + second_module._acquire_target_migration_coordinator( + second_job, + target_database, + settings, + ) + assert second_job["target_migration_coordinator"]["migration_id"] == second_job["id"] \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_manifest.py b/functional_tests/test_data_management_migration_manifest.py new file mode 100644 index 000000000..09b19cab8 --- /dev/null +++ b/functional_tests/test_data_management_migration_manifest.py @@ -0,0 +1,174 @@ +# test_data_management_migration_manifest.py +""" +Functional test for durable Data Management migration manifests. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures per-item migration outcomes are written in bounded, sanitized +job-scoped batches and can be filtered for retry/export workflows. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + + +class FakeManifestContainer: + """Persist manifest batches and expose Cosmos-like filtered iteration.""" + + def __init__(self): + self.items = [] + + def create_item(self, body): + self.items.append(copy.deepcopy(body)) + return copy.deepcopy(body) + + def query_items(self, **kwargs): + parameters = { + parameter["name"]: parameter["value"] + for parameter in kwargs.get("parameters") or [] + } + return iter([ + copy.deepcopy(item) + for item in self.items + if item.get("job_id") == parameters.get("@job_id") + and item.get("type") == parameters.get("@type") + ]) + + def read_item(self, item, partition_key): + assert item == partition_key + return { + "id": item, + "type": "data_management_job", + "operation": "migration", + } + + +class FakeJobContainer: + """Provide the unrelated job storage imports required by the module.""" + + def upsert_item(self, body): + return copy.deepcopy(body) + + def read_item(self, item, partition_key): + assert item == partition_key + return { + "id": item, + "type": "data_management_job", + "operation": "migration", + } + + +def load_data_management_module(monkeypatch, manifest_container): + """Load the production helpers with lightweight storage dependencies.""" + job_container = FakeJobContainer() + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.077" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = manifest_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + throughput_module.CosmosThroughputError = RuntimeError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_manifest_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_migration_manifest_batches_sanitize_and_filter(monkeypatch): + """Write bounded item outcomes and filter retryable failures safely.""" + manifest_container = FakeManifestContainer() + module = load_data_management_module(monkeypatch, manifest_container) + module.DATA_MANAGEMENT_MIGRATION_MANIFEST_BATCH_SIZE = 2 + append_entry, flush_entries = module._create_migration_manifest_writer( + "11111111-1111-1111-1111-111111111111", + "source_blobs:selected_documents", + ) + + append_entry({ + "service": "source_blobs", + "source_identity": "sha256:source-one", + "destination_identity": "sha256:destination-one", + "status": "copied", + "bytes": 10, + "connection_string": "must-not-persist", + "_locator": { + "service": "source_blobs", + "resource_name": "source_blobs:selected_documents", + "container_name": "user-documents", + "blob_name": "source-one.txt", + }, + }) + append_entry({ + "service": "source_blobs", + "source_identity": "sha256:source-two", + "destination_identity": "sha256:destination-two", + "status": "missing", + "error": "missing source", + "_locator": { + "service": "source_blobs", + "resource_name": "source_blobs:selected_documents", + "container_name": "user-documents", + "blob_name": "source-two.txt", + }, + }) + append_entry({ + "service": "source_blobs", + "source_identity": "sha256:source-three", + "destination_identity": "sha256:destination-three", + "status": "failed", + "error": "write failed", + }) + flush_entries() + + assert [item["entry_count"] for item in manifest_container.items] == [2, 1] + assert all( + "connection_string" not in entry + for item in manifest_container.items + for entry in item["entries"] + ) + retryable = list(module.iter_data_management_migration_manifest_entries( + "11111111-1111-1111-1111-111111111111", + statuses={"missing", "failed"}, + )) + assert [entry["status"] for entry in retryable] == ["missing", "failed"] + exported = module.export_data_management_migration_manifest( + "11111111-1111-1111-1111-111111111111", + statuses={"missing", "failed"}, + ) + assert exported["entry_count"] == 2 + assert exported["statuses"] == ["failed", "missing"] + assert "must-not-persist" not in exported["content"] + assert "source-two.txt" not in exported["content"] + missing_entry = retryable[0] + assert missing_entry["item_ref"] + resolved = module.resolve_data_management_migration_manifest_item( + "11111111-1111-1111-1111-111111111111", + missing_entry["item_ref"], + ) + assert resolved["container_name"] == "user-documents" + assert resolved["blob_name"] == "source-two.txt" diff --git a/functional_tests/test_data_management_migration_orchestration.py b/functional_tests/test_data_management_migration_orchestration.py new file mode 100644 index 000000000..cf57eddca --- /dev/null +++ b/functional_tests/test_data_management_migration_orchestration.py @@ -0,0 +1,368 @@ +# test_data_management_migration_orchestration.py +""" +Functional test for Data Management migration orchestration. +Version: 0.250.071 +Implemented in: 0.250.075 +Updated in: 0.250.071 + +This test ensures a migration job uses one provenance ID through preflight, +capacity preparation, all copy surfaces, and destination capacity restoration. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import initialize_migration_state + + +class FakeJobContainer: + """Persist migration state as the production job container would.""" + + def __init__(self): + self.job = None + + def upsert_item(self, body): + self.job = copy.deepcopy(body) + return copy.deepcopy(body) + + def read_item(self, item, partition_key): + assert item == partition_key == self.job["id"] + return copy.deepcopy(self.job) + + +def load_data_management_module(monkeypatch, job_container): + """Load the production module with job-only fake dependencies.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_orchestration_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_data_management_migration_uses_one_provenance_context_across_surfaces(monkeypatch): + """Verify job orchestration hands the same durable context to each surface.""" + migration_id = "11111111-1111-1111-1111-111111111111" + module = load_data_management_module(monkeypatch, FakeJobContainer()) + call_order = [] + received_contexts = [] + progress_updates = [] + + def preflight(job, state, _settings, _plan): + call_order.append("preflight") + state["preflight"] = {"status": "completed"} + return state + + def preview(_settings, _options, **_kwargs): + call_order.append("preview") + return { + "estimated_outcomes": { + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "delete_count": 0, + "not_applicable_count": 0, + "missing_count": 0, + "conflict_count": 0, + "failed_count": 0, + }, + } + + def capacity(job, state, _settings, _plan): + call_order.append("capacity") + state["capacity"] = {"status": "boosted", "restore_pending": True} + return state + + def copy_cosmos(_database, target_type, _selection, job, state, context, _settings): + call_order.append(f"cosmos:{target_type}") + received_contexts.append(context["migration_id"]) + job["migration_state"] = state + return [{"name": "user_settings", "type": "cosmos_container", "item_count": 1}] + + def copy_search(_settings, _plan, job, state, context, target_search_write_fence=None): + call_order.append("search") + received_contexts.append(context["migration_id"]) + assert target_search_write_fence == ("target-search-gate", {"fence_token": "fence-token"}) + job["migration_state"] = state + return [{"name": "personal_ai_search", "type": "ai_search_documents", "item_count": 1}] + + def copy_blobs(_settings, _plan, job, state, context): + call_order.append("blobs") + received_contexts.append(context["migration_id"]) + job["migration_state"] = state + return [{"name": "source_blobs", "type": "source_blobs", "blob_count": 1, "bytes": 10}] + + def reconcile(_settings, _plan, job, state, context, _database, **_kwargs): + call_order.append("reconciliation") + received_contexts.append(context["migration_id"]) + job["migration_state"] = state + return { + "name": "migration_reconciliation", + "type": "migration_reconciliation", + "readiness": "ready", + } + + def restore(job, state, _settings, **_kwargs): + call_order.append("restore") + state["capacity"] = {"status": "restored", "restore_pending": False} + job["migration_state"] = state + return [], state + + monkeypatch.setattr(module, "_run_data_management_migration_preflight", preflight) + monkeypatch.setattr(module, "preview_data_management_migration_plan", preview) + monkeypatch.setattr(module, "_apply_temporary_destination_capacity", capacity) + monkeypatch.setattr(module, "_get_target_cosmos_database", lambda _settings: object()) + monkeypatch.setattr( + module, + "_acquire_target_migration_coordinator", + lambda job, _database, _settings: call_order.append("target-coordinator-acquire") or job.update({ + "target_migration_coordinator": {"lock_token": "target-token"}, + }), + ) + monkeypatch.setattr( + module, + "_release_target_migration_coordinator", + lambda _job: call_order.append("target-coordinator-release") or True, + ) + monkeypatch.setattr( + module, + "_get_target_data_management_search_write_gate_container", + lambda _database: "target-search-gate", + ) + monkeypatch.setattr( + module, + "acquire_data_management_search_write_fence", + lambda *_args, **_kwargs: call_order.append("search-fence-acquire") or {"fence_token": "fence-token"}, + ) + monkeypatch.setattr( + module, + "release_data_management_search_write_fence", + lambda *_args, **_kwargs: call_order.append("search-fence-release") or True, + ) + monkeypatch.setattr(module, "_copy_cosmos_records_to_target", copy_cosmos) + monkeypatch.setattr(module, "_copy_ai_search_to_target", copy_search) + monkeypatch.setattr(module, "_copy_source_blobs_to_target", copy_blobs) + monkeypatch.setattr(module, "_run_data_management_migration_reconciliation", reconcile) + monkeypatch.setattr(module, "_restore_temporary_destination_capacity", restore) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + def capture_progress(job, message, completed_steps, total_steps, current_step=None, **_kwargs): + progress_updates.append((completed_steps, total_steps, current_step, message)) + job["progress"] = { + "completed_steps": completed_steps, + "total_steps": total_steps, + "current_step": current_step, + "percent_complete": int((completed_steps / total_steps) * 100), + } + return job + + monkeypatch.setattr(module, "_set_job_progress", capture_progress) + + job = { + "id": migration_id, + "options": { + "migration_plan": { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + "target_ai_search_writes_frozen": True, + "include_source_blobs": True, + }, + }, + } + settings = { + "migration_max_parallel_operations": 8, + "migration_retry_count": 5, + "migration_skip_recent_within_hours": 0, + "data_management_job_lease_seconds": 900, + } + + result = module.execute_migration_job(job, settings) + + assert result["migration_id"] == migration_id + assert received_contexts == [migration_id, migration_id, migration_id, migration_id] + assert call_order == [ + "target-coordinator-acquire", + "preview", + "preflight", + "capacity", + "cosmos:users", + "search-fence-acquire", + "search", + "search-fence-release", + "blobs", + "reconciliation", + "restore", + ] + assert result["migration_state"]["status"] == "completed" + assert result["migration_state"]["capacity"]["status"] == "restored" + assert progress_updates[0][:3] == (0, 10, "plan") + assert progress_updates[-1][:3] == (10, 10, "complete") + assert (8, 10, "capacity_restore") in [update[:3] for update in progress_updates] + assert [update[0] for update in progress_updates] == sorted( + update[0] for update in progress_updates + ) + + +def test_lightweight_progress_retains_last_progress_and_capacity_restore(monkeypatch): + """Expose liveness/progress clocks and the active capacity restoration stage.""" + migration_id = "22222222-2222-2222-2222-222222222222" + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, job_container) + job_container.job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "running", + "updated_at": "2026-07-29T12:00:02+00:00", + "last_heartbeat_at": "2026-07-29T12:00:02+00:00", + "last_progress_at": "2026-07-29T12:00:01+00:00", + "progress": { + "total_steps": 10, + "completed_steps": 8, + "current_step": "capacity_restore", + "percent_complete": 80, + }, + "migration_state": initialize_migration_state( + None, + migration_id, + {"test": "progress"}, + ), + } + + progress = module.get_data_management_job_progress(migration_id) + + assert progress["last_progress_at"] == "2026-07-29T12:00:01+00:00" + assert progress["progress"]["current_step"] == "capacity_restore" + assert progress["progress"]["completed_steps"] == 8 + assert progress["progress"]["total_steps"] == 10 + + +def test_process_migration_preserves_ten_stage_completion(monkeypatch): + """Keep 10/10 migration progress after the generic job wrapper completes.""" + migration_id = "33333333-3333-3333-3333-333333333333" + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, job_container) + claimed_job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "running", + "warnings": [], + "progress": { + "total_steps": 10, + "completed_steps": 8, + "current_step": "capacity_restore", + "percent_complete": 80, + }, + } + + def execute_migration(job, _settings): + job["progress"] = { + "total_steps": 10, + "completed_steps": 10, + "current_step": "complete", + "percent_complete": 100, + } + return {"warnings": [], "artifacts": []} + + monkeypatch.setattr(module, "get_data_management_settings", lambda: {}) + monkeypatch.setattr(module, "_try_claim_data_management_job", lambda *_args, **_kwargs: claimed_job) + monkeypatch.setattr(module, "execute_migration_job", execute_migration) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + released_target_coordinators = [] + monkeypatch.setattr( + module, + "_release_target_migration_coordinator", + lambda job: released_target_coordinators.append(job["id"]) or True, + ) + + completed = module.process_data_management_job(migration_id) + + assert completed["progress"] == { + "total_steps": 10, + "completed_steps": 10, + "current_step": "complete", + "percent_complete": 100, + } + assert released_target_coordinators == [migration_id] + + +def test_process_migration_retains_lock_after_uncertain_failure(monkeypatch): + """Keep the coordinator quarantine when a migration exits unexpectedly.""" + migration_id = "44444444-4444-4444-4444-444444444444" + job_container = FakeJobContainer() + module = load_data_management_module(monkeypatch, job_container) + claimed_job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "running", + "warnings": [], + "progress": {}, + "migration_coordinator_lock": { + "id": "data_management_migration_lock_global", + "lock_token": "lock-token", + }, + } + released_jobs = [] + + monkeypatch.setattr(module, "get_data_management_settings", lambda: {}) + monkeypatch.setattr(module, "_try_claim_data_management_job", lambda *_args, **_kwargs: claimed_job) + monkeypatch.setattr( + module, + "execute_migration_job", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("uncertain request")), + ) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + module, + "_release_migration_destination_lock", + lambda job: released_jobs.append(job["id"]), + ) + released_target_coordinators = [] + monkeypatch.setattr( + module, + "_release_target_migration_coordinator", + lambda job: released_target_coordinators.append(job["id"]) or True, + ) + + failed_job = module.process_data_management_job(migration_id) + + assert failed_job["status"] == "failed" + assert released_jobs == [] + assert released_target_coordinators == [] \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_preflight.py b/functional_tests/test_data_management_migration_preflight.py new file mode 100644 index 000000000..1f796b3a5 --- /dev/null +++ b/functional_tests/test_data_management_migration_preflight.py @@ -0,0 +1,260 @@ +# test_data_management_migration_preflight.py +""" +Functional test for Data Management migration destination preflight. +Version: 0.250.071 +Implemented in: 0.250.075 +Updated in: 0.250.071 + +This test ensures migration preflight proves target Search and Blob data-plane +write/read/delete access and rejects incompatible Cosmos partition keys. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + + +class FakeJobContainer: + """Provide the minimal persistence surface required by module import.""" + + def upsert_item(self, body): + return copy.deepcopy(body) + + def __init__(self): + self.items = {} + + def create_item(self, body): + self.items[body["id"]] = copy.deepcopy(body) + return copy.deepcopy(body) + + def read_item(self, item, partition_key): + assert item == partition_key + return copy.deepcopy(self.items[item]) + + def delete_item(self, item, partition_key): + assert item == partition_key + self.items.pop(item, None) + +class FakeSearchClient: + """Record Search preflight write/read/delete probes.""" + + def __init__(self, endpoint=""): + self.documents = {} + self.upload_count = 0 + self.delete_count = 0 + self._endpoint = endpoint + + def search(self, **kwargs): + filter_text = str(kwargs.get("filter") or "") + if "id eq" not in filter_text: + return iter([{"id": "source-document"}]) + matching = [ + {"id": document_id} + for document_id in self.documents + if document_id in filter_text + ] + return iter(matching) + + def upload_documents(self, documents, **_kwargs): + self.upload_count += 1 + for document in documents: + self.documents[document["id"]] = copy.deepcopy(document) + return [{"succeeded": True} for _ in documents] + + def delete_documents(self, documents, **_kwargs): + self.delete_count += 1 + for document in documents: + self.documents.pop(document["id"], None) + return [{"succeeded": True} for _ in documents] + + +class FakeContainerClient: + """Provide a lazy source Blob listing for access preflight.""" + + def list_blobs(self): + return iter([]) + + +class FakeBlobClient: + """Record target Blob preflight writes, reads, and cleanup.""" + + def __init__(self): + self.uploaded = False + self.read = False + self.deleted = False + + def upload_blob(self, **_kwargs): + self.uploaded = True + + def get_blob_properties(self): + self.read = True + return types.SimpleNamespace(metadata={}) + + def delete_blob(self): + self.deleted = True + + +class FakeBlobService: + """Expose source/target blob client APIs used by preflight.""" + + def __init__(self): + self.created_containers = [] + self.blob_clients = [] + + def get_container_client(self, _container_name): + return FakeContainerClient() + + def create_container(self, container_name): + self.created_containers.append(container_name) + + def get_blob_client(self, **_kwargs): + blob_client = FakeBlobClient() + self.blob_clients.append(blob_client) + return blob_client + + +class FakeTargetDatabase: + """Return a compatible target job container for Search write-gate preflight.""" + + def __init__(self): + self.containers = {} + + def create_container_if_not_exists(self, id=None, **_kwargs): + container = self.containers.setdefault(id, FakeJobContainer()) + return container + + +def load_data_management_module(monkeypatch, source_search_client, job_container): + """Load the production module with lightweight dependency fakes.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {"search_client_user": source_search_client} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + config_module.storage_account_user_documents_container_name = "user-documents" + config_module.storage_account_group_documents_container_name = "group-documents" + config_module.storage_account_public_documents_container_name = "public-documents" + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_preflight_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_search_and_blob_preflight_probes_destination_data_plane(monkeypatch): + """Validate Search and Blob preflight creates, reads, and removes probes.""" + source_search = FakeSearchClient() + target_search = FakeSearchClient() + source_blob = FakeBlobService() + target_blob = FakeBlobService() + module = load_data_management_module(monkeypatch, source_search, FakeJobContainer()) + monkeypatch.setattr(module, "_ensure_target_search_index", lambda *_args, **_kwargs: "ready") + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_search) + monkeypatch.setattr(module, "_get_target_cosmos_database", lambda _settings: FakeTargetDatabase()) + monkeypatch.setattr(module, "_get_source_blob_service_client", lambda: source_blob) + monkeypatch.setattr(module, "_get_target_enhanced_citations_blob_client", lambda _settings: target_blob) + + migration_plan = { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "target_ai_search_writes_frozen": True, + "include_source_blobs": True, + } + + search_result = module._preflight_target_ai_search_migration_access({}, migration_plan) + blob_result = module._preflight_target_blob_migration_access({}, migration_plan) + + assert search_result["index_count"] == 1 + assert search_result["target_search_write_gate_verified"] is True + assert target_search.upload_count == 1 + assert target_search.delete_count == 1 + assert target_search.documents == {} + assert blob_result["container_count"] == 1 + assert target_blob.created_containers == ["user-documents"] + assert len(target_blob.blob_clients) == 1 + assert target_blob.blob_clients[0].uploaded is True + assert target_blob.blob_clients[0].read is True + assert target_blob.blob_clients[0].deleted is True + + +def test_search_preflight_requires_frozen_distinct_target(monkeypatch): + """Reject Search migration writes without an explicit freeze or on the source service.""" + source_search = FakeSearchClient("https://source.search.windows.net") + module = load_data_management_module(monkeypatch, source_search, FakeJobContainer()) + migration_plan = { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + } + + try: + module._preflight_target_ai_search_migration_access( + {"target_ai_search_endpoint": "https://target.search.windows.net"}, + migration_plan, + ) + except module.DataManagementSettingsValidationError as exc: + assert "writers are frozen" in str(exc) + else: + raise AssertionError("Search migration proceeded without a target-write freeze confirmation.") + + migration_plan["target_ai_search_writes_frozen"] = True + try: + module._preflight_target_ai_search_migration_access( + {"target_ai_search_endpoint": "https://source.search.windows.net/"}, + migration_plan, + ) + except module.DataManagementSettingsValidationError as exc: + assert "must differ" in str(exc) + else: + raise AssertionError("Search migration accepted its source service as a destination.") + + +def test_partition_key_mismatch_is_rejected_before_data_copy(monkeypatch): + """Validate a pre-existing incompatible Cosmos container cannot pass preflight.""" + module = load_data_management_module(monkeypatch, FakeSearchClient(), FakeJobContainer()) + + class IncompatibleContainer: + def read(self): + return {"partitionKey": {"paths": ["/wrong"]}} + + try: + module._validate_target_cosmos_container_partition_key( + IncompatibleContainer(), + "documents", + "/id", + ) + except module.DataManagementSettingsValidationError as exc: + assert "partition key" in str(exc).lower() + else: + raise AssertionError("An incompatible target Cosmos partition key was accepted.") \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_provenance.py b/functional_tests/test_data_management_migration_provenance.py new file mode 100644 index 000000000..899f69da7 --- /dev/null +++ b/functional_tests/test_data_management_migration_provenance.py @@ -0,0 +1,114 @@ +# test_data_management_migration_provenance.py +""" +Functional test for Data Management migration provenance. +Version: 0.250.071 +Implemented in: 0.250.075 + +This test ensures the application migration contract uses one durable GUID, +preserves Blob metadata, and only bypasses successful current or recent items. +""" + +from datetime import datetime, timedelta, timezone +from pathlib import Path +import sys + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +sys.path.insert(0, str(APP_ROOT)) + +from functions_migration_provenance import ( + add_cosmos_migration_provenance, + add_search_migration_provenance, + create_migration_provenance_context, + get_blob_migration_provenance, + get_cosmos_migration_provenance, + get_search_migration_provenance, + is_successful_migration_record, + merge_blob_migration_metadata, + migration_record_matches_source, + should_skip_migration_record, +) + + +def test_data_management_migration_provenance_contract(): + """Validate provenance data and same-run or bounded recent skip behavior.""" + now = datetime(2026, 7, 24, 12, 0, tzinfo=timezone.utc) + context = create_migration_provenance_context( + migration_id="11111111-1111-1111-1111-111111111111", + migrated_at_utc=now, + skip_within_hours=24, + ) + + cosmos_document = add_cosmos_migration_provenance({"id": "document-1"}, context) + search_document = add_search_migration_provenance({"id": "document-1"}, context) + blob_metadata = merge_blob_migration_metadata({"existing": "preserved"}, context) + + cosmos_record = get_cosmos_migration_provenance(cosmos_document) + search_record = get_search_migration_provenance(search_document) + blob_record = get_blob_migration_provenance(blob_metadata) + + assert cosmos_record == search_record == blob_record + assert blob_metadata["existing"] == "preserved" + assert should_skip_migration_record(cosmos_record, context, current_time=now) + + versioned_cosmos = add_cosmos_migration_provenance( + {"id": "document-2"}, + context, + source_version="1700000000", + ) + hashed_search = add_search_migration_provenance( + {"id": "document-2"}, + context, + source_hash="sha256:abc123", + ) + versioned_blob = merge_blob_migration_metadata( + {}, + context, + source_hash="sha256:def456", + source_version="etag:123", + scope_hash="scope-hash", + ) + assert is_successful_migration_record(get_cosmos_migration_provenance(versioned_cosmos)) + assert migration_record_matches_source( + get_cosmos_migration_provenance(versioned_cosmos), + source_version="1700000000", + ) + assert migration_record_matches_source( + get_search_migration_provenance(hashed_search), + source_hash="sha256:abc123", + ) + assert migration_record_matches_source( + get_blob_migration_provenance(versioned_blob), + source_hash="sha256:def456", + ) + assert get_blob_migration_provenance(versioned_blob)["scopeHash"] == "scope-hash" + assert not migration_record_matches_source( + get_blob_migration_provenance(versioned_blob), + source_hash="sha256:changed", + ) + assert not migration_record_matches_source( + get_blob_migration_provenance(versioned_blob), + source_hash="sha256:def456", + source_version="etag:changed", + ) + + recent_record = { + "migrationId": "22222222-2222-2222-2222-222222222222", + "migratedAtUtc": (now - timedelta(hours=1)).isoformat(), + "status": "succeeded", + } + stale_record = { + "migrationId": "33333333-3333-3333-3333-333333333333", + "migratedAtUtc": (now - timedelta(hours=25)).isoformat(), + "status": "succeeded", + } + failed_record = { + "migrationId": context["migration_id"], + "migratedAtUtc": context["migrated_at_utc"], + "status": "failed", + } + + assert should_skip_migration_record(recent_record, context, current_time=now) + assert not should_skip_migration_record(stale_record, context, current_time=now) + assert not should_skip_migration_record(failed_record, context, current_time=now) \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_reconciliation.py b/functional_tests/test_data_management_migration_reconciliation.py new file mode 100644 index 000000000..75726c733 --- /dev/null +++ b/functional_tests/test_data_management_migration_reconciliation.py @@ -0,0 +1,1180 @@ +# test_data_management_migration_reconciliation.py +""" +Functional test for Data Management migration reconciliation. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures mirror reconciliation deletes only destination extras with +successful migration ownership and retains unowned or out-of-scope data. +""" + +import copy +import importlib.util +from pathlib import Path +import re +import sys +import time +import types + +import pytest +from azure.core.exceptions import ResourceNotFoundError + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_migration_provenance import create_migration_provenance_context +from functions_data_management_migration_state import initialize_migration_state + + +class FakeJobContainer: + """Provide the minimum persistence surface required by the module import.""" + + def __init__(self): + self.items = [] + self.job = None + + def upsert_item(self, body): + self.job = copy.deepcopy(body) + return copy.deepcopy(body) + + def read_item(self, item, partition_key): + assert item == partition_key + if self.job and self.job.get("id") == item: + return copy.deepcopy(self.job) + raise ResourceNotFoundError("not found") + + def create_item(self, body): + self.items.append(copy.deepcopy(body)) + return copy.deepcopy(body) + + def query_items(self, **kwargs): + parameters = { + parameter["name"]: parameter["value"] + for parameter in kwargs.get("parameters") or [] + } + return iter([ + copy.deepcopy(item) + for item in self.items + if item.get("job_id") == parameters.get("@job_id") + and item.get("type") == parameters.get("@type") + and item.get("plan_id") == parameters.get("@plan_id") + ]) + + +class FakeCosmosContainer: + """Expose source/target records and conditional deletion tracking.""" + + def __init__(self, documents): + self.documents = {document["id"]: copy.deepcopy(document) for document in documents} + self.deleted = [] + + def query_items(self, **kwargs): + selected_id = next( + ( + parameter["value"] + for parameter in (kwargs.get("parameters") or []) + if parameter["name"] == "@selected_id" + ), + None, + ) + documents = list(self.documents.values()) + blob_path = next( + ( + parameter["value"] + for parameter in (kwargs.get("parameters") or []) + if parameter["name"] == "@blob_path" + ), + None, + ) + if blob_path is not None: + documents = [ + document for document in documents + if blob_path in { + document.get("blob_path"), + document.get("archived_blob_path"), + } + ] + if selected_id is not None: + documents = [ + document for document in documents + if document.get("user_id") == selected_id + ] + return iter(copy.deepcopy(documents)) + + def read_item(self, item, partition_key): + if item != partition_key or item not in self.documents: + raise ResourceNotFoundError("not found") + return copy.deepcopy(self.documents[item]) + + def delete_item(self, item, partition_key, **kwargs): + assert item == partition_key + if self.documents[item].get("_etag"): + assert kwargs.get("etag") == self.documents[item]["_etag"] + assert kwargs.get("match_condition") is not None + self.deleted.append(item) + del self.documents[item] + + +class FakeTargetDatabase: + """Return document and dedicated target Data Management job containers.""" + + def __init__(self, container): + self.container = container + self.search_gate_container = FakeSearchGateContainer() + + def create_container_if_not_exists(self, id=None, **_kwargs): + if id == "data_management_jobs": + return self.search_gate_container + return self.container + + +class FakeSearchGateContainer: + """Persist a target Search write fence with Cosmos-like optimistic replacement.""" + + def __init__(self): + self.items = {} + self.etag_counter = 0 + + def _save(self, body): + self.etag_counter += 1 + saved = copy.deepcopy(body) + saved["_etag"] = f"gate-etag-{self.etag_counter}" + self.items[saved["id"]] = saved + return copy.deepcopy(saved) + + def read_item(self, item, partition_key): + assert item == partition_key + if item not in self.items: + raise ResourceNotFoundError("not found") + return copy.deepcopy(self.items[item]) + + def create_item(self, body): + if body["id"] in self.items: + conflict = RuntimeError("already exists") + conflict.status_code = 409 + raise conflict + return self._save(body) + + def replace_item(self, item, body, etag=None, **_kwargs): + current = self.items.get(item) + if current is None or (etag and current.get("_etag") != etag): + conflict = RuntimeError("stale gate") + conflict.status_code = 412 + raise conflict + return self._save(body) + + +class FakeSearchClient: + """Provide deterministic keyset pages and per-document deletes.""" + + def __init__(self, documents): + self.documents = {document["id"]: copy.deepcopy(document) for document in documents} + self.deleted = [] + + def search(self, **kwargs): + documents = sorted(self.documents.values(), key=lambda document: document["id"]) + cursor_match = re.search(r"id gt '([^']+)'", str(kwargs.get("filter") or "")) + if cursor_match: + documents = [ + document for document in documents + if document["id"] > cursor_match.group(1) + ] + return iter(copy.deepcopy(documents[:int(kwargs.get("top") or len(documents))])) + + def delete_documents(self, documents, **_kwargs): + for document in documents: + self.deleted.append(document["id"]) + self.documents.pop(document["id"], None) + return [{"succeeded": True} for _ in documents] + + +class GeneratedReconciliationSearchClient: + """Generate bounded ordered Search pages without materializing a large index.""" + + def __init__(self, document_count, decorate_document=None): + self.document_count = document_count + self.decorate_document = decorate_document + self.queries = [] + + def search(self, **kwargs): + assert "skip" not in kwargs + assert kwargs.get("order_by") == ["id asc"] + page_size = int(kwargs.get("top") or 0) + assert 0 < page_size <= 1000 + self.queries.append(copy.deepcopy(kwargs)) + cursor_match = re.search( + r"id gt 'document-(\d+)'", + str(kwargs.get("filter") or ""), + ) + start_index = int(cursor_match.group(1)) + 1 if cursor_match else 0 + end_index = min(self.document_count, start_index + page_size) + documents = [] + for index in range(start_index, end_index): + document = { + "id": f"document-{index:06d}", + "user_id": "user-1", + } + if self.decorate_document: + document = self.decorate_document(document) + documents.append(document) + return iter(documents) + + +class FakeTargetBlob: + """Record ETag-guarded blob deletion.""" + + def __init__(self, container, name): + self.container = container + self.name = name + + def get_blob_properties(self): + blob = self.container.blobs[self.name] + return types.SimpleNamespace( + metadata=copy.deepcopy(blob.get("metadata") or {}), + etag=blob.get("etag"), + size=blob.get("size", 0), + content_settings=blob.get("content_settings"), + ) + + def delete_blob(self, **kwargs): + blob = self.container.blobs[self.name] + assert kwargs.get("etag") == blob["etag"] + assert kwargs.get("match_condition") is not None + self.container.deleted.append(self.name) + + +class FakeTargetBlobContainer: + """List target Blob metadata and expose deletion clients.""" + + def __init__(self, blobs): + self.blobs = {blob["name"]: copy.deepcopy(blob) for blob in blobs} + self.deleted = [] + + def list_blobs(self, **kwargs): + assert "metadata" in kwargs.get("include", []) + return iter(copy.deepcopy(list(self.blobs.values()))) + + def get_blob_client(self, blob_name): + return FakeTargetBlob(self, blob_name) + + +class FakeTargetBlobService: + """Return one target Blob container.""" + + def __init__(self, container): + self.container = container + + def get_container_client(self, _container_name): + return self.container + + def get_blob_client(self, container, blob): + return self.container.get_blob_client(blob) + + +class FakeSourceBlob: + """Return stable source properties or a missing-source response.""" + + def __init__(self, properties=None): + self.properties = properties + + def get_blob_properties(self): + if self.properties is None: + raise ResourceNotFoundError("not found") + return copy.deepcopy(self.properties) + + +class FakeSourceBlobService: + """Expose only source blobs that currently exist.""" + + def __init__(self, blobs): + self.blobs = copy.deepcopy(blobs) + + def get_blob_client(self, container, blob): + return FakeSourceBlob(self.blobs.get((container, blob))) + + +def load_data_management_module(monkeypatch, source_cosmos, source_search): + """Load the production module with lightweight service dependencies.""" + job_container = FakeJobContainer() + config_module = types.ModuleType("config") + config_module.CLIENTS = {"search_client_user": source_search} + config_module.VERSION = "0.250.077" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + config_module.source_container = source_cosmos + config_module.cosmos_user_documents_container = source_cosmos + config_module.target_container_name = "target-container" + config_module.storage_account_user_documents_container_name = "user-documents" + config_module.storage_account_group_documents_container_name = "group-documents" + config_module.storage_account_public_documents_container_name = "public-documents" + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + throughput_module.CosmosThroughputError = RuntimeError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_reconciliation_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_mirror_reconciliation_collectors_are_read_only(monkeypatch): + """Collect owned deletion candidates without mutating any destination service.""" + migration_id = "11111111-1111-1111-1111-111111111111" + source_cosmos = FakeCosmosContainer([ + {"id": "keep", "user_id": "user-1", "blob_path": "keep.txt", "_ts": 1}, + ]) + source_search = FakeSearchClient([{"id": "keep", "user_id": "user-1"}]) + module = load_data_management_module(monkeypatch, source_cosmos, source_search) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "documents", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "filter_field": "user_id", + "documents": True, + }], + "groups": [], + "public_workspaces": [], + } + context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + owned_cosmos = lambda document_id: module.add_cosmos_migration_provenance( + {"id": document_id, "user_id": "user-1", "_etag": f"etag-{document_id}"}, + context, + ) + target_cosmos = FakeCosmosContainer([ + owned_cosmos("keep"), + owned_cosmos("delete-owned"), + {"id": "retain-unowned", "user_id": "user-1", "_etag": "etag-unowned"}, + ]) + target_search = FakeSearchClient([ + module.add_search_migration_provenance({"id": "keep", "user_id": "user-1"}, context), + module.add_search_migration_provenance({"id": "delete-owned", "user_id": "user-1"}, context), + {"id": "retain-unowned", "user_id": "user-1"}, + ]) + user_scope_hash = module._document_migration_scope_hash({"user_id": "user-1"}) + other_scope_hash = module._document_migration_scope_hash({"user_id": "user-2"}) + owned_blob_metadata = lambda scope_hash: module.merge_blob_migration_metadata( + {}, + context, + scope_hash=scope_hash, + ) + target_blob_container = FakeTargetBlobContainer([ + {"name": "keep.txt", "metadata": owned_blob_metadata(user_scope_hash), "etag": "etag-keep"}, + {"name": "delete-owned.txt", "metadata": owned_blob_metadata(user_scope_hash), "etag": "etag-delete"}, + {"name": "other-scope.txt", "metadata": owned_blob_metadata(other_scope_hash), "etag": "etag-other"}, + {"name": "retain-unowned.txt", "metadata": {}, "etag": "etag-unowned"}, + ]) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_search) + monkeypatch.setattr( + module, + "_get_target_enhanced_citations_blob_client", + lambda *_args, **_kwargs: FakeTargetBlobService(target_blob_container), + ) + source_blob_properties = types.SimpleNamespace( + metadata={}, + size=4, + content_settings=None, + etag="source-etag", + last_modified=None, + blob_tier=None, + blob_type="BlockBlob", + tags={}, + ) + monkeypatch.setattr( + module, + "_get_source_blob_service_client", + lambda: FakeSourceBlobService({ + ("user-documents", "keep.txt"): source_blob_properties, + }), + ) + migration_plan = { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + "include_source_blobs": True, + "migration_mode": "mirror_with_deletions", + "mirror_deletions_confirmed": True, + } + migration_state = {"source_cutoff_at": "2026-07-29T12:00:00+00:00"} + + cosmos_report = module._reconcile_cosmos_migration( + FakeTargetDatabase(target_cosmos), + migration_plan, + migration_state, + context, + ) + search_report = module._reconcile_ai_search_migration({}, migration_plan) + blob_report = module._reconcile_blob_migration({}, migration_plan) + + assert target_cosmos.deleted == [] + assert target_search.deleted == [] + assert target_blob_container.deleted == [] + for report in (cosmos_report, search_report, blob_report): + assert report["deleted_count"] == 0 + assert report["remaining_destination_only_owned_count"] == 1 + assert report["delete_candidate_count"] == 1 + assert "_deletion_candidates" not in report + assert cosmos_report["destination_only_unowned_count"] == 1 + assert search_report["destination_only_unowned_count"] == 1 + assert "other-scope.txt" not in target_blob_container.deleted + assert "retain-unowned.txt" not in target_blob_container.deleted + + +def test_two_phase_mirror_deletes_only_after_clean_preview_parity(monkeypatch): + """Apply owned deletes only after read-only readiness and preview counts agree.""" + migration_id = "22222222-2222-2222-2222-222222222222" + source_cosmos = FakeCosmosContainer([ + {"id": "keep", "user_id": "user-1", "blob_path": "keep.txt", "_ts": 1}, + ]) + source_search = FakeSearchClient([{"id": "keep", "user_id": "user-1"}]) + module = load_data_management_module(monkeypatch, source_cosmos, source_search) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "documents", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "filter_field": "user_id", + "documents": True, + }], + "groups": [], + "public_workspaces": [], + } + context = create_migration_provenance_context( + migration_id=migration_id, + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + source_keep = {"id": "keep", "user_id": "user-1", "blob_path": "keep.txt"} + target_cosmos = FakeCosmosContainer([ + module.add_cosmos_migration_provenance( + {"id": "keep", "user_id": "user-1", "_etag": "etag-keep"}, + context, + source_hash=module._build_cosmos_source_hash(source_keep), + source_version="1", + ), + module.add_cosmos_migration_provenance( + {"id": "delete-owned", "user_id": "user-1", "_etag": "etag-delete"}, + context, + ), + ]) + target_search = FakeSearchClient([ + module.add_search_migration_provenance( + {"id": "keep", "user_id": "user-1"}, + context, + source_hash=module._build_search_source_hash({"id": "keep", "user_id": "user-1"}), + ), + module.add_search_migration_provenance({"id": "delete-owned", "user_id": "user-1"}, context), + ]) + user_scope_hash = module._document_migration_scope_hash({"user_id": "user-1"}) + source_blob_properties = types.SimpleNamespace( + metadata={}, + size=4, + content_settings=None, + etag="source-etag", + last_modified=None, + blob_tier=None, + blob_type="BlockBlob", + tags={}, + ) + source_blob_hash, source_blob_version = module._build_blob_source_fingerprint( + source_blob_properties + ) + target_blob_container = FakeTargetBlobContainer([ + { + "name": "keep.txt", + "metadata": module.merge_blob_migration_metadata( + {}, + context, + source_hash=source_blob_hash, + source_version=source_blob_version, + scope_hash=user_scope_hash, + ), + "etag": "etag-keep", + "size": 4, + }, + { + "name": "delete-owned.txt", + "metadata": module.merge_blob_migration_metadata({}, context, scope_hash=user_scope_hash), + "etag": "etag-delete", + }, + ]) + source_blob_service = FakeSourceBlobService({ + ("user-documents", "keep.txt"): source_blob_properties, + }) + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_search) + monkeypatch.setattr( + module, + "_get_target_enhanced_citations_blob_client", + lambda *_args, **_kwargs: FakeTargetBlobService(target_blob_container), + ) + monkeypatch.setattr(module, "_get_source_blob_service_client", lambda: source_blob_service) + monkeypatch.setattr(module, "_assert_migration_job_lease", lambda *_args, **_kwargs: None) + migration_plan = { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + "include_source_blobs": True, + "migration_mode": "mirror_with_deletions", + "mirror_deletions_confirmed": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "mirror-two-phase"}) + job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "running", + "migration_state": migration_state, + "result": {}, + } + preview = { + "estimated_outcomes": { + "create_count": 0, + "update_count": 0, + "unchanged_count": 3, + "delete_count": 3, + "not_applicable_count": 0, + "missing_count": 0, + "conflict_count": 0, + "failed_count": 0, + }, + } + artifacts = [ + {"type": "cosmos_container", "created_count": 0, "updated_count": 0, "unchanged_count": 1}, + {"type": "ai_search_documents", "created_count": 0, "updated_count": 0, "unchanged_count": 1}, + {"type": "source_blobs", "created_count": 0, "updated_count": 0, "unchanged_count": 1, "missing_count": 0}, + ] + + result = module._run_data_management_migration_reconciliation( + {}, + migration_plan, + job, + migration_state, + context, + FakeTargetDatabase(target_cosmos), + preview_snapshot=preview, + migration_artifacts=artifacts, + ) + + assert result["deletion_status"] == "completed" + assert result["readiness"] == "ready" + assert result["deleted_count"] == 3 + assert target_cosmos.deleted == ["delete-owned"] + assert target_search.deleted == ["delete-owned"] + assert target_blob_container.deleted == ["delete-owned.txt"] + deleted_entries = list(module.iter_data_management_migration_manifest_entries( + migration_id, + statuses={"deleted"}, + )) + assert len(deleted_entries) == 3 + assert all(entry["item_ref"] for entry in deleted_entries) + + +def test_mirror_blocks_all_deletes_when_current_source_record_exists_after_cutoff(monkeypatch): + """Never infer deletion when a current source identity was updated after the cutoff.""" + migration_id = "33333333-3333-3333-3333-333333333333" + source_cosmos = FakeCosmosContainer([ + {"id": "updated-after-cutoff", "user_id": "user-1", "_ts": 200}, + ]) + source_search = FakeSearchClient([]) + module = load_data_management_module(monkeypatch, source_cosmos, source_search) + module.DATA_MANAGEMENT_MIGRATION_COSMOS_CONTAINERS = { + "users": [{ + "name": "documents", + "container_attr": "source_container", + "container_name_attr": "target_container_name", + "partition_key_path": "/id", + "filter_field": "user_id", + "documents": True, + }], + "groups": [], + "public_workspaces": [], + } + context = create_migration_provenance_context(migration_id=migration_id) + target_cosmos = FakeCosmosContainer([ + module.add_cosmos_migration_provenance( + {"id": "updated-after-cutoff", "user_id": "user-1", "_etag": "etag-updated"}, + context, + ), + module.add_cosmos_migration_provenance( + {"id": "other-owned-extra", "user_id": "user-1", "_etag": "etag-extra"}, + context, + ), + ]) + monkeypatch.setattr(module, "_assert_migration_job_lease", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "_reconcile_ai_search_migration", lambda *_args, **_kwargs: { + "service": "ai_search", "resources": [], "_deletion_candidates": [], + }) + monkeypatch.setattr(module, "_reconcile_blob_migration", lambda *_args, **_kwargs: { + "service": "source_blobs", "resources": [], "_deletion_candidates": [], + }) + migration_plan = { + "users": {"mode": "selected", "ids": ["user-1"], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": False, + "include_source_blobs": False, + "migration_mode": "mirror_with_deletions", + "mirror_deletions_confirmed": True, + } + migration_state = initialize_migration_state(None, migration_id, {"test": "post-cutoff"}) + migration_state["source_cutoff_at"] = "1970-01-01T00:01:40+00:00" + job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "running", + "migration_state": migration_state, + "result": {}, + } + preview = {"estimated_outcomes": { + "create_count": 0, "update_count": 0, "unchanged_count": 0, + "delete_count": 2, "not_applicable_count": 0, "missing_count": 0, + "conflict_count": 0, "failed_count": 0, + }} + + with pytest.raises(RuntimeError, match="not ready for cutover"): + module._run_data_management_migration_reconciliation( + {}, + migration_plan, + job, + migration_state, + context, + FakeTargetDatabase(target_cosmos), + preview_snapshot=preview, + migration_artifacts=[], + ) + + result = job["migration_state"]["resources"]["reconciliation:cutover"]["result"] + assert result["deletion_status"] == "blocked" + assert result["readiness"] == "not_ready" + assert target_cosmos.deleted == [] + detail = module.get_data_management_job_detail(migration_id) + reconciliation_artifacts = [ + artifact + for artifact in detail["job"]["result"]["artifacts"] + if artifact.get("type") == "migration_reconciliation" + ] + assert reconciliation_artifacts[0]["readiness"] == "not_ready" + assert reconciliation_artifacts[0]["deletion_status"] == "blocked" + + +def test_reconciliation_persists_preview_actual_divergence(monkeypatch): + """Record actual migration outcomes and their delta from the server preview.""" + source_cosmos = FakeCosmosContainer([]) + source_search = FakeSearchClient([]) + module = load_data_management_module(monkeypatch, source_cosmos, source_search) + empty_report = { + "matched_count": 0, + "missing_count": 0, + "destination_only_owned_count": 0, + "remaining_destination_only_owned_count": 0, + "destination_only_unowned_count": 0, + "conflict_count": 0, + "deleted_count": 1, + "unresolved_scope_count": 0, + "create_count": 0, + "update_count": 0, + "unchanged_count": 0, + "delete_candidate_count": 1, + "not_applicable_count": 0, + "source_missing_count": 0, + "resources": [], + } + monkeypatch.setattr( + module, + "_reconcile_cosmos_migration", + lambda *_args, **_kwargs: {"service": "cosmos", **empty_report}, + ) + monkeypatch.setattr( + module, + "_reconcile_ai_search_migration", + lambda *_args, **_kwargs: {"service": "ai_search", **empty_report, "deleted_count": 0}, + ) + monkeypatch.setattr( + module, + "_reconcile_blob_migration", + lambda *_args, **_kwargs: {"service": "source_blobs", **empty_report, "deleted_count": 0}, + ) + migration_id = "99999999-9999-9999-9999-999999999999" + migration_state = initialize_migration_state(None, migration_id, {"test": "divergence"}) + job = {"id": migration_id, "migration_state": migration_state} + preview = { + "estimated_outcomes": { + "create_count": 4, + "update_count": 2, + "unchanged_count": 2, + "delete_count": 1, + "not_applicable_count": 1, + "missing_count": 0, + "conflict_count": 0, + "failed_count": 0, + }, + } + artifacts = [ + { + "type": "cosmos_container", + "created_count": 3, + "updated_count": 3, + "unchanged_count": 1, + "failed_count": 0, + "collision_count": 0, + }, + { + "type": "source_blobs", + "created_count": 1, + "updated_count": 0, + "unchanged_count": 1, + "not_applicable_count": 2, + "missing_count": 1, + "failed_count": 0, + "collision_count": 0, + }, + ] + + result = module._run_data_management_migration_reconciliation( + {}, + {"migration_mode": "mirror_with_deletions"}, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + object(), + preview_snapshot=preview, + migration_artifacts=artifacts, + ) + + assert result["actual_outcomes"] == { + "create_count": 4, + "update_count": 3, + "unchanged_count": 2, + "delete_count": 0, + "not_applicable_count": 2, + "missing_count": 1, + "conflict_count": 0, + "failed_count": 0, + } + assert result["preview_actual_divergence"] == { + "create_count": 0, + "update_count": 1, + "unchanged_count": 0, + "delete_count": -1, + "not_applicable_count": 1, + "missing_count": 1, + "conflict_count": 0, + "failed_count": 0, + } + + +def test_ai_search_reconciliation_merges_more_than_one_hundred_thousand_keys(monkeypatch): + """Reconcile a large index with bounded keyset pages and no full identity maps.""" + document_count = 100_001 + source_cosmos = FakeCosmosContainer([]) + placeholder_source = GeneratedReconciliationSearchClient(0) + module = load_data_management_module(monkeypatch, source_cosmos, placeholder_source) + context = create_migration_provenance_context( + migration_id="44444444-4444-4444-4444-444444444444", + migrated_at_utc="2026-07-29T12:00:00+00:00", + ) + source_client = GeneratedReconciliationSearchClient(document_count) + + def decorate_target(document): + return module.add_search_migration_provenance( + document, + context, + source_hash=module._build_search_source_hash(document), + ) + + target_client = GeneratedReconciliationSearchClient( + document_count, + decorate_document=decorate_target, + ) + module.CLIENTS["search_client_user"] = source_client + monkeypatch.setattr(module, "_get_target_search_client", lambda *_args, **_kwargs: target_client) + + report = module._reconcile_ai_search_migration( + {}, + { + "users": {"mode": "all", "ids": [], "include_documents": True}, + "groups": {"mode": "none", "ids": [], "include_documents": False}, + "public_workspaces": {"mode": "none", "ids": [], "include_documents": False}, + "include_ai_search": True, + "migration_mode": "delta_upsert", + }, + ) + + assert report["matched_count"] == document_count + assert report["unchanged_count"] == document_count + assert report["missing_count"] == 0 + assert report["stale_count"] == 0 + assert len(source_client.queries) > 100 + assert len(target_client.queries) > 100 + assert all("skip" not in query for query in source_client.queries + target_client.queries) + + +def test_reconciliation_renews_heartbeat_during_slow_service_call(monkeypatch): + """Renew the lease by time even when fewer than threshold items are returned.""" + migration_id = "55555555-5555-5555-5555-555555555555" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + heartbeat_messages = [] + original_persist_state = module._persist_migration_state + + def capture_persist_state(job, state, settings, message, **kwargs): + heartbeat_messages.append(message) + return original_persist_state(job, state, settings, message, **kwargs) + + def slow_cosmos(*_args, **_kwargs): + time.sleep(2.2) + return {"service": "cosmos", "resources": []} + + monkeypatch.setattr(module, "_persist_migration_state", capture_persist_state) + monkeypatch.setattr(module, "_assert_migration_job_lease", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "_reconcile_cosmos_migration", slow_cosmos) + monkeypatch.setattr(module, "_reconcile_ai_search_migration", lambda *_args, **_kwargs: { + "service": "ai_search", "resources": [], + }) + monkeypatch.setattr(module, "_reconcile_blob_migration", lambda *_args, **_kwargs: { + "service": "source_blobs", "resources": [], + }) + migration_state = initialize_migration_state(None, migration_id, {"test": "slow-heartbeat"}) + job = {"id": migration_id, "migration_state": migration_state} + + module._run_data_management_migration_reconciliation( + {}, + {"migration_mode": "new_only"}, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + object(), + preview_snapshot={"estimated_outcomes": {}}, + migration_artifacts=[], + ) + + assert "Reconciliation worker heartbeat renewed" in heartbeat_messages + + +def test_partial_mirror_delete_flushes_prior_success_evidence(monkeypatch): + """Persist a successful delete before a later candidate failure is raised.""" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + flushed_entries = [] + pending_entries = [] + validation_count = 0 + + def fake_validate(candidate, **_kwargs): + nonlocal validation_count + validation_count += 1 + if validation_count == 1: + candidate["target_blob"] = types.SimpleNamespace( + delete_blob=lambda **_delete_kwargs: None + ) + return True, "" + return False, "second candidate changed" + + def append_manifest(entry): + pending_entries.append(copy.deepcopy(entry)) + + def flush_manifest(): + flushed_entries.extend(pending_entries) + pending_entries.clear() + + monkeypatch.setattr(module, "_validate_mirror_deletion_candidate", fake_validate) + candidates = [ + { + "service": "source_blobs", + "target_type": "users", + "container_name": "user-documents", + "blob_name": "first.txt", + "target_etag": "etag-first", + }, + { + "service": "source_blobs", + "target_type": "users", + "container_name": "user-documents", + "blob_name": "second.txt", + "target_etag": "etag-second", + }, + ] + + with pytest.raises(module.DataManagementSettingsValidationError, match="second candidate changed"): + module._apply_mirror_deletion_candidates( + iter(candidates), + manifest_append=append_manifest, + manifest_flush=flush_manifest, + ) + + assert [entry["status"] for entry in flushed_entries] == ["deleted"] + assert flushed_entries[0]["_locator"]["blob_name"] == "first.txt" + + +def test_partial_search_delete_flushes_success_before_batch_failure(monkeypatch): + """Persist successful Search deletions from a mixed-result batch before raising.""" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + flushed_entries = [] + pending_entries = [] + + class MixedDeleteClient: + def delete_documents(self, documents, **_kwargs): + assert len(documents) == 2 + return [{"succeeded": True}, {"succeeded": False}] + + def fake_validate(candidate, **_kwargs): + candidate["target_client"] = MixedDeleteClient() + candidate["index_name"] = "simplechat-user-index" + return True, "" + + def append_manifest(entry): + pending_entries.append(copy.deepcopy(entry)) + + def flush_manifest(): + flushed_entries.extend(pending_entries) + pending_entries.clear() + + monkeypatch.setattr(module, "_validate_mirror_deletion_candidate", fake_validate) + candidates = [ + { + "service": "ai_search", + "target_type": "users", + "index_name": "simplechat-user-index", + "document_id": "first", + }, + { + "service": "ai_search", + "target_type": "users", + "index_name": "simplechat-user-index", + "document_id": "second", + }, + ] + + with pytest.raises(RuntimeError, match="failed to delete 1 document"): + module._apply_mirror_deletion_candidates( + iter(candidates), + manifest_append=append_manifest, + manifest_flush=flush_manifest, + ) + + assert [entry["status"] for entry in flushed_entries] == ["deleted"] + assert flushed_entries[0]["_locator"]["document_id"] == "first" + + +def test_slow_mirror_search_validation_heartbeats_before_delete(monkeypatch): + """Keep the target Search fence renewable while a large delete batch is revalidated.""" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + heartbeat_calls = [] + release_validation = __import__("threading").Event() + + class DeleteClient: + def __init__(self): + self.deleted = [] + + def delete_documents(self, documents, **_kwargs): + assert heartbeat_calls + self.deleted.extend(document["id"] for document in documents) + return [{"succeeded": True} for _ in documents] + + delete_client = DeleteClient() + validation_count = 0 + + def fake_validate(candidate, **_kwargs): + nonlocal validation_count + validation_count += 1 + if validation_count == 2: + assert release_validation.wait(timeout=5.0) + candidate["target_client"] = delete_client + candidate["index_name"] = "simplechat-user-index" + return True, "" + + def heartbeat(_message, _completed_count=0): + heartbeat_calls.append(time.monotonic()) + release_validation.set() + + monkeypatch.setattr(module, "_validate_mirror_deletion_candidate", fake_validate) + module._apply_mirror_deletion_candidates( + iter([ + { + "service": "ai_search", + "target_type": "users", + "index_name": "simplechat-user-index", + "document_id": "first", + }, + { + "service": "ai_search", + "target_type": "users", + "index_name": "simplechat-user-index", + "document_id": "second", + }, + ]), + heartbeat_callback=heartbeat, + ) + + assert validation_count >= 4 + assert delete_client.deleted == ["first", "second"] + + +def test_reconciliation_exception_fails_resource_checkpoint(monkeypatch): + """Persist failed reconciliation state before propagating a scan exception.""" + migration_id = "66666666-6666-6666-6666-666666666666" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + monkeypatch.setattr(module, "_assert_migration_job_lease", lambda *_args, **_kwargs: None) + monkeypatch.setattr( + module, + "_reconcile_cosmos_migration", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("scan failed")), + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "scan-failure"}) + job = {"id": migration_id, "migration_state": migration_state} + + with pytest.raises(RuntimeError, match="scan failed"): + module._run_data_management_migration_reconciliation( + {}, + {"migration_mode": "new_only"}, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + object(), + preview_snapshot={"estimated_outcomes": {}}, + migration_artifacts=[], + ) + + resource = job["migration_state"]["resources"]["reconciliation:cutover"] + assert resource["status"] == "failed" + assert resource["last_error"] == "scan failed" + assert resource["result"]["readiness"] == "not_ready" + assert resource["result"]["services_completed"] == 0 + assert resource["result"]["error"] == "scan failed" + + +def test_later_service_failure_preserves_completed_reconciliation_report(monkeypatch): + """Retain completed Cosmos evidence when a later Search scan fails.""" + migration_id = "77777777-7777-7777-7777-777777777777" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + monkeypatch.setattr(module, "_assert_migration_job_lease", lambda *_args, **_kwargs: None) + monkeypatch.setattr(module, "_reconcile_cosmos_migration", lambda *_args, **_kwargs: { + "service": "cosmos", + "matched_count": 2, + "resources": [{"container_name": "documents", "matched_count": 2}], + }) + monkeypatch.setattr( + module, + "_reconcile_ai_search_migration", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("search scan failed")), + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "search-failure"}) + job = {"id": migration_id, "migration_state": migration_state} + + with pytest.raises(RuntimeError, match="search scan failed"): + module._run_data_management_migration_reconciliation( + {}, + {"migration_mode": "new_only"}, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + object(), + preview_snapshot={"estimated_outcomes": {}}, + migration_artifacts=[], + ) + + result = job["migration_state"]["resources"]["reconciliation:cutover"]["result"] + assert result["readiness"] == "not_ready" + assert result["services_completed"] == 1 + assert result["matched_count"] == 2 + assert result["services"][0]["service"] == "cosmos" + + +def test_deletion_plan_flush_failure_still_fails_resource_with_evidence(monkeypatch): + """Record reconciliation failure even when candidate-plan persistence fails.""" + migration_id = "88888888-8888-8888-8888-888888888888" + module = load_data_management_module( + monkeypatch, + FakeCosmosContainer([]), + FakeSearchClient([]), + ) + monkeypatch.setattr(module, "_assert_migration_job_lease", lambda *_args, **_kwargs: None) + + def cosmos_report(*_args, deletion_candidate_callback=None, **_kwargs): + deletion_candidate_callback({ + "service": "cosmos", + "target_type": "users", + "container_name": "documents", + "document_id": "candidate", + "partition_key": "candidate", + "target_etag": "etag-candidate", + }) + return { + "service": "cosmos", + "destination_only_owned_count": 1, + "delete_candidate_count": 1, + "resources": [], + } + + monkeypatch.setattr(module, "_reconcile_cosmos_migration", cosmos_report) + monkeypatch.setattr( + module, + "_write_mirror_deletion_candidate_batch", + lambda *_args, **_kwargs: (_ for _ in ()).throw(RuntimeError("plan write failed")), + ) + migration_state = initialize_migration_state(None, migration_id, {"test": "plan-failure"}) + job = {"id": migration_id, "migration_state": migration_state} + + with pytest.raises(RuntimeError, match="plan write failed"): + module._run_data_management_migration_reconciliation( + {}, + {"migration_mode": "mirror_with_deletions", "mirror_deletions_confirmed": True}, + job, + migration_state, + create_migration_provenance_context(migration_id=migration_id), + object(), + preview_snapshot={"estimated_outcomes": {}}, + migration_artifacts=[], + ) + + resource = job["migration_state"]["resources"]["reconciliation:cutover"] + assert resource["status"] == "failed" + assert resource["last_error"] == "plan write failed" + assert resource["result"]["readiness"] == "not_ready" + assert resource["result"]["services_completed"] == 1 + assert resource["result"]["persistence_warnings"] \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_recovery.py b/functional_tests/test_data_management_migration_recovery.py new file mode 100644 index 000000000..c90b7843a --- /dev/null +++ b/functional_tests/test_data_management_migration_recovery.py @@ -0,0 +1,169 @@ +# test_data_management_migration_recovery.py +""" +Functional test for Data Management migration recovery scheduling. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures delayed queued and stale migration jobs are resubmitted to +the executor, including when scheduled backup processing is disabled. +""" + +import copy +import importlib.util +from datetime import datetime, timezone +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +BACKGROUND_TASKS_PATH = APP_ROOT / "background_tasks.py" +APP_PATH = APP_ROOT / "app.py" +sys.path.insert(0, str(APP_ROOT)) + + +class FakeJobContainer: + """Store multiple job records through Cosmos-like query and write methods.""" + + def __init__(self, jobs): + self.jobs = {job["id"]: copy.deepcopy(job) for job in jobs} + + def query_items(self, **_kwargs): + return iter(copy.deepcopy(list(self.jobs.values()))) + + def read_item(self, item, partition_key): + assert item == partition_key + return copy.deepcopy(self.jobs[item]) + + def upsert_item(self, body): + self.jobs[body["id"]] = copy.deepcopy(body) + return copy.deepcopy(body) + + +class FakeExecutor: + """Record executor submissions without starting worker code.""" + + def __init__(self): + self.submissions = [] + + def submit_stored(self, name, function, **kwargs): + self.submissions.append((name, function.__name__, kwargs)) + + +class FakeApp: + """Expose the Flask executor extension used by recovery submission.""" + + def __init__(self, executor): + self.extensions = {"executor": executor} + + +def load_data_management_module(monkeypatch, job_container): + """Load production recovery helpers with in-memory Cosmos dependencies.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.076" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_recovery_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def migration_job(job_id, status, updated_at, heartbeat_at=None): + """Build a minimal durable migration job candidate.""" + return { + "id": job_id, + "type": "data_management_job", + "operation": "migration", + "status": status, + "created_at": updated_at, + "updated_at": updated_at, + "last_heartbeat_at": heartbeat_at, + "progress": {}, + } + + +def test_recovery_resubmits_delayed_queued_and_stale_migrations(monkeypatch): + """Validate automatic recovery is executor-backed and leaves checkpoints durable.""" + now = datetime(2026, 7, 24, 12, 0, tzinfo=timezone.utc) + queued_job = migration_job( + "11111111-1111-1111-1111-111111111111", + "queued", + "2026-07-24T11:58:00+00:00", + ) + stale_job = migration_job( + "22222222-2222-2222-2222-222222222222", + "running", + "2026-07-24T11:00:00+00:00", + heartbeat_at="2026-07-24T11:00:00+00:00", + ) + job_container = FakeJobContainer([queued_job, stale_job]) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + executor = FakeExecutor() + + recovery = module.recover_data_management_migration_jobs( + app=FakeApp(executor), + current_time=now, + ) + + assert {(item["job_id"], item["reason"]) for item in recovery} == { + (queued_job["id"], "queued_recovery"), + (stale_job["id"], "stale_recovery"), + } + assert len(executor.submissions) == 2 + assert all(item["submitted"] is True for item in recovery) + assert job_container.jobs[queued_job["id"]]["recovery_attempt_count"] == 1 + assert job_container.jobs[stale_job["id"]]["recovery_attempt_count"] == 1 + + +def test_recovery_runs_when_scheduled_backups_are_disabled(monkeypatch): + """Validate disabled backup schedules do not suppress migration recovery.""" + module = load_data_management_module(monkeypatch, FakeJobContainer([])) + expected_recovery = [{"job_id": "recovered-job", "reason": "stale_recovery", "submitted": True}] + monkeypatch.setattr(module, "get_data_management_settings", lambda: {"enabled": False}) + monkeypatch.setattr( + module, + "recover_data_management_migration_jobs", + lambda **_kwargs: expected_recovery, + ) + + result = module.check_due_data_management_jobs_once(app=FakeApp(FakeExecutor())) + + assert result == expected_recovery + + +def test_background_scheduler_receives_flask_app_for_executor_recovery(): + """Validate the scheduler loop receives the app rather than running migrations inline.""" + background_source = BACKGROUND_TASKS_PATH.read_text(encoding="utf-8") + app_source = APP_PATH.read_text(encoding="utf-8") + + assert "def run_data_management_scheduler_loop(app=None):" in background_source + assert "check_due_data_management_jobs_once(app=app)" in background_source + assert "lambda: run_data_management_scheduler_loop(app=app)" in background_source + assert "start_background_task_threads(app=app)" in app_source \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_retry.py b/functional_tests/test_data_management_migration_retry.py new file mode 100644 index 000000000..dd7773ad2 --- /dev/null +++ b/functional_tests/test_data_management_migration_retry.py @@ -0,0 +1,258 @@ +# test_data_management_migration_retry.py +""" +Functional test for Data Management migration retry. +Version: 0.250.071 +Implemented in: 0.250.075 + +This test ensures a failed migration retry keeps its original migration GUID +and completed resource checkpoints while returning the job to the queue. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import types + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management.py" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import ( + complete_migration_resource, + initialize_migration_state, + start_migration_resource, +) + + +class FakeJobContainer: + """Store one job using Cosmos-like read and upsert methods.""" + + def __init__(self, job): + self.job = copy.deepcopy(job) + + def read_item(self, item, partition_key): + assert item == partition_key == self.job["id"] + return copy.deepcopy(self.job) + + def upsert_item(self, body): + self.job = copy.deepcopy(body) + return copy.deepcopy(self.job) + + +class FakeConditionalJobContainer: + """Capture ETag-protected job replacement requests.""" + + def __init__(self): + self.replace_calls = [] + + def replace_item(self, **kwargs): + self.replace_calls.append(copy.deepcopy(kwargs)) + saved = copy.deepcopy(kwargs["body"]) + saved["_etag"] = "next-etag" + return saved + + def upsert_item(self, _body): + raise AssertionError("An ETag-protected migration job must use replace_item.") + + +def load_data_management_module(monkeypatch, job_container): + """Load the production helper with only its direct storage dependencies.""" + config_module = types.ModuleType("config") + config_module.CLIENTS = {} + config_module.VERSION = "0.250.075" + config_module.cosmos_data_management_jobs_container = job_container + config_module.cosmos_data_management_job_items_container = job_container + config_module.cosmos_settings_container = job_container + monkeypatch.setitem(sys.modules, "config", config_module) + + appinsights_module = types.ModuleType("functions_appinsights") + appinsights_module.log_event = lambda *_args, **_kwargs: None + monkeypatch.setitem(sys.modules, "functions_appinsights", appinsights_module) + + throughput_module = types.ModuleType("functions_cosmos_throughput") + + class FakeCosmosThroughputError(Exception): + pass + + throughput_module.CosmosThroughputError = FakeCosmosThroughputError + throughput_module.get_container_throughput = lambda *_args, **_kwargs: {} + throughput_module.get_database_throughput = lambda *_args, **_kwargs: {} + throughput_module.set_database_throughput = lambda *_args, **_kwargs: {} + monkeypatch.setitem(sys.modules, "functions_cosmos_throughput", throughput_module) + + module_name = "data_management_migration_retry_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return module + + +def test_failed_migration_retry_preserves_guid_and_completed_resources(monkeypatch): + """Validate restart state remains tied to the same durable migration identity.""" + migration_id = "11111111-1111-1111-1111-111111111111" + migration_state = initialize_migration_state(None, migration_id, {"plan": "stable"}) + start_migration_resource(migration_state, "cosmos:users:user_settings") + complete_migration_resource( + migration_state, + "cosmos:users:user_settings", + result={"copied_count": 3}, + ) + migration_state.update({"status": "failed", "last_error": "interrupted"}) + job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "failed", + "completed_at": "2026-07-24T12:00:00+00:00", + "last_error": "interrupted", + "migration_state": migration_state, + "result": {"artifacts": [{"name": "partial"}]}, + } + job_container = FakeJobContainer(job) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + + retried = module.retry_data_management_migration_job(migration_id) + + assert retried["status"] == "queued" + assert retried["last_error"] is None + assert retried["completed_at"] is None + assert retried["result"] == {} + assert retried["migration_state"]["migration_id"] == migration_id + assert retried["migration_state"]["status"] == "queued" + assert retried["migration_state"]["resources"]["cosmos:users:user_settings"]["status"] == "completed" + + +def test_stale_running_migration_is_eligible_for_safe_resume(monkeypatch): + """Validate only a stale running migration is exposed as resumable.""" + migration_id = "22222222-2222-2222-2222-222222222222" + migration_state = initialize_migration_state(None, migration_id, {"plan": "stable"}) + stale_job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "running", + "updated_at": "2026-07-24T00:00:00+00:00", + "last_heartbeat_at": "2026-07-24T00:00:00+00:00", + "migration_state": migration_state, + "result": {}, + } + job_container = FakeJobContainer(stale_job) + module = load_data_management_module(monkeypatch, job_container) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + + public_job = module.sanitize_data_management_job_for_admin(stale_job) + resumed = module.retry_data_management_migration_job(migration_id) + + assert public_job["can_retry"] is True + assert resumed["status"] == "queued" + assert resumed["migration_state"]["migration_id"] == migration_id + + +def test_canceled_migration_retry_clears_active_cancellation(monkeypatch): + """Preserve cancellation audit history without immediately canceling the retry.""" + migration_id = "55555555-5555-5555-5555-555555555555" + migration_state = initialize_migration_state(None, migration_id, {"plan": "stable"}) + migration_state.update({ + "status": "canceled", + "cancel_requested_at": "2026-07-29T12:00:00+00:00", + "cancel_requested_by": "admin-1", + "cancel_reason": "maintenance", + }) + canceled_job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "canceled", + "cancel_requested_at": "2026-07-29T12:00:00+00:00", + "cancel_requested_by": "admin-1", + "cancel_requested_by_email": "admin@example.com", + "cancel_reason": "maintenance", + "migration_state": migration_state, + "result": {}, + } + module = load_data_management_module(monkeypatch, FakeJobContainer(canceled_job)) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + + retried = module.retry_data_management_migration_job(migration_id) + + assert retried["status"] == "queued" + assert retried["cancel_requested_at"] is None + assert retried["migration_state"]["cancel_requested_at"] is None + assert retried["last_cancellation"]["reason"] == "maintenance" + assert retried["migration_state"]["last_cancellation"]["reason"] == "maintenance" + + +def test_prestart_canceled_migration_can_be_requeued(monkeypatch): + """Restart a queued migration canceled before durable state initialization.""" + migration_id = "66666666-6666-6666-6666-666666666666" + canceled_job = { + "id": migration_id, + "type": "data_management_job", + "operation": "migration", + "status": "canceled", + "cancel_requested_at": "2026-07-29T12:00:00+00:00", + "migration_state": None, + "options": {"migration_plan": {"users": {"mode": "all"}}}, + "result": {}, + } + module = load_data_management_module(monkeypatch, FakeJobContainer(canceled_job)) + monkeypatch.setattr(module, "_record_data_management_job_event", lambda *_args, **_kwargs: None) + + retried = module.retry_data_management_migration_job(migration_id) + + assert retried["status"] == "queued" + assert retried["migration_state"] is None + assert retried["cancel_requested_at"] is None + + +def test_migration_job_save_uses_etag_replacement_when_available(monkeypatch): + """Validate stale workers cannot silently overwrite a newer job checkpoint.""" + job_container = FakeConditionalJobContainer() + module = load_data_management_module(monkeypatch, job_container) + job = { + "id": "33333333-3333-3333-3333-333333333333", + "_etag": "original-etag", + "status": "running", + } + + saved = module._save_data_management_job(job) + + assert len(job_container.replace_calls) == 1 + replace_call = job_container.replace_calls[0] + assert replace_call["item"] == saved["id"] + assert replace_call["etag"] == "original-etag" + assert saved["_etag"] == "next-etag" + + +def test_superseded_migration_worker_lease_is_rejected(monkeypatch): + """Validate a stale worker cannot continue after another lease holder takes over.""" + migration_id = "44444444-4444-4444-4444-444444444444" + persisted_job = { + "id": migration_id, + "operation": "migration", + "status": "running", + "lease_holder_id": "new-worker", + "lease_expires_at": "2099-01-01T00:00:00+00:00", + "_etag": "new-etag", + } + module = load_data_management_module(monkeypatch, FakeJobContainer(persisted_job)) + stale_job = { + "id": migration_id, + "operation": "migration", + "lease_holder_id": "old-worker", + "_etag": "old-etag", + } + + try: + module._assert_migration_job_lease(stale_job) + except module.DataManagementMigrationLeaseLostError as exc: + assert "superseded" in str(exc).lower() + else: + raise AssertionError("A superseded migration worker lease was accepted.") \ No newline at end of file diff --git a/functional_tests/test_data_management_migration_state.py b/functional_tests/test_data_management_migration_state.py new file mode 100644 index 000000000..85690c306 --- /dev/null +++ b/functional_tests/test_data_management_migration_state.py @@ -0,0 +1,92 @@ +# test_data_management_migration_state.py +""" +Functional test for Data Management migration checkpoint state. +Version: 0.250.071 +Implemented in: 0.250.075 + +This test ensures in-app migration checkpoints retain one migration ID, +reject changed scopes, and expose durable transfer throughput metrics. +""" + +from datetime import datetime, timedelta, timezone +from pathlib import Path +import sys + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +sys.path.insert(0, str(APP_ROOT)) + +from functions_data_management_migration_state import ( + build_transfer_metrics, + complete_migration_resource, + initialize_migration_state, + is_migration_resource_completed, + start_migration_resource, + update_migration_resource, +) + + +def test_data_management_migration_checkpoint_and_metrics(): + """Verify resumable resource state and calculated rate fields.""" + started_at = datetime(2026, 7, 24, 12, 0, tzinfo=timezone.utc) + configuration = { + "migration_plan": {"users": {"mode": "selected", "ids": ["user-1"]}}, + "parallel_operations": 8, + } + state = initialize_migration_state( + None, + "11111111-1111-1111-1111-111111111111", + configuration, + current_time=started_at, + ) + start_migration_resource(state, "cosmos:users:user_settings", current_time=started_at) + metrics = build_transfer_metrics( + state["resources"]["cosmos:users:user_settings"]["started_at"], + copied_count=8, + skipped_count=1, + byte_count=900, + request_units=45.5, + current_time=started_at + timedelta(seconds=3), + ) + update_migration_resource( + state, + "cosmos:users:user_settings", + metrics, + current_time=started_at + timedelta(seconds=3), + ) + complete_migration_resource( + state, + "cosmos:users:user_settings", + result=metrics, + current_time=started_at + timedelta(seconds=3), + ) + + assert is_migration_resource_completed(state, "cosmos:users:user_settings") + assert metrics["processed_count"] == 9 + assert metrics["items_per_second"] == 3.0 + assert metrics["bytes_per_second"] == 300.0 + assert metrics["request_units_per_second"] == round(45.5 / 3, 3) + + resumed_state = initialize_migration_state( + state, + "11111111-1111-1111-1111-111111111111", + configuration, + current_time=started_at + timedelta(minutes=1), + ) + assert resumed_state["resume_count"] == 1 + assert is_migration_resource_completed(resumed_state, "cosmos:users:user_settings") + + changed_configuration = dict(configuration) + changed_configuration["parallel_operations"] = 16 + try: + initialize_migration_state( + resumed_state, + "11111111-1111-1111-1111-111111111111", + changed_configuration, + current_time=started_at + timedelta(minutes=2), + ) + except ValueError as exc: + assert "settings changed" in str(exc).lower() + else: + raise AssertionError("Checkpoint state accepted changed migration settings.") \ No newline at end of file diff --git a/functional_tests/test_data_management_search_write_fence.py b/functional_tests/test_data_management_search_write_fence.py new file mode 100644 index 000000000..66ee05aa6 --- /dev/null +++ b/functional_tests/test_data_management_search_write_fence.py @@ -0,0 +1,206 @@ +# test_data_management_search_write_fence.py +""" +Functional test for Data Management target AI Search write fencing. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures target SimpleChat Search writes drain before a migration +freezes them and cannot resume until the owning migration releases the fence. +""" + +import copy +import importlib.util +from pathlib import Path +import sys +import threading +import time + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +APP_ROOT = REPO_ROOT / "application" / "single_app" +MODULE_PATH = APP_ROOT / "functions_data_management_search_write_fence.py" +sys.path.insert(0, str(APP_ROOT)) + + +class ConflictError(Exception): + """Minimal Cosmos conflict error for optimistic fence replacement.""" + + status_code = 412 + + +class FakeGateContainer: + """Persist one fence document with Cosmos-like ETag replacement semantics.""" + + def __init__(self): + self.items = {} + self.etag_counter = 0 + self.lock = threading.Lock() + + def _save(self, body): + self.etag_counter += 1 + saved = copy.deepcopy(body) + saved["_etag"] = f"etag-{self.etag_counter}" + self.items[saved["id"]] = saved + return copy.deepcopy(saved) + + def read_item(self, item, partition_key): + assert item == partition_key + with self.lock: + if item not in self.items: + raise KeyError(item) + return copy.deepcopy(self.items[item]) + + def create_item(self, body): + with self.lock: + if body["id"] in self.items: + conflict = ConflictError("already exists") + conflict.status_code = 409 + raise conflict + return self._save(body) + + def replace_item(self, item, body, etag=None, **_kwargs): + with self.lock: + current = self.items.get(item) + if current is None or (etag and current.get("_etag") != etag): + raise ConflictError("stale fence") + return self._save(body) + + def delete_item(self, item, partition_key, etag=None, **_kwargs): + assert item == partition_key + with self.lock: + current = self.items.get(item) + if current is None or (etag and current.get("_etag") != etag): + raise ConflictError("stale fence") + self.items.pop(item, None) + + +def load_fence_module(): + """Load the standalone fence module without the Flask application.""" + module_name = "data_management_search_write_fence_test_module" + spec = importlib.util.spec_from_file_location(module_name, MODULE_PATH) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + assert spec and spec.loader + spec.loader.exec_module(module) + return module + + +def test_target_search_write_fence_drains_and_reopens_writer_slots(): + """A migration closes new target writes, drains existing work, and reopens only its own gate.""" + module = load_fence_module() + container = FakeGateContainer() + active_writer_slot = module.acquire_data_management_search_write_slot(container) + fence_holder = {} + heartbeat_count = [] + + def acquire_fence(): + fence_holder["fence"] = module.acquire_data_management_search_write_fence( + container, + "11111111-1111-1111-1111-111111111111", + lease_seconds=150, + heartbeat_callback=lambda: heartbeat_count.append(time.monotonic()), + ) + + thread = threading.Thread(target=acquire_fence) + thread.start() + time.sleep(0.2) + assert thread.is_alive() + module.release_data_management_search_write_slot(container, active_writer_slot) + thread.join(timeout=2.0) + + assert not thread.is_alive() + fence = fence_holder["fence"] + assert heartbeat_count == [] or heartbeat_count + with pytest.raises(module.DataManagementSearchWritesFrozenError): + module.acquire_data_management_search_write_slot(container) + + renewed = module.renew_data_management_search_write_fence( + container, + fence, + lease_seconds=150, + ) + assert renewed["fence_token"] == fence["fence_token"] + assert module.release_data_management_search_write_fence(container, fence) is True + + replacement_writer_slot = module.acquire_data_management_search_write_slot(container) + assert replacement_writer_slot + assert module.release_data_management_search_write_slot(container, replacement_writer_slot) is True + + +def test_only_the_owning_migration_can_release_target_search_write_fence(): + """A stale migration token cannot reopen another migration's frozen target writer gate.""" + module = load_fence_module() + container = FakeGateContainer() + fence = module.acquire_data_management_search_write_fence( + container, + "22222222-2222-2222-2222-222222222222", + lease_seconds=150, + ) + + stale_fence = dict(fence) + stale_fence["fence_token"] = "stale-token" + assert module.release_data_management_search_write_fence(container, stale_fence) is False + with pytest.raises(module.DataManagementSearchWritesFrozenError): + module.acquire_data_management_search_write_slot(container) + + assert module.release_data_management_search_write_fence(container, fence) is True + + +def test_target_migration_coordinator_blocks_independent_source_job_stores(): + """Two independent sources sharing one target must not overlap before Cosmos copy starts.""" + module = load_fence_module() + target_jobs = FakeGateContainer() + first_lock = module.acquire_data_management_target_migration_coordinator( + target_jobs, + "33333333-3333-3333-3333-333333333333", + lease_seconds=150, + ) + + with pytest.raises(module.DataManagementTargetMigrationCoordinatorError, match="Another SimpleChat source"): + module.acquire_data_management_target_migration_coordinator( + target_jobs, + "44444444-4444-4444-4444-444444444444", + lease_seconds=150, + ) + + renewed = module.renew_data_management_target_migration_coordinator( + target_jobs, + first_lock, + lease_seconds=150, + ) + assert renewed["migration_id"] == "33333333-3333-3333-3333-333333333333" + assert module.release_data_management_target_migration_coordinator( + target_jobs, + first_lock, + ) is True + + second_lock = module.acquire_data_management_target_migration_coordinator( + target_jobs, + "44444444-4444-4444-4444-444444444444", + lease_seconds=150, + ) + assert second_lock["migration_id"] == "44444444-4444-4444-4444-444444444444" + + +def test_ambiguous_target_search_write_retains_slot_until_quarantine_expires(monkeypatch): + """A lost response must keep the target writer slot long enough to block migration fencing.""" + module = load_fence_module() + container = FakeGateContainer() + + with pytest.raises(TimeoutError): + with module.hold_data_management_search_write_slot(container): + raise TimeoutError("response lost after Search accepted the request") + + gate = container.items[module.DATA_MANAGEMENT_SEARCH_WRITE_GATE_ID] + assert gate["active_writer_count"] == 1 + assert module.DATA_MANAGEMENT_SEARCH_WRITE_SLOT_LEASE_SECONDS >= 150 + gate["writer_leases"][0]["expires_at"] = "2000-01-01T00:00:00+00:00" + + fence = module.acquire_data_management_search_write_fence( + container, + "55555555-5555-5555-5555-555555555555", + lease_seconds=150, + ) + assert fence["migration_id"] == "55555555-5555-5555-5555-555555555555" diff --git a/functional_tests/test_data_management_search_write_fence_authorization.py b/functional_tests/test_data_management_search_write_fence_authorization.py new file mode 100644 index 000000000..da8ac76bc --- /dev/null +++ b/functional_tests/test_data_management_search_write_fence_authorization.py @@ -0,0 +1,191 @@ +# test_data_management_search_write_fence_authorization.py +""" +Functional test for Data Management Search fence authorization safety. +Version: 0.250.071 +Implemented in: 0.250.071 + +This test ensures an AI Search migration fence cannot make a document-unshare +request report success while stale Search chunks still grant access. +""" + +import ast +import copy +from pathlib import Path + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +DOCUMENTS_PATH = REPO_ROOT / "application" / "single_app" / "functions_documents.py" +DOCUMENTS_ROUTE_PATH = REPO_ROOT / "application" / "single_app" / "route_backend_documents.py" + + +def load_unshare_function(): + """Load only the unshare function from the production module with test dependencies.""" + source = DOCUMENTS_PATH.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(DOCUMENTS_PATH)) + function_node = next( + node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "unshare_document_from_user" + ) + isolated_module = ast.Module(body=[function_node], type_ignores=[]) + ast.fix_missing_locations(isolated_module) + + class FakeNotFoundError(Exception): + pass + + class FakeSearchWritesFrozenError(Exception): + pass + + class FakeAclProjectionDeferredError(Exception): + pass + + namespace = { + "CosmosResourceNotFoundError": FakeNotFoundError, + "DataManagementSearchWritesFrozenError": FakeSearchWritesFrozenError, + "DocumentSearchAclProjectionDeferredError": FakeAclProjectionDeferredError, + "datetime": __import__("datetime").datetime, + "timezone": __import__("datetime").timezone, + } + exec(compile(isolated_module, str(DOCUMENTS_PATH), "exec"), namespace) + return ( + namespace["unshare_document_from_user"], + namespace, + FakeSearchWritesFrozenError, + FakeAclProjectionDeferredError, + ) + + +def load_search_write_helpers(): + """Load only the indexing result checker and gated write helper with fake dependencies.""" + source = DOCUMENTS_PATH.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(DOCUMENTS_PATH)) + helper_names = { + "_search_indexing_results_succeeded", + "_execute_document_search_write", + } + helper_nodes = [ + node for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name in helper_names + ] + isolated_module = ast.Module(body=helper_nodes, type_ignores=[]) + ast.fix_missing_locations(isolated_module) + + class FakeSlot: + def __enter__(self): + return self + + def __exit__(self, *_args): + return False + + namespace = { + "hold_data_management_search_write_slot": lambda *_args, **_kwargs: FakeSlot(), + "cosmos_data_management_jobs_container": object(), + } + exec(compile(isolated_module, str(DOCUMENTS_PATH), "exec"), namespace) + return namespace["_execute_document_search_write"] + + +class FakeDocumentContainer: + """Persist the personal document ACL record in-memory.""" + + def __init__(self, document): + self.document = copy.deepcopy(document) + self.upserts = [] + + def read_item(self, item, partition_key): + assert item == partition_key == self.document["id"] + return copy.deepcopy(self.document) + + +def test_unshare_preserves_cosmos_acl_when_search_projection_is_frozen(): + """A fenced Search ACL update must not report a revocation or commit Cosmos first.""" + ( + unshare_document, + namespace, + frozen_error, + deferred_error, + ) = load_unshare_function() + container = FakeDocumentContainer({ + "id": "document-1", + "user_id": "owner-1", + "shared_user_ids": ["viewer-1,approved"], + }) + namespace["cosmos_user_documents_container"] = container + namespace["get_all_chunks"] = lambda *_args, **_kwargs: [{"id": "chunk-1"}] + namespace["_upsert_document_and_sync_access_index"] = lambda *_args, **_kwargs: container.upserts.append(True) + namespace["update_chunk_metadata"] = lambda *_args, **_kwargs: (_ for _ in ()).throw( + frozen_error("AI Search writes are temporarily frozen while a Data Management migration is running.") + ) + + with pytest.raises(deferred_error, match="temporarily frozen"): + unshare_document("document-1", "owner-1", "viewer-1") + assert container.document["shared_user_ids"] == ["viewer-1,approved"] + assert container.upserts == [] + + +def test_unshare_commits_cosmos_acl_after_search_projection_succeeds(): + """A completed Search projection permits the authoritative Cosmos ACL revocation.""" + unshare_document, namespace, _frozen_error, _deferred_error = load_unshare_function() + container = FakeDocumentContainer({ + "id": "document-1", + "user_id": "owner-1", + "shared_user_ids": ["viewer-1,approved"], + }) + projected_shared_ids = [] + namespace["cosmos_user_documents_container"] = container + namespace["get_all_chunks"] = lambda *_args, **_kwargs: [{"id": "chunk-1"}] + namespace["update_chunk_metadata"] = lambda **kwargs: projected_shared_ids.append( + kwargs["shared_user_ids"] + ) + + def upsert(_container, document, **_kwargs): + container.document = copy.deepcopy(document) + return copy.deepcopy(document) + + namespace["_upsert_document_and_sync_access_index"] = upsert + + assert unshare_document("document-1", "owner-1", "viewer-1") is True + assert projected_shared_ids == [[]] + assert container.document["shared_user_ids"] == [] + + +def test_search_write_rejects_unsuccessful_indexing_results(): + """Do not treat a non-throwing failed Search indexing result as a completed ACL projection.""" + execute_search_write = load_search_write_helpers() + + class FailedIndexClient: + def upload_documents(self, documents, **_kwargs): + assert documents == [{"id": "chunk-1"}] + return [{"succeeded": False}] + + try: + execute_search_write( + FailedIndexClient(), + "upload_documents", + documents=[{"id": "chunk-1"}], + ) + except RuntimeError as exc: + assert "did not acknowledge" in str(exc) + else: + raise AssertionError("A failed AI Search indexing result was treated as successful.") + + +def test_unshare_route_returns_retryable_response_for_deferred_acl_projection(): + """Keep an active target migration fence visible to callers as a retryable unshare response.""" + source = DOCUMENTS_ROUTE_PATH.read_text(encoding="utf-8") + tree = ast.parse(source, filename=str(DOCUMENTS_ROUTE_PATH)) + unshare_route = next( + node for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "api_unshare_document" + ) + deferred_handler = next( + handler for handler in ast.walk(unshare_route) + if isinstance(handler, ast.ExceptHandler) and + isinstance(handler.type, ast.Name) and + handler.type.id == "DocumentSearchAclProjectionDeferredError" + ) + handler_source = ast.get_source_segment(source, deferred_handler) or "" + + assert "Retry-After" in handler_source + assert "503" in handler_source diff --git a/functional_tests/test_data_management_security_patterns.py b/functional_tests/test_data_management_security_patterns.py index 83617738f..501f284d5 100644 --- a/functional_tests/test_data_management_security_patterns.py +++ b/functional_tests/test_data_management_security_patterns.py @@ -2,9 +2,9 @@ # test_data_management_security_patterns.py """ Functional test for Data Management security patterns. -Version: 0.250.051 +Version: 0.250.071 Implemented in: 0.241.211 -Updated in: 0.250.051 +Updated in: 0.250.076 This test ensures Data Management admin routes require authenticated admin access, secrets stay redacted in frontend responses, and the admin browser @@ -17,6 +17,8 @@ partition_key kwargs into the Python Cosmos SDK replace_item call. Version 0.250.051 keeps version coverage aligned with the Cosmos editor results-pane scroll refinement. +Version 0.250.071 adds resilient Data Management migrations, provenance, checkpoints, +incremental modes, reconciliation, and target Search writer fencing. """ import ast @@ -66,7 +68,7 @@ def test_version_and_container_registration(): """Validate the Data Management version and Cosmos job container registrations.""" config_source = read_text(CONFIG_FILE) - assert 'VERSION = "0.250.051"' in config_source + assert 'VERSION = "0.250.071"' in config_source assert 'cosmos_data_management_jobs_container_name = "data_management_jobs"' in config_source assert 'partition_key=PartitionKey(path="/id")' in config_source assert 'cosmos_data_management_job_items_container_name = "data_management_job_items"' in config_source @@ -76,7 +78,7 @@ def test_version_and_container_registration(): def test_admin_routes_require_login_admin_and_swagger_security(): """Validate every Data Management route has the required admin security stack.""" routes = route_functions_with_decorators() - assert len(routes) == 18 + assert len(routes) == 23 for function_name, decorators in routes: assert "swagger_route" in decorators, f"{function_name} missing swagger_route" @@ -100,6 +102,13 @@ def test_admin_routes_require_login_admin_and_swagger_security(): assert '/api/admin/data-management/cosmos-editor/document' in source assert 'save_data_management_cosmos_editor_document(' in source assert 'current_app._get_current_object()' in source + assert '/api/admin/data-management/jobs//retry' in source + assert '/api/admin/data-management/jobs//cancel' in source + assert '/api/admin/data-management/jobs//migration-manifest' in source + assert '/api/admin/data-management/jobs//migration-manifest/items/' in source + assert '/api/admin/data-management/jobs//progress' in source + assert 'retry_data_management_migration_job(job_id)' in source + assert 'request_data_management_migration_cancellation(' in source def test_settings_secrets_are_redacted_for_frontend(): @@ -135,6 +144,7 @@ def test_settings_secrets_are_redacted_for_frontend(): assert '_copy_source_blobs_to_target' in source assert 'get_data_management_migration_catalog' in source assert 'summarize_data_management_migration_plan' in source + assert 'preview_data_management_migration_plan' in source assert 'test_target_cosmos_connection' in source assert 'test_target_search_connection' in source assert 'test_target_enhanced_citation_storage_connection' in source @@ -147,6 +157,45 @@ def test_settings_secrets_are_redacted_for_frontend(): assert 'summarize_backup_artifacts(artifacts)' in source assert 'sanitized[field_name] = DATA_MANAGEMENT_REDACTED_VALUE' in source assert 'if payload.get(secret_field) == DATA_MANAGEMENT_REDACTED_VALUE:' in source + assert '_run_data_management_migration_preflight' in source + assert '_apply_temporary_destination_capacity' in source + assert '_restore_temporary_destination_capacity' in source + assert '_preflight_target_cosmos_migration_access' in source + assert '_preflight_target_ai_search_migration_access' in source + assert '_preflight_target_blob_migration_access' in source + assert '_acquire_migration_destination_lock' in source + assert '_assert_migration_job_lease' in source + assert '_validate_target_cosmos_container_partition_key' in source + assert 'recover_data_management_migration_jobs' in source + assert 'DataManagementMigrationCanceledError' in source + assert 'contains an unowned record that conflicts' in source + assert 'contains an unowned document that conflicts' in source + assert 'Destination blob exists without successful migration provenance' in source + assert 'create_migration_provenance_context' in source + assert 'add_cosmos_migration_provenance' in source + assert 'add_search_migration_provenance' in source + assert 'merge_blob_migration_metadata' in source + assert 'migration_max_parallel_operations' in source + assert 'migration_temporary_destination_ru_enabled' in source + assert 'DATA_MANAGEMENT_MIGRATION_MAX_DESTINATION_RU = 10000' in source + assert 'DATA_MANAGEMENT_MIGRATION_MODE_DELTA_UPSERT = "delta_upsert"' in source + assert 'DATA_MANAGEMENT_MIGRATION_MODE_MIRROR = "mirror_with_deletions"' in source + assert 'DATA_MANAGEMENT_MIRROR_CONFIRMATION = "MIRROR WITH DELETIONS"' in source + assert 'DATA_MANAGEMENT_SEARCH_WRITE_FREEZE_CONFIRMATION_ERROR' in source + assert '_validate_target_ai_search_migration_write_safety' in source + assert '_get_target_data_management_search_write_gate_container' in source + assert 'acquire_data_management_search_write_fence' in source + assert 'renew_data_management_search_write_fence' in source + assert 'release_data_management_search_write_fence' in source + assert 'acquire_data_management_target_migration_coordinator' in source + assert 'renew_data_management_target_migration_coordinator' in source + assert 'release_data_management_target_migration_coordinator' in source + assert '_acquire_target_migration_coordinator' in source + assert '_iter_search_document_pages' in source + assert 'order_by=["id asc"]' not in source + assert '"order_by": ["id asc"]' in source + assert '_run_data_management_migration_reconciliation' in source + assert 'preview_actual_divergence' in source def test_cosmos_editor_backend_safety_contract(): @@ -224,8 +273,9 @@ def test_admin_javascript_uses_safe_dom_patterns(): 'buildMigrationPlan()', 'queueMigration(false)', 'loadMigrationCatalog(targetType)', - 'renderMigrationSummary(data.summary || {})', + 'renderMigrationSummary(data.summary || {}, data.preview || null)', 'testTargetCosmos', + 'testMigrationAccess', 'testTargetSearch', 'testTargetEnhancedCitationStorage', 'Migration preview refreshed.', @@ -261,6 +311,15 @@ def test_admin_javascript_uses_safe_dom_patterns(): 'saveCosmosEditorDocument', 'confirmation_phrase: cosmosEditorConfirmationPhrase', 'closest("[data-ignore-data-management-change', + 'retryMigrationJob', + 'getMigrationLiveMetrics', + 'updateMigrationCapacityVisibility', + 'updateMigrationModeVisibility', + 'updateMigrationSearchWriteFreezeVisibility', + 'createMigrationPreviewOutcomes', + 'migrationMirrorConfirmationPhrase', + 'target_ai_search_writes_frozen', + 'Collisions', ]: assert required_snippet in source @@ -293,12 +352,31 @@ def test_admin_ui_exposes_data_management_without_external_assets(): 'id="data_management_target_cosmos_endpoint"', 'id="data_management_target_cosmos_database" value="SimpleChat" readonly aria-readonly="true"', 'id="data-management-target-cosmos-key-field"', + 'id="data_management_target_cosmos_subscription_id"', + 'id="data_management_target_cosmos_resource_group"', + 'id="data-management-migration-mode-section"', + 'id="data_management_migration_mode_new_only"', + 'id="data_management_migration_mode_delta_upsert"', + 'id="data_management_migration_mode_mirror_with_deletions"', + 'id="data_management_migration_baseline_job_id"', + 'id="data_management_migration_mirror_confirmation_phrase"', + 'id="data-management-migration-search-write-freeze"', + 'id="data_management_migration_target_search_writes_frozen"', + 'I confirm external destination AI Search writers are frozen for this migration', 'id="data-management-test-target-cosmos-btn"', 'id="data-management-target-ai-search-section"', 'id="data-management-test-target-search-btn"', 'id="data-management-target-enhanced-citations-section"', 'id="data-management-test-target-ec-storage-btn"', 'id="data-management-migration-workflow-section"', + 'id="data-management-test-migration-access-btn"', + 'Validate Cosmos Migration Access', + 'id="data_management_migration_max_parallel_operations"', + 'id="data_management_migration_retry_count"', + 'id="data_management_migration_skip_recent_within_hours"', + 'id="data_management_migration_temporary_destination_ru_enabled"', + 'id="data_management_migration_temporary_destination_ru"', + 'max="10000"', 'id="data-management-migration-summary"', 'id="data-management-execute-migration-btn"', 'id="data-management-cosmos-editor-section"', diff --git a/functional_tests/test_file_sync_azure_blob_storage.py b/functional_tests/test_file_sync_azure_blob_storage.py index de9442c08..ef0a1292f 100644 --- a/functional_tests/test_file_sync_azure_blob_storage.py +++ b/functional_tests/test_file_sync_azure_blob_storage.py @@ -85,7 +85,7 @@ def test_version_and_dependency_pin(): config_text = read_text("application/single_app/config.py") requirements_text = read_text("application/single_app/requirements.txt") - assert 'VERSION = "0.250.070"' in config_text + assert 'VERSION = "0.250.071"' in config_text assert "azure-storage-blob==12.24.1" in requirements_text diff --git a/functional_tests/test_file_sync_azure_files_identity.py b/functional_tests/test_file_sync_azure_files_identity.py index ee772c3a9..c89db0796 100644 --- a/functional_tests/test_file_sync_azure_files_identity.py +++ b/functional_tests/test_file_sync_azure_files_identity.py @@ -43,7 +43,7 @@ def test_version_and_dependency_pin(): config_text = read_text("application/single_app/config.py") requirements_text = read_text("application/single_app/requirements.txt") - assert 'VERSION = "0.250.070"' in config_text + assert 'VERSION = "0.250.071"' in config_text assert "azure-storage-file-share==12.25.0" in requirements_text diff --git a/functional_tests/test_file_sync_onedrive_personal.py b/functional_tests/test_file_sync_onedrive_personal.py index 4a7688d16..96d153c6c 100644 --- a/functional_tests/test_file_sync_onedrive_personal.py +++ b/functional_tests/test_file_sync_onedrive_personal.py @@ -43,7 +43,7 @@ def test_version_and_source_defaults(): settings_text = read_text("application/single_app/functions_settings.py") file_sync_text = read_text("application/single_app/functions_file_sync.py") - assert 'VERSION = "0.250.070"' in config_text + assert 'VERSION = "0.250.071"' in config_text assert "FILE_SYNC_SOURCE_TYPE_ONEDRIVE = \"onedrive\"" in file_sync_text assert "FILE_SYNC_SOURCE_TYPE_ONEDRIVE" in file_sync_text assert "FILE_SYNC_SOURCE_TYPE_ONEDRIVE: {\"client_secret\"}" in file_sync_text diff --git a/functional_tests/test_migration_provenance.py b/functional_tests/test_migration_provenance.py new file mode 100644 index 000000000..f367efff1 --- /dev/null +++ b/functional_tests/test_migration_provenance.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +""" +Functional test for migration provenance metadata. +Version: 0.250.071 +Implemented in: 0.250.074 + +This test ensures that migration IDs, timestamps, success states, and replay +eligibility are represented consistently for Cosmos, AI Search, and Storage. +""" + +import json +import shutil +import subprocess +from pathlib import Path + + +REPO_ROOT = Path(__file__).resolve().parents[1] +PROVENANCE_SCRIPT = REPO_ROOT / "scripts" / "Migration-Provenance.ps1" +STATE_SCRIPT = REPO_ROOT / "scripts" / "Migration-State.ps1" + + +def run_provenance_script(command: str) -> dict: + """Run a focused PowerShell provenance check and parse its JSON result.""" + powershell = shutil.which("pwsh") or shutil.which("powershell") + assert powershell, "PowerShell 7 or Windows PowerShell is required for this test." + + completed = subprocess.run( + [ + powershell, + "-NoProfile", + "-NonInteractive", + "-ExecutionPolicy", + "Bypass", + "-Command", + command, + ], + cwd=REPO_ROOT, + capture_output=True, + check=False, + text=True, + ) + assert completed.returncode == 0, completed.stderr + return json.loads(completed.stdout) + + +def test_migration_provenance_metadata_and_skip_rules() -> None: + """Verify one provenance contract across all migrated resource types.""" + script_path = str(PROVENANCE_SCRIPT).replace("'", "''") + state_script_path = str(STATE_SCRIPT).replace("'", "''") + command = f""" + $ErrorActionPreference = 'Stop' + . '{script_path}' + . '{state_script_path}' + $context = New-MigrationProvenanceContext ` + -MigrationId '11111111-1111-1111-1111-111111111111' ` + -MigratedAtUtc '2026-07-24T10:00:00.0000000+00:00' ` + -SkipMigratedWithinHours 24 + $statePath = Join-Path ([IO.Path]::GetTempPath()) "simplechat-provenance-$([Guid]::NewGuid().ToString('N')).json" + $stateConfiguration = [ordered]@{{ source = 'source'; destination = 'destination' }} + $initialState = Initialize-MigrationState ` + -MigrationType 'provenance' ` + -StateFilePath $statePath ` + -Configuration $stateConfiguration ` + -MigrationId $context.MigrationId ` + -Reset + $resumedState = Initialize-MigrationState ` + -MigrationType 'provenance' ` + -StateFilePath $statePath ` + -Configuration $stateConfiguration ` + -MigrationId $context.MigrationId + $record = New-MigrationProvenanceRecord -Context $context + $cosmosDocument = [ordered]@{{ id = 'document-1' }} + Add-CosmosMigrationProvenance -Document $cosmosDocument -Context $context + $searchDocument = [pscustomobject]@{{ id = 'document-1'; '@search.action' = 'upload' }} + Add-AISearchMigrationProvenance -Document $searchDocument -Context $context + $storageMetadata = Merge-StorageMigrationProvenance ` + -Metadata ([ordered]@{{ custom = 'preserved' }}) ` + -Context $context + $recentRecord = [ordered]@{{ + migrationId = '22222222-2222-2222-2222-222222222222' + migratedAtUtc = [DateTimeOffset]::UtcNow.AddHours(-1).ToString('o') + status = 'succeeded' + }} + $staleRecord = [ordered]@{{ + migrationId = '33333333-3333-3333-3333-333333333333' + migratedAtUtc = [DateTimeOffset]::UtcNow.AddHours(-25).ToString('o') + status = 'succeeded' + }} + $failedRecord = [ordered]@{{ + migrationId = $context.MigrationId + migratedAtUtc = $context.MigratedAtUtc + status = 'failed' + }} + [ordered]@{{ + migrationId = $context.MigrationId + stateMigrationId = $initialState.Data.migrationId + resumedStateMigrationId = $resumedState.Data.migrationId + record = $record + sameMigrationSkipped = Test-MigrationProvenanceSkip -Provenance $record -Context $context + recentMigrationSkipped = Test-MigrationProvenanceSkip -Provenance $recentRecord -Context $context + staleMigrationSkipped = Test-MigrationProvenanceSkip -Provenance $staleRecord -Context $context + failedMigrationSkipped = Test-MigrationProvenanceSkip -Provenance $failedRecord -Context $context + cosmos = Get-CosmosMigrationProvenance -Document $cosmosDocument + search = Get-AISearchMigrationProvenance -Document $searchDocument + storage = Get-StorageMigrationProvenance -Metadata $storageMetadata + customStorageMetadata = $storageMetadata['custom'] + }} | ConvertTo-Json -Depth 10 -Compress + Remove-Item -LiteralPath $statePath -Force + """ + + result = run_provenance_script(command) + + assert result["migrationId"] == "11111111-1111-1111-1111-111111111111" + assert result["stateMigrationId"] == result["migrationId"] + assert result["resumedStateMigrationId"] == result["migrationId"] + assert result["record"]["status"] == "succeeded" + assert result["sameMigrationSkipped"] is True + assert result["recentMigrationSkipped"] is True + assert result["staleMigrationSkipped"] is False + assert result["failedMigrationSkipped"] is False + assert result["cosmos"] == result["record"] + assert result["search"] == result["record"] + assert result["storage"] == result["record"] + assert result["customStorageMetadata"] == "preserved" \ No newline at end of file diff --git a/functional_tests/test_startup_scheduler_support.py b/functional_tests/test_startup_scheduler_support.py index 0516fa920..042668881 100644 --- a/functional_tests/test_startup_scheduler_support.py +++ b/functional_tests/test_startup_scheduler_support.py @@ -2,7 +2,7 @@ # test_startup_scheduler_support.py """ Functional test for startup and scheduler separation support. -Version: 0.239.129 +Version: 0.250.071 Implemented in: 0.239.129 This test ensures that the web process uses shared background-task helpers, @@ -33,7 +33,7 @@ def test_startup_scheduler_support() -> bool: print('Testing startup and scheduler support...') assert_contains(APP_FILE, 'from background_tasks import start_background_task_threads') - assert_contains(APP_FILE, 'start_background_task_threads()') + assert_contains(APP_FILE, 'start_background_task_threads(app=app)') assert_contains(APP_FILE, 'def should_start_background_tasks():') assert_contains(BACKGROUND_TASKS_FILE, 'def acquire_distributed_task_lock(task_name, lease_seconds):') assert_contains(BACKGROUND_TASKS_FILE, 'def release_distributed_task_lock(lock_document):') diff --git a/functional_tests/test_storage_migration_provenance.py b/functional_tests/test_storage_migration_provenance.py new file mode 100644 index 000000000..2e1da6e8b --- /dev/null +++ b/functional_tests/test_storage_migration_provenance.py @@ -0,0 +1,224 @@ +#!/usr/bin/env python3 +""" +Functional test for Storage Account migration provenance. +Version: 0.250.071 +Implemented in: 0.250.074 + +This test ensures successful blob copies retain existing metadata, receive +migration provenance metadata, and are skipped on a later run with the same ID. +""" + +from pathlib import Path +import shutil +import subprocess +import sys +import tempfile + + +REPO_ROOT = Path(__file__).resolve().parents[1] +SCRIPT_PATH = REPO_ROOT / "scripts" / "Migration-StorageAccount.ps1" + + +def powershell_single_quote(value: Path) -> str: + """Escape a path for a PowerShell single-quoted string.""" + return str(value).replace("'", "''") + + +def test_storage_migration_provenance() -> None: + """Exercise blob metadata stamping and remote provenance replay avoidance.""" + powershell = shutil.which("pwsh") + assert powershell, "PowerShell 7 is required to test the Storage migration script." + + with tempfile.TemporaryDirectory(prefix="simplechat-storage-provenance-") as temp_dir: + state_path = Path(temp_dir) / "migration-state.json" + harness = r''' +$ErrorActionPreference = "Stop" +$scriptPath = '__SCRIPT_PATH__' +$statePath = '__STATE_PATH__' +$global:azCopyInvocationCount = 0 +$global:metadataSetCount = 0 + +$destinationCloudBlob = [pscustomobject]@{ + Name = "documents/report.txt" + Metadata = [ordered]@{ existing = "preserved" } + FetchCount = 0 +} +$destinationCloudBlob | Add-Member -MemberType ScriptMethod -Name FetchAttributes -Value { + $this.FetchCount++ +} +$destinationCloudBlob | Add-Member -MemberType ScriptMethod -Name SetMetadata -Value { + $global:metadataSetCount++ +} +$global:sourceBlob = [pscustomobject]@{ + Name = "documents/report.txt" +} +$global:destinationBlob = [pscustomobject]@{ + Name = "documents/report.txt" + ICloudBlob = $destinationCloudBlob +} + +function Connect-AzAccount { + [CmdletBinding()] + param() +} + +function Set-AzContext { + [CmdletBinding()] + param( + [string]$SubscriptionId + ) +} + +function New-AzStorageContext { + [CmdletBinding()] + param( + [string]$StorageAccountName, + [switch]$UseConnectedAccount + ) + + return [pscustomobject]@{ AccountName = $StorageAccountName } +} + +function Get-AzStorageContainer { + [CmdletBinding()] + param( + [string]$Name, + [object]$Context + ) + + return [pscustomobject]@{ Name = $Name } +} + +function New-AzStorageContainer { + [CmdletBinding()] + param( + [string]$Name, + [object]$Context + ) + + return [pscustomobject]@{ Name = $Name } +} + +function New-AzStorageContainerSASToken { + [CmdletBinding()] + param( + [string]$Name, + [object]$Context, + [string]$Permission, + [string]$Protocol, + [datetime]$StartTime, + [datetime]$ExpiryTime, + [switch]$FullUri + ) + + return "https://$($Context.AccountName).blob.core.windows.net/$Name?sig=mock" +} + +function Get-AzStorageBlob { + [CmdletBinding()] + param( + [string]$Container, + [string]$Blob, + [object]$Context + ) + + if ($Context.AccountName -eq "sourcestorage") { + return @($global:sourceBlob) + } + if ($Context.AccountName -eq "destinationstorage") { + return $global:destinationBlob + } + throw "Unexpected storage context: $($Context.AccountName)" +} + +function azcopy { + [CmdletBinding()] + param( + [Parameter(ValueFromRemainingArguments = $true)] + [object[]]$Arguments + ) + + $global:azCopyInvocationCount++ + $global:LASTEXITCODE = 0 +} + +$parameters = @{ + SourceStorageAccount = "sourcestorage" + SourceSubscriptionId = "00000000-0000-0000-0000-000000000001" + DestinationStorageAccount = "destinationstorage" + DestinationSubscriptionId = "00000000-0000-0000-0000-000000000002" + Containers = @("user-documents") + DifferentialMigration = $true + MigrationId = "11111111-1111-1111-1111-111111111111" + SkipMigratedWithinHours = 24 + ShowProgress = $false + StateFilePath = $statePath +} + +& $scriptPath @parameters | Out-Null +$firstState = [IO.File]::ReadAllText($statePath) | ConvertFrom-Json -AsHashtable -Depth 100 +$secondStatePath = "$($statePath).second" +$parameters.StateFilePath = $secondStatePath +& $scriptPath @parameters | Out-Null +$secondState = [IO.File]::ReadAllText($secondStatePath) | ConvertFrom-Json -AsHashtable -Depth 100 +$firstResult = $firstState.resources["user-documents"].result +$secondResult = $secondState.resources["user-documents"].result + +if ( + $firstState.status -ne "completed" -or + $firstState.migrationId -ne $parameters.MigrationId -or + $firstResult.SourceBlobCount -ne 1 -or + $firstResult.TaggedBlobCount -ne 1 -or + $firstResult.SkippedByProvenance -ne $false +) { + throw "Unexpected first migration state: $($firstState | ConvertTo-Json -Compress -Depth 100)" +} +if ( + $destinationCloudBlob.Metadata["existing"] -ne "preserved" -or + $destinationCloudBlob.Metadata["simplechatMigrationId"] -ne $parameters.MigrationId -or + $destinationCloudBlob.Metadata["simplechatMigrationStatus"] -ne "succeeded" -or + [string]::IsNullOrWhiteSpace([string]$destinationCloudBlob.Metadata["simplechatMigratedAtUtc"]) +) { + throw "Blob metadata was not preserved and stamped correctly: $($destinationCloudBlob.Metadata | ConvertTo-Json -Compress)" +} +if ( + $secondState.status -ne "completed" -or + $secondResult.SkippedByProvenance -ne $true -or + $secondResult.SourceBlobCount -ne 1 -or + $global:azCopyInvocationCount -ne 1 -or + $global:metadataSetCount -ne 1 +) { + throw "The same migration ID did not bypass the already marked blob: $($secondState | ConvertTo-Json -Compress -Depth 100)" +} +''' + harness = harness.replace( + "__SCRIPT_PATH__", powershell_single_quote(SCRIPT_PATH) + ).replace("__STATE_PATH__", powershell_single_quote(state_path)) + + result = subprocess.run( + [ + powershell, + "-NoLogo", + "-NoProfile", + "-ExecutionPolicy", + "Bypass", + "-Command", + harness, + ], + cwd=REPO_ROOT, + text=True, + capture_output=True, + check=False, + ) + + if result.returncode != 0: + print(result.stdout) + print(result.stderr, file=sys.stderr) + raise AssertionError( + f"Storage migration provenance harness failed with exit code {result.returncode}." + ) + + +if __name__ == "__main__": + test_storage_migration_provenance() + print("Storage migration provenance test passed.") \ No newline at end of file diff --git a/scripts/Migration-AISearch.ps1 b/scripts/Migration-AISearch.ps1 new file mode 100644 index 000000000..a0a3254fc --- /dev/null +++ b/scripts/Migration-AISearch.ps1 @@ -0,0 +1,1466 @@ +# Migration-AISearch.ps1 +#Requires -Version 7.0 + +[CmdletBinding()] +param( + [string]$SourceSearchService = "", + + [string]$SourceResourceGroup = "", + + [string]$SourceSubscriptionId = "", + + [string]$SourceAdminKey = "", + + [string]$DestinationSearchService = "", + + [string]$DestinationResourceGroup = "", + + [string]$DestinationSubscriptionId = "", + + [string]$DestinationAdminKey = "", + + [bool]$DifferentialMigration = $false, + + [string]$MigrationId = "", + + [ValidateRange(0, 8760)] + [int]$SkipMigratedWithinHours = 24, + + [bool]$ShowProgress = $true, + + [ValidateRange(1, 100000)] + [int]$ProgressUpdateInterval = 100, + + [ValidateRange(1, 1000)] + [int]$BatchSize = 100, + + [ValidateRange(1, 64)] + [int]$MaxConcurrentBatches = 30, + + [ValidateRange(1024, 16777216)] + [int]$MaxBatchBytes = 15000000, + + [ValidateRange(1, 1000)] + [int]$PageSize = 100, + + [ValidateRange(1, 10)] + [int]$MaxRetryCount = 5, + + [ValidateRange(30, 3600)] + [int]$RequestTimeoutSeconds = 300, + + [string]$ApiVersion = "2026-04-01", + + [string]$ManagementApiVersion = "2025-05-01", + + # Optional parameters for testing against Azure Government or other sovereign clouds. + [string]$SearchDnsSuffix = "search.windows.net", + + [string]$StateFilePath = "", + + [switch]$ResetState +) + +$ErrorActionPreference = "Stop" +. (Join-Path $PSScriptRoot "Migration-State.ps1") +. (Join-Path $PSScriptRoot "Migration-Provenance.ps1") + +function Get-AISearchProgressPercent { + param( + [long]$ProcessedCount, + [long]$TotalCount + ) + + if ($TotalCount -le 0) { + return 100 + } + + return [int][Math]::Min( + 100, + [Math]::Floor(([double]$ProcessedCount / [double]$TotalCount) * 100) + ) +} + +function Write-AISearchProgress { + param( + [int]$Id, + [int]$ParentId = -1, + [string]$Activity, + [string]$Status = "", + [string]$CurrentOperation = "", + [int]$PercentComplete = -1, + [switch]$Completed + ) + + if (-not $ShowProgress) { + return + } + + $progressParameters = @{ + Id = $Id + ParentId = $ParentId + Activity = $Activity + } + + if ($Completed) { + $progressParameters.Completed = $true + } + else { + $progressParameters.Status = $Status + $progressParameters.PercentComplete = $PercentComplete + if (-not [string]::IsNullOrWhiteSpace($CurrentOperation)) { + $progressParameters.CurrentOperation = $CurrentOperation + } + } + + Write-Progress @progressParameters +} + +function Write-AISearchCountProgress { + param( + [int]$Id, + [int]$ParentId = -1, + [string]$Activity, + [string]$Phase, + [long]$ProcessedCount, + [long]$TotalCount, + [string]$CurrentOperation = "" + ) + + $percentComplete = Get-AISearchProgressPercent ` + -ProcessedCount $ProcessedCount ` + -TotalCount $TotalCount + $remainingCount = [Math]::Max(0, $TotalCount - $ProcessedCount) + $status = "$Phase`: $ProcessedCount/$TotalCount | Remaining: $remainingCount | $percentComplete%" + + Write-AISearchProgress ` + -Id $Id ` + -ParentId $ParentId ` + -Activity $Activity ` + -Status $status ` + -CurrentOperation $CurrentOperation ` + -PercentComplete $percentComplete +} + +function Test-AISearchProgressCheckpoint { + param( + [long]$ProcessedCount, + [long]$TotalCount + ) + + return ( + $ProcessedCount -eq 1 -or + $ProcessedCount -eq $TotalCount -or + ($ProcessedCount % $ProgressUpdateInterval) -eq 0 + ) +} + +function Get-AISearchAdminKey { + param( + [string]$ProvidedKey, + [string]$SubscriptionId, + [string]$ResourceGroupName, + [string]$ServiceName + ) + + if (-not [string]::IsNullOrWhiteSpace($ProvidedKey)) { + return $ProvidedKey + } + + Set-AzContext -SubscriptionId $SubscriptionId | Out-Null + $encodedResourceGroupName = [Uri]::EscapeDataString($ResourceGroupName) + $encodedServiceName = [Uri]::EscapeDataString($ServiceName) + $keyPath = "/subscriptions/$SubscriptionId/resourceGroups/$encodedResourceGroupName/providers/Microsoft.Search/searchServices/$encodedServiceName/listAdminKeys?api-version=$ManagementApiVersion" + $keyResponse = Invoke-AzRestMethod -Method "POST" -Path $keyPath + + if ($keyResponse.StatusCode -ne 200) { + throw "Admin key lookup failed for AI Search service '$ServiceName' with HTTP $($keyResponse.StatusCode)." + } + + $keyPair = $keyResponse.Content | ConvertFrom-Json + if ([string]::IsNullOrWhiteSpace($keyPair.primaryKey)) { + throw "The primary admin key for AI Search service '$ServiceName' could not be resolved." + } + + return $keyPair.primaryKey +} + +function New-AISearchUri { + param( + [string]$Endpoint, + [string]$RelativePath + ) + + return "{0}/{1}?api-version={2}" -f $Endpoint.Trim().TrimEnd("/"), $RelativePath.Trim().TrimStart("/"), $ApiVersion.Trim() +} + +function Invoke-AISearchRequest { + param( + [ValidateSet("GET", "POST", "PUT")] + [string]$Method, + [string]$Uri, + [string]$AdminKey, + [AllowNull()] + [object]$Body = $null + ) + + $headers = @{ + "api-key" = $AdminKey + "Accept" = "application/json" + } + + $invokeParameters = @{ + Method = $Method + Uri = $Uri + Headers = $headers + TimeoutSec = $RequestTimeoutSeconds + ErrorAction = "Stop" + } + + if ($null -ne $Body) { + $invokeParameters.ContentType = "application/json" + $invokeParameters.Body = $Body | ConvertTo-Json -Depth 100 -Compress + } + + for ($attempt = 1; $attempt -le $MaxRetryCount; $attempt++) { + try { + return Invoke-RestMethod @invokeParameters + } + catch { + $statusCode = $null + if ($null -ne $_.Exception.Response) { + $statusCode = [int]$_.Exception.Response.StatusCode + } + + $requestTimedOut = $_.Exception.Message -match '(?i)timed out|timeout' + $retryable = $statusCode -in @(409, 422, 429, 503) -or $requestTimedOut + if (-not $retryable -or $attempt -eq $MaxRetryCount) { + throw + } + + $retryDelaySeconds = [Math]::Pow(2, $attempt - 1) + $failureDescription = if ($requestTimedOut) { "timed out" } else { "returned HTTP $statusCode" } + Write-Warning "AI Search request $failureDescription. Retrying in $retryDelaySeconds second(s)." + Start-Sleep -Seconds $retryDelaySeconds + } + } +} + +function ConvertTo-AISearchWritableDefinition { + param( + [object]$Definition + ) + + $copy = $Definition | ConvertTo-Json -Depth 100 | ConvertFrom-Json -Depth 100 + foreach ($property in @($copy.PSObject.Properties)) { + if ($property.Name.StartsWith("@odata.", [System.StringComparison]::Ordinal)) { + $copy.PSObject.Properties.Remove($property.Name) + } + } + + $definitionJson = $copy | ConvertTo-Json -Depth 100 -Compress + if ($definitionJson.Contains('""')) { + throw "The definition for '$($copy.name)' contains a redacted secret and cannot be recreated automatically." + } + + return $copy +} + +function Update-AISearchMigrationProvenanceFields { + param( + [object]$IndexDefinition + ) + + $fields = [System.Collections.Generic.List[object]]::new() + $existingFieldsByName = @{} + foreach ($field in @($IndexDefinition.fields)) { + $fields.Add($field) + $existingFieldsByName[[string]$field.name] = $field + } + + $addedFieldCount = 0 + foreach ($fieldDefinition in Get-AISearchMigrationProvenanceFieldDefinitions) { + $fieldName = [string]$fieldDefinition["name"] + if ($existingFieldsByName.ContainsKey($fieldName)) { + $existingField = $existingFieldsByName[$fieldName] + if ( + [string]$existingField.type -ne [string]$fieldDefinition["type"] -or + $existingField.filterable -ne $true + ) { + throw "Index '$($IndexDefinition.name)' has an incompatible migration provenance field '$fieldName'. It must be a filterable $($fieldDefinition["type"]) field." + } + continue + } + + $fields.Add([pscustomobject]$fieldDefinition) + $addedFieldCount++ + } + + $IndexDefinition | Add-Member ` + -MemberType NoteProperty ` + -Name "fields" ` + -Value $fields.ToArray() ` + -Force + return $addedFieldCount +} + +function Get-AISearchKeyField { + param( + [object]$IndexDefinition + ) + + $keyFields = @($IndexDefinition.fields | Where-Object { $_.key -eq $true }) + if ($keyFields.Count -ne 1) { + throw "AI Search index '$($IndexDefinition.name)' must have exactly one key field." + } + + return $keyFields[0] +} + +function Get-AISearchDocumentCount { + param( + [string]$Endpoint, + [string]$AdminKey, + [string]$IndexName + ) + + $encodedIndexName = [Uri]::EscapeDataString($IndexName) + $countUri = New-AISearchUri -Endpoint $Endpoint -RelativePath "indexes/$encodedIndexName/docs/`$count" + return [long](Invoke-AISearchRequest -Method "GET" -Uri $countUri -AdminKey $AdminKey) +} + +function Get-AISearchDocuments { + param( + [string]$Endpoint, + [string]$AdminKey, + [object]$IndexDefinition, + [switch]$KeysOnly, + [string]$FilterExpression = "", + [string[]]$SelectFields = @(), + [long]$TotalCount = 0, + [string]$ProgressActivity = "", + [string]$ProgressPhase = "Documents", + [string]$StartAfterKey = "", + [long]$InitialProcessedCount = 0 + ) + + $keyField = Get-AISearchKeyField -IndexDefinition $IndexDefinition + $encodedIndexName = [Uri]::EscapeDataString($IndexDefinition.name) + $searchUri = New-AISearchUri -Endpoint $Endpoint -RelativePath "indexes/$encodedIndexName/docs/search" + $supportsKeysetPagination = $keyField.filterable -eq $true -and $keyField.sortable -eq $true + $skip = 0 + if (-not [string]::IsNullOrEmpty($StartAfterKey) -and -not $supportsKeysetPagination) { + throw "Index '$($IndexDefinition.name)' cannot resume by key because its key field is not filterable and sortable." + } + $lastKey = if ([string]::IsNullOrEmpty($StartAfterKey)) { $null } else { $StartAfterKey } + $migrationComplete = $false + $pageNumber = 0 + $emittedCount = $InitialProcessedCount + + while (-not $migrationComplete) { + $requestBody = [ordered]@{ + search = "*" + top = $PageSize + count = $true + } + + $selectedFields = [System.Collections.Generic.List[string]]::new() + if ($KeysOnly) { + $selectedFields.Add([string]$keyField.name) + } + foreach ($fieldName in $SelectFields) { + if (-not [string]::IsNullOrWhiteSpace($fieldName) -and -not $selectedFields.Contains($fieldName)) { + $selectedFields.Add($fieldName) + } + } + if ($selectedFields.Count -gt 0) { + $requestBody.select = $selectedFields -join "," + } + + $filterExpressions = [System.Collections.Generic.List[string]]::new() + if (-not [string]::IsNullOrWhiteSpace($FilterExpression)) { + $filterExpressions.Add("($FilterExpression)") + } + if ($supportsKeysetPagination) { + $requestBody.orderby = "$($keyField.name) asc" + if ($null -ne $lastKey) { + $escapedLastKey = ([string]$lastKey).Replace("'", "''") + $filterExpressions.Add("$($keyField.name) gt '$escapedLastKey'") + } + } + else { + $requestBody.skip = $skip + } + if ($filterExpressions.Count -gt 0) { + $requestBody.filter = $filterExpressions -join " and " + } + + $requestUri = $searchUri + $logicalPageCount = 0 + $logicalPageLastKey = $lastKey + + do { + $pageNumber++ + if (-not [string]::IsNullOrWhiteSpace($ProgressActivity)) { + Write-AISearchCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $ProgressActivity ` + -Phase $ProgressPhase ` + -ProcessedCount $emittedCount ` + -TotalCount $TotalCount ` + -CurrentOperation "Requesting page $pageNumber (up to $PageSize documents; timeout: $RequestTimeoutSeconds seconds)" + } + + $requestStopwatch = [System.Diagnostics.Stopwatch]::StartNew() + $response = Invoke-AISearchRequest ` + -Method "POST" ` + -Uri $requestUri ` + -AdminKey $AdminKey ` + -Body $requestBody + $requestStopwatch.Stop() + + $documents = @($response.value) + if ($pageNumber -eq 1 -or ($pageNumber % 10) -eq 0 -or $documents.Count -lt $PageSize) { + $elapsedSeconds = [Math]::Round($requestStopwatch.Elapsed.TotalSeconds, 1) + Write-Host "Fetched $ProgressPhase page $pageNumber for index '$($IndexDefinition.name)': $($documents.Count) document(s) in $elapsedSeconds second(s)." + } + foreach ($document in $documents) { + Write-Output $document + $emittedCount++ + } + + $logicalPageCount += $documents.Count + if ($documents.Count -gt 0) { + $logicalPageLastKey = $documents[-1].PSObject.Properties[$keyField.name].Value + } + + $nextLink = $response.PSObject.Properties["@odata.nextLink"].Value + $nextPageParameters = $response.PSObject.Properties["@search.nextPageParameters"].Value + if ($null -ne $nextLink -and $null -ne $nextPageParameters) { + $requestUri = $nextLink + $requestBody = $nextPageParameters + } + else { + $requestUri = $null + } + } while ($null -ne $requestUri) + + if ($logicalPageCount -lt $PageSize) { + $migrationComplete = $true + continue + } + + if ($supportsKeysetPagination) { + if ($logicalPageLastKey -eq $lastKey) { + throw "Keyset pagination did not advance for index '$($IndexDefinition.name)'." + } + $lastKey = $logicalPageLastKey + } + else { + $skip += $logicalPageCount + if ($skip -gt 100000) { + throw "Index '$($IndexDefinition.name)' exceeds skip-based paging limits. Its key field must be filterable and sortable to migrate more than 100,000 documents." + } + } + } +} + +function ConvertTo-AISearchWriteDocument { + param( + [object]$Document, + [ValidateSet("upload", "mergeOrUpload")] + [string]$Action, + [AllowNull()] + [object]$MigrationProvenanceContext = $null + ) + + $writeDocument = [ordered]@{ + "@search.action" = $Action + } + + foreach ($property in $Document.PSObject.Properties) { + if (-not $property.Name.StartsWith("@search.", [System.StringComparison]::Ordinal)) { + $writeDocument[$property.Name] = $property.Value + } + } + + if ($null -ne $MigrationProvenanceContext) { + Add-AISearchMigrationProvenance ` + -Document $writeDocument ` + -Context $MigrationProvenanceContext + } + + return [pscustomobject]$writeDocument +} + +function Send-AISearchDocumentBatch { + param( + [string]$Endpoint, + [string]$AdminKey, + [string]$IndexName, + [System.Collections.Generic.List[object]]$Documents + ) + + if ($Documents.Count -eq 0) { + return 0 + } + + $encodedIndexName = [Uri]::EscapeDataString($IndexName) + $indexUri = New-AISearchUri -Endpoint $Endpoint -RelativePath "indexes/$encodedIndexName/docs/index" + for ($attempt = 1; $attempt -le $MaxRetryCount; $attempt++) { + $response = Invoke-AISearchRequest ` + -Method "POST" ` + -Uri $indexUri ` + -AdminKey $AdminKey ` + -Body @{ value = @($Documents) } + + $failedDocuments = @($response.value | Where-Object { $_.status -ne $true }) + if ($failedDocuments.Count -eq 0) { + return $Documents.Count + } + + $permanentFailures = @($failedDocuments | Where-Object { + [int]$_.statusCode -notin @(409, 422, 429, 503) + }) + if ($permanentFailures.Count -gt 0 -or $attempt -eq $MaxRetryCount) { + $failureSummary = $failedDocuments | ForEach-Object { + "key=$($_.key), status=$($_.statusCode), error=$($_.errorMessage)" + } + throw "Document indexing failed for index '$IndexName': $($failureSummary -join '; ')" + } + + $retryDelaySeconds = [Math]::Pow(2, $attempt - 1) + Write-Warning "Document indexing returned transient failures for '$IndexName'. Retrying the batch in $retryDelaySeconds second(s)." + Start-Sleep -Seconds $retryDelaySeconds + } +} + +function New-AISearchDocumentBatchWorkItem { + param( + [System.Collections.Generic.List[object]]$Documents, + [long]$Sequence + ) + + $documentArray = @($Documents.ToArray()) + return [pscustomobject]@{ + Sequence = $Sequence + DocumentCount = $documentArray.Count + Documents = $documentArray + Body = @{ value = $documentArray } | ConvertTo-Json -Depth 100 -Compress + } +} + +function Invoke-AISearchParallelDocumentBatches { + param( + [Parameter(ValueFromPipeline)] + [object]$WorkItem, + [string]$Endpoint, + [string]$AdminKey, + [string]$IndexName + ) + + begin { + $workItems = [System.Collections.Generic.List[object]]::new() + } + process { + $workItems.Add($WorkItem) + } + end { + if ($workItems.Count -eq 0) { + return + } + + $encodedIndexName = [Uri]::EscapeDataString($IndexName) + $indexUri = New-AISearchUri ` + -Endpoint $Endpoint ` + -RelativePath "indexes/$encodedIndexName/docs/index" + $retryLimit = $MaxRetryCount + $requestTimeout = $RequestTimeoutSeconds + $invokeRestMethodCommand = Get-Command "Invoke-RestMethod" -ErrorAction "Stop" + $invokeRestMethodOverrideDefinition = if ( + $invokeRestMethodCommand.CommandType -eq "Function" + ) { + $invokeRestMethodCommand.Definition + } + else { + "" + } + $startSleepCommand = Get-Command "Start-Sleep" -ErrorAction "Stop" + $startSleepOverrideDefinition = if ($startSleepCommand.CommandType -eq "Function") { + $startSleepCommand.Definition + } + else { + "" + } + + $workItems | ForEach-Object -Parallel { + $currentWorkItem = $_ + if (-not [string]::IsNullOrWhiteSpace($using:invokeRestMethodOverrideDefinition)) { + Set-Item ` + -Path "Function:Invoke-RestMethod" ` + -Value ([scriptblock]::Create($using:invokeRestMethodOverrideDefinition)) + } + if (-not [string]::IsNullOrWhiteSpace($using:startSleepOverrideDefinition)) { + Set-Item ` + -Path "Function:Start-Sleep" ` + -Value ([scriptblock]::Create($using:startSleepOverrideDefinition)) + } + + $headers = @{ + "api-key" = $using:AdminKey + "Accept" = "application/json" + } + for ($attempt = 1; $attempt -le $using:retryLimit; $attempt++) { + try { + $response = Invoke-RestMethod ` + -Method "POST" ` + -Uri $using:indexUri ` + -Headers $headers ` + -ContentType "application/json" ` + -Body $currentWorkItem.Body ` + -TimeoutSec $using:requestTimeout ` + -ErrorAction "Stop" + } + catch { + $statusCode = 0 + if ($null -ne $_.Exception.Response) { + try { + $statusCode = [int]$_.Exception.Response.StatusCode + } + catch { + $statusCode = 0 + } + } + + $requestTimedOut = $_.Exception.Message -match '(?i)timed out|timeout' + $retryable = $statusCode -in @(409, 422, 429, 503) -or $requestTimedOut + if ($retryable -and $attempt -lt $using:retryLimit) { + $retryDelaySeconds = [int][Math]::Pow(2, $attempt - 1) + Start-Sleep -Seconds $retryDelaySeconds + continue + } + + [pscustomobject]@{ + Sequence = $currentWorkItem.Sequence + DocumentCount = $currentWorkItem.DocumentCount + Succeeded = $false + ErrorMessage = "AI Search batch request failed after $attempt attempt(s): $($_.Exception.Message)" + } + return + } + + $failedDocuments = @($response.value | Where-Object { $_.status -ne $true }) + if ($failedDocuments.Count -eq 0) { + [pscustomobject]@{ + Sequence = $currentWorkItem.Sequence + DocumentCount = $currentWorkItem.DocumentCount + Succeeded = $true + ErrorMessage = "" + } + return + } + + $permanentFailures = @($failedDocuments | Where-Object { + [int]$_.statusCode -notin @(409, 422, 429, 503) + }) + if ($permanentFailures.Count -gt 0 -or $attempt -eq $using:retryLimit) { + $failureSummary = $failedDocuments | ForEach-Object { + "key=$($_.key), status=$($_.statusCode), error=$($_.errorMessage)" + } + [pscustomobject]@{ + Sequence = $currentWorkItem.Sequence + DocumentCount = $currentWorkItem.DocumentCount + Succeeded = $false + ErrorMessage = "Document indexing failed for index '$using:IndexName': $($failureSummary -join '; ')" + } + return + } + + $retryDelaySeconds = [int][Math]::Pow(2, $attempt - 1) + Start-Sleep -Seconds $retryDelaySeconds + } + } -ThrottleLimit $MaxConcurrentBatches + } +} + +function Send-AISearchDocumentBatchWindow { + param( + [System.Collections.Generic.List[object]]$WorkItems, + [string]$Endpoint, + [string]$AdminKey, + [string]$IndexName + ) + + if ($WorkItems.Count -eq 0) { + return 0 + } + + if ($MaxConcurrentBatches -eq 1) { + $copiedCount = 0 + foreach ($workItem in $WorkItems) { + $documents = [System.Collections.Generic.List[object]]::new() + foreach ($document in @($workItem.Documents)) { + $documents.Add($document) + } + $copiedCount += Send-AISearchDocumentBatch ` + -Endpoint $Endpoint ` + -AdminKey $AdminKey ` + -IndexName $IndexName ` + -Documents $documents + } + return $copiedCount + } + + $batchResults = @( + $WorkItems.ToArray() | + Invoke-AISearchParallelDocumentBatches ` + -Endpoint $Endpoint ` + -AdminKey $AdminKey ` + -IndexName $IndexName | + Sort-Object Sequence + ) + if ($batchResults.Count -ne $WorkItems.Count) { + throw "AI Search parallel batch processing returned $($batchResults.Count) result(s) for $($WorkItems.Count) batch(es)." + } + + $failedBatch = @($batchResults | Where-Object { -not $_.Succeeded } | Select-Object -First 1) + if ($failedBatch.Count -gt 0) { + throw $failedBatch[0].ErrorMessage + } + + return [long](($batchResults | Measure-Object -Property DocumentCount -Sum).Sum) +} + +function Copy-AISearchSynonymMaps { + param( + [string]$SourceEndpoint, + [string]$SourceKey, + [string]$DestinationEndpoint, + [string]$DestinationKey + ) + + $sourceUri = New-AISearchUri -Endpoint $SourceEndpoint -RelativePath "synonymmaps" + $destinationUri = New-AISearchUri -Endpoint $DestinationEndpoint -RelativePath "synonymmaps" + $sourceMaps = @((Invoke-AISearchRequest -Method "GET" -Uri $sourceUri -AdminKey $SourceKey).value) + $destinationMaps = @((Invoke-AISearchRequest -Method "GET" -Uri $destinationUri -AdminKey $DestinationKey).value) + $destinationMapNames = @{} + foreach ($synonymMap in $destinationMaps) { + $destinationMapNames[$synonymMap.name] = $true + } + + foreach ($synonymMap in $sourceMaps) { + if ($DifferentialMigration -and $destinationMapNames.ContainsKey($synonymMap.name)) { + Write-Host "Skipping existing synonym map: $($synonymMap.name)" + continue + } + + $writableMap = ConvertTo-AISearchWritableDefinition -Definition $synonymMap + $encodedMapName = [Uri]::EscapeDataString($synonymMap.name) + $mapUri = New-AISearchUri -Endpoint $DestinationEndpoint -RelativePath "synonymmaps/$encodedMapName" + Invoke-AISearchRequest -Method "PUT" -Uri $mapUri -AdminKey $DestinationKey -Body $writableMap | Out-Null + Write-Host "Migrated synonym map: $($synonymMap.name)" + } +} + +function Update-AISearchDocumentCheckpoint { + param( + [object]$StateContext, + [string]$ResourceName, + [object]$KeyField, + [long]$SourceDocumentCount, + [AllowNull()] + [object]$LastCommittedKey, + [long]$ProcessedCount, + [long]$CopiedCount, + [long]$SkippedCount, + [long]$BatchCount + ) + + Update-MigrationResourceCheckpoint ` + -Context $StateContext ` + -ResourceName $ResourceName ` + -Progress ([ordered]@{ + phase = "source_documents" + keyField = $KeyField.name + resumeSupported = ($KeyField.filterable -eq $true -and $KeyField.sortable -eq $true) + sourceDocumentCount = $SourceDocumentCount + lastCommittedKey = $LastCommittedKey + processedCount = $ProcessedCount + copiedCount = $CopiedCount + skippedCount = $SkippedCount + batchCount = $BatchCount + }) +} + +function Get-AISearchMigrationProvenanceSkipKeys { + param( + [object]$DestinationIndex, + [string]$DestinationEndpoint, + [string]$DestinationKey, + [object]$MigrationProvenanceContext + ) + + $escapedMigrationId = ([string]$MigrationProvenanceContext.MigrationId).Replace("'", "''") + $filterExpression = "simplechatMigrationStatus eq 'succeeded' and (simplechatMigrationId eq '$escapedMigrationId'" + if ([int]$MigrationProvenanceContext.SkipMigratedWithinHours -gt 0) { + $cutoffUtc = [DateTimeOffset]::UtcNow.AddHours( + -[int]$MigrationProvenanceContext.SkipMigratedWithinHours + ) + $cutoffText = $cutoffUtc.UtcDateTime.ToString( + "yyyy-MM-ddTHH:mm:ss.fffffff'Z'", + [System.Globalization.CultureInfo]::InvariantCulture + ) + $filterExpression += " or simplechatMigratedAtUtc ge $cutoffText" + } + $filterExpression += ")" + + $keyField = Get-AISearchKeyField -IndexDefinition $DestinationIndex + $skippedKeys = [System.Collections.Generic.HashSet[string]]::new( + [System.StringComparer]::Ordinal + ) + Get-AISearchDocuments ` + -Endpoint $DestinationEndpoint ` + -AdminKey $DestinationKey ` + -IndexDefinition $DestinationIndex ` + -KeysOnly ` + -FilterExpression $filterExpression ` + -ProgressActivity "Index '$($DestinationIndex.name)' documents" ` + -ProgressPhase "Recently migrated destination keys" | ForEach-Object { + $documentKey = [string]$_.PSObject.Properties[$keyField.name].Value + [void]$skippedKeys.Add($documentKey) + } + return ,$skippedKeys +} + +function Copy-AISearchIndexDocuments { + param( + [object]$SourceIndex, + [AllowNull()] + [object]$DestinationIndex, + [bool]$DestinationIndexExisted, + [string]$SourceEndpoint, + [string]$SourceKey, + [string]$DestinationEndpoint, + [string]$DestinationKey, + [object]$StateContext, + [string]$ResourceName, + [object]$MigrationProvenanceContext + ) + + $sourceKeyField = Get-AISearchKeyField -IndexDefinition $SourceIndex + $indexActivity = "Index '$($SourceIndex.name)' documents" + $sourceDocumentCount = Get-AISearchDocumentCount ` + -Endpoint $SourceEndpoint ` + -AdminKey $SourceKey ` + -IndexName $SourceIndex.name + $resumeSupported = $sourceKeyField.filterable -eq $true -and $sourceKeyField.sortable -eq $true + $processedCount = 0 + $copiedCount = 0 + $skippedCount = 0 + $batchCount = 0 + $lastCommittedKey = $null + $resumeAfterKey = "" + + $resourceCheckpoint = Get-MigrationResourceCheckpoint ` + -Context $StateContext ` + -ResourceName $ResourceName + if ($null -eq $resourceCheckpoint) { + throw "Migration state for index '$($SourceIndex.name)' was not initialized." + } + + $storedProgress = $resourceCheckpoint["progress"] + if ($null -ne $storedProgress -and $storedProgress.Count -gt 0) { + if ($storedProgress["keyField"] -ne $sourceKeyField.name) { + throw "Index '$($SourceIndex.name)' changed key fields after checkpointing. Use -ResetState to start over." + } + if ([bool]$storedProgress["resumeSupported"] -ne $resumeSupported) { + throw "Index '$($SourceIndex.name)' changed keyset paging capabilities after checkpointing. Use -ResetState to start over." + } + + if ([long]$storedProgress["sourceDocumentCount"] -ne $sourceDocumentCount) { + Write-Warning "Source document count changed for index '$($SourceIndex.name)'. Restarting this index from the beginning." + } + elseif ( + $resumeSupported -and + -not [string]::IsNullOrEmpty([string]$storedProgress["lastCommittedKey"]) + ) { + $lastCommittedKey = [string]$storedProgress["lastCommittedKey"] + $resumeAfterKey = $lastCommittedKey + $processedCount = [long]$storedProgress["processedCount"] + $copiedCount = [long]$storedProgress["copiedCount"] + $skippedCount = [long]$storedProgress["skippedCount"] + $batchCount = [long]$storedProgress["batchCount"] + Write-Host "Resuming index '$($SourceIndex.name)' after committed key '$lastCommittedKey' ($processedCount/$sourceDocumentCount processed)." + } + elseif ([long]$storedProgress["processedCount"] -gt 0) { + Write-Warning "Index '$($SourceIndex.name)' has no safe keyset checkpoint. Restarting this index from the beginning." + } + } + + Update-AISearchDocumentCheckpoint ` + -StateContext $StateContext ` + -ResourceName $ResourceName ` + -KeyField $sourceKeyField ` + -SourceDocumentCount $sourceDocumentCount ` + -LastCommittedKey $lastCommittedKey ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -BatchCount $batchCount + $existingKeys = [System.Collections.Generic.HashSet[string]]::new([System.StringComparer]::Ordinal) + $migrationProvenanceSkipKeys = [System.Collections.Generic.HashSet[string]]::new( + [System.StringComparer]::Ordinal + ) + + if ($DestinationIndexExisted -and -not $DifferentialMigration) { + $migrationProvenanceSkipKeys = Get-AISearchMigrationProvenanceSkipKeys ` + -DestinationIndex $DestinationIndex ` + -DestinationEndpoint $DestinationEndpoint ` + -DestinationKey $DestinationKey ` + -MigrationProvenanceContext $MigrationProvenanceContext + if ($migrationProvenanceSkipKeys.Count -gt 0) { + Write-Host "Skipping $($migrationProvenanceSkipKeys.Count) recently migrated document(s) in index '$($SourceIndex.name)'." + } + } + + if ($DifferentialMigration -and $DestinationIndexExisted) { + $destinationKeyField = Get-AISearchKeyField -IndexDefinition $DestinationIndex + if ($destinationKeyField.name -ne $sourceKeyField.name) { + throw "Index '$($SourceIndex.name)' uses different source and destination key fields." + } + + $destinationDocumentCount = Get-AISearchDocumentCount ` + -Endpoint $DestinationEndpoint ` + -AdminKey $DestinationKey ` + -IndexName $DestinationIndex.name + $destinationKeysRead = 0 + Write-Host "Reading destination keys for differential migration: $($SourceIndex.name)" + Write-AISearchCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $indexActivity ` + -Phase "Destination keys" ` + -ProcessedCount 0 ` + -TotalCount $destinationDocumentCount ` + -CurrentOperation "Comparing destination document keys" + Get-AISearchDocuments ` + -Endpoint $DestinationEndpoint ` + -AdminKey $DestinationKey ` + -IndexDefinition $DestinationIndex ` + -KeysOnly ` + -TotalCount $destinationDocumentCount ` + -ProgressActivity $indexActivity ` + -ProgressPhase "Destination keys" | ForEach-Object { + $destinationDocument = $_ + $keyValue = [string]$destinationDocument.PSObject.Properties[$destinationKeyField.name].Value + [void]$existingKeys.Add($keyValue) + $destinationKeysRead++ + if (Test-AISearchProgressCheckpoint ` + -ProcessedCount $destinationKeysRead ` + -TotalCount $destinationDocumentCount) { + Write-AISearchCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $indexActivity ` + -Phase "Destination keys" ` + -ProcessedCount $destinationKeysRead ` + -TotalCount $destinationDocumentCount ` + -CurrentOperation "Comparing destination document keys" + } + } + } + + $nonRetrievableFields = @($SourceIndex.fields | Where-Object { $_.retrievable -eq $false }) + if ($nonRetrievableFields.Count -gt 0) { + $fieldNames = $nonRetrievableFields.name -join ", " + Write-Warning "Index '$($SourceIndex.name)' has non-retrievable fields that cannot be copied: $fieldNames" + } + + $action = "upload" + $batch = [System.Collections.Generic.List[object]]::new() + $pendingBatches = [System.Collections.Generic.List[object]]::new() + $batchBytes = 12 + $pendingDocumentCount = 0 + $lastProcessedKey = $lastCommittedKey + $nextBatchSequence = $batchCount + 1 + + Write-AISearchCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $indexActivity ` + -Phase "Source documents" ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Copied: $copiedCount | Skipped: $skippedCount | Batches: $batchCount | Buffered: 0" + + Get-AISearchDocuments ` + -Endpoint $SourceEndpoint ` + -AdminKey $SourceKey ` + -IndexDefinition $SourceIndex ` + -TotalCount $sourceDocumentCount ` + -ProgressActivity $indexActivity ` + -ProgressPhase "Source documents" ` + -StartAfterKey $resumeAfterKey ` + -InitialProcessedCount $processedCount | ForEach-Object { + $sourceDocument = $_ + $keyValue = [string]$sourceDocument.PSObject.Properties[$sourceKeyField.name].Value + $sourceProvenanceSkip = Test-MigrationProvenanceSkip ` + -Provenance (Get-AISearchMigrationProvenance -Document $sourceDocument) ` + -Context $MigrationProvenanceContext + $documentAlreadyExists = ( + $sourceProvenanceSkip -or + ($DifferentialMigration -and $existingKeys.Contains($keyValue)) -or + $migrationProvenanceSkipKeys.Contains($keyValue) + ) + $writeDocument = $null + $documentBytes = 0 + if (-not $documentAlreadyExists) { + $writeDocument = ConvertTo-AISearchWriteDocument ` + -Document $sourceDocument ` + -Action $action ` + -MigrationProvenanceContext $MigrationProvenanceContext + $documentJson = $writeDocument | ConvertTo-Json -Depth 100 -Compress + $documentBytes = [System.Text.Encoding]::UTF8.GetByteCount($documentJson) + if ($documentBytes -gt $MaxBatchBytes) { + throw "Document '$keyValue' in index '$($SourceIndex.name)' exceeds the configured maximum batch payload size." + } + } + + $windowFlushed = $false + if ( + -not $documentAlreadyExists -and + $batch.Count -gt 0 -and + ($batchBytes + $documentBytes + 1) -gt $MaxBatchBytes + ) { + $pendingBatches.Add((New-AISearchDocumentBatchWorkItem ` + -Documents $batch ` + -Sequence $nextBatchSequence)) + $nextBatchSequence++ + $pendingDocumentCount += $batch.Count + $batch.Clear() + $batchBytes = 12 + if ($pendingBatches.Count -ge $MaxConcurrentBatches) { + $completedBatchCount = $pendingBatches.Count + $copiedCount += Send-AISearchDocumentBatchWindow ` + -WorkItems $pendingBatches ` + -Endpoint $DestinationEndpoint ` + -AdminKey $DestinationKey ` + -IndexName $SourceIndex.name + $batchCount += $completedBatchCount + $pendingBatches.Clear() + $pendingDocumentCount = 0 + $lastCommittedKey = $lastProcessedKey + Update-AISearchDocumentCheckpoint ` + -StateContext $StateContext ` + -ResourceName $ResourceName ` + -KeyField $sourceKeyField ` + -SourceDocumentCount $sourceDocumentCount ` + -LastCommittedKey $lastCommittedKey ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -BatchCount $batchCount + $windowFlushed = $true + } + } + + $processedCount++ + $lastProcessedKey = $keyValue + if ($documentAlreadyExists) { + $skippedCount++ + } + else { + $batch.Add($writeDocument) + $batchBytes += $documentBytes + 1 + } + + if ($batch.Count -ge $BatchSize) { + $pendingBatches.Add((New-AISearchDocumentBatchWorkItem ` + -Documents $batch ` + -Sequence $nextBatchSequence)) + $nextBatchSequence++ + $pendingDocumentCount += $batch.Count + $batch.Clear() + $batchBytes = 12 + } + + if ($pendingBatches.Count -ge $MaxConcurrentBatches) { + $completedBatchCount = $pendingBatches.Count + $copiedCount += Send-AISearchDocumentBatchWindow ` + -WorkItems $pendingBatches ` + -Endpoint $DestinationEndpoint ` + -AdminKey $DestinationKey ` + -IndexName $SourceIndex.name + $batchCount += $completedBatchCount + $pendingBatches.Clear() + $pendingDocumentCount = 0 + $lastCommittedKey = $lastProcessedKey + Update-AISearchDocumentCheckpoint ` + -StateContext $StateContext ` + -ResourceName $ResourceName ` + -KeyField $sourceKeyField ` + -SourceDocumentCount $sourceDocumentCount ` + -LastCommittedKey $lastCommittedKey ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -BatchCount $batchCount + $windowFlushed = $true + } + + $checkpointDue = Test-AISearchProgressCheckpoint ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount + if ($checkpointDue -and $batch.Count -eq 0 -and $pendingBatches.Count -eq 0) { + $lastCommittedKey = $lastProcessedKey + Update-AISearchDocumentCheckpoint ` + -StateContext $StateContext ` + -ResourceName $ResourceName ` + -KeyField $sourceKeyField ` + -SourceDocumentCount $sourceDocumentCount ` + -LastCommittedKey $lastCommittedKey ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -BatchCount $batchCount + } + + if ($windowFlushed -or $checkpointDue) { + $bufferedCount = $batch.Count + $pendingDocumentCount + Write-AISearchCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $indexActivity ` + -Phase "Source documents" ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Copied: $copiedCount | Skipped: $skippedCount | Batches: $batchCount | Buffered: $bufferedCount" + } + } + + if ($batch.Count -gt 0) { + $pendingBatches.Add((New-AISearchDocumentBatchWorkItem ` + -Documents $batch ` + -Sequence $nextBatchSequence)) + $pendingDocumentCount += $batch.Count + $batch.Clear() + $batchBytes = 12 + } + + if ($pendingBatches.Count -gt 0) { + $completedBatchCount = $pendingBatches.Count + $copiedCount += Send-AISearchDocumentBatchWindow ` + -WorkItems $pendingBatches ` + -Endpoint $DestinationEndpoint ` + -AdminKey $DestinationKey ` + -IndexName $SourceIndex.name + $batchCount += $completedBatchCount + $pendingBatches.Clear() + $pendingDocumentCount = 0 + } + + $lastCommittedKey = $lastProcessedKey + Update-AISearchDocumentCheckpoint ` + -StateContext $StateContext ` + -ResourceName $ResourceName ` + -KeyField $sourceKeyField ` + -SourceDocumentCount $sourceDocumentCount ` + -LastCommittedKey $lastCommittedKey ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -BatchCount $batchCount + + Write-AISearchCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $indexActivity ` + -Phase "Source documents" ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Copied: $copiedCount | Skipped: $skippedCount | Batches: $batchCount | Buffered: 0" + Write-AISearchProgress -Id 1 -ParentId 0 -Activity $indexActivity -Completed + Write-Host "Completed index '$($SourceIndex.name)': copied=$copiedCount, skipped=$skippedCount" + return [pscustomobject]@{ + CopiedCount = $copiedCount + SkippedCount = $skippedCount + ProcessedCount = $processedCount + TotalCount = $sourceDocumentCount + BatchCount = $batchCount + } +} + +$migrationState = $null +$activeMigrationResource = $null + +try { + $SourceSearchService = $SourceSearchService.Trim() + $SourceResourceGroup = $SourceResourceGroup.Trim() + $SourceSubscriptionId = $SourceSubscriptionId.Trim() + $DestinationSearchService = $DestinationSearchService.Trim() + $DestinationResourceGroup = $DestinationResourceGroup.Trim() + $DestinationSubscriptionId = $DestinationSubscriptionId.Trim() + $SearchDnsSuffix = $SearchDnsSuffix.Trim().TrimEnd(".") + $ApiVersion = $ApiVersion.Trim() + $ManagementApiVersion = $ManagementApiVersion.Trim() + $SourceAdminKey = $SourceAdminKey.Trim() + $DestinationAdminKey = $DestinationAdminKey.Trim() + if ([string]::IsNullOrWhiteSpace($StateFilePath)) { + $StateFilePath = Join-Path $PSScriptRoot "Migration-AISearch.state.json" + } + + foreach ($serviceName in @($SourceSearchService, $DestinationSearchService)) { + if ($serviceName -notmatch '^(?!.*--)[a-z0-9](?:[a-z0-9-]{0,58}[a-z0-9])$') { + throw "AI Search service name '$serviceName' must contain 2-60 lowercase letters, numbers, or single hyphens." + } + } + if ([string]::IsNullOrWhiteSpace($SearchDnsSuffix)) { + throw "SearchDnsSuffix cannot be empty." + } + + $sourceEndpoint = "https://$SourceSearchService.$SearchDnsSuffix" + $destinationEndpoint = "https://$DestinationSearchService.$SearchDnsSuffix" + $parsedEndpoint = $null + foreach ($endpoint in @($sourceEndpoint, $destinationEndpoint)) { + if (-not [Uri]::TryCreate($endpoint, [UriKind]::Absolute, [ref]$parsedEndpoint)) { + throw "AI Search endpoint '$endpoint' is not a valid absolute URI." + } + } + + Write-AISearchProgress ` + -Id 0 ` + -Activity "AI Search migration" ` + -Status "Connecting and discovering source resources" ` + -CurrentOperation "Resolving credentials" ` + -PercentComplete -1 + + if ([string]::Equals($sourceEndpoint, $destinationEndpoint, [System.StringComparison]::OrdinalIgnoreCase)) { + throw "Source and destination AI Search services must be different." + } + + $migrationMode = if ($DifferentialMigration) { "differential" } else { "full" } + $stateConfiguration = [ordered]@{ + sourceService = $SourceSearchService + sourceResourceGroup = $SourceResourceGroup + sourceSubscriptionId = $SourceSubscriptionId + destinationService = $DestinationSearchService + destinationResourceGroup = $DestinationResourceGroup + destinationSubscriptionId = $DestinationSubscriptionId + mode = $migrationMode + skipMigratedWithinHours = $SkipMigratedWithinHours + searchApiVersion = $ApiVersion + managementApiVersion = $ManagementApiVersion + searchDnsSuffix = $SearchDnsSuffix + } + $migrationState = Initialize-MigrationState ` + -MigrationType "ai_search" ` + -StateFilePath $StateFilePath ` + -Configuration $stateConfiguration ` + -MigrationId $MigrationId ` + -Reset:$ResetState + $migrationProvenanceContext = New-MigrationProvenanceContext ` + -MigrationId ([string]$migrationState.Data["migrationId"]) ` + -MigratedAtUtc ([string]$migrationState.Data["migrationStartedUtc"]) ` + -SkipMigratedWithinHours $SkipMigratedWithinHours + Write-Host "Migration state: $($migrationState.Path)" + Write-Host "Migration ID: $($migrationProvenanceContext.MigrationId)" + + if ([string]::IsNullOrWhiteSpace($SourceAdminKey) -or [string]::IsNullOrWhiteSpace($DestinationAdminKey)) { + foreach ($requiredCommand in @("Connect-AzAccount", "Set-AzContext", "Invoke-AzRestMethod")) { + if (-not (Get-Command $requiredCommand -ErrorAction SilentlyContinue)) { + throw "The Az.Accounts command '$requiredCommand' is required when admin keys are not supplied." + } + } + + Connect-AzAccount | Out-Null + } + + $resolvedSourceKey = Get-AISearchAdminKey ` + -ProvidedKey $SourceAdminKey ` + -SubscriptionId $SourceSubscriptionId ` + -ResourceGroupName $SourceResourceGroup ` + -ServiceName $SourceSearchService + $resolvedDestinationKey = Get-AISearchAdminKey ` + -ProvidedKey $DestinationAdminKey ` + -SubscriptionId $DestinationSubscriptionId ` + -ResourceGroupName $DestinationResourceGroup ` + -ServiceName $DestinationSearchService + + Write-Host "Starting $migrationMode AI Search migration. Destination-only indexes and documents will not be deleted." + Write-Host "Document batch concurrency: $MaxConcurrentBatches" + + $synonymMapResource = "synonymmaps" + if (Test-MigrationResourceCompleted ` + -Context $migrationState ` + -ResourceName $synonymMapResource) { + Write-Host "Skipping completed synonym maps from migration state." + } + else { + Start-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $synonymMapResource + $activeMigrationResource = $synonymMapResource + Copy-AISearchSynonymMaps ` + -SourceEndpoint $sourceEndpoint ` + -SourceKey $resolvedSourceKey ` + -DestinationEndpoint $destinationEndpoint ` + -DestinationKey $resolvedDestinationKey + Complete-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $synonymMapResource + $activeMigrationResource = $null + } + + $sourceIndexesUri = New-AISearchUri -Endpoint $sourceEndpoint -RelativePath "indexes" + $destinationIndexesUri = New-AISearchUri -Endpoint $destinationEndpoint -RelativePath "indexes" + $sourceIndexes = @((Invoke-AISearchRequest -Method "GET" -Uri $sourceIndexesUri -AdminKey $resolvedSourceKey).value) + $destinationIndexes = @((Invoke-AISearchRequest -Method "GET" -Uri $destinationIndexesUri -AdminKey $resolvedDestinationKey).value) + $destinationIndexesByName = @{} + foreach ($destinationIndex in $destinationIndexes) { + $destinationIndexesByName[$destinationIndex.name] = $destinationIndex + } + + $completedIndexCount = 0 + $totalCopiedCount = 0 + $totalSkippedCount = 0 + $indexNumber = 0 + Write-AISearchCountProgress ` + -Id 0 ` + -Activity "AI Search migration" ` + -Phase "Indexes" ` + -ProcessedCount 0 ` + -TotalCount $sourceIndexes.Count ` + -CurrentOperation "Ready to migrate $($sourceIndexes.Count) index(es)" + + foreach ($sourceIndex in $sourceIndexes) { + $indexNumber++ + $resourceName = "index:$($sourceIndex.name)" + Write-AISearchCountProgress ` + -Id 0 ` + -Activity "AI Search migration" ` + -Phase "Indexes" ` + -ProcessedCount $completedIndexCount ` + -TotalCount $sourceIndexes.Count ` + -CurrentOperation "Current index: $indexNumber/$($sourceIndexes.Count) - $($sourceIndex.name)" + + if (Test-MigrationResourceCompleted ` + -Context $migrationState ` + -ResourceName $resourceName) { + $checkpoint = Get-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $resourceName + $completedIndexCount++ + $totalCopiedCount += [long]$checkpoint["result"]["CopiedCount"] + $totalSkippedCount += [long]$checkpoint["result"]["SkippedCount"] + Write-Host "Skipping completed index from migration state: $($sourceIndex.name)" + Write-AISearchCountProgress ` + -Id 0 ` + -Activity "AI Search migration" ` + -Phase "Indexes" ` + -ProcessedCount $completedIndexCount ` + -TotalCount $sourceIndexes.Count ` + -CurrentOperation "Documents copied: $totalCopiedCount | Skipped: $totalSkippedCount" + continue + } + + Start-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $resourceName + $activeMigrationResource = $resourceName + $destinationIndexExisted = $destinationIndexesByName.ContainsKey($sourceIndex.name) + $destinationIndex = if ($destinationIndexExisted) { + $destinationIndexesByName[$sourceIndex.name] + } + else { + $null + } + + if (-not $destinationIndexExisted -or -not $DifferentialMigration) { + $writableIndex = ConvertTo-AISearchWritableDefinition -Definition $sourceIndex + $provenanceFieldsAdded = Update-AISearchMigrationProvenanceFields ` + -IndexDefinition $writableIndex + $updateIndexDefinition = $true + } + else { + $writableIndex = ConvertTo-AISearchWritableDefinition -Definition $destinationIndex + $provenanceFieldsAdded = Update-AISearchMigrationProvenanceFields ` + -IndexDefinition $writableIndex + $updateIndexDefinition = $provenanceFieldsAdded -gt 0 + } + + if ($updateIndexDefinition) { + $encodedIndexName = [Uri]::EscapeDataString($sourceIndex.name) + $destinationIndexUri = New-AISearchUri ` + -Endpoint $destinationEndpoint ` + -RelativePath "indexes/$encodedIndexName" + Invoke-AISearchRequest ` + -Method "PUT" ` + -Uri $destinationIndexUri ` + -AdminKey $resolvedDestinationKey ` + -Body $writableIndex | Out-Null + Write-Host "Migrated index definition: $($sourceIndex.name) (migration provenance fields added: $provenanceFieldsAdded)" + } + else { + Write-Host "Keeping existing destination index definition: $($sourceIndex.name)" + } + $destinationIndex = $writableIndex + + $indexResult = Copy-AISearchIndexDocuments ` + -SourceIndex $sourceIndex ` + -DestinationIndex $destinationIndex ` + -DestinationIndexExisted $destinationIndexExisted ` + -SourceEndpoint $sourceEndpoint ` + -SourceKey $resolvedSourceKey ` + -DestinationEndpoint $destinationEndpoint ` + -DestinationKey $resolvedDestinationKey ` + -StateContext $migrationState ` + -ResourceName $resourceName ` + -MigrationProvenanceContext $migrationProvenanceContext + + Complete-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $resourceName ` + -Result $indexResult + $activeMigrationResource = $null + $completedIndexCount++ + $totalCopiedCount += $indexResult.CopiedCount + $totalSkippedCount += $indexResult.SkippedCount + Write-AISearchCountProgress ` + -Id 0 ` + -Activity "AI Search migration" ` + -Phase "Indexes" ` + -ProcessedCount $completedIndexCount ` + -TotalCount $sourceIndexes.Count ` + -CurrentOperation "Documents copied: $totalCopiedCount | Skipped: $totalSkippedCount" + } + + Complete-MigrationState ` + -Context $migrationState ` + -Summary ([ordered]@{ + IndexCount = $completedIndexCount + CopiedCount = $totalCopiedCount + SkippedCount = $totalSkippedCount + }) + Write-Host "AI Search migration completed successfully. Indexes processed: $($sourceIndexes.Count), documents copied: $totalCopiedCount, documents skipped: $totalSkippedCount" +} +catch { + $migrationErrorMessage = $_.Exception.Message + if ($null -ne $migrationState) { + try { + Set-MigrationStateFailure ` + -Context $migrationState ` + -ResourceName $activeMigrationResource ` + -ErrorMessage $migrationErrorMessage + } + catch { + Write-Warning "Failed to update migration state after an error: $($_.Exception.Message)" + } + } + Write-Error "AI Search migration failed: $migrationErrorMessage" -ErrorAction Continue + throw +} +finally { + Write-AISearchProgress -Id 1 -ParentId 0 -Activity "Index documents" -Completed + Write-AISearchProgress -Id 0 -Activity "AI Search migration" -Completed +} \ No newline at end of file diff --git a/scripts/Migration-Cosmos.ps1 b/scripts/Migration-Cosmos.ps1 new file mode 100644 index 000000000..6295f21d5 --- /dev/null +++ b/scripts/Migration-Cosmos.ps1 @@ -0,0 +1,2342 @@ +# Migration-Cosmos.ps1 +#Requires -Version 7.0 + +[CmdletBinding()] +param( + [string]$SourceCosmosAccount = "", + + [string]$SourceResourceGroup = " ", + + [string]$SourceSubscriptionId = "", + + [string]$SourceDatabaseName = "SimpleChat", + + [string]$SourcePrimaryKey = "", + + [string]$DestinationCosmosAccount = "", + + [string]$DestinationResourceGroup = "", + + [string]$DestinationSubscriptionId = "", + + [string]$DestinationDatabaseName = "SimpleChat", + + [string]$DestinationPrimaryKey = "", + + [AllowEmptyCollection()] + [string[]]$Containers = @(), + + [bool]$DifferentialMigration = $true, + + [string]$MigrationId = "", + + [ValidateRange(0, 8760)] + [int]$SkipMigratedWithinHours = 24, + + [bool]$ShowProgress = $true, + + [ValidateRange(1, 100000)] + [int]$ProgressUpdateInterval = 100, + + [ValidateRange(1, 1000)] + [int]$PageSize = 100, + + [ValidateRange(1, 64)] + [int]$MaxConcurrentDocuments = 60, + + [ValidateRange(1, 10)] + [int]$MaxRetryCount = 5, + + [ValidateRange(0, 20)] + [int]$MaxThrottleRecoveryPauses = 5, + + [ValidateRange(1, 3600)] + [int]$ThrottleRecoveryPauseSeconds = 60, + + [string]$CosmosApiVersion = "2020-07-15", + + [string]$ManagementApiVersion = "2025-04-15", + + [string]$CosmosDnsSuffix = "documents.azure.com", + + [string]$StateFilePath = "", + + [switch]$ResetState +) + +$ErrorActionPreference = "Stop" +. (Join-Path $PSScriptRoot "Migration-State.ps1") +. (Join-Path $PSScriptRoot "Migration-Provenance.ps1") +$adminSettingsContainerName = "settings" +$adminSettingsDocumentId = "app_settings" +$excludedMigrationContainerNames = [string[]]@("file_processing") + +function Get-CosmosMigrationProgressPercent { + param( + [long]$ProcessedCount, + [long]$TotalCount + ) + + if ($TotalCount -le 0) { + return 100 + } + + return [int][Math]::Min( + 100, + [Math]::Floor(([double]$ProcessedCount / [double]$TotalCount) * 100) + ) +} + +function Write-CosmosMigrationProgress { + param( + [int]$Id, + [int]$ParentId = -1, + [string]$Activity, + [string]$Status = "", + [string]$CurrentOperation = "", + [int]$PercentComplete = -1, + [switch]$Completed + ) + + if (-not $ShowProgress) { + return + } + + $progressParameters = @{ + Id = $Id + ParentId = $ParentId + Activity = $Activity + } + + if ($Completed) { + $progressParameters.Completed = $true + } + else { + $progressParameters.Status = $Status + $progressParameters.PercentComplete = $PercentComplete + if (-not [string]::IsNullOrWhiteSpace($CurrentOperation)) { + $progressParameters.CurrentOperation = $CurrentOperation + } + } + + Write-Progress @progressParameters +} + +function Write-CosmosMigrationCountProgress { + param( + [int]$Id, + [int]$ParentId = -1, + [string]$Activity, + [string]$Phase, + [long]$ProcessedCount, + [long]$TotalCount, + [string]$CurrentOperation = "" + ) + + $percentComplete = Get-CosmosMigrationProgressPercent ` + -ProcessedCount $ProcessedCount ` + -TotalCount $TotalCount + $remainingCount = [Math]::Max(0, $TotalCount - $ProcessedCount) + $status = "$Phase`: $ProcessedCount/$TotalCount | Remaining: $remainingCount | $percentComplete%" + + Write-CosmosMigrationProgress ` + -Id $Id ` + -ParentId $ParentId ` + -Activity $Activity ` + -Status $status ` + -CurrentOperation $CurrentOperation ` + -PercentComplete $percentComplete +} + +function Test-CosmosMigrationProgressCheckpoint { + param( + [long]$ProcessedCount, + [long]$TotalCount = -1 + ) + + return ( + $ProcessedCount -eq 1 -or + ($TotalCount -gt 0 -and $ProcessedCount -eq $TotalCount) -or + ($ProcessedCount % $ProgressUpdateInterval) -eq 0 + ) +} + +function Get-CosmosProgressDocumentLabel { + param( + [AllowNull()] + [object]$DocumentId, + [int]$MaximumLength = 96 + ) + + $label = [string]$DocumentId + if ([string]::IsNullOrWhiteSpace($label)) { + return "" + } + if ($label.Length -le $MaximumLength) { + return $label + } + return "$($label.Substring(0, $MaximumLength - 3))..." +} + +function New-CosmosSkippedDocumentRecord { + param( + [string]$ContainerName, + [long]$Sequence, + [AllowNull()] + [object]$DocumentId, + [ValidateSet("Preparation", "Write")] + [string]$Stage, + [AllowNull()] + [object]$Attempt = $null, + [AllowNull()] + [object]$StatusCode = $null, + [string]$Reason + ) + + $normalizedDocumentId = [string]$DocumentId + if ([string]::IsNullOrWhiteSpace($normalizedDocumentId)) { + $normalizedDocumentId = "" + } + $normalizedReason = ([string]$Reason).Trim() + if ([string]::IsNullOrWhiteSpace($normalizedReason)) { + $normalizedReason = "No error details were returned." + } + if ($normalizedReason.Length -gt 2000) { + $normalizedReason = $normalizedReason.Substring(0, 2000) + } + + return [pscustomobject][ordered]@{ + ContainerName = $ContainerName + Sequence = $Sequence + DocumentId = $normalizedDocumentId + Stage = $Stage + Attempt = $Attempt + StatusCode = $StatusCode + Reason = $normalizedReason + RecordedUtc = [DateTime]::UtcNow.ToString( + "o", + [System.Globalization.CultureInfo]::InvariantCulture + ) + } +} + +function Test-CosmosDocumentWriteFailureCanBeSkipped { + param( + [int]$StatusCode + ) + + return $StatusCode -in @(400, 409, 413, 422) +} + +function Update-CosmosDocumentSkipCheckpoint { + param( + [AllowNull()] + [object]$Context, + [string]$ResourceName, + [long]$ProcessedCount, + [long]$CopiedCount, + [long]$SkippedCount, + [long]$RetryCount, + [object[]]$SkippedDocuments + ) + + if ($null -eq $Context) { + return + } + + Update-MigrationResourceCheckpoint ` + -Context $Context ` + -ResourceName $ResourceName ` + -Progress ([ordered]@{ + ProcessedCount = $ProcessedCount + CopiedCount = $CopiedCount + SkippedCount = $SkippedCount + RetryCount = $RetryCount + ErrorSkippedCount = $SkippedDocuments.Count + SkippedDocuments = @($SkippedDocuments) + }) +} + +function Get-CosmosRecordedErrorSkippedCount { + param( + [AllowNull()] + [object]$Context + ) + + if ($null -eq $Context) { + return 0 + } + + $recordedCount = 0L + foreach ($checkpoint in $Context.Data["resources"].Values) { + $countSource = if ($checkpoint["status"] -eq "completed") { + $checkpoint["result"] + } + else { + $checkpoint["progress"] + } + if ( + $countSource -is [System.Collections.IDictionary] -and + $countSource.Contains("ErrorSkippedCount") + ) { + $recordedCount += [long]$countSource["ErrorSkippedCount"] + } + } + return $recordedCount +} + +function Get-NormalizedCosmosContainerSelection { + param( + [AllowNull()] + [string[]]$ContainerNames + ) + + $normalizedNames = [System.Collections.Generic.List[string]]::new() + $seenNames = [System.Collections.Generic.HashSet[string]]::new( + [System.StringComparer]::Ordinal + ) + foreach ($containerName in @($ContainerNames)) { + $normalizedName = ([string]$containerName).Trim() + if ([string]::IsNullOrWhiteSpace($normalizedName)) { + throw "Containers cannot contain an empty container name." + } + if ( + $normalizedName.Length -gt 255 -or + $normalizedName.IndexOfAny([char[]]'\/\?#') -ge 0 + ) { + throw "Cosmos DB container name '$normalizedName' contains an unsupported character or exceeds 255 characters." + } + if (-not $seenNames.Add($normalizedName)) { + throw "Cosmos DB container '$normalizedName' is listed more than once." + } + $normalizedNames.Add($normalizedName) + } + return $normalizedNames.ToArray() +} + +function Select-CosmosMigrationContainers { + param( + [object[]]$SourceContainers, + [string[]]$RequestedContainerNames, + [string]$DatabaseName + ) + + $migratableSourceContainers = @( + $SourceContainers | Where-Object { + $excludedMigrationContainerNames -cnotcontains [string]$_.id + } + ) + $effectiveRequestedContainerNames = @( + $RequestedContainerNames | Where-Object { + $excludedMigrationContainerNames -cnotcontains $_ + } + ) + $skippedContainerNames = if ($RequestedContainerNames.Count -eq 0) { + @( + $SourceContainers | Where-Object { + $excludedMigrationContainerNames -ccontains [string]$_.id + } | ForEach-Object { [string]$_.id } + ) + } + else { + @( + $RequestedContainerNames | Where-Object { + $excludedMigrationContainerNames -ccontains $_ + } + ) + } + if ($skippedContainerNames.Count -gt 0) { + Write-Host "Skipping log-only Cosmos DB container(s): $($skippedContainerNames -join ', ')" + } + + if ($RequestedContainerNames.Count -eq 0) { + return $migratableSourceContainers + } + if ($effectiveRequestedContainerNames.Count -eq 0) { + return @() + } + + $sourceContainersByName = [System.Collections.Generic.Dictionary[string, object]]::new( + [System.StringComparer]::Ordinal + ) + foreach ($sourceContainer in $migratableSourceContainers) { + $sourceContainersByName[[string]$sourceContainer.id] = $sourceContainer + } + + $selectedContainers = [System.Collections.Generic.List[object]]::new() + $missingContainerNames = [System.Collections.Generic.List[string]]::new() + foreach ($containerName in $effectiveRequestedContainerNames) { + if ($sourceContainersByName.ContainsKey($containerName)) { + $selectedContainers.Add($sourceContainersByName[$containerName]) + } + else { + $missingContainerNames.Add($containerName) + } + } + if ($missingContainerNames.Count -gt 0) { + throw "Requested Cosmos DB container(s) were not found in source database '$DatabaseName': $($missingContainerNames -join ', ')." + } + return $selectedContainers.ToArray() +} + +function Test-CosmosMigrationConfiguration { + $configuredValues = [ordered]@{ + SourceCosmosAccount = $SourceCosmosAccount + SourceResourceGroup = $SourceResourceGroup + SourceSubscriptionId = $SourceSubscriptionId + SourceDatabaseName = $SourceDatabaseName + DestinationCosmosAccount = $DestinationCosmosAccount + DestinationResourceGroup = $DestinationResourceGroup + DestinationSubscriptionId = $DestinationSubscriptionId + DestinationDatabaseName = $DestinationDatabaseName + } + foreach ($entry in $configuredValues.GetEnumerator()) { + if ( + [string]::IsNullOrWhiteSpace($entry.Value) -or + $entry.Value -match '^<[^>]+>$' + ) { + throw "Set '$($entry.Key)' in the script or supply it as a parameter." + } + } + + foreach ($accountName in @($SourceCosmosAccount, $DestinationCosmosAccount)) { + if ($accountName -notmatch '^[a-z0-9](?:[a-z0-9-]{1,42}[a-z0-9])$') { + throw "Cosmos DB account name '$accountName' must contain 3-44 lowercase letters, numbers, or hyphens and cannot start or end with a hyphen." + } + } + + foreach ($subscriptionId in @($SourceSubscriptionId, $DestinationSubscriptionId)) { + $parsedSubscriptionId = [Guid]::Empty + if (-not [Guid]::TryParse($subscriptionId, [ref]$parsedSubscriptionId)) { + throw "Subscription ID '$subscriptionId' must be a valid GUID." + } + } + + foreach ($databaseName in @($SourceDatabaseName, $DestinationDatabaseName)) { + if ($databaseName.Length -gt 255 -or $databaseName.IndexOfAny([char[]]'\/\?#') -ge 0) { + throw "Cosmos DB database name '$databaseName' contains an unsupported character or exceeds 255 characters." + } + } + + if ( + [string]::Equals( + $SourceCosmosAccount, + $DestinationCosmosAccount, + [System.StringComparison]::OrdinalIgnoreCase + ) -and + [string]::Equals( + $SourceDatabaseName, + $DestinationDatabaseName, + [System.StringComparison]::Ordinal + ) + ) { + throw "Source and destination Cosmos DB account/database pairs must be different." + } +} + +function Get-CosmosAccountPrimaryKey { + param( + [string]$ProvidedKey, + [string]$SubscriptionId, + [string]$ResourceGroupName, + [string]$AccountName + ) + + if (-not [string]::IsNullOrWhiteSpace($ProvidedKey)) { + return $ProvidedKey.Trim() + } + + Set-AzContext -SubscriptionId $SubscriptionId | Out-Null + $encodedResourceGroupName = [Uri]::EscapeDataString($ResourceGroupName) + $encodedAccountName = [Uri]::EscapeDataString($AccountName) + $keyPath = "/subscriptions/$SubscriptionId/resourceGroups/$encodedResourceGroupName/providers/Microsoft.DocumentDB/databaseAccounts/$encodedAccountName/listKeys?api-version=$ManagementApiVersion" + $keyResponse = Invoke-AzRestMethod -Method "POST" -Path $keyPath + + if ($keyResponse.StatusCode -ne 200) { + throw "Primary key lookup failed for Cosmos DB account '$AccountName' with HTTP $($keyResponse.StatusCode)." + } + + $keyPair = $keyResponse.Content | ConvertFrom-Json + if ([string]::IsNullOrWhiteSpace($keyPair.primaryMasterKey)) { + throw "The primary key for Cosmos DB account '$AccountName' could not be resolved." + } + + return $keyPair.primaryMasterKey +} + +function New-CosmosAuthorizationHeader { + param( + [string]$Method, + [string]$ResourceType, + [string]$ResourceLink, + [string]$RequestDate, + [string]$PrimaryKey + ) + + try { + $keyBytes = [Convert]::FromBase64String($PrimaryKey) + } + catch { + throw "A supplied Cosmos DB primary key is not valid Base64." + } + + $payload = "{0}`n{1}`n{2}`n{3}`n`n" -f ` + $Method.ToLowerInvariant(), ` + $ResourceType.ToLowerInvariant(), ` + $ResourceLink, ` + $RequestDate.ToLowerInvariant() + $hmac = [System.Security.Cryptography.HMACSHA256]::new($keyBytes) + try { + $hash = $hmac.ComputeHash([System.Text.Encoding]::UTF8.GetBytes($payload)) + } + finally { + $hmac.Dispose() + } + + $signature = [Convert]::ToBase64String($hash) + return [Uri]::EscapeDataString("type=master&ver=1.0&sig=$signature") +} + +function ConvertTo-CosmosRequestPath { + param( + [string]$ResourcePath + ) + + $encodedSegments = foreach ($segment in $ResourcePath.Split('/')) { + [Uri]::EscapeDataString($segment) + } + return $encodedSegments -join '/' +} + +function Get-CosmosResponseHeaderValue { + param( + [AllowNull()] + [object]$Headers, + [string]$Name + ) + + if ($null -eq $Headers) { + return "" + } + + if ($Headers -is [System.Collections.IDictionary]) { + foreach ($headerName in $Headers.Keys) { + if ([string]::Equals([string]$headerName, $Name, [System.StringComparison]::OrdinalIgnoreCase)) { + return @($Headers[$headerName]) -join ',' + } + } + return "" + } + + foreach ($header in $Headers) { + if ([string]::Equals([string]$header.Key, $Name, [System.StringComparison]::OrdinalIgnoreCase)) { + return @($header.Value) -join ',' + } + } + return "" +} + +function Get-CosmosStatisticsDocumentCount { + param( + [AllowNull()] + [object]$Resource + ) + + if ($null -eq $Resource) { + return -1 + } + + $statisticsProperty = $Resource.PSObject.Properties["statistics"] + if ($null -eq $statisticsProperty -or $null -eq $statisticsProperty.Value) { + return -1 + } + + $documentCounts = @( + @($statisticsProperty.Value) | Where-Object { $null -ne $_ } | ForEach-Object { + $documentCountProperty = $_.PSObject.Properties["documentCount"] + if ($null -ne $documentCountProperty -and $null -ne $documentCountProperty.Value) { + [long]$documentCountProperty.Value + } + } + ) + if ($documentCounts.Count -eq 0) { + return -1 + } + return [long](($documentCounts | Measure-Object -Sum).Sum) +} + +function Get-CosmosContainerDocumentCount { + param( + [object]$ContainerDefinition, + [string]$SubscriptionId, + [string]$ResourceGroupName, + [string]$AccountName, + [string]$DatabaseName + ) + + $embeddedCount = Get-CosmosStatisticsDocumentCount -Resource $ContainerDefinition + if ($embeddedCount -ge 0) { + return $embeddedCount + } + + $encodedResourceGroupName = [Uri]::EscapeDataString($ResourceGroupName) + $encodedAccountName = [Uri]::EscapeDataString($AccountName) + $encodedDatabaseName = [Uri]::EscapeDataString($DatabaseName) + $encodedContainerName = [Uri]::EscapeDataString([string]$ContainerDefinition.id) + $managementPath = "/subscriptions/$SubscriptionId/resourceGroups/$encodedResourceGroupName/providers/Microsoft.DocumentDB/databaseAccounts/$encodedAccountName/sqlDatabases/$encodedDatabaseName/containers/${encodedContainerName}?api-version=$ManagementApiVersion" + $managementResource = $null + + if (Get-Command "Invoke-AzRestMethod" -ErrorAction "SilentlyContinue") { + try { + $response = Invoke-AzRestMethod -Method "GET" -Path $managementPath + if ($response.StatusCode -eq 200) { + $managementResource = ($response.Content | ConvertFrom-Json -Depth 100).properties.resource + } + } + catch { + Write-Verbose "PowerShell ARM document-count lookup failed for '$($ContainerDefinition.id)': $($_.Exception.Message)" + } + } + + if ($null -eq $managementResource -and (Get-Command "az" -ErrorAction "SilentlyContinue")) { + try { + $managementUrl = "https://management.azure.com$managementPath" + $azOutput = & az rest ` + --method get ` + --url $managementUrl ` + --output json 2>$null + if ($LASTEXITCODE -eq 0 -and -not [string]::IsNullOrWhiteSpace($azOutput)) { + $managementResource = ($azOutput | ConvertFrom-Json -Depth 100).properties.resource + } + } + catch { + Write-Verbose "Azure CLI document-count lookup failed for '$($ContainerDefinition.id)': $($_.Exception.Message)" + } + } + + $managementCount = Get-CosmosStatisticsDocumentCount -Resource $managementResource + if ($managementCount -lt 0) { + Write-Verbose "Document total is unavailable for container '$($ContainerDefinition.id)'." + } + return $managementCount +} + +function Write-CosmosDocumentProgress { + param( + [string]$Activity, + [long]$ProcessedCount, + [long]$TotalCount, + [string]$CurrentOperation + ) + + if ($TotalCount -ge 0) { + Write-CosmosMigrationCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $Activity ` + -Phase "Source documents" ` + -ProcessedCount $ProcessedCount ` + -TotalCount $TotalCount ` + -CurrentOperation $CurrentOperation + return + } + + Write-CosmosMigrationProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $Activity ` + -Status "Source documents completed: $ProcessedCount | Total: unavailable" ` + -PercentComplete -1 ` + -CurrentOperation $CurrentOperation +} + +function Invoke-CosmosThrottleRecoveryPause { + param( + [string]$Activity, + [string]$ContainerName, + [int]$PauseNumber, + [int]$ThrottledDocumentCount + ) + + Write-Warning "Cosmos DB returned HTTP 429 for $ThrottledDocumentCount document(s) in container '$ContainerName' after $MaxRetryCount request attempt(s). Recovery pause $PauseNumber/$MaxThrottleRecoveryPauses begins now; retrying only the throttled document(s) in $ThrottleRecoveryPauseSeconds second(s)." + for ($remainingSeconds = $ThrottleRecoveryPauseSeconds; $remainingSeconds -gt 0; $remainingSeconds--) { + Write-CosmosMigrationProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $Activity ` + -Status "Cosmos DB throttling recovery pause $PauseNumber/$MaxThrottleRecoveryPauses" ` + -CurrentOperation "$remainingSeconds second(s) until retrying $ThrottledDocumentCount throttled document(s)" ` + -PercentComplete -1 + Start-Sleep -Seconds 1 + } +} + +function New-CosmosThrottleRecoveryExhaustedError { + param( + [string]$ContainerName, + [object[]]$ThrottledWorkItems, + [int]$PauseCount + ) + + $sampleDocument = $ThrottledWorkItems[0] + return "Cosmos DB throttling recovery exhausted in container '$ContainerName'. $($ThrottledWorkItems.Count) document(s) still returned HTTP 429 after $MaxRetryCount request attempt(s) and $PauseCount recovery pause(s) of $ThrottleRecoveryPauseSeconds second(s). Example document: '$($sampleDocument.DocumentId)'. Reduce -MaxConcurrentDocuments, increase destination RU capacity, or rerun with the same state file after capacity is available." +} + +function Invoke-CosmosRestRequest { + param( + [string]$Endpoint, + [string]$PrimaryKey, + [ValidateSet("GET", "POST", "PUT")] + [string]$Method, + [string]$ResourcePath, + [string]$ResourceType, + [string]$ResourceLink, + [AllowNull()] + [string]$Body = $null, + [hashtable]$AdditionalHeaders = @{}, + [int[]]$AllowedStatusCodes = @(200), + [switch]$PreserveJsonPropertyNames + ) + + $encodedPath = ConvertTo-CosmosRequestPath -ResourcePath $ResourcePath + $requestUri = "$($Endpoint.TrimEnd('/'))/$encodedPath" + + for ($attempt = 1; $attempt -le $MaxRetryCount; $attempt++) { + $requestDate = [DateTime]::UtcNow.ToString( + "r", + [System.Globalization.CultureInfo]::InvariantCulture + ) + $requestHeaders = @{ + "Accept" = "application/json" + "Authorization" = New-CosmosAuthorizationHeader ` + -Method $Method ` + -ResourceType $ResourceType ` + -ResourceLink $ResourceLink ` + -RequestDate $requestDate ` + -PrimaryKey $PrimaryKey + "x-ms-date" = $requestDate + "x-ms-version" = $CosmosApiVersion + } + $contentType = "" + foreach ($headerName in $AdditionalHeaders.Keys) { + if ([string]::Equals( + $headerName, + "Content-Type", + [System.StringComparison]::OrdinalIgnoreCase + )) { + $contentType = [string]$AdditionalHeaders[$headerName] + continue + } + $requestHeaders[$headerName] = $AdditionalHeaders[$headerName] + } + + $invokeParameters = @{ + Method = $Method + Uri = $requestUri + Headers = $requestHeaders + SkipHttpErrorCheck = $true + ErrorAction = "Stop" + } + if ($null -ne $Body) { + $invokeParameters.Body = $Body + } + if (-not [string]::IsNullOrWhiteSpace($contentType)) { + $invokeParameters.ContentType = $contentType + } + + try { + $response = Invoke-WebRequest @invokeParameters + } + catch { + if ($attempt -eq $MaxRetryCount) { + throw + } + + $retryDelayMilliseconds = [int]([Math]::Pow(2, $attempt - 1) * 1000) + Write-Warning "Cosmos DB request failed before receiving a response. Retrying in $retryDelayMilliseconds millisecond(s)." + Start-Sleep -Milliseconds $retryDelayMilliseconds + continue + } + + $statusCode = [int]$response.StatusCode + if ($AllowedStatusCodes -contains $statusCode) { + $responseBody = $null + $responseContent = [string]$response.Content + if (-not [string]::IsNullOrWhiteSpace($responseContent)) { + $responseBody = if ($PreserveJsonPropertyNames) { + $responseContent | + ConvertFrom-Json -AsHashtable -Depth 100 + } + else { + $responseContent | ConvertFrom-Json -Depth 100 + } + } + return [pscustomobject]@{ + StatusCode = $statusCode + Headers = $response.Headers + Body = $responseBody + Content = $responseContent + } + } + + $retryable = $statusCode -in @(408, 429, 449, 500, 503) + if ($retryable -and $attempt -lt $MaxRetryCount) { + $retryAfterHeader = Get-CosmosResponseHeaderValue ` + -Headers $response.Headers ` + -Name "x-ms-retry-after-ms" + $retryAfterMilliseconds = 0 + if (-not [int]::TryParse($retryAfterHeader, [ref]$retryAfterMilliseconds)) { + $retryAfterMilliseconds = [int]([Math]::Pow(2, $attempt - 1) * 1000) + } + $retryAfterMilliseconds = [Math]::Max(1, $retryAfterMilliseconds) + Write-Warning "Cosmos DB request returned HTTP $statusCode. Retrying in $retryAfterMilliseconds millisecond(s)." + Start-Sleep -Milliseconds $retryAfterMilliseconds + continue + } + + $errorContent = [string]$response.Content + if ($errorContent.Length -gt 1000) { + $errorContent = $errorContent.Substring(0, 1000) + } + throw "Cosmos DB request '$Method $ResourcePath' failed with HTTP $statusCode. $errorContent" + } +} + +function Get-CosmosDatabase { + param( + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName + ) + + $resourceLink = "dbs/$DatabaseName" + return Invoke-CosmosRestRequest ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -Method "GET" ` + -ResourcePath $resourceLink ` + -ResourceType "dbs" ` + -ResourceLink $resourceLink ` + -AllowedStatusCodes @(200, 404) +} + +function New-CosmosDatabaseIfMissing { + param( + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName + ) + + $databaseResponse = Get-CosmosDatabase ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -DatabaseName $DatabaseName + if ($databaseResponse.StatusCode -eq 200) { + return + } + + $body = @{ id = $DatabaseName } | ConvertTo-Json -Compress + Invoke-CosmosRestRequest ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -Method "POST" ` + -ResourcePath "dbs" ` + -ResourceType "dbs" ` + -ResourceLink "" ` + -Body $body ` + -AdditionalHeaders @{ "Content-Type" = "application/json" } ` + -AllowedStatusCodes @(201) | Out-Null + Write-Host "Created destination database: $DatabaseName" +} + +function Get-CosmosContainers { + param( + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName + ) + + $databaseLink = "dbs/$DatabaseName" + $continuation = "" + do { + $headers = @{ + "x-ms-max-item-count" = [string]$PageSize + } + if (-not [string]::IsNullOrWhiteSpace($continuation)) { + $headers["x-ms-continuation"] = $continuation + } + $response = Invoke-CosmosRestRequest ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -Method "GET" ` + -ResourcePath "$databaseLink/colls" ` + -ResourceType "colls" ` + -ResourceLink $databaseLink ` + -AdditionalHeaders $headers + foreach ($container in @($response.Body.DocumentCollections)) { + Write-Output $container + } + $continuation = Get-CosmosResponseHeaderValue ` + -Headers $response.Headers ` + -Name "x-ms-continuation" + } while (-not [string]::IsNullOrWhiteSpace($continuation)) +} + +function ConvertTo-CosmosNonNullJsonValue { + param( + [AllowNull()] + [object]$Value + ) + + if ($null -eq $Value) { + return $null + } + if ( + $Value -is [string] -or + $Value -is [ValueType] + ) { + return $Value + } + if ($Value -is [System.Collections.IDictionary]) { + $convertedDictionary = [ordered]@{} + foreach ($key in $Value.Keys) { + if ($null -ne $Value[$key]) { + $convertedDictionary[[string]$key] = ConvertTo-CosmosNonNullJsonValue ` + -Value $Value[$key] + } + } + return $convertedDictionary + } + if ($Value -is [System.Collections.IEnumerable]) { + $convertedItems = [System.Collections.Generic.List[object]]::new() + foreach ($item in $Value) { + if ($null -ne $item) { + $convertedItems.Add((ConvertTo-CosmosNonNullJsonValue -Value $item)) + } + } + return ,$convertedItems.ToArray() + } + + $convertedObject = [ordered]@{} + foreach ($property in $Value.PSObject.Properties) { + if ($null -ne $property.Value) { + $convertedObject[$property.Name] = ConvertTo-CosmosNonNullJsonValue ` + -Value $property.Value + } + } + return $convertedObject +} + +function ConvertTo-CosmosWritableResource { + param( + [object]$Resource + ) + + $partitionKey = [ordered]@{ + paths = @($Resource.partitionKey.paths) + } + if (-not [string]::IsNullOrWhiteSpace([string]$Resource.partitionKey.kind)) { + $partitionKey.kind = [string]$Resource.partitionKey.kind + } + if ($null -ne $Resource.partitionKey.version) { + $partitionKey.version = $Resource.partitionKey.version + } + + $writableResource = [ordered]@{ + id = [string]$Resource.id + partitionKey = $partitionKey + } + $optionalWritableProperties = @( + "indexingPolicy", + "defaultTtl", + "uniqueKeyPolicy", + "conflictResolutionPolicy", + "analyticalStorageTtl", + "computedProperties", + "clientEncryptionPolicy", + "geospatialConfig", + "changeFeedPolicy", + "vectorEmbeddingPolicy", + "fullTextPolicy" + ) + foreach ($propertyName in $optionalWritableProperties) { + $property = $Resource.PSObject.Properties[$propertyName] + if ($null -ne $property -and $null -ne $property.Value) { + $writableResource[$propertyName] = ConvertTo-CosmosNonNullJsonValue ` + -Value $property.Value + } + } + return [pscustomobject]$writableResource +} + +function Get-CosmosPartitionKeyDefinitionJson { + param( + [object]$Container + ) + + $partitionKey = [ordered]@{ + paths = @($Container.partitionKey.paths) + kind = [string]$Container.partitionKey.kind + version = $Container.partitionKey.version + } + return $partitionKey | ConvertTo-Json -Depth 10 -Compress +} + +function Assert-CosmosContainerCompatibility { + param( + [object]$SourceContainer, + [object]$DestinationContainer + ) + + $sourcePartitionKey = Get-CosmosPartitionKeyDefinitionJson -Container $SourceContainer + $destinationPartitionKey = Get-CosmosPartitionKeyDefinitionJson -Container $DestinationContainer + if ($sourcePartitionKey -ne $destinationPartitionKey) { + throw "Container '$($SourceContainer.id)' has incompatible source and destination partition key definitions. Cosmos DB partition keys cannot be changed in place." + } +} + +function Sync-CosmosContainerDefinition { + param( + [object]$SourceContainer, + [AllowNull()] + [object]$DestinationContainer, + [string]$DestinationEndpoint, + [string]$DestinationKey, + [string]$DestinationDatabase + ) + + $databaseLink = "dbs/$DestinationDatabase" + $containerLink = "$databaseLink/colls/$($SourceContainer.id)" + $writableContainer = ConvertTo-CosmosWritableResource -Resource $SourceContainer + $containerBody = $writableContainer | ConvertTo-Json -Depth 100 -Compress + Write-Verbose "Writable container definition for '$($SourceContainer.id)': $containerBody" + + if ($null -eq $DestinationContainer) { + Invoke-CosmosRestRequest ` + -Endpoint $DestinationEndpoint ` + -PrimaryKey $DestinationKey ` + -Method "POST" ` + -ResourcePath "$databaseLink/colls" ` + -ResourceType "colls" ` + -ResourceLink $databaseLink ` + -Body $containerBody ` + -AdditionalHeaders @{ "Content-Type" = "application/json" } ` + -AllowedStatusCodes @(201) | Out-Null + Write-Host "Created destination container: $($SourceContainer.id)" + return + } + + Assert-CosmosContainerCompatibility ` + -SourceContainer $SourceContainer ` + -DestinationContainer $DestinationContainer + if ($DifferentialMigration) { + Write-Host "Keeping existing destination container definition: $($SourceContainer.id)" + return + } + + Invoke-CosmosRestRequest ` + -Endpoint $DestinationEndpoint ` + -PrimaryKey $DestinationKey ` + -Method "PUT" ` + -ResourcePath $containerLink ` + -ResourceType "colls" ` + -ResourceLink $containerLink ` + -Body $containerBody ` + -AdditionalHeaders @{ "Content-Type" = "application/json" } ` + -AllowedStatusCodes @(200) | Out-Null + Write-Host "Updated destination container definition: $($SourceContainer.id)" +} + +function Invoke-CosmosDocumentFeedPage { + param( + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName, + [string]$ContainerName, + [string]$Continuation = "" + ) + + $containerLink = "dbs/$DatabaseName/colls/$ContainerName" + $headers = @{ + "x-ms-max-item-count" = [string]$PageSize + } + if (-not [string]::IsNullOrWhiteSpace($Continuation)) { + $headers["x-ms-continuation"] = $Continuation + } + + return Invoke-CosmosRestRequest ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -Method "GET" ` + -ResourcePath "$containerLink/docs" ` + -ResourceType "docs" ` + -ResourceLink $containerLink ` + -AdditionalHeaders $headers ` + -PreserveJsonPropertyNames +} + +function Get-CosmosDocuments { + param( + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName, + [string]$ContainerName, + [bool]$ExcludeAdminSettings + ) + + $continuation = "" + + do { + $response = Invoke-CosmosDocumentFeedPage ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -DatabaseName $DatabaseName ` + -ContainerName $ContainerName ` + -Continuation $continuation + foreach ($document in @($response.Body.Documents)) { + if ( + $ExcludeAdminSettings -and + [string]::Equals( + [string]$document.id, + $adminSettingsDocumentId, + [System.StringComparison]::Ordinal + ) + ) { + continue + } + Write-Output $document + } + $continuation = Get-CosmosResponseHeaderValue ` + -Headers $response.Headers ` + -Name "x-ms-continuation" + } while (-not [string]::IsNullOrWhiteSpace($continuation)) +} + +function Test-CosmosDocumentMigrationProvenanceSkip { + param( + [object]$Document, + [object]$MigrationProvenanceContext + ) + + return Test-MigrationProvenanceSkip ` + -Provenance (Get-CosmosMigrationProvenance -Document $Document) ` + -Context $MigrationProvenanceContext +} + +function Get-CosmosPropertyPathValue { + param( + [object]$Document, + [string]$Path + ) + + $current = $Document + $segments = $Path.TrimStart('/').Split('/') + foreach ($encodedSegment in $segments) { + $segment = $encodedSegment.Replace('~1', '/').Replace('~0', '~') + if ($current -is [System.Collections.IDictionary]) { + if (-not $current.Contains($segment)) { + return [pscustomobject]@{ Exists = $false; Value = $null } + } + $current = $current[$segment] + continue + } + + $property = $current.PSObject.Properties[$segment] + if ($null -eq $property) { + return [pscustomobject]@{ Exists = $false; Value = $null } + } + $current = $property.Value + } + + return [pscustomobject]@{ Exists = $true; Value = $current } +} + +function Get-CosmosPartitionKeyHeader { + param( + [object]$Document, + [object]$ContainerDefinition + ) + + $partitionKeyValues = [System.Collections.Generic.List[object]]::new() + foreach ($path in @($ContainerDefinition.partitionKey.paths)) { + $pathValue = Get-CosmosPropertyPathValue -Document $Document -Path $path + if ($pathValue.Exists) { + $partitionKeyValues.Add($pathValue.Value) + } + else { + $partitionKeyValues.Add([pscustomobject]@{}) + } + } + return ConvertTo-Json -InputObject $partitionKeyValues.ToArray() -Depth 100 -Compress +} + +function Get-CosmosDocumentIdentity { + param( + [object]$Document, + [object]$ContainerDefinition + ) + + $documentId = [string](Get-MigrationProvenancePropertyValue ` + -Source $Document ` + -Name "id") + if ([string]::IsNullOrWhiteSpace($documentId)) { + return "" + } + + $partitionKeyHeader = Get-CosmosPartitionKeyHeader ` + -Document $Document ` + -ContainerDefinition $ContainerDefinition + $identityText = "$documentId`n$partitionKeyHeader" + return [Convert]::ToBase64String( + [System.Text.Encoding]::UTF8.GetBytes($identityText) + ) +} + +function Get-CosmosMigrationProvenanceSkipDocumentIdentities { + param( + [object]$ContainerDefinition, + [string]$DestinationEndpoint, + [string]$DestinationKey, + [string]$DestinationDatabase, + [object]$MigrationProvenanceContext + ) + + $migrationId = [string]$MigrationProvenanceContext.MigrationId + $queryParameters = [System.Collections.Generic.List[object]]::new() + $queryParameters.Add([ordered]@{ name = "@status"; value = "succeeded" }) + $queryParameters.Add([ordered]@{ name = "@migrationId"; value = $migrationId }) + $provenanceFilter = "c.simplechatMigration.migrationId = @migrationId" + if ([int]$MigrationProvenanceContext.SkipMigratedWithinHours -gt 0) { + $cutoffUtc = [DateTimeOffset]::UtcNow.AddHours( + -[int]$MigrationProvenanceContext.SkipMigratedWithinHours + ).ToString("o", [System.Globalization.CultureInfo]::InvariantCulture) + $queryParameters.Add([ordered]@{ name = "@cutoffUtc"; value = $cutoffUtc }) + $provenanceFilter += " OR c.simplechatMigration.migratedAtUtc >= @cutoffUtc" + } + $query = [ordered]@{ + query = "SELECT * FROM c WHERE c.simplechatMigration.status = @status AND ($provenanceFilter)" + parameters = $queryParameters.ToArray() + } | ConvertTo-Json -Depth 20 -Compress + + $containerLink = "dbs/$DestinationDatabase/colls/$($ContainerDefinition.id)" + $continuation = "" + $documentIdentities = [System.Collections.Generic.HashSet[string]]::new( + [System.StringComparer]::Ordinal + ) + do { + $headers = @{ + "Content-Type" = "application/query+json" + "x-ms-documentdb-isquery" = "true" + "x-ms-documentdb-query-enablecrosspartition" = "true" + "x-ms-max-item-count" = [string]$PageSize + } + if (-not [string]::IsNullOrWhiteSpace($continuation)) { + $headers["x-ms-continuation"] = $continuation + } + $response = Invoke-CosmosRestRequest ` + -Endpoint $DestinationEndpoint ` + -PrimaryKey $DestinationKey ` + -Method "POST" ` + -ResourcePath "$containerLink/docs" ` + -ResourceType "docs" ` + -ResourceLink $containerLink ` + -Body $query ` + -AdditionalHeaders $headers ` + -PreserveJsonPropertyNames + foreach ($destinationDocument in @($response.Body.Documents)) { + $provenance = Get-CosmosMigrationProvenance -Document $destinationDocument + if (Test-MigrationProvenanceSkip ` + -Provenance $provenance ` + -Context $MigrationProvenanceContext) { + $documentIdentity = Get-CosmosDocumentIdentity ` + -Document $destinationDocument ` + -ContainerDefinition $ContainerDefinition + if (-not [string]::IsNullOrWhiteSpace($documentIdentity)) { + [void]$documentIdentities.Add($documentIdentity) + } + } + } + $continuation = Get-CosmosResponseHeaderValue ` + -Headers $response.Headers ` + -Name "x-ms-continuation" + } while (-not [string]::IsNullOrWhiteSpace($continuation)) + + return ,$documentIdentities +} + +function ConvertTo-CosmosWritableDocument { + param( + [object]$Document, + [AllowNull()] + [object]$MigrationProvenanceContext = $null + ) + + $copy = $Document | + ConvertTo-Json -Depth 100 | + ConvertFrom-Json -AsHashtable -Depth 100 + foreach ($propertyName in @("_rid", "_self", "_etag", "_attachments", "_ts")) { + [void]$copy.Remove($propertyName) + } + if ($null -ne $MigrationProvenanceContext) { + Add-CosmosMigrationProvenance ` + -Document $copy ` + -Context $MigrationProvenanceContext + } + return $copy +} + +function Send-CosmosDocument { + param( + [object]$Document, + [object]$ContainerDefinition, + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName, + [AllowNull()] + [object]$MigrationProvenanceContext = $null + ) + + $writableDocument = ConvertTo-CosmosWritableDocument ` + -Document $Document ` + -MigrationProvenanceContext $MigrationProvenanceContext + if ([string]::IsNullOrWhiteSpace([string]$writableDocument.id)) { + throw "Container '$($ContainerDefinition.id)' contains a document without a valid id." + } + + $containerLink = "dbs/$DatabaseName/colls/$($ContainerDefinition.id)" + $headers = @{ + "Content-Type" = "application/json" + "x-ms-documentdb-partitionkey" = Get-CosmosPartitionKeyHeader ` + -Document $writableDocument ` + -ContainerDefinition $ContainerDefinition + } + $allowedStatusCodes = @(201) + if (-not $DifferentialMigration) { + $headers["x-ms-documentdb-is-upsert"] = "true" + $allowedStatusCodes = @(200, 201) + } + + $response = Invoke-CosmosRestRequest ` + -Endpoint $Endpoint ` + -PrimaryKey $PrimaryKey ` + -Method "POST" ` + -ResourcePath "$containerLink/docs" ` + -ResourceType "docs" ` + -ResourceLink $containerLink ` + -Body ($writableDocument | ConvertTo-Json -Depth 100 -Compress) ` + -AdditionalHeaders $headers ` + -AllowedStatusCodes ($allowedStatusCodes + @(409)) ` + -PreserveJsonPropertyNames + if ($response.StatusCode -eq 409) { + if (-not $DifferentialMigration) { + throw "Full migration received an unexpected HTTP 409 conflict for document '$($writableDocument.id)' in container '$($ContainerDefinition.id)'." + } + return "Skipped" + } + return "Copied" +} + +function New-CosmosDocumentWriteWorkItem { + param( + [object]$Document, + [object]$ContainerDefinition, + [long]$Sequence, + [AllowNull()] + [object]$MigrationProvenanceContext = $null + ) + + $writableDocument = ConvertTo-CosmosWritableDocument ` + -Document $Document ` + -MigrationProvenanceContext $MigrationProvenanceContext + if ([string]::IsNullOrWhiteSpace([string]$writableDocument.id)) { + throw "Container '$($ContainerDefinition.id)' contains a document without a valid id." + } + + return [pscustomobject]@{ + Sequence = $Sequence + DocumentId = [string]$writableDocument.id + DocumentLabel = Get-CosmosProgressDocumentLabel -DocumentId $writableDocument.id + Body = $writableDocument | ConvertTo-Json -Depth 100 -Compress + PartitionKeyHeader = Get-CosmosPartitionKeyHeader ` + -Document $writableDocument ` + -ContainerDefinition $ContainerDefinition + } +} + +function Invoke-CosmosParallelDocumentWrites { + param( + [Parameter(ValueFromPipeline)] + [object]$WorkItem, + [string]$Endpoint, + [string]$PrimaryKey, + [string]$DatabaseName, + [string]$ContainerName + ) + + begin { + $workItems = [System.Collections.Generic.List[object]]::new() + } + process { + $workItems.Add($WorkItem) + } + end { + if ($workItems.Count -eq 0) { + return + } + + $resourceLink = "dbs/$DatabaseName/colls/$ContainerName" + $encodedDatabaseName = [Uri]::EscapeDataString($DatabaseName) + $encodedContainerName = [Uri]::EscapeDataString($ContainerName) + $requestUri = "$($Endpoint.TrimEnd('/'))/dbs/$encodedDatabaseName/colls/$encodedContainerName/docs" + $keyBytes = [Convert]::FromBase64String($PrimaryKey) + $apiVersion = $CosmosApiVersion + $retryLimit = $MaxRetryCount + $useUpsert = -not $DifferentialMigration + $invokeWebRequestCommand = Get-Command "Invoke-WebRequest" -ErrorAction "Stop" + $invokeWebRequestOverrideDefinition = if ($invokeWebRequestCommand.CommandType -eq "Function") { + $invokeWebRequestCommand.Definition + } + else { + "" + } + + $workItems | ForEach-Object -Parallel { + $currentWorkItem = $_ + if (-not [string]::IsNullOrWhiteSpace($using:invokeWebRequestOverrideDefinition)) { + Set-Item ` + -Path "Function:Invoke-WebRequest" ` + -Value ([scriptblock]::Create($using:invokeWebRequestOverrideDefinition)) + } + [pscustomobject]@{ + EventType = "Started" + Sequence = $currentWorkItem.Sequence + DocumentId = $currentWorkItem.DocumentId + DocumentLabel = $currentWorkItem.DocumentLabel + Attempt = 1 + StatusCode = 0 + RetryDelayMilliseconds = 0 + Result = "" + ErrorMessage = "" + } + + for ($attempt = 1; $attempt -le $using:retryLimit; $attempt++) { + $requestDate = [DateTime]::UtcNow.ToString( + "r", + [System.Globalization.CultureInfo]::InvariantCulture + ) + $signaturePayload = "post`ndocs`n$using:resourceLink`n$($requestDate.ToLowerInvariant())`n`n" + $hmac = [System.Security.Cryptography.HMACSHA256]::new($using:keyBytes) + try { + $hash = $hmac.ComputeHash( + [System.Text.Encoding]::UTF8.GetBytes($signaturePayload) + ) + } + finally { + $hmac.Dispose() + } + $signature = [Convert]::ToBase64String($hash) + $authorization = [Uri]::EscapeDataString( + "type=master&ver=1.0&sig=$signature" + ) + $headers = @{ + "Accept" = "application/json" + "Authorization" = $authorization + "x-ms-date" = $requestDate + "x-ms-version" = $using:apiVersion + "x-ms-documentdb-partitionkey" = $currentWorkItem.PartitionKeyHeader + } + if ($using:useUpsert) { + $headers["x-ms-documentdb-is-upsert"] = "true" + } + + try { + $response = Invoke-WebRequest ` + -Method "POST" ` + -Uri $using:requestUri ` + -Headers $headers ` + -ContentType "application/json" ` + -Body $currentWorkItem.Body ` + -SkipHttpErrorCheck ` + -ErrorAction "Stop" + } + catch { + if ($attempt -eq $using:retryLimit) { + [pscustomobject]@{ + EventType = "Failed" + Sequence = $currentWorkItem.Sequence + DocumentId = $currentWorkItem.DocumentId + DocumentLabel = $currentWorkItem.DocumentLabel + Attempt = $attempt + StatusCode = 0 + RetryDelayMilliseconds = 0 + Result = "" + ErrorMessage = $_.Exception.Message + } + return + } + + $retryDelayMilliseconds = [int]([Math]::Pow(2, $attempt - 1) * 1000) + [pscustomobject]@{ + EventType = "Retrying" + Sequence = $currentWorkItem.Sequence + DocumentId = $currentWorkItem.DocumentId + DocumentLabel = $currentWorkItem.DocumentLabel + Attempt = $attempt + StatusCode = 0 + RetryDelayMilliseconds = $retryDelayMilliseconds + Result = "" + ErrorMessage = $_.Exception.Message + } + Start-Sleep -Milliseconds $retryDelayMilliseconds + continue + } + + $statusCode = [int]$response.StatusCode + $successfulWrite = $statusCode -in @(200, 201) + $differentialConflict = -not $using:useUpsert -and $statusCode -eq 409 + if ($successfulWrite -or $differentialConflict) { + [pscustomobject]@{ + EventType = "Completed" + Sequence = $currentWorkItem.Sequence + DocumentId = $currentWorkItem.DocumentId + DocumentLabel = $currentWorkItem.DocumentLabel + Attempt = $attempt + StatusCode = $statusCode + RetryDelayMilliseconds = 0 + Result = if ($differentialConflict) { "Skipped" } else { "Copied" } + ErrorMessage = "" + } + return + } + + $retryable = $statusCode -in @(408, 429, 449, 500, 503) + if ($retryable -and $attempt -lt $using:retryLimit) { + $retryAfterHeader = @($response.Headers["x-ms-retry-after-ms"]) -join "," + $retryDelayMilliseconds = 0 + if (-not [int]::TryParse($retryAfterHeader, [ref]$retryDelayMilliseconds)) { + $retryDelayMilliseconds = [int]([Math]::Pow(2, $attempt - 1) * 1000) + } + $retryDelayMilliseconds = [Math]::Max(1, $retryDelayMilliseconds) + [pscustomobject]@{ + EventType = "Retrying" + Sequence = $currentWorkItem.Sequence + DocumentId = $currentWorkItem.DocumentId + DocumentLabel = $currentWorkItem.DocumentLabel + Attempt = $attempt + StatusCode = $statusCode + RetryDelayMilliseconds = $retryDelayMilliseconds + Result = "" + ErrorMessage = "" + } + Start-Sleep -Milliseconds $retryDelayMilliseconds + continue + } + + $errorContent = [string]$response.Content + if ($errorContent.Length -gt 1000) { + $errorContent = $errorContent.Substring(0, 1000) + } + [pscustomobject]@{ + EventType = "Failed" + Sequence = $currentWorkItem.Sequence + DocumentId = $currentWorkItem.DocumentId + DocumentLabel = $currentWorkItem.DocumentLabel + Attempt = $attempt + StatusCode = $statusCode + RetryDelayMilliseconds = 0 + Result = "" + ErrorMessage = "HTTP $statusCode. $errorContent" + } + return + } + } -ThrottleLimit $MaxConcurrentDocuments + } +} + +function Invoke-CosmosParallelDocumentBatch { + param( + [object[]]$WorkItems, + [string]$Activity, + [string]$DestinationEndpoint, + [string]$DestinationKey, + [string]$DestinationDatabase, + [string]$ContainerName, + [long]$SourceDocumentCount, + [long]$InitialProcessedCount, + [long]$InitialCopiedCount, + [long]$InitialSkippedCount, + [long]$InitialRetryCount, + [int]$InitialThrottleRecoveryPauseCount = 0, + [object[]]$InitialSkippedDocuments = @(), + [AllowNull()] + [object]$MigrationState, + [string]$MigrationResourceName + ) + + $processedCount = $InitialProcessedCount + $copiedCount = $InitialCopiedCount + $skippedCount = $InitialSkippedCount + $retryCount = $InitialRetryCount + $throttleRecoveryPauseCount = $InitialThrottleRecoveryPauseCount + $inFlightCount = 0 + $skippedDocuments = [System.Collections.Generic.List[object]]::new() + foreach ($skippedDocument in $InitialSkippedDocuments) { + $skippedDocuments.Add($skippedDocument) + } + + $workItemsBySequence = @{} + foreach ($workItem in $WorkItems) { + $workItemsBySequence[[long]$workItem.Sequence] = $workItem + } + $pendingWorkItems = @($WorkItems) + + while ($pendingWorkItems.Count -gt 0) { + $throttledWorkItems = [System.Collections.Generic.List[object]]::new() + $pendingWorkItems | + Invoke-CosmosParallelDocumentWrites ` + -Endpoint $DestinationEndpoint ` + -PrimaryKey $DestinationKey ` + -DatabaseName $DestinationDatabase ` + -ContainerName $ContainerName | + ForEach-Object { + $writeEvent = $_ + $updateDocumentProgress = Test-CosmosMigrationProgressCheckpoint ` + -ProcessedCount $writeEvent.Sequence + switch ($writeEvent.EventType) { + "Started" { + $inFlightCount++ + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $Activity ` + -ProcessedCount $processedCount ` + -TotalCount $SourceDocumentCount ` + -CurrentOperation "Copying document $($writeEvent.Sequence)`: '$($writeEvent.DocumentLabel)' | Copied: $copiedCount | Skipped: $skippedCount" + } + } + "Retrying" { + $retryCount++ + } + "Completed" { + $inFlightCount = [Math]::Max(0, $inFlightCount - 1) + $processedCount++ + if ($writeEvent.Result -eq "Skipped") { + $skippedCount++ + } + else { + $copiedCount++ + } + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $Activity ` + -ProcessedCount $processedCount ` + -TotalCount $SourceDocumentCount ` + -CurrentOperation "$($writeEvent.Result) document $($writeEvent.Sequence)`: '$($writeEvent.DocumentLabel)' | Copied: $copiedCount | Skipped: $skippedCount" + } + } + "Failed" { + if ($writeEvent.StatusCode -eq 429) { + $inFlightCount = [Math]::Max(0, $inFlightCount - 1) + $throttledWorkItems.Add( + $workItemsBySequence[[long]$writeEvent.Sequence] + ) + return + } + if (-not (Test-CosmosDocumentWriteFailureCanBeSkipped ` + -StatusCode $writeEvent.StatusCode)) { + throw "Document '$($writeEvent.DocumentId)' in container '$ContainerName' failed after $($writeEvent.Attempt) attempt(s): $($writeEvent.ErrorMessage)" + } + $inFlightCount = [Math]::Max(0, $inFlightCount - 1) + $processedCount++ + $skippedCount++ + $skippedDocument = New-CosmosSkippedDocumentRecord ` + -ContainerName $ContainerName ` + -Sequence $writeEvent.Sequence ` + -DocumentId $writeEvent.DocumentId ` + -Stage "Write" ` + -Attempt $writeEvent.Attempt ` + -StatusCode $writeEvent.StatusCode ` + -Reason $writeEvent.ErrorMessage + $skippedDocuments.Add($skippedDocument) + Update-CosmosDocumentSkipCheckpoint ` + -Context $MigrationState ` + -ResourceName $MigrationResourceName ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -RetryCount $retryCount ` + -SkippedDocuments $skippedDocuments.ToArray() + Write-Warning "Skipped document '$($writeEvent.DocumentLabel)' in container '$ContainerName' after a write error. Details were recorded in migration state." + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $Activity ` + -ProcessedCount $processedCount ` + -TotalCount $SourceDocumentCount ` + -CurrentOperation "Skipped document $($writeEvent.Sequence) after a write error: '$($writeEvent.DocumentLabel)' | Copied: $copiedCount | Skipped: $skippedCount" + } + } + default { + throw "Unexpected parallel document event '$($writeEvent.EventType)'." + } + } + } + + if ($throttledWorkItems.Count -eq 0) { + break + } + if ($throttleRecoveryPauseCount -ge $MaxThrottleRecoveryPauses) { + throw (New-CosmosThrottleRecoveryExhaustedError ` + -ContainerName $ContainerName ` + -ThrottledWorkItems $throttledWorkItems.ToArray() ` + -PauseCount $throttleRecoveryPauseCount) + } + + $throttleRecoveryPauseCount++ + Invoke-CosmosThrottleRecoveryPause ` + -Activity $Activity ` + -ContainerName $ContainerName ` + -PauseNumber $throttleRecoveryPauseCount ` + -ThrottledDocumentCount $throttledWorkItems.Count + $pendingWorkItems = $throttledWorkItems.ToArray() + } + + return [pscustomobject]@{ + ProcessedCount = $processedCount + CopiedCount = $copiedCount + SkippedCount = $skippedCount + RetryCount = $retryCount + ThrottleRecoveryPauseCount = $throttleRecoveryPauseCount + ErrorSkippedCount = $skippedDocuments.Count + SkippedDocuments = $skippedDocuments.ToArray() + } +} + +function Copy-CosmosContainerDocuments { + param( + [object]$SourceContainer, + [string]$SourceEndpoint, + [string]$SourceKey, + [string]$SourceDatabase, + [string]$DestinationEndpoint, + [string]$DestinationKey, + [string]$DestinationDatabase, + [AllowNull()] + [object]$MigrationState, + [string]$MigrationResourceName, + [object]$MigrationProvenanceContext + ) + + $excludeAdminSettings = [string]::Equals( + $SourceContainer.id, + $adminSettingsContainerName, + [System.StringComparison]::Ordinal + ) + $activity = "Container '$($SourceContainer.id)' documents" + $copiedCount = 0 + $skippedCount = 0 + $processedCount = 0 + $retryCount = 0 + $throttleRecoveryPauseCount = 0 + $skippedDocuments = [System.Collections.Generic.List[object]]::new() + $sourceDocumentCount = Get-CosmosContainerDocumentCount ` + -ContainerDefinition $SourceContainer ` + -SubscriptionId $SourceSubscriptionId ` + -ResourceGroupName $SourceResourceGroup ` + -AccountName $SourceCosmosAccount ` + -DatabaseName $SourceDatabase + if ($excludeAdminSettings -and $sourceDocumentCount -gt 0) { + $sourceDocumentCount-- + } + $migrationProvenanceSkipDocumentIdentities = [System.Collections.Generic.HashSet[string]]::new( + [System.StringComparer]::Ordinal + ) + if (-not $DifferentialMigration) { + $migrationProvenanceSkipDocumentIdentities = Get-CosmosMigrationProvenanceSkipDocumentIdentities ` + -ContainerDefinition $SourceContainer ` + -DestinationEndpoint $DestinationEndpoint ` + -DestinationKey $DestinationKey ` + -DestinationDatabase $DestinationDatabase ` + -MigrationProvenanceContext $MigrationProvenanceContext + if ($migrationProvenanceSkipDocumentIdentities.Count -gt 0) { + Write-Host "Skipping $($migrationProvenanceSkipDocumentIdentities.Count) recently migrated document(s) in container '$($SourceContainer.id)'." + } + } + + Update-CosmosDocumentSkipCheckpoint ` + -Context $MigrationState ` + -ResourceName $MigrationResourceName ` + -ProcessedCount 0 ` + -CopiedCount 0 ` + -SkippedCount 0 ` + -RetryCount 0 ` + -SkippedDocuments @() + + Write-CosmosDocumentProgress ` + -Activity $activity ` + -ProcessedCount 0 ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Copied: 0 | Skipped: 0" + + if ($MaxConcurrentDocuments -eq 1) { + Get-CosmosDocuments ` + -Endpoint $SourceEndpoint ` + -PrimaryKey $SourceKey ` + -DatabaseName $SourceDatabase ` + -ContainerName $SourceContainer.id ` + -ExcludeAdminSettings $excludeAdminSettings | + ForEach-Object { + $document = $_ + $nextDocumentNumber = $processedCount + 1 + $documentLabel = Get-CosmosProgressDocumentLabel -DocumentId $document.id + $updateDocumentProgress = Test-CosmosMigrationProgressCheckpoint ` + -ProcessedCount $nextDocumentNumber ` + -TotalCount $sourceDocumentCount + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $activity ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Copying document $nextDocumentNumber`: '$documentLabel' | Copied: $copiedCount | Skipped: $skippedCount" + } + + $documentIdentity = Get-CosmosDocumentIdentity ` + -Document $document ` + -ContainerDefinition $SourceContainer + $sourceProvenanceSkip = Test-CosmosDocumentMigrationProvenanceSkip ` + -Document $document ` + -MigrationProvenanceContext $MigrationProvenanceContext + $destinationProvenanceSkip = ( + -not [string]::IsNullOrWhiteSpace($documentIdentity) -and + $migrationProvenanceSkipDocumentIdentities.Contains($documentIdentity) + ) + if ($sourceProvenanceSkip -or $destinationProvenanceSkip) { + $processedCount++ + $skippedCount++ + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $activity ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Skipped previously migrated document $processedCount`: '$documentLabel' | Copied: $copiedCount | Skipped: $skippedCount" + } + return + } + + try { + [void](New-CosmosDocumentWriteWorkItem ` + -Document $document ` + -ContainerDefinition $SourceContainer ` + -Sequence $nextDocumentNumber ` + -MigrationProvenanceContext $MigrationProvenanceContext) + } + catch { + $skippedDocument = New-CosmosSkippedDocumentRecord ` + -ContainerName $SourceContainer.id ` + -Sequence $nextDocumentNumber ` + -DocumentId $document.id ` + -Stage "Preparation" ` + -Reason $_.Exception.Message + $skippedDocuments.Add($skippedDocument) + $processedCount++ + $skippedCount++ + Update-CosmosDocumentSkipCheckpoint ` + -Context $MigrationState ` + -ResourceName $MigrationResourceName ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -RetryCount $retryCount ` + -SkippedDocuments $skippedDocuments.ToArray() + Write-Warning "Skipped document '$documentLabel' in container '$($SourceContainer.id)' after a preparation error. Details were recorded in migration state." + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $activity ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "Skipped document $processedCount after a preparation error: '$documentLabel' | Copied: $copiedCount | Skipped: $skippedCount" + } + return + } + + $result = "" + do { + try { + $result = Send-CosmosDocument ` + -Document $document ` + -ContainerDefinition $SourceContainer ` + -Endpoint $DestinationEndpoint ` + -PrimaryKey $DestinationKey ` + -DatabaseName $DestinationDatabase ` + -MigrationProvenanceContext $MigrationProvenanceContext + } + catch { + $errorMessage = $_.Exception.Message + $statusCode = $null + if ($errorMessage -match '\bHTTP (?[1-5][0-9]{2})\b') { + $statusCode = [int]$Matches.StatusCode + } + if ($statusCode -eq 429) { + if ($throttleRecoveryPauseCount -ge $MaxThrottleRecoveryPauses) { + throw (New-CosmosThrottleRecoveryExhaustedError ` + -ContainerName $SourceContainer.id ` + -ThrottledWorkItems @([pscustomobject]@{ + DocumentId = $document.id + }) ` + -PauseCount $throttleRecoveryPauseCount) + } + $throttleRecoveryPauseCount++ + Invoke-CosmosThrottleRecoveryPause ` + -Activity $activity ` + -ContainerName $SourceContainer.id ` + -PauseNumber $throttleRecoveryPauseCount ` + -ThrottledDocumentCount 1 + continue + } + if ( + $null -eq $statusCode -or + -not (Test-CosmosDocumentWriteFailureCanBeSkipped ` + -StatusCode $statusCode) + ) { + throw + } + $skippedDocument = New-CosmosSkippedDocumentRecord ` + -ContainerName $SourceContainer.id ` + -Sequence $nextDocumentNumber ` + -DocumentId $document.id ` + -Stage "Write" ` + -Attempt 1 ` + -StatusCode $statusCode ` + -Reason $errorMessage + $skippedDocuments.Add($skippedDocument) + Update-CosmosDocumentSkipCheckpoint ` + -Context $MigrationState ` + -ResourceName $MigrationResourceName ` + -ProcessedCount ($processedCount + 1) ` + -CopiedCount $copiedCount ` + -SkippedCount ($skippedCount + 1) ` + -RetryCount $retryCount ` + -SkippedDocuments $skippedDocuments.ToArray() + Write-Warning "Skipped document '$documentLabel' in container '$($SourceContainer.id)' after a write error. Details were recorded in migration state." + $result = "Skipped after a write error" + } + } while ([string]::IsNullOrWhiteSpace($result)) + + $processedCount++ + if ($result -like "Skipped*") { + $skippedCount++ + } + else { + $copiedCount++ + } + + if ($updateDocumentProgress) { + Write-CosmosDocumentProgress ` + -Activity $activity ` + -ProcessedCount $processedCount ` + -TotalCount $sourceDocumentCount ` + -CurrentOperation "$result document $processedCount`: '$documentLabel' | Copied: $copiedCount | Skipped: $skippedCount" + } + } + } + else { + $scheduledCount = 0 + $writeBatchSize = [Math]::Min( + $PageSize, + [Math]::Max($MaxConcurrentDocuments, $MaxConcurrentDocuments * 4) + ) + $workItems = [System.Collections.Generic.List[object]]::new() + + Get-CosmosDocuments ` + -Endpoint $SourceEndpoint ` + -PrimaryKey $SourceKey ` + -DatabaseName $SourceDatabase ` + -ContainerName $SourceContainer.id ` + -ExcludeAdminSettings $excludeAdminSettings | + ForEach-Object { + $document = $_ + $scheduledCount++ + $documentIdentity = Get-CosmosDocumentIdentity ` + -Document $document ` + -ContainerDefinition $SourceContainer + $sourceProvenanceSkip = Test-CosmosDocumentMigrationProvenanceSkip ` + -Document $document ` + -MigrationProvenanceContext $MigrationProvenanceContext + $destinationProvenanceSkip = ( + -not [string]::IsNullOrWhiteSpace($documentIdentity) -and + $migrationProvenanceSkipDocumentIdentities.Contains($documentIdentity) + ) + if ($sourceProvenanceSkip -or $destinationProvenanceSkip) { + $processedCount++ + $skippedCount++ + continue + } + try { + $workItems.Add((New-CosmosDocumentWriteWorkItem ` + -Document $document ` + -ContainerDefinition $SourceContainer ` + -Sequence $scheduledCount ` + -MigrationProvenanceContext $MigrationProvenanceContext)) + } + catch { + $processedCount++ + $skippedCount++ + $documentLabel = Get-CosmosProgressDocumentLabel -DocumentId $document.id + $skippedDocument = New-CosmosSkippedDocumentRecord ` + -ContainerName $SourceContainer.id ` + -Sequence $scheduledCount ` + -DocumentId $document.id ` + -Stage "Preparation" ` + -Reason $_.Exception.Message + $skippedDocuments.Add($skippedDocument) + Update-CosmosDocumentSkipCheckpoint ` + -Context $MigrationState ` + -ResourceName $MigrationResourceName ` + -ProcessedCount $processedCount ` + -CopiedCount $copiedCount ` + -SkippedCount $skippedCount ` + -RetryCount $retryCount ` + -SkippedDocuments $skippedDocuments.ToArray() + Write-Warning "Skipped document '$documentLabel' in container '$($SourceContainer.id)' after a preparation error. Details were recorded in migration state." + } + + if ($workItems.Count -ge $writeBatchSize) { + $batchResult = Invoke-CosmosParallelDocumentBatch ` + -WorkItems $workItems.ToArray() ` + -Activity $activity ` + -DestinationEndpoint $DestinationEndpoint ` + -DestinationKey $DestinationKey ` + -DestinationDatabase $DestinationDatabase ` + -ContainerName $SourceContainer.id ` + -SourceDocumentCount $sourceDocumentCount ` + -InitialProcessedCount $processedCount ` + -InitialCopiedCount $copiedCount ` + -InitialSkippedCount $skippedCount ` + -InitialRetryCount $retryCount ` + -InitialThrottleRecoveryPauseCount $throttleRecoveryPauseCount ` + -InitialSkippedDocuments $skippedDocuments.ToArray() ` + -MigrationState $MigrationState ` + -MigrationResourceName $MigrationResourceName + $processedCount = $batchResult.ProcessedCount + $copiedCount = $batchResult.CopiedCount + $skippedCount = $batchResult.SkippedCount + $retryCount = $batchResult.RetryCount + $throttleRecoveryPauseCount = $batchResult.ThrottleRecoveryPauseCount + $skippedDocuments.Clear() + foreach ($batchSkippedDocument in @($batchResult.SkippedDocuments)) { + $skippedDocuments.Add($batchSkippedDocument) + } + $workItems.Clear() + } + } + + if ($workItems.Count -gt 0) { + $batchResult = Invoke-CosmosParallelDocumentBatch ` + -WorkItems $workItems.ToArray() ` + -Activity $activity ` + -DestinationEndpoint $DestinationEndpoint ` + -DestinationKey $DestinationKey ` + -DestinationDatabase $DestinationDatabase ` + -ContainerName $SourceContainer.id ` + -SourceDocumentCount $sourceDocumentCount ` + -InitialProcessedCount $processedCount ` + -InitialCopiedCount $copiedCount ` + -InitialSkippedCount $skippedCount ` + -InitialRetryCount $retryCount ` + -InitialThrottleRecoveryPauseCount $throttleRecoveryPauseCount ` + -InitialSkippedDocuments $skippedDocuments.ToArray() ` + -MigrationState $MigrationState ` + -MigrationResourceName $MigrationResourceName + $processedCount = $batchResult.ProcessedCount + $copiedCount = $batchResult.CopiedCount + $skippedCount = $batchResult.SkippedCount + $retryCount = $batchResult.RetryCount + $throttleRecoveryPauseCount = $batchResult.ThrottleRecoveryPauseCount + $skippedDocuments.Clear() + foreach ($batchSkippedDocument in @($batchResult.SkippedDocuments)) { + $skippedDocuments.Add($batchSkippedDocument) + } + } + } + + $finalTotalCount = if ($sourceDocumentCount -ge 0) { + $sourceDocumentCount + } + else { + $processedCount + } + Write-CosmosMigrationCountProgress ` + -Id 1 ` + -ParentId 0 ` + -Activity $activity ` + -Phase "Source documents" ` + -ProcessedCount $processedCount ` + -TotalCount $finalTotalCount ` + -CurrentOperation "Copied: $copiedCount | Skipped: $skippedCount" + Write-CosmosMigrationProgress -Id 1 -ParentId 0 -Activity $activity -Completed + Write-Host "Completed container '$($SourceContainer.id)': copied=$copiedCount, skipped=$skippedCount, throttling recovery pauses=$throttleRecoveryPauseCount" + return [pscustomobject]@{ + CopiedCount = $copiedCount + SkippedCount = $skippedCount + ProcessedCount = $processedCount + TotalCount = $finalTotalCount + RetryCount = $retryCount + ThrottleRecoveryPauseCount = $throttleRecoveryPauseCount + ErrorSkippedCount = $skippedDocuments.Count + SkippedDocuments = $skippedDocuments.ToArray() + } +} + +$SourceCosmosAccount = $SourceCosmosAccount.Trim() +$SourceResourceGroup = $SourceResourceGroup.Trim() +$SourceSubscriptionId = $SourceSubscriptionId.Trim() +$SourceDatabaseName = $SourceDatabaseName.Trim() +$DestinationCosmosAccount = $DestinationCosmosAccount.Trim() +$DestinationResourceGroup = $DestinationResourceGroup.Trim() +$DestinationSubscriptionId = $DestinationSubscriptionId.Trim() +$DestinationDatabaseName = $DestinationDatabaseName.Trim() +$CosmosApiVersion = $CosmosApiVersion.Trim() +$ManagementApiVersion = $ManagementApiVersion.Trim() +$CosmosDnsSuffix = $CosmosDnsSuffix.Trim().Trim('.') +$requestedContainerNames = @(Get-NormalizedCosmosContainerSelection -ContainerNames $Containers) +if ([string]::IsNullOrWhiteSpace($StateFilePath)) { + $StateFilePath = Join-Path $PSScriptRoot "Migration-Cosmos.state.json" +} +$sourceEndpoint = "https://$SourceCosmosAccount.$CosmosDnsSuffix" +$destinationEndpoint = "https://$DestinationCosmosAccount.$CosmosDnsSuffix" +$migrationMode = if ($DifferentialMigration) { "differential" } else { "full" } +$migrationState = $null +$activeMigrationResource = $null + +try { + Test-CosmosMigrationConfiguration + + Write-CosmosMigrationProgress ` + -Id 0 ` + -Activity "Cosmos DB migration" ` + -Status "Connecting and discovering source resources" ` + -CurrentOperation "Resolving credentials" ` + -PercentComplete -1 + + if ([string]::IsNullOrWhiteSpace($SourcePrimaryKey) -or [string]::IsNullOrWhiteSpace($DestinationPrimaryKey)) { + foreach ($requiredCommand in @("Connect-AzAccount", "Set-AzContext", "Invoke-AzRestMethod")) { + if (-not (Get-Command $requiredCommand -ErrorAction SilentlyContinue)) { + throw "The Az.Accounts command '$requiredCommand' is required when primary keys are not supplied." + } + } + Connect-AzAccount | Out-Null + } + + $resolvedSourceKey = Get-CosmosAccountPrimaryKey ` + -ProvidedKey $SourcePrimaryKey ` + -SubscriptionId $SourceSubscriptionId ` + -ResourceGroupName $SourceResourceGroup ` + -AccountName $SourceCosmosAccount + $resolvedDestinationKey = Get-CosmosAccountPrimaryKey ` + -ProvidedKey $DestinationPrimaryKey ` + -SubscriptionId $DestinationSubscriptionId ` + -ResourceGroupName $DestinationResourceGroup ` + -AccountName $DestinationCosmosAccount + + $sourceDatabase = Get-CosmosDatabase ` + -Endpoint $sourceEndpoint ` + -PrimaryKey $resolvedSourceKey ` + -DatabaseName $SourceDatabaseName + if ($sourceDatabase.StatusCode -ne 200) { + throw "Source database '$SourceDatabaseName' does not exist in Cosmos DB account '$SourceCosmosAccount'." + } + + $allSourceContainers = @(Get-CosmosContainers ` + -Endpoint $sourceEndpoint ` + -PrimaryKey $resolvedSourceKey ` + -DatabaseName $SourceDatabaseName) + $sourceContainers = @(Select-CosmosMigrationContainers ` + -SourceContainers $allSourceContainers ` + -RequestedContainerNames $requestedContainerNames ` + -DatabaseName $SourceDatabaseName) + $stateContainerNames = [string[]]@($sourceContainers | ForEach-Object { [string]$_.id }) + [Array]::Sort($stateContainerNames, [System.StringComparer]::Ordinal) + $selectionMode = if ($requestedContainerNames.Count -eq 0) { "all" } else { "explicit" } + $stateConfiguration = [ordered]@{ + sourceAccount = $SourceCosmosAccount + sourceResourceGroup = $SourceResourceGroup + sourceSubscriptionId = $SourceSubscriptionId + sourceDatabase = $SourceDatabaseName + destinationAccount = $DestinationCosmosAccount + destinationResourceGroup = $DestinationResourceGroup + destinationSubscriptionId = $DestinationSubscriptionId + destinationDatabase = $DestinationDatabaseName + containers = @($stateContainerNames) + containerSelectionMode = $selectionMode + mode = $migrationMode + skipMigratedWithinHours = $SkipMigratedWithinHours + cosmosApiVersion = $CosmosApiVersion + cosmosDnsSuffix = $CosmosDnsSuffix + excludedAdminSettingsDocument = "$adminSettingsContainerName/$adminSettingsDocumentId" + } + $migrationState = Initialize-MigrationState ` + -MigrationType "cosmos" ` + -StateFilePath $StateFilePath ` + -Configuration $stateConfiguration ` + -MigrationId $MigrationId ` + -Reset:$ResetState + $migrationProvenanceContext = New-MigrationProvenanceContext ` + -MigrationId ([string]$migrationState.Data["migrationId"]) ` + -MigratedAtUtc ([string]$migrationState.Data["migrationStartedUtc"]) ` + -SkipMigratedWithinHours $SkipMigratedWithinHours + Write-Host "Migration state: $($migrationState.Path)" + Write-Host "Migration ID: $($migrationProvenanceContext.MigrationId)" + + New-CosmosDatabaseIfMissing ` + -Endpoint $destinationEndpoint ` + -PrimaryKey $resolvedDestinationKey ` + -DatabaseName $DestinationDatabaseName + + Write-Host "Starting $migrationMode Cosmos DB migration. Destination-only containers and documents will not be deleted." + Write-Host "Document write concurrency: $MaxConcurrentDocuments" + Write-Host "Preserving destination admin app settings. '$adminSettingsContainerName/$adminSettingsDocumentId' will not be migrated." + if ($selectionMode -eq "explicit") { + Write-Host "Selected containers: $($stateContainerNames -join ', ')" + } + $destinationContainers = @(Get-CosmosContainers ` + -Endpoint $destinationEndpoint ` + -PrimaryKey $resolvedDestinationKey ` + -DatabaseName $DestinationDatabaseName) + $destinationContainersByName = @{} + foreach ($destinationContainer in $destinationContainers) { + $destinationContainersByName[$destinationContainer.id] = $destinationContainer + } + + $completedContainerCount = 0 + $totalCopiedCount = 0 + $totalSkippedCount = 0 + $totalErrorSkippedCount = 0 + $totalThrottleRecoveryPauseCount = 0 + $containerNumber = 0 + Write-CosmosMigrationCountProgress ` + -Id 0 ` + -Activity "Cosmos DB migration" ` + -Phase "Containers" ` + -ProcessedCount 0 ` + -TotalCount $sourceContainers.Count ` + -CurrentOperation "Ready to migrate $($sourceContainers.Count) container(s)" + + foreach ($sourceContainer in $sourceContainers) { + $containerNumber++ + $resourceName = [string]$sourceContainer.id + Write-CosmosMigrationCountProgress ` + -Id 0 ` + -Activity "Cosmos DB migration" ` + -Phase "Containers" ` + -ProcessedCount $completedContainerCount ` + -TotalCount $sourceContainers.Count ` + -CurrentOperation "Current container: $containerNumber/$($sourceContainers.Count) - $($sourceContainer.id)" + + if (Test-MigrationResourceCompleted ` + -Context $migrationState ` + -ResourceName $resourceName) { + $checkpoint = Get-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $resourceName + $completedContainerCount++ + $totalCopiedCount += [long]$checkpoint["result"]["CopiedCount"] + $totalSkippedCount += [long]$checkpoint["result"]["SkippedCount"] + $totalErrorSkippedCount += [long]$checkpoint["result"]["ErrorSkippedCount"] + $totalThrottleRecoveryPauseCount += [long]$checkpoint["result"]["ThrottleRecoveryPauseCount"] + Write-Host "Skipping completed container from migration state: $resourceName" + Write-CosmosMigrationCountProgress ` + -Id 0 ` + -Activity "Cosmos DB migration" ` + -Phase "Containers" ` + -ProcessedCount $completedContainerCount ` + -TotalCount $sourceContainers.Count ` + -CurrentOperation "Documents copied: $totalCopiedCount | Skipped: $totalSkippedCount" + continue + } + + Start-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $resourceName + $activeMigrationResource = $resourceName + + $destinationContainer = if ($destinationContainersByName.ContainsKey($sourceContainer.id)) { + $destinationContainersByName[$sourceContainer.id] + } + else { + $null + } + Sync-CosmosContainerDefinition ` + -SourceContainer $sourceContainer ` + -DestinationContainer $destinationContainer ` + -DestinationEndpoint $destinationEndpoint ` + -DestinationKey $resolvedDestinationKey ` + -DestinationDatabase $DestinationDatabaseName + + $containerResult = Copy-CosmosContainerDocuments ` + -SourceContainer $sourceContainer ` + -SourceEndpoint $sourceEndpoint ` + -SourceKey $resolvedSourceKey ` + -SourceDatabase $SourceDatabaseName ` + -DestinationEndpoint $destinationEndpoint ` + -DestinationKey $resolvedDestinationKey ` + -DestinationDatabase $DestinationDatabaseName ` + -MigrationState $migrationState ` + -MigrationResourceName $resourceName ` + -MigrationProvenanceContext $migrationProvenanceContext + Complete-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $resourceName ` + -Result $containerResult + $activeMigrationResource = $null + $completedContainerCount++ + $totalCopiedCount += $containerResult.CopiedCount + $totalSkippedCount += $containerResult.SkippedCount + $totalErrorSkippedCount += $containerResult.ErrorSkippedCount + $totalThrottleRecoveryPauseCount += $containerResult.ThrottleRecoveryPauseCount + + Write-CosmosMigrationCountProgress ` + -Id 0 ` + -Activity "Cosmos DB migration" ` + -Phase "Containers" ` + -ProcessedCount $completedContainerCount ` + -TotalCount $sourceContainers.Count ` + -CurrentOperation "Documents copied: $totalCopiedCount | Skipped: $totalSkippedCount" + } + + Complete-MigrationState ` + -Context $migrationState ` + -Summary ([ordered]@{ + ContainerCount = $completedContainerCount + CopiedCount = $totalCopiedCount + SkippedCount = $totalSkippedCount + ErrorSkippedCount = $totalErrorSkippedCount + ThrottleRecoveryPauseCount = $totalThrottleRecoveryPauseCount + }) + Write-Host "Cosmos DB migration completed. Containers processed: $completedContainerCount, documents copied: $totalCopiedCount, documents skipped: $totalSkippedCount, throttling recovery pauses: $totalThrottleRecoveryPauseCount" + if ($totalErrorSkippedCount -gt 0) { + Write-Warning "Admin review required: $totalErrorSkippedCount document(s) were skipped because of preparation or write errors. Review each container's result.SkippedDocuments entries in '$($migrationState.Path)'." + } +} +catch { + $migrationErrorMessage = $_.Exception.Message + if ($null -ne $migrationState) { + try { + Set-MigrationStateFailure ` + -Context $migrationState ` + -ResourceName $activeMigrationResource ` + -ErrorMessage $migrationErrorMessage + } + catch { + Write-Warning "Failed to update migration state after an error: $($_.Exception.Message)" + } + } + $recordedErrorSkippedCount = Get-CosmosRecordedErrorSkippedCount ` + -Context $migrationState + if ($recordedErrorSkippedCount -gt 0) { + Write-Warning "Admin review required: $recordedErrorSkippedCount document(s) were skipped before the migration stopped. Review completed containers' result.SkippedDocuments and interrupted containers' progress.SkippedDocuments entries in '$($migrationState.Path)'." + } + Write-Error "Cosmos DB migration failed: $migrationErrorMessage" -ErrorAction Continue + throw +} +finally { + Write-CosmosMigrationProgress -Id 1 -ParentId 0 -Activity "Container documents" -Completed + Write-CosmosMigrationProgress -Id 0 -Activity "Cosmos DB migration" -Completed +} \ No newline at end of file diff --git a/scripts/Migration-Provenance.ps1 b/scripts/Migration-Provenance.ps1 new file mode 100644 index 000000000..e07353061 --- /dev/null +++ b/scripts/Migration-Provenance.ps1 @@ -0,0 +1,260 @@ +# Migration-Provenance.ps1 + +function Get-MigrationProvenancePropertyValue { + param( + [AllowNull()] + [object]$Source, + [string]$Name + ) + + if ($null -eq $Source) { + return $null + } + + if ($Source -is [System.Collections.IDictionary]) { + foreach ($key in $Source.Keys) { + if ([string]::Equals([string]$key, $Name, [System.StringComparison]::OrdinalIgnoreCase)) { + return $Source[$key] + } + } + return $null + } + + foreach ($property in $Source.PSObject.Properties) { + if ([string]::Equals($property.Name, $Name, [System.StringComparison]::OrdinalIgnoreCase)) { + return $property.Value + } + } + return $null +} + +function New-MigrationProvenanceContext { + param( + [string]$MigrationId = "", + [string]$MigratedAtUtc = "", + [ValidateRange(0, 8760)] + [int]$SkipMigratedWithinHours = 24 + ) + + $parsedMigrationId = [Guid]::Empty + if ([string]::IsNullOrWhiteSpace($MigrationId)) { + $parsedMigrationId = [Guid]::NewGuid() + } + elseif (-not [Guid]::TryParse($MigrationId, [ref]$parsedMigrationId)) { + throw "MigrationId must be a valid GUID." + } + + [DateTimeOffset]$parsedMigratedAtUtc = [DateTimeOffset]::UtcNow + if (-not [string]::IsNullOrWhiteSpace($MigratedAtUtc)) { + if (-not [DateTimeOffset]::TryParse($MigratedAtUtc, [ref]$parsedMigratedAtUtc)) { + throw "MigratedAtUtc must be a valid ISO 8601 timestamp." + } + $parsedMigratedAtUtc = $parsedMigratedAtUtc.ToUniversalTime() + } + + return [pscustomobject]@{ + MigrationId = $parsedMigrationId.ToString("D") + MigratedAtUtc = $parsedMigratedAtUtc.ToString( + "o", + [System.Globalization.CultureInfo]::InvariantCulture + ) + SkipMigratedWithinHours = $SkipMigratedWithinHours + } +} + +function New-MigrationProvenanceRecord { + param( + [object]$Context + ) + + return [ordered]@{ + migrationId = [string]$Context.MigrationId + migratedAtUtc = [string]$Context.MigratedAtUtc + status = "succeeded" + } +} + +function Test-MigrationProvenanceSkip { + param( + [AllowNull()] + [object]$Provenance, + [object]$Context + ) + + if ($null -eq $Provenance) { + return $false + } + + $status = [string](Get-MigrationProvenancePropertyValue ` + -Source $Provenance ` + -Name "status") + if (-not [string]::Equals($status, "succeeded", [System.StringComparison]::OrdinalIgnoreCase)) { + return $false + } + + $migrationId = [string](Get-MigrationProvenancePropertyValue ` + -Source $Provenance ` + -Name "migrationId") + if ([string]::Equals($migrationId, [string]$Context.MigrationId, [System.StringComparison]::OrdinalIgnoreCase)) { + return $true + } + + if ([int]$Context.SkipMigratedWithinHours -le 0) { + return $false + } + + $migratedAtUtc = [string](Get-MigrationProvenancePropertyValue ` + -Source $Provenance ` + -Name "migratedAtUtc") + [DateTimeOffset]$parsedMigratedAtUtc = [DateTimeOffset]::MinValue + if (-not [DateTimeOffset]::TryParse($migratedAtUtc, [ref]$parsedMigratedAtUtc)) { + return $false + } + + $cutoffUtc = [DateTimeOffset]::UtcNow.AddHours(-[int]$Context.SkipMigratedWithinHours) + return $parsedMigratedAtUtc.ToUniversalTime() -ge $cutoffUtc +} + +function Add-CosmosMigrationProvenance { + param( + [System.Collections.IDictionary]$Document, + [object]$Context + ) + + $Document["simplechatMigration"] = New-MigrationProvenanceRecord -Context $Context +} + +function Get-CosmosMigrationProvenance { + param( + [AllowNull()] + [object]$Document + ) + + return Get-MigrationProvenancePropertyValue ` + -Source $Document ` + -Name "simplechatMigration" +} + +function Get-AISearchMigrationProvenanceFieldDefinitions { + return @( + [ordered]@{ + name = "simplechatMigrationId" + type = "Edm.String" + searchable = $false + filterable = $true + sortable = $true + facetable = $false + retrievable = $true + }, + [ordered]@{ + name = "simplechatMigratedAtUtc" + type = "Edm.DateTimeOffset" + filterable = $true + sortable = $true + facetable = $false + retrievable = $true + }, + [ordered]@{ + name = "simplechatMigrationStatus" + type = "Edm.String" + searchable = $false + filterable = $true + sortable = $true + facetable = $false + retrievable = $true + } + ) +} + +function Add-AISearchMigrationProvenance { + param( + [object]$Document, + [object]$Context + ) + + $provenance = New-MigrationProvenanceRecord -Context $Context + $fieldValues = [ordered]@{ + simplechatMigrationId = $provenance["migrationId"] + simplechatMigratedAtUtc = $provenance["migratedAtUtc"] + simplechatMigrationStatus = $provenance["status"] + } + + foreach ($fieldName in $fieldValues.Keys) { + if ($Document -is [System.Collections.IDictionary]) { + $Document[$fieldName] = $fieldValues[$fieldName] + } + else { + $Document | Add-Member ` + -MemberType NoteProperty ` + -Name $fieldName ` + -Value $fieldValues[$fieldName] ` + -Force + } + } +} + +function Get-AISearchMigrationProvenance { + param( + [AllowNull()] + [object]$Document + ) + + return [ordered]@{ + migrationId = Get-MigrationProvenancePropertyValue ` + -Source $Document ` + -Name "simplechatMigrationId" + migratedAtUtc = Get-MigrationProvenancePropertyValue ` + -Source $Document ` + -Name "simplechatMigratedAtUtc" + status = Get-MigrationProvenancePropertyValue ` + -Source $Document ` + -Name "simplechatMigrationStatus" + } +} + +function Merge-StorageMigrationProvenance { + param( + [AllowNull()] + [object]$Metadata, + [object]$Context + ) + + $mergedMetadata = [ordered]@{} + if ($null -ne $Metadata) { + if ($Metadata -is [System.Collections.IDictionary]) { + foreach ($key in $Metadata.Keys) { + $mergedMetadata[[string]$key] = $Metadata[$key] + } + } + else { + foreach ($property in $Metadata.PSObject.Properties) { + $mergedMetadata[$property.Name] = $property.Value + } + } + } + + $provenance = New-MigrationProvenanceRecord -Context $Context + $mergedMetadata["simplechatMigrationId"] = $provenance["migrationId"] + $mergedMetadata["simplechatMigratedAtUtc"] = $provenance["migratedAtUtc"] + $mergedMetadata["simplechatMigrationStatus"] = $provenance["status"] + return $mergedMetadata +} + +function Get-StorageMigrationProvenance { + param( + [AllowNull()] + [object]$Metadata + ) + + return [ordered]@{ + migrationId = Get-MigrationProvenancePropertyValue ` + -Source $Metadata ` + -Name "simplechatMigrationId" + migratedAtUtc = Get-MigrationProvenancePropertyValue ` + -Source $Metadata ` + -Name "simplechatMigratedAtUtc" + status = Get-MigrationProvenancePropertyValue ` + -Source $Metadata ` + -Name "simplechatMigrationStatus" + } +} \ No newline at end of file diff --git a/scripts/Migration-State.ps1 b/scripts/Migration-State.ps1 new file mode 100644 index 000000000..90a2be201 --- /dev/null +++ b/scripts/Migration-State.ps1 @@ -0,0 +1,368 @@ +# Migration-State.ps1 + +function Get-MigrationStateUtcTimestamp { + return [DateTime]::UtcNow.ToString( + "o", + [System.Globalization.CultureInfo]::InvariantCulture + ) +} + +function Get-MigrationStateFingerprint { + param( + [System.Collections.IDictionary]$Configuration + ) + + $configurationJson = ConvertTo-Json ` + -InputObject $Configuration ` + -Depth 100 ` + -Compress + $sha256 = [System.Security.Cryptography.SHA256]::Create() + try { + $hash = $sha256.ComputeHash( + [System.Text.Encoding]::UTF8.GetBytes($configurationJson) + ) + } + finally { + $sha256.Dispose() + } + return [Convert]::ToHexString($hash).ToLowerInvariant() +} + +function ConvertTo-MigrationStateResult { + param( + [AllowNull()] + [object]$Result + ) + + $converted = [ordered]@{} + if ($null -eq $Result) { + return $converted + } + + if ($Result -is [System.Collections.IDictionary]) { + foreach ($key in $Result.Keys) { + $converted[[string]$key] = $Result[$key] + } + return $converted + } + + foreach ($property in $Result.PSObject.Properties) { + $converted[$property.Name] = $property.Value + } + return $converted +} + +function Save-MigrationState { + param( + [object]$Context + ) + + $Context.Data["updatedUtc"] = Get-MigrationStateUtcTimestamp + $stateDirectory = [IO.Path]::GetDirectoryName($Context.Path) + if (-not [string]::IsNullOrWhiteSpace($stateDirectory)) { + [IO.Directory]::CreateDirectory($stateDirectory) | Out-Null + } + + $temporaryPath = "$($Context.Path).$PID.tmp" + $stateJson = ConvertTo-Json ` + -InputObject $Context.Data ` + -Depth 100 + try { + [IO.File]::WriteAllText( + $temporaryPath, + "$stateJson`n", + [System.Text.UTF8Encoding]::new($false) + ) + [IO.File]::Move($temporaryPath, $Context.Path, $true) + } + finally { + if ([IO.File]::Exists($temporaryPath)) { + [IO.File]::Delete($temporaryPath) + } + } +} + +function Initialize-MigrationState { + param( + [string]$MigrationType, + [string]$StateFilePath, + [System.Collections.IDictionary]$Configuration, + [string]$MigrationId = "", + [switch]$Reset + ) + + if ([string]::IsNullOrWhiteSpace($StateFilePath)) { + throw "StateFilePath cannot be empty when migration state tracking is enabled." + } + + $requestedMigrationId = [Guid]::Empty + if ( + -not [string]::IsNullOrWhiteSpace($MigrationId) -and + -not [Guid]::TryParse($MigrationId, [ref]$requestedMigrationId) + ) { + throw "MigrationId must be a valid GUID." + } + + $absoluteStatePath = [IO.Path]::GetFullPath( + [Environment]::ExpandEnvironmentVariables($StateFilePath.Trim()) + ) + if ($Reset -and [IO.File]::Exists($absoluteStatePath)) { + [IO.File]::Delete($absoluteStatePath) + } + + $fingerprint = Get-MigrationStateFingerprint -Configuration $Configuration + if ([IO.File]::Exists($absoluteStatePath)) { + try { + $state = [IO.File]::ReadAllText($absoluteStatePath) | + ConvertFrom-Json -AsHashtable -Depth 100 + } + catch { + throw "Migration state file '$absoluteStatePath' is not valid JSON. Use -ResetState after preserving it for investigation." + } + + if ([int]$state["schemaVersion"] -ne 1) { + throw "Migration state file '$absoluteStatePath' uses an unsupported schema version." + } + if (-not [string]::Equals( + [string]$state["migrationType"], + $MigrationType, + [System.StringComparison]::Ordinal + )) { + throw "Migration state file '$absoluteStatePath' belongs to a different migration type." + } + if (-not [string]::Equals( + [string]$state["configurationFingerprint"], + $fingerprint, + [System.StringComparison]::Ordinal + )) { + throw "Migration state file '$absoluteStatePath' does not match the current source, destination, or mode. Use a different -StateFilePath or pass -ResetState." + } + if ($null -eq $state["resources"]) { + $state["resources"] = [ordered]@{} + } + $storedMigrationId = [Guid]::Empty + if ( + -not [string]::IsNullOrWhiteSpace([string]$state["migrationId"]) -and + -not [Guid]::TryParse([string]$state["migrationId"], [ref]$storedMigrationId) + ) { + throw "Migration state file '$absoluteStatePath' contains an invalid migrationId. Use -ResetState after preserving it for investigation." + } + if ($storedMigrationId -eq [Guid]::Empty) { + $storedMigrationId = if ($requestedMigrationId -eq [Guid]::Empty) { + [Guid]::NewGuid() + } + else { + $requestedMigrationId + } + $state["migrationId"] = $storedMigrationId.ToString("D") + } + elseif ( + $requestedMigrationId -ne [Guid]::Empty -and + $storedMigrationId -ne $requestedMigrationId + ) { + throw "Migration state file '$absoluteStatePath' belongs to migration ID '$($storedMigrationId.ToString("D"))'. Use a different -StateFilePath or pass -ResetState." + } + else { + $state["migrationId"] = $storedMigrationId.ToString("D") + } + if ([string]::IsNullOrWhiteSpace([string]$state["migrationStartedUtc"])) { + $state["migrationStartedUtc"] = [string]$state["createdUtc"] + } + $state["status"] = "in_progress" + $state["lastError"] = $null + $state["currentResource"] = $null + $state["resumeCount"] = [int]$state["resumeCount"] + 1 + } + else { + $timestamp = Get-MigrationStateUtcTimestamp + $generatedMigrationId = if ($requestedMigrationId -eq [Guid]::Empty) { + [Guid]::NewGuid() + } + else { + $requestedMigrationId + } + $state = [ordered]@{ + schemaVersion = 1 + migrationType = $MigrationType + migrationId = $generatedMigrationId.ToString("D") + migrationStartedUtc = $timestamp + configurationFingerprint = $fingerprint + configuration = $Configuration + status = "in_progress" + createdUtc = $timestamp + updatedUtc = $timestamp + completedUtc = $null + resumeCount = 0 + currentResource = $null + lastError = $null + resources = [ordered]@{} + summary = [ordered]@{} + } + } + + $context = [pscustomobject]@{ + Path = $absoluteStatePath + Data = $state + } + Save-MigrationState -Context $context + return $context +} + +function Get-MigrationResourceCheckpoint { + param( + [object]$Context, + [string]$ResourceName + ) + + $resources = $Context.Data["resources"] + if ($resources.Contains($ResourceName)) { + return $resources[$ResourceName] + } + return $null +} + +function Test-MigrationResourceCompleted { + param( + [object]$Context, + [string]$ResourceName + ) + + $checkpoint = Get-MigrationResourceCheckpoint ` + -Context $Context ` + -ResourceName $ResourceName + return ( + $null -ne $checkpoint -and + [string]::Equals( + [string]$checkpoint["status"], + "completed", + [System.StringComparison]::Ordinal + ) + ) +} + +function Start-MigrationResourceCheckpoint { + param( + [object]$Context, + [string]$ResourceName + ) + + $existingCheckpoint = Get-MigrationResourceCheckpoint ` + -Context $Context ` + -ResourceName $ResourceName + $attempt = if ($null -eq $existingCheckpoint) { + 1 + } + else { + [int]$existingCheckpoint["attempt"] + 1 + } + $progress = if ( + $null -ne $existingCheckpoint -and + $null -ne $existingCheckpoint["progress"] + ) { + $existingCheckpoint["progress"] + } + else { + [ordered]@{} + } + $timestamp = Get-MigrationStateUtcTimestamp + $Context.Data["resources"][$ResourceName] = [ordered]@{ + status = "in_progress" + attempt = $attempt + startedUtc = $timestamp + updatedUtc = $timestamp + completedUtc = $null + lastError = $null + progress = $progress + result = [ordered]@{} + } + $Context.Data["currentResource"] = $ResourceName + Save-MigrationState -Context $Context +} + +function Update-MigrationResourceCheckpoint { + param( + [object]$Context, + [string]$ResourceName, + [System.Collections.IDictionary]$Progress + ) + + $checkpoint = Get-MigrationResourceCheckpoint ` + -Context $Context ` + -ResourceName $ResourceName + if ($null -eq $checkpoint) { + throw "Migration resource '$ResourceName' was updated before it was started." + } + if ($checkpoint["status"] -eq "completed") { + throw "Migration resource '$ResourceName' cannot be updated after completion." + } + + $checkpoint["progress"] = ConvertTo-MigrationStateResult -Result $Progress + $checkpoint["updatedUtc"] = Get-MigrationStateUtcTimestamp + Save-MigrationState -Context $Context +} + +function Complete-MigrationResourceCheckpoint { + param( + [object]$Context, + [string]$ResourceName, + [AllowNull()] + [object]$Result = $null + ) + + $checkpoint = Get-MigrationResourceCheckpoint ` + -Context $Context ` + -ResourceName $ResourceName + if ($null -eq $checkpoint) { + throw "Migration resource '$ResourceName' was completed before it was started." + } + + $timestamp = Get-MigrationStateUtcTimestamp + $checkpoint["status"] = "completed" + $checkpoint["updatedUtc"] = $timestamp + $checkpoint["completedUtc"] = $timestamp + $checkpoint["lastError"] = $null + $checkpoint["result"] = ConvertTo-MigrationStateResult -Result $Result + $Context.Data["currentResource"] = $null + Save-MigrationState -Context $Context +} + +function Set-MigrationStateFailure { + param( + [object]$Context, + [AllowNull()] + [string]$ResourceName, + [string]$ErrorMessage + ) + + $timestamp = Get-MigrationStateUtcTimestamp + if (-not [string]::IsNullOrWhiteSpace($ResourceName)) { + $checkpoint = Get-MigrationResourceCheckpoint ` + -Context $Context ` + -ResourceName $ResourceName + if ($null -ne $checkpoint -and $checkpoint["status"] -ne "completed") { + $checkpoint["status"] = "failed" + $checkpoint["updatedUtc"] = $timestamp + $checkpoint["lastError"] = $ErrorMessage + } + } + $Context.Data["status"] = "failed" + $Context.Data["currentResource"] = $ResourceName + $Context.Data["lastError"] = $ErrorMessage + Save-MigrationState -Context $Context +} + +function Complete-MigrationState { + param( + [object]$Context, + [AllowNull()] + [object]$Summary = $null + ) + + $timestamp = Get-MigrationStateUtcTimestamp + $Context.Data["status"] = "completed" + $Context.Data["completedUtc"] = $timestamp + $Context.Data["currentResource"] = $null + $Context.Data["lastError"] = $null + $Context.Data["summary"] = ConvertTo-MigrationStateResult -Result $Summary + Save-MigrationState -Context $Context +} \ No newline at end of file diff --git a/scripts/Migration-StorageAccount.ps1 b/scripts/Migration-StorageAccount.ps1 new file mode 100644 index 000000000..f017d6645 --- /dev/null +++ b/scripts/Migration-StorageAccount.ps1 @@ -0,0 +1,482 @@ +# Migration-StorageAccount.ps1 +#Requires -Version 7.0 + +[CmdletBinding()] +param( + [string]$SourceStorageAccount = "", + + [string]$SourceSubscriptionId = "", + + [string]$DestinationStorageAccount = "", + + [string]$DestinationSubscriptionId = "", + + [ValidateNotNullOrEmpty()] + [string[]]$Containers = @( + "user-documents", + "group-documents", + "public-documents", + "group-chat", + "personal-chat" + ), + + [bool]$DifferentialMigration = $true, + + [string]$MigrationId = "", + + [ValidateRange(0, 8760)] + [int]$SkipMigratedWithinHours = 24, + + [bool]$ShowProgress = $true, + + [ValidateRange(1, 168)] + [int]$SasExpiryHours = 48, + + [string]$StateFilePath = "", + + [switch]$ResetState +) + +$ErrorActionPreference = "Stop" +. (Join-Path $PSScriptRoot "Migration-State.ps1") +. (Join-Path $PSScriptRoot "Migration-Provenance.ps1") + +function Write-StorageMigrationProgress { + param( + [int]$CompletedCount, + [int]$TotalCount, + [string]$CurrentContainer = "", + [switch]$Completed + ) + + if (-not $ShowProgress) { + return + } + + if ($Completed) { + Write-Progress -Id 0 -Activity "Storage account migration" -Completed + return + } + + $percentComplete = if ($TotalCount -eq 0) { + 100 + } + else { + [int][Math]::Floor(($CompletedCount / $TotalCount) * 100) + } + $remainingCount = [Math]::Max(0, $TotalCount - $CompletedCount) + $progressParameters = @{ + Id = 0 + Activity = "Storage account migration" + Status = "Containers: $CompletedCount/$TotalCount | Remaining: $remainingCount | $percentComplete%" + PercentComplete = $percentComplete + } + if (-not [string]::IsNullOrWhiteSpace($CurrentContainer)) { + $progressParameters.CurrentOperation = "Current container: $CurrentContainer" + } + + Write-Progress @progressParameters +} + +function Test-StorageMigrationConfiguration { + $configuredValues = [ordered]@{ + SourceStorageAccount = $SourceStorageAccount + SourceSubscriptionId = $SourceSubscriptionId + DestinationStorageAccount = $DestinationStorageAccount + DestinationSubscriptionId = $DestinationSubscriptionId + } + foreach ($entry in $configuredValues.GetEnumerator()) { + if ( + [string]::IsNullOrWhiteSpace($entry.Value) -or + $entry.Value -match '^<[^>]+>$' + ) { + throw "Set '$($entry.Key)' in the script or supply it as a parameter." + } + } + + foreach ($accountName in @($SourceStorageAccount, $DestinationStorageAccount)) { + if ($accountName -notmatch '^[a-z0-9]{3,24}$') { + throw "Storage account name '$accountName' must contain 3-24 lowercase letters or numbers." + } + } + + foreach ($subscriptionId in @($SourceSubscriptionId, $DestinationSubscriptionId)) { + $parsedSubscriptionId = [Guid]::Empty + if (-not [Guid]::TryParse($subscriptionId, [ref]$parsedSubscriptionId)) { + throw "Subscription ID '$subscriptionId' must be a valid GUID." + } + } + + if ([string]::Equals( + $SourceStorageAccount, + $DestinationStorageAccount, + [System.StringComparison]::OrdinalIgnoreCase + )) { + throw "Source and destination storage accounts must be different." + } + + $containerNames = [System.Collections.Generic.HashSet[string]]::new( + [System.StringComparer]::Ordinal + ) + foreach ($container in $Containers) { + if ( + [string]::IsNullOrWhiteSpace($container) -or + $container -notmatch '^(?!.*--)[a-z0-9](?:[a-z0-9-]{1,61}[a-z0-9])$' + ) { + throw "Container name '$container' is not a valid Azure Blob container name." + } + if (-not $containerNames.Add($container)) { + throw "Container '$container' is listed more than once." + } + } +} + +function Get-StorageBlobCloudBlob { + param( + [object]$StorageBlob + ) + + $cloudBlob = $StorageBlob.PSObject.Properties["ICloudBlob"].Value + if ($null -eq $cloudBlob) { + throw "Blob '$($StorageBlob.Name)' does not expose an ICloudBlob metadata object." + } + return $cloudBlob +} + +function Get-StorageBlobMigrationProvenance { + param( + [object]$StorageBlob + ) + + $cloudBlob = Get-StorageBlobCloudBlob -StorageBlob $StorageBlob + if ($null -ne $cloudBlob.PSObject.Methods["FetchAttributes"]) { + $cloudBlob.FetchAttributes() + } + return Get-StorageMigrationProvenance -Metadata $cloudBlob.Metadata +} + +function Set-StorageBlobMigrationProvenance { + param( + [object]$StorageBlob, + [object]$MigrationProvenanceContext + ) + + $cloudBlob = Get-StorageBlobCloudBlob -StorageBlob $StorageBlob + if ($null -ne $cloudBlob.PSObject.Methods["FetchAttributes"]) { + $cloudBlob.FetchAttributes() + } + if ($null -eq $cloudBlob.Metadata) { + throw "Blob '$($StorageBlob.Name)' does not expose mutable metadata." + } + if ($null -eq $cloudBlob.PSObject.Methods["SetMetadata"]) { + throw "Blob '$($StorageBlob.Name)' does not support metadata updates." + } + + $mergedMetadata = Merge-StorageMigrationProvenance ` + -Metadata $cloudBlob.Metadata ` + -Context $MigrationProvenanceContext + $cloudBlob.Metadata.Clear() + foreach ($metadataKey in $mergedMetadata.Keys) { + $cloudBlob.Metadata[[string]$metadataKey] = [string]$mergedMetadata[$metadataKey] + } + $cloudBlob.SetMetadata() +} + +function Test-StorageContainerMigrationProvenanceSkip { + param( + [string]$ContainerName, + [object]$SourceContext, + [object]$DestinationContext, + [object]$MigrationProvenanceContext + ) + + $sourceBlobCount = 0 + $allSourceBlobsMarked = $true + Get-AzStorageBlob ` + -Container $ContainerName ` + -Context $SourceContext | ForEach-Object { + $sourceBlob = $_ + $sourceBlobCount++ + $destinationBlob = Get-AzStorageBlob ` + -Container $ContainerName ` + -Blob $sourceBlob.Name ` + -Context $DestinationContext ` + -ErrorAction SilentlyContinue + if ($null -eq $destinationBlob) { + $allSourceBlobsMarked = $false + return + } + + $destinationProvenance = Get-StorageBlobMigrationProvenance ` + -StorageBlob $destinationBlob + if (-not (Test-MigrationProvenanceSkip ` + -Provenance $destinationProvenance ` + -Context $MigrationProvenanceContext)) { + $allSourceBlobsMarked = $false + } + } + + return [pscustomobject]@{ + SourceBlobCount = $sourceBlobCount + ShouldSkip = ($sourceBlobCount -gt 0 -and $allSourceBlobsMarked) + } +} + +function Set-StorageContainerMigrationProvenance { + param( + [string]$ContainerName, + [object]$SourceContext, + [object]$DestinationContext, + [object]$MigrationProvenanceContext + ) + + $taggedBlobCount = 0 + Get-AzStorageBlob ` + -Container $ContainerName ` + -Context $SourceContext | ForEach-Object { + $sourceBlob = $_ + $destinationBlob = Get-AzStorageBlob ` + -Container $ContainerName ` + -Blob $sourceBlob.Name ` + -Context $DestinationContext ` + -ErrorAction SilentlyContinue + if ($null -eq $destinationBlob) { + throw "Destination blob '$($sourceBlob.Name)' was not found after AzCopy completed." + } + Set-StorageBlobMigrationProvenance ` + -StorageBlob $destinationBlob ` + -MigrationProvenanceContext $MigrationProvenanceContext + $taggedBlobCount++ + } + return $taggedBlobCount +} + +$migrationMode = if ($DifferentialMigration) { "differential" } else { "full" } +if ([string]::IsNullOrWhiteSpace($StateFilePath)) { + $StateFilePath = Join-Path $PSScriptRoot "Migration-StorageAccount.state.json" +} +$migrationState = $null +$activeMigrationResource = $null + +try { + Test-StorageMigrationConfiguration + + $stateConfiguration = [ordered]@{ + sourceAccount = $SourceStorageAccount + sourceSubscriptionId = $SourceSubscriptionId + destinationAccount = $DestinationStorageAccount + destinationSubscriptionId = $DestinationSubscriptionId + containers = @($Containers) + mode = $migrationMode + skipMigratedWithinHours = $SkipMigratedWithinHours + } + $migrationState = Initialize-MigrationState ` + -MigrationType "storage" ` + -StateFilePath $StateFilePath ` + -Configuration $stateConfiguration ` + -MigrationId $MigrationId ` + -Reset:$ResetState + $migrationProvenanceContext = New-MigrationProvenanceContext ` + -MigrationId ([string]$migrationState.Data["migrationId"]) ` + -MigratedAtUtc ([string]$migrationState.Data["migrationStartedUtc"]) ` + -SkipMigratedWithinHours $SkipMigratedWithinHours + Write-Host "Migration state: $($migrationState.Path)" + Write-Host "Migration ID: $($migrationProvenanceContext.MigrationId)" + + foreach ($requiredCommand in @( + "Connect-AzAccount", + "Set-AzContext", + "New-AzStorageContext", + "Get-AzStorageContainer", + "New-AzStorageContainer", + "New-AzStorageContainerSASToken", + "azcopy" + )) { + if (-not (Get-Command $requiredCommand -ErrorAction SilentlyContinue)) { + throw "Required command '$requiredCommand' is not available." + } + } + + Connect-AzAccount | Out-Null + + Set-AzContext -SubscriptionId $SourceSubscriptionId | Out-Null + $sourceContext = New-AzStorageContext ` + -StorageAccountName $SourceStorageAccount ` + -UseConnectedAccount + + Set-AzContext -SubscriptionId $DestinationSubscriptionId | Out-Null + $destinationContext = New-AzStorageContext ` + -StorageAccountName $DestinationStorageAccount ` + -UseConnectedAccount + + $completedContainerCount = 0 + Write-Host "Starting $migrationMode storage account migration. Destination-only blobs will not be deleted." + Write-StorageMigrationProgress ` + -CompletedCount 0 ` + -TotalCount $Containers.Count + + foreach ($container in $Containers) { + Write-StorageMigrationProgress ` + -CompletedCount $completedContainerCount ` + -TotalCount $Containers.Count ` + -CurrentContainer $container + Write-Host "Migrating container in $migrationMode mode: $container" + + if (Test-MigrationResourceCompleted ` + -Context $migrationState ` + -ResourceName $container) { + $completedContainerCount++ + Write-Host "Skipping completed container from migration state: $container" + Write-StorageMigrationProgress ` + -CompletedCount $completedContainerCount ` + -TotalCount $Containers.Count ` + -CurrentContainer $container + continue + } + + Start-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $container + $activeMigrationResource = $container + + Set-AzContext -SubscriptionId $SourceSubscriptionId | Out-Null + $sourceContainer = Get-AzStorageContainer ` + -Name $container ` + -Context $sourceContext + if ($null -eq $sourceContainer) { + throw "Source container '$container' does not exist in storage account '$SourceStorageAccount'." + } + + $sasStartTime = (Get-Date).ToUniversalTime().AddMinutes(-15) + $sasExpiryTime = (Get-Date).ToUniversalTime().AddHours($SasExpiryHours) + $sourceSasUrl = New-AzStorageContainerSASToken ` + -Name $container ` + -Context $sourceContext ` + -Permission "rlt" ` + -Protocol HttpsOnly ` + -StartTime $sasStartTime ` + -ExpiryTime $sasExpiryTime ` + -FullUri + + Set-AzContext -SubscriptionId $DestinationSubscriptionId | Out-Null + $destinationContainer = Get-AzStorageContainer ` + -Name $container ` + -Context $destinationContext + if ($null -eq $destinationContainer) { + New-AzStorageContainer ` + -Name $container ` + -Context $destinationContext | Out-Null + Write-Host "Created destination container: $container" + } + + $containerProvenance = Test-StorageContainerMigrationProvenanceSkip ` + -ContainerName $container ` + -SourceContext $sourceContext ` + -DestinationContext $destinationContext ` + -MigrationProvenanceContext $migrationProvenanceContext + if ($containerProvenance.ShouldSkip) { + Complete-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $container ` + -Result ([ordered]@{ + Mode = $migrationMode + SourceBlobCount = $containerProvenance.SourceBlobCount + TaggedBlobCount = 0 + SkippedByProvenance = $true + }) + $activeMigrationResource = $null + $completedContainerCount++ + Write-Host "Skipping container '$container' because all $($containerProvenance.SourceBlobCount) blob(s) were recently migrated." + Write-StorageMigrationProgress ` + -CompletedCount $completedContainerCount ` + -TotalCount $Containers.Count ` + -CurrentContainer $container + continue + } + + $destinationSasUrl = New-AzStorageContainerSASToken ` + -Name $container ` + -Context $destinationContext ` + -Permission "racwdlt" ` + -Protocol HttpsOnly ` + -StartTime $sasStartTime ` + -ExpiryTime $sasExpiryTime ` + -FullUri + + if ($DifferentialMigration) { + $azCopyArguments = @( + "sync" + $sourceSasUrl + $destinationSasUrl + "--recursive=true" + "--delete-destination=false" + "--s2s-preserve-blob-tags=true" + ) + } + else { + $azCopyArguments = @( + "copy" + $sourceSasUrl + $destinationSasUrl + "--recursive=true" + "--overwrite=true" + "--s2s-preserve-properties=true" + "--s2s-preserve-blob-tags=true" + ) + } + + & azcopy @azCopyArguments + + if ($LASTEXITCODE -ne 0) { + throw "AzCopy failed for container '$container' with exit code $LASTEXITCODE." + } + + $taggedBlobCount = Set-StorageContainerMigrationProvenance ` + -ContainerName $container ` + -SourceContext $sourceContext ` + -DestinationContext $destinationContext ` + -MigrationProvenanceContext $migrationProvenanceContext + + Complete-MigrationResourceCheckpoint ` + -Context $migrationState ` + -ResourceName $container ` + -Result ([ordered]@{ + Mode = $migrationMode + SourceBlobCount = $containerProvenance.SourceBlobCount + TaggedBlobCount = $taggedBlobCount + SkippedByProvenance = $false + }) + $activeMigrationResource = $null + $completedContainerCount++ + Write-StorageMigrationProgress ` + -CompletedCount $completedContainerCount ` + -TotalCount $Containers.Count ` + -CurrentContainer $container + } + + Complete-MigrationState ` + -Context $migrationState ` + -Summary ([ordered]@{ + ContainerCount = $completedContainerCount + }) + Write-Host "Storage account migration completed successfully. Containers processed: $completedContainerCount" +} +catch { + $migrationErrorMessage = $_.Exception.Message + if ($null -ne $migrationState) { + try { + Set-MigrationStateFailure ` + -Context $migrationState ` + -ResourceName $activeMigrationResource ` + -ErrorMessage $migrationErrorMessage + } + catch { + Write-Warning "Failed to update migration state after an error: $($_.Exception.Message)" + } + } + Write-Error "Storage account migration failed: $migrationErrorMessage" -ErrorAction Continue + throw +} +finally { + Write-StorageMigrationProgress -CompletedCount 0 -TotalCount 0 -Completed +} \ No newline at end of file diff --git a/ui_tests/test_admin_data_management_settings_ui.py b/ui_tests/test_admin_data_management_settings_ui.py index 16399adb9..6f82e927a 100644 --- a/ui_tests/test_admin_data_management_settings_ui.py +++ b/ui_tests/test_admin_data_management_settings_ui.py @@ -1,10 +1,9 @@ # test_admin_data_management_settings_ui.py """ UI test for Admin Settings Data Management controls. -Version: 0.250.051 +Version: 0.250.071 Implemented in: 0.241.211 Updated in: 0.241.221 -Updated in: 0.250.051 This test ensures admins can discover the Data Management tab, see the operational-business-hours warning, and access the backup, encryption, @@ -12,6 +11,8 @@ Version 0.250.049 moves query results and document editing into a scrollable modal. Version 0.250.050 keeps this coverage aligned with the Cosmos editor save-path fix. Version 0.250.051 verifies the Cosmos editor results list scrolls independently. +Version 0.250.071 adds resilient migration provenance, incremental modes, cutover +reconciliation, and the external target Search writer freeze acknowledgement. """ import os @@ -74,6 +75,8 @@ def test_admin_data_management_controls_render_from_template(): "data_management_target_cosmos_endpoint", "data_management_target_cosmos_database", "data-management-target-cosmos-key-field", + "data_management_target_cosmos_subscription_id", + "data_management_target_cosmos_resource_group", "data-management-test-target-cosmos-btn", "data-management-target-ai-search-section", "data_management_target_ai_search_auth", @@ -89,6 +92,24 @@ def test_admin_data_management_controls_render_from_template(): "data_management_target_ec_connection_string", "data-management-test-target-ec-storage-btn", "data-management-migration-workflow-section", + "data-management-test-migration-access-btn", + "data-management-migration-mode-section", + "data_management_migration_mode_new_only", + "data_management_migration_mode_delta_upsert", + "data_management_migration_mode_mirror_with_deletions", + "data-management-migration-mode-description", + "data-management-migration-baseline-field", + "data_management_migration_baseline_job_id", + "data-management-migration-mirror-confirmation", + "data_management_migration_mirror_confirmation_phrase", + "data-management-migration-search-write-freeze", + "data_management_migration_target_search_writes_frozen", + "data_management_migration_max_parallel_operations", + "data_management_migration_retry_count", + "data_management_migration_skip_recent_within_hours", + "data_management_migration_temporary_destination_ru_enabled", + "data-management-migration-temporary-ru-field", + "data_management_migration_temporary_destination_ru", "data_management_migration_users_mode", "data-management-migration-users-available", "data-management-migration-users-selected", @@ -139,10 +160,14 @@ def test_admin_data_management_controls_render_from_template(): "data-management-job-detail-modal", "data-management-job-detail-refresh-state", "data-management-job-detail-progress", + "data-management-job-detail-actions", "data-management-job-items-tbody", "data-management-job-artifacts-tbody", "data-management-job-manifest-detail", "data-management-job-warnings", + "data-management-migration-cancel-modal", + "data-management-migration-cancel-message", + "data-management-confirm-migration-cancel-btn", ] for element_id in required_ids: @@ -187,6 +212,11 @@ def test_admin_data_management_controls_render_from_template(): assert "Paste a connection string to save or replace it" in template assert 'id="data-management-connection-string-status"' in template assert 'id="data_management_target_cosmos_database" value="SimpleChat" readonly aria-readonly="true"' in template + assert 'max="10000"' in template + assert 'Validate Cosmos Migration Access' in template + assert 'role="radiogroup" aria-label="Migration synchronization mode"' in template + assert "MIRROR WITH DELETIONS" in template + assert "I confirm external destination AI Search writers are frozen for this migration" in template assert 'setStorageAuthVisibility' in js_source assert 'updateConnectionStringStatus' in js_source assert 'updateSourceBlobBackupAvailability' in js_source @@ -206,6 +236,25 @@ def test_admin_data_management_controls_render_from_template(): assert 'The edit was recorded in Activity Logs.' in js_source assert 'closest("[data-ignore-data-management-change' in js_source assert 'testTargetCosmos' in js_source + assert 'testMigrationAccess' in js_source + assert 'retryMigrationJob' in js_source + assert 'openMigrationCancellationModal' in js_source + assert 'requestMigrationCancellation' in js_source + assert 'getMigrationLiveMetrics' in js_source + assert 'jobDetailRefreshIntervalMs = 2000' in js_source + assert '/progress`' in js_source + assert 'Observed transferred' in js_source + assert 'Running - alive, no recent progress' in js_source + assert 'updateMigrationCapacityVisibility' in js_source + assert 'updateMigrationModeVisibility' in js_source + assert 'updateMigrationSearchWriteFreezeVisibility' in js_source + assert 'target_ai_search_writes_frozen' in js_source + assert 'createMigrationPreviewOutcomes' in js_source + assert 'include_inventory: Boolean(showSuccess)' in js_source + assert 'Stage ${formatNumber(displayedStage)} of ${formatNumber(totalSteps)}' in js_source + assert 'Active stage' in js_source + assert 'Migration stage is active; measured throughput is shown below.' in js_source + assert 'Migration cancellation requested. The worker will stop at its next durable checkpoint.' in js_source assert 'testTargetSearch' in js_source assert 'testTargetEnhancedCitationStorage' in js_source assert 'Migration preview refreshed.' in js_source @@ -215,6 +264,8 @@ def test_admin_data_management_controls_render_from_template(): assert 'loadDataManagementBackups' in js_source assert 'loadDataManagementJobDetail' in js_source assert 'renderJobArtifacts' in js_source + assert 'renderJobDetailActions' in js_source + assert 'createMigrationManifestDownloadLink' in js_source assert 'createDetailChipGroup' in js_source assert 'startJobDetailAutoRefresh' in js_source assert 'Live updates on -' in js_source @@ -263,6 +314,17 @@ def test_admin_data_management_tab_browser_workflow(): expect(page.get_by_role("button", name="Advanced backup scope")).to_be_visible() expect(page.locator("#data_management_target_cosmos_database")).to_have_value("SimpleChat") expect(page.locator("#data_management_target_cosmos_database")).to_have_attribute("readonly", "") + expect(page.locator("#data_management_migration_max_parallel_operations")).to_be_visible() + expect(page.locator("#data_management_migration_retry_count")).to_be_visible() + expect(page.locator("#data_management_migration_skip_recent_within_hours")).to_be_visible() + expect(page.get_by_label("New only")).to_be_checked() + expect(page.locator("#data-management-migration-baseline-field")).to_have_class(re.compile(r"\bd-none\b")) + page.get_by_label("Delta / upsert").check() + expect(page.locator("#data-management-migration-baseline-field")).not_to_have_class(re.compile(r"\bd-none\b")) + page.get_by_label("Mirror with deletions").check() + expect(page.locator("#data-management-migration-mirror-confirmation")).to_be_visible() + expect(page.get_by_label("Temporarily increase destination Cosmos capacity during this migration")).to_be_visible() + expect(page.get_by_role("button", name="Validate Cosmos Migration Access")).to_be_visible() expect(page.locator("#data-management-target-ai-search-section")).to_be_visible() expect(page.locator("#data-management-test-target-search-btn")).to_be_visible() expect(page.locator("#data-management-migration-workflow-section")).to_be_visible() @@ -286,6 +348,7 @@ def test_admin_data_management_tab_browser_workflow(): expect(page.locator("#data-management-backups-tbody")).to_be_visible() expect(page.locator("#data-management-jobs-tbody")).to_be_visible() expect(page.locator("#data-management-job-detail-modal")).to_be_attached() + expect(page.locator("#data-management-migration-cancel-modal")).to_be_attached() finally: context.close() browser.close()