diff --git a/.github/scripts/check_env_drift.py b/.github/scripts/check_env_drift.py index fdc0eba5..9d13a00f 100644 --- a/.github/scripts/check_env_drift.py +++ b/.github/scripts/check_env_drift.py @@ -41,6 +41,13 @@ PLACEHOLDER = re.compile(r"<[^<>]+>") +# Typed wrappers around os.getenv (e.g. `_get_bool_env`, +# `_get_bounded_int_env` in hastegeo.core.config). They take the variable name +# as the first argument and supply their own default in the signature, so a +# call is a genuine read even when no default is passed at the call site. +# Without this the scanner sees no reader and reports the setting as dead. +ENV_HELPER = re.compile(r"^_get_[a-z0-9_]*env$") + # Variables that are genuinely optional for an Azure deployment, with the reason # each one is exempt. Anything not listed here that the code marks required must # be emitted by both deploy paths. @@ -126,6 +133,11 @@ def scan_code() -> tuple[dict[str, set[Path]], set[str]]: and func.attr in ("getenv", "get") and node.args ) + is_helper = ( + isinstance(func, ast.Name) + and ENV_HELPER.match(func.id) + and node.args + ) if is_getenv: target = ast.unparse(func) if target.endswith( @@ -135,6 +147,13 @@ def scan_code() -> tuple[dict[str, set[Path]], set[str]]: has_default = len(node.args) > 1 if has_default: default = _literal(node.args[1]) + elif is_helper: + name = _literal(node.args[0]) + # The wrapper defines its own default, so the read is + # optional even with no default at the call site. + has_default = True + if len(node.args) > 1: + default = _literal(node.args[1]) elif isinstance(node, ast.Subscript): value = node.value diff --git a/.github/scripts/deploy_apps.sh b/.github/scripts/deploy_apps.sh index c439f485..e6d55abe 100644 --- a/.github/scripts/deploy_apps.sh +++ b/.github/scripts/deploy_apps.sh @@ -56,6 +56,16 @@ BATCH_INFERENCE_POOL_IDS="${BATCH_INFERENCE_POOL_IDS:-}" BATCH_IMAGERYPREP_POOL_IDS="${BATCH_IMAGERYPREP_POOL_IDS:-}" BATCH_USE_SAS="${BATCH_USE_SAS:-false}" BATCH_MANAGE_POOLS="${BATCH_MANAGE_POOLS:-true}" +# Data publishing feature flag. Mirrors the `publishingEnabled` param in +# infra/modules/functions.bicep so both deploy paths agree; defaults off. +PUBLISHING_ENABLED="${PUBLISHING_ENABLED:-false}" +# Planetary Computer publishing target. Mirrors the pc* params in +# infra/modules/functions.bicep; default off / unset. +PC_PROVIDER_ENABLED="${PC_PROVIDER_ENABLED:-false}" +PC_GEOCATALOG_URL="${PC_GEOCATALOG_URL:-}" +PC_EXPLORER_URL="${PC_EXPLORER_URL:-}" +PC_INGESTION_SOURCE="${PC_INGESTION_SOURCE:-}" +PC_COLLECTION_PREFIX="${PC_COLLECTION_PREFIX:-haste-}" MAPS_ACCOUNT="${RESOURCE_PREFIX}haste${RANDOM_SUFFIX}maps" API_MANAGEMENT="${RESOURCE_PREFIX}-haste-${RANDOM_SUFFIX}-apim" FIXED_TAGS="project=haste created_by=deploy_apps" @@ -135,6 +145,12 @@ deploy_function() { "STATIC_APP_DOMAIN=${STATIC_APP_DOMAIN}" \ "EMAIL_CONNECTION_STRING=${EMAIL_CONNECTION_STRING}" \ "EMAIL_SENDER=${EMAIL_SENDER}" \ + "PUBLISHING_ENABLED=${PUBLISHING_ENABLED}" \ + "PC_PROVIDER_ENABLED=${PC_PROVIDER_ENABLED}" \ + "PC_GEOCATALOG_URL=${PC_GEOCATALOG_URL}" \ + "PC_EXPLORER_URL=${PC_EXPLORER_URL}" \ + "PC_INGESTION_SOURCE=${PC_INGESTION_SOURCE}" \ + "PC_COLLECTION_PREFIX=${PC_COLLECTION_PREFIX}" \ --output none fi diff --git a/.github/workflows/deploy-apps.yml b/.github/workflows/deploy-apps.yml index d5dec0f3..d13b0c34 100644 --- a/.github/workflows/deploy-apps.yml +++ b/.github/workflows/deploy-apps.yml @@ -111,6 +111,10 @@ jobs: # Non-sensitive UI feature flag, baked into the Vite bundle at build time. # Set as a GitHub Environment variable; defaults to false when unset. VITE_SHOW_FOOTER: ${{ vars.VITE_SHOW_FOOTER }} + # Non-sensitive data-publishing feature flag, mirroring the + # `publishingEnabled` param on the Bicep path. Set as a GitHub + # Environment variable; defaults to false when unset. + PUBLISHING_ENABLED: ${{ vars.PUBLISHING_ENABLED }} # hastegeo wheel pinned into the function-app requirements before # `func publish` (deploy_apps.sh); the editable default can't resolve # on Azure's remote build. diff --git a/api/hastefuncapi/function_app.py b/api/hastefuncapi/function_app.py index 7683072b..40161864 100644 --- a/api/hastefuncapi/function_app.py +++ b/api/hastefuncapi/function_app.py @@ -22,6 +22,12 @@ ModelArtifacts, Project, ) +from hastegeo.core.models.publishing import ( + ArtifactKind, + PublishRequest, + PublishStatus, + PublishTarget, +) from hastegeo.core.models.stats import ( ImageLayerStats, ProjectsSummary, @@ -32,13 +38,42 @@ from hastegeo.core.models.users import User from hastegeo.core.models.visualizer import Imagery, Visualizer from hastegeo.core.processors.artifacts import ArtifactProcessor +from hastegeo.core.processors.assessment import ( + AssessmentReportProcessor, + AssessmentSizeLimitError, +) from hastegeo.core.processors.embedding import EmbeddingPreprocessor from hastegeo.core.processors.imagery import ImageryPreProcessor from hastegeo.core.processors.inference import InferencePreprocessor from hastegeo.core.processors.metadata import MetadataProcessor +from hastegeo.core.processors.publishing import ( + PublishingDependencyError, + PublishingDisabledError, + PublishingPermissionError, + PublishingProcessor, + PublishingSizeLimitError, + PublishingStateConflictError, +) from hastegeo.core.processors.stats import StatsPreProcessor from hastegeo.core.processors.train import TrainPreprocessor from hastegeo.core.processors.uploader import FileUploader +from hastegeo.core.publishing.lease import LeaseUnavailableError +from hastegeo.core.publishing.registry import ( + ProviderUnavailableError, + PublishingProviderRegistry, +) +from hastegeo.core.publishing.repository import ( + PublishedDatasetsExistError, + PublishingConflictError, + PublishingRepository, + StaleRevisionError, +) +from hastegeo.core.publishing.source import ( + PublishingArtifactUnavailableError, + PublishingSourceNotEligibleError, + PublishingSourceNotFoundError, + PublishingSourceResolver, +) from hastegeo.core.utils.blob import ( download_blob_to_tempfile, parse_byte_range, @@ -52,7 +87,6 @@ validate_image_layer_imagery_urls, validate_image_layer_user_footprints_url, ) -from hastegeo.core.utils.user import InvitationManager, UserManager from pydantic import ValidationError # type: ignore config = Config() @@ -87,6 +121,7 @@ # to leave room for the field to grow without ever admitting an unbounded # string into log lines or blob paths. _SHORT_INT_ID_RE = re.compile(r"^[0-9]{1,8}$") +_PUBLISH_ASSESSMENT_MAX_TOTAL_BYTES = 512 * 1024**2 def _require_guid_param(req: func.HttpRequest, name: str) -> str: @@ -180,6 +215,163 @@ def _require_roles( return None +async def _get_active_publishing_caller( + req: func.HttpRequest, +) -> tuple[dict | None, func.HttpResponse | None]: + """Return the trusted active HASTE caller used by publishing routes.""" + principal = _decode_client_principal(req) + if DEVELOPMENT_MODE: + principal = principal or { + "userId": "development@local", + "userDetails": "development@local", + "userRoles": ["authenticated", "contributors", "administrators"], + } + roles = { + role.lower().strip() + for role in principal.get("userRoles", []) + if isinstance(role, str) + } + caller_id = ( + principal.get("userId") + or principal.get("userDetails") + or "development@local" + ) + return {"id": str(caller_id).lower(), "roles": roles, "name": principal.get("userDetails")}, None + + if principal is None: + return None, _publishing_error_response( + "UNAUTHENTICATED", "Authentication is required.", 401 + ) + + principal_id = principal.get("userId") + user_details = principal.get("userDetails") + if not principal_id and not user_details: + return None, _publishing_error_response( + "UNAUTHENTICATED", "Authentication is required.", 401 + ) + + try: + raw_users = await asyncio.to_thread( + MetadataProcessor( + data_type=config.get_metadata_types().USERS.value + ).load, + "acl", + ) + except FileNotFoundError: + return None, _publishing_error_response( + "FORBIDDEN", "An active HASTE user is required.", 403 + ) + + users = [User(**user) for user in raw_users] + active_user = next( + ( + user + for user in users + if ( + user.userId in {principal_id, user_details} + or user.objectId == principal_id + ) + and user.status == config.get_user_statuses().ACTIVE.value + and not user.deleted + ), + None, + ) + if active_user is None: + return None, _publishing_error_response( + "FORBIDDEN", "An active HASTE user is required.", 403 + ) + + roles = { + role.lower().strip() + for role in principal.get("userRoles", []) + if isinstance(role, str) + } + caller_id = principal_id or user_details + # Persist the email/login as the publisher identifier, never the display + # name (privacy: display names are resolved from Entra at read time). + return {"id": str(caller_id).lower(), "roles": roles, "name": (active_user.email or user_details)}, None + + +def _publishing_json_response( + payload: dict, status_code: int = 200 +) -> func.HttpResponse: + return func.HttpResponse( + json.dumps(payload), + status_code=status_code, + mimetype="application/json", + ) + + +def _publishing_error_response( + code: str, message: str, status_code: int +) -> func.HttpResponse: + return _publishing_json_response( + {"error": {"code": code, "message": message}}, status_code + ) + + +def _publishing_exception_response(error: Exception) -> func.HttpResponse: + if isinstance(error, ValidationError): + return _publishing_error_response( + "VALIDATION_ERROR", "Invalid publishing request.", 400 + ) + if isinstance(error, PublishingPermissionError): + return _publishing_error_response("FORBIDDEN", str(error), 403) + if isinstance(error, PublishingSizeLimitError): + return _publishing_error_response( + "PUBLISH_SIZE_LIMIT_EXCEEDED", str(error), 413 + ) + if isinstance( + error, + ( + PublishingConflictError, + PublishingStateConflictError, + PublishingSourceNotEligibleError, + StaleRevisionError, + LeaseUnavailableError, + ), + ): + return _publishing_error_response("CONFLICT", str(error), 409) + if isinstance( + error, + ( + PublishingSourceNotFoundError, + PublishingArtifactUnavailableError, + FileNotFoundError, + ), + ): + return _publishing_error_response("NOT_FOUND", str(error), 404) + if isinstance( + error, + ( + PublishingDependencyError, + PublishingDisabledError, + ProviderUnavailableError, + ), + ): + return _publishing_error_response( + "PUBLISHING_UNAVAILABLE", str(error), 503 + ) + if isinstance(error, ValueError): + return _publishing_error_response( + "VALIDATION_ERROR", str(error), 400 + ) + logger.error( + "Publishing request failed with %s", type(error).__name__ + ) + return _publishing_error_response( + "INTERNAL_ERROR", "Publishing request failed.", 500 + ) + + +def _publishing_mutation_authorized(caller: dict) -> bool: + return bool(caller["roles"].intersection({"contributors", "administrators"})) + + +def _publishing_processor() -> PublishingProcessor: + return PublishingProcessor(config=config) + + def add_cors_headers(response: func.HttpResponse) -> func.HttpResponse: """Add CORS headers to the response - handled by nginx proxy in local dev.""" # CORS headers are now handled by nginx reverse proxy @@ -772,11 +964,18 @@ async def DeleteProject(req: func.HttpRequest) -> func.HttpResponse: except ValueError as ve: return _bad_request(f"DeleteProject: {ve}") - await asyncio.to_thread( + repository = PublishingRepository(config=config) + + def delete_project_metadata() -> None: MetadataProcessor( data_type=config.get_metadata_types().PROJECT.value, partition_key=project_id, - ).delete_all_from_partition + ).delete_all_from_partition() + + await asyncio.to_thread( + repository.delete_project_if_unpublished, + project_id, + delete_project_metadata, ) request = StatsPreProcessor( @@ -794,6 +993,16 @@ async def DeleteProject(req: func.HttpRequest) -> func.HttpResponse: status_code=200, ) + except PublishedDatasetsExistError as e: + return _publishing_error_response( + "PUBLISHED_DATASETS_EXIST", str(e), 409 + ) + except LeaseUnavailableError: + return _publishing_error_response( + "PROJECT_PUBLISHING_ACTIVE", + "A publishing operation is active for this project.", + 409, + ) except FileNotFoundError as e: logger.error(f"Project not found: {e}\n{traceback.format_exc()}") return func.HttpResponse("Project not found.", status_code=404) @@ -1472,6 +1681,8 @@ async def PutAdminSettings(req: func.HttpRequest) -> func.HttpResponse: @app.route(route="GetUsers", auth_level=AUTH_LEVEL, methods=["GET"]) async def GetUsers(req: func.HttpRequest) -> func.HttpResponse: + from hastegeo.core.utils.user import UserManager + logger.info("GetUsers HTTP trigger function processed a request.") auth_error = _require_roles(req, {"administrators"}) if auth_error: @@ -1584,6 +1795,8 @@ async def GetUsers(req: func.HttpRequest) -> func.HttpResponse: @app.route(route="PutUser", auth_level=AUTH_LEVEL, methods=["PUT"]) async def PutUser(req: func.HttpRequest) -> func.HttpResponse: + from hastegeo.core.utils.user import InvitationManager + logger.info("PutUser HTTP trigger function processed a request.") try: req_body = req.get_json() @@ -1776,6 +1989,8 @@ def roles_changed( @app.route(route="DeleteUser", auth_level=AUTH_LEVEL, methods=["DELETE"]) async def DeleteUser(req: func.HttpRequest) -> func.HttpResponse: + from hastegeo.core.utils.user import UserManager + logger.info("DeleteUser HTTP trigger function processed a request.") auth_error = _require_roles(req, {"administrators"}) if auth_error: @@ -1828,6 +2043,8 @@ async def DeleteUser(req: func.HttpRequest) -> func.HttpResponse: @app.route(route="GetUserById", auth_level=AUTH_LEVEL, methods=["GET"]) async def GetUserById(req: func.HttpRequest) -> func.HttpResponse: + from hastegeo.core.utils.user import UserManager + logger.info("GetUser HTTP trigger function processed a request.") try: user_id = req.params.get("userId") @@ -4215,3 +4432,255 @@ async def GetAssessmentReport(req: func.HttpRequest) -> func.HttpResponse: return func.HttpResponse( "Error generating assessment report.", status_code=500 ) +@app.route( + route="GetPublishingProviders", + auth_level=AUTH_LEVEL, + methods=["GET"], +) +async def GetPublishingProviders(req: func.HttpRequest) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + registry = PublishingProviderRegistry(config=config) + return _publishing_json_response( + { + "publishingEnabled": config.publishing_config[ + "publishing_enabled" + ], + "providers": [ + info.model_dump(mode="json") for info in registry.list_infos() + ], + } + ) + + +@app.route( + route="GetPublishDatasetOptions", + auth_level=AUTH_LEVEL, + methods=["GET"], +) +async def GetPublishDatasetOptions(req: func.HttpRequest) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + if not _publishing_mutation_authorized(caller): + return _publishing_error_response( + "FORBIDDEN", "Contributor role required.", 403 + ) + try: + if not config.publishing_config["publishing_enabled"]: + raise PublishingDisabledError("Publishing is disabled") + project_id = _require_guid_param(req, "projectId") + image_layer_id = _require_guid_param(req, "imageLayerId") + model_id = _require_short_int_id_param(req, "modelId") + options = await asyncio.to_thread( + PublishingSourceResolver(config=config).resolve_options, + project_id, + image_layer_id, + model_id, + ) + return _publishing_json_response( + {"publishDatasetOptions": options.model_dump(mode="json")} + ) + except Exception as error: + return _publishing_exception_response(error) + + +@app.route( + route="GetPublishedDatasets", + auth_level=AUTH_LEVEL, + methods=["GET"], +) +async def GetPublishedDatasets(req: func.HttpRequest) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + try: + try: + page = int(req.params.get("page", "1")) + page_size = int(req.params.get("pageSize", "20")) + except ValueError as error: + raise ValueError("page and pageSize must be integers") from error + project_id = req.params.get("projectId") + if project_id and not _GUID_RE.match(project_id): + raise ValueError("Invalid projectId") + search = req.params.get("search", "").strip() + if len(search) > 200: + raise ValueError("search must be at most 200 characters") + if search and len(search) < 3: + raise ValueError("search must be at least 3 characters") + target = ( + PublishTarget(req.params["target"]) + if req.params.get("target") + else None + ) + status = ( + PublishStatus(req.params["status"]) + if req.params.get("status") + else None + ) + records, total_count = await asyncio.to_thread( + PublishingRepository(config=config).list_page, + page=page, + page_size=page_size, + project_id=project_id, + target=target, + status=status, + search=search, + sort_key=req.params.get("sortKey", "publishedDate"), + sort_direction=req.params.get("sortDirection", "desc"), + ) + return _publishing_json_response( + { + "publishedDatasets": [ + record.model_dump(mode="json") for record in records + ], + "pagination": { + "page": page, + "pageSize": page_size, + "totalCount": total_count, + }, + } + ) + except Exception as error: + return _publishing_exception_response(error) + + +@app.route( + route="GetPublishedDataset", + auth_level=AUTH_LEVEL, + methods=["GET"], +) +async def GetPublishedDataset(req: func.HttpRequest) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + try: + project_id = _require_guid_param(req, "projectId") + dataset_id = _require_guid_param(req, "datasetId") + processor = _publishing_processor() + record = await asyncio.to_thread( + processor.get_dataset, project_id, dataset_id + ) + download_urls = await asyncio.to_thread( + processor.get_download_urls, project_id, dataset_id + ) + return _publishing_json_response( + { + "publishedDataset": record.model_dump(mode="json"), + "downloadUrls": download_urls, + } + ) + except Exception as error: + return _publishing_exception_response(error) + + +@app.route( + route="PutPublishDatasetQueueMessage", + auth_level=AUTH_LEVEL, + methods=["PUT"], +) +async def PutPublishDatasetQueueMessage( + req: func.HttpRequest, +) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + if not _publishing_mutation_authorized(caller): + return _publishing_error_response( + "FORBIDDEN", "Contributor role required.", 403 + ) + try: + request = PublishRequest(**req.get_json()) + processor = _publishing_processor() + prepared = await asyncio.to_thread( + processor.prepare_create, + request, + caller["id"], + caller.get("name"), + ) + if prepared.existing is not None: + record = prepared.existing + else: + try: + assessment_summary = await AssessmentReportProcessor( + config=config + ).generate( + str(request.projectId), + request.imageLayerId, + request.modelId, + max_total_bytes=_PUBLISH_ASSESSMENT_MAX_TOTAL_BYTES, + ) + except Exception as assessment_error: + logger.warning( + "Assessment snapshot unavailable for publish: %s", + type(assessment_error).__name__, + ) + assessment_summary = {} + record = await asyncio.to_thread( + processor.create_prepared, + prepared, + assessment_summary, + ) + return _publishing_json_response( + {"publishedDataset": record.model_dump(mode="json")}, 202 + ) + except Exception as error: + return _publishing_exception_response(error) + + +@app.route( + route="PutRetryPublishedDatasetQueueMessage", + auth_level=AUTH_LEVEL, + methods=["PUT"], +) +async def PutRetryPublishedDatasetQueueMessage( + req: func.HttpRequest, +) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + try: + body = req.get_json() + project_id = str(body.get("projectId", "")) + dataset_id = str(body.get("datasetId", "")) + if not _GUID_RE.match(project_id) or not _GUID_RE.match(dataset_id): + raise ValueError("Valid projectId and datasetId are required") + record = await asyncio.to_thread( + _publishing_processor().retry, + project_id, + dataset_id, + caller["id"], + "administrators" in caller["roles"], + ) + return _publishing_json_response( + {"publishedDataset": record.model_dump(mode="json")}, 202 + ) + except Exception as error: + return _publishing_exception_response(error) + + +@app.route( + route="DeletePublishedDataset", + auth_level=AUTH_LEVEL, + methods=["DELETE"], +) +async def DeletePublishedDataset(req: func.HttpRequest) -> func.HttpResponse: + caller, auth_error = await _get_active_publishing_caller(req) + if auth_error: + return auth_error + try: + project_id = _require_guid_param(req, "projectId") + dataset_id = _require_guid_param(req, "datasetId") + record = await asyncio.to_thread( + _publishing_processor().request_unpublish, + project_id, + dataset_id, + caller["id"], + "administrators" in caller["roles"], + ) + return _publishing_json_response( + {"publishedDataset": record.model_dump(mode="json")}, 202 + ) + except Exception as error: + return _publishing_exception_response(error) diff --git a/api/hastefuncapi/requirements.txt b/api/hastefuncapi/requirements.txt index 4c3ac866..0ace9aa2 100644 --- a/api/hastefuncapi/requirements.txt +++ b/api/hastefuncapi/requirements.txt @@ -23,6 +23,13 @@ opencv-python==4.10.0.84 GDAL @ https://github.com/microsoft/haste/releases/download/haste-binaries/GDAL-3.9.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; sys_platform == 'linux' shapely==2.0.6 pyyaml==6.0.2 +# Planetary Computer publishing (vendored REST client) — geo stack used by the +# PC provider's STAC builder (lazy-loaded when the API resolves the PC target +# to validate a publish request). +pystac[validation]==1.11.0 +geopandas==1.0.1 +pyproj==3.6.1 +pyogrio==0.10.0 boto3==1.36.20 tensorboard==2.19.0 tenacity==9.1.2 diff --git a/api/hastefuncapi/tests/__init__.py b/api/hastefuncapi/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/api/hastefuncapi/tests/test_publishing_routes.py b/api/hastefuncapi/tests/test_publishing_routes.py new file mode 100644 index 00000000..7a5024bf --- /dev/null +++ b/api/hastefuncapi/tests/test_publishing_routes.py @@ -0,0 +1,627 @@ +import base64 +import io +import json +import os +import unittest +import uuid +from contextlib import redirect_stderr +from unittest.mock import AsyncMock, Mock, patch + +import azure.functions as func + +os.environ.setdefault("DEVELOPMENT_MODE", "true") +os.environ.setdefault("PUBLISHING_ENABLED", "true") +os.environ.setdefault("METADATA_STORAGE_TYPE", "local") +os.environ.setdefault("ARTIFACT_STORAGE_TYPE", "local") +os.environ.setdefault("DATA_PATH", "/tmp/haste-publishing-api-tests") +os.environ.setdefault("TEMP_DATA_PATH", "/tmp/haste-publishing-api-tests") + +with redirect_stderr(io.StringIO()): + from api.hastefuncapi import function_app + +from hastegeo.core.models.publishing import ( # noqa: E402 + PublishDatasetOptions, + PublishedDataset, +) +from hastegeo.core.processors.publishing import ( # noqa: E402 + PublishingDependencyError, + PublishingDisabledError, + PublishingPermissionError, +) +from hastegeo.core.publishing.repository import ( # noqa: E402 + PublishedDatasetsExistError, +) + +PROJECT_ID = "123e4567-e89b-12d3-a456-426614174000" +DATASET_ID = "3e8d5e90-f2fc-5412-9f97-a52c07815f0b" +REQUEST_ID = "550e8400-e29b-41d4-a716-446655440000" + + +def make_request( + method: str = "GET", + params: dict | None = None, + body: dict | None = None, + headers: dict | None = None, +) -> func.HttpRequest: + encoded_body = ( + json.dumps(body).encode("utf-8") if body is not None else b"" + ) + return func.HttpRequest( + method=method, + url="http://localhost/api/publishing", + headers=headers or {}, + params=params or {}, + route_params={}, + body=encoded_body, + ) + + +def response_json(response: func.HttpResponse) -> dict: + return json.loads(response.get_body().decode("utf-8")) + + +def make_dataset(status: str = "PENDING") -> PublishedDataset: + return PublishedDataset( + datasetId=DATASET_ID, + requestId=REQUEST_ID, + requestFingerprint="a" * 64, + name="Published damage assessment", + projectId=PROJECT_ID, + imageLayerId="layer-1", + modelId="42", + target="local", + status=status, + publishedByUser="publisher-object-id", + createdDate="2026-08-06T00:00:00Z", + updatedDate="2026-08-06T00:00:00Z", + ) + + +class TestPublishingRoutes(unittest.IsolatedAsyncioTestCase): + async def test_inference_launch_rejects_client_runtime_state(self) -> None: + response = await function_app.PutRunInferenceQueueMessage( + make_request( + method="PUT", + body={ + "projectId": PROJECT_ID, + "modelId": "42", + "gpkgUrl": "https://storage.example/forged.gpkg", + }, + ) + ) + + self.assertEqual(response.status_code, 400) + + async def test_layer_request_rejects_workflow_owned_artifact_paths( + self, + ) -> None: + response = await function_app.PutLayer( + make_request( + method="PUT", + body={ + "projectId": PROJECT_ID, + "name": "New layer", + "postEventProcessedImageryUrl": ( + "https://attacker.example/admin.tif" + ), + "buildingFootprintsUrl": ( + "https://attacker.example/users_acl.json" + ), + "status": "Processed", + }, + ) + ) + + self.assertEqual(response.status_code, 400) + + async def test_trusted_principal_maps_to_active_haste_user(self) -> None: + principal = { + "userId": "OBJECT-ID", + "userDetails": "publisher@example.com", + "userRoles": ["authenticated", "contributors"], + } + encoded = base64.b64encode( + json.dumps(principal).encode("utf-8") + ).decode("ascii") + metadata = Mock() + metadata.load.return_value = [ + { + "userId": "publisher@example.com", + "objectId": "OBJECT-ID", + "status": function_app.config.get_user_statuses().ACTIVE.value, + "deleted": False, + } + ] + with patch.object( + function_app, "DEVELOPMENT_MODE", False + ), patch.object( + function_app, "MetadataProcessor", return_value=metadata + ): + caller, error = await function_app._get_active_publishing_caller( + make_request(headers={"x-ms-client-principal": encoded}) + ) + + self.assertIsNone(error) + self.assertEqual(caller["id"], "object-id") + self.assertEqual(caller["roles"], {"authenticated", "contributors"}) + + async def test_invalid_principal_header_is_unauthenticated(self) -> None: + with patch.object(function_app, "DEVELOPMENT_MODE", False): + caller, error = await function_app._get_active_publishing_caller( + make_request( + headers={"x-ms-client-principal": "not-valid-base64"} + ) + ) + + self.assertIsNone(caller) + self.assertEqual(error.status_code, 401) + + async def test_provider_list_allows_active_viewer(self) -> None: + caller = {"id": "viewer", "roles": {"authenticated"}} + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.dict( + function_app.config.publishing_config, + {"publishing_enabled": True}, + ): + response = await function_app.GetPublishingProviders( + make_request() + ) + + payload = response_json(response) + self.assertEqual(response.status_code, 200) + self.assertTrue(payload["publishingEnabled"]) + self.assertEqual(payload["providers"][0]["id"], "local") + + async def test_options_reject_viewer_without_mutation_role(self) -> None: + caller = {"id": "viewer", "roles": {"authenticated"}} + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ): + response = await function_app.GetPublishDatasetOptions( + make_request( + params={ + "projectId": PROJECT_ID, + "imageLayerId": str(uuid.uuid4()), + "modelId": "42", + } + ) + ) + + self.assertEqual(response.status_code, 403) + self.assertEqual(response_json(response)["error"]["code"], "FORBIDDEN") + + async def test_options_return_server_verified_artifacts(self) -> None: + caller = {"id": "contributor", "roles": {"contributors"}} + options = PublishDatasetOptions( + projectId=PROJECT_ID, + projectName="Project", + imageLayerId=str(uuid.uuid4()), + imageLayerName="Layer", + modelId="42", + modelName="Model", + defaultName="Project - Layer", + availableArtifacts=[ + { + "kind": "gpkg", + "sourcePath": "damage.gpkg", + "mediaType": "application/geopackage+sqlite3", + "sizeBytes": 10, + "sourceEtag": "etag-1", + } + ], + ) + resolver = Mock() + resolver.resolve_options.return_value = options + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "PublishingSourceResolver", + return_value=resolver, + ): + response = await function_app.GetPublishDatasetOptions( + make_request( + params={ + "projectId": PROJECT_ID, + "imageLayerId": str(options.imageLayerId), + "modelId": "42", + } + ) + ) + + self.assertEqual(response.status_code, 200) + payload = response_json(response)["publishDatasetOptions"] + self.assertEqual(payload["availableArtifacts"][0]["kind"], "gpkg") + + async def test_catalog_returns_bounded_pagination_metadata(self) -> None: + caller = {"id": "viewer", "roles": {"authenticated"}} + repository = Mock() + repository.list_page.return_value = ([make_dataset("PUBLISHED")], 45) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, "PublishingRepository", return_value=repository + ): + response = await function_app.GetPublishedDatasets( + make_request( + params={ + "page": "2", + "pageSize": "20", + "target": "local", + "search": "damage", + "sortKey": "name", + "sortDirection": "asc", + } + ) + ) + + payload = response_json(response) + self.assertEqual(response.status_code, 200) + self.assertEqual( + payload["pagination"], + {"page": 2, "pageSize": 20, "totalCount": 45}, + ) + repository.list_page.assert_called_once_with( + page=2, + page_size=20, + project_id=None, + target=function_app.PublishTarget.LOCAL, + status=None, + search="damage", + sort_key="name", + sort_direction="asc", + ) + + async def test_catalog_rejects_unbounded_page_size(self) -> None: + caller = {"id": "viewer", "roles": {"authenticated"}} + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ): + response = await function_app.GetPublishedDatasets( + make_request(params={"pageSize": "101"}) + ) + + self.assertEqual(response.status_code, 400) + + async def test_catalog_rejects_one_character_search(self) -> None: + caller = {"id": "viewer", "roles": {"authenticated"}} + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ): + response = await function_app.GetPublishedDatasets( + make_request(params={"search": "a"}) + ) + + self.assertEqual(response.status_code, 400) + + async def test_publish_uses_trusted_caller_identity(self) -> None: + caller = { + "id": "publisher-object-id", + "roles": {"contributors"}, + } + processor = Mock() + prepared = Mock(existing=None) + processor.prepare_create.return_value = prepared + processor.create_prepared.return_value = make_dataset() + assessment = Mock() + assessment.generate = AsyncMock(return_value={"predictedDamaged": 5}) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ), patch.object( + function_app, + "AssessmentReportProcessor", + return_value=assessment, + ): + response = await function_app.PutPublishDatasetQueueMessage( + make_request( + method="PUT", + body={ + "requestId": REQUEST_ID, + "projectId": PROJECT_ID, + "imageLayerId": "layer-1", + "modelId": "42", + "name": "Published damage assessment", + "target": "local", + "artifacts": ["gpkg"], + }, + ) + ) + + self.assertEqual(response.status_code, 202) + self.assertEqual( + processor.prepare_create.call_args.args[1], caller["id"] + ) + self.assertEqual( + processor.create_prepared.call_args.args, + (prepared, {"predictedDamaged": 5}), + ) + assessment.generate.assert_awaited_once_with( + PROJECT_ID, + "layer-1", + "42", + max_total_bytes=function_app._PUBLISH_ASSESSMENT_MAX_TOTAL_BYTES, + ) + + async def test_publish_replay_skips_assessment_generation(self) -> None: + caller = { + "id": "publisher-object-id", + "roles": {"contributors"}, + } + processor = Mock() + processor.prepare_create.return_value = Mock( + existing=make_dataset("PUBLISHED") + ) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ), patch.object( + function_app, "AssessmentReportProcessor" + ) as assessment_type: + response = await function_app.PutPublishDatasetQueueMessage( + make_request( + method="PUT", + body={ + "requestId": REQUEST_ID, + "projectId": PROJECT_ID, + "imageLayerId": "layer-1", + "modelId": "42", + "name": "Published damage assessment", + "target": "local", + "artifacts": ["gpkg"], + }, + ) + ) + + self.assertEqual(response.status_code, 202) + assessment_type.assert_not_called() + processor.create_prepared.assert_not_called() + + async def test_disabled_publish_stops_before_assessment(self) -> None: + caller = { + "id": "publisher-object-id", + "roles": {"contributors"}, + } + processor = Mock() + processor.prepare_create.side_effect = PublishingDisabledError( + "Publishing is disabled" + ) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ), patch.object( + function_app, "AssessmentReportProcessor" + ) as assessment_type: + response = await function_app.PutPublishDatasetQueueMessage( + make_request( + method="PUT", + body={ + "requestId": REQUEST_ID, + "projectId": PROJECT_ID, + "imageLayerId": "layer-1", + "modelId": "42", + "name": "Published damage assessment", + "target": "local", + "artifacts": ["gpkg"], + }, + ) + ) + + self.assertEqual(response.status_code, 503) + assessment_type.assert_not_called() + + async def test_publish_rejects_forged_body_identity(self) -> None: + caller = { + "id": "publisher-object-id", + "roles": {"contributors"}, + } + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ): + response = await function_app.PutPublishDatasetQueueMessage( + make_request( + method="PUT", + body={ + "requestId": REQUEST_ID, + "projectId": PROJECT_ID, + "imageLayerId": "layer-1", + "modelId": "42", + "name": "Published damage assessment", + "target": "local", + "artifacts": ["gpkg"], + "publishedByUser": "attacker", + }, + ) + ) + + self.assertEqual(response.status_code, 400) + self.assertEqual( + response_json(response)["error"]["code"], "VALIDATION_ERROR" + ) + + async def test_retry_returns_structured_disabled_response(self) -> None: + caller = { + "id": "publisher-object-id", + "roles": {"contributors"}, + } + processor = Mock() + processor.retry.side_effect = PublishingDisabledError( + "Publishing is disabled" + ) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ): + response = await function_app.PutRetryPublishedDatasetQueueMessage( + make_request( + method="PUT", + body={ + "projectId": PROJECT_ID, + "datasetId": DATASET_ID, + }, + ) + ) + + self.assertEqual(response.status_code, 503) + self.assertEqual( + response_json(response)["error"]["code"], + "PUBLISHING_UNAVAILABLE", + ) + + async def test_retry_returns_structured_dependency_response(self) -> None: + caller = { + "id": "publisher-object-id", + "roles": {"contributors"}, + } + processor = Mock() + processor.retry.side_effect = PublishingDependencyError( + "Unable to enqueue publishing retry" + ) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ): + response = await function_app.PutRetryPublishedDatasetQueueMessage( + make_request( + method="PUT", + body={ + "projectId": PROJECT_ID, + "datasetId": DATASET_ID, + }, + ) + ) + + self.assertEqual(response.status_code, 503) + self.assertEqual( + response_json(response)["error"]["code"], + "PUBLISHING_UNAVAILABLE", + ) + + async def test_unpublish_returns_structured_permission_response( + self, + ) -> None: + caller = {"id": "other-user", "roles": {"contributors"}} + processor = Mock() + processor.request_unpublish.side_effect = PublishingPermissionError( + "Only the publisher may unpublish" + ) + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ): + response = await function_app.DeletePublishedDataset( + make_request( + method="DELETE", + params={ + "projectId": PROJECT_ID, + "datasetId": DATASET_ID, + }, + ) + ) + + self.assertEqual(response.status_code, 403) + self.assertEqual(response_json(response)["error"]["code"], "FORBIDDEN") + + async def test_detail_returns_fresh_urls_separate_from_metadata( + self, + ) -> None: + caller = {"id": "viewer", "roles": {"authenticated"}} + processor = Mock() + processor.get_dataset.return_value = make_dataset("PUBLISHED") + processor.get_download_urls.return_value = { + "gpkg": "https://storage/blob.gpkg?short-lived" + } + with patch.object( + function_app, + "_get_active_publishing_caller", + new=AsyncMock(return_value=(caller, None)), + ), patch.object( + function_app, + "_publishing_processor", + return_value=processor, + ): + response = await function_app.GetPublishedDataset( + make_request( + params={ + "projectId": PROJECT_ID, + "datasetId": DATASET_ID, + } + ) + ) + + payload = response_json(response) + self.assertEqual(response.status_code, 200) + self.assertEqual( + payload["downloadUrls"]["gpkg"], + "https://storage/blob.gpkg?short-lived", + ) + self.assertNotIn("downloadUrls", payload["publishedDataset"]) + + async def test_project_delete_conflicts_when_publications_exist( + self, + ) -> None: + repository = Mock() + repository.delete_project_if_unpublished.side_effect = ( + PublishedDatasetsExistError("Unpublish datasets first") + ) + with patch.object( + function_app, "PublishingRepository", return_value=repository + ): + response = await function_app.DeleteProject( + make_request(method="DELETE", params={"projectId": PROJECT_ID}) + ) + + self.assertEqual(response.status_code, 409) + self.assertEqual( + response_json(response)["error"]["code"], + "PUBLISHED_DATASETS_EXIST", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/api/hastefuncqueues/function_app.py b/api/hastefuncqueues/function_app.py index d3bdc86b..829e0cd5 100644 --- a/api/hastefuncqueues/function_app.py +++ b/api/hastefuncqueues/function_app.py @@ -15,6 +15,7 @@ ModelArtifacts, Project, ) +from hastegeo.core.models.publishing import PublishQueueMessage from hastegeo.core.models.stats import ProjectsSummary, StatsRequest from hastegeo.core.models.training import ExperimentConfig from hastegeo.core.processors.artifacts import ArtifactProcessor @@ -26,6 +27,7 @@ ) from hastegeo.core.processors.labels import LabelTaskGenerator from hastegeo.core.processors.metadata import MetadataProcessor +from hastegeo.core.processors.publishing import PublishingProcessor from hastegeo.core.processors.stats import StatsPostProcessor from hastegeo.core.processors.train import TrainPostprocessor from hastegeo.core.utils.data import convert_json_to_geojson @@ -988,3 +990,75 @@ async def GetArtifactsZipQueueMessage(msg: func.QueueMessage) -> None: f"ArtifactsZipQueueTrigger: Error saving failed status: {inner_e}\n{traceback.format_exc()}", stack_info=True, ) + + +@app.function_name(name="PublishDatasetQueueTrigger") +@app.queue_trigger( + arg_name="msg", + queue_name=config.get_queue_config()["publish_queue_name"], + connection="AzureWebJobsStorage", +) +async def GetPublishDatasetQueueMessage(msg: func.QueueMessage) -> None: + try: + message = PublishQueueMessage( + **json.loads(msg.get_body().decode("utf-8")) + ) + await asyncio.to_thread(PublishingProcessor(config=config).run_step, message) + except Exception as error: + logger.error( + "PublishDatasetQueueTrigger failed with %s", + type(error).__name__, + ) + raise RuntimeError( + f"Publishing queue step failed: {type(error).__name__}" + ) from None + + +@app.function_name(name="PublishDatasetPoisonQueueTrigger") +@app.queue_trigger( + arg_name="msg", + queue_name=f'{config.get_queue_config()["publish_queue_name"]}-poison', + connection="AzureWebJobsStorage", +) +async def GetPublishDatasetPoisonQueueMessage(msg: func.QueueMessage) -> None: + try: + message = PublishQueueMessage( + **json.loads(msg.get_body().decode("utf-8")) + ) + await asyncio.to_thread( + PublishingProcessor(config=config).mark_poisoned, message + ) + except FileNotFoundError: + logger.info("Ignoring poison message for a removed published dataset") + except Exception as error: + logger.error( + "PublishDatasetPoisonQueueTrigger failed with %s", + type(error).__name__, + ) + raise RuntimeError( + f"Publishing poison step failed: {type(error).__name__}" + ) from None + + +@app.function_name(name="ReconcilePublishingOperations") +@app.timer_trigger( + arg_name="timer", + schedule="0 */5 * * * *", + run_on_startup=False, + use_monitor=True, +) +async def ReconcilePublishingOperations(timer: func.TimerRequest) -> None: + try: + requeued = await asyncio.to_thread( + PublishingProcessor(config=config).reconcile_stale + ) + if requeued: + logger.info("Requeued %s stale publishing operations", requeued) + except Exception as error: + logger.error( + "ReconcilePublishingOperations failed with %s", + type(error).__name__, + ) + raise RuntimeError( + f"Publishing reconciliation failed: {type(error).__name__}" + ) from None diff --git a/api/hastefuncqueues/requirements.txt b/api/hastefuncqueues/requirements.txt index 4c3ac866..958635c9 100644 --- a/api/hastefuncqueues/requirements.txt +++ b/api/hastefuncqueues/requirements.txt @@ -12,7 +12,7 @@ azure-storage-queue==12.12.0 azure-mgmt-containerregistry==10.3.0 azure-mgmt-web azure-batch==14.2.0 -azure-core==1.38.0 +azure-core==1.39.0 azure-identity==1.18.0 email_validator==2.3.0 psycopg2-binary==2.9.9 @@ -23,6 +23,12 @@ opencv-python==4.10.0.84 GDAL @ https://github.com/microsoft/haste/releases/download/haste-binaries/GDAL-3.9.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl ; sys_platform == 'linux' shapely==2.0.6 pyyaml==6.0.2 +# Planetary Computer publishing (vendored REST client) — geo stack used by the +# PC provider's STAC builder (lazy-loaded when resolving the PC target). +pystac[validation]==1.11.0 +geopandas==1.0.1 +pyproj==3.6.1 +pyogrio==0.10.0 boto3==1.36.20 tensorboard==2.19.0 tenacity==9.1.2 diff --git a/api/hastefuncqueues/tests/__init__.py b/api/hastefuncqueues/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/api/hastefuncqueues/tests/test_publishing_handlers.py b/api/hastefuncqueues/tests/test_publishing_handlers.py new file mode 100644 index 00000000..b304b11e --- /dev/null +++ b/api/hastefuncqueues/tests/test_publishing_handlers.py @@ -0,0 +1,108 @@ +import json +import os +import unittest +from unittest.mock import Mock, patch + +import azure.functions as func + +os.environ.setdefault("DEVELOPMENT_MODE", "true") +os.environ.setdefault("METADATA_STORAGE_TYPE", "local") +os.environ.setdefault("ARTIFACT_STORAGE_TYPE", "local") +os.environ.setdefault("DATA_PATH", "/tmp/haste-publishing-queue-tests") +os.environ.setdefault("TEMP_DATA_PATH", "/tmp/haste-publishing-queue-tests") + +from api.hastefuncqueues import function_app # noqa: E402 + +PROJECT_ID = "123e4567-e89b-12d3-a456-426614174000" +DATASET_ID = "3e8d5e90-f2fc-5412-9f97-a52c07815f0b" + + +def queue_message(operation: str = "publish") -> func.QueueMessage: + return func.QueueMessage( + id="message-id", + body=json.dumps( + { + "projectId": PROJECT_ID, + "datasetId": DATASET_ID, + "operation": operation, + "attempt": 1, + } + ), + ) + + +class TestPublishingQueueHandlers(unittest.IsolatedAsyncioTestCase): + async def test_publish_handler_decodes_and_delegates(self) -> None: + processor = Mock() + with patch.object( + function_app, "PublishingProcessor", return_value=processor + ): + await function_app.GetPublishDatasetQueueMessage(queue_message()) + + message = processor.run_step.call_args.args[0] + self.assertEqual(str(message.projectId), PROJECT_ID) + self.assertEqual(str(message.datasetId), DATASET_ID) + self.assertEqual(message.operation.value, "publish") + + async def test_publish_handler_rejects_malformed_message(self) -> None: + message = func.QueueMessage(id="message-id", body="not-json") + + with self.assertRaisesRegex(RuntimeError, "JSONDecodeError") as raised: + await function_app.GetPublishDatasetQueueMessage(message) + + self.assertNotIn("not-json", str(raised.exception)) + + async def test_publish_handler_redacts_processor_exception_message( + self, + ) -> None: + processor = Mock() + processor.run_step.side_effect = RuntimeError( + "https://storage/blob?sig=secret" + ) + with patch.object( + function_app, "PublishingProcessor", return_value=processor + ): + with self.assertRaisesRegex( + RuntimeError, "RuntimeError" + ) as raised: + await function_app.GetPublishDatasetQueueMessage( + queue_message() + ) + + self.assertNotIn("secret", str(raised.exception)) + + async def test_poison_handler_delegates_current_operation(self) -> None: + processor = Mock() + with patch.object( + function_app, "PublishingProcessor", return_value=processor + ): + await function_app.GetPublishDatasetPoisonQueueMessage( + queue_message("unpublish") + ) + + message = processor.mark_poisoned.call_args.args[0] + self.assertEqual(message.operation.value, "unpublish") + + async def test_poison_handler_ignores_removed_dataset(self) -> None: + processor = Mock() + processor.mark_poisoned.side_effect = FileNotFoundError(DATASET_ID) + with patch.object( + function_app, "PublishingProcessor", return_value=processor + ): + await function_app.GetPublishDatasetPoisonQueueMessage( + queue_message() + ) + + async def test_timer_delegates_reconciliation(self) -> None: + processor = Mock() + processor.reconcile_stale.return_value = 2 + with patch.object( + function_app, "PublishingProcessor", return_value=processor + ): + await function_app.ReconcilePublishingOperations(Mock()) + + processor.reconcile_stale.assert_called_once_with() + + +if __name__ == "__main__": + unittest.main() diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 44680e30..923dce16 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -106,6 +106,10 @@ services: ZIP_QUEUE_NAME: "local-zip-queue" INFERENCE_QUEUE_NAME: "local-inference-queue" EMBEDDING_QUEUE_NAME: "local-embedding-queue" + PUBLISH_QUEUE_NAME: "local-publish-queue" + PUBLISHING_ENABLED: "true" + PC_PROVIDER_ENABLED: "false" + PUBLISHING_LOCK_CONTAINER: "publishing-locks" BLOB_CONTAINER: "data" METADATA_STORAGE_TYPE: "blob" IMAGERY_STORAGE_TYPE: "blob" @@ -150,6 +154,10 @@ services: ZIP_QUEUE_NAME: "local-zip-queue" INFERENCE_QUEUE_NAME: "local-inference-queue" EMBEDDING_QUEUE_NAME: "local-embedding-queue" + PUBLISH_QUEUE_NAME: "local-publish-queue" + PUBLISHING_ENABLED: "true" + PC_PROVIDER_ENABLED: "false" + PUBLISHING_LOCK_CONTAINER: "publishing-locks" BLOB_CONTAINER: "data" METADATA_STORAGE_TYPE: "blob" IMAGERY_STORAGE_TYPE: "blob" diff --git a/hastelib/pyproject.toml b/hastelib/pyproject.toml index 7ee97cba..c398a923 100644 --- a/hastelib/pyproject.toml +++ b/hastelib/pyproject.toml @@ -60,6 +60,19 @@ dependencies = [ "typing-extensions" ] +[project.optional-dependencies] +# Planetary Computer publishing target. Installed only where the PC provider is +# enabled; uses a vendored REST GeoCatalog client (no azure-planetarycomputer SDK). +planetary-computer = [ + "azure-identity>=1.17", + "requests>=2.31", + "pystac[validation]==1.11.0", + "geopandas>=1.0", + "pyogrio>=0.9", + "pyproj>=3.6", + "shapely>=2.0", +] + [project.urls] Documentation = "https://microsoft.github.io/haste" Issues = "https://github.com/microsoft/haste/issues" diff --git a/hastelib/src/hastegeo/core/artifact_storage/abstract_artifact_storage.py b/hastelib/src/hastegeo/core/artifact_storage/abstract_artifact_storage.py index 5b030166..439c7032 100644 --- a/hastelib/src/hastegeo/core/artifact_storage/abstract_artifact_storage.py +++ b/hastelib/src/hastegeo/core/artifact_storage/abstract_artifact_storage.py @@ -50,6 +50,41 @@ def store_artifact( def get_base_url(self): pass + @abstractmethod + def resolve_artifact_path(self, location: str) -> str: + pass + + @abstractmethod + def copy_artifact( + self, + source_path: str, + destination_path: str, + source_etag: str, + ) -> str: + pass + + @abstractmethod + def delete_prefix(self, prefix: str) -> int: + pass + + @abstractmethod + def artifact_exists(self, artifact_path: str) -> bool: + pass + + @abstractmethod + def get_artifact_size(self, artifact_path: str) -> int: + pass + + @abstractmethod + def get_artifact_etag(self, artifact_path: str) -> str: + pass + + @abstractmethod + def get_scoped_download_url( + self, artifact_path: str, expires_minutes: int = 15 + ) -> str: + pass + def is_json(self, data): """Check if data can be serialized as JSON. diff --git a/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py b/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py index c3376a2f..16e4c5cf 100644 --- a/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py +++ b/hastelib/src/hastegeo/core/artifact_storage/azure_blob_artifact_storage.py @@ -2,22 +2,33 @@ # Licensed under the MIT License. import json import os +import time from datetime import datetime, timedelta, timezone +from hashlib import sha256 +from pathlib import PurePosixPath +from threading import Lock +from urllib.parse import unquote, urlparse import yaml +from azure.core import MatchConditions from azure.core.exceptions import ResourceExistsError from azure.identity import DefaultAzureCredential # type: ignore from azure.storage.blob import BlobClient # type: ignore from azure.storage.blob import ( AccessPolicy, + BlobSasPermissions, BlobServiceClient, ContainerSasPermissions, + generate_blob_sas, generate_container_sas, ) from hastegeo.core.utils.logs import Logger from .abstract_artifact_storage import AbstractArtifactStorage +_INITIALIZED_CONTAINERS = set() +_INITIALIZED_CONTAINERS_LOCK = Lock() + class AzureBlobArtifactStorage(AbstractArtifactStorage): def __init__( @@ -38,11 +49,20 @@ def __init__( ) self.user_delegation_key = None self.account_key = self.blob_service_client.credential.account_key + self.identity_blob_service_client = ( + BlobServiceClient( + account_url=account_url, + credential=DefaultAzureCredential(), + ) + if account_url and urlparse(account_url).scheme == "https" + else None + ) else: credential = DefaultAzureCredential() self.blob_service_client = BlobServiceClient( account_url=account_url, credential=credential ) + self.identity_blob_service_client = self.blob_service_client self.user_delegation_key = ( self.blob_service_client.get_user_delegation_key( datetime.now(timezone.utc), @@ -51,22 +71,28 @@ def __init__( ) self.account_key = None - try: - # Create the container - self.container_client = self.blob_service_client.create_container( - container - ) - self.logger.info(f"Container '{container}' created successfully.") - except ResourceExistsError: - self.logger.info(f"Container '{container}' already exists.") - self.container_client = ( - self.blob_service_client.get_container_client(container) - ) - self.container_read_policy = container_read_policy_name self.blob_read_policy = blob_read_policy_name self.sas_expiration_days = 90 - self._create_or_update_managed_access_policy() + self.container_client = self.blob_service_client.get_container_client( + container + ) + cache_key = ( + self.blob_service_client.url, + container, + self.container_read_policy, + ) + with _INITIALIZED_CONTAINERS_LOCK: + if cache_key not in _INITIALIZED_CONTAINERS: + try: + self.container_client.create_container() + self.logger.info( + f"Container '{container}' created successfully." + ) + except ResourceExistsError: + self.logger.info(f"Container '{container}' already exists.") + self._create_or_update_managed_access_policy() + _INITIALIZED_CONTAINERS.add(cache_key) def _create_or_update_managed_access_policy(self): expiration_date = datetime.now(timezone.utc) + timedelta( @@ -213,12 +239,12 @@ def store_artifact( FileNotFoundError: If src_path is provided but doesn't exist. """ # Validate inputs early - if not src_path and not data: + if src_path is None and data is None: raise ValueError( "Either src_path or data must be provided to store the artifact." ) - if src_path and not os.path.exists(src_path): + if src_path is not None and not os.path.exists(src_path): raise FileNotFoundError(f"Source path {src_path} does not exist.") # Get destination path @@ -228,7 +254,7 @@ def store_artifact( blob_client = self.container_client.get_blob_client(dst_path) try: - if src_path: + if src_path is not None: with open(src_path, "rb") as file_data: blob_client.upload_blob(file_data, overwrite=True) self.logger.info( @@ -262,3 +288,166 @@ def store_artifact( def get_base_url(self): return f"https://{self.container_client.account_name}.blob.core.windows.net/{self.container_client.container_name}" + + def resolve_artifact_path(self, location: str) -> str: + parsed = urlparse(location) + if parsed.scheme: + container_url = urlparse(self.container_client.url) + if parsed.netloc.lower() != container_url.netloc.lower(): + raise ValueError("Artifact URL does not belong to configured storage") + container_path = container_url.path.rstrip("/") + "/" + if not parsed.path.startswith(container_path): + raise ValueError("Artifact URL does not belong to configured container") + location = unquote(parsed.path[len(container_path) :]) + + normalized = str(PurePosixPath(location.lstrip("/"))) + if not normalized or normalized == "." or ".." in PurePosixPath(normalized).parts: + raise ValueError("Invalid artifact path") + return normalized + + def copy_artifact( + self, + source_path: str, + destination_path: str, + source_etag: str, + ) -> str: + source_relative = self.resolve_artifact_path(source_path) + destination_relative = self.resolve_artifact_path(destination_path) + source_client = self.container_client.get_blob_client(source_relative) + destination_client = self.container_client.get_blob_client( + destination_relative + ) + if not source_client.exists(): + raise FileNotFoundError(source_path) + + current_source_etag = str(source_client.get_blob_properties().etag) + if current_source_etag != source_etag: + raise RuntimeError("Source artifact changed before copy") + + if destination_client.exists(): + destination_properties = destination_client.get_blob_properties() + copy_status = getattr(destination_properties.copy, "status", None) + if copy_status == "pending": + self._wait_for_copy(destination_client, destination_relative) + + source_url = self.get_scoped_download_url( + source_relative, expires_minutes=15 + ) + source_etag_digest = sha256(source_etag.encode("utf-8")).hexdigest() + destination_client.start_copy_from_url( + source_url, + metadata={"hastesourceetag": source_etag_digest}, + source_etag=source_etag, + source_match_condition=MatchConditions.IfNotModified, + ) + return self._wait_for_copy( + destination_client, + destination_relative, + expected_source_etag_digest=source_etag_digest, + ) + + def _wait_for_copy( + self, + blob_client, + destination_path: str, + timeout_seconds: int = 60, + expected_source_etag_digest: str = None, + ) -> str: + deadline = time.monotonic() + timeout_seconds + while time.monotonic() < deadline: + properties = blob_client.get_blob_properties() + status = getattr(properties.copy, "status", None) + if status == "success": + if ( + expected_source_etag_digest is not None + and (properties.metadata or {}).get("hastesourceetag") + != expected_source_etag_digest + ): + raise RuntimeError( + "Published artifact source revision is unverified" + ) + return destination_path + if status in {"failed", "aborted"}: + description = getattr( + properties.copy, "status_description", "unknown error" + ) + raise RuntimeError(f"Blob copy failed: {description}") + time.sleep(0.25) + raise TimeoutError(f"Blob copy did not finish: {destination_path}") + + def delete_prefix(self, prefix: str) -> int: + relative_prefix = self.resolve_artifact_path(prefix).rstrip("/") + "/" + blob_names = list( + self.container_client.list_blob_names( + name_starts_with=relative_prefix + ) + ) + for blob_name in blob_names: + self.container_client.delete_blob(blob_name) + return len(blob_names) + + def artifact_exists(self, artifact_path: str) -> bool: + relative_path = self.resolve_artifact_path(artifact_path) + return self.container_client.get_blob_client(relative_path).exists() + + def get_artifact_size(self, artifact_path: str) -> int: + relative_path = self.resolve_artifact_path(artifact_path) + blob_client = self.container_client.get_blob_client(relative_path) + if not blob_client.exists(): + raise FileNotFoundError(artifact_path) + return blob_client.get_blob_properties().size + + def get_artifact_etag(self, artifact_path: str) -> str: + relative_path = self.resolve_artifact_path(artifact_path) + blob_client = self.container_client.get_blob_client(relative_path) + if not blob_client.exists(): + raise FileNotFoundError(artifact_path) + return str(blob_client.get_blob_properties().etag) + + def get_scoped_download_url( + self, artifact_path: str, expires_minutes: int = 15 + ) -> str: + if expires_minutes < 5 or expires_minutes > 60: + raise ValueError("expires_minutes must be between 5 and 60") + relative_path = self.resolve_artifact_path(artifact_path) + blob_client = self.container_client.get_blob_client(relative_path) + if not blob_client.exists(): + raise FileNotFoundError(artifact_path) + + now = datetime.now(timezone.utc) + expiry = now + timedelta(minutes=expires_minutes) + is_emulator = ( + self.container_client.account_name == "devstoreaccount1" + or urlparse(self.container_client.url).hostname + in {"127.0.0.1", "localhost", "azurite"} + ) + user_delegation_key = None + account_key = self.account_key + delegation_client = self.blob_service_client + if account_key and not is_emulator: + delegation_client = self.identity_blob_service_client + if delegation_client is None: + raise RuntimeError( + "Managed identity is required for published downloads" + ) + account_key = None + if account_key is None: + user_delegation_key = ( + delegation_client.get_user_delegation_key( + now - timedelta(minutes=5), + expiry + timedelta(minutes=5), + ) + ) + + sas_token = generate_blob_sas( + account_name=self.container_client.account_name, + container_name=self.container_client.container_name, + blob_name=relative_path, + permission=BlobSasPermissions(read=True), + start=now - timedelta(minutes=5), + expiry=expiry, + account_key=account_key, + user_delegation_key=user_delegation_key, + protocol="https,http" if is_emulator else "https", + ) + return f"{blob_client.url}?{sas_token}" diff --git a/hastelib/src/hastegeo/core/artifact_storage/local_file_system_artifact_storage.py b/hastelib/src/hastegeo/core/artifact_storage/local_file_system_artifact_storage.py index b3ab340e..b05737c5 100644 --- a/hastelib/src/hastegeo/core/artifact_storage/local_file_system_artifact_storage.py +++ b/hastelib/src/hastegeo/core/artifact_storage/local_file_system_artifact_storage.py @@ -1,7 +1,14 @@ # Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. +import json import os import shutil +import tempfile +from hashlib import sha256 +from pathlib import Path +from urllib.parse import unquote, urlparse + +from hastegeo.core.utils.logs import Logger from .abstract_artifact_storage import AbstractArtifactStorage @@ -9,6 +16,7 @@ class LocalFileSystemArtifactStorage(AbstractArtifactStorage): def __init__(self, partition_key=None, **kwargs): super().__init__(partition_key) + self.logger = Logger.get_logger(__name__) if partition_key is None: partition_key = "" directory = os.path.join(kwargs.pop("directory"), partition_key) @@ -85,12 +93,12 @@ def store_artifact( FileNotFoundError: If src_path is provided but doesn't exist. """ # Validate inputs early - if not src_path and not data: + if src_path is None and data is None: raise ValueError( "Either src_path or data must be provided to store the artifact." ) - if src_path and not os.path.exists(src_path): + if src_path is not None and not os.path.exists(src_path): raise FileNotFoundError(f"Source path {src_path} does not exist.") # Get destination path @@ -104,27 +112,121 @@ def store_artifact( os.makedirs(dst_dir, exist_ok=True) try: - if src_path: + if src_path is not None: # Handle file/directory copying if os.path.isdir(src_path): if os.path.exists(dst_path): shutil.rmtree(dst_path) # Remove existing directory shutil.copytree(src_path, dst_path) - print(f"Copied directory '{src_path}' to '{dst_path}'") + self.logger.info( + f"Copied directory '{src_path}' to '{dst_path}'" + ) else: shutil.copy2(src_path, dst_path) - print(f"Copied file '{src_path}' to '{dst_path}'") + self.logger.info( + f"Copied file '{src_path}' to '{dst_path}'" + ) else: - # Handle string data writing with open(dst_path, "w", encoding="utf-8") as file: - file.write(data) - print(f"Wrote data to file '{dst_path}'") + if isinstance(data, str): + file.write(data) + else: + json.dump(data, file) + self.logger.info(f"Wrote data to file '{dst_path}'") return dst_path except Exception as e: - print(f"Failed to store artifact at '{dst_path}': {e}") + self.logger.error(f"Failed to store artifact at '{dst_path}': {e}") raise def get_base_url(self): return self.directory + + def resolve_artifact_path(self, location: str) -> str: + parsed = urlparse(location) + raw_path = unquote(parsed.path) if parsed.scheme == "file" else location + candidate = Path(raw_path) + if not candidate.is_absolute(): + candidate = Path(self.directory, candidate) + resolved = candidate.resolve() + root = Path(self.directory).resolve() + if resolved != root and root not in resolved.parents: + raise ValueError("Artifact path escapes the configured storage root") + return str(resolved.relative_to(root)) + + def copy_artifact( + self, + source_path: str, + destination_path: str, + source_etag: str, + ) -> str: + source_relative = self.resolve_artifact_path(source_path) + destination_relative = self.resolve_artifact_path(destination_path) + source = Path(self.directory, source_relative) + destination = Path(self.directory, destination_relative) + if not source.is_file(): + raise FileNotFoundError(source_path) + if self.get_artifact_etag(source_relative) != source_etag: + raise RuntimeError("Source artifact changed before copy") + destination.parent.mkdir(parents=True, exist_ok=True) + descriptor, temporary_name = tempfile.mkstemp( + dir=destination.parent, + prefix=f".{destination.name}.", + ) + os.close(descriptor) + temporary_path = Path(temporary_name) + try: + shutil.copy2(source, temporary_path) + if self._hash_file(temporary_path) != source_etag: + raise RuntimeError("Source artifact changed during copy") + os.replace(temporary_path, destination) + finally: + temporary_path.unlink(missing_ok=True) + return destination_relative + + def delete_prefix(self, prefix: str) -> int: + relative_prefix = self.resolve_artifact_path(prefix) + target = Path(self.directory, relative_prefix) + if target.is_file(): + target.unlink() + return 1 + if not target.exists(): + return 0 + deleted = sum(1 for path in target.rglob("*") if path.is_file()) + shutil.rmtree(target) + return deleted + + def artifact_exists(self, artifact_path: str) -> bool: + relative_path = self.resolve_artifact_path(artifact_path) + return Path(self.directory, relative_path).is_file() + + def get_artifact_size(self, artifact_path: str) -> int: + relative_path = self.resolve_artifact_path(artifact_path) + path = Path(self.directory, relative_path) + if not path.is_file(): + raise FileNotFoundError(artifact_path) + return path.stat().st_size + + def get_artifact_etag(self, artifact_path: str) -> str: + relative_path = self.resolve_artifact_path(artifact_path) + path = Path(self.directory, relative_path) + if not path.is_file(): + raise FileNotFoundError(artifact_path) + return self._hash_file(path) + + @staticmethod + def _hash_file(path: Path) -> str: + digest = sha256() + with path.open("rb") as source: + for chunk in iter(lambda: source.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + def get_scoped_download_url( + self, artifact_path: str, expires_minutes: int = 15 + ) -> str: + relative_path = self.resolve_artifact_path(artifact_path) + if not self.artifact_exists(relative_path): + raise FileNotFoundError(artifact_path) + return Path(self.directory, relative_path).resolve().as_uri() diff --git a/hastelib/src/hastegeo/core/artifact_storage/unified_artifact_storage.py b/hastelib/src/hastegeo/core/artifact_storage/unified_artifact_storage.py index f65f5375..be97db91 100644 --- a/hastelib/src/hastegeo/core/artifact_storage/unified_artifact_storage.py +++ b/hastelib/src/hastegeo/core/artifact_storage/unified_artifact_storage.py @@ -134,3 +134,35 @@ def store_artifact( def get_base_url(self): return self.artifact_storage.get_base_url() + + def resolve_artifact_path(self, location: str) -> str: + return self.artifact_storage.resolve_artifact_path(location) + + def copy_artifact( + self, + source_path: str, + destination_path: str, + source_etag: str, + ) -> str: + return self.artifact_storage.copy_artifact( + source_path, destination_path, source_etag + ) + + def delete_prefix(self, prefix: str) -> int: + return self.artifact_storage.delete_prefix(prefix) + + def artifact_exists(self, artifact_path: str) -> bool: + return self.artifact_storage.artifact_exists(artifact_path) + + def get_artifact_size(self, artifact_path: str) -> int: + return self.artifact_storage.get_artifact_size(artifact_path) + + def get_artifact_etag(self, artifact_path: str) -> str: + return self.artifact_storage.get_artifact_etag(artifact_path) + + def get_scoped_download_url( + self, artifact_path: str, expires_minutes: int = 15 + ) -> str: + return self.artifact_storage.get_scoped_download_url( + artifact_path, expires_minutes=expires_minutes + ) diff --git a/hastelib/src/hastegeo/core/config.py b/hastelib/src/hastegeo/core/config.py index 30a64fb5..350c7f97 100644 --- a/hastelib/src/hastegeo/core/config.py +++ b/hastelib/src/hastegeo/core/config.py @@ -15,6 +15,21 @@ _SCHEME_RE = re.compile(r"^https?://", re.IGNORECASE) +def _get_bool_env(name, default=False): + value = os.getenv(name) + if value is None: + return default + return value.strip().lower() in {"1", "true", "yes", "on"} + + +def _get_bounded_int_env(name, default, minimum, maximum=None): + value = int(os.getenv(name, str(default))) + if value < minimum or (maximum is not None and value > maximum): + upper = f" and {maximum}" if maximum is not None else "" + raise ValueError(f"{name} must be between {minimum}{upper}") + return value + + def _strip_scheme(value): """Reduce a registry URL to the bare login server. @@ -250,6 +265,7 @@ def __init__(self, env=None): }, } self.queue_config = self.get_queue_config() + self.publishing_config = self.get_publishing_config() self.storage_type = os.getenv("METADATA_STORAGE_TYPE", "local").lower() self.artifact_storage_type = os.getenv( "ARTIFACT_STORAGE_TYPE", "local" @@ -302,6 +318,44 @@ def get_queue_config(): "embedding_queue_name": os.getenv( "EMBEDDING_QUEUE_NAME", "embedding-queue" ), + "publish_queue_name": os.getenv( + "PUBLISH_QUEUE_NAME", "publish-queue" + ), + } + + @staticmethod + def get_publishing_config(): + """Get publishing feature and provider configuration.""" + return { + "publishing_enabled": _get_bool_env( + "PUBLISHING_ENABLED", True + ), + "pc_provider_enabled": _get_bool_env( + "PC_PROVIDER_ENABLED", False + ), + "max_total_bytes": _get_bounded_int_env( + "PUBLISH_MAX_TOTAL_BYTES", 5 * 1024**3, 1 + ), + "download_sas_minutes": _get_bounded_int_env( + "PUBLISHED_DOWNLOAD_SAS_MINUTES", 15, 5, 60 + ), + "pc_geocatalog_url": os.getenv("PC_GEOCATALOG_URL", ""), + "pc_ingestion_source": os.getenv("PC_INGESTION_SOURCE", ""), + "pc_collection_prefix": os.getenv( + "PC_COLLECTION_PREFIX", "haste-" + ), + "pc_explorer_url": os.getenv("PC_EXPLORER_URL", ""), + "pc_publishing_license": os.getenv( + "PC_PUBLISHING_LICENSE", "CC-BY-4.0" + ), + "pc_verify_attempts": _get_bounded_int_env( + "PC_VERIFY_ATTEMPTS", 20, 1, 60 + ), + "lease_connection_string": os.getenv("AzureWebJobsStorage"), + "lease_account_url": os.getenv("BLOB_ACCOUNT_URL"), + "lease_container": os.getenv( + "PUBLISHING_LOCK_CONTAINER", "publishing-locks" + ), } @staticmethod @@ -325,6 +379,7 @@ class DataTypes(Enum): CONFIG = "config" MODEL = "model" MODEL_CATALOG = "model_catalog" + PUBLISHED_DATASET = "published_dataset" MODEL_ARTIFACTS = "artifacts_model" VISUALIZER = "visualizer_imagery" TRAIN_LABELS = "train_labels" diff --git a/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py b/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py index ea50cfc3..b6f743fb 100644 --- a/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/abstract_data_layer.py @@ -184,6 +184,14 @@ def load_all_from_partition(self, data_type, data_format="json"): """ pass + def load_bounded( + self, data_type, max_records, data_format="json" + ): + """Load no more than ``max_records`` or fail before full materialization.""" + raise NotImplementedError( + f"{self.__class__.__name__} does not support bounded reads" + ) + @abstractmethod def delete(self, identifier, data_type, data_format="json"): """Delete a specific data record from the storage backend. diff --git a/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py index d9bff6e4..582bcc2c 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_blob_storage_data_layer.py @@ -4,11 +4,12 @@ import json import logging import os +from concurrent.futures import ThreadPoolExecutor from datetime import datetime, timedelta, timezone +from threading import Lock import yaml -from abstract_data_layer import AbstractDataLayer -from azure.core.exceptions import ResourceExistsError +from azure.core.exceptions import ResourceExistsError, ResourceNotFoundError from azure.identity import DefaultAzureCredential # type: ignore from azure.storage.blob import AccessPolicy # type: ignore from azure.storage.blob import BlobBlock # type: ignore @@ -18,6 +19,11 @@ generate_container_sas, ) +from .abstract_data_layer import AbstractDataLayer + +_INITIALIZED_CONTAINERS = set() +_INITIALIZED_CONTAINERS_LOCK = Lock() + class AzureBlobStorageDataLayer(AbstractDataLayer): def __init__( @@ -49,20 +55,26 @@ def __init__( ) self.account_key = None - try: - # Create the container - self.container_client = self.blob_service_client.create_container( - container - ) - logging.info(f"Container '{container}' created successfully.") - except ResourceExistsError: - logging.info(f"Container '{container}' already exists.") - self.container_client = ( - self.blob_service_client.get_container_client(container) - ) - self.container_read_policy = container_read_policy_name - self._create_or_update_managed_access_policy() + self.container_client = self.blob_service_client.get_container_client( + container + ) + cache_key = ( + self.blob_service_client.url, + container, + self.container_read_policy, + ) + with _INITIALIZED_CONTAINERS_LOCK: + if cache_key not in _INITIALIZED_CONTAINERS: + try: + self.container_client.create_container() + logging.info( + f"Container '{container}' created successfully." + ) + except ResourceExistsError: + logging.info(f"Container '{container}' already exists.") + self._create_or_update_managed_access_policy() + _INITIALIZED_CONTAINERS.add(cache_key) def _create_or_update_managed_access_policy(self): expiration_days = 90 @@ -177,7 +189,11 @@ def save( if self.is_bytes(data): blob_client.upload_blob(data, overwrite=True) elif self.is_json(data): - blob_client.upload_blob(json.dumps(data), overwrite=True) + blob_client.upload_blob( + json.dumps(data), + overwrite=True, + metadata=self._index_metadata(data_type, data), + ) elif self.is_yaml(data): logging.info("data is yaml, dumping to blob") blob_client.upload_blob( @@ -346,6 +362,142 @@ def load_all_from_partition(self, data_type, data_format="json"): data.append(yaml.safe_load(downloader.readall())) return data + def load_bounded(self, data_type, max_records, data_format="json"): + records, _ = self.load_page( + data_type=data_type, + page=1, + page_size=max_records, + data_format=data_format, + max_records=max_records, + ) + return records + + def load_page( + self, + data_type, + page, + page_size, + data_format="json", + target=None, + status=None, + project_id=None, + max_records=None, + ): + if page < 1 or page_size < 1: + raise ValueError("page and page_size must be positive") + if max_records is not None and max_records < 1: + raise ValueError("max_records must be positive") + + prefix = ( + f"{self.partition_key}/{data_type}_" + if self.partition_key + else None + ) + matching_blobs = [] + catalog_record_count = 0 + scanned_blob_count = 0 + scan_limit = max_records * 10 if max_records is not None else None + list_options = { + "name_starts_with": prefix, + "include": ["metadata"], + } + if max_records is not None: + list_options["results_per_page"] = max_records + 1 + blobs = self.container_client.list_blobs(**list_options) + pages = blobs.by_page() + for blob_page in pages: + for blob in blob_page: + scanned_blob_count += 1 + if scan_limit is not None and scanned_blob_count > scan_limit: + raise ValueError( + "Catalog storage scan exceeds the bounded envelope" + ) + parts = blob.name.split("/") + if "stats" in blob.name or len(parts) > 2: + continue + if self.partition_key: + matches = blob.name.startswith( + f"{self.partition_key}/{data_type}_" + ) + else: + matches = parts[-1].startswith(f"{data_type}_") + if not matches or not blob.name.endswith(f".{data_format}"): + continue + catalog_record_count += 1 + if ( + max_records is not None + and catalog_record_count > max_records + ): + raise ValueError( + f"Catalog exceeds the {max_records:,}-record limit" + ) + metadata = blob.metadata or {} + if target and metadata.get("hastetarget") != target: + continue + if status and metadata.get("hastestatus") != status: + continue + if project_id and metadata.get("hasteproject") != project_id: + continue + matching_blobs.append(blob) + + matching_blobs.sort( + key=lambda blob: ( + (blob.metadata or {}).get("hastesort") + or blob.last_modified.isoformat(), + blob.name, + ), + reverse=True, + ) + total_count = len(matching_blobs) + start = (page - 1) * page_size + page_names = [ + blob.name for blob in matching_blobs[start : start + page_size] + ] + records = self._load_blob_names(page_names, data_format) + total_count -= len(page_names) - len(records) + return records, total_count + + @staticmethod + def _index_metadata(data_type, data): + if data_type != "published_dataset" or not isinstance(data, dict): + return None + return { + "hastesort": str( + data.get("publishedDate") or data.get("createdDate") or "" + ), + "hastetarget": str(data.get("target") or ""), + "hastestatus": str(data.get("status") or ""), + "hasteproject": str(data.get("projectId") or ""), + } + + def _load_blob_names(self, blob_names, data_format): + if not blob_names: + return [] + + missing_blob = object() + + def load_blob(blob_name): + try: + downloader = self.container_client.get_blob_client( + blob_name + ).download_blob() + except ResourceNotFoundError: + return missing_blob + contents = downloader.readall() + if data_format == "json": + contents = json.loads(contents) + if isinstance(contents, str): + contents = json.loads(contents) + return contents + if data_format == "yaml": + return yaml.safe_load(contents) + raise ValueError(f"Unsupported data format: {data_format}") + + workers = min(32, len(blob_names)) + with ThreadPoolExecutor(max_workers=workers) as executor: + records = list(executor.map(load_blob, blob_names)) + return [record for record in records if record is not missing_blob] + def delete(self, identifier, data_type, data_format="json"): blob_name = self.get_file_path(identifier, data_type, data_format) blob_client = self.container_client.get_blob_client(blob_name) diff --git a/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py index 188cbf9c..29aeba16 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_cosmos_db_data_layer.py @@ -2,10 +2,11 @@ # Licensed under the MIT License. import re -from abstract_data_layer import AbstractDataLayer from azure.cosmos import CosmosClient, exceptions # type: ignore from azure.identity import DefaultAzureCredential # type: ignore +from .abstract_data_layer import AbstractDataLayer + class AzureCosmosDBDataLayer(AbstractDataLayer): def __init__(self, endpoint, database, container, partition_key=None): @@ -124,6 +125,34 @@ def load_all_from_partition(self, data_type): ) return items + def load_bounded(self, data_type, max_records, data_format="json"): + if ( + data_format != "json" + or not isinstance(max_records, int) + or not 1 <= max_records <= 10000 + ): + raise ValueError("Invalid bounded Cosmos DB read") + id_prefix = self._id_prefix(data_type) + query = ( + f"SELECT TOP {max_records + 1} * FROM c " + "WHERE STARTSWITH(c.id, @id_prefix)" + ) + items = list( + self.container.query_items( + query=query, + parameters=[ + {"name": "@id_prefix", "value": id_prefix}, + ], + enable_cross_partition_query=True, + max_item_count=max_records + 1, + ) + ) + if len(items) > max_records: + raise ValueError( + f"Metadata exceeds the {max_records:,}-record limit" + ) + return items + def delete(self, identifier, data_type, data_format="json"): partition_key = ( self.partition_key if self.partition_key else identifier diff --git a/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py index 749458a6..5cf4cfa7 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_data_lake_data_layer.py @@ -2,10 +2,11 @@ # Licensed under the MIT License. import json -from abstract_data_layer import AbstractDataLayer from azure.identity import DefaultAzureCredential # type: ignore from azure.storage.filedatalake import DataLakeServiceClient # type: ignore +from .abstract_data_layer import AbstractDataLayer + class AzureDataLakeDataLayer(AbstractDataLayer): def __init__(self, account_url, file_system, partition_key=None): @@ -138,6 +139,31 @@ def load_all_from_partition(self, data_type): data = self.load_all(data_type) return data + def load_bounded(self, data_type, max_records, data_format="json"): + if data_format != "json" or max_records < 1: + raise ValueError("Invalid bounded Data Lake read") + data = [] + scanned_paths = 0 + scan_limit = max_records * 10 + for path in self.file_system_client.get_paths(): + scanned_paths += 1 + if scanned_paths > scan_limit: + raise ValueError("Metadata scan exceeds the bounded envelope") + parts = path.name.split("/") + if len(parts) > 2 or not parts[-1].startswith(f"{data_type}_"): + continue + file_contents = ( + self.file_system_client.get_file_client(path.name) + .download_file() + .readall() + ) + data.append(json.loads(file_contents)) + if len(data) > max_records: + raise ValueError( + f"Metadata exceeds the {max_records:,}-record limit" + ) + return data + def delete(self, identifier, data_type, data_format="json"): file_name = self.get_file_path( identifier, data_type, data_format=data_format diff --git a/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py b/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py index b54267aa..9481dcd8 100644 --- a/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/azure_postgresql_data_layer.py @@ -5,11 +5,12 @@ import re import psycopg2 # type: ignore -from abstract_data_layer import AbstractDataLayer from azure.identity import DefaultAzureCredential # type: ignore from psycopg2 import sql # type: ignore from rasterio.io import MemoryFile +from .abstract_data_layer import AbstractDataLayer + class AzurePostgreSQLDataLayer(AbstractDataLayer): def __init__(self, host, database, table, partition_key=None, user=None): @@ -229,6 +230,25 @@ def load_all_from_partition(self, data_type): results = cursor.fetchall() return [json.loads(result[0]) for result in results] + def load_bounded(self, data_type, max_records, data_format="json"): + if data_format != "json" or max_records < 1: + raise ValueError("Invalid bounded PostgreSQL read") + connection_string = f"host={self.server_name} dbname={self.database_name} user={self.postgres_user} password={self.token} sslmode=require" + with psycopg2.connect(connection_string) as connection: + with connection.cursor() as cursor: + cursor.execute( + sql.SQL( + "SELECT data FROM {} WHERE data_type = %s LIMIT %s" + ).format(self._table_identifier()), + (data_type, max_records + 1), + ) + results = cursor.fetchall() + if len(results) > max_records: + raise ValueError( + f"Metadata exceeds the {max_records:,}-record limit" + ) + return [json.loads(result[0]) for result in results] + def delete(self, identifier, data_type): partition_key = ( self.partition_key if self.partition_key else identifier diff --git a/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py b/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py index 9bb54d89..a33ff2c0 100644 --- a/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py +++ b/hastelib/src/hastegeo/core/data_layer/local_file_system_data_layer.py @@ -5,7 +5,8 @@ import shutil import yaml -from abstract_data_layer import AbstractDataLayer + +from .abstract_data_layer import AbstractDataLayer class LocalFileSystemDataLayer(AbstractDataLayer): @@ -290,6 +291,49 @@ def load_all_from_partition(self, data_type, data_format="json"): data = self.load_all(data_type=data_type, data_format=data_format) return data + def load_bounded(self, data_type, max_records, data_format="json"): + if max_records < 1 or data_format not in {"json", "yaml"}: + raise ValueError("Invalid bounded local metadata read") + records = [] + scanned_paths = 0 + scan_limit = max_records * 10 + directories = [self.directory] + with os.scandir(self.directory) as root_entries: + for entry in root_entries: + scanned_paths += 1 + if scanned_paths > scan_limit: + raise ValueError( + "Metadata scan exceeds the bounded envelope" + ) + if entry.is_dir(follow_symlinks=False): + directories.append(entry.path) + + for directory in directories: + with os.scandir(directory) as entries: + for entry in entries: + scanned_paths += 1 + if scanned_paths > scan_limit: + raise ValueError( + "Metadata scan exceeds the bounded envelope" + ) + if not entry.is_file(follow_symlinks=False): + continue + if not entry.name.startswith(f"{data_type}_") or not entry.name.endswith( + f".{data_format}" + ): + continue + with open(entry.path, "r") as file: + records.append( + json.load(file) + if data_format == "json" + else yaml.safe_load(file) + ) + if len(records) > max_records: + raise ValueError( + f"Metadata exceeds the {max_records:,}-record limit" + ) + return records + def delete(self, identifier, data_type, data_format="json"): file_path = self.get_file_path(identifier, data_type, data_format) if not os.path.exists(file_path): diff --git a/hastelib/src/hastegeo/core/data_layer/unified.py b/hastelib/src/hastegeo/core/data_layer/unified.py index 2e1752e9..20e35fe1 100644 --- a/hastelib/src/hastegeo/core/data_layer/unified.py +++ b/hastelib/src/hastegeo/core/data_layer/unified.py @@ -35,7 +35,9 @@ def __init__(self, storage_type, partition_key=None, **kwargs): if storage_type in storage_class_map: module_name, class_name = storage_class_map[storage_type] - module = importlib.import_module(module_name) + module = importlib.import_module( + f"{__package__}.{module_name}" + ) data_layer_class = getattr(module, class_name) self.data_layer = data_layer_class( partition_key=self.partition_key, **kwargs @@ -123,6 +125,35 @@ def load_all_from_partition(self, data_type, data_format="json"): data_type, data_format=data_format ) + def load_bounded(self, data_type, max_records, data_format="json"): + return self.data_layer.load_bounded( + data_type=data_type, + max_records=max_records, + data_format=data_format, + ) + + def load_page( + self, + data_type, + page, + page_size, + data_format="json", + target=None, + status=None, + project_id=None, + max_records=None, + ): + return self.data_layer.load_page( + data_type=data_type, + page=page, + page_size=page_size, + data_format=data_format, + target=target, + status=status, + project_id=project_id, + max_records=max_records, + ) + def delete(self, identifier, data_type, data_format="json"): self.data_layer.delete(identifier, data_type, data_format=data_format) diff --git a/hastelib/src/hastegeo/core/models/publishing.py b/hastelib/src/hastegeo/core/models/publishing.py new file mode 100644 index 00000000..347374bb --- /dev/null +++ b/hastelib/src/hastegeo/core/models/publishing.py @@ -0,0 +1,229 @@ +import hashlib +import json +import unicodedata +import uuid +from enum import Enum +from typing import Any, Dict, List, Optional + +from pydantic import BaseModel, ConfigDict, Field, field_validator + +PUBLISHING_UUID_NAMESPACE = uuid.NAMESPACE_URL +PUBLISHING_UUID_NAME_PREFIX = ( + "https://github.com/microsoft/haste/data-publishing" +) + + +class PublishTarget(str, Enum): + LOCAL = "local" + PLANETARY_COMPUTER = "planetary_computer" + + +class PublishStatus(str, Enum): + PENDING = "PENDING" + IN_PROGRESS = "IN_PROGRESS" + PUBLISHED = "PUBLISHED" + FAILED = "FAILED" + UNPUBLISH_PENDING = "UNPUBLISH_PENDING" + UNPUBLISHING = "UNPUBLISHING" + UNPUBLISH_FAILED = "UNPUBLISH_FAILED" + + +class PublishOperation(str, Enum): + PUBLISH = "publish" + UNPUBLISH = "unpublish" + + +class ArtifactKind(str, Enum): + GPKG = "gpkg" + VALID_MASK = "valid_mask" + FOOTPRINTS = "footprints" + PROCESSED_COG = "processed_cog" + + +class ProviderConfigField(BaseModel): + key: str + label: str + required: bool + secret: bool = False + + +class ProviderInfo(BaseModel): + id: str + displayName: str + description: str = "" + isEnabled: bool + isConfigured: bool + disabledReason: Optional[str] = None + supportsAsync: bool = True + supportedArtifactKinds: List[ArtifactKind] = Field(default_factory=list) + requiredSupportingArtifactKinds: List[ArtifactKind] = Field( + default_factory=list + ) + configRequirements: List[ProviderConfigField] = Field(default_factory=list) + + +class SourceArtifact(BaseModel): + kind: ArtifactKind + sourcePath: str + mediaType: str + sizeBytes: Optional[int] = Field(default=None, ge=0) + sourceEtag: str = Field(min_length=1, max_length=256) + + +class PublishedArtifact(SourceArtifact): + publishedPath: str + + +class PublishDatasetOptions(BaseModel): + projectId: uuid.UUID + projectName: str + imageLayerId: str + imageLayerName: str + modelId: str + modelName: str + defaultName: str + availableArtifacts: List[SourceArtifact] = Field(default_factory=list) + + +class ArtifactBundle(BaseModel): + selectedArtifacts: List[SourceArtifact] = Field(default_factory=list) + supportingArtifacts: List[SourceArtifact] = Field(default_factory=list) + # Optional post-event preview used as a STAC thumbnail. Stored as the URL + # HASTE holds for the image layer's preview (may carry a SAS); the PC + # provider resolves it to a plain blob path, copies it into the published + # prefix, and attaches it as a best-effort thumbnail asset. + thumbnailUrl: Optional[str] = None + + def get(self, kind: ArtifactKind) -> Optional[SourceArtifact]: + for artifact in self.selectedArtifacts + self.supportingArtifacts: + if artifact.kind == kind: + return artifact + return None + + +class PublishRequest(BaseModel): + model_config = ConfigDict(extra="forbid") + + requestId: uuid.UUID + projectId: uuid.UUID + imageLayerId: str = Field( + min_length=1, + max_length=64, + pattern=r"^[A-Za-z0-9_-]+$", + ) + modelId: str = Field( + min_length=1, + max_length=8, + pattern=r"^[0-9]+$", + ) + name: str = Field(min_length=1, max_length=200) + description: Optional[str] = Field(default=None, max_length=4000) + target: PublishTarget + artifacts: List[ArtifactKind] = Field(min_length=1) + + @field_validator("imageLayerId", "modelId", "name") + @classmethod + def normalize_required_text(cls, value: str) -> str: + normalized = unicodedata.normalize("NFC", value).strip() + if not normalized: + raise ValueError("value must not be empty") + return normalized + + @field_validator("description") + @classmethod + def normalize_description(cls, value: Optional[str]) -> Optional[str]: + if value is None: + return None + return unicodedata.normalize("NFC", value).strip() + + @field_validator("artifacts") + @classmethod + def normalize_artifacts( + cls, value: List[ArtifactKind] + ) -> List[ArtifactKind]: + return sorted(set(value), key=lambda artifact: artifact.value) + + +class PublishResult(BaseModel): + artifacts: List[PublishedArtifact] = Field(default_factory=list) + links: Dict[str, str] = Field(default_factory=dict) + providerMetadata: Dict[str, Any] = Field(default_factory=dict) + continuationToken: Optional[str] = None + isComplete: bool = True + + +class PublishedDataset(BaseModel): + schemaVersion: int = 1 + revision: int = Field(default=1, ge=1) + datasetId: uuid.UUID + requestId: uuid.UUID + requestFingerprint: str = Field(pattern=r"^[0-9a-f]{64}$") + name: str + description: str = "" + projectId: uuid.UUID + projectName: str = "" + imageLayerId: str + imageLayerName: str = "" + modelId: str + modelName: str = "" + target: PublishTarget + status: PublishStatus + statusMessage: str = "" + lastOperation: PublishOperation = PublishOperation.PUBLISH + attempt: int = Field(default=1, ge=1) + queueDispatchedAt: Optional[str] = None + reconciledAttempt: Optional[int] = Field(default=None, ge=1) + publishedByUser: str + publishedByName: Optional[str] = None + createdDate: str + updatedDate: str + publishedDate: Optional[str] = None + artifacts: List[PublishedArtifact] = Field(default_factory=list) + selectedArtifactKinds: List[ArtifactKind] = Field(default_factory=list) + sourceArtifacts: List[SourceArtifact] = Field(default_factory=list) + links: Dict[str, str] = Field(default_factory=dict) + providerMetadata: Dict[str, Any] = Field(default_factory=dict) + assessmentSummary: Dict[str, Any] = Field(default_factory=dict) + + +class PublishQueueMessage(BaseModel): + datasetId: uuid.UUID + projectId: uuid.UUID + operation: PublishOperation + attempt: int = Field(ge=1) + + +def derive_dataset_id( + project_id: uuid.UUID, request_id: uuid.UUID +) -> uuid.UUID: + name = ( + f"{PUBLISHING_UUID_NAME_PREFIX}/" + f"{str(project_id).lower()}/{str(request_id).lower()}" + ) + return uuid.uuid5(PUBLISHING_UUID_NAMESPACE, name) + + +def compute_request_fingerprint( + request: PublishRequest, publisher_id: str +) -> str: + normalized_publisher = publisher_id.strip().lower() + if not normalized_publisher: + raise ValueError("publisher_id must not be empty") + + canonical_request = { + "projectId": str(request.projectId).lower(), + "imageLayerId": request.imageLayerId, + "modelId": request.modelId, + "name": request.name, + "description": request.description or "", + "target": request.target.value, + "artifacts": [artifact.value for artifact in request.artifacts], + "publisherId": normalized_publisher, + } + encoded = json.dumps( + canonical_request, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + ).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() diff --git a/hastelib/src/hastegeo/core/processors/assessment.py b/hastelib/src/hastegeo/core/processors/assessment.py new file mode 100644 index 00000000..720301d1 --- /dev/null +++ b/hastelib/src/hastegeo/core/processors/assessment.py @@ -0,0 +1,145 @@ +import asyncio +import os +from typing import Any, Callable, Dict, Optional + +from ..config import Config +from ..utils.assessment import ( + build_assessment_inputs_from_gpkgs, + compute_assessment_report, +) +from ..utils.blob import download_blob_to_tempfile +from .metadata import MetadataProcessor + + +class AssessmentSizeLimitError(ValueError): + """Raised when assessment inputs exceed the configured byte budget.""" + + +class AssessmentReportProcessor: + """Load HASTE artifacts and compute one server-side assessment report.""" + + def __init__( + self, + config: Optional[Config] = None, + processor_factory: Callable[ + ..., MetadataProcessor + ] = MetadataProcessor, + downloader: Callable[..., Any] = download_blob_to_tempfile, + ) -> None: + self.config = config or Config() + self.processor_factory = processor_factory + self.downloader = downloader + + async def generate( + self, + project_id: str, + image_layer_id: str, + model_id: str, + threshold: float = 0.1, + min_area_m2: float = 50, + max_total_bytes: Optional[int] = None, + ) -> Dict[str, Any]: + if max_total_bytes is not None and max_total_bytes < 1: + raise ValueError("max_total_bytes must be positive") + metadata_types = self.config.get_metadata_types() + model_data = await asyncio.to_thread( + self.processor_factory( + data_type=metadata_types.MODEL.value, + partition_key=project_id, + config=self.config, + ).load, + model_id, + ) + gpkg_url = model_data.get("gpkgUrl") + if not gpkg_url: + raise FileNotFoundError( + "No inference results available for this model" + ) + + image_layer_data = await asyncio.to_thread( + self.processor_factory( + data_type=metadata_types.IMAGELAYER.value, + partition_key=project_id, + config=self.config, + ).load, + image_layer_id, + ) + footprints_url = image_layer_data.get("buildingFootprintsUrl") + if not footprints_url: + raise FileNotFoundError( + "No building footprints available for this image layer" + ) + + try: + validation_data = await asyncio.to_thread( + self.processor_factory( + data_type=metadata_types.VALIDATION.value, + partition_key=project_id, + config=self.config, + ).load, + image_layer_id, + ) + labels_dict = validation_data.get("labels") or {} + except FileNotFoundError: + labels_dict = {} + labels = [ + (building_id, value.get("label")) + for building_id, value in labels_dict.items() + if value.get("label") + ] + + downloaded_paths = [] + try: + try: + footprints_path = await self.downloader( + footprints_url, + suffix=".gpkg", + max_bytes=max_total_bytes, + ) + except ValueError as error: + if "allowed download size" in str(error): + raise AssessmentSizeLimitError( + "Assessment inputs exceed the allowed size" + ) from error + raise + downloaded_paths.append(footprints_path) + remaining_bytes = ( + None + if max_total_bytes is None + else max_total_bytes - os.path.getsize(footprints_path) + ) + if remaining_bytes is not None and remaining_bytes < 1: + raise AssessmentSizeLimitError( + "Assessment inputs exceed the allowed size" + ) + try: + gpkg_path = await self.downloader( + gpkg_url, + suffix=".gpkg", + max_bytes=remaining_bytes, + ) + except ValueError as error: + if "allowed download size" in str(error): + raise AssessmentSizeLimitError( + "Assessment inputs exceed the allowed size" + ) from error + raise + downloaded_paths.append(gpkg_path) + inputs = await asyncio.to_thread( + build_assessment_inputs_from_gpkgs, + footprints_path, + gpkg_path, + labels=labels, + ) + return await asyncio.to_thread( + compute_assessment_report, + inputs, + threshold=threshold, + min_area_m2=min_area_m2, + ) + finally: + for path in downloaded_paths: + try: + os.unlink(path) + except OSError: + pass diff --git a/hastelib/src/hastegeo/core/processors/metadata.py b/hastelib/src/hastegeo/core/processors/metadata.py index d2bdbede..d8f42102 100644 --- a/hastelib/src/hastegeo/core/processors/metadata.py +++ b/hastelib/src/hastegeo/core/processors/metadata.py @@ -125,6 +125,36 @@ def load_all_from_partition(self, data_format="json"): metadata.append(each_metadata) return metadata + def load_bounded( + self, max_records: int, data_format: str = "json" + ) -> list[dict]: + return self.storage.load_bounded( + data_type=self.data_type, + max_records=max_records, + data_format=data_format, + ) + + def load_page( + self, + page: int, + page_size: int, + data_format: str = "json", + target: str = None, + status: str = None, + project_id: str = None, + max_records: int = None, + ) -> tuple[list[dict], int]: + return self.storage.load_page( + data_type=self.data_type, + page=page, + page_size=page_size, + data_format=data_format, + target=target, + status=status, + project_id=project_id, + max_records=max_records, + ) + def load_and_combine_sub_data_types(self, key, data_types): """ Load and combine metadata from multiple data types. diff --git a/hastelib/src/hastegeo/core/processors/publishing.py b/hastelib/src/hastegeo/core/processors/publishing.py new file mode 100644 index 00000000..f64b931a --- /dev/null +++ b/hastelib/src/hastegeo/core/processors/publishing.py @@ -0,0 +1,880 @@ +import re +from dataclasses import dataclass +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, Optional + +from ..config import Config +from ..models.publishing import ( + ArtifactBundle, + PublishDatasetOptions, + PublishedDataset, + PublishOperation, + PublishQueueMessage, + PublishRequest, + PublishStatus, + compute_request_fingerprint, + derive_dataset_id, +) +from ..publishing.lease import LeaseUnavailableError +from ..publishing.registry import PublishingProviderRegistry +from ..publishing.repository import ( + PublishingConflictError, + PublishingRepository, +) +from ..publishing.source import PublishingSourceResolver +from ..utils.logs import Logger + + +class PublishingDisabledError(RuntimeError): + pass + + +class PublishingPermissionError(PermissionError): + pass + + +class PublishingStateConflictError(RuntimeError): + pass + + +class PublishingSizeLimitError(ValueError): + pass + + +class PublishingDependencyError(RuntimeError): + pass + + +@dataclass(frozen=True) +class PreparedPublish: + request: PublishRequest + publisher_id: str + dataset_id: Any + request_fingerprint: str + publisher_name: Optional[str] = None + existing: Optional[PublishedDataset] = None + options: Optional[PublishDatasetOptions] = None + bundle: Optional[ArtifactBundle] = None + + +def _utc_timestamp(now: Optional[datetime] = None) -> str: + timestamp = now or datetime.now(timezone.utc) + return ( + timestamp.astimezone(timezone.utc).isoformat().replace("+00:00", "Z") + ) + + +def _sanitize_status_message(error: Exception) -> str: + message = re.sub(r"https?://\S+", "[url redacted]", str(error)) + message = re.sub( + r"(?i)(authorization\s*:\s*bearer)\s+[^\s,;]+", + r"\1 [redacted]", + message, + ) + message = re.sub( + r"""(?ix) + (["']?(?:access_token|refresh_token|client_secret)["']? + \s*[:=]\s*["']?)[^"'\s,}]+ + """, + r"\1[redacted]", + message, + ) + message = re.sub( + r"(?i)(sig|token|secret|key)=[^&\s]+", + r"\1=[redacted]", + message, + ) + return message[:500] or type(error).__name__ + + +class PublishingProcessor: + """Orchestrate the generic publishing lifecycle in bounded queue steps.""" + + def __init__( + self, + config: Optional[Config] = None, + repository: Optional[PublishingRepository] = None, + source_resolver: Optional[PublishingSourceResolver] = None, + registry: Optional[PublishingProviderRegistry] = None, + queue_handler: Any = None, + ) -> None: + self.config = config or Config() + self.repository = repository or PublishingRepository(self.config) + self.source_resolver = source_resolver or PublishingSourceResolver( + self.config + ) + if registry is None: + from ..publishing.local_provider import LocalPublishingProvider + + def planetary_computer_factory(): + from ..publishing.planetary_computer_provider import ( + PlanetaryComputerPublishingProvider, + ) + + return PlanetaryComputerPublishingProvider( + config=self.config, + artifact_storage=self.source_resolver.artifact_storage, + ) + + registry = PublishingProviderRegistry( + self.config, + factories={ + "local": lambda: LocalPublishingProvider( + config=self.config, + artifact_storage=self.source_resolver.artifact_storage, + ), + "planetary_computer": planetary_computer_factory, + }, + ) + self.registry = registry + self.queue_handler = queue_handler + self.logger = Logger.get_logger(__name__) + + def create( + self, + request: PublishRequest, + publisher_id: str, + assessment_summary: Optional[Dict[str, Any]] = None, + publisher_name: Optional[str] = None, + ) -> PublishedDataset: + prepared = self.prepare_create( + request, publisher_id, publisher_name + ) + return self.create_prepared(prepared, assessment_summary) + + def prepare_create( + self, + request: PublishRequest, + publisher_id: str, + publisher_name: Optional[str] = None, + ) -> PreparedPublish: + self._require_enabled() + dataset_id = derive_dataset_id(request.projectId, request.requestId) + request_fingerprint = compute_request_fingerprint( + request, publisher_id + ) + try: + existing = self.repository.load( + str(request.projectId), str(dataset_id) + ) + except FileNotFoundError: + existing = None + if existing is not None: + if existing.requestFingerprint != request_fingerprint: + raise PublishingConflictError( + "requestId was already used with different publish values" + ) + self._audit(existing, "publish_replayed", publisher_id) + return PreparedPublish( + request=request, + publisher_id=publisher_id, + dataset_id=dataset_id, + request_fingerprint=request_fingerprint, + publisher_name=publisher_name, + existing=existing, + ) + + provider_info = self.registry.get_info(request.target.value) + if not provider_info.isEnabled or not provider_info.isConfigured: + raise PublishingDisabledError( + provider_info.disabledReason or "Provider is unavailable" + ) + unsupported = set(request.artifacts) - set( + provider_info.supportedArtifactKinds + ) + if unsupported: + kinds = ", ".join(sorted(kind.value for kind in unsupported)) + raise ValueError( + f"Provider {request.target.value} does not support: {kinds}" + ) + + options = self.source_resolver.resolve_options( + str(request.projectId), request.imageLayerId, request.modelId + ) + bundle = self.source_resolver.resolve_bundle( + request, + supporting_kinds=provider_info.requiredSupportingArtifactKinds, + options=options, + ) + total_bytes = sum( + artifact.sizeBytes or 0 for artifact in bundle.selectedArtifacts + ) + if total_bytes > self.config.publishing_config["max_total_bytes"]: + raise PublishingSizeLimitError( + "Selected artifacts exceed PUBLISH_MAX_TOTAL_BYTES" + ) + + provider = self.registry.resolve(request.target.value) + provider.validate(request, bundle) + return PreparedPublish( + request=request, + publisher_id=publisher_id, + dataset_id=dataset_id, + request_fingerprint=request_fingerprint, + publisher_name=publisher_name, + options=options, + bundle=bundle, + ) + + def create_prepared( + self, + prepared: PreparedPublish, + assessment_summary: Optional[Dict[str, Any]] = None, + ) -> PublishedDataset: + self._require_enabled() + if prepared.existing is not None: + return prepared.existing + if prepared.options is None or prepared.bundle is None: + raise ValueError("Prepared publish is incomplete") + + request = prepared.request + options = prepared.options + bundle = prepared.bundle + dataset_id = prepared.dataset_id + request_fingerprint = prepared.request_fingerprint + now = _utc_timestamp() + dataset = PublishedDataset( + datasetId=dataset_id, + requestId=request.requestId, + requestFingerprint=request_fingerprint, + name=request.name, + description=request.description or "", + projectId=request.projectId, + projectName=options.projectName, + imageLayerId=request.imageLayerId, + imageLayerName=options.imageLayerName, + modelId=request.modelId, + modelName=options.modelName, + target=request.target, + status=PublishStatus.PENDING, + publishedByUser=prepared.publisher_id.strip().lower(), + publishedByName=prepared.publisher_name, + createdDate=now, + updatedDate=now, + queueDispatchedAt=now, + selectedArtifactKinds=request.artifacts, + sourceArtifacts=( + bundle.selectedArtifacts + bundle.supportingArtifacts + ), + assessmentSummary=assessment_summary or {}, + ) + try: + with self.repository.project_lock(str(request.projectId)): + self.source_resolver.ensure_project_exists( + str(request.projectId) + ) + stored, created = self.repository.create_or_replay_locked( + dataset + ) + except LeaseUnavailableError as error: + try: + existing = self.repository.load( + str(request.projectId), str(dataset_id) + ) + except FileNotFoundError: + raise PublishingStateConflictError( + "Project publishing state is changing; retry the request" + ) from error + if existing.requestFingerprint != request_fingerprint: + raise PublishingConflictError( + "requestId was already used with different publish values" + ) from error + return existing + if not created: + return stored + try: + self._enqueue(stored, visibility_timeout=5) + self._audit(stored, "publish_queued", prepared.publisher_id) + return stored + except Exception as error: + failed = stored.model_copy( + update={ + "status": PublishStatus.FAILED, + "statusMessage": _sanitize_status_message(error), + "updatedDate": _utc_timestamp(), + } + ) + try: + self.repository.update( + failed, expected_revision=stored.revision + ) + except Exception as persistence_error: + self.logger.error( + "Failed to persist publish dispatch failure for %s: %s", + stored.datasetId, + type(persistence_error).__name__, + ) + raise PublishingDependencyError( + "Unable to enqueue publishing operation" + ) from error + + def retry( + self, + project_id: str, + dataset_id: str, + caller_id: str, + is_admin: bool = False, + ) -> PublishedDataset: + self._require_enabled() + dataset = self.repository.load(project_id, dataset_id) + self._require_owner(dataset, caller_id, is_admin) + if dataset.status == PublishStatus.FAILED: + status = PublishStatus.PENDING + operation = PublishOperation.PUBLISH + elif dataset.status == PublishStatus.UNPUBLISH_FAILED: + status = PublishStatus.UNPUBLISH_PENDING + operation = PublishOperation.UNPUBLISH + else: + raise PublishingStateConflictError( + f"Cannot retry dataset in status {dataset.status.value}" + ) + provider = self.registry.resolve(dataset.target.value) + prepare_retry = getattr(provider, "prepare_retry", None) + provider_metadata = ( + prepare_retry(dataset, operation) + if callable(prepare_retry) + else dict(dataset.providerMetadata) + ) + pending = dataset.model_copy( + update={ + "status": status, + "lastOperation": operation, + "attempt": dataset.attempt + 1, + "statusMessage": "", + "queueDispatchedAt": _utc_timestamp(), + "reconciledAttempt": None, + "updatedDate": _utc_timestamp(), + "providerMetadata": dict(provider_metadata), + } + ) + updated = self.repository.update( + pending, expected_revision=dataset.revision + ) + try: + self._enqueue(updated, visibility_timeout=5) + self._audit(updated, "publish_retry_queued", caller_id) + except Exception as error: + failed_status = ( + PublishStatus.FAILED + if operation == PublishOperation.PUBLISH + else PublishStatus.UNPUBLISH_FAILED + ) + self._persist_dispatch_failure(updated, failed_status) + raise PublishingDependencyError( + "Unable to enqueue publishing retry" + ) from error + return updated + + def request_unpublish( + self, + project_id: str, + dataset_id: str, + caller_id: str, + is_admin: bool = False, + ) -> PublishedDataset: + self._require_enabled() + dataset = self.repository.load(project_id, dataset_id) + self._require_owner(dataset, caller_id, is_admin) + if dataset.status in { + PublishStatus.PENDING, + PublishStatus.IN_PROGRESS, + PublishStatus.UNPUBLISH_PENDING, + PublishStatus.UNPUBLISHING, + }: + raise PublishingStateConflictError( + f"Cannot unpublish dataset in status {dataset.status.value}" + ) + pending = dataset.model_copy( + update={ + "status": PublishStatus.UNPUBLISH_PENDING, + "lastOperation": PublishOperation.UNPUBLISH, + "attempt": dataset.attempt + 1, + "statusMessage": "", + "queueDispatchedAt": _utc_timestamp(), + "reconciledAttempt": None, + "updatedDate": _utc_timestamp(), + } + ) + updated = self.repository.update( + pending, expected_revision=dataset.revision + ) + try: + self._enqueue(updated, visibility_timeout=5) + self._audit(updated, "unpublish_queued", caller_id) + except Exception as error: + self._persist_dispatch_failure( + updated, PublishStatus.UNPUBLISH_FAILED + ) + raise PublishingDependencyError( + "Unable to enqueue unpublish operation" + ) from error + return updated + + def run_step( + self, message: PublishQueueMessage + ) -> Optional[PublishedDataset]: + project_id = str(message.projectId) + dataset_id = str(message.datasetId) + try: + with self.repository.operation_lock( + project_id, dataset_id + ) as lease: + dataset = self.repository.load(project_id, dataset_id) + if ( + dataset.attempt != message.attempt + or dataset.lastOperation != message.operation + ): + return dataset + if dataset.reconciledAttempt == message.attempt: + dataset = self.repository.update_locked( + dataset.model_copy( + update={ + "reconciledAttempt": None, + "updatedDate": _utc_timestamp(), + } + ), + dataset.revision, + ) + if message.operation == PublishOperation.PUBLISH: + return self._run_publish_step(dataset, lease) + return self._run_unpublish_step(dataset, lease) + except LeaseUnavailableError: + self.logger.info( + "Publishing operation already claimed for dataset %s", + dataset_id, + ) + return None + + def list_datasets( + self, + project_id: Optional[str] = None, + target=None, + status=None, + ) -> list[PublishedDataset]: + return self.repository.list_all( + project_id=project_id, + target=target, + status=status, + ) + + def list_datasets_page(self, **kwargs): + return self.repository.list_page(**kwargs) + + def get_dataset( + self, project_id: str, dataset_id: str + ) -> PublishedDataset: + return self.repository.load(project_id, dataset_id) + + def _run_publish_step( + self, dataset: PublishedDataset, lease: Any + ) -> PublishedDataset: + if dataset.status not in { + PublishStatus.PENDING, + PublishStatus.IN_PROGRESS, + }: + return dataset + try: + provider = self.registry.resolve(dataset.target.value) + request = self._request_from_dataset(dataset) + info = self.registry.get_info(dataset.target.value) + bundle = self.source_resolver.resolve_bundle( + request, + supporting_kinds=info.requiredSupportingArtifactKinds, + ) + self._validate_worker_bundle(dataset, bundle) + if dataset.status == PublishStatus.PENDING: + current = self.repository.update_locked( + dataset.model_copy( + update={ + "status": PublishStatus.IN_PROGRESS, + "updatedDate": _utc_timestamp(), + } + ), + dataset.revision, + ) + self._renew_lease(lease) + else: + current = dataset + if current.target.value == "planetary_computer": + try: + with self.repository.project_lock(str(current.projectId)): + result = self._run_provider_publish_step( + provider, + current, + bundle, + ) + except LeaseUnavailableError: + self.logger.info( + "Planetary Computer collection update is busy for " + "project %s", + current.projectId, + ) + self._enqueue(current, visibility_timeout=30) + return current + else: + result = self._run_provider_publish_step( + provider, + current, + bundle, + ) + self._renew_lease(lease) + if not result.isComplete: + continued = current.model_copy( + update={ + "providerMetadata": { + **current.providerMetadata, + **result.providerMetadata, + "continuationToken": result.continuationToken, + }, + "updatedDate": _utc_timestamp(), + } + ) + updated = self.repository.update_locked( + continued, current.revision + ) + self._enqueue(updated, visibility_timeout=30) + return updated + completed = current.model_copy( + update={ + "status": PublishStatus.PUBLISHED, + "statusMessage": "", + "publishedDate": _utc_timestamp(), + "updatedDate": _utc_timestamp(), + "artifacts": result.artifacts, + "links": result.links, + "providerMetadata": result.providerMetadata, + } + ) + completed = self.repository.update_locked( + completed, current.revision + ) + self._audit(completed, "publish_completed") + return completed + except Exception as error: + latest = self.repository.load( + str(dataset.projectId), str(dataset.datasetId) + ) + failed = latest.model_copy( + update={ + "status": PublishStatus.FAILED, + "statusMessage": _sanitize_status_message(error), + "updatedDate": _utc_timestamp(), + } + ) + failed = self.repository.update_locked(failed, latest.revision) + self._audit(failed, "publish_failed") + raise + + @staticmethod + def _run_provider_publish_step( + provider: Any, + dataset: PublishedDataset, + bundle: ArtifactBundle, + ) -> Any: + if dataset.providerMetadata.get("continuationToken"): + return provider.continue_publish(dataset, bundle) + return provider.start_publish(dataset, bundle) + + def _run_unpublish_step( + self, dataset: PublishedDataset, lease: Any + ) -> Optional[PublishedDataset]: + if dataset.status not in { + PublishStatus.UNPUBLISH_PENDING, + PublishStatus.UNPUBLISHING, + }: + return dataset + try: + provider = self.registry.resolve(dataset.target.value) + if dataset.status == PublishStatus.UNPUBLISH_PENDING: + current = self.repository.update_locked( + dataset.model_copy( + update={ + "status": PublishStatus.UNPUBLISHING, + "updatedDate": _utc_timestamp(), + } + ), + dataset.revision, + ) + self._renew_lease(lease) + result = provider.start_unpublish(current) + else: + current = dataset + if current.providerMetadata.get("continuationToken"): + result = provider.continue_unpublish(current) + else: + result = provider.start_unpublish(current) + self._renew_lease(lease) + if not result.isComplete: + continued = current.model_copy( + update={ + "providerMetadata": { + **current.providerMetadata, + **result.providerMetadata, + "continuationToken": result.continuationToken, + }, + "updatedDate": _utc_timestamp(), + } + ) + updated = self.repository.update_locked( + continued, current.revision + ) + self._enqueue(updated, visibility_timeout=30) + return updated + self.repository.delete_locked( + str(current.projectId), str(current.datasetId) + ) + self._audit(current, "unpublish_completed") + return None + except Exception as error: + latest = self.repository.load( + str(dataset.projectId), str(dataset.datasetId) + ) + failed = latest.model_copy( + update={ + "status": PublishStatus.UNPUBLISH_FAILED, + "statusMessage": _sanitize_status_message(error), + "updatedDate": _utc_timestamp(), + } + ) + failed = self.repository.update_locked(failed, latest.revision) + self._audit(failed, "unpublish_failed") + raise + + def mark_poisoned(self, message: PublishQueueMessage) -> PublishedDataset: + project_id = str(message.projectId) + dataset_id = str(message.datasetId) + with self.repository.operation_lock(project_id, dataset_id): + dataset = self.repository.load(project_id, dataset_id) + expected_statuses = ( + {PublishStatus.PENDING, PublishStatus.IN_PROGRESS} + if message.operation == PublishOperation.PUBLISH + else { + PublishStatus.UNPUBLISH_PENDING, + PublishStatus.UNPUBLISHING, + } + ) + if ( + dataset.attempt != message.attempt + or dataset.lastOperation != message.operation + or dataset.status not in expected_statuses + ): + return dataset + status = ( + PublishStatus.FAILED + if message.operation == PublishOperation.PUBLISH + else PublishStatus.UNPUBLISH_FAILED + ) + failed = dataset.model_copy( + update={ + "status": status, + "statusMessage": "Publishing operation moved to poison queue", + "updatedDate": _utc_timestamp(), + } + ) + failed = self.repository.update_locked(failed, dataset.revision) + self._audit(failed, "operation_poisoned") + return failed + + def reconcile_stale( + self, + now: Optional[datetime] = None, + stale_after: timedelta = timedelta(minutes=2), + ) -> int: + current_time = now or datetime.now(timezone.utc) + nonterminal = { + PublishStatus.PENDING, + PublishStatus.IN_PROGRESS, + PublishStatus.UNPUBLISH_PENDING, + PublishStatus.UNPUBLISHING, + } + requeued = 0 + for candidate in self.repository.list_for_reconciliation(): + if candidate.status not in nonterminal: + continue + updated_at = datetime.fromisoformat( + candidate.updatedDate.replace("Z", "+00:00") + ) + if current_time - updated_at < stale_after: + continue + project_id = str(candidate.projectId) + dataset_id = str(candidate.datasetId) + try: + with self.repository.operation_lock(project_id, dataset_id): + dataset = self.repository.load(project_id, dataset_id) + updated_at = datetime.fromisoformat( + dataset.updatedDate.replace("Z", "+00:00") + ) + if ( + dataset.status not in nonterminal + or current_time - updated_at < stale_after + or dataset.reconciledAttempt == dataset.attempt + ): + continue + self._enqueue(dataset, visibility_timeout=5) + reconciled = dataset.model_copy( + update={ + "queueDispatchedAt": current_time.isoformat(), + "reconciledAttempt": dataset.attempt, + "updatedDate": current_time.isoformat(), + } + ) + reconciled = self.repository.update_locked( + reconciled, dataset.revision + ) + self._audit(reconciled, "operation_reconciled") + requeued += 1 + except LeaseUnavailableError: + self.logger.info( + "Skipping reconciliation for claimed dataset %s", + dataset_id, + ) + return requeued + + def get_download_urls( + self, project_id: str, dataset_id: str + ) -> Dict[str, str]: + dataset = self.repository.load(project_id, dataset_id) + if ( + dataset.target.value != "local" + or dataset.status != PublishStatus.PUBLISHED + ): + return {} + ttl = self.config.publishing_config["download_sas_minutes"] + return { + artifact.kind.value: self.source_resolver.artifact_storage.get_scoped_download_url( + artifact.publishedPath, expires_minutes=ttl + ) + for artifact in dataset.artifacts + } + + def _enqueue( + self, dataset: PublishedDataset, visibility_timeout: int + ) -> None: + message = PublishQueueMessage( + datasetId=dataset.datasetId, + projectId=dataset.projectId, + operation=dataset.lastOperation, + attempt=dataset.attempt, + ) + self._get_queue_handler().put_message( + message.model_dump_json(), + visibility_timeout=visibility_timeout, + ) + + def _get_queue_handler(self): + if self.queue_handler is None: + from ..utils.queues import AzureQueueHandler + + queue = self.config.queue_config + self.queue_handler = AzureQueueHandler( + connection_string=queue["queue_connection_string"], + queue_name=queue["publish_queue_name"], + account_url=queue["queue_account_url"], + ) + return self.queue_handler + + @staticmethod + def _request_from_dataset(dataset: PublishedDataset) -> PublishRequest: + return PublishRequest( + requestId=dataset.requestId, + projectId=dataset.projectId, + imageLayerId=dataset.imageLayerId, + modelId=dataset.modelId, + name=dataset.name, + description=dataset.description, + target=dataset.target, + artifacts=dataset.selectedArtifactKinds, + ) + + @staticmethod + def _renew_lease(lease: Any) -> None: + renew = getattr(lease, "renew", None) + if callable(renew): + renew() + + def _require_enabled(self) -> None: + if not self.config.publishing_config["publishing_enabled"]: + raise PublishingDisabledError("Publishing is disabled") + + def _audit( + self, + dataset: PublishedDataset, + event: str, + actor: Optional[str] = None, + ) -> None: + self.logger.info( + "Publishing audit event=%s dataset=%s project=%s target=%s " + "operation=%s attempt=%s actor=%s status=%s", + event, + dataset.datasetId, + dataset.projectId, + dataset.target.value, + dataset.lastOperation.value, + dataset.attempt, + (actor or "system").strip().lower(), + dataset.status.value, + ) + + def _persist_dispatch_failure( + self, + dataset: PublishedDataset, + status: PublishStatus, + ) -> PublishedDataset: + failed = dataset.model_copy( + update={ + "status": status, + "statusMessage": "Publishing queue is unavailable", + "queueDispatchedAt": None, + "updatedDate": _utc_timestamp(), + } + ) + return self.repository.update( + failed, expected_revision=dataset.revision + ) + + def _validate_worker_bundle( + self, dataset: PublishedDataset, bundle: Any + ) -> None: + total_bytes = sum( + artifact.sizeBytes or 0 for artifact in bundle.selectedArtifacts + ) + if total_bytes > self.config.publishing_config["max_total_bytes"]: + raise PublishingSizeLimitError( + "Selected artifacts exceed PUBLISH_MAX_TOTAL_BYTES" + ) + + expected = { + artifact.kind: artifact for artifact in dataset.sourceArtifacts + } + if not expected: + return + current_artifacts = ( + bundle.selectedArtifacts + bundle.supportingArtifacts + ) + current = {artifact.kind: artifact for artifact in current_artifacts} + if set(expected) != set(current): + raise PublishingStateConflictError( + "Source artifact selection changed after publishing was queued" + ) + for kind, expected_artifact in expected.items(): + current_artifact = current[kind] + if ( + expected_artifact.sourcePath != current_artifact.sourcePath + or expected_artifact.sizeBytes != current_artifact.sizeBytes + or expected_artifact.mediaType != current_artifact.mediaType + or expected_artifact.sourceEtag != current_artifact.sourceEtag + ): + raise PublishingStateConflictError( + f"Source artifact changed after publishing was queued: {kind.value}" + ) + + @staticmethod + def _require_owner( + dataset: PublishedDataset, caller_id: str, is_admin: bool + ) -> None: + if is_admin: + return + if dataset.publishedByUser != caller_id.strip().lower(): + raise PublishingPermissionError( + "Only the publisher or an administrator may perform this action" + ) diff --git a/hastelib/src/hastegeo/core/publishing/__init__.py b/hastelib/src/hastegeo/core/publishing/__init__.py new file mode 100644 index 00000000..f83ba3e2 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/__init__.py @@ -0,0 +1,21 @@ +from .lease import ( + BlobLeaseCoordinator, + LeaseRenewalError, + LeaseUnavailableError, +) +from .repository import ( + PublishedDatasetsExistError, + PublishingConflictError, + PublishingRepository, + StaleRevisionError, +) + +__all__ = [ + "BlobLeaseCoordinator", + "LeaseRenewalError", + "LeaseUnavailableError", + "PublishedDatasetsExistError", + "PublishingConflictError", + "PublishingRepository", + "StaleRevisionError", +] diff --git a/hastelib/src/hastegeo/core/publishing/base.py b/hastelib/src/hastegeo/core/publishing/base.py new file mode 100644 index 00000000..4b711e74 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/base.py @@ -0,0 +1,53 @@ +from abc import ABC, abstractmethod + +from ..models.publishing import ( + ArtifactBundle, + ProviderInfo, + PublishedDataset, + PublishOperation, + PublishRequest, + PublishResult, +) + + +class PublishingProvider(ABC): + """Provider contract for bounded, replay-safe publishing steps.""" + + @property + @abstractmethod + def info(self) -> ProviderInfo: + pass # pragma: no cover + + @abstractmethod + def validate( + self, request: PublishRequest, source: ArtifactBundle + ) -> None: + pass # pragma: no cover + + def prepare_retry( + self, + dataset: PublishedDataset, + operation: PublishOperation, + ) -> dict: + """Return provider metadata safe for a new operation attempt.""" + return dict(dataset.providerMetadata) + + @abstractmethod + def start_publish( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> PublishResult: + pass # pragma: no cover + + @abstractmethod + def continue_publish( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> PublishResult: + pass # pragma: no cover + + @abstractmethod + def start_unpublish(self, dataset: PublishedDataset) -> PublishResult: + pass # pragma: no cover + + @abstractmethod + def continue_unpublish(self, dataset: PublishedDataset) -> PublishResult: + pass # pragma: no cover diff --git a/hastelib/src/hastegeo/core/publishing/geocatalog_client.py b/hastelib/src/hastegeo/core/publishing/geocatalog_client.py new file mode 100644 index 00000000..f8422542 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/geocatalog_client.py @@ -0,0 +1,139 @@ +"""Minimal, hardened REST client for the Planetary Computer Pro GeoCatalog. + +Adapted from the AI for Good Lab ``pcpublish`` reference client, trimmed to the +calls the publishing transport needs and hardened for server-side use: + +- Entra ID auth via ``DefaultAzureCredential`` (scope + ``https://geocatalog.spatio.azure.com/.default``), token cached with an expiry + skew. ``azure-identity`` is imported lazily so importing this module does not + require the ``planetary-computer`` extra. +- Redirects are never followed and every request carries explicit + (connect, read) timeouts — the caller pins operation URLs to the GeoCatalog + origin (see the transport adapter), so redirect-following would be an SSRF + vector. +- ``GeoCatalogError`` never embeds server response bodies (the ``pcpublish`` + reference does); it carries only a sanitized message and the HTTP status. + +This client returns raw ``requests.Response`` objects; the transport adapter +converts the async 202 + operation-location flow into resumable steps. +""" + +from __future__ import annotations + +import time +from typing import Any, Iterable, Optional + +import requests + +API_VERSION = "2026-04-15" +_TOKEN_SCOPE = "https://geocatalog.spatio.azure.com/.default" +# Refresh a little before expiry so long ingestions do not fail mid-flight. +_EXPIRY_SKEW_SECONDS = 300 + + +class GeoCatalogError(RuntimeError): + """Raised when the GeoCatalog API returns an unexpected response. + + Carries only the HTTP status code; server response text is never embedded, + to avoid leaking tokens/SAS present in error bodies. + """ + + def __init__( + self, message: str, *, status_code: Optional[int] = None + ) -> None: + super().__init__(message) + self.status_code = status_code + + +class GeoCatalogAuth: + """Caches an Entra ID access token for the GeoCatalog data plane.""" + + def __init__(self, credential: Any = None) -> None: + self._credential = credential + self._token: Optional[str] = None + self._expires_on: float = 0.0 + + def token(self) -> str: + now = time.time() + if self._token and now < self._expires_on - _EXPIRY_SKEW_SECONDS: + return self._token + credential = self._credential or self._default_credential() + self._credential = credential + access_token = credential.get_token(_TOKEN_SCOPE) + self._token = access_token.token + self._expires_on = float(access_token.expires_on) + return self._token + + def headers(self) -> dict[str, str]: + return {"Authorization": f"Bearer {self.token()}"} + + @staticmethod + def _default_credential() -> Any: + # DefaultAzureCredential already chains managed identity, env, and the + # az CLI, so no separate CLI-subprocess fallback is needed. + from azure.identity import DefaultAzureCredential + + return DefaultAzureCredential() + + +class GeoCatalogClient: + """Thin, hardened REST wrapper over the GeoCatalog STAC/ingestion APIs.""" + + def __init__( + self, + endpoint: str, + *, + auth: Optional[GeoCatalogAuth] = None, + api_version: str = API_VERSION, + connection_timeout: int = 10, + read_timeout: int = 30, + ) -> None: + self.endpoint = endpoint.rstrip("/") + self.auth = auth or GeoCatalogAuth() + self.api_version = api_version + self.timeout = (connection_timeout, read_timeout) + self._session = requests.Session() + # Never auto-follow redirects; operation URLs are origin-pinned by the + # caller, and following a server-supplied redirect would defeat that. + self._session.max_redirects = 0 + + def request( + self, + method: str, + url: str, + *, + json: Any = None, + params: Optional[dict[str, Any]] = None, + expected: Iterable[int] = (200, 201, 202, 204), + absolute: bool = False, + ) -> requests.Response: + target = url if absolute else f"{self.endpoint}{url}" + query = dict(params or {}) + # Absolute operation/continuation URLs already carry api-version in their + # query string; adding it again would produce a duplicate parameter. + if "api-version" not in query and "api-version=" not in target: + query["api-version"] = self.api_version + headers = self.auth.headers() + try: + response = self._session.request( + method, + target, + params=query, + json=json, + headers=headers, + timeout=self.timeout, + allow_redirects=False, + ) + except requests.RequestException as error: + raise GeoCatalogError( + f"{method} request to GeoCatalog failed" + ) from error + if response.status_code not in tuple(expected): + raise GeoCatalogError( + f"{method} {url} returned HTTP {response.status_code}", + status_code=response.status_code, + ) + return response + + def close(self) -> None: + self._session.close() diff --git a/hastelib/src/hastegeo/core/publishing/lease.py b/hastelib/src/hastegeo/core/publishing/lease.py new file mode 100644 index 00000000..736288a4 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/lease.py @@ -0,0 +1,136 @@ +from contextlib import contextmanager +from threading import Event, Thread +from time import monotonic, sleep +from typing import Any, Iterator, Optional + +from azure.core.exceptions import HttpResponseError, ResourceExistsError +from azure.identity import DefaultAzureCredential +from azure.storage.blob import BlobServiceClient + +from ..utils.metadata import MetadataUtils + + +class LeaseUnavailableError(RuntimeError): + """Raised when another operation owns a dataset lease.""" + + +class LeaseRenewalError(RuntimeError): + """Raised when a held dataset lease cannot be renewed.""" + + +class BlobLeaseCoordinator: + """Serialize publishing operations with per-dataset Azure Blob leases.""" + + def __init__( + self, + connection_string: Optional[str], + account_url: Optional[str], + container_name: str = "publishing-locks", + blob_service_client: Any = None, + renewal_interval_seconds: Optional[float] = None, + ) -> None: + if ( + renewal_interval_seconds is not None + and renewal_interval_seconds <= 0 + ): + raise ValueError("renewal_interval_seconds must be positive") + self.renewal_interval_seconds = renewal_interval_seconds + if blob_service_client is not None: + self.blob_service_client = blob_service_client + elif connection_string: + self.blob_service_client = ( + BlobServiceClient.from_connection_string(connection_string) + ) + elif account_url: + self.blob_service_client = BlobServiceClient( + account_url=account_url, + credential=DefaultAzureCredential(), + ) + else: + raise ValueError( + "A publishing lease connection string or account URL is required" + ) + + try: + self.container_client = self.blob_service_client.create_container( + container_name + ) + except ResourceExistsError: + self.container_client = ( + self.blob_service_client.get_container_client(container_name) + ) + + @contextmanager + def acquire( + self, + project_id: str, + dataset_id: str, + lease_duration: int = 60, + wait_timeout_seconds: float = 0, + retry_interval_seconds: float = 0.05, + ) -> Iterator[Any]: + if lease_duration < 15 or lease_duration > 60: + raise ValueError( + "lease_duration must be between 15 and 60 seconds" + ) + if wait_timeout_seconds < 0 or retry_interval_seconds <= 0: + raise ValueError("Lease wait values must be positive") + + lock_name = ( + f"{MetadataUtils.hash_string(project_id)}/{dataset_id}.lock" + ) + blob_client = self.container_client.get_blob_client(lock_name) + try: + blob_client.upload_blob(b"", overwrite=False) + except ResourceExistsError: + pass + + deadline = monotonic() + wait_timeout_seconds + while True: + try: + lease = blob_client.acquire_lease( + lease_duration=lease_duration + ) + break + except HttpResponseError as error: + if error.status_code != 409: + raise + remaining = deadline - monotonic() + if remaining <= 0: + raise LeaseUnavailableError( + f"Publishing operation already active for {dataset_id}" + ) from error + sleep(min(retry_interval_seconds, remaining)) + + renewal_stop = Event() + renewal_errors: list[Exception] = [] + renewal_interval = ( + self.renewal_interval_seconds + if self.renewal_interval_seconds is not None + else max(5.0, lease_duration / 3) + ) + + def renew_lease() -> None: + while not renewal_stop.wait(renewal_interval): + try: + lease.renew() + except Exception as error: + renewal_errors.append(error) + return + + renewal_thread = Thread( + target=renew_lease, + name=f"publishing-lease-{dataset_id}", + daemon=True, + ) + renewal_thread.start() + try: + yield lease + finally: + renewal_stop.set() + renewal_thread.join(timeout=5) + lease.release() + if renewal_errors: + raise LeaseRenewalError( + f"Publishing lease renewal failed for {dataset_id}" + ) from renewal_errors[0] diff --git a/hastelib/src/hastegeo/core/publishing/local_provider.py b/hastelib/src/hastegeo/core/publishing/local_provider.py new file mode 100644 index 00000000..7fe375e0 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/local_provider.py @@ -0,0 +1,115 @@ +from pathlib import PurePosixPath +from typing import Optional + +from ..artifact_storage.unified_artifact_storage import UnifiedArtifactStorage +from ..config import Config +from ..models.publishing import ( + ArtifactBundle, + ArtifactKind, + ProviderInfo, + PublishedArtifact, + PublishedDataset, + PublishRequest, + PublishResult, +) +from .base import PublishingProvider + + +class LocalPublishingProvider(PublishingProvider): + """Publish immutable copies into HASTE-managed artifact storage.""" + + def __init__( + self, + config: Optional[Config] = None, + artifact_storage: Optional[UnifiedArtifactStorage] = None, + ) -> None: + self.config = config or Config() + self.artifact_storage = artifact_storage or UnifiedArtifactStorage( + storage_type=self.config.artifact_storage_type, + **self.config.artifact_storage_config, + ) + + @property + def info(self) -> ProviderInfo: + enabled = self.config.publishing_config["publishing_enabled"] + return ProviderInfo( + id="local", + displayName="Local (HASTE storage)", + description="Immutable copy in HASTE-managed storage", + isEnabled=enabled, + isConfigured=True, + disabledReason=None if enabled else "Publishing is disabled", + supportedArtifactKinds=list(ArtifactKind), + ) + + def validate( + self, request: PublishRequest, source: ArtifactBundle + ) -> None: + if not source.selectedArtifacts: + raise ValueError("Select at least one artifact to publish") + supported = set(self.info.supportedArtifactKinds) + unsupported = set(request.artifacts) - supported + if unsupported: + names = ", ".join(sorted(kind.value for kind in unsupported)) + raise ValueError(f"Unsupported Local artifacts: {names}") + + def start_publish( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> PublishResult: + prefix = self._dataset_prefix(dataset) + published_artifacts = [] + for artifact in source.selectedArtifacts: + file_name = PurePosixPath(artifact.sourcePath).name + destination = f"{prefix}/{artifact.kind.value}_{file_name}" + published_path = self.artifact_storage.copy_artifact( + artifact.sourcePath, + destination, + artifact.sourceEtag, + ) + published_artifacts.append( + PublishedArtifact( + **artifact.model_dump(), + publishedPath=published_path, + ) + ) + + report_name = f"assessment_report_{dataset.datasetId}.json" + report_path = self.artifact_storage.store_artifact( + artifact_name=report_name, + data=dataset.assessmentSummary, + namespace=prefix.split("/"), + ) + return PublishResult( + artifacts=published_artifacts, + providerMetadata={"assessmentReportPath": report_path}, + isComplete=True, + ) + + def continue_publish( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> PublishResult: + raise RuntimeError( + "Local publishing does not have a continuation step" + ) + + def start_unpublish(self, dataset: PublishedDataset) -> PublishResult: + deleted_count = self.artifact_storage.delete_prefix( + self._dataset_prefix(dataset) + ) + return PublishResult( + providerMetadata={"deletedArtifactCount": deleted_count}, + isComplete=True, + ) + + def continue_unpublish(self, dataset: PublishedDataset) -> PublishResult: + return self.start_unpublish(dataset) + + @staticmethod + def _dataset_prefix(dataset: PublishedDataset) -> str: + # Flat `published/` prefix (datasetId is a UUID, so globally + # unique). Keeping the blob path to three in-container segments + # (published//) lets the UI serve downloads through the + # existing managed-identity storage proxy (get-artifacts), which the + # VNet-only storage account requires — a direct blob SAS from the browser + # is denied by the storage firewall. + return f"published/{dataset.datasetId}" diff --git a/hastelib/src/hastegeo/core/publishing/planetary_computer_provider.py b/hastelib/src/hastegeo/core/publishing/planetary_computer_provider.py new file mode 100644 index 00000000..0c21fb93 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/planetary_computer_provider.py @@ -0,0 +1,1732 @@ +import json +import math +import re +import tempfile +from contextlib import contextmanager +from enum import Enum +from http.client import HTTPSConnection +from pathlib import Path, PurePosixPath +from typing import Any, Callable, Iterator, Mapping, Optional +from urllib.parse import quote, unquote, urlparse + +from pyproj import CRS +from pyproj.exceptions import CRSError +from shapely.geometry import shape + +from ..artifact_storage.unified_artifact_storage import UnifiedArtifactStorage +from ..config import Config +from ..models.publishing import ( + ArtifactBundle, + ArtifactKind, + ProviderInfo, + PublishedArtifact, + PublishedDataset, + PublishOperation, + PublishRequest, + PublishResult, + SourceArtifact, +) +from ..utils.gdal_security import harden_gdal +from ..utils.logs import Logger +from .base import PublishingProvider +from .planetary_computer_transport import ( + PlanetaryComputerOperationError, + PlanetaryComputerRestAdapter, +) +from .stac import ( + ASSET_KEYS, + DEFAULT_THUMBNAIL_MEDIA_TYPE, + THUMBNAIL_ASSET_KEY, + StacObjects, + build_collection_id, + build_item_id, + build_stac_objects, + rebuild_collection_after_removal, + serialize_stac_objects, + validate_stac_objects, +) + +PC_API_VERSION = "2026-04-15" +MAX_VALID_MASK_BYTES = 32 * 1024 * 1024 +AZURE_BLOB_HOST_SUFFIXES = ( + ".blob.core.windows.net", + ".blob.core.usgovcloudapi.net", +) +_THUMBNAIL_MEDIA_TYPES = { + ".png": "image/png", + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".webp": "image/webp", +} + + +class PlanetaryComputerProviderError(RuntimeError): + """Raised when Planetary Computer publishing cannot safely continue.""" + + +class PlanetaryComputerPhase(str, Enum): + COLLECTION_OPERATION = "collection_operation" + COLLECTION_VERIFY = "collection_verify" + ITEM_OPERATION = "item_operation" + ITEM_VERIFY = "item_verify" + ITEM_REPLACE_DELETE_OPERATION = "item_replace_delete_operation" + ITEM_REPLACE_DELETE_VERIFY = "item_replace_delete_verify" + DELETE_OPERATION = "delete_operation" + DELETE_VERIFY = "delete_verify" + DRAIN_COLLECTION = "drain_collection" + DRAIN_ITEM = "drain_item" + DRAIN_DELETE = "drain_delete" + DELETE_DISCOVER = "delete_discover" + DELETE_COLLECTION_OPERATION = "delete_collection_operation" + DELETE_COLLECTION_VERIFY = "delete_collection_verify" + + +class PlanetaryComputerPublishingProvider(PublishingProvider): + """Publish HASTE vector outputs into a Planetary Computer GeoCatalog.""" + + def __init__( + self, + config: Optional[Config] = None, + artifact_storage: Optional[UnifiedArtifactStorage] = None, + sdk_adapter: Optional[PlanetaryComputerRestAdapter] = None, + json_reader: Optional[ + Callable[[SourceArtifact], Mapping[str, Any]] + ] = None, + projection_resolver: Optional[Callable[[SourceArtifact], str]] = None, + stac_validator: Callable[[StacObjects], None] = validate_stac_objects, + asset_reachability_checker: Optional[Callable[[str], None]] = None, + ) -> None: + self.config = config or Config() + settings = self.config.publishing_config + self.endpoint = str(settings.get("pc_geocatalog_url") or "").rstrip( + "/" + ) + self.ingestion_source = str(settings.get("pc_ingestion_source") or "") + self.collection_prefix = str( + settings.get("pc_collection_prefix") or "haste-" + ) + self.explorer_url = str( + settings.get("pc_explorer_url") or self.endpoint + ).rstrip("/") + self.license_id = str( + settings.get("pc_publishing_license") or "CC-BY-4.0" + ) + self.max_verify_attempts = int(settings.get("pc_verify_attempts") or 5) + self.artifact_storage = artifact_storage or UnifiedArtifactStorage( + storage_type=self.config.artifact_storage_type, + **self.config.artifact_storage_config, + ) + self.sdk = sdk_adapter or PlanetaryComputerRestAdapter( + self.endpoint + ) + self.json_reader = json_reader or self._read_json_artifact + self.projection_resolver = ( + projection_resolver or self._read_projection_code + ) + self.stac_validator = stac_validator + self.asset_reachability_checker = ( + asset_reachability_checker or self._read_signed_asset + ) + self.logger = Logger.get_logger(__name__) + + @property + def info(self) -> ProviderInfo: + settings = self.config.publishing_config + enabled = bool(settings.get("pc_provider_enabled")) + # The GeoCatalog URL is what makes the target configurable. An ingestion + # source is only required for private containers; public containers + # publish without one (verified against a live GeoCatalog). + configured = bool(self.endpoint) + if not enabled: + disabled_reason = "Disabled by the operator" + elif not configured: + disabled_reason = "Planetary Computer is not configured" + else: + disabled_reason = None + return ProviderInfo( + id="planetary_computer", + displayName="Planetary Computer", + description="STAC discovery and vector downloads", + isEnabled=enabled, + isConfigured=configured, + disabledReason=disabled_reason, + supportedArtifactKinds=[ + ArtifactKind.GPKG, + ArtifactKind.VALID_MASK, + ArtifactKind.FOOTPRINTS, + ], + requiredSupportingArtifactKinds=[ArtifactKind.VALID_MASK], + ) + + def validate( + self, request: PublishRequest, source: ArtifactBundle + ) -> None: + self._require_configuration() + self._validate_bundle(source) + requested = set(request.artifacts) + selected = {artifact.kind for artifact in source.selectedArtifacts} + if requested != selected: + raise ValueError("Selected Planetary Computer artifacts changed") + for artifact in source.selectedArtifacts: + self._artifact_href(artifact) + + def prepare_retry( + self, + dataset: PublishedDataset, + operation: PublishOperation, + ) -> dict[str, Any]: + metadata = self._stable_metadata(dataset) + phase = dataset.providerMetadata.get("phase") + discovery_phases = { + PlanetaryComputerPhase.COLLECTION_OPERATION.value, + PlanetaryComputerPhase.COLLECTION_VERIFY.value, + PlanetaryComputerPhase.ITEM_OPERATION.value, + PlanetaryComputerPhase.ITEM_VERIFY.value, + PlanetaryComputerPhase.ITEM_REPLACE_DELETE_OPERATION.value, + PlanetaryComputerPhase.ITEM_REPLACE_DELETE_VERIFY.value, + PlanetaryComputerPhase.DRAIN_COLLECTION.value, + PlanetaryComputerPhase.DRAIN_ITEM.value, + PlanetaryComputerPhase.DRAIN_DELETE.value, + PlanetaryComputerPhase.DELETE_DISCOVER.value, + } + requires_discovery = bool( + dataset.providerMetadata.get("cleanupDiscoveryRequired") + or phase in discovery_phases + ) + if operation == PublishOperation.UNPUBLISH and requires_discovery: + metadata.update( + { + "phase": PlanetaryComputerPhase.DELETE_DISCOVER.value, + "verificationAttempts": 0, + "cleanupDiscoveryRequired": True, + } + ) + return metadata + + def start_publish( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> PublishResult: + self._require_configuration() + self._validate_bundle(source) + if self.ingestion_source: + self._validate_ingestion_source() + projection_codes = self._projection_codes(dataset, source) + collection_id = build_collection_id(dataset, self.collection_prefix) + item_id = build_item_id(dataset) + existing_collection = self.sdk.get_collection(collection_id) + documents = self._build_documents( + dataset, + source, + projection_codes, + existing_collection, + ) + metadata = self._operation_metadata( + dataset, + projection_codes, + phase=None, + ) + + if existing_collection is None: + try: + operation = self.sdk.start_create_collection( + collection_id, + documents.collection, + ) + except Exception as error: + if not self._is_status(error, 409): + raise + return self._verification_pending( + metadata, + PlanetaryComputerPhase.COLLECTION_VERIFY, + self._collection_href(collection_id), + ) + if operation.is_complete: + return self._collection_ready_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + return self._operation_pending( + metadata, + PlanetaryComputerPhase.COLLECTION_OPERATION, + operation.continuation_token, + ) + + self.sdk.replace_collection( + collection_id, + documents.collection, + ) + return self._start_or_complete_item( + dataset, + source, + documents.item, + metadata, + collection_id, + item_id, + ) + + def continue_publish( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> PublishResult: + self._require_configuration() + self._validate_bundle(source) + metadata = dataset.providerMetadata + phase = self._phase(metadata) + token = self._continuation_token(metadata) + collection_id, item_id = self._ids(dataset) + projection_codes = self._projection_codes(dataset, source) + + if phase == PlanetaryComputerPhase.COLLECTION_OPERATION: + operation = self.sdk.continue_create_collection( + collection_id, + token, + ) + if not operation.is_complete: + return self._operation_pending( + metadata, + phase, + operation.continuation_token, + count_attempt=True, + ) + return self._collection_ready_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.COLLECTION_VERIFY: + return self._collection_ready_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.ITEM_OPERATION: + operation = self.sdk.continue_create_item( + collection_id, + item_id, + token, + ) + if not operation.is_complete: + return self._operation_pending( + metadata, + phase, + operation.continuation_token, + count_attempt=True, + ) + return self._item_complete_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.ITEM_VERIFY: + return self._item_complete_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.ITEM_REPLACE_DELETE_OPERATION: + operation = self.sdk.continue_delete_item( + collection_id, + item_id, + token, + ) + if not operation.is_complete: + return self._operation_pending( + metadata, + phase, + operation.continuation_token, + count_attempt=True, + ) + return self._replace_delete_complete_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.ITEM_REPLACE_DELETE_VERIFY: + return self._replace_delete_complete_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + raise PlanetaryComputerProviderError( + "Planetary Computer publish phase is invalid" + ) + + def start_unpublish(self, dataset: PublishedDataset) -> PublishResult: + self._require_configuration() + metadata = dataset.providerMetadata + phase_value = metadata.get("phase") + if phase_value in { + PlanetaryComputerPhase.DELETE_OPERATION.value, + PlanetaryComputerPhase.DELETE_VERIFY.value, + PlanetaryComputerPhase.DRAIN_COLLECTION.value, + PlanetaryComputerPhase.DRAIN_ITEM.value, + PlanetaryComputerPhase.DRAIN_DELETE.value, + PlanetaryComputerPhase.DELETE_DISCOVER.value, + PlanetaryComputerPhase.DELETE_COLLECTION_OPERATION.value, + PlanetaryComputerPhase.DELETE_COLLECTION_VERIFY.value, + }: + return self.continue_unpublish(dataset) + + collection_id, item_id = self._ids(dataset) + if phase_value == PlanetaryComputerPhase.COLLECTION_OPERATION.value: + return self._drain_collection_operation( + dataset, + metadata, + collection_id, + item_id, + ) + if phase_value == PlanetaryComputerPhase.COLLECTION_VERIFY.value: + return self._start_delete_or_discover( + dataset, + self._cleanup_discovery_metadata(metadata), + collection_id, + item_id, + ) + if phase_value == PlanetaryComputerPhase.ITEM_OPERATION.value: + return self._drain_item_operation( + dataset, + metadata, + collection_id, + item_id, + ) + if phase_value == PlanetaryComputerPhase.ITEM_VERIFY.value: + return self._start_delete_or_discover( + dataset, + self._cleanup_discovery_metadata(metadata), + collection_id, + item_id, + ) + if ( + phase_value + == PlanetaryComputerPhase.ITEM_REPLACE_DELETE_OPERATION.value + ): + return self._drain_delete_operation( + dataset, + metadata, + collection_id, + item_id, + ) + if ( + phase_value + == PlanetaryComputerPhase.ITEM_REPLACE_DELETE_VERIFY.value + ): + return self._start_delete_or_discover( + dataset, + self._cleanup_discovery_metadata(metadata), + collection_id, + item_id, + ) + stable_metadata = self._stable_metadata(dataset) + if metadata.get("assetsCopiedToManagedStorage") is True: + return self._start_delete_or_complete( + dataset, + stable_metadata, + collection_id, + item_id, + ) + return self._start_delete_or_discover( + dataset, + self._cleanup_discovery_metadata(stable_metadata), + collection_id, + item_id, + ) + + def continue_unpublish(self, dataset: PublishedDataset) -> PublishResult: + self._require_configuration() + metadata = dataset.providerMetadata + phase = self._phase(metadata) + collection_id, item_id = self._ids(dataset) + + if phase == PlanetaryComputerPhase.DELETE_OPERATION: + operation = self.sdk.continue_delete_item( + collection_id, + item_id, + self._continuation_token(metadata), + ) + if not operation.is_complete: + return self._operation_pending( + metadata, + phase, + operation.continuation_token, + count_attempt=True, + ) + return self._delete_complete_or_pending( + dataset, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.DELETE_VERIFY: + return self._delete_complete_or_pending( + dataset, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.DRAIN_COLLECTION: + return self._drain_collection_operation( + dataset, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.DRAIN_ITEM: + return self._drain_item_operation( + dataset, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.DRAIN_DELETE: + return self._drain_delete_operation( + dataset, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.DELETE_DISCOVER: + return self._start_delete_or_discover( + dataset, + metadata, + collection_id, + item_id, + ) + + if phase == PlanetaryComputerPhase.DELETE_COLLECTION_OPERATION: + try: + operation = self.sdk.continue_delete_collection( + collection_id, + self._continuation_token(metadata), + ) + except PlanetaryComputerOperationError: + return self._collection_delete_complete_or_pending( + dataset, metadata, collection_id + ) + if not operation.is_complete: + return self._operation_pending( + metadata, + PlanetaryComputerPhase.DELETE_COLLECTION_OPERATION, + operation.continuation_token, + count_attempt=True, + ) + return self._collection_delete_complete_or_pending( + dataset, metadata, collection_id + ) + + if phase == PlanetaryComputerPhase.DELETE_COLLECTION_VERIFY: + return self._collection_delete_complete_or_pending( + dataset, metadata, collection_id + ) + + raise PlanetaryComputerProviderError( + "Planetary Computer unpublish phase is invalid" + ) + + def _drain_collection_operation( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + cleanup_metadata = self._cleanup_discovery_metadata(metadata) + try: + operation = self.sdk.continue_create_collection( + collection_id, + self._continuation_token(metadata), + ) + except PlanetaryComputerOperationError: + return self._start_delete_or_discover( + dataset, + cleanup_metadata, + collection_id, + item_id, + ) + if not operation.is_complete: + return self._operation_pending( + cleanup_metadata, + PlanetaryComputerPhase.DRAIN_COLLECTION, + operation.continuation_token, + count_attempt=True, + ) + return self._start_delete_or_discover( + dataset, + cleanup_metadata, + collection_id, + item_id, + ) + + def _drain_item_operation( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + cleanup_metadata = self._cleanup_discovery_metadata(metadata) + try: + operation = self.sdk.continue_create_item( + collection_id, + item_id, + self._continuation_token(metadata), + ) + if not operation.is_complete: + return self._operation_pending( + cleanup_metadata, + PlanetaryComputerPhase.DRAIN_ITEM, + operation.continuation_token, + count_attempt=True, + ) + except PlanetaryComputerOperationError: + pass + return self._start_delete_or_discover( + dataset, + cleanup_metadata, + collection_id, + item_id, + ) + + def _drain_delete_operation( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + cleanup_metadata = self._cleanup_discovery_metadata(metadata) + try: + operation = self.sdk.continue_delete_item( + collection_id, + item_id, + self._continuation_token(metadata), + ) + if not operation.is_complete: + return self._operation_pending( + cleanup_metadata, + PlanetaryComputerPhase.DRAIN_DELETE, + operation.continuation_token, + count_attempt=True, + ) + except PlanetaryComputerOperationError: + pass + return self._start_delete_or_discover( + dataset, + cleanup_metadata, + collection_id, + item_id, + ) + + def _start_delete_or_complete( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + if self.sdk.get_item(collection_id, item_id) is None: + return self._cleanup_collection_or_complete(dataset, collection_id) + try: + operation = self.sdk.start_delete_item(collection_id, item_id) + except Exception as error: + if self._is_status(error, 404): + return self._cleanup_collection_or_complete( + dataset, collection_id + ) + if self._is_status(error, 409): + return self._verification_pending( + metadata, + PlanetaryComputerPhase.DELETE_VERIFY, + self._item_href(collection_id, item_id), + ) + raise + return self._operation_pending( + metadata, + PlanetaryComputerPhase.DELETE_OPERATION, + operation.continuation_token, + ) + + def _start_delete_or_discover( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + if self.sdk.get_item(collection_id, item_id) is not None: + return self._start_delete_or_complete( + dataset, + metadata, + collection_id, + item_id, + ) + attempts = int(metadata.get("verificationAttempts") or 0) + 1 + if attempts > self.max_verify_attempts: + raise PlanetaryComputerProviderError( + "Planetary Computer cleanup verification timed out" + ) + updated = dict(metadata) + updated["phase"] = PlanetaryComputerPhase.DELETE_DISCOVER.value + updated["verificationAttempts"] = attempts + updated["operationAttempts"] = 0 + return PublishResult( + providerMetadata=updated, + continuationToken=self._item_href(collection_id, item_id), + isComplete=False, + ) + + @staticmethod + def _cleanup_discovery_metadata( + metadata: Mapping[str, Any] + ) -> dict[str, Any]: + updated = dict(metadata) + updated["verificationAttempts"] = 0 + updated["cleanupDiscoveryRequired"] = True + return updated + + def _cleanup_collection_or_complete( + self, + dataset: PublishedDataset, + collection_id: str, + ) -> PublishResult: + # The dataset's item is gone. Collections are per-project and may still + # hold other datasets, so only delete the collection once no items + # remain; otherwise refresh its rolling summary to drop this dataset. + existing = self.sdk.get_collection(collection_id) + if existing is None: + return PublishResult( + providerMetadata=self._stable_metadata(dataset) + ) + if self.sdk.list_item_ids(collection_id): + self.sdk.replace_collection( + collection_id, + rebuild_collection_after_removal(existing, dataset), + ) + return PublishResult( + providerMetadata=self._stable_metadata(dataset) + ) + return self._start_delete_collection_or_complete( + dataset, + self._stable_metadata(dataset), + collection_id, + ) + + def _start_delete_collection_or_complete( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + ) -> PublishResult: + try: + operation = self.sdk.start_delete_collection(collection_id) + except Exception as error: + if self._is_status(error, 404): + return PublishResult( + providerMetadata=self._stable_metadata(dataset) + ) + if self._is_status(error, 409): + return self._verification_pending( + metadata, + PlanetaryComputerPhase.DELETE_COLLECTION_VERIFY, + self._collection_href(collection_id), + ) + raise + if operation.is_complete: + return self._collection_delete_complete_or_pending( + dataset, metadata, collection_id + ) + return self._operation_pending( + metadata, + PlanetaryComputerPhase.DELETE_COLLECTION_OPERATION, + operation.continuation_token, + ) + + def _collection_delete_complete_or_pending( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + ) -> PublishResult: + if self.sdk.get_collection(collection_id) is None: + return PublishResult( + providerMetadata=self._stable_metadata(dataset) + ) + return self._verification_pending( + metadata, + PlanetaryComputerPhase.DELETE_COLLECTION_VERIFY, + self._collection_href(collection_id), + ) + + def _collection_ready_or_pending( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + projection_codes: Mapping[str, str], + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + existing_collection = self.sdk.get_collection(collection_id) + if existing_collection is None: + return self._verification_pending( + metadata, + PlanetaryComputerPhase.COLLECTION_VERIFY, + self._collection_href(collection_id), + ) + documents = self._build_documents( + dataset, + source, + projection_codes, + existing_collection, + ) + self.sdk.replace_collection( + collection_id, + documents.collection, + ) + return self._start_or_complete_item( + dataset, + source, + documents.item, + metadata, + collection_id, + item_id, + ) + + def _start_or_complete_item( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + expected_item: Mapping[str, Any], + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + existing_item = self.sdk.get_item(collection_id, item_id) + if existing_item is not None: + try: + self._verify_item( + source, + expected_item, + existing_item, + verify_reachability=False, + ) + except PlanetaryComputerProviderError: + return self._start_replace_item_delete( + dataset, + source, + expected_item, + metadata, + collection_id, + item_id, + ) + return self._completed_publish( + dataset, + source, + expected_item, + existing_item, + ) + return self._start_item_operation( + expected_item, + metadata, + collection_id, + item_id, + ) + + def _start_item_operation( + self, + expected_item: Mapping[str, Any], + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + try: + operation = self.sdk.start_create_item( + collection_id, + item_id, + expected_item, + ) + except Exception as error: + if not self._is_status(error, 409): + raise + return self._verification_pending( + metadata, + PlanetaryComputerPhase.ITEM_VERIFY, + self._item_href(collection_id, item_id), + ) + return self._operation_pending( + metadata, + PlanetaryComputerPhase.ITEM_OPERATION, + operation.continuation_token, + ) + + def _start_replace_item_delete( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + expected_item: Mapping[str, Any], + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + try: + operation = self.sdk.start_delete_item(collection_id, item_id) + except Exception as error: + if self._is_status(error, 404): + return self._start_item_operation( + expected_item, + metadata, + collection_id, + item_id, + ) + if not self._is_status(error, 409): + raise + return self._verification_pending( + metadata, + PlanetaryComputerPhase.ITEM_REPLACE_DELETE_VERIFY, + self._item_href(collection_id, item_id), + ) + return self._operation_pending( + metadata, + PlanetaryComputerPhase.ITEM_REPLACE_DELETE_OPERATION, + operation.continuation_token, + ) + + def _replace_delete_complete_or_pending( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + projection_codes: Mapping[str, str], + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + if self.sdk.get_item(collection_id, item_id) is not None: + return self._verification_pending( + metadata, + PlanetaryComputerPhase.ITEM_REPLACE_DELETE_VERIFY, + self._item_href(collection_id, item_id), + ) + return self._collection_ready_or_pending( + dataset, + source, + projection_codes, + metadata, + collection_id, + item_id, + ) + + def _item_complete_or_pending( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + projection_codes: Mapping[str, str], + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + actual_item = self.sdk.get_item(collection_id, item_id) + if actual_item is None: + return self._verification_pending( + metadata, + PlanetaryComputerPhase.ITEM_VERIFY, + self._item_href(collection_id, item_id), + ) + collection = self.sdk.get_collection(collection_id) + if collection is None: + raise PlanetaryComputerProviderError( + "Planetary Computer collection disappeared after ingestion" + ) + documents = self._build_documents( + dataset, + source, + projection_codes, + collection, + ) + return self._completed_publish( + dataset, + source, + documents.item, + actual_item, + ) + + def _delete_complete_or_pending( + self, + dataset: PublishedDataset, + metadata: Mapping[str, Any], + collection_id: str, + item_id: str, + ) -> PublishResult: + if self.sdk.get_item(collection_id, item_id) is None: + return self._cleanup_collection_or_complete(dataset, collection_id) + return self._verification_pending( + metadata, + PlanetaryComputerPhase.DELETE_VERIFY, + self._item_href(collection_id, item_id), + ) + + def _completed_publish( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + expected_item: Mapping[str, Any], + actual_item: Mapping[str, Any], + ) -> PublishResult: + published_artifacts = self._verify_item( + source, + expected_item, + actual_item, + ) + collection_id, _ = self._ids(dataset) + metadata = self._stable_metadata(dataset) + metadata["assetsCopiedToManagedStorage"] = True + return PublishResult( + artifacts=published_artifacts, + links={ + "stac_collection": self._collection_href(collection_id), + "explorer": self.explorer_url, + }, + providerMetadata=metadata, + ) + + def _verify_item( + self, + source: ArtifactBundle, + expected_item: Mapping[str, Any], + actual_item: Mapping[str, Any], + *, + verify_reachability: bool = True, + ) -> list[PublishedArtifact]: + if actual_item.get("id") != expected_item.get("id"): + raise PlanetaryComputerProviderError( + "Planetary Computer item ID verification failed" + ) + if actual_item.get("collection") != expected_item.get("collection"): + raise PlanetaryComputerProviderError( + "Planetary Computer collection verification failed" + ) + expected_geometry = expected_item.get("geometry") + actual_geometry = actual_item.get("geometry") + if not expected_geometry or not actual_geometry: + raise PlanetaryComputerProviderError( + "Planetary Computer item geometry is missing" + ) + if not shape(expected_geometry).equals(shape(actual_geometry)): + raise PlanetaryComputerProviderError( + "Planetary Computer item geometry verification failed" + ) + self._verify_item_bbox(expected_item, actual_item) + self._verify_item_properties(expected_item, actual_item) + + expected_assets = expected_item.get("assets") or {} + actual_assets = actual_item.get("assets") or {} + if not isinstance(expected_assets, Mapping) or not isinstance( + actual_assets, Mapping + ): + raise PlanetaryComputerProviderError( + "Planetary Computer item assets are invalid" + ) + # The thumbnail is a best-effort preview: don't fail verification (or + # the whole publish) if the GeoCatalog drops or adds it. + expected_asset_keys = set(expected_assets) - {THUMBNAIL_ASSET_KEY} + actual_asset_keys = set(actual_assets) - {THUMBNAIL_ASSET_KEY} + if actual_asset_keys != expected_asset_keys: + raise PlanetaryComputerProviderError( + "Planetary Computer selected assets changed" + ) + published_artifacts = [] + for artifact in source.selectedArtifacts: + asset_key = ASSET_KEYS[artifact.kind] + expected_asset = expected_assets.get(asset_key) or {} + actual_asset = actual_assets.get(asset_key) or {} + if not isinstance(actual_asset, Mapping): + raise PlanetaryComputerProviderError( + f"Planetary Computer asset is missing: {asset_key}" + ) + for field in ("type", "title", "proj:code"): + if actual_asset.get(field) != expected_asset.get(field): + raise PlanetaryComputerProviderError( + f"Planetary Computer asset {field} changed: " + f"{asset_key}" + ) + expected_roles = set(expected_asset.get("roles") or []) + actual_roles = set(actual_asset.get("roles") or []) + if expected_roles != actual_roles: + raise PlanetaryComputerProviderError( + f"Planetary Computer asset roles changed: {asset_key}" + ) + managed_href = self._managed_asset_href(actual_asset.get("href")) + if self._same_url(managed_href, expected_asset.get("href")): + raise PlanetaryComputerProviderError( + f"Planetary Computer did not copy asset: {asset_key}" + ) + if verify_reachability: + self._verify_managed_asset_reachable(managed_href) + published_artifacts.append( + PublishedArtifact( + **artifact.model_dump(), + publishedPath=managed_href, + ) + ) + return published_artifacts + + @staticmethod + def _verify_item_bbox( + expected_item: Mapping[str, Any], + actual_item: Mapping[str, Any], + ) -> None: + expected_bbox = expected_item.get("bbox") + actual_bbox = actual_item.get("bbox") + if not ( + isinstance(expected_bbox, list) + and isinstance(actual_bbox, list) + and len(expected_bbox) == len(actual_bbox) == 4 + ): + raise PlanetaryComputerProviderError( + "Planetary Computer item bbox is invalid" + ) + try: + matches = all( + math.isclose( + float(expected), + float(actual), + rel_tol=0, + abs_tol=1e-9, + ) + for expected, actual in zip(expected_bbox, actual_bbox) + ) + except (TypeError, ValueError): + matches = False + if not matches: + raise PlanetaryComputerProviderError( + "Planetary Computer item bbox verification failed" + ) + + @staticmethod + def _verify_item_properties( + expected_item: Mapping[str, Any], + actual_item: Mapping[str, Any], + ) -> None: + expected = expected_item.get("properties") + actual = actual_item.get("properties") + if not isinstance(expected, Mapping) or not isinstance( + actual, Mapping + ): + raise PlanetaryComputerProviderError( + "Planetary Computer item properties are invalid" + ) + for key, expected_value in expected.items(): + if actual.get(key) != expected_value: + raise PlanetaryComputerProviderError( + f"Planetary Computer item property changed: {key}" + ) + + def _build_documents( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + projection_codes: Mapping[str, str], + existing_collection: Optional[Mapping[str, Any]], + ): + mask = source.get(ArtifactKind.VALID_MASK) + if mask is None: + raise ValueError("Planetary Computer requires a valid-area mask") + valid_mask = self.json_reader(mask) + asset_hrefs = { + artifact.sourcePath: self._artifact_href(artifact) + for artifact in source.selectedArtifacts + } + collection_id = build_collection_id(dataset, self.collection_prefix) + thumbnail_href, thumbnail_media_type = self._resolve_thumbnail_href( + dataset, source + ) + objects = build_stac_objects( + dataset, + source, + valid_mask, + asset_hrefs, + projection_codes, + self._collection_href(collection_id), + valid_mask_crs=self._valid_mask_crs(valid_mask), + existing_collection=existing_collection, + collection_prefix=self.collection_prefix, + license_id=self.license_id, + thumbnail_href=thumbnail_href, + thumbnail_media_type=thumbnail_media_type, + ) + self.stac_validator(objects) + return serialize_stac_objects(objects) + + def _resolve_thumbnail_href( + self, dataset: PublishedDataset, source: ArtifactBundle + ) -> tuple[Optional[str], str]: + # Copy the post-event preview into the dataset's published prefix (the + # ingestion-source container the GeoCatalog can read) and return a plain + # blob href for it. Best-effort: any failure yields no thumbnail rather + # than breaking the publish. + preview_url = getattr(source, "thumbnailUrl", None) + if not preview_url: + return None, DEFAULT_THUMBNAIL_MEDIA_TYPE + try: + relative_path = self._preview_relative_path(preview_url) + if not relative_path or not self.artifact_storage.artifact_exists( + relative_path + ): + return None, DEFAULT_THUMBNAIL_MEDIA_TYPE + extension = PurePosixPath(relative_path).suffix.lower() or ".png" + media_type = _THUMBNAIL_MEDIA_TYPES.get( + extension, DEFAULT_THUMBNAIL_MEDIA_TYPE + ) + destination = ( + f"published/{dataset.datasetId}/thumbnail{extension}" + ) + etag = self.artifact_storage.get_artifact_etag(relative_path) + published_path = self.artifact_storage.copy_artifact( + relative_path, destination, etag + ) + resolved = self.artifact_storage.resolve_artifact_path( + published_path + ) + base_url = str(self.artifact_storage.get_base_url()).rstrip("/") + href = f"{base_url}/{quote(resolved, safe='/-_.~')}" + return href, media_type + except Exception as error: + self.logger.warning( + "Skipping Planetary Computer thumbnail: %s", + type(error).__name__, + ) + return None, DEFAULT_THUMBNAIL_MEDIA_TYPE + + def _preview_relative_path(self, preview_url: str) -> Optional[str]: + # A preview URL is only usable if it lives in the same account+container + # as the published artifacts (the container the ingestion source reads). + base = urlparse(str(self.artifact_storage.get_base_url())) + parsed = urlparse(preview_url) + if ( + parsed.scheme != "https" + or not parsed.hostname + or parsed.hostname != base.hostname + ): + return None + base_path = base.path.rstrip("/") + if base_path: + if not parsed.path.startswith(base_path + "/"): + return None + relative = parsed.path[len(base_path) + 1 :] + else: + relative = parsed.path.lstrip("/") + return unquote(relative) or None + + @staticmethod + def _valid_mask_crs(valid_mask: Mapping[str, Any]) -> str: + crs_value = valid_mask.get("crs") + if crs_value is None: + return "EPSG:4326" + if not isinstance(crs_value, Mapping): + raise ValueError("Valid-area mask CRS is invalid") + properties = crs_value.get("properties") + if not isinstance(properties, Mapping): + raise ValueError("Valid-area mask CRS is invalid") + name = properties.get("name") + if crs_value.get("type") != "name" or not isinstance(name, str): + raise ValueError("Valid-area mask CRS is invalid") + try: + parsed = CRS.from_user_input(name) + except CRSError as error: + raise ValueError("Valid-area mask CRS is invalid") from error + return parsed.to_string() + + def _projection_codes( + self, + dataset: PublishedDataset, + source: ArtifactBundle, + ) -> dict[str, str]: + persisted = dataset.providerMetadata.get("projectionCodes") or {} + if not isinstance(persisted, Mapping): + raise PlanetaryComputerProviderError( + "Planetary Computer projection metadata is invalid" + ) + projection_codes = {} + for artifact in source.selectedArtifacts: + if artifact.kind == ArtifactKind.VALID_MASK: + continue + projection_code = persisted.get(artifact.sourcePath) + if not projection_code: + projection_code = self.projection_resolver(artifact) + projection_codes[artifact.sourcePath] = str(projection_code) + return projection_codes + + def _validate_bundle(self, source: ArtifactBundle) -> None: + if not source.selectedArtifacts: + raise ValueError("Select at least one artifact to publish") + supported = set(self.info.supportedArtifactKinds) + unsupported = { + artifact.kind + for artifact in source.selectedArtifacts + if artifact.kind not in supported + } + if unsupported: + names = ", ".join(sorted(kind.value for kind in unsupported)) + raise ValueError( + f"Unsupported Planetary Computer artifacts: {names}" + ) + mask = source.get(ArtifactKind.VALID_MASK) + if mask is None: + raise ValueError("Planetary Computer requires a valid-area mask") + if mask.sizeBytes and mask.sizeBytes > MAX_VALID_MASK_BYTES: + raise ValueError("Valid-area mask is too large") + + def _require_configuration(self) -> None: + if not self.info.isEnabled or not self.info.isConfigured: + raise PlanetaryComputerProviderError( + self.info.disabledReason + or "Planetary Computer provider is unavailable" + ) + if self.ingestion_source and not re.fullmatch( + r"[A-Za-z0-9._-]{1,256}", self.ingestion_source + ): + raise PlanetaryComputerProviderError( + "Planetary Computer ingestion source ID is invalid" + ) + if not 1 <= self.max_verify_attempts <= 60: + raise PlanetaryComputerProviderError( + "Planetary Computer verification limit is invalid" + ) + self._safe_https_url(self.endpoint, allow_path=False) + # The Explorer URL is a display-only link surfaced on the dataset row. + # MPC Pro's real Explorer links carry a query string + # (e.g. ?geocatalogname=...&c=...&z=...), so allow query/fragment here + # while keeping the GeoCatalog and asset URLs strict. + self._safe_https_url( + self.explorer_url, allow_path=True, allow_query=True + ) + + def _validate_ingestion_source(self) -> None: + source = self.sdk.get_ingestion_source(self.ingestion_source) + if source is None: + raise PlanetaryComputerProviderError( + "Planetary Computer ingestion source was not found" + ) + kind = self._mapping_value(source, "kind") + if kind not in {"BlobManagedIdentity", "SasToken"}: + raise PlanetaryComputerProviderError( + "Planetary Computer ingestion source type is unsupported" + ) + connection = self._mapping_value( + source, + "connectionInfo", + "connection_info", + ) + # MPC Pro returns the blob container as `containerUri` (camelCase); + # accept the other spellings defensively. + container_url = self._mapping_value( + connection, + "containerUri", + "containerUrl", + "container_uri", + "container_url", + ) + expected_url = str(self.artifact_storage.get_base_url()).rstrip("/") + if not isinstance(container_url, str) or not self._same_url( + expected_url, + container_url.rstrip("/"), + ): + raise PlanetaryComputerProviderError( + "Planetary Computer ingestion source container does not " + "match HASTE storage" + ) + + def _artifact_href(self, artifact: SourceArtifact) -> str: + base_url = str(self.artifact_storage.get_base_url()).rstrip("/") + parsed = self._safe_https_url(base_url, allow_path=True) + if str(parsed.hostname).lower().startswith("devstoreaccount1"): + raise ValueError( + "Planetary Computer cannot ingest from the storage emulator" + ) + if not self._is_azure_blob_host(str(parsed.hostname)): + raise ValueError( + "Planetary Computer source must use Azure Blob Storage" + ) + relative_path = self.artifact_storage.resolve_artifact_path( + artifact.sourcePath + ) + encoded_path = quote(relative_path, safe="/-_.~") + return f"{base_url}/{encoded_path}" + + def _managed_asset_href(self, value: Any) -> str: + if not isinstance(value, str): + raise PlanetaryComputerProviderError( + "Planetary Computer managed asset URL is missing" + ) + parsed = self._safe_https_url(value, allow_path=True) + if parsed.query: + raise PlanetaryComputerProviderError( + "Planetary Computer managed asset URL must not contain a token" + ) + if not self._is_azure_blob_host(str(parsed.hostname)): + raise PlanetaryComputerProviderError( + "Planetary Computer managed asset URL is not Azure Blob " + "Storage" + ) + return value + + def _verify_managed_asset_reachable(self, managed_href: str) -> None: + try: + signed_href = self.sdk.get_signed_asset_url(managed_href) + self._validate_signed_asset_url(managed_href, signed_href) + self.asset_reachability_checker(signed_href) + except PlanetaryComputerProviderError: + raise + except Exception as error: + raise PlanetaryComputerProviderError( + "Planetary Computer managed asset is not reachable" + ) from error + + @classmethod + def _validate_signed_asset_url( + cls, + managed_href: str, + signed_href: str, + ) -> None: + if not isinstance(signed_href, str) or len(signed_href) > 8192: + raise PlanetaryComputerProviderError( + "Planetary Computer signed asset URL is invalid" + ) + unsigned = urlparse(managed_href) + signed = urlparse(signed_href) + try: + signed_port = signed.port + except ValueError as error: + raise PlanetaryComputerProviderError( + "Planetary Computer signed asset URL is invalid" + ) from error + if ( + signed.scheme != "https" + or not signed.hostname + or signed.username is not None + or signed.password is not None + or signed.fragment + or not signed.query + or (signed_port is not None and not 1 <= signed_port <= 65535) + or signed.hostname.lower() != str(unsigned.hostname).lower() + or signed_port != unsigned.port + or signed.path != unsigned.path + ): + raise PlanetaryComputerProviderError( + "Planetary Computer signed asset URL is invalid" + ) + if not cls._is_azure_blob_host(signed.hostname): + raise PlanetaryComputerProviderError( + "Planetary Computer signed asset URL is not Azure Blob " + "Storage" + ) + + @staticmethod + def _read_signed_asset(signed_href: str) -> None: + parsed = urlparse(signed_href) + connection = HTTPSConnection( + str(parsed.hostname), + parsed.port or 443, + timeout=30, + ) + target = parsed.path or "/" + if parsed.query: + target = f"{target}?{parsed.query}" + try: + connection.request( + "GET", + target, + headers={"Range": "bytes=0-0"}, + ) + response = connection.getresponse() + if response.status not in {200, 206}: + raise PlanetaryComputerProviderError( + "Planetary Computer managed asset is not reachable" + ) + response.read(1) + except PlanetaryComputerProviderError: + raise + except Exception as error: + raise PlanetaryComputerProviderError( + "Planetary Computer managed asset is not reachable" + ) from error + finally: + connection.close() + + @staticmethod + def _mapping_value(value: Any, *names: str) -> Any: + if isinstance(value, Mapping): + for name in names: + if name in value: + return value[name] + for name in names: + if hasattr(value, name): + return getattr(value, name) + return None + + @staticmethod + def _same_url(first: Any, second: Any) -> bool: + if not isinstance(first, str) or not isinstance(second, str): + return False + first_url = urlparse(first) + second_url = urlparse(second) + try: + first_port = first_url.port + second_port = second_url.port + except ValueError: + return False + return ( + first_url.scheme.lower() == second_url.scheme.lower() + and str(first_url.hostname).lower() + == str(second_url.hostname).lower() + and first_port == second_port + and first_url.path == second_url.path + and first_url.params == second_url.params + and first_url.query == second_url.query + and first_url.fragment == second_url.fragment + ) + + @staticmethod + def _is_azure_blob_host(hostname: str) -> bool: + host = hostname.lower().rstrip(".") + return any( + host.endswith(suffix) for suffix in AZURE_BLOB_HOST_SUFFIXES + ) + + @staticmethod + def _safe_https_url( + value: str, *, allow_path: bool, allow_query: bool = False + ): + parsed = urlparse(value) + try: + port = parsed.port + except ValueError as error: + raise ValueError("Planetary Computer URL is invalid") from error + if ( + parsed.scheme != "https" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or (not allow_query and (parsed.fragment or parsed.query)) + or (port is not None and not 1 <= port <= 65535) + or (not allow_path and parsed.path not in {"", "/"}) + ): + raise ValueError("Planetary Computer URL must use HTTPS") + return parsed + + def _operation_metadata( + self, + dataset: PublishedDataset, + projection_codes: Mapping[str, str], + phase: Optional[PlanetaryComputerPhase], + ) -> dict[str, Any]: + metadata = self._stable_metadata(dataset) + metadata["projectionCodes"] = dict(projection_codes) + metadata["verificationAttempts"] = 0 + if phase is not None: + metadata["phase"] = phase.value + return metadata + + def _stable_metadata(self, dataset: PublishedDataset) -> dict[str, Any]: + collection_id = str( + dataset.providerMetadata.get("collectionId") + or build_collection_id(dataset, self.collection_prefix) + ) + item_ids = dataset.providerMetadata.get("itemIds") or [ + build_item_id(dataset) + ] + if not isinstance(item_ids, list) or len(item_ids) != 1: + raise PlanetaryComputerProviderError( + "Planetary Computer item metadata is invalid" + ) + return { + "collectionId": collection_id, + "itemIds": [str(item_ids[0])], + "apiVersion": PC_API_VERSION, + "ingestionSource": self.ingestion_source, + } + + def _operation_pending( + self, + metadata: Mapping[str, Any], + phase: PlanetaryComputerPhase, + continuation_token: Optional[str], + *, + count_attempt: bool = False, + ) -> PublishResult: + if not continuation_token: + raise PlanetaryComputerProviderError( + "Planetary Computer operation URL is missing" + ) + updated = dict(metadata) + updated["phase"] = phase.value + updated["verificationAttempts"] = 0 + attempts = int(metadata.get("operationAttempts") or 0) + if count_attempt: + attempts += 1 + if attempts > self.max_verify_attempts: + raise PlanetaryComputerProviderError( + "Planetary Computer ingestion timed out" + ) + else: + attempts = 0 + updated["operationAttempts"] = attempts + return PublishResult( + providerMetadata=updated, + continuationToken=continuation_token, + isComplete=False, + ) + + def _verification_pending( + self, + metadata: Mapping[str, Any], + phase: PlanetaryComputerPhase, + continuation_token: str, + ) -> PublishResult: + attempts = int(metadata.get("verificationAttempts") or 0) + 1 + if attempts > self.max_verify_attempts: + raise PlanetaryComputerProviderError( + "Planetary Computer verification timed out" + ) + updated = dict(metadata) + updated["phase"] = phase.value + updated["verificationAttempts"] = attempts + updated["operationAttempts"] = 0 + return PublishResult( + providerMetadata=updated, + continuationToken=continuation_token, + isComplete=False, + ) + + @staticmethod + def _phase(metadata: Mapping[str, Any]) -> PlanetaryComputerPhase: + try: + return PlanetaryComputerPhase(str(metadata.get("phase"))) + except ValueError as error: + raise PlanetaryComputerProviderError( + "Planetary Computer operation phase is missing" + ) from error + + @staticmethod + def _continuation_token(metadata: Mapping[str, Any]) -> str: + token = metadata.get("continuationToken") + if not isinstance(token, str) or not token: + raise PlanetaryComputerProviderError( + "Planetary Computer operation URL is missing" + ) + return token + + def _ids(self, dataset: PublishedDataset) -> tuple[str, str]: + metadata = self._stable_metadata(dataset) + return metadata["collectionId"], metadata["itemIds"][0] + + def _collection_href(self, collection_id: str) -> str: + encoded = quote(collection_id, safe="-_.") + return f"{self.endpoint}/stac/collections/{encoded}" + + def _item_href(self, collection_id: str, item_id: str) -> str: + encoded_collection = quote(collection_id, safe="-_.") + encoded_item = quote(item_id, safe="-_+,.()") + return ( + f"{self.endpoint}/stac/collections/{encoded_collection}/items/" + f"{encoded_item}" + ) + + @staticmethod + def _is_status(error: Exception, status_code: int) -> bool: + direct_status = getattr(error, "status_code", None) + response = getattr(error, "response", None) + response_status = getattr(response, "status_code", None) + return direct_status == status_code or response_status == status_code + + def _read_json_artifact( + self, artifact: SourceArtifact + ) -> Mapping[str, Any]: + with self._materialized_artifact(artifact) as local_path: + with open(local_path, encoding="utf-8") as source_file: + value = json.load(source_file) + if not isinstance(value, Mapping): + raise ValueError("Valid-area mask JSON must be an object") + return value + + def _read_projection_code(self, artifact: SourceArtifact) -> str: + harden_gdal() + import fiona + + with self._materialized_artifact(artifact) as local_path: + layers = fiona.listlayers(local_path) + if not layers: + raise ValueError("GeoPackage has no vector layers") + source_crs_values = [] + for layer in layers: + with fiona.open(local_path, layer=layer) as source: + source_crs_values.append(source.crs_wkt or source.crs) + if any(not value for value in source_crs_values): + raise ValueError( + f"Artifact has no CRS: {PurePosixPath(artifact.sourcePath).name}" + ) + try: + epsg_codes = { + CRS.from_user_input(value).to_epsg() + for value in source_crs_values + } + except CRSError as error: + raise ValueError("Artifact CRS is invalid") from error + if None in epsg_codes: + raise ValueError("Artifact CRS has no EPSG authority code") + if len(epsg_codes) != 1: + raise ValueError("GeoPackage layers use different CRS values") + epsg = epsg_codes.pop() + return f"EPSG:{epsg}" + + @contextmanager + def _materialized_artifact( + self, artifact: SourceArtifact + ) -> Iterator[Path]: + relative_path = self.artifact_storage.resolve_artifact_path( + artifact.sourcePath + ) + with tempfile.TemporaryDirectory() as directory: + self.artifact_storage.fetch_artifact( + src_path=relative_path, + dst_path=directory, + ) + local_path = Path(directory, relative_path) + if not local_path.is_file(): + matches = list( + Path(directory).rglob(PurePosixPath(relative_path).name) + ) + if len(matches) != 1 or not matches[0].is_file(): + raise FileNotFoundError(artifact.sourcePath) + local_path = matches[0] + yield local_path diff --git a/hastelib/src/hastegeo/core/publishing/planetary_computer_transport.py b/hastelib/src/hastegeo/core/publishing/planetary_computer_transport.py new file mode 100644 index 00000000..054a3905 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/planetary_computer_transport.py @@ -0,0 +1,510 @@ +"""REST transport adapter for the Planetary Computer Pro GeoCatalog. + +Presents the resumable, one-request-per-step interface the publishing provider +depends on (start/continue create-collection, create-item, delete-item; plus +get-collection/item/ingestion-source and asset signing), backed by the vendored +``GeoCatalogClient`` REST client instead of the ``azure-planetarycomputer`` SDK. + +The GeoCatalog ingests items asynchronously: ``POST .../items`` returns 202 with +an ``operation-location`` header that is polled to a terminal state. Collection +creation is synchronous (201). ``_start_from_response`` handles both: a 202 +yields an incomplete step carrying the (origin-pinned) operation URL; a 2xx +without a 202 yields a completed step. + +Security controls carried over from the SDK adapter (and absent from the +reference REST client): operation-URL SSRF pinning to the GeoCatalog origin, +endpoint validation, sanitized error/status text (no server bodies), and +strict failed-item accounting. Redirect suppression and per-request timeouts +live in ``GeoCatalogClient``. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Mapping, Optional +from urllib.parse import parse_qsl, urlparse + +from .geocatalog_client import GeoCatalogAuth, GeoCatalogClient + + +class PlanetaryComputerOperationKind(str, Enum): + CREATE_COLLECTION = "create_collection" + CREATE_ITEM = "create_item" + DELETE_ITEM = "delete_item" + DELETE_COLLECTION = "delete_collection" + + +class PlanetaryComputerOperationError(RuntimeError): + """Raised when a GeoCatalog long-running operation fails.""" + + +@dataclass(frozen=True) +class PlanetaryComputerOperationStep: + kind: PlanetaryComputerOperationKind + collection_id: str + item_id: Optional[str] + continuation_token: Optional[str] = field(repr=False) + is_complete: bool + + +class PlanetaryComputerRestAdapter: + """Run one authenticated GeoCatalog REST request per adapter step.""" + + _IN_PROGRESS_STATUSES = { + "inprogress", + "notstarted", + "pending", + "running", + "accepted", + } + # Includes "finished": the GeoCatalog reports terminal success as either + # Succeeded or Finished; omitting the latter would fail a good ingest. + _SUCCESS_STATUSES = {"completed", "succeeded", "success", "finished"} + _FAILURE_STATUSES = {"cancelled", "canceled", "failed"} + + def __init__( + self, + endpoint: str, + *, + client: Any = None, + credential: Any = None, + connection_timeout: int = 10, + read_timeout: int = 30, + ) -> None: + self.endpoint = self._validate_endpoint(endpoint) + if connection_timeout < 1 or read_timeout < 1: + raise ValueError("Planetary Computer timeouts must be positive") + self.connection_timeout = connection_timeout + self.read_timeout = read_timeout + self._credential = credential + self._client = client + + @property + def client(self) -> GeoCatalogClient: + if self._client is None: + self._client = GeoCatalogClient( + self.endpoint, + auth=GeoCatalogAuth(self._credential), + connection_timeout=self.connection_timeout, + read_timeout=self.read_timeout, + ) + return self._client + + # ------------------------------------------------------------ collections + + def get_collection( + self, collection_id: str + ) -> Optional[Mapping[str, Any]]: + response = self.client.request( + "GET", + f"/stac/collections/{collection_id}", + expected=(200, 404), + ) + if response.status_code == 404: + return None + return self._as_mapping(response.json()) + + def start_create_collection( + self, + collection_id: str, + body: Mapping[str, Any], + ) -> PlanetaryComputerOperationStep: + response = self.client.request( + "POST", + "/stac/collections", + json=dict(body), + expected=(200, 201, 202), + ) + return self._start_from_response( + response, + PlanetaryComputerOperationKind.CREATE_COLLECTION, + collection_id, + ) + + def replace_collection( + self, + collection_id: str, + body: Mapping[str, Any], + ) -> None: + self.client.request( + "PUT", + f"/stac/collections/{collection_id}", + json=dict(body), + expected=(200, 201, 202, 204), + ) + + def continue_create_collection( + self, + collection_id: str, + continuation_token: str, + ) -> PlanetaryComputerOperationStep: + return self._continue( + PlanetaryComputerOperationKind.CREATE_COLLECTION, + collection_id, + None, + continuation_token, + ) + + def start_delete_collection( + self, + collection_id: str, + ) -> PlanetaryComputerOperationStep: + response = self.client.request( + "DELETE", + f"/stac/collections/{collection_id}", + expected=(200, 202, 204), + ) + return self._start_from_response( + response, + PlanetaryComputerOperationKind.DELETE_COLLECTION, + collection_id, + ) + + def continue_delete_collection( + self, + collection_id: str, + continuation_token: str, + ) -> PlanetaryComputerOperationStep: + return self._continue( + PlanetaryComputerOperationKind.DELETE_COLLECTION, + collection_id, + None, + continuation_token, + ) + + def list_item_ids( + self, collection_id: str, limit: int = 100 + ) -> list: + """Return the ids of items currently in the collection (bounded).""" + response = self.client.request( + "GET", + f"/stac/collections/{collection_id}/items", + params={"limit": limit}, + expected=(200, 404), + ) + if response.status_code == 404: + return [] + payload = response.json() + features = [] + if isinstance(payload, Mapping): + features = payload.get("features") or [] + return [ + str(feature["id"]) + for feature in features + if isinstance(feature, Mapping) and feature.get("id") + ] + + # ------------------------------------------------------------------ items + + def get_item( + self, + collection_id: str, + item_id: str, + ) -> Optional[Mapping[str, Any]]: + response = self.client.request( + "GET", + f"/stac/collections/{collection_id}/items/{item_id}", + expected=(200, 404), + ) + if response.status_code == 404: + return None + return self._as_mapping(response.json()) + + def start_create_item( + self, + collection_id: str, + item_id: str, + body: Mapping[str, Any], + ) -> PlanetaryComputerOperationStep: + response = self.client.request( + "POST", + f"/stac/collections/{collection_id}/items", + json=dict(body), + expected=(200, 201, 202), + ) + return self._start_from_response( + response, + PlanetaryComputerOperationKind.CREATE_ITEM, + collection_id, + item_id, + ) + + def continue_create_item( + self, + collection_id: str, + item_id: str, + continuation_token: str, + ) -> PlanetaryComputerOperationStep: + return self._continue( + PlanetaryComputerOperationKind.CREATE_ITEM, + collection_id, + item_id, + continuation_token, + ) + + def start_delete_item( + self, + collection_id: str, + item_id: str, + ) -> PlanetaryComputerOperationStep: + response = self.client.request( + "DELETE", + f"/stac/collections/{collection_id}/items/{item_id}", + expected=(200, 202, 204), + ) + return self._start_from_response( + response, + PlanetaryComputerOperationKind.DELETE_ITEM, + collection_id, + item_id, + ) + + def continue_delete_item( + self, + collection_id: str, + item_id: str, + continuation_token: str, + ) -> PlanetaryComputerOperationStep: + return self._continue( + PlanetaryComputerOperationKind.DELETE_ITEM, + collection_id, + item_id, + continuation_token, + ) + + # ------------------------------------------------------- ingestion / SAS + + def get_ingestion_source( + self, source_id: str + ) -> Optional[Mapping[str, Any]]: + # The REST API exposes only a list endpoint; find the source by id. + response = self.client.request( + "GET", "/inma/ingestion-sources", expected=(200,) + ) + payload = response.json() + if isinstance(payload, Mapping): + sources = payload.get("value", []) + elif isinstance(payload, list): + sources = payload + else: + sources = [] + for source in sources: + if isinstance(source, Mapping) and str(source.get("id")) == str( + source_id + ): + return self._as_mapping(source) + return None + + def get_signed_asset_url(self, href: str) -> str: + response = self.client.request( + "GET", "/sas/sign", params={"href": href}, expected=(200,) + ) + signed_link = self._as_mapping(response.json()) + signed_href = signed_link.get("href") + if not isinstance(signed_href, str) or not signed_href: + raise PlanetaryComputerOperationError( + "Planetary Computer returned an invalid signed asset URL" + ) + return signed_href + + def close(self) -> None: + if self._client is not None: + close = getattr(self._client, "close", None) + if close is not None: + close() + + # ------------------------------------------------------------- internals + + def _start_from_response( + self, + response: Any, + kind: PlanetaryComputerOperationKind, + collection_id: str, + item_id: Optional[str] = None, + ) -> PlanetaryComputerOperationStep: + if response.status_code == 202: + headers = { + str(key).lower(): value + for key, value in response.headers.items() + } + operation_url = headers.get("operation-location") or headers.get( + "location" + ) + operation_url = self._validate_operation_url(operation_url) + return PlanetaryComputerOperationStep( + kind=kind, + collection_id=collection_id, + item_id=item_id, + continuation_token=operation_url, + is_complete=False, + ) + # Synchronous completion (e.g. 201 Created for a collection). + return PlanetaryComputerOperationStep( + kind=kind, + collection_id=collection_id, + item_id=item_id, + continuation_token=None, + is_complete=True, + ) + + def _continue( + self, + kind: PlanetaryComputerOperationKind, + collection_id: str, + item_id: Optional[str], + continuation_token: str, + ) -> PlanetaryComputerOperationStep: + operation_url = self._validate_operation_url(continuation_token) + response = self.client.request( + "GET", operation_url, expected=(200,), absolute=True + ) + payload = response.json() + if not isinstance(payload, Mapping): + raise PlanetaryComputerOperationError( + "Planetary Computer returned an invalid operation response" + ) + status = str(payload.get("status") or "").replace("_", "").lower() + if status in self._IN_PROGRESS_STATUSES: + return PlanetaryComputerOperationStep( + kind=kind, + collection_id=collection_id, + item_id=item_id, + continuation_token=operation_url, + is_complete=False, + ) + if status in self._SUCCESS_STATUSES: + failed_items = self._failed_item_count(payload) + if failed_items: + raise PlanetaryComputerOperationError( + "Planetary Computer operation failed " + f"({failed_items} items)" + ) + return PlanetaryComputerOperationStep( + kind=kind, + collection_id=collection_id, + item_id=item_id, + continuation_token=None, + is_complete=True, + ) + if status in self._FAILURE_STATUSES: + error = payload.get("error") or {} + if not isinstance(error, Mapping): + error = {} + code = re.sub( + r"[^A-Za-z0-9_.-]", "", str(error.get("code") or "") + ) + detail = f" ({code})" if code else "" + raise PlanetaryComputerOperationError( + f"Planetary Computer operation {status}{detail}" + ) + raise PlanetaryComputerOperationError( + "Planetary Computer returned an unknown operation status" + ) + + @staticmethod + def _failed_item_count(payload: Mapping[str, Any]) -> int: + value = None + candidates = [payload] + for key in ("additionalInformation", "additional_information"): + additional = payload.get(key) + if isinstance(additional, Mapping): + candidates.append(additional) + for candidate in candidates: + for key, candidate_value in candidate.items(): + normalized_key = str(key).replace("_", "").lower() + if normalized_key == "totalfaileditems": + value = candidate_value + break + if value is not None: + break + if value is None: + return 0 + if isinstance(value, bool): + raise PlanetaryComputerOperationError( + "Planetary Computer returned an invalid failed-item count" + ) + try: + count = int(value) + except (TypeError, ValueError) as error: + raise PlanetaryComputerOperationError( + "Planetary Computer returned an invalid failed-item count" + ) from error + if count < 0: + raise PlanetaryComputerOperationError( + "Planetary Computer returned an invalid failed-item count" + ) + return count + + @staticmethod + def _validate_endpoint(endpoint: str) -> str: + parsed = urlparse(endpoint) + try: + port = parsed.port + except ValueError as error: + raise ValueError( + "Planetary Computer endpoint is invalid" + ) from error + if port is not None and not 1 <= port <= 65535: + raise ValueError("Planetary Computer endpoint is invalid") + if ( + parsed.scheme != "https" + or not parsed.hostname + or parsed.username is not None + or parsed.password is not None + or parsed.query + or parsed.fragment + or parsed.path not in {"", "/"} + ): + raise ValueError( + "Planetary Computer endpoint must be an HTTPS origin" + ) + return endpoint.rstrip("/") + + @classmethod + def _as_mapping(cls, value: Any) -> Mapping[str, Any]: + if isinstance(value, Mapping): + return dict(value) + raise TypeError("Planetary Computer response is not a mapping") + + def _validate_operation_url(self, value: Optional[str]) -> str: + if not value or len(value) > 4096: + raise ValueError("Planetary Computer operation URL is invalid") + endpoint = urlparse(self.endpoint) + operation = urlparse(value) + try: + operation_port = operation.port + except ValueError as error: + raise ValueError( + "Planetary Computer operation URL is invalid" + ) from error + endpoint_port = endpoint.port + if operation_port is not None and not 1 <= operation_port <= 65535: + raise ValueError("Planetary Computer operation URL is invalid") + if ( + operation.scheme != "https" + or operation.hostname != endpoint.hostname + or (operation_port if operation_port is not None else 443) + != (endpoint_port if endpoint_port is not None else 443) + or operation.username is not None + or operation.password is not None + or operation.fragment + or not operation.path + ): + raise ValueError( + "Planetary Computer operation URL must use the " + "GeoCatalog origin" + ) + query_names = { + name.lower() + for name, _ in parse_qsl( + operation.query, + keep_blank_values=True, + ) + } + if not query_names.issubset({"api-version"}): + raise ValueError( + "Planetary Computer operation URL has unsupported query fields" + ) + return value diff --git a/hastelib/src/hastegeo/core/publishing/registry.py b/hastelib/src/hastegeo/core/publishing/registry.py new file mode 100644 index 00000000..00557617 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/registry.py @@ -0,0 +1,121 @@ +from typing import Callable, Dict, Optional + +from ..config import Config +from ..models.publishing import ArtifactKind, ProviderConfigField, ProviderInfo +from .base import PublishingProvider + + +class ProviderUnavailableError(RuntimeError): + """Raised when a provider cannot currently accept publishing work.""" + + +class PublishingProviderRegistry: + """Expose known provider capabilities and lazily resolve implementations.""" + + def __init__( + self, + config: Optional[Config] = None, + factories: Optional[ + Dict[str, Callable[[], PublishingProvider]] + ] = None, + ) -> None: + self.config = config or Config() + self.factories = factories or {} + + def _provider_infos(self) -> Dict[str, ProviderInfo]: + settings = self.config.publishing_config + publishing_enabled = settings["publishing_enabled"] + pc_enabled = settings["pc_provider_enabled"] + # The GeoCatalog URL is what makes the target configurable; the + # ingestion source is only needed for private containers (public + # containers publish without one), so it must not be required here. + pc_configured = bool(settings["pc_geocatalog_url"]) + return { + "local": ProviderInfo( + id="local", + displayName="Local (HASTE storage)", + description="Immutable copy in HASTE-managed storage", + isEnabled=publishing_enabled, + isConfigured=True, + disabledReason=( + None if publishing_enabled else "Publishing is disabled" + ), + supportedArtifactKinds=list(ArtifactKind), + ), + "planetary_computer": ProviderInfo( + id="planetary_computer", + displayName="Planetary Computer", + description="STAC discovery and vector downloads", + isEnabled=pc_enabled, + isConfigured=pc_configured, + disabledReason=self._pc_disabled_reason( + pc_enabled, pc_configured + ), + supportedArtifactKinds=[ + ArtifactKind.GPKG, + ArtifactKind.VALID_MASK, + ArtifactKind.FOOTPRINTS, + ], + requiredSupportingArtifactKinds=[ArtifactKind.VALID_MASK], + configRequirements=[ + ProviderConfigField( + key="geocatalog_url", + label="GeoCatalog URL", + required=True, + ), + ProviderConfigField( + key="ingestion_source", + label="Ingestion source", + required=True, + ), + ], + ), + } + + @staticmethod + def _pc_disabled_reason(enabled: bool, configured: bool) -> Optional[str]: + if not enabled: + return "Disabled by the operator" + if not configured: + return "Planetary Computer is not configured" + return None + + def list_infos(self) -> list[ProviderInfo]: + return list(self._provider_infos().values()) + + def get_info(self, provider_id: str) -> ProviderInfo: + try: + return self._provider_infos()[provider_id] + except KeyError as error: + raise ProviderUnavailableError( + f"Unknown publishing provider: {provider_id}" + ) from error + + def resolve(self, provider_id: str) -> PublishingProvider: + info = self.get_info(provider_id) + if not info.isEnabled or not info.isConfigured: + raise ProviderUnavailableError( + info.disabledReason or "Publishing provider is unavailable" + ) + factory = self.factories.get(provider_id) + if factory is None and provider_id == "local": + from .local_provider import LocalPublishingProvider + + def local_factory() -> PublishingProvider: + return LocalPublishingProvider(config=self.config) + + factory = local_factory + if factory is None and provider_id == "planetary_computer": + from .planetary_computer_provider import ( + PlanetaryComputerPublishingProvider, + ) + + def planetary_computer_factory() -> PublishingProvider: + return PlanetaryComputerPublishingProvider(config=self.config) + + factory = planetary_computer_factory + if factory is None: + raise ProviderUnavailableError( + f"No implementation registered for provider: {provider_id}" + ) + return factory() diff --git a/hastelib/src/hastegeo/core/publishing/repository.py b/hastelib/src/hastegeo/core/publishing/repository.py new file mode 100644 index 00000000..9927ff04 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/repository.py @@ -0,0 +1,322 @@ +from contextlib import AbstractContextManager, nullcontext +from typing import Any, Callable, List, Optional, Tuple + +from ..config import Config +from ..models.publishing import PublishedDataset, PublishStatus, PublishTarget +from ..processors.metadata import MetadataProcessor +from .lease import BlobLeaseCoordinator + +MAX_CATALOG_SCAN_RECORDS = 1000 + + +class PublishingConflictError(RuntimeError): + """Raised when an idempotency key is reused with a different request.""" + + +class PublishedDatasetsExistError(RuntimeError): + """Raised when project deletion is blocked by published datasets.""" + + +class StaleRevisionError(RuntimeError): + """Raised when a caller tries to update an outdated dataset revision.""" + + +class PublishingRepository: + """Persist independently updateable published-dataset records.""" + + def __init__( + self, + config: Optional[Config] = None, + processor_factory: Callable[ + ..., MetadataProcessor + ] = MetadataProcessor, + lease_coordinator: Optional[BlobLeaseCoordinator] = None, + ) -> None: + self.config = config or Config() + self.processor_factory = processor_factory + self.lease_coordinator = lease_coordinator + self.data_type = ( + self.config.get_metadata_types().PUBLISHED_DATASET.value + ) + + def _processor(self, project_id: str) -> MetadataProcessor: + return self.processor_factory( + data_type=self.data_type, + partition_key=project_id, + config=self.config, + ) + + def create_or_replay( + self, dataset: PublishedDataset + ) -> Tuple[PublishedDataset, bool]: + project_id = str(dataset.projectId) + dataset_id = str(dataset.datasetId) + with self.operation_lock(project_id, dataset_id): + return self.create_or_replay_locked(dataset) + + def create_or_replay_locked( + self, dataset: PublishedDataset + ) -> Tuple[PublishedDataset, bool]: + project_id = str(dataset.projectId) + dataset_id = str(dataset.datasetId) + try: + existing = self.load(project_id, dataset_id) + except FileNotFoundError: + self._processor(project_id).save( + dataset_id, dataset.model_dump(mode="json") + ) + return dataset, True + + if existing.requestFingerprint != dataset.requestFingerprint: + raise PublishingConflictError( + "requestId was already used with different publish values" + ) + return existing, False + + def load(self, project_id: str, dataset_id: str) -> PublishedDataset: + data = self._processor(project_id).load(dataset_id) + return PublishedDataset(**data) + + def list_all( + self, + project_id: Optional[str] = None, + target: Optional[PublishTarget] = None, + status: Optional[PublishStatus] = None, + ) -> List[PublishedDataset]: + if project_id: + raw_records = self._processor(project_id).load_all_from_partition() + else: + raw_records = self.processor_factory( + data_type=self.data_type, + config=self.config, + ).load_all() + + records = [PublishedDataset(**record) for record in raw_records] + if target is not None: + records = [record for record in records if record.target == target] + if status is not None: + records = [record for record in records if record.status == status] + return sorted( + records, + key=lambda record: record.publishedDate or record.createdDate, + reverse=True, + ) + + def list_for_reconciliation(self) -> List[PublishedDataset]: + if getattr(self.config, "storage_type", None) == "blob": + raw_records, _ = self.processor_factory( + data_type=self.data_type, + config=self.config, + ).load_page( + page=1, + page_size=MAX_CATALOG_SCAN_RECORDS, + max_records=MAX_CATALOG_SCAN_RECORDS, + ) + else: + raw_records = self.processor_factory( + data_type=self.data_type, + config=self.config, + ).load_bounded(max_records=MAX_CATALOG_SCAN_RECORDS) + return [PublishedDataset(**record) for record in raw_records] + + def list_page( + self, + page: int, + page_size: int, + project_id: Optional[str] = None, + target: Optional[PublishTarget] = None, + status: Optional[PublishStatus] = None, + search: str = "", + sort_key: str = "publishedDate", + sort_direction: str = "desc", + ) -> Tuple[List[PublishedDataset], int]: + if page < 1 or page_size < 1 or page_size > 100: + raise ValueError("Invalid publishing page request") + sortable_fields = { + "name", + "projectName", + "target", + "status", + "publishedByUser", + "publishedDate", + } + if sort_key not in sortable_fields or sort_direction not in { + "asc", + "desc", + }: + raise ValueError("Invalid publishing sort request") + + normalized_search = search.strip().lower() + if ( + getattr(self.config, "storage_type", None) == "blob" + and not normalized_search + and sort_key == "publishedDate" + and sort_direction == "desc" + ): + processor = self.processor_factory( + data_type=self.data_type, + config=self.config, + ) + raw_records, total_count = processor.load_page( + page=page, + page_size=page_size, + target=target.value if target else None, + status=status.value if status else None, + project_id=project_id, + max_records=MAX_CATALOG_SCAN_RECORDS, + ) + records = [PublishedDataset(**record) for record in raw_records] + return records, total_count + + if getattr(self.config, "storage_type", None) == "blob": + processor = self.processor_factory( + data_type=self.data_type, + config=self.config, + ) + raw_records, scan_count = processor.load_page( + page=1, + page_size=MAX_CATALOG_SCAN_RECORDS, + project_id=project_id, + max_records=MAX_CATALOG_SCAN_RECORDS, + ) + if scan_count > MAX_CATALOG_SCAN_RECORDS: + raise ValueError( + "Search and custom sorting are limited to 1,000 records" + ) + records = [PublishedDataset(**record) for record in raw_records] + if project_id is not None: + records = [ + record + for record in records + if str(record.projectId) == project_id + ] + if target is not None: + records = [ + record for record in records if record.target == target + ] + if status is not None: + records = [ + record for record in records if record.status == status + ] + else: + records = self.list_all( + project_id=project_id, + target=target, + status=status, + ) + if normalized_search: + records = [ + record + for record in records + if normalized_search + in " ".join( + [ + record.name, + record.description, + record.projectName, + record.imageLayerName, + record.publishedByUser, + record.target.value, + record.status.value, + ] + ).lower() + ] + + def sort_value(record: PublishedDataset) -> str: + value = getattr(record, sort_key) + if hasattr(value, "value"): + value = value.value + if sort_key == "publishedDate": + value = value or record.createdDate + return str(value or "").lower() + + records.sort( + key=sort_value, + reverse=sort_direction == "desc", + ) + total_count = len(records) + start = (page - 1) * page_size + return records[start : start + page_size], total_count + + def operation_lock( + self, project_id: str, dataset_id: str + ) -> AbstractContextManager[Any]: + return self._get_lease_coordinator().acquire(project_id, dataset_id) + + def project_lock(self, project_id: str) -> AbstractContextManager[Any]: + if not self.config.publishing_config.get("publishing_enabled", False): + return nullcontext() + return self._get_lease_coordinator().acquire( + project_id, + "project-publishing", + wait_timeout_seconds=2, + retry_interval_seconds=0.02, + ) + + def delete_project_if_unpublished( + self, project_id: str, delete_action: Callable[[], None] + ) -> None: + with self.project_lock(project_id): + if self.list_all(project_id=project_id): + raise PublishedDatasetsExistError( + "Unpublish all project datasets before deleting the project" + ) + delete_action() + + def _get_lease_coordinator(self) -> BlobLeaseCoordinator: + if self.lease_coordinator is None: + publishing = self.config.publishing_config + self.lease_coordinator = BlobLeaseCoordinator( + connection_string=publishing["lease_connection_string"], + account_url=publishing["lease_account_url"], + container_name=publishing["lease_container"], + ) + return self.lease_coordinator + + def update( + self, + dataset: PublishedDataset, + expected_revision: int, + ) -> PublishedDataset: + project_id = str(dataset.projectId) + dataset_id = str(dataset.datasetId) + with self.operation_lock(project_id, dataset_id): + return self.update_locked(dataset, expected_revision) + + def update_locked( + self, + dataset: PublishedDataset, + expected_revision: int, + ) -> PublishedDataset: + project_id = str(dataset.projectId) + dataset_id = str(dataset.datasetId) + current = self.load(project_id, dataset_id) + if current.revision != expected_revision: + raise StaleRevisionError( + f"Expected revision {expected_revision}, found {current.revision}" + ) + updated = dataset.model_copy( + update={"revision": expected_revision + 1} + ) + self._processor(project_id).save( + dataset_id, updated.model_dump(mode="json") + ) + return updated + + def delete( + self, + project_id: str, + dataset_id: str, + expected_revision: Optional[int] = None, + ) -> None: + with self.operation_lock(project_id, dataset_id): + if expected_revision is not None: + current = self.load(project_id, dataset_id) + if current.revision != expected_revision: + raise StaleRevisionError( + f"Expected revision {expected_revision}, found {current.revision}" + ) + self.delete_locked(project_id, dataset_id) + + def delete_locked(self, project_id: str, dataset_id: str) -> None: + self._processor(project_id).delete(dataset_id) diff --git a/hastelib/src/hastegeo/core/publishing/schemas/item-assets-v1.0.0.json b/hastelib/src/hastegeo/core/publishing/schemas/item-assets-v1.0.0.json new file mode 100644 index 00000000..5fc0fa42 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/schemas/item-assets-v1.0.0.json @@ -0,0 +1,35 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "$id": "https://stac-extensions.github.io/item-assets/v1.0.0/schema.json#", + "title": "Item Assets Definition Extension", + "description": "STAC Item Assets Definition Extension for STAC Collections.", + "type": "object", + "required": ["stac_extensions", "type", "item_assets"], + "properties": { + "stac_extensions": { + "type": "array", + "contains": { + "const": "https://stac-extensions.github.io/item-assets/v1.0.0/schema.json" + } + }, + "type": {"const": "Collection"}, + "item_assets": { + "type": "object", + "additionalProperties": { + "type": "object", + "minProperties": 2, + "properties": { + "href": {"title": "Disallow href", "not": {}}, + "title": {"title": "Asset title", "type": "string"}, + "description": {"title": "Asset description", "type": "string"}, + "type": {"title": "Asset type", "type": "string"}, + "roles": { + "title": "Asset roles", + "type": "array", + "items": {"type": "string"} + } + } + } + } + } +} \ No newline at end of file diff --git a/hastelib/src/hastegeo/core/publishing/schemas/projection-v2.0.0.json b/hastelib/src/hastegeo/core/publishing/schemas/projection-v2.0.0.json new file mode 100644 index 00000000..62eea5d9 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/schemas/projection-v2.0.0.json @@ -0,0 +1,113 @@ +{ + "$schema": "http://json-schema.org/draft-07/schema#", + "$id": "https://stac-extensions.github.io/projection/v2.0.0/schema.json", + "title": "Projection Extension", + "description": "STAC Projection Extension for STAC Items.", + "$comment": "This schema succeeds if the proj: fields are not used at all, please keep this in mind.", + "oneOf": [ + { + "$comment": "This is the schema for STAC Items.", + "allOf": [ + {"$ref": "#/definitions/stac_extensions"}, + { + "type": "object", + "required": ["type", "properties", "assets"], + "properties": { + "type": {"const": "Feature"}, + "properties": {"$ref": "#/definitions/fields"}, + "assets": { + "type": "object", + "additionalProperties": {"$ref": "#/definitions/fields"} + } + } + } + ] + }, + { + "$comment": "This is the schema for STAC Collections.", + "allOf": [ + { + "type": "object", + "required": ["type"], + "properties": { + "type": {"const": "Collection"}, + "assets": { + "type": "object", + "additionalProperties": {"$ref": "#/definitions/fields"} + }, + "item_assets": { + "type": "object", + "additionalProperties": {"$ref": "#/definitions/fields"} + } + } + }, + {"$ref": "#/definitions/stac_extensions"} + ] + } + ], + "definitions": { + "stac_extensions": { + "type": "object", + "required": ["stac_extensions"], + "properties": { + "stac_extensions": { + "type": "array", + "contains": { + "const": "https://stac-extensions.github.io/projection/v2.0.0/schema.json" + } + } + } + }, + "fields": { + "type": "object", + "properties": { + "proj:code": {"title": "Projection code", "type": ["string", "null"]}, + "proj:wkt2": {"title": "Coordinate Reference System in WKT2 format", "type": ["string", "null"]}, + "proj:projjson": { + "title": "Coordinate Reference System in PROJJSON format", + "oneOf": [ + {"$ref": "https://proj.org/schemas/v0.7/projjson.schema.json"}, + {"type": "null"} + ] + }, + "proj:geometry": {"$ref": "https://geojson.org/schema/Geometry.json"}, + "proj:bbox": { + "title": "Extent", + "type": "array", + "oneOf": [ + {"minItems": 4, "maxItems": 4}, + {"minItems": 6, "maxItems": 6} + ], + "items": {"type": "number"} + }, + "proj:centroid": { + "title": "Centroid", + "type": "object", + "required": ["lat", "lon"], + "properties": { + "lat": {"type": "number", "minimum": -90, "maximum": 90}, + "lon": {"type": "number", "minimum": -180, "maximum": 180} + } + }, + "proj:shape": { + "title": "Shape", + "type": "array", + "minItems": 2, + "maxItems": 2, + "items": {"type": "integer"} + }, + "proj:transform": { + "title": "Transform", + "type": "array", + "oneOf": [ + {"minItems": 6, "maxItems": 6}, + {"minItems": 9, "maxItems": 9} + ], + "items": {"type": "number"} + } + }, + "patternProperties": {"^(?!proj:)": {}}, + "additionalProperties": false + } + } +} \ No newline at end of file diff --git a/hastelib/src/hastegeo/core/publishing/source.py b/hastelib/src/hastegeo/core/publishing/source.py new file mode 100644 index 00000000..4b9fa8e9 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/source.py @@ -0,0 +1,459 @@ +from pathlib import PurePosixPath +from typing import Callable, Dict, Iterable, Optional, Set, Tuple + +from ..artifact_storage.unified_artifact_storage import UnifiedArtifactStorage +from ..config import ArtifactTypes, Config +from ..models.projects import ImageLayer, Model, Project +from ..models.publishing import ( + ArtifactBundle, + ArtifactKind, + PublishDatasetOptions, + PublishRequest, + SourceArtifact, +) +from ..processors.metadata import MetadataProcessor +from ..utils.metadata import MetadataUtils + + +class PublishingSourceNotFoundError(FileNotFoundError): + """Raised when a requested project, layer, or model does not exist.""" + + +class PublishingSourceNotEligibleError(RuntimeError): + """Raised when inference has not produced a publishable result.""" + + +class PublishingArtifactUnavailableError(RuntimeError): + """Raised when a requested or supporting artifact is unavailable.""" + + +ARTIFACT_FIELDS: Dict[ArtifactKind, Tuple[str, str, str]] = { + ArtifactKind.GPKG: ( + "model", + "gpkgUrl", + "application/geopackage+sqlite3", + ), + ArtifactKind.VALID_MASK: ( + "layer", + "validAreaMaskUrl", + "application/geo+json", + ), + ArtifactKind.FOOTPRINTS: ( + "layer", + "buildingFootprintsUrl", + "application/geopackage+sqlite3", + ), + ArtifactKind.PROCESSED_COG: ( + "layer", + "postEventProcessedImageryUrl", + "image/tiff; application=geotiff; profile=cloud-optimized", + ), +} + + +class PublishingSourceResolver: + """Resolve publishable source metadata into verified storage artifacts.""" + + def __init__( + self, + config: Optional[Config] = None, + processor_factory: Callable[ + ..., MetadataProcessor + ] = MetadataProcessor, + artifact_storage: Optional[UnifiedArtifactStorage] = None, + ) -> None: + self.config = config or Config() + self.processor_factory = processor_factory + self.artifact_storage = artifact_storage or UnifiedArtifactStorage( + storage_type=self.config.artifact_storage_type, + **self.config.artifact_storage_config, + ) + + def _load_source( + self, + project_id: str, + image_layer_id: str, + model_id: str, + ) -> Tuple[Project, ImageLayer, Model]: + metadata_types = self.config.get_metadata_types() + try: + project = Project( + **self.processor_factory( + data_type=metadata_types.PROJECT.value, + partition_key=project_id, + config=self.config, + ).load(project_id) + ) + image_layer = ImageLayer( + **self.processor_factory( + data_type=metadata_types.IMAGELAYER.value, + partition_key=project_id, + config=self.config, + ).load(image_layer_id) + ) + model = Model( + **self.processor_factory( + data_type=metadata_types.MODEL.value, + partition_key=project_id, + config=self.config, + ).load(model_id) + ) + except FileNotFoundError as error: + raise PublishingSourceNotFoundError(str(error)) from error + + if image_layer.projectId != project_id: + raise PublishingSourceNotFoundError( + "Image layer does not belong to the requested project" + ) + if ( + model.projectId != project_id + or model.imageLayerId != image_layer_id + ): + raise PublishingSourceNotFoundError( + "Model does not belong to the requested project and image layer" + ) + completed = self.config.get_status_types().COMPLETED.value + # Embedding models signal completion via `status`; trained/inference + # models via `inferenceStatus`. Gate on the field that actually applies. + if model.modelType == "embedding": + is_complete = model.status == completed + else: + is_complete = model.inferenceStatus == completed + if not is_complete: + raise PublishingSourceNotEligibleError( + "Model must be Processed before publishing" + ) + return project, image_layer, model + + def ensure_project_exists(self, project_id: str) -> None: + metadata_types = self.config.get_metadata_types() + try: + self.processor_factory( + data_type=metadata_types.PROJECT.value, + partition_key=project_id, + config=self.config, + ).load(project_id) + except FileNotFoundError as error: + raise PublishingSourceNotFoundError(str(error)) from error + + def _available_artifacts( + self, image_layer: ImageLayer, model: Model + ) -> list[SourceArtifact]: + sources = {"layer": image_layer, "model": model} + expected_paths = self._expected_artifact_paths(image_layer, model) + artifacts = [] + for kind, ( + source_name, + field_name, + media_type, + ) in ARTIFACT_FIELDS.items(): + location = getattr(sources[source_name], field_name) + if not location: + continue + try: + source_path = self.artifact_storage.resolve_artifact_path( + location + ) + if source_path not in expected_paths.get(kind, set()): + continue + if not self.artifact_storage.artifact_exists(source_path): + continue + size_bytes = self.artifact_storage.get_artifact_size( + source_path + ) + source_etag = self.artifact_storage.get_artifact_etag( + source_path + ) + except (FileNotFoundError, ValueError): + continue + artifacts.append( + SourceArtifact( + kind=kind, + sourcePath=source_path, + mediaType=media_type, + sizeBytes=size_bytes, + sourceEtag=source_etag, + ) + ) + return artifacts + + @staticmethod + def _completed_preprocess_task(image_layer: ImageLayer) -> Optional[str]: + job = image_layer.preprocessJob + if ( + image_layer.status != "Processed" + or job is None + or job.status != "Processed" + or job.projectId != image_layer.projectId + or job.imageLayerId != image_layer.imageLayerId + or not job.taskId + ): + return None + return str(job.taskId) + + @classmethod + def _expected_layer_output_path( + cls, + image_layer: ImageLayer, + artifact_type, + extension: str, + ) -> Optional[str]: + task_id = cls._completed_preprocess_task(image_layer) + if task_id is None: + return None + project_id = str(image_layer.projectId) + file_name = ( + artifact_type.value.substitute( + projectId=project_id, + imageLayerId=str(image_layer.imageLayerId), + ) + + extension + ) + return str( + PurePosixPath( + MetadataUtils.hash_string(project_id), + task_id, + file_name, + ) + ) + + def resolve_layer_output( + self, + image_layer: ImageLayer, + location: Optional[str], + artifact_type, + extension: str, + ) -> str: + expected_path = self._expected_layer_output_path( + image_layer, artifact_type, extension + ) + if expected_path is None: + raise PublishingSourceNotEligibleError( + "Image layer preprocessing must be Processed" + ) + if not location: + raise PublishingArtifactUnavailableError( + "Required image layer output is unavailable" + ) + source_path = self.artifact_storage.resolve_artifact_path(location) + if source_path != expected_path: + raise PublishingArtifactUnavailableError( + "Image layer output does not match its preprocessing job" + ) + if not self.artifact_storage.artifact_exists(source_path): + raise PublishingArtifactUnavailableError( + "Required image layer output is unavailable" + ) + return source_path + + def resolve_layer_artifact( + self, image_layer: ImageLayer, kind: ArtifactKind + ) -> SourceArtifact: + if kind not in { + ArtifactKind.VALID_MASK, + ArtifactKind.FOOTPRINTS, + ArtifactKind.PROCESSED_COG, + }: + raise ValueError("Artifact kind is not owned by an image layer") + _, field_name, media_type = ARTIFACT_FIELDS[kind] + artifact_type_and_extension = { + ArtifactKind.VALID_MASK: ( + ArtifactTypes.VALID_AREA_MASK, + ".geojson", + ), + ArtifactKind.FOOTPRINTS: ( + ArtifactTypes.BUILDING_FOOTPRINTS, + ".gpkg", + ), + ArtifactKind.PROCESSED_COG: ( + ArtifactTypes.POST_EVENT_PROCESSED_COG, + ".tif", + ), + } + artifact_type, extension = artifact_type_and_extension[kind] + source_path = self.resolve_layer_output( + image_layer, + getattr(image_layer, field_name), + artifact_type, + extension, + ) + return SourceArtifact( + kind=kind, + sourcePath=source_path, + mediaType=media_type, + sizeBytes=self.artifact_storage.get_artifact_size(source_path), + sourceEtag=self.artifact_storage.get_artifact_etag(source_path), + ) + + @staticmethod + def _expected_artifact_paths( + image_layer: ImageLayer, model: Model + ) -> Dict[ArtifactKind, Set[str]]: + project_id = str(image_layer.projectId) + project_prefix = MetadataUtils.hash_string(project_id) + image_layer_id = str(image_layer.imageLayerId) + expected: Dict[ArtifactKind, Set[str]] = { + kind: set() for kind in ArtifactKind + } + + preprocess_task_id = ( + PublishingSourceResolver._completed_preprocess_task(image_layer) + ) + if preprocess_task_id: + task_prefix = PurePosixPath(project_prefix, preprocess_task_id) + layer_artifacts = { + ArtifactKind.VALID_MASK: ( + ArtifactTypes.VALID_AREA_MASK.value.substitute( + projectId=project_id, + imageLayerId=image_layer_id, + ) + + ".geojson" + ), + ArtifactKind.FOOTPRINTS: ( + ArtifactTypes.BUILDING_FOOTPRINTS.value.substitute( + projectId=project_id, + imageLayerId=image_layer_id, + ) + + ".gpkg" + ), + ArtifactKind.PROCESSED_COG: ( + ArtifactTypes.POST_EVENT_PROCESSED_COG.value.substitute( + projectId=project_id, + imageLayerId=image_layer_id, + ) + + ".tif" + ), + } + for kind, file_name in layer_artifacts.items(): + expected[kind].add(str(task_prefix / file_name)) + + if ( + model.modelType == "embedding" + and model.status == "Processed" + and model.embeddingJob is not None + and model.embeddingJob.status == "Processed" + and model.embeddingJob.projectId == project_id + and str(model.embeddingJob.modelId) == str(model.modelId) + and model.embeddingJob.taskId + ): + file_name = ( + ArtifactTypes.BUILDING_PREDICTIONS_GPKG.value.substitute( + modelName=str(model.modelId) + ) + + ".gpkg" + ) + expected[ArtifactKind.GPKG].add( + str(PurePosixPath(project_prefix, file_name)) + ) + elif model.currentInferenceTaskId: + current_jobs = [ + job + for job in model.inferenceJobs or [] + if job.taskId == model.currentInferenceTaskId + and job.projectId == project_id + and str(job.modelId) == str(model.modelId) + and job.status == "Processed" + ] + expected_output_path = str( + PurePosixPath(project_prefix, model.currentInferenceTaskId) + ) + if ( + len(current_jobs) == 1 + and model.inferenceOutputPath == expected_output_path + ): + file_name = ( + ArtifactTypes.INFERENCE_GPKG.value.substitute( + modelName=str(model.name) + ) + + ".gpkg" + ) + if PurePosixPath(file_name).name == file_name: + expected[ArtifactKind.GPKG].update( + { + str( + PurePosixPath(expected_output_path, file_name) + ), + str( + PurePosixPath( + expected_output_path, + "inference", + file_name, + ) + ), + } + ) + + return expected + + def resolve_options( + self, + project_id: str, + image_layer_id: str, + model_id: str, + ) -> PublishDatasetOptions: + project, image_layer, model = self._load_source( + project_id, image_layer_id, model_id + ) + available_artifacts = self._available_artifacts(image_layer, model) + if not available_artifacts: + raise PublishingSourceNotEligibleError( + "Model has no available publishable artifacts" + ) + project_name = project.name or project_id + image_layer_name = image_layer.name or image_layer_id + return PublishDatasetOptions( + projectId=project_id, + projectName=project_name, + imageLayerId=image_layer_id, + imageLayerName=image_layer_name, + modelId=model_id, + modelName=model.name or model_id, + defaultName=f"{project_name} – {image_layer_name}", + availableArtifacts=available_artifacts, + ) + + def resolve_bundle( + self, + request: PublishRequest, + supporting_kinds: Iterable[ArtifactKind] = (), + options: Optional[PublishDatasetOptions] = None, + ) -> ArtifactBundle: + options = options or self.resolve_options( + str(request.projectId), request.imageLayerId, request.modelId + ) + available = { + artifact.kind: artifact for artifact in options.availableArtifacts + } + selected = [] + for kind in request.artifacts: + if kind not in available: + raise PublishingArtifactUnavailableError( + f"Requested artifact is unavailable: {kind.value}" + ) + selected.append(available[kind]) + + selected_kinds = set(request.artifacts) + supporting = [] + for kind in sorted(set(supporting_kinds), key=lambda item: item.value): + if kind not in available: + raise PublishingArtifactUnavailableError( + f"Required supporting artifact is unavailable: {kind.value}" + ) + if kind not in selected_kinds: + supporting.append(available[kind]) + _, image_layer, _ = self._load_source( + str(request.projectId), request.imageLayerId, request.modelId + ) + return ArtifactBundle( + selectedArtifacts=selected, + supportingArtifacts=supporting, + thumbnailUrl=self._first_preview_url(image_layer), + ) + + @staticmethod + def _first_preview_url(image_layer: ImageLayer) -> Optional[str]: + urls = getattr(image_layer, "postEventPreviewUrls", None) or [] + for url in urls: + if isinstance(url, str) and url.strip(): + return url.strip() + return None diff --git a/hastelib/src/hastegeo/core/publishing/stac.py b/hastelib/src/hastegeo/core/publishing/stac.py new file mode 100644 index 00000000..858ad062 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/stac.py @@ -0,0 +1,765 @@ +import json +import math +import re +import unicodedata +from dataclasses import dataclass +from datetime import datetime, timezone +from importlib import resources +from typing import Any, Dict, Mapping, Optional, Sequence, Tuple +from urllib.parse import urlparse + +import geopandas as gpd +from pyproj import CRS +from pyproj.exceptions import CRSError +from shapely.geometry import mapping, shape +from shapely.ops import unary_union + +from ..models.publishing import ( + ArtifactBundle, + ArtifactKind, + PublishedDataset, + SourceArtifact, +) + +COLLECTION_ID_MAX_LENGTH = 242 +ITEM_ID_MAX_LENGTH = 149 +ASSET_KEY_MAX_LENGTH = 255 +STAC_VERSION = "1.0.0" +ITEM_ASSETS_EXTENSION = ( + "https://stac-extensions.github.io/item-assets/v1.0.0/schema.json" +) +PROJECTION_EXTENSION = ( + "https://stac-extensions.github.io/projection/v2.0.0/schema.json" +) +EXTENSION_SCHEMA_FILES = { + ITEM_ASSETS_EXTENSION: "item-assets-v1.0.0.json", + PROJECTION_EXTENSION: "projection-v2.0.0.json", +} + +ASSET_KEYS = { + ArtifactKind.GPKG: "damage", + ArtifactKind.VALID_MASK: "aoi", + ArtifactKind.FOOTPRINTS: "footprints", +} +ASSET_ROLES = { + ArtifactKind.GPKG: ["data"], + ArtifactKind.VALID_MASK: ["metadata"], + ArtifactKind.FOOTPRINTS: ["data"], +} +ASSET_TITLES = { + ArtifactKind.GPKG: "Damage assessment", + ArtifactKind.VALID_MASK: "Valid assessment area", + ArtifactKind.FOOTPRINTS: "Building footprints", +} + +# Compact per-dataset summaries persisted on the (project-level) collection so +# its description can be rendered as a rolling summary of every dataset it +# holds. Read back from the existing collection on the next publish/unpublish. +COLLECTION_DATASETS_FIELD = "ai4g:datasets" + +# Best-effort preview asset (post-event imagery) so the GeoCatalog shows a +# thumbnail for the otherwise download-only vector item. +THUMBNAIL_ASSET_KEY = "thumbnail" +THUMBNAIL_ROLES = ["thumbnail"] +DEFAULT_THUMBNAIL_MEDIA_TYPE = "image/png" + + +@dataclass(frozen=True) +class ValidMaskGeometry: + geometry: Dict[str, Any] + bbox: Sequence[float] + area_square_kilometers: float + + +@dataclass(frozen=True) +class StacObjects: + collection: Any + item: Any + + +@dataclass(frozen=True) +class StacDocuments: + collection: Dict[str, Any] + item: Dict[str, Any] + + +def _sanitize_identifier( + value: str, + *, + punctuation: str, + max_length: int, +) -> str: + normalized = ( + unicodedata.normalize("NFKD", value) + .encode("ascii", "ignore") + .decode("ascii") + .strip() + ) + normalized = re.sub(r"\s+", "-", normalized) + allowed = re.escape(punctuation) + normalized = re.sub( + rf"[^A-Za-z0-9{allowed}]", + "-", + normalized, + ) + normalized = re.sub(r"-{2,}", "-", normalized) + normalized = normalized[:max_length] + if not re.search(r"[A-Za-z0-9]", normalized): + raise ValueError("STAC identifier must contain a letter or digit") + return normalized + + +def sanitize_collection_id(value: str) -> str: + return _sanitize_identifier( + value, + punctuation="-_.", + max_length=COLLECTION_ID_MAX_LENGTH, + ) + + +def sanitize_item_id(value: str) -> str: + return _sanitize_identifier( + value, + punctuation="-_+,.()", + max_length=ITEM_ID_MAX_LENGTH, + ) + + +def sanitize_asset_key(value: str) -> str: + return _sanitize_identifier( + value, + punctuation="-_+,.()", + max_length=ASSET_KEY_MAX_LENGTH, + ) + + +def build_collection_id( + dataset: PublishedDataset, + collection_prefix: str = "haste-", +) -> str: + project_id = str(dataset.projectId).lower() + prefix = _sanitize_identifier( + collection_prefix, + punctuation="-_.", + max_length=COLLECTION_ID_MAX_LENGTH - len(project_id), + ) + return sanitize_collection_id(f"{prefix}{project_id}") + + +def build_item_id(dataset: PublishedDataset) -> str: + return sanitize_item_id(str(dataset.datasetId)) + + +def resolve_valid_mask_geometry( + valid_mask: Mapping[str, Any], + source_crs: str = "EPSG:4326", +) -> ValidMaskGeometry: + if valid_mask.get("type") != "FeatureCollection": + raise ValueError("Valid-area mask must be a GeoJSON FeatureCollection") + features = valid_mask.get("features") + if not isinstance(features, list) or not features: + raise ValueError("Valid-area mask must contain at least one feature") + + geometries = [] + for feature in features: + if not isinstance(feature, Mapping) or not feature.get("geometry"): + raise ValueError("Valid-area mask feature is missing geometry") + geometry = shape(feature["geometry"]) + if geometry.is_empty: + raise ValueError("Valid-area mask geometry must not be empty") + geometries.append(geometry) + + try: + source_geometries = gpd.GeoSeries(geometries, crs=source_crs) + wgs84_geometries = source_geometries.to_crs("EPSG:4326") + except (CRSError, TypeError, ValueError) as error: + raise ValueError("Valid-area mask CRS is invalid") from error + + combined = unary_union(wgs84_geometries.tolist()) + if combined.geom_type not in {"Polygon", "MultiPolygon"}: + raise ValueError("Valid-area mask must contain polygon geometry") + if not combined.is_valid: + raise ValueError("Valid-area mask geometry is invalid") + + bounds = list(combined.bounds) + if not all(math.isfinite(value) for value in bounds): + raise ValueError("Valid-area mask bounds must be finite") + if not ( + -180 <= bounds[0] <= bounds[2] <= 180 + and -90 <= bounds[1] <= bounds[3] <= 90 + ): + raise ValueError("Valid-area mask bounds must be within EPSG:4326") + + projected = gpd.GeoSeries([combined], crs="EPSG:4326").to_crs("EPSG:6933") + return ValidMaskGeometry( + geometry=mapping(combined), + bbox=bounds, + area_square_kilometers=float(projected.area.iloc[0] / 1_000_000), + ) + + +def build_vector_item( + dataset: PublishedDataset, + source: ArtifactBundle, + valid_mask: Mapping[str, Any], + asset_hrefs: Mapping[str, str], + projection_codes: Mapping[str, str], + collection_href: str, + *, + valid_mask_crs: str = "EPSG:4326", + collection_prefix: str = "haste-", + license_id: str = "proprietary", + thumbnail_href: Optional[str] = None, + thumbnail_media_type: str = DEFAULT_THUMBNAIL_MEDIA_TYPE, +) -> Any: + pystac = _load_pystac() + if not source.selectedArtifacts: + raise ValueError("Select at least one artifact to publish") + mask_geometry = resolve_valid_mask_geometry(valid_mask, valid_mask_crs) + collection_id = build_collection_id(dataset, collection_prefix) + properties = _build_item_properties( + dataset, + mask_geometry.area_square_kilometers, + license_id, + ) + item = pystac.Item( + id=build_item_id(dataset), + geometry=mask_geometry.geometry, + bbox=list(mask_geometry.bbox), + datetime=_parse_timestamp(dataset.createdDate), + properties=properties, + stac_extensions=[PROJECTION_EXTENSION], + collection=collection_id, + ) + item.add_link( + pystac.Link( + rel=pystac.RelType.COLLECTION, + target=_require_https_url( + collection_href, "GeoCatalog collection" + ), + media_type=pystac.MediaType.JSON, + ) + ) + + seen_keys = set() + selected_projections = {} + for artifact in source.selectedArtifacts: + if artifact.kind not in ASSET_KEYS: + raise ValueError( + f"Unsupported Planetary Computer artifact: " + f"{artifact.kind.value}" + ) + asset_key = sanitize_asset_key(ASSET_KEYS[artifact.kind]) + if asset_key in seen_keys: + raise ValueError(f"Duplicate STAC asset key: {asset_key}") + seen_keys.add(asset_key) + href = _require_https_href( + asset_hrefs.get(artifact.sourcePath), artifact + ) + projection_code = _projection_code( + artifact, + projection_codes, + valid_mask_crs, + ) + selected_projections[artifact.kind] = projection_code + item.add_asset( + asset_key, + pystac.Asset( + href=href, + media_type=artifact.mediaType, + title=ASSET_TITLES[artifact.kind], + roles=ASSET_ROLES[artifact.kind], + extra_fields={"proj:code": projection_code}, + ), + ) + + for kind in ( + ArtifactKind.GPKG, + ArtifactKind.FOOTPRINTS, + ArtifactKind.VALID_MASK, + ): + if kind in selected_projections: + item.properties["proj:code"] = selected_projections[kind] + break + + if thumbnail_href: + item.add_asset( + THUMBNAIL_ASSET_KEY, + pystac.Asset( + href=_require_https_url(thumbnail_href, "thumbnail"), + media_type=thumbnail_media_type, + title="Preview", + roles=list(THUMBNAIL_ROLES), + ), + ) + return item + + +def _collection_dataset_entry( + dataset: PublishedDataset, item: Any +) -> Dict[str, Any]: + """Compact summary of one dataset for the collection's rolling description.""" + properties = getattr(item, "properties", None) or {} + entry: Dict[str, Any] = { + "id": str(dataset.datasetId), + "name": dataset.name, + } + for source_key, target_key in ( + ("ai4g:buildings_damaged", "buildings_damaged"), + ("ai4g:buildings_total", "buildings_total"), + ("ai4g:aoi_area_km2", "area_km2"), + ): + value = properties.get(source_key) + if value is not None: + entry[target_key] = value + return entry + + +def merge_collection_datasets( + existing_collection: Optional[Mapping[str, Any]], + entry: Mapping[str, Any], +) -> list: + """Upsert ``entry`` into the datasets persisted on the existing collection.""" + existing: list = [] + if existing_collection: + stored = existing_collection.get(COLLECTION_DATASETS_FIELD) + if isinstance(stored, list): + existing = [ + dict(item) + for item in stored + if isinstance(item, Mapping) and item.get("id") + ] + merged = [item for item in existing if item.get("id") != entry.get("id")] + merged.append(dict(entry)) + merged.sort( + key=lambda item: (str(item.get("name") or ""), str(item.get("id"))) + ) + return merged + + +def _format_count(value: Any) -> Optional[str]: + try: + return f"{int(value):,}" + except (TypeError, ValueError): + return None + + +def render_collection_description( + dataset: PublishedDataset, entries: Sequence[Mapping[str, Any]] +) -> str: + """Render a rolling summary of every dataset held by the collection.""" + project = dataset.projectName or str(dataset.projectId) + count = len(entries) + noun = "dataset" if count == 1 else "datasets" + lines = [ + f"HASTE disaster assessment for {project}. " + f"This collection contains {count} published {noun}." + ] + for entry in entries: + name = entry.get("name") or entry.get("id") + details = [] + damaged = _format_count(entry.get("buildings_damaged")) + total = _format_count(entry.get("buildings_total")) + area = entry.get("area_km2") + if damaged is not None and total is not None: + details.append(f"{damaged} of {total} buildings assessed as damaged") + elif damaged is not None: + details.append(f"{damaged} buildings assessed as damaged") + if area is not None: + try: + details.append(f"{float(area):.1f} km² assessed") + except (TypeError, ValueError): + pass + suffix = f" — {'; '.join(details)}" if details else "" + lines.append(f"- {name}{suffix}.") + return "\n".join(lines) + + +def rebuild_collection_after_removal( + existing_collection: Mapping[str, Any], + dataset: PublishedDataset, +) -> Dict[str, Any]: + """Drop ``dataset`` from the collection's rolling summary (for unpublish). + + Returns a copy of the existing collection document with this dataset removed + from ``ai4g:datasets`` and the description re-rendered from what remains. + """ + updated = dict(existing_collection) + stored = updated.get(COLLECTION_DATASETS_FIELD) + entries: list = [] + if isinstance(stored, list): + entries = [ + dict(entry) + for entry in stored + if isinstance(entry, Mapping) + and entry.get("id") != str(dataset.datasetId) + ] + updated[COLLECTION_DATASETS_FIELD] = entries + updated["description"] = render_collection_description(dataset, entries) + return updated + + +def build_collection( + dataset: PublishedDataset, + item: Any, + collection_href: str, + *, + existing_collection: Optional[Mapping[str, Any]] = None, + collection_prefix: str = "haste-", + license_id: str = "proprietary", + thumbnail_href: Optional[str] = None, + thumbnail_media_type: str = DEFAULT_THUMBNAIL_MEDIA_TYPE, +) -> Any: + pystac = _load_pystac() + collection_id = build_collection_id(dataset, collection_prefix) + spatial_bbox, temporal_interval = _merge_collection_extent( + collection_id, + item, + existing_collection, + ) + datasets = merge_collection_datasets( + existing_collection, + _collection_dataset_entry(dataset, item), + ) + collection = pystac.Collection( + id=collection_id, + title=dataset.projectName or collection_id, + description=render_collection_description(dataset, datasets), + license=license_id, + keywords=[ + "HASTE", + "disaster assessment", + "building damage", + ], + providers=[ + pystac.Provider( + name="Microsoft AI for Good Lab", + roles=[pystac.ProviderRole.PRODUCER], + ) + ], + extent=pystac.Extent( + spatial=pystac.SpatialExtent([spatial_bbox]), + temporal=pystac.TemporalExtent([temporal_interval]), + ), + summaries=pystac.Summaries( + {"ai4g:project_id": [str(dataset.projectId)]} + ), + stac_extensions=[ITEM_ASSETS_EXTENSION], + extra_fields={ + "item_assets": _collection_item_assets(), + COLLECTION_DATASETS_FIELD: datasets, + }, + ) + collection.add_link( + pystac.Link( + rel=pystac.RelType.SELF, + target=_require_https_url( + collection_href, "GeoCatalog collection" + ), + media_type=pystac.MediaType.JSON, + ) + ) + if thumbnail_href: + collection.add_asset( + THUMBNAIL_ASSET_KEY, + pystac.Asset( + href=_require_https_url(thumbnail_href, "thumbnail"), + media_type=thumbnail_media_type, + title="Preview", + roles=list(THUMBNAIL_ROLES), + ), + ) + return collection + + +def build_stac_objects( + dataset: PublishedDataset, + source: ArtifactBundle, + valid_mask: Mapping[str, Any], + asset_hrefs: Mapping[str, str], + projection_codes: Mapping[str, str], + collection_href: str, + *, + valid_mask_crs: str = "EPSG:4326", + existing_collection: Optional[Mapping[str, Any]] = None, + collection_prefix: str = "haste-", + license_id: str = "proprietary", + thumbnail_href: Optional[str] = None, + thumbnail_media_type: str = DEFAULT_THUMBNAIL_MEDIA_TYPE, +) -> StacObjects: + item = build_vector_item( + dataset, + source, + valid_mask, + asset_hrefs, + projection_codes, + collection_href, + valid_mask_crs=valid_mask_crs, + collection_prefix=collection_prefix, + license_id=license_id, + thumbnail_href=thumbnail_href, + thumbnail_media_type=thumbnail_media_type, + ) + collection = build_collection( + dataset, + item, + collection_href, + existing_collection=existing_collection, + collection_prefix=collection_prefix, + license_id=license_id, + thumbnail_href=thumbnail_href, + thumbnail_media_type=thumbnail_media_type, + ) + return StacObjects(collection=collection, item=item) + + +def validate_stac_objects(objects: StacObjects, validator: Any = None) -> None: + """Validate the exact serialized STAC 1.0 documents.""" + pystac = _load_pystac() + validator = validator or offline_stac_validator() + documents = serialize_stac_objects(objects) + pystac.validation.validate_dict( + documents.collection, + stac_object_type=pystac.STACObjectType.COLLECTION, + stac_version=STAC_VERSION, + extensions=documents.collection.get("stac_extensions", []), + validator=validator, + ) + pystac.validation.validate_dict( + documents.item, + stac_object_type=pystac.STACObjectType.ITEM, + stac_version=STAC_VERSION, + extensions=documents.item.get("stac_extensions", []), + validator=validator, + ) + + +def offline_stac_validator() -> Any: + pystac = _load_pystac() + validator = pystac.validation.JsonSchemaSTACValidator() + schema_directory = resources.files(__package__).joinpath("schemas") + for uri, file_name in EXTENSION_SCHEMA_FILES.items(): + schema = json.loads( + schema_directory.joinpath(file_name).read_text(encoding="utf-8") + ) + validator.schema_cache[uri] = schema + return validator + + +def serialize_stac_objects(objects: StacObjects) -> StacDocuments: + collection = objects.collection.to_dict() + item = objects.item.to_dict() + if ( + collection.get("stac_version") != STAC_VERSION + or item.get("stac_version") != STAC_VERSION + ): + raise RuntimeError( + f"Planetary Computer publishing requires STAC {STAC_VERSION}" + ) + return StacDocuments(collection=collection, item=item) + + +def _merge_collection_extent( + collection_id: str, + item: Any, + existing_collection: Optional[Mapping[str, Any]], +) -> Tuple[list[float], list[datetime]]: + bboxes = [list(item.bbox)] + starts = [item.datetime] + ends = [item.datetime] + if existing_collection is not None: + if existing_collection.get("id") != collection_id: + raise ValueError("Existing STAC collection ID does not match") + summaries = existing_collection.get("summaries") or {} + project_ids = summaries.get("ai4g:project_id") or [] + if project_ids != [str(item.properties["ai4g:project_id"])]: + raise ValueError( + "Existing STAC collection project provenance does not match" + ) + extent = existing_collection.get("extent") or {} + for bbox in (extent.get("spatial") or {}).get("bbox") or []: + if not isinstance(bbox, list) or len(bbox) != 4: + raise ValueError("Existing STAC collection bbox is invalid") + normalized_bbox = [float(value) for value in bbox] + if not all(math.isfinite(value) for value in normalized_bbox): + raise ValueError("Existing STAC collection bbox is invalid") + if not ( + -180 <= normalized_bbox[0] <= normalized_bbox[2] <= 180 + and -90 <= normalized_bbox[1] <= normalized_bbox[3] <= 90 + ): + raise ValueError("Existing STAC collection bbox is invalid") + bboxes.append(normalized_bbox) + for interval in (extent.get("temporal") or {}).get("interval") or []: + if not isinstance(interval, list) or len(interval) != 2: + raise ValueError( + "Existing STAC collection interval is invalid" + ) + start = ( + _parse_timestamp(interval[0]) + if interval[0] is not None + else None + ) + end = ( + _parse_timestamp(interval[1]) + if interval[1] is not None + else None + ) + if start is not None and end is not None and start > end: + raise ValueError( + "Existing STAC collection interval is invalid" + ) + if start is not None: + starts.append(start) + if end is not None: + ends.append(end) + + spatial_bbox = [ + min(bbox[0] for bbox in bboxes), + min(bbox[1] for bbox in bboxes), + max(bbox[2] for bbox in bboxes), + max(bbox[3] for bbox in bboxes), + ] + return spatial_bbox, [min(starts), max(ends)] + + +def _load_pystac() -> Any: + try: + import pystac + except ImportError as error: + raise RuntimeError( + "Planetary Computer publishing requires the " + "hastegeo[planetary-computer] extra" + ) from error + return pystac + + +def _parse_timestamp(value: str) -> datetime: + normalized = value.replace("Z", "+00:00") + timestamp = datetime.fromisoformat(normalized) + if timestamp.tzinfo is None: + raise ValueError("STAC datetime must include a UTC offset") + return timestamp.astimezone(timezone.utc) + + +def _require_https_href(value: str | None, artifact: SourceArtifact) -> str: + if not value: + raise ValueError( + f"Missing STAC HREF for {artifact.kind.value} artifact" + ) + return _require_https_url(value, "Planetary Computer asset HREF") + + +def _require_https_url(value: str, label: str) -> str: + parsed = urlparse(value) + if parsed.scheme != "https" or not parsed.netloc: + raise ValueError(f"{label} must use HTTPS") + return value + + +def _projection_code( + artifact: SourceArtifact, + projection_codes: Mapping[str, str], + valid_mask_crs: str, +) -> str: + if artifact.kind == ArtifactKind.VALID_MASK: + projection_code = valid_mask_crs + else: + projection_code = projection_codes.get(artifact.sourcePath) + if not projection_code: + raise ValueError( + f"Missing projection code for {artifact.kind.value} artifact" + ) + if not re.fullmatch(r"EPSG:[1-9][0-9]*", projection_code): + raise ValueError("Projection code must use the EPSG: form") + try: + CRS.from_epsg(int(projection_code.split(":", 1)[1])) + except CRSError as error: + raise ValueError("Projection code is not a known EPSG CRS") from error + return projection_code + + +def _build_item_properties( + dataset: PublishedDataset, + area_square_kilometers: float, + license_id: str, +) -> Dict[str, Any]: + summary = dataset.assessmentSummary + predictions = summary.get("predictions") or {} + metrics = summary.get("metrics") or {} + population = summary.get("populationEstimate") or {} + properties = { + "title": dataset.name, + "description": dataset.description, + "license": license_id, + "ai4g:project_id": str(dataset.projectId), + "ai4g:image_layer_id": dataset.imageLayerId, + "ai4g:model_id": dataset.modelId, + "ai4g:aoi_area_km2": round(area_square_kilometers, 6), + "ai4g:buildings_total": _first_present( + predictions, "total", fallback=summary.get("buildingsTotal") + ), + "ai4g:buildings_cloud": _first_present( + predictions, "cloudy", fallback=summary.get("buildingsCloud") + ), + "ai4g:buildings_clear": _first_present( + predictions, + "knownNonCloudy", + fallback=summary.get("buildingsClear"), + ), + "ai4g:buildings_damaged": _first_present( + predictions, + "predictedDamaged", + fallback=summary.get("predictedDamaged"), + ), + "ai4g:damaged_pct_of_clear": _first_present( + predictions, + "predictedDamagedPctOfKnown", + fallback=summary.get("damagedPctOfClear"), + ), + "ai4g:validation_precision": _first_present( + metrics, "precision", fallback=summary.get("precision") + ), + "ai4g:validation_recall": _first_present( + metrics, "recall", fallback=summary.get("recall") + ), + "ai4g:validation_extrapolated_damaged": _first_present( + population, + "estimatedDamaged", + fallback=summary.get("estimatedDamaged"), + ), + "ai4g:validation_ci_lower": _first_present( + population, "ciLower", fallback=summary.get("ciLower") + ), + "ai4g:validation_ci_upper": _first_present( + population, "ciUpper", fallback=summary.get("ciUpper") + ), + } + return { + key: value for key, value in properties.items() if value is not None + } + + +def _first_present( + values: Mapping[str, Any], key: str, *, fallback: Any = None +) -> Any: + return values[key] if key in values else fallback + + +def _collection_item_assets() -> Dict[str, Dict[str, Any]]: + return { + "damage": { + "title": ASSET_TITLES[ArtifactKind.GPKG], + "type": "application/geopackage+sqlite3", + "roles": ASSET_ROLES[ArtifactKind.GPKG], + }, + "aoi": { + "title": ASSET_TITLES[ArtifactKind.VALID_MASK], + "type": "application/geo+json", + "roles": ASSET_ROLES[ArtifactKind.VALID_MASK], + }, + "footprints": { + "title": ASSET_TITLES[ArtifactKind.FOOTPRINTS], + "type": "application/geopackage+sqlite3", + "roles": ASSET_ROLES[ArtifactKind.FOOTPRINTS], + }, + } diff --git a/hastelib/src/hastegeo/core/utils/blob.py b/hastelib/src/hastegeo/core/utils/blob.py index c000946b..b3522dc1 100644 --- a/hastelib/src/hastegeo/core/utils/blob.py +++ b/hastelib/src/hastegeo/core/utils/blob.py @@ -60,7 +60,11 @@ def split_blob_url(url: str) -> Tuple[str, str]: return parts[1], "/".join(parts[2:]) -async def download_blob_to_tempfile(url: str, suffix: str = "") -> str: +async def download_blob_to_tempfile( + url: str, + suffix: str = "", + max_bytes: Optional[int] = None, +) -> str: """Download the blob at ``url`` to a NamedTemporaryFile and return the path. Routes the download through the function-app-internal @@ -77,15 +81,45 @@ async def download_blob_to_tempfile(url: str, suffix: str = "") -> str: conn_str = os.environ.get("BLOB_CONNECTION_STRING", "") container_name, blob_name = split_blob_url(url) bsc = BlobServiceClient.from_connection_string(conn_str) - blob_bytes = await asyncio.to_thread( - lambda: bsc.get_container_client(container_name) - .get_blob_client(blob_name) - .download_blob() - .readall() - ) - with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp: - tmp.write(blob_bytes) - return tmp.name + if max_bytes is not None and max_bytes < 1: + raise ValueError("max_bytes must be positive") + + def download() -> str: + blob_client = bsc.get_container_client(container_name).get_blob_client( + blob_name + ) + if max_bytes is not None: + size = blob_client.get_blob_properties().size + if size > max_bytes: + raise ValueError("Blob exceeds the allowed download size") + + temp_path = None + try: + with tempfile.NamedTemporaryFile( + suffix=suffix, delete=False + ) as temp_file: + temp_path = temp_file.name + downloaded_bytes = 0 + for chunk in blob_client.download_blob().chunks(): + downloaded_bytes += len(chunk) + if ( + max_bytes is not None + and downloaded_bytes > max_bytes + ): + raise ValueError( + "Blob exceeds the allowed download size" + ) + temp_file.write(chunk) + return temp_path + except Exception: + if temp_path: + try: + os.unlink(temp_path) + except OSError: + pass + raise + + return await asyncio.to_thread(download) class BlobRange(NamedTuple): diff --git a/hastelib/tests/core/publishing/test_lease.py b/hastelib/tests/core/publishing/test_lease.py new file mode 100644 index 00000000..f840b32d --- /dev/null +++ b/hastelib/tests/core/publishing/test_lease.py @@ -0,0 +1,182 @@ +import threading +import unittest +from unittest.mock import patch + +from azure.core.exceptions import HttpResponseError +from hastegeo.core.publishing.lease import ( + BlobLeaseCoordinator, + LeaseRenewalError, + LeaseUnavailableError, +) + + +class FakeLease: + def __init__(self, fail_renewal: bool = False) -> None: + self.fail_renewal = fail_renewal + self.renewed = threading.Event() + self.released = False + + def renew(self) -> None: + self.renewed.set() + if self.fail_renewal: + raise RuntimeError("renewal failed") + + def release(self) -> None: + self.released = True + + +class FakeBlobClient: + def __init__(self, lease=None, acquire_error=None) -> None: + self.lease = lease or FakeLease() + self.acquire_error = acquire_error + + def upload_blob(self, data, overwrite=False) -> None: + return None + + def acquire_lease(self, lease_duration: int): + if self.acquire_error is not None: + raise self.acquire_error + return self.lease + + +class RetryOnceBlobClient(FakeBlobClient): + def __init__(self, lease: FakeLease) -> None: + super().__init__(lease=lease) + self.attempts = 0 + + def acquire_lease(self, lease_duration: int): + self.attempts += 1 + if self.attempts == 1: + conflict = HttpResponseError("conflict") + conflict.status_code = 409 + raise conflict + return self.lease + + +class FakeContainerClient: + def __init__(self, blob_client: FakeBlobClient) -> None: + self.blob_client = blob_client + + def get_blob_client(self, name: str) -> FakeBlobClient: + return self.blob_client + + +class FakeBlobServiceClient: + def __init__(self, container_client: FakeContainerClient) -> None: + self.container_client = container_client + + def create_container(self, name: str) -> FakeContainerClient: + return self.container_client + + +class ExistingContainerService(FakeBlobServiceClient): + def create_container(self, name: str) -> FakeContainerClient: + from azure.core.exceptions import ResourceExistsError + + raise ResourceExistsError("exists") + + def get_container_client(self, name: str) -> FakeContainerClient: + return self.container_client + + +def build_coordinator(blob_client: FakeBlobClient) -> BlobLeaseCoordinator: + container = FakeContainerClient(blob_client) + service = FakeBlobServiceClient(container) + return BlobLeaseCoordinator( + connection_string=None, + account_url=None, + blob_service_client=service, + renewal_interval_seconds=0.01, + ) + + +class TestBlobLeaseCoordinator(unittest.TestCase): + def test_constructor_rejects_invalid_configuration(self) -> None: + with self.assertRaisesRegex(ValueError, "positive"): + BlobLeaseCoordinator(None, None, renewal_interval_seconds=0) + with self.assertRaisesRegex(ValueError, "connection string"): + BlobLeaseCoordinator(None, None) + + def test_constructor_uses_connection_string_and_existing_container( + self, + ) -> None: + container = FakeContainerClient(FakeBlobClient()) + service = ExistingContainerService(container) + with patch( + "hastegeo.core.publishing.lease.BlobServiceClient.from_connection_string", + return_value=service, + ) as factory: + coordinator = BlobLeaseCoordinator("connection", None) + + self.assertIs(coordinator.container_client, container) + factory.assert_called_once_with("connection") + + def test_acquire_rejects_invalid_wait_and_duration(self) -> None: + coordinator = build_coordinator(FakeBlobClient()) + + with self.assertRaisesRegex(ValueError, "between 15 and 60"): + with coordinator.acquire("project", "dataset", lease_duration=10): + pass + with self.assertRaisesRegex(ValueError, "wait values"): + with coordinator.acquire( + "project", "dataset", retry_interval_seconds=0 + ): + pass + + def test_non_conflict_acquire_error_is_preserved(self) -> None: + error = HttpResponseError("service unavailable") + error.status_code = 500 + coordinator = build_coordinator(FakeBlobClient(acquire_error=error)) + + with self.assertRaises(HttpResponseError): + with coordinator.acquire("project", "dataset"): + pass + + def test_held_lease_is_renewed_and_released(self) -> None: + lease = FakeLease() + coordinator = build_coordinator(FakeBlobClient(lease=lease)) + + with coordinator.acquire("project", "dataset", lease_duration=15): + self.assertTrue(lease.renewed.wait(timeout=1)) + + self.assertTrue(lease.released) + + def test_renewal_failure_is_reported_after_operation(self) -> None: + lease = FakeLease(fail_renewal=True) + coordinator = build_coordinator(FakeBlobClient(lease=lease)) + + with self.assertRaisesRegex(LeaseRenewalError, "dataset"): + with coordinator.acquire("project", "dataset", lease_duration=15): + self.assertTrue(lease.renewed.wait(timeout=1)) + + self.assertTrue(lease.released) + + def test_conflicting_lease_maps_to_unavailable(self) -> None: + conflict = HttpResponseError("conflict") + conflict.status_code = 409 + coordinator = build_coordinator(FakeBlobClient(acquire_error=conflict)) + + with self.assertRaises(LeaseUnavailableError): + with coordinator.acquire("project", "dataset", lease_duration=15): + self.fail("lease should not be acquired") + + def test_waiting_claim_retries_contention(self) -> None: + lease = FakeLease() + blob_client = RetryOnceBlobClient(lease) + coordinator = build_coordinator(blob_client) + + with coordinator.acquire( + "project", + "dataset", + lease_duration=15, + wait_timeout_seconds=0.1, + retry_interval_seconds=0.001, + ): + pass + + self.assertEqual(blob_client.attempts, 2) + self.assertTrue(lease.released) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/publishing/test_local_provider.py b/hastelib/tests/core/publishing/test_local_provider.py new file mode 100644 index 00000000..7a4e02e4 --- /dev/null +++ b/hastelib/tests/core/publishing/test_local_provider.py @@ -0,0 +1,146 @@ +import json +import tempfile +import unittest +import uuid +from pathlib import Path + +from hastegeo.core.artifact_storage.unified_artifact_storage import ( + UnifiedArtifactStorage, +) +from hastegeo.core.models.publishing import ( + ArtifactBundle, + PublishedDataset, + PublishRequest, + SourceArtifact, +) +from hastegeo.core.publishing.local_provider import LocalPublishingProvider + + +class FakeConfig: + artifact_storage_type = "local" + + def __init__(self, directory: str) -> None: + self.artifact_storage_config = {"directory": directory} + self.publishing_config = {"publishing_enabled": True} + + +class TestLocalPublishingProvider(unittest.TestCase): + def setUp(self) -> None: + self.temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self.temporary_directory.cleanup) + self.config = FakeConfig(self.temporary_directory.name) + self.storage = UnifiedArtifactStorage( + storage_type="local", + directory=self.temporary_directory.name, + ) + Path(self.temporary_directory.name, "damage.gpkg").write_bytes( + b"damage" + ) + Path(self.temporary_directory.name, "mask.geojson").write_text( + "{}", encoding="utf-8" + ) + self.provider = LocalPublishingProvider( + config=self.config, + artifact_storage=self.storage, + ) + self.project_id = uuid.uuid4() + self.dataset = PublishedDataset( + datasetId=uuid.uuid4(), + requestId=uuid.uuid4(), + requestFingerprint="a" * 64, + name="Dataset", + projectId=self.project_id, + imageLayerId="layer", + modelId="42", + target="local", + status="IN_PROGRESS", + publishedByUser="publisher", + createdDate="2026-08-06T00:00:00Z", + updatedDate="2026-08-06T00:00:00Z", + assessmentSummary={"predictedDamaged": 5}, + ) + self.request = PublishRequest( + requestId=self.dataset.requestId, + projectId=self.project_id, + imageLayerId="layer", + modelId="42", + name="Dataset", + target="local", + artifacts=["gpkg"], + ) + damage_etag = self.storage.get_artifact_etag("damage.gpkg") + mask_etag = self.storage.get_artifact_etag("mask.geojson") + self.bundle = ArtifactBundle( + selectedArtifacts=[ + SourceArtifact( + kind="gpkg", + sourcePath="damage.gpkg", + mediaType="application/geopackage+sqlite3", + sizeBytes=6, + sourceEtag=damage_etag, + ) + ], + supportingArtifacts=[ + SourceArtifact( + kind="valid_mask", + sourcePath="mask.geojson", + mediaType="application/geo+json", + sizeBytes=2, + sourceEtag=mask_etag, + ) + ], + ) + + def test_publish_copies_only_selected_artifacts_and_provenance( + self, + ) -> None: + self.provider.validate(self.request, self.bundle) + + result = self.provider.start_publish(self.dataset, self.bundle) + + self.assertEqual(len(result.artifacts), 1) + self.assertEqual(result.artifacts[0].kind.value, "gpkg") + prefix = f"published/{self.dataset.datasetId}" + self.assertTrue( + self.storage.artifact_exists(result.artifacts[0].publishedPath) + ) + self.assertFalse( + self.storage.artifact_exists(f"{prefix}/valid_mask_mask.geojson") + ) + report_path = result.providerMetadata["assessmentReportPath"] + with open( + Path(self.temporary_directory.name, report_path), + encoding="utf-8", + ) as report_file: + self.assertEqual(json.load(report_file), {"predictedDamaged": 5}) + + def test_publish_and_unpublish_are_idempotent(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + second = self.provider.start_publish(self.dataset, self.bundle) + + self.assertEqual( + first.artifacts[0].publishedPath, + second.artifacts[0].publishedPath, + ) + self.provider.start_unpublish(self.dataset) + repeated = self.provider.start_unpublish(self.dataset) + + self.assertEqual(repeated.providerMetadata["deletedArtifactCount"], 0) + + def test_validate_rejects_empty_bundle(self) -> None: + with self.assertRaisesRegex(ValueError, "at least one"): + self.provider.validate(self.request, ArtifactBundle()) + + def test_disabled_info_and_invalid_continuation_are_explicit(self) -> None: + self.config.publishing_config["publishing_enabled"] = False + self.assertFalse(self.provider.info.isEnabled) + self.assertEqual( + self.provider.info.disabledReason, "Publishing is disabled" + ) + + with self.assertRaisesRegex(RuntimeError, "continuation"): + self.provider.continue_publish(self.dataset, self.bundle) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/publishing/test_planetary_computer_provider.py b/hastelib/tests/core/publishing/test_planetary_computer_provider.py new file mode 100644 index 00000000..90742684 --- /dev/null +++ b/hastelib/tests/core/publishing/test_planetary_computer_provider.py @@ -0,0 +1,901 @@ +import copy +import unittest +import uuid + +from hastegeo.core.models.publishing import ( + ArtifactBundle, + ArtifactKind, + PublishedDataset, + PublishOperation, + PublishRequest, + SourceArtifact, +) +from hastegeo.core.publishing.planetary_computer_provider import ( + PlanetaryComputerPhase, + PlanetaryComputerProviderError, + PlanetaryComputerPublishingProvider, +) +from hastegeo.core.publishing.planetary_computer_transport import ( + PlanetaryComputerOperationKind, + PlanetaryComputerOperationStep, +) + + +class FakeConfig: + artifact_storage_type = "blob" + artifact_storage_config = {} + + def __init__(self) -> None: + self.publishing_config = { + "pc_provider_enabled": True, + "pc_geocatalog_url": "https://catalog.test", + "pc_ingestion_source": "haste-source", + "pc_collection_prefix": "haste-", + "pc_explorer_url": "https://catalog.test", + "pc_verify_attempts": 2, + } + + +class FakeArtifactStorage: + def __init__( + self, + base_url: str = "https://source.blob.core.windows.net/container", + ) -> None: + self.base_url = base_url + self.blobs: set = set() + self.copied: list = [] + + def get_base_url(self) -> str: + return self.base_url + + def resolve_artifact_path(self, path: str) -> str: + return path + + def artifact_exists(self, path: str) -> bool: + return path in self.blobs + + def get_artifact_etag(self, path: str) -> str: + return f"etag-{path}" + + def copy_artifact(self, source: str, destination: str, etag: str) -> str: + self.copied.append((source, destination)) + self.blobs.add(destination) + return destination + + +class FakeSdkAdapter: + def __init__(self) -> None: + self.collections = {} + self.items = {} + self.pending_collection = None + self.pending_item = None + self.pending_delete = None + self.create_collection_calls = 0 + self.create_item_calls = 0 + self.delete_item_calls = 0 + self.replace_calls = 0 + self.materialize_items = True + self.complete_collection_operations = True + self.ingestion_source_calls = 0 + self.delete_collection_calls = 0 + self.pending_collection_delete = None + + def get_ingestion_source(self, source_id): + self.ingestion_source_calls += 1 + return { + "id": source_id, + "kind": "BlobManagedIdentity", + # MPC Pro returns the container as `containerUri` (camelCase). + "connectionInfo": { + "containerUri": "https://source.blob.core.windows.net/container" + }, + } + + @staticmethod + def get_signed_asset_url(href): + return f"{href}?sv=test&sig=secret" + + def get_collection(self, collection_id): + return copy.deepcopy(self.collections.get(collection_id)) + + def replace_collection(self, collection_id, body): + self.replace_calls += 1 + self.collections[collection_id] = copy.deepcopy(body) + + def start_create_collection(self, collection_id, body): + self.create_collection_calls += 1 + self.pending_collection = (collection_id, copy.deepcopy(body)) + return self._operation( + PlanetaryComputerOperationKind.CREATE_COLLECTION, + collection_id, + None, + "https://catalog.test/operations/collection", + False, + ) + + def continue_create_collection(self, collection_id, token): + if not self.complete_collection_operations: + return self._operation( + PlanetaryComputerOperationKind.CREATE_COLLECTION, + collection_id, + None, + token, + False, + ) + pending_id, body = self.pending_collection + self.collections[pending_id] = body + return self._operation( + PlanetaryComputerOperationKind.CREATE_COLLECTION, + collection_id, + None, + None, + True, + ) + + def get_item(self, collection_id, item_id): + return copy.deepcopy(self.items.get((collection_id, item_id))) + + def start_create_item(self, collection_id, item_id, body): + self.create_item_calls += 1 + self.pending_item = (collection_id, item_id, copy.deepcopy(body)) + return self._operation( + PlanetaryComputerOperationKind.CREATE_ITEM, + collection_id, + item_id, + "https://catalog.test/operations/item", + False, + ) + + def continue_create_item(self, collection_id, item_id, token): + if self.materialize_items: + pending_collection, pending_item, body = self.pending_item + self.items[ + (pending_collection, pending_item) + ] = self._managed_item(body) + return self._operation( + PlanetaryComputerOperationKind.CREATE_ITEM, + collection_id, + item_id, + None, + True, + ) + + def start_delete_item(self, collection_id, item_id): + self.delete_item_calls += 1 + self.pending_delete = (collection_id, item_id) + return self._operation( + PlanetaryComputerOperationKind.DELETE_ITEM, + collection_id, + item_id, + "https://catalog.test/operations/delete", + False, + ) + + def continue_delete_item(self, collection_id, item_id, token): + self.items.pop(self.pending_delete, None) + return self._operation( + PlanetaryComputerOperationKind.DELETE_ITEM, + collection_id, + item_id, + None, + True, + ) + + def list_item_ids(self, collection_id, limit=100): + return [ + item_id + for (cid, item_id) in self.items + if cid == collection_id + ] + + def start_delete_collection(self, collection_id): + self.delete_collection_calls += 1 + self.pending_collection_delete = collection_id + return self._operation( + PlanetaryComputerOperationKind.DELETE_COLLECTION, + collection_id, + None, + "https://catalog.test/operations/delete-collection", + False, + ) + + def continue_delete_collection(self, collection_id, token): + self.collections.pop(self.pending_collection_delete, None) + return self._operation( + PlanetaryComputerOperationKind.DELETE_COLLECTION, + collection_id, + None, + None, + True, + ) + + @staticmethod + def _operation(kind, collection_id, item_id, token, complete): + return PlanetaryComputerOperationStep( + kind=kind, + collection_id=collection_id, + item_id=item_id, + continuation_token=token, + is_complete=complete, + ) + + @staticmethod + def _managed_item(body): + item = copy.deepcopy(body) + for key, asset in item["assets"].items(): + asset["href"] = ( + "https://managed.blob.core.windows.net/" + f"collection/{key}.data" + ) + return item + + +class TestPlanetaryComputerPublishingProvider(unittest.TestCase): + def setUp(self) -> None: + self.config = FakeConfig() + self.storage = FakeArtifactStorage() + self.sdk = FakeSdkAdapter() + self.valid_mask = { + "type": "FeatureCollection", + "features": [ + { + "type": "Feature", + "properties": {}, + "geometry": { + "type": "Polygon", + "coordinates": [ + [ + [-67.1, 10.4], + [-67.0, 10.4], + [-67.0, 10.5], + [-67.1, 10.5], + [-67.1, 10.4], + ] + ], + }, + } + ], + } + self.damage = SourceArtifact( + kind=ArtifactKind.GPKG, + sourcePath="project/damage.gpkg", + mediaType="application/geopackage+sqlite3", + sizeBytes=100, + sourceEtag="damage-etag", + ) + self.mask = SourceArtifact( + kind=ArtifactKind.VALID_MASK, + sourcePath="project/mask.geojson", + mediaType="application/geo+json", + sizeBytes=100, + sourceEtag="mask-etag", + ) + self.bundle = ArtifactBundle( + selectedArtifacts=[self.damage], + supportingArtifacts=[self.mask], + ) + self.dataset = PublishedDataset( + datasetId=uuid.UUID("11111111-1111-4111-8111-111111111111"), + requestId=uuid.UUID("22222222-2222-4222-8222-222222222222"), + requestFingerprint="a" * 64, + name="Caracas damage assessment", + description="Post-event damage", + projectId=uuid.UUID("33333333-3333-4333-8333-333333333333"), + projectName="Caracas", + imageLayerId="layer-1", + imageLayerName="Post-event", + modelId="42", + modelName="Damage model", + target="planetary_computer", + status="IN_PROGRESS", + publishedByUser="publisher", + createdDate="2026-08-08T00:00:00Z", + updatedDate="2026-08-08T00:00:00Z", + selectedArtifactKinds=[ArtifactKind.GPKG], + sourceArtifacts=[self.damage], + assessmentSummary={"predictedDamaged": 10}, + ) + self.request = PublishRequest( + requestId=self.dataset.requestId, + projectId=self.dataset.projectId, + imageLayerId=self.dataset.imageLayerId, + modelId=self.dataset.modelId, + name=self.dataset.name, + target="planetary_computer", + artifacts=[ArtifactKind.GPKG], + ) + self.provider = PlanetaryComputerPublishingProvider( + config=self.config, + artifact_storage=self.storage, + sdk_adapter=self.sdk, + json_reader=lambda artifact: self.valid_mask, + projection_resolver=lambda artifact: "EPSG:4326", + asset_reachability_checker=self._record_reachable_asset, + ) + self.reachable_assets = [] + + def _record_reachable_asset(self, href: str) -> None: + self.reachable_assets.append(href) + + @staticmethod + def _continued(dataset, result): + return dataset.model_copy( + update={ + "providerMetadata": { + **dataset.providerMetadata, + **result.providerMetadata, + "continuationToken": result.continuationToken, + } + } + ) + + def test_new_collection_publish_runs_all_bounded_phases(self) -> None: + self.provider.validate(self.request, self.bundle) + self.assertEqual(self.sdk.ingestion_source_calls, 0) + + collection_pending = self.provider.start_publish( + self.dataset, + self.bundle, + ) + collection_dataset = self._continued( + self.dataset, + collection_pending, + ) + item_pending = self.provider.continue_publish( + collection_dataset, + self.bundle, + ) + item_dataset = self._continued(collection_dataset, item_pending) + completed = self.provider.continue_publish( + item_dataset, + self.bundle, + ) + + self.assertFalse(collection_pending.isComplete) + self.assertEqual( + collection_pending.providerMetadata["phase"], + PlanetaryComputerPhase.COLLECTION_OPERATION.value, + ) + self.assertEqual( + item_pending.providerMetadata["phase"], + PlanetaryComputerPhase.ITEM_OPERATION.value, + ) + self.assertTrue(completed.isComplete) + self.assertEqual( + completed.links["stac_collection"], + "https://catalog.test/stac/collections/" + "haste-33333333-3333-4333-8333-333333333333", + ) + self.assertEqual(completed.links["explorer"], "https://catalog.test") + self.assertEqual( + completed.artifacts[0].publishedPath, + "https://managed.blob.core.windows.net/collection/damage.data", + ) + self.assertTrue( + completed.providerMetadata["assetsCopiedToManagedStorage"] + ) + self.assertNotIn("phase", completed.providerMetadata) + self.assertEqual(self.sdk.create_collection_calls, 1) + self.assertEqual(self.sdk.create_item_calls, 1) + self.assertEqual(self.sdk.ingestion_source_calls, 1) + self.assertEqual(len(self.reachable_assets), 1) + self.assertIn("sig=secret", self.reachable_assets[0]) + + def test_existing_item_replay_is_verified_without_duplicate_ingest( + self, + ) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + collection_dataset = self._continued(self.dataset, first) + second = self.provider.continue_publish( + collection_dataset, + self.bundle, + ) + item_dataset = self._continued(collection_dataset, second) + completed = self.provider.continue_publish( + item_dataset, + self.bundle, + ) + + replayed = self.provider.start_publish(self.dataset, self.bundle) + + self.assertTrue(completed.isComplete) + self.assertTrue(replayed.isComplete) + self.assertEqual(self.sdk.create_item_calls, 1) + self.assertGreaterEqual(self.sdk.replace_calls, 2) + + def test_retry_replaces_invalid_deterministic_item(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + completed = self.provider.continue_publish(current, self.bundle) + collection_id = completed.providerMetadata["collectionId"] + item_id = completed.providerMetadata["itemIds"][0] + self.sdk.items[(collection_id, item_id)]["assets"]["damage"][ + "type" + ] = "text/plain" + + delete_pending = self.provider.start_publish( + self.dataset, + self.bundle, + ) + deleting = self._continued(self.dataset, delete_pending) + ingest_pending = self.provider.continue_publish( + deleting, + self.bundle, + ) + ingesting = self._continued(deleting, ingest_pending) + retried = self.provider.continue_publish( + ingesting, + self.bundle, + ) + + self.assertEqual( + delete_pending.providerMetadata["phase"], + PlanetaryComputerPhase.ITEM_REPLACE_DELETE_OPERATION.value, + ) + self.assertEqual( + ingest_pending.providerMetadata["phase"], + PlanetaryComputerPhase.ITEM_OPERATION.value, + ) + self.assertTrue(retried.isComplete) + self.assertEqual(self.sdk.delete_item_calls, 1) + self.assertEqual(self.sdk.create_item_calls, 2) + + def test_item_verification_is_bounded(self) -> None: + self.sdk.materialize_items = False + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + verification = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, verification) + verification = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, verification) + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "verification timed out", + ): + self.provider.continue_publish(current, self.bundle) + + def test_operation_polling_is_bounded_and_retry_resets_state(self) -> None: + self.sdk.complete_collection_operations = False + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + for _ in range(self.provider.max_verify_attempts): + pending = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, pending) + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "ingestion timed out", + ): + self.provider.continue_publish(current, self.bundle) + + reset = self.provider.prepare_retry( + current, + PublishOperation.PUBLISH, + ) + self.assertNotIn("phase", reset) + self.assertNotIn("continuationToken", reset) + self.assertNotIn("operationAttempts", reset) + + def test_unpublish_deletes_only_item_and_is_idempotent(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + completed = self.provider.continue_publish(current, self.bundle) + published = self.dataset.model_copy( + update={"providerMetadata": completed.providerMetadata} + ) + collection_id = completed.providerMetadata["collectionId"] + + delete_pending = self.provider.start_unpublish(published) + deleting = self._continued(published, delete_pending) + item_deleted = self.provider.continue_unpublish(deleting) + collection_deleting = self._continued(deleting, item_deleted) + deleted = self.provider.continue_unpublish(collection_deleting) + repeated = self.provider.start_unpublish(published) + + self.assertFalse(delete_pending.isComplete) + # Item delete completes, then the now-empty collection is deleted. + self.assertFalse(item_deleted.isComplete) + self.assertEqual( + item_deleted.providerMetadata["phase"], + PlanetaryComputerPhase.DELETE_COLLECTION_OPERATION.value, + ) + self.assertTrue(deleted.isComplete) + self.assertTrue(repeated.isComplete) + self.assertNotIn(collection_id, self.sdk.collections) + self.assertEqual(self.sdk.delete_item_calls, 1) + self.assertEqual(self.sdk.delete_collection_calls, 1) + + def test_unpublish_keeps_collection_with_remaining_datasets(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + completed = self.provider.continue_publish(current, self.bundle) + published = self.dataset.model_copy( + update={"providerMetadata": completed.providerMetadata} + ) + collection_id = completed.providerMetadata["collectionId"] + # A second dataset still lives in the same (project-level) collection. + self.sdk.items[(collection_id, "other-item")] = {"id": "other-item"} + + delete_pending = self.provider.start_unpublish(published) + deleting = self._continued(published, delete_pending) + deleted = self.provider.continue_unpublish(deleting) + + self.assertTrue(deleted.isComplete) + self.assertEqual(self.sdk.delete_item_calls, 1) + # Collection survives because another dataset remains; no delete. + self.assertEqual(self.sdk.delete_collection_calls, 0) + self.assertIn(collection_id, self.sdk.collections) + # Its rolling summary drops the unpublished dataset. + remaining = self.sdk.collections[collection_id]["ai4g:datasets"] + self.assertNotIn( + str(self.dataset.datasetId), + [entry["id"] for entry in remaining], + ) + + def test_unpublish_drains_inflight_item_before_deleting_it(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + item_pending = self.provider.continue_publish(current, self.bundle) + failed = self._continued(current, item_pending) + + delete_pending = self.provider.start_unpublish(failed) + deleting = self._continued(failed, delete_pending) + item_deleted = self.provider.continue_unpublish(deleting) + collection_deleting = self._continued(deleting, item_deleted) + deleted = self.provider.continue_unpublish(collection_deleting) + + self.assertEqual( + delete_pending.providerMetadata["phase"], + PlanetaryComputerPhase.DELETE_OPERATION.value, + ) + self.assertTrue(deleted.isComplete) + self.assertEqual(self.sdk.delete_item_calls, 1) + self.assertEqual(self.sdk.delete_collection_calls, 1) + collection_id = failed.providerMetadata["collectionId"] + item_id = failed.providerMetadata["itemIds"][0] + self.assertNotIn((collection_id, item_id), self.sdk.items) + self.assertNotIn(collection_id, self.sdk.collections) + + def test_unpublish_collection_crash_window_requires_item_discovery( + self, + ) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + stale = self._continued(self.dataset, first) + collection_id = stale.providerMetadata["collectionId"] + item_id = stale.providerMetadata["itemIds"][0] + self.sdk.continue_create_collection( + collection_id, + stale.providerMetadata["continuationToken"], + ) + documents = self.provider._build_documents( + self.dataset, + self.bundle, + self.provider._projection_codes(self.dataset, self.bundle), + self.sdk.get_collection(collection_id), + ) + self.sdk.start_create_item(collection_id, item_id, documents.item) + + discovery = self.provider.start_unpublish(stale) + + self.assertFalse(discovery.isComplete) + self.assertEqual( + discovery.providerMetadata["phase"], + PlanetaryComputerPhase.DELETE_DISCOVER.value, + ) + + def test_unpublish_discovery_timeout_fails_closed(self) -> None: + metadata = self.provider._cleanup_discovery_metadata( + self.provider._stable_metadata(self.dataset) + ) + current = self.dataset.model_copy( + update={"providerMetadata": metadata} + ) + for _ in range(self.provider.max_verify_attempts): + pending = self.provider._start_delete_or_discover( + current, + current.providerMetadata, + metadata["collectionId"], + metadata["itemIds"][0], + ) + current = self._continued(current, pending) + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "cleanup verification timed out", + ): + self.provider._start_delete_or_discover( + current, + current.providerMetadata, + metadata["collectionId"], + metadata["itemIds"][0], + ) + + def test_failed_phase_less_unpublish_requires_discovery(self) -> None: + failed = self.dataset.model_copy(update={"providerMetadata": {}}) + + discovery = self.provider.start_unpublish(failed) + + self.assertFalse(discovery.isComplete) + self.assertEqual( + discovery.providerMetadata["phase"], + PlanetaryComputerPhase.DELETE_DISCOVER.value, + ) + + def test_duplicate_delete_enters_bounded_verification(self) -> None: + class ConflictError(RuntimeError): + status_code = 409 + + metadata = self.provider._stable_metadata(self.dataset) + metadata["assetsCopiedToManagedStorage"] = True + collection_id = metadata["collectionId"] + item_id = metadata["itemIds"][0] + self.sdk.items[(collection_id, item_id)] = { + "id": item_id, + "collection": collection_id, + } + self.sdk.start_delete_item = lambda *args: (_ for _ in ()).throw( + ConflictError("already deleting") + ) + published = self.dataset.model_copy( + update={"providerMetadata": metadata} + ) + + pending = self.provider.start_unpublish(published) + + self.assertFalse(pending.isComplete) + self.assertEqual( + pending.providerMetadata["phase"], + PlanetaryComputerPhase.DELETE_VERIFY.value, + ) + + def test_validation_rejects_emulator_and_missing_mask(self) -> None: + emulator_provider = PlanetaryComputerPublishingProvider( + config=self.config, + artifact_storage=FakeArtifactStorage( + "https://devstoreaccount1.blob.core.windows.net/data" + ), + sdk_adapter=self.sdk, + json_reader=lambda artifact: self.valid_mask, + projection_resolver=lambda artifact: "EPSG:4326", + asset_reachability_checker=lambda href: None, + ) + + with self.assertRaisesRegex(ValueError, "storage emulator"): + emulator_provider.validate(self.request, self.bundle) + with self.assertRaisesRegex(ValueError, "valid-area mask"): + self.provider.validate( + self.request, + ArtifactBundle(selectedArtifacts=[self.damage]), + ) + + def test_explorer_url_allows_query_but_geocatalog_stays_strict( + self, + ) -> None: + # MPC Pro Explorer links carry a query string + # (?geocatalogname=...&c=...&z=...); the display-only Explorer URL must + # accept it, while the GeoCatalog URL stays scheme + host only. + explorer = ( + "https://explorer.geocatalog.spatio.azure.com/explorer" + "?geocatalogname=damage-assessment-gc.eastus&c=30.05%2C29.99&z=2" + ) + config = FakeConfig() + config.publishing_config["pc_explorer_url"] = explorer + provider = PlanetaryComputerPublishingProvider( + config=config, + artifact_storage=self.storage, + sdk_adapter=self.sdk, + json_reader=lambda artifact: self.valid_mask, + projection_resolver=lambda artifact: "EPSG:4326", + asset_reachability_checker=self._record_reachable_asset, + ) + + provider.validate(self.request, self.bundle) + self.assertEqual(provider.explorer_url, explorer) + + # A query string (or path) on the GeoCatalog URL is still rejected. + strict_config = FakeConfig() + strict_config.publishing_config["pc_geocatalog_url"] = ( + "https://catalog.test/stac?foo=bar" + ) + strict_provider = PlanetaryComputerPublishingProvider( + config=strict_config, + artifact_storage=self.storage, + sdk_adapter=self.sdk, + json_reader=lambda artifact: self.valid_mask, + projection_resolver=lambda artifact: "EPSG:4326", + asset_reachability_checker=self._record_reachable_asset, + ) + with self.assertRaisesRegex(ValueError, "must use HTTPS"): + strict_provider.validate(self.request, self.bundle) + + def test_configured_license_is_applied_to_stac_documents(self) -> None: + self.config.publishing_config["pc_publishing_license"] = "CC-BY-SA-4.0" + provider = PlanetaryComputerPublishingProvider( + config=self.config, + artifact_storage=self.storage, + sdk_adapter=self.sdk, + json_reader=lambda artifact: self.valid_mask, + projection_resolver=lambda artifact: "EPSG:4326", + asset_reachability_checker=lambda href: None, + ) + + documents = provider._build_documents( + self.dataset, + self.bundle, + provider._projection_codes(self.dataset, self.bundle), + None, + ) + + self.assertEqual(documents.collection["license"], "CC-BY-SA-4.0") + self.assertEqual( + documents.item["properties"]["license"], "CC-BY-SA-4.0" + ) + # Default when unset is CC-BY-4.0 (not the old hardcoded "proprietary"). + self.assertEqual(self.provider.license_id, "CC-BY-4.0") + + def test_thumbnail_is_copied_into_published_prefix(self) -> None: + self.storage.blobs.add("hash/task/preview_post_event.png") + bundle = ArtifactBundle( + selectedArtifacts=[self.damage], + supportingArtifacts=[self.mask], + thumbnailUrl=( + "https://source.blob.core.windows.net/container/" + "hash/task/preview_post_event.png" + ), + ) + + href, media_type = self.provider._resolve_thumbnail_href( + self.dataset, bundle + ) + + self.assertEqual(media_type, "image/png") + self.assertEqual( + href, + "https://source.blob.core.windows.net/container/" + f"published/{self.dataset.datasetId}/thumbnail.png", + ) + self.assertIn( + ( + "hash/task/preview_post_event.png", + f"published/{self.dataset.datasetId}/thumbnail.png", + ), + self.storage.copied, + ) + + def test_thumbnail_from_a_foreign_container_is_skipped(self) -> None: + bundle = ArtifactBundle( + selectedArtifacts=[self.damage], + supportingArtifacts=[self.mask], + thumbnailUrl="https://other.blob.core.windows.net/data/x.png", + ) + + href, _ = self.provider._resolve_thumbnail_href(self.dataset, bundle) + + self.assertIsNone(href) + self.assertEqual(self.storage.copied, []) + + def test_valid_mask_crs_is_explicit_and_validated(self) -> None: + projected = copy.deepcopy(self.valid_mask) + projected["crs"] = { + "type": "name", + "properties": {"name": "EPSG:3857"}, + } + + self.assertEqual( + self.provider._valid_mask_crs(projected), + "EPSG:3857", + ) + with self.assertRaisesRegex(ValueError, "CRS is invalid"): + self.provider._valid_mask_crs( + {**projected, "crs": {"type": "name"}} + ) + + def test_verification_rejects_changed_asset_metadata(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + collection_id = current.providerMetadata["collectionId"] + item_id = current.providerMetadata["itemIds"][0] + _, _, pending_body = self.sdk.pending_item + actual = self.sdk._managed_item(pending_body) + actual["assets"]["damage"]["type"] = "text/plain" + self.sdk.items[(collection_id, item_id)] = actual + self.sdk.materialize_items = False + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "asset type changed", + ): + self.provider.continue_publish(current, self.bundle) + + def test_verification_rejects_uncopied_source_asset(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + collection_id = current.providerMetadata["collectionId"] + item_id = current.providerMetadata["itemIds"][0] + _, _, pending_body = self.sdk.pending_item + self.sdk.items[(collection_id, item_id)] = copy.deepcopy(pending_body) + self.sdk.materialize_items = False + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "did not copy asset", + ): + self.provider.continue_publish(current, self.bundle) + + def test_verification_rejects_unsigned_or_wrong_host_asset(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + collection_id = current.providerMetadata["collectionId"] + item_id = current.providerMetadata["itemIds"][0] + _, _, pending_body = self.sdk.pending_item + actual = self.sdk._managed_item(pending_body) + actual["assets"]["damage"][ + "href" + ] = "https://attacker.test/damage.gpkg" + self.sdk.items[(collection_id, item_id)] = actual + self.sdk.materialize_items = False + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "not Azure Blob Storage", + ): + self.provider.continue_publish(current, self.bundle) + + def test_verification_rejects_unexpected_asset(self) -> None: + first = self.provider.start_publish(self.dataset, self.bundle) + current = self._continued(self.dataset, first) + second = self.provider.continue_publish(current, self.bundle) + current = self._continued(current, second) + collection_id = current.providerMetadata["collectionId"] + item_id = current.providerMetadata["itemIds"][0] + _, _, pending_body = self.sdk.pending_item + actual = self.sdk._managed_item(pending_body) + actual["assets"]["unexpected"] = { + "href": "https://managed.blob.core.windows.net/collection/extra", + "type": "application/octet-stream", + "roles": ["data"], + } + self.sdk.items[(collection_id, item_id)] = actual + self.sdk.materialize_items = False + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "selected assets changed", + ): + self.provider.continue_publish(current, self.bundle) + + def test_ingestion_source_must_match_haste_container(self) -> None: + self.sdk.get_ingestion_source = lambda source_id: { + "id": source_id, + "kind": "BlobManagedIdentity", + "connectionInfo": { + "containerUrl": ( + "https://other.blob.core.windows.net/container" + ) + }, + } + + with self.assertRaisesRegex( + PlanetaryComputerProviderError, + "does not match HASTE storage", + ): + self.provider.start_publish(self.dataset, self.bundle) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/publishing/test_planetary_computer_transport.py b/hastelib/tests/core/publishing/test_planetary_computer_transport.py new file mode 100644 index 00000000..6aae8141 --- /dev/null +++ b/hastelib/tests/core/publishing/test_planetary_computer_transport.py @@ -0,0 +1,222 @@ +import unittest +from typing import Any, Optional + +from hastegeo.core.publishing.geocatalog_client import ( + GeoCatalogAuth, + GeoCatalogClient, + GeoCatalogError, +) +from hastegeo.core.publishing.planetary_computer_transport import ( + PlanetaryComputerOperationError, + PlanetaryComputerOperationKind, + PlanetaryComputerRestAdapter, +) + +ENDPOINT = "https://cat.example.geocatalog.spatio.azure.com" +OP_URL = f"{ENDPOINT}/inma/operations/op-1?api-version=2026-04-15" + + +class FakeResponse: + def __init__(self, status_code, headers=None, payload=None): + self.status_code = status_code + self.headers = headers or {} + self._payload = payload + + def json(self): + return self._payload + + +class FakeClient: + """Stands in for GeoCatalogClient; replays queued responses and mirrors its + raise-on-unexpected-status contract so 4xx branches are exercised.""" + + def __init__(self, handler): + self.handler = handler + self.calls = [] + + def request( + self, + method, + url, + *, + json=None, + params=None, + expected=(200, 201, 202, 204), + absolute=False, + ): + self.calls.append( + {"method": method, "url": url, "json": json, "absolute": absolute} + ) + response = self.handler(method, url) + if response.status_code not in tuple(expected): + raise GeoCatalogError( + f"{method} {url} -> {response.status_code}", + status_code=response.status_code, + ) + return response + + +def adapter(handler): + return PlanetaryComputerRestAdapter(ENDPOINT, client=FakeClient(handler)) + + +class TestGeoCatalogClient(unittest.TestCase): + def _client(self, response): + auth = GeoCatalogAuth(credential=_FakeCredential()) + client = GeoCatalogClient(ENDPOINT, auth=auth) + client._session = _FakeSession(response) + return client + + def test_builds_url_adds_api_version_and_auth(self): + client = self._client(FakeResponse(200, payload={"ok": True})) + client.request("GET", "/stac/collections/x") + call = client._session.last + self.assertEqual( + call["url"], f"{ENDPOINT}/stac/collections/x" + ) + self.assertEqual(call["params"]["api-version"], "2026-04-15") + self.assertEqual( + call["headers"]["Authorization"], "Bearer test-token" + ) + self.assertFalse(call["allow_redirects"]) + + def test_absolute_url_not_prefixed(self): + client = self._client(FakeResponse(200, payload={})) + client.request("GET", OP_URL, absolute=True) + self.assertEqual(client._session.last["url"], OP_URL) + + def test_unexpected_status_raises_without_body(self): + client = self._client(FakeResponse(409, payload={"secret": "x"})) + with self.assertRaises(GeoCatalogError) as ctx: + client.request("POST", "/stac/collections", expected=(201,)) + self.assertEqual(ctx.exception.status_code, 409) + self.assertNotIn("secret", str(ctx.exception)) + + +class TestRestAdapter(unittest.TestCase): + def test_start_collection_async_202_pins_operation_url(self): + step = adapter( + lambda m, u: FakeResponse(202, {"operation-location": OP_URL}) + ).start_create_collection("c", {"id": "c"}) + self.assertFalse(step.is_complete) + self.assertEqual(step.continuation_token, OP_URL) + self.assertEqual( + step.kind, PlanetaryComputerOperationKind.CREATE_COLLECTION + ) + + def test_start_collection_sync_201_is_complete(self): + step = adapter( + lambda m, u: FakeResponse(201, payload={"id": "c"}) + ).start_create_collection("c", {"id": "c"}) + self.assertTrue(step.is_complete) + self.assertIsNone(step.continuation_token) + + def test_start_rejects_offorigin_operation_url(self): + evil = "https://evil.example.com/inma/operations/op-1" + with self.assertRaises(ValueError): + adapter( + lambda m, u: FakeResponse(202, {"operation-location": evil}) + ).start_create_item("c", "i", {"id": "i"}) + + def test_continue_in_progress_stays_pending(self): + step = adapter( + lambda m, u: FakeResponse(200, payload={"status": "Running"}) + ).continue_create_item("c", "i", OP_URL) + self.assertFalse(step.is_complete) + self.assertEqual(step.continuation_token, OP_URL) + + def test_continue_finished_is_success(self): + # Regression: "Finished" must be treated as terminal success. + step = adapter( + lambda m, u: FakeResponse(200, payload={"status": "Finished"}) + ).continue_create_item("c", "i", OP_URL) + self.assertTrue(step.is_complete) + self.assertIsNone(step.continuation_token) + + def test_continue_succeeded_with_failed_items_raises(self): + payload = { + "status": "Succeeded", + "additionalInformation": {"totalFailedItems": 2}, + } + with self.assertRaises(PlanetaryComputerOperationError): + adapter(lambda m, u: FakeResponse(200, payload=payload)).\ + continue_create_item("c", "i", OP_URL) + + def test_continue_failed_sanitizes_error_code(self): + payload = { + "status": "Failed", + "error": {"code": "Bad Code