From 013f66611760244ae68eaa1700ce9acb55355845 Mon Sep 17 00:00:00 2001 From: Anthony Cintron Roman Date: Sun, 9 Aug 2026 10:17:03 -0400 Subject: [PATCH 01/17] =?UTF-8?q?feat(publishing):=20core=20library=20?= =?UTF-8?q?=E2=80=94=20Local=20provider,=20repository,=20source,=20lease?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bucket A (additive) from the -sol audit split, Phase 1 (Local only; no PC): - core/publishing/: base ABC + registry (PC lazy-loaded, disabled by default), repository (per-dataset docs + optimistic concurrency), source (eligibility + artifact resolution), lease (blob-lease coordinator), local_provider (immutable copy to published/{datasetId}/). - models/publishing.py; processors/publishing.py + assessment.py. - Additive infra: config publishing block + PUBLISHED_DATASET type + publish queue; data-layer load_bounded/load_page/update_locked; artifact-storage copy/scoped-download/etag/size/exists/delete_prefix; blob.py max_bytes. - hastefuncqueues publish trigger. - Phase-1 unit tests (registry/repository/source/local_provider/lease): 45 pass. Excludes Bucket B: reverted the high-risk load_all/load_all_from_partition rewrite to main's versions; endpoint hardening + DTOs left for a separate -hardening PR. See spec/features/data-publishing/sol-classification.md. --- api/hastefuncqueues/function_app.py | 74 ++ api/hastefuncqueues/requirements.txt | 4 +- .../abstract_artifact_storage.py | 35 + .../azure_blob_artifact_storage.py | 221 ++++- .../local_file_system_artifact_storage.py | 120 ++- .../unified_artifact_storage.py | 32 + hastelib/src/hastegeo/core/config.py | 52 ++ .../core/data_layer/abstract_data_layer.py | 8 + .../azure_blob_storage_data_layer.py | 184 +++- .../data_layer/azure_cosmos_db_data_layer.py | 31 +- .../data_layer/azure_data_lake_data_layer.py | 28 +- .../data_layer/azure_postgresql_data_layer.py | 22 +- .../local_file_system_data_layer.py | 46 +- .../src/hastegeo/core/data_layer/unified.py | 33 +- .../src/hastegeo/core/models/publishing.py | 223 +++++ .../hastegeo/core/processors/assessment.py | 145 +++ .../src/hastegeo/core/processors/metadata.py | 30 + .../hastegeo/core/processors/publishing.py | 872 ++++++++++++++++++ .../src/hastegeo/core/publishing/__init__.py | 21 + hastelib/src/hastegeo/core/publishing/base.py | 53 ++ .../src/hastegeo/core/publishing/lease.py | 136 +++ .../core/publishing/local_provider.py | 111 +++ .../src/hastegeo/core/publishing/registry.py | 120 +++ .../hastegeo/core/publishing/repository.py | 322 +++++++ .../src/hastegeo/core/publishing/source.py | 443 +++++++++ hastelib/src/hastegeo/core/utils/blob.py | 54 +- hastelib/tests/core/publishing/test_lease.py | 182 ++++ .../core/publishing/test_local_provider.py | 150 +++ .../tests/core/publishing/test_registry.py | 104 +++ .../tests/core/publishing/test_repository.py | 427 +++++++++ hastelib/tests/core/publishing/test_source.py | 371 ++++++++ 31 files changed, 4597 insertions(+), 57 deletions(-) create mode 100644 hastelib/src/hastegeo/core/models/publishing.py create mode 100644 hastelib/src/hastegeo/core/processors/assessment.py create mode 100644 hastelib/src/hastegeo/core/processors/publishing.py create mode 100644 hastelib/src/hastegeo/core/publishing/__init__.py create mode 100644 hastelib/src/hastegeo/core/publishing/base.py create mode 100644 hastelib/src/hastegeo/core/publishing/lease.py create mode 100644 hastelib/src/hastegeo/core/publishing/local_provider.py create mode 100644 hastelib/src/hastegeo/core/publishing/registry.py create mode 100644 hastelib/src/hastegeo/core/publishing/repository.py create mode 100644 hastelib/src/hastegeo/core/publishing/source.py create mode 100644 hastelib/tests/core/publishing/test_lease.py create mode 100644 hastelib/tests/core/publishing/test_local_provider.py create mode 100644 hastelib/tests/core/publishing/test_registry.py create mode 100644 hastelib/tests/core/publishing/test_repository.py create mode 100644 hastelib/tests/core/publishing/test_source.py 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..dc86ca81 100644 --- a/api/hastefuncqueues/requirements.txt +++ b/api/hastefuncqueues/requirements.txt @@ -12,8 +12,9 @@ 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 +azure-planetarycomputer==1.0.0 email_validator==2.3.0 psycopg2-binary==2.9.9 requests==2.33.0 @@ -23,6 +24,7 @@ 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 +pystac[validation]==1.11.0 boto3==1.36.20 tensorboard==2.19.0 tenacity==9.1.2 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..fa83cf10 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,41 @@ 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", False + ), + "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_verify_attempts": _get_bounded_int_env( + "PC_VERIFY_ATTEMPTS", 5, 1, 20 + ), + "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 +376,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..517cf4b2 --- /dev/null +++ b/hastelib/src/hastegeo/core/models/publishing.py @@ -0,0 +1,223 @@ +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) + + 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 + 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..e1ff3c49 --- /dev/null +++ b/hastelib/src/hastegeo/core/processors/publishing.py @@ -0,0 +1,872 @@ +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 + 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, + ) -> PublishedDataset: + prepared = self.prepare_create(request, publisher_id) + return self.create_prepared(prepared, assessment_summary) + + def prepare_create( + self, + request: PublishRequest, + publisher_id: str, + ) -> 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, + 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, + 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(), + 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/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..a43163d6 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/local_provider.py @@ -0,0 +1,111 @@ +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 ..utils.metadata import MetadataUtils +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: + project_hash = MetadataUtils.hash_string(str(dataset.projectId)) + return f"{project_hash}/published/{dataset.datasetId}" diff --git a/hastelib/src/hastegeo/core/publishing/registry.py b/hastelib/src/hastegeo/core/publishing/registry.py new file mode 100644 index 00000000..7538744d --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/registry.py @@ -0,0 +1,120 @@ +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"] + pc_configured = bool( + settings["pc_geocatalog_url"] and settings["pc_ingestion_source"] + ) + 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/source.py b/hastelib/src/hastegeo/core/publishing/source.py new file mode 100644 index 00000000..6118d064 --- /dev/null +++ b/hastelib/src/hastegeo/core/publishing/source.py @@ -0,0 +1,443 @@ +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" + ) + if ( + model.inferenceStatus + != self.config.get_status_types().COMPLETED.value + ): + raise PublishingSourceNotEligibleError( + "Model inference 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]) + return ArtifactBundle( + selectedArtifacts=selected, + supportingArtifacts=supporting, + ) 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..f684c319 --- /dev/null +++ b/hastelib/tests/core/publishing/test_local_provider.py @@ -0,0 +1,150 @@ +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 +from hastegeo.core.utils.metadata import MetadataUtils + + +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"{MetadataUtils.hash_string(str(self.project_id))}/published/" + f"{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_registry.py b/hastelib/tests/core/publishing/test_registry.py new file mode 100644 index 00000000..eb182b37 --- /dev/null +++ b/hastelib/tests/core/publishing/test_registry.py @@ -0,0 +1,104 @@ +import unittest +from unittest.mock import patch + +from hastegeo.core.models.publishing import ArtifactKind +from hastegeo.core.publishing.registry import ( + ProviderUnavailableError, + PublishingProviderRegistry, +) + + +class FakeConfig: + def __init__( + self, enabled: bool, pc_enabled: bool, configured: bool + ) -> None: + self.publishing_config = { + "publishing_enabled": enabled, + "pc_provider_enabled": pc_enabled, + "pc_geocatalog_url": "https://catalog" if configured else None, + "pc_ingestion_source": "source" if configured else None, + } + + +class TestPublishingProviderRegistry(unittest.TestCase): + def test_lists_all_known_provider_descriptors(self) -> None: + registry = PublishingProviderRegistry( + config=FakeConfig(True, False, False) + ) + + infos = registry.list_infos() + + self.assertEqual( + [info.id for info in infos], ["local", "planetary_computer"] + ) + + def test_unknown_provider_is_rejected(self) -> None: + registry = PublishingProviderRegistry( + config=FakeConfig(True, False, False) + ) + + with self.assertRaisesRegex(ProviderUnavailableError, "Unknown"): + registry.get_info("unknown") + + def test_default_local_factory_is_loaded_lazily(self) -> None: + provider = object() + registry = PublishingProviderRegistry( + config=FakeConfig(True, False, False) + ) + + with patch( + "hastegeo.core.publishing.local_provider.LocalPublishingProvider", + return_value=provider, + ) as provider_type: + resolved = registry.resolve("local") + + self.assertIs(resolved, provider) + provider_type.assert_called_once_with(config=registry.config) + + def test_disabled_pc_descriptor_does_not_invoke_factory(self) -> None: + invoked = [] + registry = PublishingProviderRegistry( + config=FakeConfig(True, False, True), + factories={"planetary_computer": lambda: invoked.append(True)}, + ) + + info = registry.get_info("planetary_computer") + + self.assertFalse(info.isEnabled) + self.assertTrue(info.isConfigured) + self.assertEqual( + info.supportedArtifactKinds, + [ + ArtifactKind.GPKG, + ArtifactKind.VALID_MASK, + ArtifactKind.FOOTPRINTS, + ], + ) + with self.assertRaises(ProviderUnavailableError): + registry.resolve("planetary_computer") + self.assertEqual(invoked, []) + + def test_enabled_unconfigured_pc_is_independently_disabled(self) -> None: + registry = PublishingProviderRegistry( + config=FakeConfig(True, True, False) + ) + + info = registry.get_info("planetary_computer") + + self.assertTrue(info.isEnabled) + self.assertFalse(info.isConfigured) + + def test_enabled_configured_pc_invokes_injected_factory(self) -> None: + provider = object() + registry = PublishingProviderRegistry( + config=FakeConfig(True, True, True), + factories={"planetary_computer": lambda: provider}, + ) + + resolved = registry.resolve("planetary_computer") + + self.assertIs(resolved, provider) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/publishing/test_repository.py b/hastelib/tests/core/publishing/test_repository.py new file mode 100644 index 00000000..c8131e28 --- /dev/null +++ b/hastelib/tests/core/publishing/test_repository.py @@ -0,0 +1,427 @@ +import threading +import unittest +import uuid +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from copy import deepcopy +from typing import Dict, Iterator, Tuple +from unittest.mock import Mock + +from hastegeo.core.models.publishing import PublishedDataset, PublishStatus +from hastegeo.core.publishing.repository import ( + PublishedDatasetsExistError, + PublishingConflictError, + PublishingRepository, + StaleRevisionError, +) + + +class FakeMetadataProcessor: + records: Dict[Tuple[str, str], dict] = {} + records_lock = threading.Lock() + + def __init__( + self, + data_type: str, + partition_key: str = None, + config=None, + ) -> None: + self.partition_key = partition_key or "" + + def save( + self, key: str, metadata: dict, data_format: str = "json" + ) -> None: + with self.records_lock: + self.records[(self.partition_key, key)] = deepcopy(metadata) + + def load(self, key: str, data_format: str = "json") -> dict: + with self.records_lock: + try: + return deepcopy(self.records[(self.partition_key, key)]) + except KeyError as error: + raise FileNotFoundError(key) from error + + def load_all(self, data_format: str = "json") -> list[dict]: + with self.records_lock: + return [deepcopy(record) for record in self.records.values()] + + def load_all_from_partition(self, data_format: str = "json") -> list[dict]: + with self.records_lock: + return [ + deepcopy(record) + for (partition, _), record in self.records.items() + if partition == self.partition_key + ] + + def delete(self, key: str, data_format: str = "json") -> None: + with self.records_lock: + del self.records[(self.partition_key, key)] + + def load_page( + self, + page, + page_size, + target=None, + status=None, + **kwargs, + ): + records = ( + self.load_all_from_partition() + if self.partition_key + else self.load_all() + ) + if target: + records = [ + record for record in records if record["target"] == target + ] + if status: + records = [ + record for record in records if record["status"] == status + ] + start = (page - 1) * page_size + return records[start : start + page_size], len(records) + + def load_bounded(self, max_records, data_format="json"): + records = self.load_all(data_format=data_format) + if len(records) > max_records: + raise ValueError("record limit") + return records + + +class FakeLeaseCoordinator: + def __init__(self) -> None: + self.guard = threading.Lock() + self.locks: Dict[Tuple[str, str], threading.Lock] = {} + + @contextmanager + def acquire( + self, + project_id: str, + dataset_id: str, + lease_duration: int = 60, + **kwargs, + ) -> Iterator[None]: + key = (project_id, dataset_id) + with self.guard: + operation_lock = self.locks.setdefault(key, threading.Lock()) + with operation_lock: + yield + + +class FakeConfig: + publishing_config = {"publishing_enabled": True} + + @staticmethod + def get_metadata_types(): + class Types: + class PUBLISHED_DATASET: + value = "published_dataset" + + return Types + + +class TestPublishingRepository(unittest.TestCase): + def setUp(self) -> None: + FakeMetadataProcessor.records = {} + self.repository = PublishingRepository( + config=FakeConfig(), + processor_factory=FakeMetadataProcessor, + lease_coordinator=FakeLeaseCoordinator(), + ) + self.dataset = PublishedDataset( + datasetId=uuid.uuid4(), + requestId=uuid.uuid4(), + requestFingerprint="a" * 64, + name="Dataset", + projectId=uuid.uuid4(), + imageLayerId="layer", + modelId="model", + target="local", + status="PENDING", + publishedByUser="publisher", + createdDate="2026-08-06T00:00:00Z", + updatedDate="2026-08-06T00:00:00Z", + ) + + def test_concurrent_replays_create_one_record(self) -> None: + with ThreadPoolExecutor(max_workers=20) as executor: + results = list( + executor.map( + lambda _: self.repository.create_or_replay(self.dataset), + range(50), + ) + ) + + self.assertEqual(sum(created for _, created in results), 1) + self.assertEqual(len(self.repository.list_all()), 1) + + def test_concurrent_distinct_creates_preserve_fifty_records(self) -> None: + datasets = [ + self.dataset.model_copy( + update={ + "datasetId": uuid.uuid4(), + "requestId": uuid.uuid4(), + "requestFingerprint": f"{index:064x}", + } + ) + for index in range(50) + ] + + with ThreadPoolExecutor(max_workers=20) as executor: + results = list( + executor.map(self.repository.create_or_replay, datasets) + ) + + self.assertTrue(all(created for _, created in results)) + self.assertEqual(len(self.repository.list_all()), 50) + for dataset in datasets: + loaded = self.repository.load( + str(dataset.projectId), str(dataset.datasetId) + ) + self.assertEqual(loaded.requestId, dataset.requestId) + + def test_reused_request_id_with_changed_payload_conflicts(self) -> None: + self.repository.create_or_replay(self.dataset) + conflicting = self.dataset.model_copy( + update={"requestFingerprint": "b" * 64} + ) + + with self.assertRaises(PublishingConflictError): + self.repository.create_or_replay(conflicting) + + def test_update_rejects_stale_revision(self) -> None: + stored, _ = self.repository.create_or_replay(self.dataset) + changed = stored.model_copy( + update={"status": PublishStatus.IN_PROGRESS} + ) + updated = self.repository.update(changed, expected_revision=1) + + self.assertEqual(updated.revision, 2) + with self.assertRaises(StaleRevisionError): + self.repository.update(changed, expected_revision=1) + + def test_project_delete_action_is_blocked_when_dataset_exists( + self, + ) -> None: + self.repository.create_or_replay(self.dataset) + deleted = [] + + with self.assertRaises(PublishedDatasetsExistError): + self.repository.delete_project_if_unpublished( + str(self.dataset.projectId), lambda: deleted.append(True) + ) + + self.assertEqual(deleted, []) + + def test_project_delete_action_runs_under_empty_project_guard( + self, + ) -> None: + deleted = [] + + self.repository.delete_project_if_unpublished( + str(self.dataset.projectId), lambda: deleted.append(True) + ) + + self.assertEqual(deleted, [True]) + + def test_disabled_project_guard_does_not_require_lease_storage( + self, + ) -> None: + config = FakeConfig() + config.publishing_config = {"publishing_enabled": False} + repository = PublishingRepository( + config=config, + processor_factory=FakeMetadataProcessor, + ) + deleted = [] + + repository.delete_project_if_unpublished( + str(self.dataset.projectId), lambda: deleted.append(True) + ) + + self.assertEqual(deleted, [True]) + + def test_list_orders_by_published_date_with_created_fallback(self) -> None: + older = self.dataset.model_copy( + update={ + "createdDate": "2026-08-06T02:00:00Z", + "publishedDate": None, + } + ) + republished = self.dataset.model_copy( + update={ + "datasetId": uuid.uuid4(), + "requestId": uuid.uuid4(), + "createdDate": "2026-08-06T01:00:00Z", + "publishedDate": "2026-08-06T03:00:00Z", + } + ) + self.repository.create_or_replay(older) + self.repository.create_or_replay(republished) + + records = self.repository.list_all() + + self.assertEqual(records[0].datasetId, republished.datasetId) + + def test_list_page_returns_total_and_exact_search_results(self) -> None: + for index, name in enumerate( + ("Alpha damage", "Bravo flood", "Charlie") + ): + dataset = self.dataset.model_copy( + update={ + "datasetId": uuid.uuid4(), + "requestId": uuid.uuid4(), + "requestFingerprint": f"{index + 10:064x}", + "name": name, + } + ) + self.repository.create_or_replay(dataset) + + records, total = self.repository.list_page( + page=1, + page_size=2, + search="flood", + sort_key="name", + sort_direction="asc", + ) + + self.assertEqual(total, 1) + self.assertEqual([record.name for record in records], ["Bravo flood"]) + + def test_list_page_rejects_unbounded_page_size(self) -> None: + with self.assertRaises(ValueError): + self.repository.list_page(page=1, page_size=101) + + def test_blob_list_page_uses_indexed_storage_path(self) -> None: + config = FakeConfig() + config.storage_type = "blob" + repository = PublishingRepository( + config=config, + processor_factory=FakeMetadataProcessor, + lease_coordinator=FakeLeaseCoordinator(), + ) + repository.create_or_replay(self.dataset) + + records, total = repository.list_page( + page=1, + page_size=20, + target=self.dataset.target, + status=self.dataset.status, + ) + + self.assertEqual(total, 1) + self.assertEqual(records[0].datasetId, self.dataset.datasetId) + + def test_blob_search_rejects_catalog_above_scan_limit(self) -> None: + config = FakeConfig() + config.storage_type = "blob" + processor = Mock() + processor.load_page.return_value = ([], 1001) + repository = PublishingRepository( + config=config, + processor_factory=Mock(return_value=processor), + lease_coordinator=FakeLeaseCoordinator(), + ) + + with self.assertRaisesRegex(ValueError, "1,000"): + repository.list_page( + page=1, + page_size=20, + search="damage", + ) + + processor.load_page.assert_called_once() + + def test_blob_search_uses_bounded_unfiltered_page(self) -> None: + config = FakeConfig() + config.storage_type = "blob" + processor = Mock() + processor.load_page.return_value = ( + [self.dataset.model_dump(mode="json")], + 1, + ) + repository = PublishingRepository( + config=config, + processor_factory=Mock(return_value=processor), + lease_coordinator=FakeLeaseCoordinator(), + ) + + records, total = repository.list_page( + page=1, + page_size=20, + target=self.dataset.target, + search="dataset", + sort_key="name", + ) + + self.assertEqual(total, 1) + self.assertEqual(records[0].datasetId, self.dataset.datasetId) + processor.load_page.assert_called_once_with( + page=1, + page_size=1000, + project_id=None, + max_records=1000, + ) + processor.load_all.assert_not_called() + + def test_blob_reconciliation_uses_bounded_page(self) -> None: + config = FakeConfig() + config.storage_type = "blob" + processor = Mock() + processor.load_page.return_value = ( + [self.dataset.model_dump(mode="json")], + 1, + ) + repository = PublishingRepository( + config=config, + processor_factory=Mock(return_value=processor), + lease_coordinator=FakeLeaseCoordinator(), + ) + + records = repository.list_for_reconciliation() + + self.assertEqual(records[0].datasetId, self.dataset.datasetId) + processor.load_page.assert_called_once_with( + page=1, + page_size=1000, + max_records=1000, + ) + processor.load_all.assert_not_called() + + def test_non_blob_reconciliation_uses_bounded_read(self) -> None: + processor = Mock() + processor.load_bounded.return_value = [ + self.dataset.model_dump(mode="json") + ] + repository = PublishingRepository( + config=FakeConfig(), + processor_factory=Mock(return_value=processor), + lease_coordinator=FakeLeaseCoordinator(), + ) + + records = repository.list_for_reconciliation() + + self.assertEqual(records[0].datasetId, self.dataset.datasetId) + processor.load_bounded.assert_called_once_with(max_records=1000) + processor.load_all.assert_not_called() + + def test_delete_with_revision_guard(self) -> None: + stored, _ = self.repository.create_or_replay(self.dataset) + + with self.assertRaises(StaleRevisionError): + self.repository.delete( + str(stored.projectId), + str(stored.datasetId), + expected_revision=2, + ) + self.repository.delete( + str(stored.projectId), + str(stored.datasetId), + expected_revision=1, + ) + with self.assertRaises(FileNotFoundError): + self.repository.load(str(stored.projectId), str(stored.datasetId)) + + +if __name__ == "__main__": + unittest.main() diff --git a/hastelib/tests/core/publishing/test_source.py b/hastelib/tests/core/publishing/test_source.py new file mode 100644 index 00000000..5c36f8aa --- /dev/null +++ b/hastelib/tests/core/publishing/test_source.py @@ -0,0 +1,371 @@ +import unittest +import uuid + +from hastegeo.core.models.publishing import ArtifactKind, PublishRequest +from hastegeo.core.publishing.source import ( + PublishingArtifactUnavailableError, + PublishingSourceNotEligibleError, + PublishingSourceNotFoundError, + PublishingSourceResolver, +) +from hastegeo.core.utils.metadata import MetadataUtils + + +class FakeTypes: + class PROJECT: + value = "project" + + class IMAGELAYER: + value = "imagelayer" + + class MODEL: + value = "model" + + +class FakeConfig: + artifact_storage_type = "local" + artifact_storage_config = {} + + @staticmethod + def get_metadata_types(): + return FakeTypes + + @staticmethod + def get_status_types(): + class Types: + class COMPLETED: + value = "Processed" + + return Types + + +class FakeMetadataProcessor: + records = {} + + def __init__( + self, + data_type: str, + partition_key: str = None, + config=None, + ) -> None: + self.data_type = data_type + self.partition_key = partition_key + + def load(self, key: str) -> dict: + try: + return self.records[(self.data_type, self.partition_key, key)] + except KeyError as error: + raise FileNotFoundError(key) from error + + +class FakeArtifactStorage: + def __init__(self, artifacts: dict[str, int]) -> None: + self.artifacts = artifacts + + def resolve_artifact_path(self, location: str) -> str: + return location.removeprefix("storage://") + + def artifact_exists(self, artifact_path: str) -> bool: + return artifact_path in self.artifacts + + def get_artifact_size(self, artifact_path: str) -> int: + return self.artifacts[artifact_path] + + def get_artifact_etag(self, artifact_path: str) -> str: + return f"etag-{artifact_path}-{self.artifacts[artifact_path]}" + + +class TestPublishingSourceResolver(unittest.TestCase): + def setUp(self) -> None: + self.project_id = str(uuid.uuid4()) + self.layer_id = "layer-1" + self.model_id = "42" + self.project_prefix = MetadataUtils.hash_string(self.project_id) + self.imagery_task_id = "img-task" + self.inference_task_id = "inf-task" + self.paths = { + "damage": ( + f"{self.project_prefix}/{self.inference_task_id}/" + "predicted_damage_Model.gpkg" + ), + "mask": ( + f"{self.project_prefix}/{self.imagery_task_id}/" + f"valid_area_mask_{self.project_id}_{self.layer_id}.geojson" + ), + "footprints": ( + f"{self.project_prefix}/{self.imagery_task_id}/" + f"building_footprints_{self.project_id}_{self.layer_id}.gpkg" + ), + "image": ( + f"{self.project_prefix}/{self.imagery_task_id}/" + f"processed_imagery_post_event_cog_{self.project_id}_" + f"{self.layer_id}.tif" + ), + } + FakeMetadataProcessor.records = { + ("project", self.project_id, self.project_id): { + "projectId": self.project_id, + "name": "Project", + }, + ("imagelayer", self.project_id, self.layer_id): { + "imageLayerId": self.layer_id, + "projectId": self.project_id, + "name": "Layer", + "status": "Processed", + "preprocessJob": { + "taskId": self.imagery_task_id, + "jobId": "imagery-job", + "imageLayerId": self.layer_id, + "projectId": self.project_id, + "status": "Processed", + }, + "validAreaMaskUrl": f"storage://{self.paths['mask']}", + "buildingFootprintsUrl": ( + f"storage://{self.paths['footprints']}" + ), + "postEventProcessedImageryUrl": ( + f"storage://{self.paths['image']}" + ), + }, + ("model", self.project_id, self.model_id): { + "modelId": self.model_id, + "projectId": self.project_id, + "imageLayerId": self.layer_id, + "name": "Model", + "inferenceStatus": "Processed", + "inferenceOutputPath": ( + f"{self.project_prefix}/{self.inference_task_id}" + ), + "currentInferenceTaskId": self.inference_task_id, + "inferenceJobs": [ + { + "jobId": "inference-job", + "taskId": self.inference_task_id, + "modelId": self.model_id, + "projectId": self.project_id, + "status": "Processed", + } + ], + "gpkgUrl": f"storage://{self.paths['damage']}", + }, + } + self.artifact_storage = FakeArtifactStorage( + dict(zip(self.paths.values(), (10, 20, 30, 40))) + ) + self.resolver = PublishingSourceResolver( + config=FakeConfig(), + processor_factory=FakeMetadataProcessor, + artifact_storage=self.artifact_storage, + ) + + def build_request(self, artifacts: list[str]) -> PublishRequest: + return PublishRequest( + requestId=uuid.uuid4(), + projectId=self.project_id, + imageLayerId=self.layer_id, + modelId=self.model_id, + name="Dataset", + target="local", + artifacts=artifacts, + ) + + def test_options_return_only_verified_artifacts(self) -> None: + del self.artifact_storage.artifacts[self.paths["image"]] + + options = self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + + self.assertEqual(options.defaultName, "Project – Layer") + self.assertEqual( + {artifact.kind for artifact in options.availableArtifacts}, + { + ArtifactKind.GPKG, + ArtifactKind.VALID_MASK, + ArtifactKind.FOOTPRINTS, + }, + ) + + def test_bundle_separates_selected_and_supporting_artifacts(self) -> None: + request = self.build_request(["gpkg"]) + + bundle = self.resolver.resolve_bundle( + request, supporting_kinds=[ArtifactKind.VALID_MASK] + ) + + self.assertEqual( + [artifact.kind for artifact in bundle.selectedArtifacts], + [ArtifactKind.GPKG], + ) + self.assertEqual( + [artifact.kind for artifact in bundle.supportingArtifacts], + [ArtifactKind.VALID_MASK], + ) + + def test_bundle_rejects_unavailable_requested_artifact(self) -> None: + del self.artifact_storage.artifacts[self.paths["footprints"]] + request = self.build_request(["footprints"]) + + with self.assertRaisesRegex( + PublishingArtifactUnavailableError, "footprints" + ): + self.resolver.resolve_bundle(request) + + def test_options_reject_incomplete_inference(self) -> None: + FakeMetadataProcessor.records[ + ("model", self.project_id, self.model_id) + ]["inferenceStatus"] = "InProgress" + + with self.assertRaisesRegex( + PublishingSourceNotEligibleError, "Processed" + ): + self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + + def test_options_reject_missing_and_mismatched_sources(self) -> None: + project_key = ("project", self.project_id, self.project_id) + project = FakeMetadataProcessor.records.pop(project_key) + with self.assertRaises(PublishingSourceNotFoundError): + self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + FakeMetadataProcessor.records[project_key] = project + FakeMetadataProcessor.records[ + ("imagelayer", self.project_id, self.layer_id) + ]["projectId"] = str(uuid.uuid4()) + with self.assertRaisesRegex(FileNotFoundError, "does not belong"): + self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + + def test_options_reject_when_no_artifacts_exist(self) -> None: + self.artifact_storage.artifacts.clear() + + with self.assertRaisesRegex( + PublishingSourceNotEligibleError, "no available" + ): + self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + + def test_bundle_rejects_missing_supporting_artifact(self) -> None: + del self.artifact_storage.artifacts[self.paths["mask"]] + request = self.build_request(["gpkg"]) + + with self.assertRaisesRegex( + PublishingArtifactUnavailableError, "supporting" + ): + self.resolver.resolve_bundle( + request, supporting_kinds=[ArtifactKind.VALID_MASK] + ) + + def test_options_reject_tampered_same_container_source_paths(self) -> None: + model_record = FakeMetadataProcessor.records[ + ("model", self.project_id, self.model_id) + ] + model_record["gpkgUrl"] = "storage://users_acl.json" + self.artifact_storage.artifacts["users_acl.json"] = 100 + + layer_record = FakeMetadataProcessor.records[ + ("imagelayer", self.project_id, self.layer_id) + ] + sibling_path = self.paths["mask"].replace( + self.project_prefix, + MetadataUtils.hash_string(str(uuid.uuid4())), + 1, + ) + layer_record["validAreaMaskUrl"] = f"storage://{sibling_path}" + self.artifact_storage.artifacts[sibling_path] = 100 + + options = self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + + self.assertNotIn( + ArtifactKind.GPKG, + {artifact.kind for artifact in options.availableArtifacts}, + ) + self.assertNotIn( + ArtifactKind.VALID_MASK, + {artifact.kind for artifact in options.availableArtifacts}, + ) + + def test_options_reject_output_not_owned_by_completed_current_job( + self, + ) -> None: + model_record = FakeMetadataProcessor.records[ + ("model", self.project_id, self.model_id) + ] + model_record[ + "inferenceOutputPath" + ] = f"{self.project_prefix}/forged-task" + forged_path = ( + f"{self.project_prefix}/forged-task/predicted_damage_Model.gpkg" + ) + model_record["gpkgUrl"] = f"storage://{forged_path}" + self.artifact_storage.artifacts[forged_path] = 10 + + options = self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + + self.assertNotIn( + ArtifactKind.GPKG, + {artifact.kind for artifact in options.availableArtifacts}, + ) + + def test_embedding_gpkg_requires_completed_owned_job(self) -> None: + model_record = FakeMetadataProcessor.records[ + ("model", self.project_id, self.model_id) + ] + embedding_path = ( + f"{self.project_prefix}/building_predictions_{self.model_id}.gpkg" + ) + model_record.update( + { + "modelType": "embedding", + "status": "Processed", + "gpkgUrl": f"storage://{embedding_path}", + "inferenceOutputPath": None, + "currentInferenceTaskId": None, + "inferenceJobs": [], + "embeddingJob": None, + } + ) + self.artifact_storage.artifacts[embedding_path] = 10 + + without_job = self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + self.assertNotIn( + ArtifactKind.GPKG, + {artifact.kind for artifact in without_job.availableArtifacts}, + ) + + model_record["embeddingJob"] = { + "jobId": "embedding-job", + "taskId": "embedding-task", + "modelId": self.model_id, + "projectId": self.project_id, + "status": "Processed", + } + with_job = self.resolver.resolve_options( + self.project_id, self.layer_id, self.model_id + ) + self.assertIn( + ArtifactKind.GPKG, + {artifact.kind for artifact in with_job.availableArtifacts}, + ) + + def test_ensure_project_exists_maps_missing_record(self) -> None: + del FakeMetadataProcessor.records[ + ("project", self.project_id, self.project_id) + ] + + with self.assertRaises(FileNotFoundError): + self.resolver.ensure_project_exists(self.project_id) + + +if __name__ == "__main__": + unittest.main() From 2410b3a09d4ad708385148a6194f2dad95f781dc Mon Sep 17 00:00:00 2001 From: Anthony Cintron Roman Date: Sun, 9 Aug 2026 10:22:47 -0400 Subject: [PATCH 02/17] feat(publishing): API routes + queue trigger (Local publishing) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bucket A only. Adds to api/hastefuncapi/function_app.py: publishing imports, active-caller/authorization helpers, the DeleteProject published-dataset guard, and 7 publishing routes (GetPublishingProviders, GetPublishDatasetOptions, GetPublishedDatasets, GetPublishedDataset, PutPublishDatasetQueueMessage, PutRetryPublishedDatasetQueueMessage, DeletePublishedDataset). All existing endpoints reverted to main's versions — the Bucket-B endpoint hardening (IDOR checks, strict *Request DTOs, server-managed-field enforcement, assessment size caps) is deferred to a separate -hardening PR. Tests: 19 publishing route tests + 6 queue-handler tests pass. Existing core suite: 191 pass (3 failures are pre-existing on main — clipBbox mocks). --- api/hastefuncapi/function_app.py | 475 ++++++++++++- api/hastefuncapi/tests/__init__.py | 0 .../tests/test_publishing_routes.py | 627 ++++++++++++++++++ api/hastefuncqueues/tests/__init__.py | 0 .../tests/test_publishing_handlers.py | 108 +++ 5 files changed, 1206 insertions(+), 4 deletions(-) create mode 100644 api/hastefuncapi/tests/__init__.py create mode 100644 api/hastefuncapi/tests/test_publishing_routes.py create mode 100644 api/hastefuncqueues/tests/__init__.py create mode 100644 api/hastefuncqueues/tests/test_publishing_handlers.py diff --git a/api/hastefuncapi/function_app.py b/api/hastefuncapi/function_app.py index 7683072b..61a16d18 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,41 @@ 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 +86,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 +120,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 +214,161 @@ 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}, 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 + return {"id": str(caller_id).lower(), "roles": roles}, 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 +961,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 +990,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 +1678,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 +1792,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 +1986,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 +2040,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") @@ -2738,6 +2952,8 @@ async def PutArtifactsZipQueueMessage( async def PutCancelModelQueueMessage( req: func.HttpRequest, ) -> func.HttpResponse: + from hastegeo.core.processors.train import TrainPreprocessor + logger.info( "PutCancelModelQueueMessage HTTP trigger function processed a request." ) @@ -4215,3 +4431,254 @@ 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"], + ) + 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/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/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() From c656bc62648c17275b6a1b67179c5a7e2265bb6a Mon Sep 17 00:00:00 2001 From: Anthony Cintron Roman Date: Sun, 9 Aug 2026 10:28:11 -0400 Subject: [PATCH 03/17] feat(publishing): Published Datasets UI (section, dialog, entry points) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bucket A UI, gated behind the publishingEnabled flag: - New: PublishedDatasets section, PublishedDatasetRow, PublishDatasetModal (name/description prefill, asset checklist, provider-driven target dropdown), util/publishing.js, util/assessmentSummary.js (extracted from AssessmentReportModal). - Additive wiring: App/AppContext (publishingEnabled + provider fetch), AppBody route + AppSidebar nav, "Publish dataset..." in ModelResultsButton + EmbeddingModelRow, Database icon. - api.js: apiPut/apiDelete now accept any 2xx (async 202 queue responses) — minimal additive change; the shared-helper ApiError rewrite is NOT taken (deferred to -hardening). Existing 200 callers unaffected. Validated: full `vite build` compiles cleanly (all imports resolve). Vitest unit tests + config are preserved on -sol/archive and land in a follow-up UI-test-infra slice (needs vitest devDeps + package-lock regen). --- ui/src/App.jsx | 14 +- ui/src/AppContext.jsx | 22 +- ui/src/Components/AppBody.jsx | 17 +- ui/src/Components/AppSidebar.jsx | 27 +- .../AssessmentReportModal.jsx | 52 +-- .../ProjectManagement/EmbeddingModelRow.jsx | 50 ++- .../ProjectManagement/ModelResultsButton.jsx | 48 ++- ui/src/Components/PublishDatasetModal.jsx | 394 ++++++++++++++++++ ui/src/Components/PublishedDatasetRow.jsx | 231 ++++++++++ ui/src/Components/PublishedDatasets.jsx | 298 +++++++++++++ ui/src/util/api.js | 9 +- ui/src/util/assessmentSummary.js | 44 ++ ui/src/util/icons.jsx | 11 + ui/src/util/publishing.js | 32 ++ 14 files changed, 1161 insertions(+), 88 deletions(-) create mode 100644 ui/src/Components/PublishDatasetModal.jsx create mode 100644 ui/src/Components/PublishedDatasetRow.jsx create mode 100644 ui/src/Components/PublishedDatasets.jsx create mode 100644 ui/src/util/assessmentSummary.js create mode 100644 ui/src/util/publishing.js diff --git a/ui/src/App.jsx b/ui/src/App.jsx index a77a92f2..a13de10c 100644 --- a/ui/src/App.jsx +++ b/ui/src/App.jsx @@ -13,8 +13,7 @@ import { DialogActions, } from "@fluentui/react-components"; import { AppContext } from "./AppContext"; -import { apiValidateUser, apiLogout, apiGet } from "./util/api"; -import { upsertUser } from "./AppHelper"; +import { apiValidateUser, apiGet } from "./util/api"; import { useTheme } from "./util/ThemeContext"; import { getPalette } from "./util/theme"; @@ -56,6 +55,16 @@ function App() { const validateUser = async () => { setIsLoading(true); await apiValidateUser(setAppParams); + try { + const publishing = await apiGet("GetPublishingProviders"); + setAppParams((previous) => ({ + ...previous, + publishingEnabled: !!publishing.publishingEnabled, + publishingProviders: publishing.providers || [], + })); + } catch (error) { + console.error("Error loading publishing capabilities:", error); + } setIsLoading(false); }; @@ -93,6 +102,7 @@ function App() { useEffect(() => { if (isMobileNav) { + // eslint-disable-next-line react-hooks/set-state-in-effect setNavCollapsed(true); } }, [isMobileNav]); diff --git a/ui/src/AppContext.jsx b/ui/src/AppContext.jsx index 1e454762..73ed8e8c 100644 --- a/ui/src/AppContext.jsx +++ b/ui/src/AppContext.jsx @@ -1,16 +1,13 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import { createContext, useState } from "react"; +import { createContext, useCallback, useState } from "react"; import PropTypes from "prop-types"; +// eslint-disable-next-line react-refresh/only-export-components export const AppContext = createContext(); export const AppProvider = ({ children }) => { - AppProvider.propTypes = { - children: PropTypes.node.isRequired, - }; - const guidedTourList = [ { name: "dashboardGuide", @@ -257,6 +254,9 @@ export const AppProvider = ({ children }) => { userRoles: null, userSettings: null, userStatus: null, + identityId: null, + publishingEnabled: false, + publishingProviders: [], isLoading: false, loadingMessage: "", appHeaderRightButtons: [], @@ -273,20 +273,20 @@ export const AppProvider = ({ children }) => { })); } - function setIsLoading(isLoading, message = "Loading...") { + const setIsLoading = useCallback((isLoading, message = "Loading...") => { setAppParams((prevParams) => ({ ...prevParams, isLoading: isLoading, loadingMessage: message, })); - } + }, []); - function setAppHeaderRightButtons(appHeaderRightButtons) { + const setAppHeaderRightButtons = useCallback((appHeaderRightButtons) => { setAppParams((prevParams) => ({ ...prevParams, appHeaderRightButtons: appHeaderRightButtons, })); - } + }, []); function setCurrentTour(currentTourName) { const currentTourTemp = guidedTourList.find( @@ -341,3 +341,7 @@ export const AppProvider = ({ children }) => { ); }; + +AppProvider.propTypes = { + children: PropTypes.node.isRequired, +}; diff --git a/ui/src/Components/AppBody.jsx b/ui/src/Components/AppBody.jsx index b5a98f6b..3f0edb5f 100644 --- a/ui/src/Components/AppBody.jsx +++ b/ui/src/Components/AppBody.jsx @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. -import { Route, Routes, useLocation } from "react-router-dom"; -import { useContext, useEffect } from "react"; +import { Route, Routes } from "react-router-dom"; +import { useContext } from "react"; import Loading from "./OtherComponents/Loading"; import Error404 from "./Error404"; @@ -14,28 +14,27 @@ import BuildingValidation from "./BuildingValidation/BuildingValidation"; import InteractiveLabeler from "./InteractiveLabeler/InteractiveLabeler"; import Visualizer from "./Visualizer/Visualizer"; import ModelCatalog from "./ModelCatalog"; +import PublishedDatasets from "./PublishedDatasets"; import AdminUsers from "./AdminUsers"; import AdminSourceTypes from "./AdminSourceTypes"; import AdminLabelingTool from "./AdminLabelingTool"; import CreateEditImageLayerForm from "./CreateEditImageLayerForm"; import HelpDocs from "./HelpDocs"; -import { apiValidateUser, apiLogout } from "../util/api"; import PropType from "prop-types"; import { AppContext } from "../AppContext"; const AppBody = ({ setModalComponent }) => { - AppBody.propTypes = { - setModalComponent: PropType.func.isRequired, - }; - const { appParams } = useContext(AppContext); return (
{appParams.isLoading && } + {appParams.userRoles !== null && appParams.publishingEnabled && ( + } /> + )} {appParams.userRoles !== null && (appParams.userRoles.includes("administrators") || appParams.userRoles.includes("contributors")) && ( @@ -107,4 +106,8 @@ const AppBody = ({ setModalComponent }) => { ); }; +AppBody.propTypes = { + setModalComponent: PropType.func.isRequired, +}; + export default AppBody; diff --git a/ui/src/Components/AppSidebar.jsx b/ui/src/Components/AppSidebar.jsx index a34805cd..d71bac93 100644 --- a/ui/src/Components/AppSidebar.jsx +++ b/ui/src/Components/AppSidebar.jsx @@ -14,14 +14,6 @@ const AppSidebar = ({ open, onItemSelected, }) => { - AppSidebar.propTypes = { - setModalComponent: PropTypes.func.isRequired, - collapsed: PropTypes.bool, - mobile: PropTypes.bool, - open: PropTypes.bool, - onItemSelected: PropTypes.func, - }; - const { appParams } = useContext(AppContext); const navigate = useNavigate(); const location = useLocation(); @@ -82,6 +74,17 @@ const AppSidebar = ({ pathname.startsWith("/projects") || pathname.startsWith("/project/"), }, + ...(appParams.publishingEnabled + ? [ + { + key: "published-datasets", + label: "Published Datasets", + icon: "Database", + onClick: () => handleNavigate("/published-datasets"), + active: pathname.startsWith("/published-datasets"), + }, + ] + : []), ], }, ...(isAdmin @@ -167,4 +170,12 @@ const AppSidebar = ({ ); }; +AppSidebar.propTypes = { + setModalComponent: PropTypes.func.isRequired, + collapsed: PropTypes.bool, + mobile: PropTypes.bool, + open: PropTypes.bool, + onItemSelected: PropTypes.func, +}; + export default AppSidebar; diff --git a/ui/src/Components/BuildingValidation/AssessmentReportModal.jsx b/ui/src/Components/BuildingValidation/AssessmentReportModal.jsx index 4d8a66cc..1cc6f278 100644 --- a/ui/src/Components/BuildingValidation/AssessmentReportModal.jsx +++ b/ui/src/Components/BuildingValidation/AssessmentReportModal.jsx @@ -18,6 +18,7 @@ import { import { FluentIcon } from "../../util/icons"; import PropTypes from "prop-types"; import { apiGet } from "../../util/api"; +import { buildAssessmentSummary } from "../../util/assessmentSummary"; /* ── Theme-aware design tokens (follow light/dark via Fluent) ─── */ const tokens = { @@ -163,14 +164,6 @@ const AssessmentReportModal = ({ modelName, onDismiss, }) => { - AssessmentReportModal.propTypes = { - projectId: PropTypes.string.isRequired, - imageLayerId: PropTypes.string.isRequired, - modelId: PropTypes.string.isRequired, - modelName: PropTypes.string, - onDismiss: PropTypes.func.isRequired, - }; - const [report, setReport] = useState(null); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); @@ -194,7 +187,10 @@ const AssessmentReportModal = ({ }; useEffect(() => { + // State updates occur after the report request resolves. + // eslint-disable-next-line react-hooks/set-state-in-effect fetchReport(); + // eslint-disable-next-line react-hooks/exhaustive-deps }, [projectId, imageLayerId, modelId]); const preds = report?.predictions; @@ -203,37 +199,7 @@ const AssessmentReportModal = ({ const sample = report?.evaluationSample; const hasLabels = (report?.matched ?? 0) > 0; - function buildSummarySentence() { - if (!report || !preds) return ""; - let s = - `Out of a total of ${int(preds.total)} building footprints in the study area, ` + - `${int(preds.cloudy)} were obscured by clouds; of the remaining ` + - `${int(preds.knownNonCloudy)} non-cloudy footprints, the model ` + - `predicted that ${int(preds.predictedDamaged)} ` + - `(${preds.predictedDamagedPctOfKnown ?? 0}%) were damaged to some extent.`; - if (!hasLabels) { - return ( - s + - " No human validation labels are available for this image layer yet — " + - "labeling some via the Building Validation tool will populate the " + - "metrics and population estimate." - ); - } - s += - ` We independently labeled ${int(report.totalLabels)} footprints; ` + - `${int(report.sureLabels)} were sure-labeled. Estimated recall ` + - `${pct(metrics?.recall)} and precision ${pct(metrics?.precision)}.`; - if (pop && pop.N > 0) { - s += - ` Extrapolating to all ${int(pop.N)} buildings with area > ` + - `${pop.minAreaM2.toFixed(0)} m², we estimate ` + - `${int(pop.estimatedDamaged)} damaged buildings ` + - `(${pct(pop.pHat)}) with a 95% CI of ` + - `[${int(pop.ciLower)}, ${int(pop.ciUpper)}].`; - } - return s; - } - const summarySentence = buildSummarySentence(); + const summarySentence = buildAssessmentSummary(report); return ( { - EmbeddingModelRow.propTypes = { - model: PropTypes.object.isRequired, - projectId: PropTypes.string.isRequired, - imageLayerId: PropTypes.string.isRequired, - index: PropTypes.number.isRequired, - fetchProjectDetails: PropTypes.func.isRequired, - validationLabelCount: PropTypes.number, - mobile: PropTypes.bool, - }; - - const { setDialog, setIsLoading } = useContext(AppContext); + const { appParams, setDialog, setIsLoading } = useContext(AppContext); const navigate = useNavigate(); const [showValidationReport, setShowValidationReport] = useState(false); const [showAssessmentReport, setShowAssessmentReport] = useState(false); + const [showPublishDataset, setShowPublishDataset] = useState(false); const isProcessed = model.status === "Processed"; const hasPredictions = !!model.gpkgUrl; @@ -155,6 +147,18 @@ const EmbeddingModelRow = ({ disabled: !hasPredictions, onClick: () => setShowAssessmentReport(true), }, + ...(appParams.publishingEnabled + ? [ + { + key: "publishDataset", + text: "Publish dataset…", + icon: , + disabled: + model.inferenceStatus !== "Processed" || !hasPredictions, + onClick: () => setShowPublishDataset(true), + }, + ] + : []), ], }; @@ -204,6 +208,20 @@ const EmbeddingModelRow = ({ onDismiss={() => setShowAssessmentReport(false)} /> )} + {showPublishDataset && ( + setShowPublishDataset(false)} + onStarted={() => + setDialog( + "Publishing started", + "Track progress in Published Datasets.", + ) + } + /> + )} ); @@ -452,4 +470,14 @@ const EmbeddingModelRow = ({ ); }; +EmbeddingModelRow.propTypes = { + model: PropTypes.object.isRequired, + projectId: PropTypes.string.isRequired, + imageLayerId: PropTypes.string.isRequired, + index: PropTypes.number.isRequired, + fetchProjectDetails: PropTypes.func.isRequired, + validationLabelCount: PropTypes.number, + mobile: PropTypes.bool, +}; + export default EmbeddingModelRow; diff --git a/ui/src/Components/ProjectManagement/ModelResultsButton.jsx b/ui/src/Components/ProjectManagement/ModelResultsButton.jsx index e01d347d..d89fbd58 100644 --- a/ui/src/Components/ProjectManagement/ModelResultsButton.jsx +++ b/ui/src/Components/ProjectManagement/ModelResultsButton.jsx @@ -18,6 +18,7 @@ import { AppContext } from "../../AppContext"; import ModelResultsStatusIndicator from "../OtherComponents/ModelResultsStatusIndicator"; import ValidationReportModal from "../BuildingValidation/ValidationReportModal"; import AssessmentReportModal from "../BuildingValidation/AssessmentReportModal"; +import PublishDatasetModal from "../PublishDatasetModal"; function formatFileSize(bytes) { @@ -31,18 +32,11 @@ function formatFileSize(bytes) { const ModelResultsButton = ({ model, projectId, imageLayerId, index, validationLabelCount }) => { - ModelResultsButton.propTypes = { - model: PropTypes.object.isRequired, - projectId: PropTypes.string.isRequired, - imageLayerId: PropTypes.string.isRequired, - index: PropTypes.number.isRequired, - validationLabelCount: PropTypes.number, - }; - - const { setDialog } = useContext(AppContext); + const { appParams, setDialog } = useContext(AppContext); const navigate = useNavigate(); const [showValidationReport, setShowValidationReport] = useState(false); const [showAssessmentReport, setShowAssessmentReport] = useState(false); + const [showPublishDataset, setShowPublishDataset] = useState(false); function evaluateViewResultsButtonState(model) { // Results button must be enabled if inference jobs exist and status is processed @@ -67,7 +61,7 @@ const ModelResultsButton = ({ model, projectId, imageLayerId, index, validationL ); } fileDownload(url, setDialog); - } catch (error) { + } catch { setDialog({ title: "Download Error", message: "An error occurred while downloading. Please try again later.", @@ -146,6 +140,18 @@ const ModelResultsButton = ({ model, projectId, imageLayerId, index, validationL disabled: model.inferenceStatus !== "Processed", onClick: () => setShowAssessmentReport(true), }, + ...(appParams.publishingEnabled + ? [ + { + key: "publishDataset", + text: "Publish dataset…", + icon: , + disabled: + model.inferenceStatus !== "Processed" || !model.gpkgUrl, + onClick: () => setShowPublishDataset(true), + }, + ] + : []), ], }); @@ -204,8 +210,30 @@ const ModelResultsButton = ({ model, projectId, imageLayerId, index, validationL onDismiss={() => setShowAssessmentReport(false)} /> )} + {showPublishDataset && ( + setShowPublishDataset(false)} + onStarted={() => + setDialog( + "Publishing started", + "Track progress in Published Datasets.", + ) + } + /> + )} ); }; +ModelResultsButton.propTypes = { + model: PropTypes.object.isRequired, + projectId: PropTypes.string.isRequired, + imageLayerId: PropTypes.string.isRequired, + index: PropTypes.number.isRequired, + validationLabelCount: PropTypes.number, +}; + export default ModelResultsButton; diff --git a/ui/src/Components/PublishDatasetModal.jsx b/ui/src/Components/PublishDatasetModal.jsx new file mode 100644 index 00000000..a801052f --- /dev/null +++ b/ui/src/Components/PublishDatasetModal.jsx @@ -0,0 +1,394 @@ +import { useEffect, useId, useRef, useState } from "react"; +import { + Button, + Checkbox, + Dialog, + DialogActions, + DialogBody, + DialogContent, + DialogSurface, + DialogTitle, + Dropdown, + Field, + Input, + Label, + MessageBar, + MessageBarBody, + Option, + Spinner, + Text, + Textarea, + makeStyles, + tokens, +} from "@fluentui/react-components"; +import PropTypes from "prop-types"; +import { v4 as uuidv4 } from "uuid"; + +import { apiGet, apiPut } from "../util/api"; +import { buildAssessmentSummary } from "../util/assessmentSummary"; +import { FluentIcon } from "../util/icons"; +import { selectSupportedArtifacts } from "../util/publishing"; + + +const ARTIFACT_LABELS = { + gpkg: "Damage GeoPackage (.gpkg)", + valid_mask: "Valid-area mask (.geojson)", + footprints: "Building footprints (.gpkg)", + processed_cog: "Processed image (COG .tif)", +}; + +const useStyles = makeStyles({ + surface: { + width: "min(680px, 94vw)", + maxWidth: "94vw", + }, + content: { + display: "flex", + flexDirection: "column", + gap: tokens.spacingVerticalL, + }, + title: { + display: "flex", + alignItems: "center", + gap: tokens.spacingHorizontalS, + }, + assets: { + display: "flex", + flexDirection: "column", + gap: tokens.spacingVerticalXS, + }, + assetsField: { + display: "flex", + flexDirection: "column", + gap: tokens.spacingVerticalXS, + }, + assetsValidation: { + color: tokens.colorStatusDangerForeground1, + }, + assetLabel: { + display: "flex", + justifyContent: "space-between", + gap: tokens.spacingHorizontalM, + width: "100%", + }, + size: { + color: tokens.colorNeutralForeground3, + whiteSpace: "nowrap", + }, + loading: { + minHeight: "180px", + display: "flex", + alignItems: "center", + justifyContent: "center", + }, +}); + +function formatFileSize(bytes) { + if (bytes == null) return ""; + if (bytes < 1024) return `${bytes} B`; + if (bytes < 1024 ** 2) return `${Math.round(bytes / 1024)} KB`; + if (bytes < 1024 ** 3) return `${(bytes / 1024 ** 2).toFixed(1)} MB`; + return `${(bytes / 1024 ** 3).toFixed(1)} GB`; +} + +const PublishDatasetModal = ({ + projectId, + imageLayerId, + modelId, + onDismiss, + onStarted, +}) => { + const styles = useStyles(); + const assetsLabelId = useId(); + const assetsValidationId = useId(); + const descriptionTouched = useRef(false); + const requestId = useRef(uuidv4()); + const [loading, setLoading] = useState(true); + const [descriptionLoading, setDescriptionLoading] = useState(true); + const [options, setOptions] = useState(null); + const [providers, setProviders] = useState([]); + const [name, setName] = useState(""); + const [description, setDescription] = useState(""); + const [target, setTarget] = useState(""); + const [selectedArtifacts, setSelectedArtifacts] = useState([]); + const [submitting, setSubmitting] = useState(false); + const [error, setError] = useState(""); + + useEffect(() => { + let cancelled = false; + const query = + `projectId=${encodeURIComponent(projectId)}` + + `&imageLayerId=${encodeURIComponent(imageLayerId)}` + + `&modelId=${encodeURIComponent(modelId)}`; + + Promise.all([ + apiGet(`GetPublishDatasetOptions?${query}`), + apiGet("GetPublishingProviders"), + ]) + .then(([optionsResponse, providerResponse]) => { + if (cancelled) return; + const resolvedOptions = optionsResponse.publishDatasetOptions; + const resolvedProviders = providerResponse.providers || []; + const initialProvider = + resolvedProviders.find( + (provider) => + provider.id === "local" && + provider.isEnabled && + provider.isConfigured, + ) || + resolvedProviders.find( + (provider) => provider.isEnabled && provider.isConfigured, + ); + setOptions(resolvedOptions); + setProviders(resolvedProviders); + setName(resolvedOptions.defaultName || ""); + if (initialProvider) { + setTarget(initialProvider.id); + setSelectedArtifacts( + selectSupportedArtifacts(resolvedOptions, initialProvider), + ); + } + }) + .catch((loadError) => { + if (!cancelled) setError(loadError.message || "Unable to load publishing options."); + }) + .finally(() => { + if (!cancelled) setLoading(false); + }); + + apiGet(`GetAssessmentReport?${query}`) + .then((report) => { + if (cancelled || descriptionTouched.current) return; + setDescription(buildAssessmentSummary(report)); + }) + .catch(() => {}) + .finally(() => { + if (!cancelled) setDescriptionLoading(false); + }); + + return () => { + cancelled = true; + }; + }, [projectId, imageLayerId, modelId]); + + const selectedProvider = providers.find( + (provider) => provider.id === target, + ); + const availableArtifacts = options?.availableArtifacts || []; + const supportedKinds = new Set( + selectedProvider?.supportedArtifactKinds || [], + ); + const effectiveSelectedArtifacts = selectedArtifacts.filter((kind) => + supportedKinds.has(kind), + ); + const canSubmit = + !loading && + !submitting && + !!name.trim() && + !!selectedProvider?.isEnabled && + !!selectedProvider?.isConfigured && + effectiveSelectedArtifacts.length > 0; + + function handleTargetChange(_, data) { + const providerId = data.optionValue || data.selectedOptions?.[0]; + const provider = providers.find((item) => item.id === providerId); + if (!provider || !provider.isEnabled || !provider.isConfigured) return; + setTarget(provider.id); + setSelectedArtifacts(selectSupportedArtifacts(options, provider)); + setError(""); + } + + function handleArtifactChange(kind, checked) { + setSelectedArtifacts((current) => + checked + ? [...new Set([...current, kind])] + : current.filter((value) => value !== kind), + ); + } + + async function handleSubmit(event) { + event.preventDefault(); + if (!canSubmit) return; + setSubmitting(true); + setError(""); + try { + const response = await apiPut("PutPublishDatasetQueueMessage", { + requestId: requestId.current, + projectId, + imageLayerId, + modelId, + name: name.trim(), + description: description.trim(), + target, + artifacts: effectiveSelectedArtifacts, + }); + onStarted?.(response.publishedDataset); + onDismiss(); + } catch (submitError) { + setError(submitError.message || "Unable to start publishing."); + } finally { + setSubmitting(false); + } + } + + return ( + !data.open && onDismiss()}> + +
+ + } + onClick={onDismiss} + /> + } + > + + + Publish dataset + + + + {loading ? ( +
+ +
+ ) : ( + <> + {error && ( + + {error} + + )} + + setName(data.value)} + disabled={submitting} + /> + + +