From bde3974c8f124e351cc781e37e9c170d4176807f Mon Sep 17 00:00:00 2001 From: JonnyTran Date: Tue, 18 Aug 2026 23:43:45 -0700 Subject: [PATCH 1/3] refactor: remove S3 native object versioning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Object versioning was enabled on every workspace bucket but was never a working feature, and in several places it actively misled callers: - `GET /api/v1/file/...` accepted `version_id` and never forwarded it to S3, so asking for an old version silently returned the current one. - The handler never emitted the `Version-Id`/`Is-Latest`/`Version-Tag` headers, because `FileObjectResponse.http_headers` had no caller. The frontend read exactly those headers, so `is_latest` was permanently null and the "Fetch latest schema" button could never render. - Deletion was modelled with delete markers that three separate call sites reconstructed from the same `etag == "" and size == 0` heuristic; the CLI integration test for upload/download/delete has been skipped ever since with "buckets with versioning enabled still list deleted files". It is un-skipped here. - `delete_bucket` drained with `list_objects_v2` and issued unversioned deletes, so noncurrent versions survived and `DeleteBucket` 500'd with BucketNotEmpty. - `SchemaVersion.object_version_id` was written best-effort and never read. Schema versioning is unaffected: a version's identity is `(dataset_id, version)`, allocated under a row lock so each version lands on its own key, and its integrity is `checksum`. Neither depended on the object store. `ListObjectsResponse.objects` becomes `list` rather than `Iterable` — the deleted `assign_version_id` validator was the only thing coercing it, and both `__len__` and `__getitem__` already assumed a list. Breaking: `include_version` and `version_id` query params are gone (422 on old SDKs), as are `is_latest`/`version_id`/`version_tag` on `ObjectMetadata` and the `--versions`/`--version-id` CLI flags. --- .../base/base-render-table/RenderTable.vue | 10 +- .../useSchemaTableViewModel.ts | 8 +- .../v1/domain/entities/IAnswer.ts | 5 - .../v1/domain/entities/table/Schema.ts | 17 +- .../get-extraction-schema-use-case.ts | 22 +-- .../repositories/SchemaRepository.test.ts | 1 - .../repositories/SchemaRepository.ts | 1 - ..._drop_schema_versions_object_version_id.py | 29 +++ .../extralit_server/api/handlers/v1/files.py | 8 +- .../api/handlers/v1/workspaces.py | 39 +--- .../extralit_server/api/schemas/v1/files.py | 56 +----- .../api/schemas/v1/schema_versions.py | 1 - .../src/extralit_server/contexts/files.py | 174 +++--------------- .../contexts/schema_versions.py | 1 - .../src/extralit_server/models/database.py | 7 +- extralit-server/tests/factories.py | 8 +- .../v1/datasets/test_schema_versions.py | 1 - .../tests/unit/api/handlers/v1/test_files.py | 66 +------ .../v1/workspaces/test_workspace_doctor.py | 8 +- .../unit/contexts/test_schema_versions.py | 1 - extralit/src/extralit/_api/_workspaces.py | 60 ++---- extralit/src/extralit/_models/_files.py | 4 - extralit/src/extralit/cli/files/delete.py | 5 +- extralit/src/extralit/cli/files/download.py | 3 +- extralit/src/extralit/cli/files/list.py | 7 +- extralit/src/extralit/cli/rich.py | 10 +- extralit/src/extralit/workspaces/_resource.py | 15 +- .../tests/integration/test_cli_commands.py | 1 - .../unit/api/test_workspace_files_api.py | 15 +- .../unit/api/test_workspace_schemas_api.py | 4 - 30 files changed, 106 insertions(+), 481 deletions(-) create mode 100644 extralit-server/src/extralit_server/alembic/versions/a1c9f4e2b703_drop_schema_versions_object_version_id.py diff --git a/extralit-frontend/components/base/base-render-table/RenderTable.vue b/extralit-frontend/components/base/base-render-table/RenderTable.vue index f0034a504..6dddda673 100644 --- a/extralit-frontend/components/base/base-render-table/RenderTable.vue +++ b/extralit-frontend/components/base/base-render-table/RenderTable.vue @@ -65,19 +65,13 @@ visibleCheckdropdown = !visibleCheckdropdown; " > - Check data ! + Check data @@ -175,7 +169,7 @@ export default { handler(newValidation, oldValidation) { if (this.isLoaded) { if (this.editable) - console.warn("Changes validation", this.tableJSON.schema.schemaName, this.tableJSON.schema.version_tag); + console.warn("Changes validation", this.tableJSON.schema.schemaName); this.tabulator?.setColumns(this.columnsConfig); this.validateTable(); } diff --git a/extralit-frontend/components/base/base-render-table/useSchemaTableViewModel.ts b/extralit-frontend/components/base/base-render-table/useSchemaTableViewModel.ts index 4c7eb27e2..7adaad381 100644 --- a/extralit-frontend/components/base/base-render-table/useSchemaTableViewModel.ts +++ b/extralit-frontend/components/base/base-render-table/useSchemaTableViewModel.ts @@ -43,15 +43,14 @@ export const useSchemaTableViewModel = (props: { ); const groupbyColumns = ref(refColumns.value || null); - const fetchValidation = async ({ latest = false }: { latest?: boolean } = {}) => { + const fetchValidation = async () => { var schemaName: string = props.tableJSON?.schema?.schemaName || props.tableJSON?.validation?.name; if (!schemaName) { return; } - var version_id: string = latest ? null : props.tableJSON.schema?.version_id; await waitForAsyncValue(() => dataset.workspaceName); - const [schema, fileMetadata] = await getSchema.fetch(dataset.workspaceName, schemaName, version_id); + const [schema, fileMetadata] = await getSchema.fetch(dataset.workspaceName, schemaName); // const schemaMetadataUpdate = { // ...props.tableJSON.schema, @@ -61,8 +60,7 @@ export const useSchemaTableViewModel = (props: { props.tableJSON?.schema?.fields || [], props.tableJSON?.schema?.primaryKey || [], fileMetadata, - schemaName, - version_id + schemaName ); if (!isEqual(props.tableJSON?.schema, schemaMetadataUpdate)) { diff --git a/extralit-frontend/v1/domain/entities/IAnswer.ts b/extralit-frontend/v1/domain/entities/IAnswer.ts index bd392912c..bf874667b 100644 --- a/extralit-frontend/v1/domain/entities/IAnswer.ts +++ b/extralit-frontend/v1/domain/entities/IAnswer.ts @@ -24,14 +24,9 @@ export type TableAnswer = { metadata?: { schemaName?: string; etag?: string; - version_id?: string; last_modified?: Date; - version_tag?: string; - is_latest?: boolean; }; schemaName?: string; - version_id?: string; - is_latest?: boolean; }; reference?: string; validation?: { diff --git a/extralit-frontend/v1/domain/entities/table/Schema.ts b/extralit-frontend/v1/domain/entities/table/Schema.ts index cafdc924a..ccbbf77a9 100644 --- a/extralit-frontend/v1/domain/entities/table/Schema.ts +++ b/extralit-frontend/v1/domain/entities/table/Schema.ts @@ -7,10 +7,7 @@ export interface DataFrameField { export interface FileMetadata { schemaName?: string; etag?: string; - version_id?: string; last_modified?: Date; - version_tag?: string; - is_latest?: boolean; } export class DataFrameSchema { @@ -18,24 +15,12 @@ export class DataFrameSchema { fields: DataFrameField[]; metadata?: FileMetadata; schemaName?: string; - version_id?: string; - version_tag?: string; - is_latest?: boolean; - constructor( - fields: DataFrameField[] = [], - primaryKey: string[] = [], - metadata?: FileMetadata, - schemaName?: string, - version_id?: string, - is_latest?: boolean - ) { + constructor(fields: DataFrameField[] = [], primaryKey: string[] = [], metadata?: FileMetadata, schemaName?: string) { this.fields = fields; this.primaryKey = primaryKey; this.metadata = metadata; this.schemaName = schemaName; - this.version_id = version_id; - this.is_latest = is_latest; } get fieldNames(): string[] { diff --git a/extralit-frontend/v1/domain/usecases/get-extraction-schema-use-case.ts b/extralit-frontend/v1/domain/usecases/get-extraction-schema-use-case.ts index 505c91bdf..9fe23876a 100644 --- a/extralit-frontend/v1/domain/usecases/get-extraction-schema-use-case.ts +++ b/extralit-frontend/v1/domain/usecases/get-extraction-schema-use-case.ts @@ -9,34 +9,16 @@ const FILES_API_ERRORS = { export class GetExtractionSchemaUseCase { constructor(private readonly axios: AxiosInstance) {} - async fetch( - workspaceName: string, - schemaName: string, - versionId?: string - ): Promise<[ValidationSchema, FileMetadata]> { + async fetch(workspaceName: string, schemaName: string): Promise<[ValidationSchema, FileMetadata]> { try { const url = `/v1/file/${workspaceName}/schemas/${schemaName}`; - const response = await this.axios.get(url, { - params: { - version_id: versionId, - }, - }); + const response = await this.axios.get(url); const headers = response.headers as AxiosHeaders; const schema = response.data; - let isLatest = null; - const headerValue = headers.get("is-latest"); - if (headerValue === "true") { - isLatest = true; - } else if (headerValue === "false") { - isLatest = false; - } const SchemaMetadata: FileMetadata = { schemaName, etag: headers.get("etag") as string, - version_id: headers.get("version-id") as string, - version_tag: headers.get("version-tag") as string, - is_latest: isLatest, last_modified: new Date((headers.get("last-modified") as string) || ""), }; diff --git a/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.test.ts b/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.test.ts index e3406c4a6..1f97f36a2 100644 --- a/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.test.ts +++ b/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.test.ts @@ -58,7 +58,6 @@ describe("SchemaRepository", () => { dataset_id: "s-1", version: 1, object_key: "k", - object_version_id: null, etag: "e", checksum: "c", parent_version_id: null, diff --git a/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.ts b/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.ts index bcc2d156d..0fe73117a 100644 --- a/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.ts +++ b/extralit-frontend/v1/infrastructure/repositories/SchemaRepository.ts @@ -29,7 +29,6 @@ interface BackendSchemaVersion { dataset_id: string; version: number; object_key: string; - object_version_id: string | null; etag: string; checksum: string; parent_version_id: string | null; diff --git a/extralit-server/src/extralit_server/alembic/versions/a1c9f4e2b703_drop_schema_versions_object_version_id.py b/extralit-server/src/extralit_server/alembic/versions/a1c9f4e2b703_drop_schema_versions_object_version_id.py new file mode 100644 index 000000000..3fda366ae --- /dev/null +++ b/extralit-server/src/extralit_server/alembic/versions/a1c9f4e2b703_drop_schema_versions_object_version_id.py @@ -0,0 +1,29 @@ +"""drop schema_versions.object_version_id + +Revision ID: a1c9f4e2b703 +Revises: 13da2d87e660 +Create Date: 2026-08-18 09:00:00.000000 + +The column recorded the S3 native object version of a published schema body. It was +written best-effort and never read: a version's identity is `(dataset_id, version)`, +allocated under a row lock so every version lands on its own key, and its integrity is +`checksum`. Object-store versioning is being removed wholesale, so the column goes with it. +""" + +import sqlalchemy as sa +from alembic import op + +revision = "a1c9f4e2b703" +down_revision = "13da2d87e660" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + with op.batch_alter_table("schema_versions") as batch_op: + batch_op.drop_column("object_version_id") + + +def downgrade() -> None: + with op.batch_alter_table("schema_versions") as batch_op: + batch_op.add_column(sa.Column("object_version_id", sa.Text(), nullable=True)) diff --git a/extralit-server/src/extralit_server/api/handlers/v1/files.py b/extralit-server/src/extralit_server/api/handlers/v1/files.py index 070324ad1..21a6416ec 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/files.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/files.py @@ -21,7 +21,6 @@ async def get_file( *, bucket: str, object: str, - version_id: str | None = None, range_header: str | None = Header(None, alias="range"), if_none_match: str | None = Header(None, alias="if-none-match"), s3_client=Depends(files.get_s3_client), @@ -158,8 +157,7 @@ async def list_objects_endpoint( *, bucket: str, prefix: str, - include_version=True, - recursive=True, + recursive: bool = True, start_after: str | None = None, s3_client=Depends(files.get_s3_client), current_user: User = Security(auth.get_optional_current_user), @@ -171,7 +169,6 @@ async def list_objects_endpoint( s3_client, bucket, prefix=prefix, - include_version=include_version, recursive=recursive, start_after=start_after, ) @@ -188,14 +185,13 @@ async def delete_files( *, bucket: str, object: str, - version_id: str | None = None, s3_client=Depends(files.get_s3_client), current_user: User = Security(auth.get_current_user), ): await authorize(current_user, FilePolicy.delete(bucket)) try: - await files.delete_object(s3_client, bucket, object, version_id=version_id) + await files.delete_object(s3_client, bucket, object) return {"message": "File deleted"} except HTTPException: raise diff --git a/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py b/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py index 14c3b8ec4..92b22e0ae 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py @@ -183,7 +183,6 @@ async def workspace_doctor( Checks: - S3 bucket exists (can auto-fix) - - Bucket has proper versioning policy (informational) - RQ worker pool connectivity (informational) """ await authorize(current_user, WorkspacePolicy.get(workspace_id)) @@ -233,39 +232,7 @@ async def workspace_doctor( ) ) - # Check 2: Bucket versioning policy - if bucket_exists or any(check.check_name == "s3_bucket" and check.fixed for check in checks): - versioning = await files.get_bucket_versioning(s3_client, workspace.name) - if versioning: - if versioning["status"] == "Enabled": - checks.append( - WorkspaceDoctorCheckResult( - check_name="bucket_versioning", - status="ok", - message=f"Bucket versioning is enabled (Status: {versioning['status']})", - fixed=False, - ) - ) - else: - checks.append( - WorkspaceDoctorCheckResult( - check_name="bucket_versioning", - status="warning", - message=f"Bucket versioning is not enabled (Status: {versioning['status']})", - fixed=False, - ) - ) - else: - checks.append( - WorkspaceDoctorCheckResult( - check_name="bucket_versioning", - status="warning", - message="Could not retrieve bucket versioning configuration", - fixed=False, - ) - ) - - # Check 3: RQ worker pool connectivity + # Check 2: RQ worker pool connectivity try: from extralit_server.jobs.queues import DEFAULT_QUEUE @@ -291,7 +258,7 @@ async def workspace_doctor( ) ) - # Check 4: Elasticsearch indexes for datasets (informational only) + # Check 3: Elasticsearch indexes for datasets (informational only) try: # Get datasets for this workspace from sqlalchemy import select @@ -345,7 +312,7 @@ async def workspace_doctor( ) ) - # Check 5: Database connections health with autofix + # Check 4: Database connections health with autofix try: import asyncio diff --git a/extralit-server/src/extralit_server/api/schemas/v1/files.py b/extralit-server/src/extralit_server/api/schemas/v1/files.py index 71ae72782..af681ccc2 100644 --- a/extralit-server/src/extralit_server/api/schemas/v1/files.py +++ b/extralit-server/src/extralit_server/api/schemas/v1/files.py @@ -1,5 +1,3 @@ -from collections import defaultdict -from collections.abc import Iterable from datetime import datetime from typing import Any @@ -11,12 +9,9 @@ class ObjectMetadata(BaseModel): bucket_name: str object_name: str last_modified: datetime | None = None - is_latest: bool | None = None etag: str | None = None size: int | None = None content_type: str | None = None - version_id: str | None = None - version_tag: str | None = None metadata: dict[str, Any] | None = None @field_validator("metadata", mode="before") @@ -29,38 +24,17 @@ def parse_metadata(cls, v): class ListObjectsResponse(BaseModel): - objects: Iterable[ObjectMetadata] = Field(default_factory=list) + objects: list[ObjectMetadata] = Field(default_factory=list) def __len__(self) -> int: - return len(self.objects) # type: ignore + return len(self.objects) def __getitem__(self, index) -> ObjectMetadata: - return self.objects[index] # type: ignore + return self.objects[index] def __iter__(self): return iter(self.objects) - @field_validator("objects") - def assign_version_id(cls, objects: list[ObjectMetadata]) -> list[ObjectMetadata]: - # Group objects by object_name - grouped_objects = defaultdict(list) - for obj in objects: - grouped_objects[obj.object_name].append(obj) - - # Assign version_id based on last_modified - for _object_name, object_list in grouped_objects.items(): - sorted_objects = sorted(object_list, key=lambda o: o.last_modified or datetime.min) - - for i, obj in enumerate(sorted_objects): - obj.version_tag = f"v{i + 1}" - if obj.is_latest is None: - obj.is_latest = i == len(sorted_objects) - 1 - - # Flatten the list of objects - objects = [obj for object_list in grouped_objects.values() for obj in object_list] - - return objects - class FileObjectResponse(BaseModel): # The S3 body is a streaming object whose concrete type varies by client/version @@ -69,31 +43,10 @@ class FileObjectResponse(BaseModel): # annotation permissive — a strict type triggers pydantic is_instance_of 422s. response: Any metadata: ObjectMetadata - versions: ListObjectsResponse | None class Config: arbitrary_types_allowed = True - @property - def version_tag(self) -> str | None: - if not self.metadata or not self.versions: - return "" - else: - for version in self.versions: - if version.version_id == self.metadata.version_id: - return version.version_tag - return "" - - @property - def is_latest(self) -> bool | None: - if not self.metadata or not self.versions: - return None - else: - for version in self.versions: - if version.version_id == self.metadata.version_id: - return version.is_latest - return None - @property def http_headers(self) -> dict[str, str]: if not self.metadata: @@ -102,12 +55,9 @@ def http_headers(self) -> dict[str, str]: headers = { "Content-Type": str(self.metadata.content_type) if self.metadata.content_type else "", "ETag": str(self.metadata.etag) if self.metadata.etag else "", - "Version-Id": str(self.metadata.version_id) if self.metadata.version_id else "", "Last-Modified": self.metadata.last_modified.strftime("%Y-%m-%dT%H:%M:%SZ") if self.metadata.last_modified else "", - "Is-Latest": str(self.is_latest).lower() if self.is_latest is not None else "", - "Version-Tag": self.version_tag, } headers = {key: value for key, value in headers.items() if value} return headers diff --git a/extralit-server/src/extralit_server/api/schemas/v1/schema_versions.py b/extralit-server/src/extralit_server/api/schemas/v1/schema_versions.py index c8e514488..15b12587f 100644 --- a/extralit-server/src/extralit_server/api/schemas/v1/schema_versions.py +++ b/extralit-server/src/extralit_server/api/schemas/v1/schema_versions.py @@ -27,7 +27,6 @@ class SchemaVersionRead(BaseModel): dataset_id: UUID version: int object_key: str - object_version_id: str | None etag: str checksum: str parent_version_id: UUID | None diff --git a/extralit-server/src/extralit_server/contexts/files.py b/extralit-server/src/extralit_server/contexts/files.py index e60db50a4..6c2308e17 100644 --- a/extralit-server/src/extralit_server/contexts/files.py +++ b/extralit-server/src/extralit_server/contexts/files.py @@ -14,9 +14,6 @@ from types_aiobotocore_s3.client import S3Client CHUNK_LENGTH_MB = 10 * 1024 * 1024 -# Layout datasets rewrite whole files on every commit; keeping their noncurrent versions grows -# without bound and nothing reads them (the canonical JSON is the history). -LAYOUT_NONCURRENT_EXPIRATION_DAYS = 1 _LOGGER = logging.getLogger(__name__) @@ -169,7 +166,6 @@ async def list_objects( s3_client: "S3Client", bucket: str, prefix: str | None = None, - include_version=True, recursive=True, start_after: str | None = None, ) -> ListObjectsResponse: @@ -180,72 +176,23 @@ async def list_objects( kwargs["Prefix"] = prefix if not recursive: kwargs["Delimiter"] = "/" - - objects = [] - - if include_version: - # Use list_object_versions to get all versions of objects - version_kwargs = {"Bucket": bucket} - if prefix: - version_kwargs["Prefix"] = prefix - if start_after: - version_kwargs["KeyMarker"] = start_after - if not recursive: - version_kwargs["Delimiter"] = "/" - - response = await s3_client.list_object_versions(**version_kwargs) - - # Process versions - for version in response.get("Versions", []): - objects.append( - ObjectMetadata( - bucket_name=bucket, - object_name=version.get("Key") or "", - etag=version.get("ETag", "").strip('"'), - size=version.get("Size"), - last_modified=version.get("LastModified"), - content_type="application/octet-stream", # Default, would need head_object for actual - version_id=version.get("VersionId"), - is_latest=version.get("IsLatest", False), - metadata={}, - ) - ) - - # Process delete markers if needed - for delete_marker in response.get("DeleteMarkers", []): - objects.append( - ObjectMetadata( - bucket_name=bucket, - object_name=delete_marker.get("Key") or "", - etag="", # Delete markers don't have ETags - size=0, - last_modified=delete_marker.get("LastModified"), - content_type="application/octet-stream", - version_id=delete_marker.get("VersionId"), - is_latest=delete_marker.get("IsLatest", False), - metadata={}, - ) - ) - else: - # Use list_objects_v2 for current versions only - if start_after: - kwargs["StartAfter"] = start_after - - response = await s3_client.list_objects_v2(**kwargs) - - for obj in response.get("Contents", []): - objects.append( - ObjectMetadata( - bucket_name=bucket, - object_name=obj.get("Key") or "", - etag=obj.get("ETag", "").strip('"'), - size=obj.get("Size"), - last_modified=obj.get("LastModified"), - content_type="application/octet-stream", # Default, would need head_object for actual - is_latest=True, # All objects from list_objects_v2 are latest versions - metadata={}, - ) - ) + if start_after: + kwargs["StartAfter"] = start_after + + response = await s3_client.list_objects_v2(**kwargs) + + objects = [ + ObjectMetadata( + bucket_name=bucket, + object_name=obj.get("Key") or "", + etag=obj.get("ETag", "").strip('"'), + size=obj.get("Size"), + last_modified=obj.get("LastModified"), + content_type="application/octet-stream", # Default, would need head_object for actual + metadata={}, + ) + for obj in response.get("Contents", []) + ] return ListObjectsResponse(objects=objects) except ClientError as e: @@ -257,22 +204,11 @@ async def get_object( s3_client: "S3Client", bucket: str, object: str, - version_id: str | None = None, - include_versions=False, ) -> FileObjectResponse: """Get object from S3 and return as FileObjectResponse.""" try: - # Get object metadata first - head_kwargs = {"Bucket": bucket, "Key": object} - if version_id: - head_kwargs["VersionId"] = version_id - head_response = await s3_client.head_object(**head_kwargs) - - # Get the actual object - get_kwargs = {"Bucket": bucket, "Key": object} - if version_id: - get_kwargs["VersionId"] = version_id - get_response = await s3_client.get_object(**get_kwargs) + head_response = await s3_client.head_object(Bucket=bucket, Key=object) + get_response = await s3_client.get_object(Bucket=bucket, Key=object) metadata = ObjectMetadata( bucket_name=bucket, @@ -281,19 +217,10 @@ async def get_object( size=head_response["ContentLength"], last_modified=head_response["LastModified"], content_type=head_response.get("ContentType", "application/octet-stream"), - version_id=head_response.get("VersionId") or version_id, metadata=head_response.get("Metadata", {}), ) - versions = None - if include_versions: - versions = await list_objects(s3_client, bucket, prefix=object, include_version=include_versions) - - return FileObjectResponse( - response=get_response["Body"], - metadata=metadata, - versions=versions, - ) + return FileObjectResponse(response=get_response["Body"], metadata=metadata) except ClientError as e: if e.response["Error"]["Code"] == "NoSuchKey": @@ -333,9 +260,6 @@ async def put_object( size=head_response["ContentLength"], last_modified=head_response["LastModified"], content_type=head_response.get("ContentType", content_type), - # Propagate the S3 object version (mirrors get_object) so callers can pin the - # immutable version; otherwise schema_versions.object_version_id is always NULL. - version_id=head_response.get("VersionId"), metadata=head_response.get("Metadata", {}), ) @@ -347,13 +271,10 @@ async def put_object( raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") -async def delete_object(s3_client, bucket: str, object: str, version_id: str | None = None): +async def delete_object(s3_client, bucket: str, object: str): """Delete object from S3.""" try: - kwargs = {"Bucket": bucket, "Key": object} - if version_id: - kwargs["VersionId"] = version_id - await s3_client.delete_object(**kwargs) + await s3_client.delete_object(Bucket=bucket, Key=object) except ClientError as e: _LOGGER.error(f"Error deleting object {object} from bucket {bucket}: {e}") raise HTTPException(status_code=500, detail=f"Error deleting file: {e!s}") @@ -400,59 +321,16 @@ async def bucket_exists(s3_client: "S3Client", bucket_name: str) -> bool: return False -async def get_bucket_versioning(s3_client: "S3Client", bucket_name: str) -> dict[str, str] | None: - """Get bucket versioning configuration.""" - try: - response = await s3_client.get_bucket_versioning(Bucket=bucket_name) - return { - "status": response.get("Status", "Disabled"), - "mfa_delete": response.get("MFADelete", "Disabled"), - } - except ClientError as e: - _LOGGER.error(f"Error getting bucket versioning for {bucket_name}: {e}") - return None - except Exception as e: - _LOGGER.error(f"Unexpected error getting bucket versioning for {bucket_name}: {e}") - return None - - async def create_bucket(s3_client: "S3Client", workspace_name: str): - """Create the workspace's bucket, versioned, with layout noncurrent versions expiring.""" - bucket, prefix = workspace_root(workspace_name) + """Create the workspace's bucket if it does not already exist.""" + bucket, _prefix = workspace_root(workspace_name) try: try: await s3_client.create_bucket(Bucket=bucket) except ClientError as e: - # An existing bucket must still pick up versioning and the lifecycle rule. if e.response["Error"]["Code"] not in ["BucketAlreadyOwnedByYou", "BucketAlreadyExists"]: raise - await s3_client.put_bucket_versioning( - Bucket=bucket, - VersioningConfiguration={ - "Status": "Enabled", - "MFADelete": "Disabled", - }, - ) - - try: - await s3_client.put_bucket_lifecycle_configuration( - Bucket=bucket, - LifecycleConfiguration={ - "Rules": [ - { - "ID": "expire-noncurrent-layout-versions", - "Status": "Enabled", - "Filter": {"Prefix": f"{prefix}layout/"}, - "NoncurrentVersionExpiration": {"NoncurrentDays": LAYOUT_NONCURRENT_EXPIRATION_DAYS}, - } - ] - }, - ) - except Exception as e: - # A backend without lifecycle support costs storage, never correctness. - _LOGGER.warning(f"Could not set the layout lifecycle rule on bucket {bucket}: {e}") - except ClientError as e: _LOGGER.error(f"Error creating bucket {bucket}: {e}") raise HTTPException(status_code=500, detail=f"Error creating bucket: {e!s}") @@ -488,9 +366,7 @@ async def put_document_file( # Check if file already exists with same hash try: - existing_files = await list_objects( - s3_client, workspace_name, prefix=object_path, include_version=False, recursive=False - ) + existing_files = await list_objects(s3_client, workspace_name, prefix=object_path, recursive=False) should_upload = True if existing_files.objects: diff --git a/extralit-server/src/extralit_server/contexts/schema_versions.py b/extralit-server/src/extralit_server/contexts/schema_versions.py index 2a65bc7b1..25e48cdd3 100644 --- a/extralit-server/src/extralit_server/contexts/schema_versions.py +++ b/extralit-server/src/extralit_server/contexts/schema_versions.py @@ -165,7 +165,6 @@ async def publish_version( dataset_id=dataset.id, version=next_version, object_key=key, - object_version_id=getattr(metadata, "version_id", None), etag=metadata.etag, checksum=files_ctx.compute_hash(body.encode("utf-8")), parent_version_id=parent_id, diff --git a/extralit-server/src/extralit_server/models/database.py b/extralit-server/src/extralit_server/models/database.py index f2c141c0a..940058a48 100644 --- a/extralit-server/src/extralit_server/models/database.py +++ b/extralit-server/src/extralit_server/models/database.py @@ -542,8 +542,10 @@ class SchemaVersion(DatabaseModel): """An immutable, object-store-backed Pandera schema body for a dataset. The body itself lives in the workspace bucket at `object_key`; this row is the - pointer plus integrity metadata. The column manifest derived from the body is - materialized as `Field` rows on the dataset, so there is no cached copy here. + pointer plus integrity metadata. Immutability comes from `version` (allocated under + a row lock, so each version gets its own key) plus `checksum` -- not from object-store + versioning. The column manifest derived from the body is materialized as `Field` rows + on the dataset, so there is no cached copy here. """ __tablename__ = "schema_versions" @@ -551,7 +553,6 @@ class SchemaVersion(DatabaseModel): dataset_id: Mapped[UUID] = mapped_column(ForeignKey("datasets.id", ondelete="CASCADE"), index=True) version: Mapped[int] = mapped_column(index=True) object_key: Mapped[str] = mapped_column(Text) - object_version_id: Mapped[str | None] = mapped_column(Text, nullable=True) etag: Mapped[str] = mapped_column(String) checksum: Mapped[str] = mapped_column(String) parent_version_id: Mapped[UUID | None] = mapped_column( diff --git a/extralit-server/tests/factories.py b/extralit-server/tests/factories.py index e53a457ec..cce4ebd72 100644 --- a/extralit-server/tests/factories.py +++ b/extralit-server/tests/factories.py @@ -612,10 +612,7 @@ class Meta: etag = None size = 0 content_type = "application/octet-stream" - version_id = None - is_latest = True metadata = None - version_tag = factory.LazyAttribute(lambda o: f"v{factory.Faker('pyint', min_value=1, max_value=5).generate()}") @classmethod def attributes(cls, **kwargs): @@ -626,10 +623,7 @@ def attributes(cls, **kwargs): "etag": kwargs.get("etag", None), "size": kwargs.get("size", 0), "content_type": kwargs.get("content_type", "application/octet-stream"), - "version_id": kwargs.get("version_id", None), - "is_latest": kwargs.get("is_latest", True), "metadata": kwargs.get("metadata", None), - "version_tag": kwargs.get("version_tag", "v1"), } @classmethod @@ -655,7 +649,7 @@ def mock_put_object(bucket_name, object_name, data, content_type=None, metadata= return file # Mock get_object to return file data - def mock_get_object(bucket_name, object_name, version_id=None): + def mock_get_object(bucket_name, object_name): response = MagicMock() response.data = b"test data" return response diff --git a/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py b/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py index d8cdf7acb..729b2bc9c 100644 --- a/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py +++ b/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py @@ -80,7 +80,6 @@ def _mock_put_object(): bucket_name="workspace", object_name="schemas/dataset/v1.json", etag="etag", - version_id="v1", ) yield mock_put_object diff --git a/extralit-server/tests/unit/api/handlers/v1/test_files.py b/extralit-server/tests/unit/api/handlers/v1/test_files.py index 71b282c60..32a0921eb 100644 --- a/extralit-server/tests/unit/api/handlers/v1/test_files.py +++ b/extralit-server/tests/unit/api/handlers/v1/test_files.py @@ -1,7 +1,6 @@ import io -import os from typing import TYPE_CHECKING -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest @@ -26,7 +25,7 @@ async def test_put_file(async_client: "AsyncClient", owner_auth_header: dict): # Mock the Minio client and the response with patch("extralit_server.contexts.files.put_object") as mock_put_object: - mock_response = ObjectMetadata(bucket_name=bucket_name, object_name=object_name, is_latest=True) + mock_response = ObjectMetadata(bucket_name=bucket_name, object_name=object_name) mock_put_object.return_value = mock_response response = await async_client.post( @@ -37,7 +36,6 @@ async def test_put_file(async_client: "AsyncClient", owner_auth_header: dict): assert response.status_code == 200 assert response.json()["object_name"] == mock_response.object_name - assert response.json()["is_latest"] == mock_response.is_latest # Verify put_object was called correctly mock_put_object.assert_called_once() @@ -80,66 +78,6 @@ async def test_list_objects(async_client: "AsyncClient", owner_auth_header: dict assert response.json() == mock_response.dict() -@pytest.mark.asyncio -async def test_list_objects_with_versions(async_client: "AsyncClient", owner_auth_header: dict): - bucket_name = "workspace-files" - prefix = "schemas" - object_name = os.path.join(prefix, "test") - - # Mock get_s3_client and bucket_exists - with ( - patch("extralit_server.contexts.files.get_s3_client") as mock_get_s3_client, - patch("extralit_server.contexts.files.delete_bucket"), - patch("extralit_server.contexts.files.list_objects") as mock_list_objects, - ): - # Setup mocks - mock_client = MagicMock() - mock_client.bucket_exists.return_value = True - mock_get_s3_client.return_value = mock_client - - # Create workspace and user - workspace_a = await WorkspaceFactory.create(name=bucket_name) - user_a = await UserFactory.create(username="username-a") - await WorkspaceUserFactory.create(workspace_id=workspace_a.id, user_id=user_a.id) - - # Create mock objects for response - file1 = MinioFileFactory.build( - bucket_name=bucket_name, object_name=object_name, version_tag="v1", is_latest=False - ) - file2 = MinioFileFactory.build( - bucket_name=bucket_name, object_name=object_name, version_tag="v2", is_latest=True - ) - - # Set up list_objects mock - mock_list_objects.return_value = ListObjectsResponse( - objects=[ - ObjectMetadata(**file1.dict()), - ObjectMetadata(**file2.dict()), - ] - ) - - response = await async_client.get( - f"/api/v1/files/{bucket_name}/{prefix}", headers={API_KEY_HEADER_NAME: user_a.api_key} - ) - - assert response.status_code == 200 - - # Check versions - response_objects = response.json()["objects"] - assert len(response_objects) == 2 - - # Check version tags - response_version_tags = {item["version_tag"] for item in response_objects} - assert response_version_tags == {"v1", "v2"} - - # Check latest flag - for item in response_objects: - if item["version_tag"] == "v2": - assert item["is_latest"] is True - else: - assert item["is_latest"] is False - - @pytest.mark.asyncio async def test_delete_file(async_client: "AsyncClient", owner_auth_header: dict): bucket_name = "workspace" diff --git a/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py b/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py index 2886ce359..6f3205191 100644 --- a/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py +++ b/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py @@ -18,11 +18,9 @@ async def test_workspace_doctor_healthy(self, db: AsyncSession, async_client: As # Mock S3 client and Redis connection with ( patch("extralit_server.contexts.files.bucket_exists") as mock_bucket_exists, - patch("extralit_server.contexts.files.get_bucket_versioning") as mock_get_versioning, patch("extralit_server.jobs.queues.DEFAULT_QUEUE") as mock_queue, ): mock_bucket_exists.return_value = True - mock_get_versioning.return_value = {"status": "Enabled", "mfa_delete": "Disabled"} mock_queue.connection.ping.return_value = True response = await async_client.post( @@ -37,7 +35,7 @@ async def test_workspace_doctor_healthy(self, db: AsyncSession, async_client: As assert data["workspace_id"] == str(workspace.id) assert data["workspace_name"] == workspace.name assert data["overall_status"] == "healthy" - assert len(data["checks"]) >= 3 # At least bucket, versioning, and RQ checks + assert len(data["checks"]) >= 2 # At least the bucket and RQ checks # Check that bucket check passed bucket_check = next((c for c in data["checks"] if c["check_name"] == "s3_bucket"), None) @@ -53,12 +51,10 @@ async def test_workspace_doctor_missing_bucket_with_autofix( with ( patch("extralit_server.contexts.files.bucket_exists") as mock_bucket_exists, patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket, - patch("extralit_server.contexts.files.get_bucket_versioning") as mock_get_versioning, patch("extralit_server.jobs.queues.DEFAULT_QUEUE") as mock_queue, ): mock_bucket_exists.return_value = False mock_create_bucket.return_value = None - mock_get_versioning.return_value = {"status": "Enabled", "mfa_delete": "Disabled"} mock_queue.connection.ping.return_value = True response = await async_client.post( @@ -90,11 +86,9 @@ async def test_workspace_doctor_missing_bucket_without_autofix( with ( patch("extralit_server.contexts.files.bucket_exists") as mock_bucket_exists, patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket, - patch("extralit_server.contexts.files.get_bucket_versioning") as mock_get_versioning, patch("extralit_server.jobs.queues.DEFAULT_QUEUE") as mock_queue, ): mock_bucket_exists.return_value = False - mock_get_versioning.return_value = None mock_queue.connection.ping.return_value = True response = await async_client.post( diff --git a/extralit-server/tests/unit/contexts/test_schema_versions.py b/extralit-server/tests/unit/contexts/test_schema_versions.py index 40e4bb4f7..1916bb9ee 100644 --- a/extralit-server/tests/unit/contexts/test_schema_versions.py +++ b/extralit-server/tests/unit/contexts/test_schema_versions.py @@ -45,7 +45,6 @@ def _s3_client() -> AsyncMock: "ContentLength": 0, "LastModified": None, "ContentType": "application/json", - "VersionId": "v1", "Metadata": {}, } return client diff --git a/extralit/src/extralit/_api/_workspaces.py b/extralit/src/extralit/_api/_workspaces.py index b2e47b12c..5a17df593 100644 --- a/extralit/src/extralit/_api/_workspaces.py +++ b/extralit/src/extralit/_api/_workspaces.py @@ -132,16 +132,13 @@ def doctor(self, workspace_id: "UUID", autofix: bool = True) -> "WorkspaceDoctor #################### @api_error_handler - def list_files( - self, workspace_name: str, path: str, recursive: bool = True, include_version: bool = True - ) -> "ListObjectsResponse": + def list_files(self, workspace_name: str, path: str, recursive: bool = True) -> "ListObjectsResponse": """List files in a workspace. Args: workspace_name: The name of the workspace. path: The path to list files from. recursive: Whether to list files recursively. - include_version: Whether to include version information. Returns: A list of files. @@ -156,10 +153,7 @@ def list_files( logger.info(f"Listing files in workspace '{workspace_name}' with path '{path}'") url = f"/api/v1/files/{workspace_name}/{path}" - params = { - "recursive": recursive, - "include_version": include_version, - } + params = {"recursive": recursive} try: response = self.http_client.get(url=url, params=params) @@ -173,17 +167,6 @@ def list_files( result = ListObjectsResponse(**response.json()) - # Filter out files that have deletion markers - # First, find all files that have deletion markers (etag="" and size=0) - deleted_files = {obj.object_name for obj in result.objects if obj.etag == "" and obj.size == 0} - - # Then filter out ALL entries for those files (both original and deletion marker) - valid_objects = [obj for obj in result.objects if obj.object_name not in deleted_files] - - if len(valid_objects) < len(result.objects): - logger.info(f"Filtered out {len(result.objects) - len(valid_objects)} deleted files") - result.objects = valid_objects - logger.info(f"Found {len(result.objects)} files in workspace '{workspace_name}'") return result except httpx.HTTPStatusError as e: @@ -197,13 +180,12 @@ def list_files( raise @api_error_handler - def get_file(self, workspace_name: str, path: str, version_id: Optional[str] = None) -> "FileObjectResponse": + def get_file(self, workspace_name: str, path: str) -> "FileObjectResponse": """Get a file from a workspace. Args: workspace_name: The name of the workspace. path: The path of the file. - version_id: The version ID of the file. Returns: The file content and metadata. @@ -223,10 +205,9 @@ def get_file(self, workspace_name: str, path: str, version_id: Optional[str] = N logger.info(f"Getting file '{path}' from workspace '{workspace_name}'") url = f"/api/v1/file/{workspace_name}/{path}" - params = {"version_id": version_id} if version_id else {} try: - response = self.http_client.get(url=url, params=params) + response = self.http_client.get(url=url) # Handle 404 directly if response.status_code == 404: @@ -235,25 +216,15 @@ def get_file(self, workspace_name: str, path: str, version_id: Optional[str] = N response.raise_for_status() - # Create a FileObjectResponse with the content - file_response = FileObjectResponse(content=response.content) - - # Check if the file is a deleted file (empty content with headers) - if len(response.content) == 0 and "X-Amz-Meta-Version-Tag" in response.headers: - logger.warning(f"File '{path}' exists but has empty content (likely deleted)") - raise FileNotFoundError(f"File '{path}' exists but has empty content (likely deleted)") - - # Get metadata if available - if "X-Amz-Meta-Version-Tag" in response.headers: - metadata = ObjectMetadata( + file_response = FileObjectResponse( + content=response.content, + metadata=ObjectMetadata( bucket_name=workspace_name, object_name=path, content_type=response.headers.get("Content-Type"), etag=response.headers.get("ETag"), - version_id=version_id or response.headers.get("X-Amz-Version-Id"), - version_tag=response.headers.get("X-Amz-Meta-Version-Tag"), - ) - file_response.metadata = metadata + ), + ) logger.info(f"Successfully retrieved file '{path}' from workspace '{workspace_name}'") return file_response @@ -328,13 +299,12 @@ def put_file(self, workspace_name: str, path: str, file_path: Path) -> "ObjectMe raise @api_error_handler - def delete_file(self, workspace_name: str, path: str, version_id: Optional[str] = None) -> None: + def delete_file(self, workspace_name: str, path: str) -> None: """Delete a file from a workspace. Args: workspace_name: The name of the workspace. path: The path of the file to delete. - version_id: The version ID of the file. Raises: ExtralitAPIError: If the API request fails. @@ -351,10 +321,9 @@ def delete_file(self, workspace_name: str, path: str, version_id: Optional[str] logger.info(f"Deleting file '{path}' from workspace '{workspace_name}'") url = f"/api/v1/file/{workspace_name}/{path}" - params = {"version_id": version_id} if version_id else {} try: - response = self.http_client.delete(url=url, params=params) + response = self.http_client.delete(url=url) response.raise_for_status() logger.info(f"Successfully deleted file '{path}' from workspace '{workspace_name}'") except httpx.HTTPStatusError as e: @@ -426,12 +395,7 @@ def list_schemas( ) try: - schema_files = self.list_files(workspace_name, path=prefix, recursive=True, include_version=True) - - # Filter out deleted schema versions (those with etag=None and size=None) - schema_files.objects = [ - obj for obj in schema_files.objects if obj.etag is not None and obj.size is not None - ] + schema_files = self.list_files(workspace_name, path=prefix, recursive=True) logger.info(f"Found {len(schema_files.objects)} valid schema files in workspace '{workspace_name}'") diff --git a/extralit/src/extralit/_models/_files.py b/extralit/src/extralit/_models/_files.py index 180c0ceb5..afe24db93 100644 --- a/extralit/src/extralit/_models/_files.py +++ b/extralit/src/extralit/_models/_files.py @@ -10,12 +10,9 @@ class ObjectMetadata(BaseModel): bucket_name: str object_name: str last_modified: Optional[datetime] = None - is_latest: Optional[bool] = None etag: Optional[str] = None size: Optional[int] = None content_type: Optional[str] = None - version_id: Optional[str] = None - version_tag: Optional[str] = None metadata: Optional[dict[str, Any]] = None @@ -39,4 +36,3 @@ class FileObjectResponse(BaseModel): content: bytes metadata: Optional[ObjectMetadata] = None - versions: Optional[ListObjectsResponse] = None diff --git a/extralit/src/extralit/cli/files/delete.py b/extralit/src/extralit/cli/files/delete.py index 6bfa42f85..ada03e547 100644 --- a/extralit/src/extralit/cli/files/delete.py +++ b/extralit/src/extralit/cli/files/delete.py @@ -1,5 +1,3 @@ -from typing import Optional - import typer from rich.console import Console @@ -10,7 +8,6 @@ def delete_file( remote_path: str = typer.Argument(..., help="Remote path of the file to delete"), workspace: str = typer.Option(..., "--workspace", "-w", help="Workspace name"), - version_id: Optional[str] = typer.Option(None, "--version-id", "-v", help="Version ID of the file to delete"), force: bool = typer.Option(False, "--force", "-f", help="Force deletion without confirmation"), ) -> None: """Delete a file from a workspace.""" @@ -48,7 +45,7 @@ def delete_file( return # Delete the file - workspace_obj.delete_file(remote_path, version_id=version_id) + workspace_obj.delete_file(remote_path) # Print a success message panel = get_themed_panel( diff --git a/extralit/src/extralit/cli/files/download.py b/extralit/src/extralit/cli/files/download.py index 7921f72e4..a7c842ed0 100644 --- a/extralit/src/extralit/cli/files/download.py +++ b/extralit/src/extralit/cli/files/download.py @@ -16,7 +16,6 @@ def download_file( output_path: Optional[Path] = typer.Option( None, "--output", "-o", help="Local path to save the file (default: same as remote filename)" ), - version_id: Optional[str] = typer.Option(None, "--version-id", "-v", help="Version ID of the file to download"), overwrite: bool = typer.Option(False, "--overwrite", help="Overwrite existing file"), ) -> None: """Download a file from a workspace.""" @@ -62,7 +61,7 @@ def download_file( task = progress.add_task(f"Downloading {remote_path} from {workspace}...", total=None) # Download the file - file_response = workspace_obj.get_file(remote_path, version_id=version_id) + file_response = workspace_obj.get_file(remote_path) # Save the file with open(output_path, "wb") as f: diff --git a/extralit/src/extralit/cli/files/list.py b/extralit/src/extralit/cli/files/list.py index e5bb447f2..bad2aab44 100644 --- a/extralit/src/extralit/cli/files/list.py +++ b/extralit/src/extralit/cli/files/list.py @@ -8,7 +8,6 @@ def list_files( workspace: str = typer.Option(..., "--workspace", "-w", help="Workspace name"), path: str = typer.Option("", "--path", "-p", help="Path prefix to filter files"), recursive: bool = typer.Option(True, "--recursive/--no-recursive", help="List files recursively"), - include_version: bool = typer.Option(False, "--versions/--no-versions", help="Include version information"), ) -> None: from rich.console import Console @@ -29,11 +28,7 @@ def list_files( console.print(panel) raise typer.Exit(code=1) - files = workspace_obj.list_files(path, recursive=recursive, include_version=include_version) - - # Filter out files that have deletion markers - remove ALL entries for files with deletion markers - deleted_files = {obj.object_name for obj in files.objects if obj.etag == "" and obj.size == 0} - files.objects = [obj for obj in files.objects if obj.object_name not in deleted_files] + files = workspace_obj.list_files(path, recursive=recursive) if not files.objects: panel = get_themed_panel( diff --git a/extralit/src/extralit/cli/rich.py b/extralit/src/extralit/cli/rich.py index 31bfee417..197af68ec 100644 --- a/extralit/src/extralit/cli/rich.py +++ b/extralit/src/extralit/cli/rich.py @@ -135,9 +135,9 @@ def print_rich_table( "Name": lambda r: r.name, "Workspace": lambda r: r.workspace.name, "Creation Date": lambda r: r.inserted_at.isoformat(sep=" ") if r.inserted_at else "", - "Last Activity Date": lambda r: r._model.last_activity_at.isoformat(sep=" ") - if r._model.last_activity_at - else "", + "Last Activity Date": lambda r: ( + r._model.last_activity_at.isoformat(sep=" ") if r._model.last_activity_at else "" + ), }, "styles": { "ID": "cyan", @@ -182,19 +182,17 @@ def print_rich_table( }, }, "ObjectMetadata": { - "columns": ["Object Name", "Size", "Last Modified", "Version ID", "Content Type"], + "columns": ["Object Name", "Size", "Last Modified", "Content Type"], "getters": { "Object Name": lambda r: r.object_name, "Size": lambda r: r.size, "Last Modified": lambda r: r.last_modified.isoformat() if r.last_modified else "", - "Version ID": lambda r: r.version_id, "Content Type": lambda r: r.content_type, }, "styles": { "Object Name": "cyan", "Size": "green", "Last Modified": "yellow", - "Version ID": "magenta", "Content Type": "blue", }, }, diff --git a/extralit/src/extralit/workspaces/_resource.py b/extralit/src/extralit/workspaces/_resource.py index 363a9901f..80a875c93 100644 --- a/extralit/src/extralit/workspaces/_resource.py +++ b/extralit/src/extralit/workspaces/_resource.py @@ -101,30 +101,28 @@ def list_datasets(self) -> list["Dataset"]: # File methods # #################### - def list_files(self, path: str, recursive: bool = True, include_version: bool = True) -> "ListObjectsResponse": + def list_files(self, path: str, recursive: bool = True) -> "ListObjectsResponse": """List files in the workspace. Args: path: The path to list files from. recursive: Whether to list files recursively. - include_version: Whether to include version information. Returns: A list of files. """ - return self._api.list_files(self.name, path, recursive, include_version) + return self._api.list_files(self.name, path, recursive) - def get_file(self, path: str, version_id: Optional[str] = None) -> "FileObjectResponse": + def get_file(self, path: str) -> "FileObjectResponse": """Get a file from the workspace. Args: path: The path of the file. - version_id: The version ID of the file. Returns: The file content and metadata. """ - return self._api.get_file(self.name, path, version_id) + return self._api.get_file(self.name, path) def put_file(self, path: str, file_path: Union[str, Path]) -> "ObjectMetadata": """Upload a file to the workspace. @@ -140,14 +138,13 @@ def put_file(self, path: str, file_path: Union[str, Path]) -> "ObjectMetadata": file_path = Path(file_path) return self._api.put_file(self.name, path, file_path) - def delete_file(self, path: str, version_id: Optional[str] = None) -> None: + def delete_file(self, path: str) -> None: """Delete a file from the workspace. Args: path: The path of the file to delete. - version_id: The version ID of the file. """ - self._api.delete_file(self.name, path, version_id) + self._api.delete_file(self.name, path) #################### # Document methods # diff --git a/extralit/tests/integration/test_cli_commands.py b/extralit/tests/integration/test_cli_commands.py index 37067839a..263a326ef 100644 --- a/extralit/tests/integration/test_cli_commands.py +++ b/extralit/tests/integration/test_cli_commands.py @@ -75,7 +75,6 @@ def test_files_upload_and_list_command(self, test_workspace): except Exception: pass - @pytest.mark.skip(reason="buckets with versioning enabled still list deleted files, needs further investigation") def test_files_upload_download_and_delete_command(self, test_workspace): """Test the 'files upload', 'files download', and 'files delete' commands.""" with tempfile.NamedTemporaryFile(delete=False, suffix=".txt") as temp_file: diff --git a/extralit/tests/unit/api/test_workspace_files_api.py b/extralit/tests/unit/api/test_workspace_files_api.py index 752b7a039..fa98f1a85 100644 --- a/extralit/tests/unit/api/test_workspace_files_api.py +++ b/extralit/tests/unit/api/test_workspace_files_api.py @@ -21,12 +21,9 @@ def test_list_files(workspace_api: WorkspacesAPI): "bucket_name": "test-workspace", "object_name": "test-file.txt", "last_modified": "2023-01-01T00:00:00Z", - "is_latest": True, "etag": "test-etag", "size": 100, "content_type": "text/plain", - "version_id": "test-version-id", - "version_tag": "test-version-tag", "metadata": {}, } ] @@ -41,7 +38,7 @@ def test_list_files(workspace_api: WorkspacesAPI): assert result.objects[0].object_name == "test-file.txt" workspace_api.http_client.get.assert_called_once_with( # type: ignore - url="/api/v1/files/test-workspace/test-path", params={"recursive": True, "include_version": True} + url="/api/v1/files/test-workspace/test-path", params={"recursive": True} ) @@ -53,7 +50,6 @@ def test_get_file(workspace_api: WorkspacesAPI): mock_response.headers = { "Content-Type": "text/plain", "ETag": "test-etag", - "X-Amz-Meta-Version-Tag": "test-version-tag", } workspace_api.http_client.get.return_value = mock_response @@ -65,10 +61,9 @@ def test_get_file(workspace_api: WorkspacesAPI): assert result.metadata.object_name == "test-file.txt" assert result.metadata.content_type == "text/plain" assert result.metadata.etag == "test-etag" - assert result.metadata.version_tag == "test-version-tag" # Verify the API call - workspace_api.http_client.get.assert_called_once_with(url="/api/v1/file/test-workspace/test-file.txt", params={}) + workspace_api.http_client.get.assert_called_once_with(url="/api/v1/file/test-workspace/test-file.txt") def test_put_file(workspace_api, tmp_path): @@ -82,12 +77,9 @@ def test_put_file(workspace_api, tmp_path): "bucket_name": "test-workspace", "object_name": "test-file.txt", "last_modified": "2023-01-01T00:00:00Z", - "is_latest": True, "etag": "test-etag", "size": 100, "content_type": "text/plain", - "version_id": "test-version-id", - "version_tag": "test-version-tag", "metadata": {}, } workspace_api.http_client.post.return_value = mock_response # type: ignore @@ -98,7 +90,6 @@ def test_put_file(workspace_api, tmp_path): assert result.bucket_name == "test-workspace" assert result.object_name == "test-file.txt" assert result.etag == "test-etag" - assert result.version_id == "test-version-id" workspace_api.http_client.post.assert_called_once() assert workspace_api.http_client.post.call_args.kwargs["url"] == "/api/v1/file/test-workspace/test-file.txt" @@ -114,4 +105,4 @@ def test_delete_file(workspace_api: WorkspacesAPI): workspace_api.delete_file("test-workspace", "test-file.txt") # Verify the API call - workspace_api.http_client.delete.assert_called_once_with(url="/api/v1/file/test-workspace/test-file.txt", params={}) # type: ignore + workspace_api.http_client.delete.assert_called_once_with(url="/api/v1/file/test-workspace/test-file.txt") # type: ignore diff --git a/extralit/tests/unit/api/test_workspace_schemas_api.py b/extralit/tests/unit/api/test_workspace_schemas_api.py index 84fa9ff29..451435e07 100644 --- a/extralit/tests/unit/api/test_workspace_schemas_api.py +++ b/extralit/tests/unit/api/test_workspace_schemas_api.py @@ -61,8 +61,6 @@ def test_get_schemas(workspace_api): object_name="schemas/test_schema", content_type="application/json", etag="test-etag", - version_id="test-version-id", - version_tag="test-version-tag", ) ] ) @@ -108,8 +106,6 @@ def test_update_schemas(workspace_api, mock_schema_structure): object_name="schemas/test_schema", content_type="application/json", etag="test-etag", - version_id="test-version-id", - version_tag="test-version-tag", ) ] ) From 0bb5b8fc280eefd7a98e81f0f920ef8275ef6869 Mon Sep 17 00:00:00 2001 From: JonnyTran Date: Wed, 19 Aug 2026 00:14:02 -0700 Subject: [PATCH 2/3] refactor(server): port object storage to obstore Replaces the aioboto3/botocore dict-shape contract and the hand-rolled local backend with one obstore port. `contexts/files.py` now talks to `S3Store` in production and `LocalStore` in local dev, behind an `ObjectStorage` that resolves a bucket name to a store (obstore binds a store to one bucket at construction, so the per-call `bucket` argument becomes a lookup). What this deletes: - `aiofiles`, entirely. Its only user was `LocalFileClient`, ~430 lines subclassing aiobotocore's generated client to fake S3 over a directory tree. It was untested, and derived paths with `Path.with_suffix`, so any key containing a dot did not round-trip (`pdf/a.b.pdf` came back as `pdf/a`). It also silently lacked `head_bucket`, so `bucket_exists` was always false locally. - The read-whole-file-into-RAM path in `GET /file`. `head_object` + a full `get_object` + a 10MB chunking branch collapse into one `get_async` and `StreamingResponse(result.stream())`. `aioboto3` survives only in the new `contexts/buckets.py`: obstore binds to an existing bucket and has no admin API, so CreateBucket/DeleteBucket/HeadBucket cannot move. That is 3 functions and 4 call sites rather than the app-wide contract it was. Details worth knowing: - obstore raises the builtin `FileNotFoundError` for a missing key, not `obstore.exceptions.NotFoundError`. - An HTTP range is inclusive of its end and obstore's is exclusive, so the range handler passes `(start, end + 1)`. `Range: bytes=0-0` must return one byte. - `LocalStore` rejects `put` attributes outright, so local mode recovers the content type from the key. The keys we mint (`pdf/`, `thumbnails/`) carry no extension, hence the prefix map -- without it a local-dev PDF is served as octet-stream and the viewer breaks. - `ObjectMetadata.metadata` no longer strips an `x-amz-meta-` prefix. obstore returns user attributes under their own keys, so that stripping discarded every one of them. Verified against MinIO. Also drops the ES mapping properties for the removed version fields. The `table` field object sets `dynamic: false`, which overrides the index-level `dynamic: strict`, so a legacy response still carrying those keys is ignored rather than rejected -- no reindex needed. tests/unit/contexts/test_files_store.py is new and drives a real `LocalStore`. Nothing previously exercised a storage backend end to end; every other test patches `contexts.files.*`. --- extralit-server/pyproject.toml | 2 - extralit-server/src/extralit_server/_app.py | 16 +- .../handlers/v1/datasets/schema_versions.py | 4 +- .../api/handlers/v1/documents.py | 18 +- .../extralit_server/api/handlers/v1/files.py | 53 +- .../api/handlers/v1/workspaces.py | 16 +- .../extralit_server/api/schemas/v1/files.py | 11 +- .../cli/database/users/create.py | 8 +- .../cli/database/users/create_default.py | 8 +- .../src/extralit_server/contexts/buckets.py | 117 ++++ .../src/extralit_server/contexts/files.py | 460 +++++++-------- .../src/extralit_server/contexts/imports.py | 6 +- .../extralit_server/contexts/ocr/storage.py | 20 +- .../contexts/schema_versions.py | 10 +- .../src/extralit_server/helpers.py | 522 ------------------ .../src/extralit_server/jobs/document_jobs.py | 8 +- .../src/extralit_server/jobs/ocr_jobs.py | 8 +- .../src/extralit_server/jobs/preload.py | 1 - .../extralit_server/search_engine/commons.py | 5 - extralit-server/tests/factories.py | 15 +- .../v1/datasets/test_schema_versions.py | 4 +- .../api/handlers/v1/test_document_layout.py | 6 +- .../v1/workspaces/test_create_workspace.py | 6 +- .../v1/workspaces/test_workspace_doctor.py | 10 +- .../tests/unit/contexts/test_files_store.py | 185 +++++++ .../unit/contexts/test_schema_versions.py | 76 ++- .../tests/unit/jobs/test_document_jobs.py | 2 +- .../tests/unit/jobs/test_ocr_jobs.py | 4 +- extralit-server/uv.lock | 16 - 29 files changed, 639 insertions(+), 978 deletions(-) create mode 100644 extralit-server/src/extralit_server/contexts/buckets.py create mode 100644 extralit-server/tests/unit/contexts/test_files_store.py diff --git a/extralit-server/pyproject.toml b/extralit-server/pyproject.toml index a398f14b8..23fe52e2e 100644 --- a/extralit-server/pyproject.toml +++ b/extralit-server/pyproject.toml @@ -32,7 +32,6 @@ dependencies = [ # Async SQLite "aiosqlite == 0.20.0", # Statics server - "aiofiles >= 24.1.0", "PyYAML >= 5.4.1,< 6.1.0", # security dependencies "python-jose[cryptography] >= 3.3.0", @@ -71,7 +70,6 @@ dependencies = [ "Jinja2>=3.1.4", # Used by huggingface-hub to render dataset card templates # For file storage "aioboto3>=13.1.1", - "types-aiobotocore-s3==2.24.2", # For document processing "ocrmypdf>=16.11.0", "pdf2image>=1.17.0", diff --git a/extralit-server/src/extralit_server/_app.py b/extralit-server/src/extralit_server/_app.py index 5ab6d335b..9049cc05e 100644 --- a/extralit-server/src/extralit_server/_app.py +++ b/extralit-server/src/extralit_server/_app.py @@ -25,9 +25,9 @@ from extralit_server._version import __version__ as extralit_version from extralit_server.api.routes import api_v1 from extralit_server.constants import DEFAULT_API_KEY, DEFAULT_PASSWORD, DEFAULT_USERNAME -from extralit_server.contexts import accounts, files +from extralit_server.contexts import accounts, buckets, files from extralit_server.database import get_async_db -from extralit_server.helpers import create_s3_client, shared_resources +from extralit_server.helpers import shared_resources from extralit_server.jobs.queues import REDIS_CONNECTION from extralit_server.logging import configure_logging from extralit_server.models import User, Workspace @@ -47,14 +47,12 @@ async def app_lifespan(app: FastAPI): configure_redis() try: - await create_s3_client() track_server_startup() yield finally: - # Clean up S3 client if it exists - s3_client = shared_resources.get("s3_client") - if s3_client: - await s3_client.__aexit__(None, None, None) + storage = shared_resources.get("storage") + if storage: + await storage.aclose() shared_resources.clear() @@ -304,8 +302,8 @@ async def _create_oauth_allowed_workspaces(db: AsyncSession): if await Workspace.get_by(db, name=allowed_workspace.name) is None: _LOGGER.info(f"Creating workspace with name {allowed_workspace.name!r}") try: - client = await files.get_s3_client() - await files.create_bucket(client, allowed_workspace.name) + storage = await files.get_storage() + await buckets.create(storage, allowed_workspace.name) except Exception as e: _LOGGER.error(f"Failed to create bucket for workspace {allowed_workspace.name!r}: {e}") diff --git a/extralit-server/src/extralit_server/api/handlers/v1/datasets/schema_versions.py b/extralit-server/src/extralit_server/api/handlers/v1/datasets/schema_versions.py index 8140dcd06..28d9d9e63 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/datasets/schema_versions.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/datasets/schema_versions.py @@ -27,7 +27,7 @@ async def publish_schema_version( dataset_id: UUID, version_create: SchemaVersionCreate, db: Annotated[AsyncSession, Depends(get_async_db)], - s3_client=Depends(files_ctx.get_s3_client), + storage=Depends(files_ctx.get_storage), current_user: Annotated[User, Security(auth.get_current_user)], ): dataset = await Dataset.get_or_raise(db, dataset_id, options=[selectinload(Dataset.workspace)]) @@ -35,7 +35,7 @@ async def publish_schema_version( return await schema_versions.publish_version( db, - s3_client, + storage, dataset, body=version_create.body, # One bucket per workspace, named exactly Workspace.name — contexts/files.py:381. diff --git a/extralit-server/src/extralit_server/api/handlers/v1/documents.py b/extralit-server/src/extralit_server/api/handlers/v1/documents.py index a3987e9d9..fd8823df3 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/documents.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/documents.py @@ -14,7 +14,7 @@ from extralit_server.api.schemas.v1.documents import DocumentCreate, DocumentDelete, DocumentListItem, DocumentUpdate from extralit_server.api.schemas.v1.imports import DocumentsBulkCreate, DocumentsBulkResponse from extralit_server.contexts import files, imports -from extralit_server.contexts.ocr import storage +from extralit_server.contexts.ocr import storage as layout_storage from extralit_server.contexts.ocr.projection import project_layout from extralit_server.database import get_async_db from extralit_server.models import User, Workspace @@ -36,7 +36,7 @@ async def add_document( document_create: Annotated[str, Form()], file_data: Annotated[UploadFile | None, File()] = None, db: AsyncSession = Depends(get_async_db), - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_current_user), ): await authorize(current_user, DocumentPolicy.create()) @@ -80,7 +80,7 @@ async def add_document( document_new.file_name = file_data.filename file_url = await files.put_document_file( - s3_client=s3_client, + storage=storage, workspace_name=workspace.name, document_id=document_new.id, # type: ignore[arg-type] file_data=await file_data.read(), @@ -129,7 +129,7 @@ async def get_document( doi: Annotated[str | None, Query(description="DOI")] = None, limit: Annotated[int | None, Query(description="Maximum number of documents to return")] = None, db: AsyncSession = Depends(get_async_db), - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_current_user), ) -> list[DocumentListItem]: await authorize(current_user, DocumentPolicy.get()) @@ -208,7 +208,7 @@ async def delete_documents_by_workspace_id( workspace_id: UUID, document_delete: Annotated[DocumentDelete | None, Body()] = None, db: AsyncSession = Depends(get_async_db), - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_current_user), ): await authorize(current_user, DocumentPolicy.delete(workspace_id)) @@ -231,7 +231,7 @@ async def delete_documents_by_workspace_id( _LOGGER.info(f"Deleting {len(documents)} documents") for document in documents: - await files.delete_document_artifacts(s3_client, workspace.name, document.id) + await files.delete_document_artifacts(storage, workspace.name, document.id) return len(documents) @@ -261,7 +261,7 @@ async def get_document_layout( document_id: Annotated[UUID, Path(title="The UUID of the document whose layout will be retrieved")], pages: Annotated[list[int] | None, Query(description="1-indexed pages to include")] = None, labels: Annotated[list[str] | None, Query(description="DocItemLabels to include, e.g. `table`")] = None, - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_current_user), ) -> DocumentLayoutOut: await authorize(current_user, DocumentPolicy.get()) @@ -291,8 +291,8 @@ async def get_document_layout( ) try: - doc = await storage.load_layout( - s3_client, + doc = await layout_storage.load_layout( + storage, workspace.name, document_id, object_path=layout_metadata.get("layout_url"), diff --git a/extralit-server/src/extralit_server/api/handlers/v1/files.py b/extralit-server/src/extralit_server/api/handlers/v1/files.py index 21a6416ec..a78846cf0 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/files.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/files.py @@ -1,7 +1,6 @@ import logging from typing import Annotated -from botocore.exceptions import ClientError from fastapi import APIRouter, Depends, File, Header, HTTPException, Security, UploadFile from fastapi.responses import Response, StreamingResponse @@ -23,18 +22,19 @@ async def get_file( object: str, range_header: str | None = Header(None, alias="range"), if_none_match: str | None = Header(None, alias="if-none-match"), - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User | None = Security(auth.get_optional_current_user), ): if current_user is not None: await authorize(current_user, FilePolicy.get(bucket)) + store = storage.store_for(bucket) + try: - # Get object metadata first - head_response = await s3_client.head_object(Bucket=bucket, Key=object) - content_length = head_response["ContentLength"] - etag = head_response["ETag"].strip('"') - content_type = head_response.get("ContentType", "application/octet-stream") + head = await store.get_async(object, options={"head": True}) + content_length = head.meta["size"] + etag = (head.meta["e_tag"] or "").strip('"') + content_type = files.content_type_of(object, head.attributes) # Handle ETag for caching if if_none_match and etag and if_none_match.strip('"') == etag: @@ -73,15 +73,15 @@ async def get_file( headers["Content-Range"] = f"bytes */{content_length}" return Response(status_code=416, headers=headers) - # Get object with range - response = await s3_client.get_object(Bucket=bucket, Key=object, Range=f"bytes={start}-{end}") + # An HTTP range is inclusive of `end`; obstore's is exclusive. + result = await store.get_async(object, options={"range": (start, end + 1)}) # Update headers for partial content headers["Content-Range"] = f"bytes {start}-{end}/{content_length}" headers["Content-Length"] = str(end - start + 1) return StreamingResponse( - response["Body"], + result.stream(), status_code=206, media_type=content_type, headers=headers, @@ -90,23 +90,12 @@ async def get_file( # Invalid range header, serve full content pass - # Get full object - response = await s3_client.get_object(Bucket=bucket, Key=object) - - # Use chunked streaming for large files - if content_length > files.CHUNK_LENGTH_MB: - file_chunks = files.get_file_chunk(s3_client, bucket, object, files.CHUNK_LENGTH_MB) - return StreamingResponse(file_chunks, media_type=content_type, headers=headers) - - return StreamingResponse(response["Body"], media_type=content_type, headers=headers) + result = await store.get_async(object) + return StreamingResponse(result.stream(), media_type=content_type, headers=headers) - except ClientError as e: - if e.response["Error"]["Code"] in {"NoSuchKey", "404"}: - _LOGGER.error(f"Object '{bucket}/{object}' not found") - raise HTTPException(status_code=404, detail=f"No object at path '{object}' was found") - else: - _LOGGER.error(f"Error getting object '{bucket}/{object}': {e.response['Error']}", exc_info=True) - raise HTTPException(status_code=500, detail=str(e)) + except FileNotFoundError: + _LOGGER.error(f"Object '{bucket}/{object}' not found") + raise HTTPException(status_code=404, detail=f"No object at path '{object}' was found") except Exception as e: _LOGGER.error(f"Error getting object '{bucket}/{object}': {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e)) @@ -132,7 +121,7 @@ async def put_file( bucket: str, object: str, file: Annotated[UploadFile, File()], - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_current_user), ): await authorize(current_user, FilePolicy.put_object(bucket)) @@ -140,7 +129,7 @@ async def put_file( try: file_data = await file.read() response = await files.put_object( - s3_client, + storage, bucket, object, data=file_data, @@ -159,14 +148,14 @@ async def list_objects_endpoint( prefix: str, recursive: bool = True, start_after: str | None = None, - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_optional_current_user), ): await authorize(current_user, FilePolicy.list(bucket)) try: objects = await files.list_objects( - s3_client, + storage, bucket, prefix=prefix, recursive=recursive, @@ -185,13 +174,13 @@ async def delete_files( *, bucket: str, object: str, - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), current_user: User = Security(auth.get_current_user), ): await authorize(current_user, FilePolicy.delete(bucket)) try: - await files.delete_object(s3_client, bucket, object) + await files.delete_object(storage, bucket, object) return {"message": "File deleted"} except HTTPException: raise diff --git a/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py b/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py index 92b22e0ae..207737987 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py @@ -17,7 +17,7 @@ Workspaces, WorkspaceUserCreate, ) -from extralit_server.contexts import accounts, files +from extralit_server.contexts import accounts, buckets, files from extralit_server.database import get_async_db from extralit_server.errors import GenericServerError from extralit_server.errors.future import NotFoundError, NotUniqueError, UnprocessableEntityError @@ -46,12 +46,12 @@ async def create_workspace( db: Annotated[AsyncSession, Depends(get_async_db)], workspace_create: WorkspaceCreate, current_user: Annotated[User, Security(auth.get_current_user)], - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), ): await authorize(current_user, WorkspacePolicy.create) try: - await files.create_bucket(s3_client, workspace_create.name) + await buckets.create(storage, workspace_create.name) except Exception as e: raise GenericServerError(e) @@ -69,7 +69,7 @@ async def delete_workspace( db: Annotated[AsyncSession, Depends(get_async_db)], workspace_id: UUID, current_user: Annotated[User, Security(auth.get_current_user)], - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), ): await authorize(current_user, WorkspacePolicy.delete) @@ -79,7 +79,7 @@ async def delete_workspace( raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=str(e)) try: - await files.delete_bucket(s3_client, workspace.name) + await buckets.delete(storage, workspace.name) except Exception as e: # Log the error but continue with workspace deletion print(f"Error deleting bucket for workspace {workspace.name}: {e!s}") @@ -175,7 +175,7 @@ async def workspace_doctor( db: Annotated[AsyncSession, Depends(get_async_db)], workspace_id: UUID, current_user: Annotated[User, Security(auth.get_current_user)], - s3_client=Depends(files.get_s3_client), + storage=Depends(files.get_storage), autofix: bool = True, ): """ @@ -191,7 +191,7 @@ async def workspace_doctor( checks = [] # Check 1: S3 bucket exists - bucket_exists = await files.bucket_exists(s3_client, workspace.name) + bucket_exists = await buckets.exists(storage, workspace.name) if bucket_exists: checks.append( WorkspaceDoctorCheckResult( @@ -204,7 +204,7 @@ async def workspace_doctor( else: if autofix: try: - await files.create_bucket(s3_client, workspace.name) + await buckets.create(storage, workspace.name) checks.append( WorkspaceDoctorCheckResult( check_name="s3_bucket", diff --git a/extralit-server/src/extralit_server/api/schemas/v1/files.py b/extralit-server/src/extralit_server/api/schemas/v1/files.py index af681ccc2..6db7f9ebc 100644 --- a/extralit-server/src/extralit_server/api/schemas/v1/files.py +++ b/extralit-server/src/extralit_server/api/schemas/v1/files.py @@ -2,7 +2,6 @@ from typing import Any from pydantic import BaseModel, Field, field_validator -from urllib3._collections import HTTPHeaderDict class ObjectMetadata(BaseModel): @@ -12,16 +11,10 @@ class ObjectMetadata(BaseModel): etag: str | None = None size: int | None = None content_type: str | None = None + # obstore returns user attributes under their own keys; the botocore-era `x-amz-meta-` + # prefix stripping this used to do would now discard every one of them. metadata: dict[str, Any] | None = None - @field_validator("metadata", mode="before") - def parse_metadata(cls, v): - if v and isinstance(v, HTTPHeaderDict | dict): - v = {key[11:]: value for key, value in v.items() if key.lower().startswith("x-amz-meta-")} - else: - v = None - return v - class ListObjectsResponse(BaseModel): objects: list[ObjectMetadata] = Field(default_factory=list) diff --git a/extralit-server/src/extralit_server/cli/database/users/create.py b/extralit-server/src/extralit_server/cli/database/users/create.py index d6a76f694..8601b900e 100644 --- a/extralit-server/src/extralit_server/cli/database/users/create.py +++ b/extralit-server/src/extralit_server/cli/database/users/create.py @@ -6,7 +6,7 @@ from extralit_server.api.schemas.v1.users import USER_PASSWORD_MIN_LENGTH, UserCreate from extralit_server.api.schemas.v1.workspaces import WorkspaceCreate -from extralit_server.contexts import accounts, files +from extralit_server.contexts import accounts, buckets, files from extralit_server.database import AsyncSessionLocal from extralit_server.models import User, UserRole @@ -88,11 +88,11 @@ async def _create( # Create S3 buckets for each workspace if they don't exist if workspace: - s3_client = await files.get_s3_client() - if s3_client is not None: + storage = await files.get_storage() + if storage is not None: for workspace_name in workspace: try: - await files.create_bucket(s3_client, workspace_name) + await buckets.create(storage, workspace_name) typer.echo(f"✓ Created/verified bucket for workspace: {workspace_name}") except Exception as e: typer.echo(f"⚠ Warning: Failed to create bucket for workspace {workspace_name}: {e}") diff --git a/extralit-server/src/extralit_server/cli/database/users/create_default.py b/extralit-server/src/extralit_server/cli/database/users/create_default.py index f3e4720dc..c4b3f7dee 100644 --- a/extralit-server/src/extralit_server/cli/database/users/create_default.py +++ b/extralit-server/src/extralit_server/cli/database/users/create_default.py @@ -3,7 +3,7 @@ import typer from extralit_server.constants import DEFAULT_API_KEY, DEFAULT_PASSWORD, DEFAULT_USERNAME -from extralit_server.contexts import accounts, files +from extralit_server.contexts import accounts, buckets, files from extralit_server.database import AsyncSessionLocal from extralit_server.models import User, UserRole @@ -32,11 +32,11 @@ async def _create_default(api_key: str, password: str, quiet: bool): ) if workspaces: - s3_client = await files.get_s3_client() - if s3_client is not None: + storage = await files.get_storage() + if storage is not None: for workspace in workspaces: try: - await files.create_bucket(s3_client, workspace.name) + await buckets.create(storage, workspace.name) typer.echo(f"✓ Created/verified bucket for workspace: {workspace.name}") if not quiet else None except Exception as e: typer.echo( diff --git a/extralit-server/src/extralit_server/contexts/buckets.py b/extralit-server/src/extralit_server/contexts/buckets.py new file mode 100644 index 000000000..4c47cf51f --- /dev/null +++ b/extralit-server/src/extralit_server/contexts/buckets.py @@ -0,0 +1,117 @@ +"""Bucket lifecycle, the one thing obstore cannot do. + +An obstore store binds to an existing bucket and exposes no admin API, so CreateBucket / +DeleteBucket / HeadBucket stay on aioboto3. This module is the whole of that surface; everything +else in the object-storage path goes through `contexts.files`. +""" + +import logging +import shutil +from pathlib import Path + +import aioboto3 +from botocore.exceptions import ClientError +from fastapi import HTTPException + +from extralit_server.contexts.files import ObjectStorage, list_objects, workspace_root +from extralit_server.settings import settings + +_LOGGER = logging.getLogger(__name__) + + +def _remote() -> bool: + return all([settings.s3_endpoint, settings.s3_access_key, settings.s3_secret_key]) + + +def _local_root(bucket: str) -> Path: + return Path(settings.home_path) / bucket + + +def _client(): + endpoint = settings.s3_endpoint or "" + session = aioboto3.Session( + aws_access_key_id=settings.s3_access_key, + aws_secret_access_key=settings.s3_secret_key, + region_name=settings.s3_region or "us-east-1", + ) + return session.client("s3", endpoint_url=endpoint, use_ssl=endpoint.startswith("https://")) + + +async def exists(storage: ObjectStorage, workspace_name: str) -> bool: + """Check whether the workspace's bucket exists.""" + bucket, _prefix = workspace_root(workspace_name) + + if not _remote(): + return _local_root(bucket).is_dir() + + try: + async with _client() as client: + await client.head_bucket(Bucket=bucket) + return True + except ClientError as e: + if e.response["Error"]["Code"] in ["404", "NoSuchBucket"]: + return False + # For other errors (like permissions), log and return False + _LOGGER.warning(f"Error checking bucket {bucket}: {e}") + return False + except Exception as e: + _LOGGER.warning(f"Unexpected error checking bucket {bucket}: {e}") + return False + + +async def create(storage: ObjectStorage, workspace_name: str) -> None: + """Create the workspace's bucket if it does not already exist.""" + bucket, _prefix = workspace_root(workspace_name) + + if not _remote(): + _local_root(bucket).mkdir(parents=True, exist_ok=True) + return + + try: + async with _client() as client: + try: + await client.create_bucket(Bucket=bucket) + except ClientError as e: + if e.response["Error"]["Code"] not in ["BucketAlreadyOwnedByYou", "BucketAlreadyExists"]: + raise + except ClientError as e: + _LOGGER.error(f"Error creating bucket {bucket}: {e}") + raise HTTPException(status_code=500, detail=f"Error creating bucket: {e!s}") + except Exception as e: + _LOGGER.error(f"Error creating bucket {workspace_name}: {e}") + raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") + + +async def delete(storage: ObjectStorage, workspace_name: str) -> None: + """Drain the workspace's bucket and delete it.""" + bucket, _prefix = workspace_root(workspace_name) + + if not _remote(): + shutil.rmtree(_local_root(bucket), ignore_errors=True) + return + + store = storage.store_for(bucket) + try: + listing = await list_objects(storage, bucket) + except HTTPException: + return # No such bucket; nothing to drain or delete. + + for object in listing.objects: + try: + await store.delete_async(object.object_name) + except Exception as remove_err: + _LOGGER.warning(f"Error removing object {object.object_name} during bucket delete: {remove_err}") + + try: + async with _client() as client: + await client.delete_bucket(Bucket=bucket) + _LOGGER.info(f"Successfully deleted bucket: {bucket}") + except ClientError as e: + if e.response["Error"]["Code"] in ["NoSuchBucket", "NotImplemented"]: + pass # Bucket doesn't exist, that's fine + else: + _LOGGER.error(f"Error deleting S3 bucket {bucket}: {e}") + raise HTTPException(status_code=500, detail=f"Error deleting bucket: {e!s}") + except Exception as e: + _LOGGER.error(f"Error deleting S3 bucket {bucket}: {e}") + raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") diff --git a/extralit-server/src/extralit_server/contexts/files.py b/extralit-server/src/extralit_server/contexts/files.py index 6c2308e17..626f3821d 100644 --- a/extralit-server/src/extralit_server/contexts/files.py +++ b/extralit-server/src/extralit_server/contexts/files.py @@ -1,19 +1,29 @@ import hashlib import logging -from collections.abc import AsyncGenerator -from typing import TYPE_CHECKING, Any, BinaryIO +import mimetypes +from datetime import timedelta +from pathlib import Path +from typing import Any, BinaryIO from uuid import UUID -from botocore.exceptions import ClientError +import obstore from fastapi import HTTPException +from obstore.exceptions import BaseError as ObjectStoreError +from obstore.store import LocalStore, S3Store from extralit_server.api.schemas.v1.files import FileObjectResponse, ListObjectsResponse, ObjectMetadata from extralit_server.helpers import shared_resources - -if TYPE_CHECKING: - from types_aiobotocore_s3.client import S3Client - -CHUNK_LENGTH_MB = 10 * 1024 * 1024 +from extralit_server.settings import settings + +# LocalStore cannot persist attributes (obstore raises NotImplementedError for `put_opts` with +# attributes), so in local mode the content type has to be recovered from the key. These are the +# extension-less prefixes minted below; everything else is guessed from the filename. +_CONTENT_TYPE_BY_PREFIX = { + "pdf/": "application/pdf", + "thumbnails/": "image/png", + "layout/": "application/json", + "schemas/": "application/json", +} _LOGGER = logging.getLogger(__name__) @@ -24,7 +34,7 @@ def workspace_root(workspace_name: str) -> tuple[str, str]: The single place that knows how a workspace maps onto storage: PDFs, thumbnails, layout JSON and the Lance datasets all address through it, so moving to one bucket with `{org}/{workspace}/` prefixes is a change here rather than at every call site. Today it is bucket-per-workspace, - which is why the prefix is empty and `files` can still pass `Bucket=workspace_name` directly. + which is why the prefix is empty and `files` can still pass the workspace name as the bucket. """ if not workspace_name: raise ValueError("workspace_name cannot be empty") @@ -32,55 +42,86 @@ def workspace_root(workspace_name: str) -> tuple[str, str]: return workspace_name, "" -async def get_s3_client() -> "S3Client": - """Dependency function to get shared S3 client.""" - s3_client = shared_resources.get("s3_client") - if s3_client is None: - from extralit_server.helpers import create_s3_client +class ObjectStorage: + """Resolves a bucket name to an obstore store, caching one store per bucket. - try: - s3_client = await create_s3_client() - shared_resources["s3_client"] = s3_client - except ValueError as e: - raise HTTPException(status_code=500, detail=str(e)) from e + obstore binds a store to a single bucket at construction, so the per-call `bucket` argument + every function in this module takes is a lookup here rather than a request parameter. + """ - return s3_client + def __init__(self) -> None: + self._stores: dict[str, S3Store | LocalStore] = {} + self._remote = all([settings.s3_endpoint, settings.s3_access_key, settings.s3_secret_key]) + + @property + def signable(self) -> bool: + """Only S3 can presign; `LocalStore` is not a `SignCapableStore`.""" + return self._remote + + def store_for(self, bucket: str) -> S3Store | LocalStore: + store = self._stores.get(bucket) + if store is None: + store = self._build(bucket) + self._stores[bucket] = store + return store + + def _build(self, bucket: str) -> S3Store | LocalStore: + if not self._remote: + return LocalStore(prefix=Path(settings.home_path) / bucket, mkdir=True) + + endpoint = settings.s3_endpoint or "" + return S3Store( + bucket, + endpoint=endpoint, + access_key_id=settings.s3_access_key, + secret_access_key=settings.s3_secret_key, + region=settings.s3_region or "us-east-1", + virtual_hosted_style_request=False, + client_options={"allow_http": not endpoint.startswith("https://")}, + ) + async def aclose(self) -> None: + self._stores.clear() -async def get_file_chunk( - s3_client: "S3Client", bucket_name: str, key: str, chunk_length: int -) -> AsyncGenerator[bytes, None]: - """Async generator to get file chunks for streaming.""" - head = await s3_client.head_object(Bucket=bucket_name, Key=key) - content_length = head["ContentLength"] - for offset in range(0, content_length, chunk_length): - end = min(offset + chunk_length - 1, content_length - 1) - s3_file = await s3_client.get_object(Bucket=bucket_name, Key=key, Range=f"bytes={offset}-{end}") +async def get_storage() -> ObjectStorage: + """Dependency function to get the shared object storage.""" + storage = shared_resources.get("storage") + if storage is None: + storage = ObjectStorage() + shared_resources["storage"] = storage - async with s3_file["Body"] as stream: - yield await stream.read() + return storage -async def _put_object_to_s3( - s3_client: "S3Client", - bucket: str, - key: str, - data: BinaryIO | bytes, - content_type: str, - metadata: dict[str, Any] | None = None, -): - """Put object to S3.""" - kwargs = { - "Bucket": bucket, - "Key": key, - "Body": data, - "ContentType": content_type, - } - if metadata: - kwargs["Metadata"] = metadata +def content_type_of(key: str, attributes: Any = None) -> str: + declared = dict(attributes or {}).get("Content-Type") + if declared: + return declared - return await s3_client.put_object(**kwargs) + for prefix, content_type in _CONTENT_TYPE_BY_PREFIX.items(): + if key.startswith(prefix): + return content_type + + guessed, _ = mimetypes.guess_type(key) + return guessed or "application/octet-stream" + + +def _user_metadata(attributes: Any = None) -> dict[str, str]: + return {key: value for key, value in dict(attributes or {}).items() if key != "Content-Type"} + + +def _object_metadata(bucket: str, meta: Any, attributes: Any = None) -> ObjectMetadata: + key = meta["path"] + return ObjectMetadata( + bucket_name=bucket, + object_name=key, + etag=(meta["e_tag"] or "").strip('"') or None, + size=meta["size"], + last_modified=meta["last_modified"], + content_type=content_type_of(key, attributes), + metadata=_user_metadata(attributes), + ) def compute_hash(data: bytes) -> str: @@ -124,146 +165,129 @@ def get_proxy_document_url(bucket_name: str, object_path: str) -> str: return f"/api/v1/file/{bucket_name}/{object_path}" -async def get_presigned_url_from_document_url(s3_client, document_url: str, expires: int = 3600) -> str: - """ - Generate a presigned URL from a document URL by parsing the bucket_name and object_path. +def split_document_url(document_url: str) -> tuple[str, str]: + """Split `/api/v1/file/{bucket}/{object}` back into its bucket and key.""" + prefix = "/api/v1/file/" + if not document_url.startswith(prefix): + raise ValueError(f"Invalid document URL format: {document_url}") - Args: - s3_client: aioboto3 S3 client - document_url: URL in format "/api/v1/file/{bucket_name}/{object_path}" - expires: Expiration time in seconds (default: 1 hour) + parts = document_url[len(prefix) :].split("/", 1) + if len(parts) != 2: + raise ValueError(f"Invalid document URL format: {document_url}") - Returns: - Presigned URL if successful, original URL if parsing fails + return parts[0], parts[1] + + +async def get_presigned_url_from_document_url(storage: ObjectStorage, document_url: str, expires: int = 3600) -> str: + """Presign a `/api/v1/file/{bucket}/{object}` URL, valid for `expires` seconds. + + Local storage cannot sign, so it keeps serving through the proxy route. """ try: - # Parse the URL to extract bucket_name and object_path - # Expected format: "/api/v1/file/{bucket_name}/{object_path}" - if not document_url.startswith("/api/v1/file/"): - _LOGGER.warning(f"Invalid document URL format: {document_url}") - return document_url - - path_parts = document_url[13:].split("/", 1) # 13 = len("/api/v1/file/") - if len(path_parts) != 2: - _LOGGER.warning(f"Invalid document URL format: {document_url}") - return document_url - - bucket_name, object_path = path_parts - - presigned_url = await s3_client.generate_presigned_url( - "get_object", - Params={"Bucket": bucket_name, "Key": object_path}, - ExpiresIn=expires, - ) - return presigned_url + bucket, object_path = split_document_url(document_url) + except ValueError: + _LOGGER.warning(f"Invalid document URL format: {document_url}") + return document_url + + if not storage.signable: + return document_url + try: + return await obstore.sign_async(storage.store_for(bucket), "GET", object_path, timedelta(seconds=expires)) except Exception as e: _LOGGER.error(f"Error generating presigned URL from document URL {document_url}: {e}") return document_url async def list_objects( - s3_client: "S3Client", + storage: ObjectStorage, bucket: str, prefix: str | None = None, recursive=True, start_after: str | None = None, ) -> ListObjectsResponse: - """List objects in S3 bucket and return as ListObjectsResponse.""" + """List objects in a bucket and return as ListObjectsResponse.""" + store = storage.store_for(bucket) try: - kwargs = {"Bucket": bucket} - if prefix: - kwargs["Prefix"] = prefix - if not recursive: - kwargs["Delimiter"] = "/" - if start_after: - kwargs["StartAfter"] = start_after - - response = await s3_client.list_objects_v2(**kwargs) - - objects = [ - ObjectMetadata( - bucket_name=bucket, - object_name=obj.get("Key") or "", - etag=obj.get("ETag", "").strip('"'), - size=obj.get("Size"), - last_modified=obj.get("LastModified"), - content_type="application/octet-stream", # Default, would need head_object for actual - metadata={}, - ) - for obj in response.get("Contents", []) - ] - - return ListObjectsResponse(objects=objects) - except ClientError as e: + if recursive: + metas = [] + async for batch in store.list(prefix, offset=start_after): + metas.extend(batch) + else: + result = await store.list_with_delimiter_async(prefix) + metas = list(result["objects"]) + if start_after: + metas = [meta for meta in metas if meta["path"] > start_after] + + # Attributes are not returned by listing APIs, so the content type is key-derived here. + return ListObjectsResponse(objects=[_object_metadata(bucket, meta) for meta in metas]) + except FileNotFoundError: + _LOGGER.error(f"Bucket '{bucket}' not found") + raise HTTPException(status_code=404, detail=f"Bucket '{bucket}' not found") + except ObjectStoreError as e: _LOGGER.error(f"Error listing objects in bucket {bucket}: {e}") raise HTTPException(status_code=404, detail=f"Bucket '{bucket}' not found") -async def get_object( - s3_client: "S3Client", - bucket: str, - object: str, -) -> FileObjectResponse: - """Get object from S3 and return as FileObjectResponse.""" +async def get_object(storage: ObjectStorage, bucket: str, object: str) -> FileObjectResponse: + """Get an object and return it as a FileObjectResponse whose `response` streams.""" try: - head_response = await s3_client.head_object(Bucket=bucket, Key=object) - get_response = await s3_client.get_object(Bucket=bucket, Key=object) - - metadata = ObjectMetadata( - bucket_name=bucket, - object_name=object, - etag=head_response["ETag"].strip('"'), - size=head_response["ContentLength"], - last_modified=head_response["LastModified"], - content_type=head_response.get("ContentType", "application/octet-stream"), - metadata=head_response.get("Metadata", {}), + result = await storage.store_for(bucket).get_async(object) + return FileObjectResponse( + response=result, + metadata=_object_metadata(bucket, result.meta, result.attributes), ) + except FileNotFoundError: + _LOGGER.error(f"Object {object} not found in bucket {bucket}") + raise HTTPException(status_code=404, detail=f"Object {object} not found in bucket {bucket}") + except ObjectStoreError as e: + _LOGGER.error(f"Error getting object {object} from bucket {bucket}: {e}") + raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") - return FileObjectResponse(response=get_response["Body"], metadata=metadata) - except ClientError as e: - if e.response["Error"]["Code"] == "NoSuchKey": - _LOGGER.error(f"Object {object} not found in bucket {bucket}") - raise HTTPException(status_code=404, detail=f"Object {object} not found in bucket {bucket}") - else: - _LOGGER.error(f"Error getting object {object} from bucket {bucket}: {e}") - raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") +async def _put( + storage: ObjectStorage, + bucket: str, + key: str, + data: bytes, + content_type: str, + metadata: dict[str, Any] | None = None, +): + attributes = {"Content-Type": content_type, **{k: str(v) for k, v in (metadata or {}).items()}} + store = storage.store_for(bucket) + if isinstance(store, LocalStore): + # LocalStore rejects attributes outright; the content type is recovered from the key. + return await store.put_async(key, data) + + return await store.put_async(key, data, attributes=attributes) async def put_object( - s3_client: "S3Client", + storage: ObjectStorage, bucket: str, object: str, data: BinaryIO | bytes | str, content_type: str = "application/octet-stream", metadata: dict[str, Any] | None = None, ) -> ObjectMetadata: - """Put object to S3 and return ObjectMetadata.""" - try: - # Prepare data - if isinstance(data, str): - data = data.encode("utf-8") - elif hasattr(data, "read"): # File-like object - data = data.read() - - # Upload to S3 - await _put_object_to_s3(s3_client, bucket, object, data, content_type, metadata) + """Put an object and return its ObjectMetadata.""" + if isinstance(data, str): + data = data.encode("utf-8") + elif hasattr(data, "read"): + data = data.read() - # Get metadata for response - head_response = await s3_client.head_object(Bucket=bucket, Key=object) + try: + result = await _put(storage, bucket, object, data, content_type, metadata) return ObjectMetadata( bucket_name=bucket, object_name=object, - etag=head_response["ETag"].strip('"'), - size=head_response["ContentLength"], - last_modified=head_response["LastModified"], - content_type=head_response.get("ContentType", content_type), - metadata=head_response.get("Metadata", {}), + etag=(result["e_tag"] or "").strip('"') or None, + size=len(data), + content_type=content_type, + metadata=metadata or {}, ) - - except ClientError as e: + except ObjectStoreError as e: _LOGGER.error(f"Error putting object {object} in bucket {bucket}: {e}") raise HTTPException(status_code=500, detail=f"Error uploading file: {e!s}") except Exception as e: @@ -271,11 +295,13 @@ async def put_object( raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") -async def delete_object(s3_client, bucket: str, object: str): - """Delete object from S3.""" +async def delete_object(storage: ObjectStorage, bucket: str, object: str): + """Delete an object. Deleting a key that is already gone is not an error.""" try: - await s3_client.delete_object(Bucket=bucket, Key=object) - except ClientError as e: + await storage.store_for(bucket).delete_async(object) + except FileNotFoundError: + pass + except ObjectStoreError as e: _LOGGER.error(f"Error deleting object {object} from bucket {bucket}: {e}") raise HTTPException(status_code=500, detail=f"Error deleting file: {e!s}") except Exception as e: @@ -283,7 +309,7 @@ async def delete_object(s3_client, bucket: str, object: str): raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") -async def delete_document_artifacts(s3_client: "S3Client", workspace_name: str, document_id: UUID | str) -> None: +async def delete_document_artifacts(storage: ObjectStorage, workspace_name: str, document_id: UUID | str) -> None: """Remove every artifact of a document: PDF, thumbnail, layout JSON and layout rows. Best effort — the DB rows are already gone by the time this runs, so a storage hiccup leaves a @@ -291,56 +317,22 @@ async def delete_document_artifacts(s3_client: "S3Client", workspace_name: str, be ignored: the objects are unreachable and layout rows only survive until the document is re-parsed or a sweeper runs. """ - from extralit_server.contexts.ocr import storage + from extralit_server.contexts.ocr import storage as layout_storage for object_path in (get_pdf_s3_object_path(document_id), get_thumbnail_s3_object_path(document_id)): try: - await delete_object(s3_client, workspace_name, object_path) + await delete_object(storage, workspace_name, object_path) except Exception as e: _LOGGER.warning(f"Could not delete {object_path} for document {document_id}: {e}") try: - await storage.delete_layout(s3_client, workspace_name, document_id) + await layout_storage.delete_layout(storage, workspace_name, document_id) except Exception as e: _LOGGER.warning(f"Could not delete layout artifacts for document {document_id}: {e}") -async def bucket_exists(s3_client: "S3Client", bucket_name: str) -> bool: - """Check if S3 bucket exists.""" - try: - await s3_client.head_bucket(Bucket=bucket_name) - return True - except ClientError as e: - if e.response["Error"]["Code"] in ["404", "NoSuchBucket"]: - return False - # For other errors (like permissions), log and return False - _LOGGER.warning(f"Error checking bucket {bucket_name}: {e}") - return False - except Exception as e: - _LOGGER.warning(f"Unexpected error checking bucket {bucket_name}: {e}") - return False - - -async def create_bucket(s3_client: "S3Client", workspace_name: str): - """Create the workspace's bucket if it does not already exist.""" - bucket, _prefix = workspace_root(workspace_name) - try: - try: - await s3_client.create_bucket(Bucket=bucket) - except ClientError as e: - if e.response["Error"]["Code"] not in ["BucketAlreadyOwnedByYou", "BucketAlreadyExists"]: - raise - - except ClientError as e: - _LOGGER.error(f"Error creating bucket {bucket}: {e}") - raise HTTPException(status_code=500, detail=f"Error creating bucket: {e!s}") - except Exception as e: - _LOGGER.error(f"Error creating bucket {workspace_name}: {e}") - raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") - - async def put_document_file( - s3_client: "S3Client", + storage: ObjectStorage, workspace_name: str, document_id: UUID, file_data: bytes, @@ -349,24 +341,15 @@ async def put_document_file( metadata: dict[str, Any] | None = None, ) -> str | None: """ - Upload a document file to S3 with deduplication. - - Args: - s3_client: aioboto3 S3 client - workspace_name: Name of the workspace bucket - document_id: UUID of the document - file_data: File data as bytes - filename: Original filename - metadata: Optional metadata to store with the file + Upload a document file with deduplication. Returns: - S3 object URL if file was uploaded, None if file already exists with same hash + The proxy object URL if the file was uploaded, None if an identical file already exists. """ object_path = get_pdf_s3_object_path(document_id) - # Check if file already exists with same hash try: - existing_files = await list_objects(s3_client, workspace_name, prefix=object_path, recursive=False) + existing_files = await list_objects(storage, workspace_name, prefix=object_path, recursive=False) should_upload = True if existing_files.objects: @@ -379,14 +362,7 @@ async def put_document_file( should_upload = False if should_upload: - await _put_object_to_s3( - s3_client, - workspace_name, - object_path, - file_data, - content_type, - metadata, - ) + await _put(storage, workspace_name, object_path, file_data, content_type, metadata) return get_proxy_document_url(workspace_name, object_path) @@ -397,58 +373,16 @@ async def put_document_file( raise HTTPException(status_code=500, detail=f"Error uploading document: {e!s}") -async def download_file_content(s3_client, document_url: str) -> bytes: - """ - Download file content from a document URL. - - Args: - s3_client: aioboto3 S3 client - document_url: URL in format "/api/v1/file/{bucket_name}/{object_path}" - - Returns: - File content as bytes - """ - # Parse URL to get bucket and object path - if not document_url.startswith("/api/v1/file/"): - raise ValueError(f"Invalid document URL format: {document_url}") - - url_parts = document_url.replace("/api/v1/file/", "").split("/", 1) - if len(url_parts) != 2: - raise ValueError(f"Invalid document URL format: {document_url}") - - bucket_name, object_path = url_parts +async def download_file_content(storage: ObjectStorage, document_url: str) -> bytes: + """Download the whole body behind a `/api/v1/file/{bucket}/{object}` URL.""" + bucket_name, object_path = split_document_url(document_url) try: - response = await s3_client.get_object(Bucket=bucket_name, Key=object_path) - return await response["Body"].read() - except ClientError as e: + result = await storage.store_for(bucket_name).get_async(object_path) + return bytes(await result.bytes_async()) + except FileNotFoundError: + _LOGGER.error(f"File not found: {document_url}") + raise HTTPException(status_code=404, detail=f"File not found: {document_url}") + except ObjectStoreError as e: _LOGGER.error(f"Error downloading file content from {document_url}: {e}") raise HTTPException(status_code=404, detail=f"File not found: {document_url}") - - -async def delete_bucket(s3_client, workspace_name: str): - """Delete S3 bucket and all its contents.""" - try: - # First, delete all objects in the bucket - response = await s3_client.list_objects_v2(Bucket=workspace_name) - - if "Contents" in response: - for obj in response["Contents"]: - try: - await s3_client.delete_object(Bucket=workspace_name, Key=obj["Key"]) - except ClientError as remove_err: - _LOGGER.warning(f"Error removing object {obj['Key']} during bucket delete: {remove_err}") - - # Then delete the bucket itself - await s3_client.delete_bucket(Bucket=workspace_name) - _LOGGER.info(f"Successfully deleted bucket: {workspace_name}") - - except ClientError as e: - if e.response["Error"]["Code"] in ["NoSuchBucket", "NotImplemented"]: - pass # Bucket doesn't exist, that's fine - else: - _LOGGER.error(f"Error deleting S3 bucket {workspace_name}: {e}") - raise HTTPException(status_code=500, detail=f"Error deleting bucket: {e!s}") - except Exception as e: - _LOGGER.error(f"Error deleting S3 bucket {workspace_name}: {e}") - raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") diff --git a/extralit-server/src/extralit_server/contexts/imports.py b/extralit-server/src/extralit_server/contexts/imports.py index 25f54cc0c..695d0c9cd 100644 --- a/extralit-server/src/extralit_server/contexts/imports.py +++ b/extralit-server/src/extralit_server/contexts/imports.py @@ -374,8 +374,8 @@ async def process_bulk_upload( ) reference_to_doc[doc.reference] = doc - s3_client = shared_resources.get("s3_client") - if s3_client is None: + storage = shared_resources.get("storage") + if storage is None: raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="S3 client not available") # Process each reference: upload files to S3, create documents, start workflows @@ -432,7 +432,7 @@ async def process_bulk_upload( # Upload file to S3 file_url = await file_context.put_document_file( - s3_client=s3_client, + storage=storage, workspace_name=workspace.name, document_id=document_new.id, # type: ignore[arg-type] file_data=await file.read(), diff --git a/extralit-server/src/extralit_server/contexts/ocr/storage.py b/extralit-server/src/extralit_server/contexts/ocr/storage.py index 6f71b7ddd..087c5da61 100644 --- a/extralit-server/src/extralit_server/contexts/ocr/storage.py +++ b/extralit-server/src/extralit_server/contexts/ocr/storage.py @@ -9,19 +9,17 @@ import json import logging -from typing import TYPE_CHECKING, Any, Optional +from typing import Any, Optional from uuid import UUID from anyio import to_thread from docling_core.types.doc import DoclingDocument from extralit_server.contexts import files +from extralit_server.contexts.files import ObjectStorage from extralit_server.contexts.ocr.arrow import items_table, pages_table from extralit_server.contexts.ocr.layout_store import LAYOUT_PREFIX, LayoutStore -if TYPE_CHECKING: - from types_aiobotocore_s3 import S3Client - _LOGGER = logging.getLogger("extralit_server.contexts.ocr.storage") @@ -30,7 +28,7 @@ def layout_object_path(document_id: UUID | str) -> str: async def store_layout( - s3_client: S3Client, + storage: ObjectStorage, workspace_name: str, document_id: UUID | str, doc: DoclingDocument, @@ -46,7 +44,7 @@ async def store_layout( store = store or LayoutStore.for_workspace(workspace_name) await files.put_object( - s3_client, + storage, workspace_name, layout_url, json.dumps(doc.export_to_dict(), ensure_ascii=False), @@ -69,7 +67,7 @@ async def store_layout( async def delete_layout( - s3_client: S3Client, + storage: ObjectStorage, workspace_name: str, document_id: UUID | str, store: Optional[LayoutStore] = None, @@ -77,7 +75,7 @@ async def delete_layout( """Drop both artifacts. Orphaned rows would skew workspace aggregates, but a failure here is negligible: they are superseded on the next parse of that document.""" try: - await files.delete_object(s3_client, workspace_name, layout_object_path(document_id)) + await files.delete_object(storage, workspace_name, layout_object_path(document_id)) except Exception as error: _LOGGER.warning(f"Could not delete layout JSON for document {document_id}: {error}") @@ -90,7 +88,7 @@ async def delete_layout( async def load_layout( - s3_client: S3Client, + storage: ObjectStorage, workspace_name: str, document_id: UUID | str, object_path: Optional[str] = None, @@ -101,6 +99,6 @@ async def load_layout( demands an equal major and a minor no higher than the SDK's. """ key = object_path or layout_object_path(document_id) - response = await s3_client.get_object(Bucket=workspace_name, Key=key) - raw = await response["Body"].read() + file = await files.get_object(storage, workspace_name, key) + raw = bytes(await file.response.bytes_async()) return DoclingDocument.model_validate(json.loads(raw)) diff --git a/extralit-server/src/extralit_server/contexts/schema_versions.py b/extralit-server/src/extralit_server/contexts/schema_versions.py index 25e48cdd3..56414a95b 100644 --- a/extralit-server/src/extralit_server/contexts/schema_versions.py +++ b/extralit-server/src/extralit_server/contexts/schema_versions.py @@ -6,7 +6,7 @@ table is the queryable column manifest and there is no cached copy of it. """ -from typing import TYPE_CHECKING, Any +from typing import Any from uuid import UUID import pandera.pandas as pa @@ -14,15 +14,13 @@ from sqlalchemy.ext.asyncio import AsyncSession from extralit_server.contexts import files as files_ctx +from extralit_server.contexts.files import ObjectStorage from extralit_server.enums import FieldType from extralit_server.errors.future import UnprocessableEntityError from extralit_server.models.database import Dataset, Field, SchemaVersion from extralit_server.webhooks.v1.datasets import notify_dataset_event as notify_dataset_event_v1 from extralit_server.webhooks.v1.enums import DatasetEvent -if TYPE_CHECKING: - from types_aiobotocore_s3.client import S3Client - def object_key_for(dataset_id: UUID, version: int) -> str: return f"schemas/{dataset_id}/v{version}.json" @@ -140,7 +138,7 @@ async def _reject_incompatible_columns( async def publish_version( db: AsyncSession, - s3_client: "S3Client", + storage: ObjectStorage, dataset: Dataset, *, body: str, @@ -156,7 +154,7 @@ async def publish_version( next_version = await _next_version_number(db, dataset.id) key = object_key_for(dataset.id, next_version) - metadata = await files_ctx.put_object(s3_client, bucket, key, body, content_type="application/json") + metadata = await files_ctx.put_object(storage, bucket, key, body, content_type="application/json") parent_id = dataset.current_schema_version_id diff --git a/extralit-server/src/extralit_server/helpers.py b/extralit-server/src/extralit_server/helpers.py index fcc239ed6..3d270a51d 100644 --- a/extralit-server/src/extralit_server/helpers.py +++ b/extralit-server/src/extralit_server/helpers.py @@ -2,534 +2,12 @@ Common helper functions """ -import hashlib -import io -import json import logging -import os -import uuid -from datetime import datetime -from pathlib import Path -from typing import Any - -import aioboto3 -import aiofiles -import aiofiles.os -from botocore.exceptions import ClientError -from types_aiobotocore_s3.client import S3Client - -from extralit_server.settings import settings _LOGGER = logging.getLogger("extralit_server") shared_resources = {} -def _compute_hash(data: bytes) -> str: - """Compute MD5 hash for data.""" - return hashlib.md5(data).hexdigest() - - -class LocalFileClient(S3Client): - """Local file storage implementation that mimics S3Client interface.""" - - def __init__(self, base_dir: str | Path): - self.base_dir = Path(base_dir) - - async def _ensure_base_dir(self): - """Ensure base directory exists.""" - await aiofiles.os.makedirs(self.base_dir, exist_ok=True) - - def _get_bucket_path(self, bucket_name: str) -> Path: - """Get bucket directory path.""" - return self.base_dir / bucket_name - - def _get_object_path(self, bucket_name: str, key: str) -> Path: - """Get object file path.""" - bucket_path = self._get_bucket_path(bucket_name) - return bucket_path / key - - def _get_version_path(self, bucket_name: str, key: str) -> Path: - """Get versions directory path for an object.""" - bucket_path = self._get_bucket_path(bucket_name) - return bucket_path / ".versions" / key - - def _get_metadata_path(self, bucket_name: str, key: str) -> Path: - """Get metadata file path for an object.""" - object_path = self._get_object_path(bucket_name, key) - return object_path.with_suffix(object_path.suffix + ".metadata.json") - - async def create_bucket(self, Bucket: str, **kwargs) -> dict[str, Any]: - """Create a bucket (directory).""" - bucket_path = self._get_bucket_path(Bucket) - await aiofiles.os.makedirs(bucket_path, exist_ok=True) - # Create versions directory - versions_path = bucket_path / ".versions" - await aiofiles.os.makedirs(versions_path, exist_ok=True) - return {} - - async def put_bucket_versioning(self, Bucket: str, VersioningConfiguration: dict, **kwargs) -> dict[str, Any]: - """Enable bucket versioning (just ensure versions directory exists).""" - bucket_path = self._get_bucket_path(Bucket) - versions_path = bucket_path / ".versions" - await aiofiles.os.makedirs(versions_path, exist_ok=True) - return {} - - async def head_object(self, Bucket: str, Key: str, VersionId: str | None = None, **kwargs) -> dict[str, Any]: - """Get object metadata.""" - try: - if VersionId: - version_path = self._get_version_path(Bucket, Key).with_suffix(f".{VersionId}") - if not version_path.exists(): - raise ClientError( - {"Error": {"Code": "NoSuchKey", "Message": "The specified version does not exist"}}, - "HeadObject", - ) - file_path = version_path - else: - object_path = self._get_object_path(Bucket, Key) - if not object_path.exists(): - raise ClientError( - {"Error": {"Code": "NoSuchKey", "Message": "The specified key does not exist"}}, "HeadObject" - ) - file_path = object_path - - # Get file stats - stat_result = await aiofiles.os.stat(file_path) - - # Get metadata - metadata_path = self._get_metadata_path(Bucket, Key) - if metadata_path.exists(): - async with aiofiles.open(metadata_path) as f: - metadata = json.loads(await f.read()) - else: - metadata = {} - - return { - "ContentLength": stat_result.st_size, - "LastModified": datetime.fromtimestamp(stat_result.st_mtime), - "ETag": f'"{metadata.get("etag", "")}"', - "VersionId": VersionId or metadata.get("version_id"), - "ContentType": metadata.get("content_type", "application/octet-stream"), - "Metadata": metadata.get("metadata", {}), - } - - except (FileNotFoundError, OSError): - raise ClientError( - {"Error": {"Code": "NoSuchKey", "Message": "The specified key does not exist"}}, "HeadObject" - ) - - async def get_object( - self, Bucket: str, Key: str, VersionId: str | None = None, Range: str | None = None, **kwargs - ) -> dict[str, Any]: - """Get object content.""" - try: - if VersionId: - version_path = self._get_version_path(Bucket, Key).with_suffix(f".{VersionId}") - if not version_path.exists(): - raise ClientError( - {"Error": {"Code": "NoSuchKey", "Message": "The specified version does not exist"}}, "GetObject" - ) - file_path = version_path - else: - object_path = self._get_object_path(Bucket, Key) - if not object_path.exists(): - raise ClientError( - {"Error": {"Code": "NoSuchKey", "Message": "The specified key does not exist"}}, "GetObject" - ) - file_path = object_path - - # Read file content - async with aiofiles.open(file_path, "rb") as f: - if Range: - # Parse range like "bytes=0-1023" - range_match = Range.replace("bytes=", "").split("-") - start = int(range_match[0]) if range_match[0] else 0 - end = int(range_match[1]) if len(range_match) > 1 and range_match[1] else None - - await f.seek(start) - if end is not None: - content = await f.read(end - start + 1) - else: - content = await f.read() - else: - content = await f.read() - - # Get metadata for response - head_info = await self.head_object(Bucket, Key, VersionId) - - # Create a mock response body that can be read - body = MockAsyncStreamingBody(content) - - return { - "Body": body, - "ContentLength": len(content), - "LastModified": head_info["LastModified"], - "ETag": head_info["ETag"], - "ContentType": head_info["ContentType"], - "VersionId": head_info.get("VersionId"), - "Metadata": head_info.get("Metadata", {}), - } - - except (FileNotFoundError, OSError): - raise ClientError( - {"Error": {"Code": "NoSuchKey", "Message": "The specified key does not exist"}}, "GetObject" - ) - - async def put_object( - self, - Bucket: str, - Key: str, - Body: Any, - ContentType: str = "application/octet-stream", - Metadata: dict | None = None, - **kwargs, - ) -> dict[str, Any]: - """Put object to storage.""" - await self._ensure_base_dir() - bucket_path = self._get_bucket_path(Bucket) - await aiofiles.os.makedirs(bucket_path, exist_ok=True) - - # Convert body to bytes - if hasattr(Body, "read"): - if hasattr(Body, "seek"): - Body.seek(0) - data_bytes = Body.read() - elif isinstance(Body, str): - data_bytes = Body.encode("utf-8") - else: - data_bytes = Body - - # Generate version ID and hash - content_hash = _compute_hash(data_bytes) - version_id = str(uuid.uuid4()) - - # Ensure versions directory exists - version_dir = self._get_version_path(Bucket, Key).parent - await aiofiles.os.makedirs(version_dir, exist_ok=True) - - # Write to version file - version_path = self._get_version_path(Bucket, Key).with_suffix(f".{version_id}") - async with aiofiles.open(version_path, "wb") as f: - await f.write(data_bytes) - - # Update main object path (symlink or copy) - object_path = self._get_object_path(Bucket, Key) - await aiofiles.os.makedirs(object_path.parent, exist_ok=True) - - # Remove existing file/symlink - if object_path.exists(): - await aiofiles.os.remove(object_path) - - # Create symlink to version file - try: - object_path.symlink_to(version_path) - except OSError: - # Fallback to copy if symlink fails - async with aiofiles.open(version_path, "rb") as src, aiofiles.open(object_path, "wb") as dst: - content = await src.read() - await dst.write(content) - - # Save metadata - metadata_info = { - "etag": content_hash, - "content_type": ContentType, - "version_id": version_id, - "metadata": Metadata or {}, - } - - metadata_path = self._get_metadata_path(Bucket, Key) - async with aiofiles.open(metadata_path, "w") as f: - await f.write(json.dumps(metadata_info, default=str)) - - return { - "ETag": f'"{content_hash}"', - "VersionId": version_id, - } - - async def delete_object(self, Bucket: str, Key: str, VersionId: str | None = None, **kwargs) -> dict[str, Any]: - """Delete object or specific version.""" - if VersionId: - version_path = self._get_version_path(Bucket, Key).with_suffix(f".{VersionId}") - if version_path.exists(): - await aiofiles.os.remove(version_path) - else: - object_path = self._get_object_path(Bucket, Key) - if object_path.exists(): - await aiofiles.os.remove(object_path) - - # Remove metadata - metadata_path = self._get_metadata_path(Bucket, Key) - if metadata_path.exists(): - await aiofiles.os.remove(metadata_path) - - return {} - - async def list_objects_v2( - self, - Bucket: str, - Prefix: str | None = None, - Delimiter: str | None = None, - StartAfter: str | None = None, - **kwargs, - ) -> dict[str, Any]: - """List current objects in bucket.""" - bucket_path = self._get_bucket_path(Bucket) - if not bucket_path.exists(): - return {"Contents": []} - - contents = [] - - try: - # Always search from bucket root and filter by prefix - search_path = bucket_path - - if search_path.exists(): - # Use Path.rglob for recursive search (ignoring delimiter for now) - files = [f for f in search_path.rglob("*") if f.is_file()] - - for file_path in files: - # Skip metadata files and version files - if file_path.name.endswith(".metadata.json") or ".versions" in str( - file_path.relative_to(bucket_path) - ): - continue - - # Get the key relative to bucket path and normalize path separators - key = str(file_path.relative_to(bucket_path)).replace("\\", "/") - - # Apply prefix filter (ensure prefix is normalized) - if Prefix: - normalized_prefix = Prefix.replace("\\", "/").rstrip("/") - if not key.startswith(normalized_prefix): - continue - - # Apply start_after filter - if StartAfter and key <= StartAfter: - continue - - try: - stat_result = await aiofiles.os.stat(file_path) - - # Try to get metadata - metadata_path = self._get_metadata_path(Bucket, key) - etag = "" - if metadata_path.exists(): - try: - async with aiofiles.open(metadata_path) as f: - metadata = json.loads(await f.read()) - etag = metadata.get("etag", "") - except (json.JSONDecodeError, OSError): - # Ignore metadata read errors - pass - - contents.append( - { - "Key": key, - "LastModified": datetime.fromtimestamp(stat_result.st_mtime), - "ETag": f'"{etag}"', - "Size": stat_result.st_size, - } - ) - except (OSError, json.JSONDecodeError): - continue - - except OSError: - pass - - # Sort by key - contents.sort(key=lambda x: x["Key"]) - - return {"Contents": contents} - - async def list_object_versions( - self, - Bucket: str, - Prefix: str | None = None, - Delimiter: str | None = None, - KeyMarker: str | None = None, - **kwargs, - ) -> dict[str, Any]: - """List all versions of objects.""" - bucket_path = self._get_bucket_path(Bucket) - versions_path = bucket_path / ".versions" - - versions = [] - delete_markers = [] - - if not versions_path.exists(): - return {"Versions": versions, "DeleteMarkers": delete_markers} - - try: - # Find all version files - for version_file in versions_path.rglob("*"): - if not version_file.is_file(): - continue - - # Parse version file name (key.version_id) - relative_path = version_file.relative_to(versions_path) - key_parts = str(relative_path).rsplit(".", 1) - if len(key_parts) != 2: - continue - - key, version_id = key_parts - - # Apply prefix filter - if Prefix and not key.startswith(Prefix): - continue - - # Apply key marker filter - if KeyMarker and key <= KeyMarker: - continue - - try: - stat_result = await aiofiles.os.stat(version_file) - - # Check if this is the latest version - current_object = self._get_object_path(Bucket, key) - is_latest = False - if current_object.exists(): - # Get current version from metadata - metadata_path = self._get_metadata_path(Bucket, key) - if metadata_path.exists(): - async with aiofiles.open(metadata_path) as f: - metadata = json.loads(await f.read()) - is_latest = metadata.get("version_id") == version_id - - # Get etag from version metadata (or compute from file) - etag = "" - version_metadata_path = version_file.with_suffix(version_file.suffix + ".metadata.json") - if version_metadata_path.exists(): - async with aiofiles.open(version_metadata_path) as f: - metadata = json.loads(await f.read()) - etag = metadata.get("etag", "") - else: - # Fallback: compute from file content - async with aiofiles.open(version_file, "rb") as f: - content = await f.read() - etag = _compute_hash(content) - - versions.append( - { - "Key": key, - "VersionId": version_id, - "IsLatest": is_latest, - "LastModified": datetime.fromtimestamp(stat_result.st_mtime), - "ETag": f'"{etag}"', - "Size": stat_result.st_size, - } - ) - - except (OSError, json.JSONDecodeError): - continue - - except OSError: - pass - - # Sort by key and version - versions.sort(key=lambda x: (x["Key"], x["LastModified"]), reverse=True) - - return {"Versions": versions, "DeleteMarkers": delete_markers} - - async def delete_bucket(self, Bucket: str, **kwargs) -> dict[str, Any]: - """Delete bucket and all its contents.""" - bucket_path = self._get_bucket_path(Bucket) - - if bucket_path.exists(): - # Remove all files recursively - import shutil - - shutil.rmtree(bucket_path) - - return {} - - async def generate_presigned_url( - self, ClientMethod: str, Params: dict | None = None, ExpiresIn: int = 3600, **kwargs - ) -> str: - """Generate a presigned URL (return local file path for local storage).""" - if not Params: - return "" - - bucket = Params.get("Bucket", "") - key = Params.get("Key", "") - - # For local files, return a proxy URL that matches the expected format - return f"/api/v1/file/{bucket}/{key}" - - # Additional methods needed for compatibility - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - pass - - -class MockAsyncStreamingBody: - """Mock streaming body for local file content.""" - - def __init__(self, content: bytes, chunk_size: int = 8192): - self._content = content - self._stream = io.BytesIO(content) - self._chunk_size = chunk_size - self._position = 0 - - async def read(self, amt: int | None = None) -> bytes: - """Read content.""" - if amt is None: - # For full reads, always return the complete content - return self._content - else: - # For partial reads, use the stream position - return self._stream.read(amt) - - def __aiter__(self): - """Async iterator for streaming.""" - self._position = 0 - return self - - async def __anext__(self) -> bytes: - """Return next chunk of data.""" - if self._position >= len(self._content): - raise StopAsyncIteration - - chunk = self._content[self._position : self._position + self._chunk_size] - self._position += len(chunk) - return chunk - - async def __aenter__(self): - return self - - async def __aexit__(self, exc_type, exc_val, exc_tb): - pass - - -async def create_s3_client() -> "S3Client": - """Initialize S3 client with settings configuration or LocalFileClient as fallback.""" - # Check if S3 is configured - if not all([settings.s3_endpoint, settings.s3_access_key, settings.s3_secret_key]): - # Use local file storage as fallback - _LOGGER.info("S3 not configured, using local file storage at %s", settings.home_path) - local_client = LocalFileClient(settings.home_path or os.path.expanduser("~/.extralit")) - await local_client._ensure_base_dir() - shared_resources["s3_client"] = local_client - return local_client - - # Use real S3 client - session = aioboto3.Session( - aws_access_key_id=settings.s3_access_key, - aws_secret_access_key=settings.s3_secret_key, - region_name=settings.s3_region or "us-east-1", - ) - - s3_client = await session.client( # pyright: ignore[reportGeneralTypeIssues] - "s3", - endpoint_url=settings.s3_endpoint, - use_ssl=settings.s3_endpoint.startswith("https://") if settings.s3_endpoint else True, - ).__aenter__() - - shared_resources["s3_client"] = s3_client - return s3_client - - def replace_string_in_file(filename: str, string: str, replace_by: str, encoding: str = "utf-8"): # TODO Move where is used """Read a file and replace an old value in file by a new one""" diff --git a/extralit-server/src/extralit_server/jobs/document_jobs.py b/extralit-server/src/extralit_server/jobs/document_jobs.py index 7f92c6add..b0df00a70 100644 --- a/extralit-server/src/extralit_server/jobs/document_jobs.py +++ b/extralit-server/src/extralit_server/jobs/document_jobs.py @@ -59,8 +59,8 @@ async def analysis_and_preprocess_job( current_job.save_meta() try: - s3_client = await files.get_s3_client() - pdf_data = await files.download_file_content(s3_client, s3_url) + storage = await files.get_storage() + pdf_data = await files.download_file_content(storage, s3_url) filename = s3_url.split("/")[-1] triage = triage_pdf(pdf_data) @@ -95,7 +95,7 @@ async def analysis_and_preprocess_job( if thumbnail_data is not None: try: await files.put_object( - s3_client, + storage, workspace_name, files.get_thumbnail_s3_object_path(document_id), thumbnail_data, @@ -109,7 +109,7 @@ async def analysis_and_preprocess_job( _LOGGER.warning(f"No thumbnail data available for document {document_id}") await files.put_object( - s3_client, + storage, workspace_name, object_path, processing_response.processed_data, diff --git a/extralit-server/src/extralit_server/jobs/ocr_jobs.py b/extralit-server/src/extralit_server/jobs/ocr_jobs.py index 456470163..28db93dcb 100644 --- a/extralit-server/src/extralit_server/jobs/ocr_jobs.py +++ b/extralit-server/src/extralit_server/jobs/ocr_jobs.py @@ -11,7 +11,7 @@ from extralit_server.api.schemas.v1.document.metadata import LayoutMetadata from extralit_server.contexts import files from extralit_server.contexts.document.metadata import update_processing_metadata -from extralit_server.contexts.ocr import storage +from extralit_server.contexts.ocr import storage as layout_storage from extralit_server.contexts.ocr.layout_store import LayoutStore from extralit_server.contexts.ocr.parsers.pdf_inspector import classify from extralit_server.contexts.ocr.parsers.registry import default_parser_name, get_parser @@ -75,8 +75,8 @@ async def async_document_layout_job( try: # Shared client — do not enter it as a context manager, that would close it for everyone. - s3_client = await files.get_s3_client() - pdf_bytes = await files.download_file_content(s3_client, s3_url) + storage = await files.get_storage() + pdf_bytes = await files.download_file_content(storage, s3_url) routed, classification = route_parser(pdf_bytes) parser_name = parser or routed @@ -100,7 +100,7 @@ async def async_document_layout_job( _LOGGER.info(f"Layout for document {document_id} was not stored: {skip}") return {"document_id": str(document_id), "parser": parser_name, "skipped": skip} - paths = await storage.store_layout(s3_client, workspace_name, document_id, doc, store=store) + paths = await layout_storage.store_layout(storage, workspace_name, document_id, doc, store=store) layout = LayoutMetadata( **paths, diff --git a/extralit-server/src/extralit_server/jobs/preload.py b/extralit-server/src/extralit_server/jobs/preload.py index d45809665..fbccfd0c1 100644 --- a/extralit-server/src/extralit_server/jobs/preload.py +++ b/extralit-server/src/extralit_server/jobs/preload.py @@ -11,7 +11,6 @@ from extralit_server.contexts.document.preprocessing import PDFPreprocessingSettings, PDFPreprocessor # noqa: F401 from extralit_server.contexts.ocr.triage import triage_pdf # noqa: F401 from extralit_server.database import AsyncSessionLocal, async_engine # noqa: F401 -from extralit_server.helpers import create_s3_client # noqa: F401 from extralit_server.jobs import ( # noqa: F401 dataset_jobs, document_jobs, diff --git a/extralit-server/src/extralit_server/search_engine/commons.py b/extralit-server/src/extralit_server/search_engine/commons.py index 962d3a3af..c6dae6bce 100644 --- a/extralit-server/src/extralit_server/search_engine/commons.py +++ b/extralit-server/src/extralit_server/search_engine/commons.py @@ -240,15 +240,10 @@ def es_mapping_for_field(field: Field) -> dict: "properties": { "schemaName": {"type": "keyword"}, "etag": {"type": "keyword"}, - "version_id": {"type": "keyword"}, "last_modified": {"type": "date"}, - "version_tag": {"type": "keyword"}, - "is_latest": {"type": "boolean"}, }, }, "schemaName": {"type": "keyword"}, - "version_id": {"type": "keyword"}, - "is_latest": {"type": "boolean"}, }, }, "reference": {"type": "keyword"}, diff --git a/extralit-server/tests/factories.py b/extralit-server/tests/factories.py index cce4ebd72..aef8422d1 100644 --- a/extralit-server/tests/factories.py +++ b/extralit-server/tests/factories.py @@ -8,7 +8,7 @@ from factory.builder import BuildStep, StepBuilder, parse_declarations from sqlalchemy.ext.asyncio import async_object_session -from extralit_server.contexts.files import ObjectMetadata, get_s3_client +from extralit_server.contexts.files import ObjectMetadata, get_storage from extralit_server.enums import ( DatasetDistributionStrategy, FieldType, @@ -161,11 +161,10 @@ class Meta: @classmethod async def create_with_s3(cls, **kwargs): workspace = await cls.create(**kwargs) - s3_client = await get_s3_client() - try: - await s3_client.make_bucket(workspace.name) - except Exception as e: - print(f"Error creating bucket for workspace {workspace.name}: {e!s}") + from extralit_server.contexts import buckets + + storage = await get_storage() + await buckets.create(storage, workspace.name) return workspace @@ -634,11 +633,11 @@ def build(cls, **kwargs): @classmethod def create(cls, **kwargs): """Create a MinioFile and mock the put_object and get_object methods to return it.""" - from extralit_server.contexts.files import get_s3_client + from extralit_server.contexts.files import get_storage file = cls.build(**kwargs) - client = get_s3_client() + client = get_storage() # Store original methods getattr(client, "put_object", None) diff --git a/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py b/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py index 729b2bc9c..e44ced3e3 100644 --- a/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py +++ b/extralit-server/tests/unit/api/handlers/v1/datasets/test_schema_versions.py @@ -69,8 +69,8 @@ async def index_exists(self, dataset) -> bool: @pytest.fixture(autouse=True) def _mock_put_object(): # `publish_version` (Task 6) calls the real `files_ctx.put_object`, which would - # otherwise hit real object storage (or the LocalFileClient fallback under - # ~/.extralit) through the `files_ctx.get_s3_client` dependency. Stub at the + # otherwise hit real object storage (or the local object store under + # ~/.extralit) through the `files_ctx.get_storage` dependency. Stub at the # `put_object` call site rather than the dependency itself, matching the existing # convention in tests/unit/api/handlers/v1/test_files.py (`test_put_file` patches # `extralit_server.contexts.files.put_object`). This keeps the stub local to this diff --git a/extralit-server/tests/unit/api/handlers/v1/test_document_layout.py b/extralit-server/tests/unit/api/handlers/v1/test_document_layout.py index aaa3aad11..9a42613e5 100644 --- a/extralit-server/tests/unit/api/handlers/v1/test_document_layout.py +++ b/extralit-server/tests/unit/api/handlers/v1/test_document_layout.py @@ -70,7 +70,7 @@ async def make_document(db: AsyncSession, layout_metadata: dict | None) -> Docum @pytest.fixture def load_layout(): with patch( - "extralit_server.api.handlers.v1.documents.storage.load_layout", + "extralit_server.api.handlers.v1.documents.layout_storage.load_layout", new=AsyncMock(return_value=build_layout()), ) as mock: yield mock @@ -178,7 +178,7 @@ async def test_unreadable_stored_layout_is_404(self, async_client: AsyncClient, document = await make_document(db, LAYOUT_METADATA) with patch( - "extralit_server.api.handlers.v1.documents.storage.load_layout", + "extralit_server.api.handlers.v1.documents.layout_storage.load_layout", new=AsyncMock(side_effect=RuntimeError("object missing")), ): response = await async_client.get(f"/api/v1/documents/{document.id}/layout", headers=owner_auth_header) @@ -192,7 +192,7 @@ async def test_layout_from_a_newer_docling_is_409(self, async_client: AsyncClien error = ValidationError.from_exception_data("DoclingDocument", []) with patch( - "extralit_server.api.handlers.v1.documents.storage.load_layout", + "extralit_server.api.handlers.v1.documents.layout_storage.load_layout", new=AsyncMock(side_effect=error), ): response = await async_client.get(f"/api/v1/documents/{document.id}/layout", headers=owner_auth_header) diff --git a/extralit-server/tests/unit/api/handlers/v1/workspaces/test_create_workspace.py b/extralit-server/tests/unit/api/handlers/v1/workspaces/test_create_workspace.py index 4f099d1fe..5b69e8e63 100644 --- a/extralit-server/tests/unit/api/handlers/v1/workspaces/test_create_workspace.py +++ b/extralit-server/tests/unit/api/handlers/v1/workspaces/test_create_workspace.py @@ -41,7 +41,7 @@ async def test_create_workspace_with_predefined_id( self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict ): # Mock the create_bucket function directly - with patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket: + with patch("extralit_server.contexts.buckets.create") as mock_create_bucket: # Mock create_bucket as async coroutine mock_create_bucket.return_value = None @@ -68,7 +68,7 @@ async def test_create_workspace_with_none_id( self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict ): # Mock the create_bucket function directly - with patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket: + with patch("extralit_server.contexts.buckets.create") as mock_create_bucket: # Mock create_bucket as async coroutine mock_create_bucket.return_value = None @@ -107,7 +107,7 @@ async def test_create_workspace_with_existing_id( self, db: AsyncSession, async_client: AsyncClient, owner_auth_header: dict ): # Mock the create_bucket function directly - with patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket: + with patch("extralit_server.contexts.buckets.create") as mock_create_bucket: # Mock create_bucket as async coroutine mock_create_bucket.return_value = None diff --git a/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py b/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py index 6f3205191..15967ad2a 100644 --- a/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py +++ b/extralit-server/tests/unit/api/handlers/v1/workspaces/test_workspace_doctor.py @@ -17,7 +17,7 @@ async def test_workspace_doctor_healthy(self, db: AsyncSession, async_client: As # Mock S3 client and Redis connection with ( - patch("extralit_server.contexts.files.bucket_exists") as mock_bucket_exists, + patch("extralit_server.contexts.buckets.exists") as mock_bucket_exists, patch("extralit_server.jobs.queues.DEFAULT_QUEUE") as mock_queue, ): mock_bucket_exists.return_value = True @@ -49,8 +49,8 @@ async def test_workspace_doctor_missing_bucket_with_autofix( workspace = await WorkspaceFactory.create() with ( - patch("extralit_server.contexts.files.bucket_exists") as mock_bucket_exists, - patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket, + patch("extralit_server.contexts.buckets.exists") as mock_bucket_exists, + patch("extralit_server.contexts.buckets.create") as mock_create_bucket, patch("extralit_server.jobs.queues.DEFAULT_QUEUE") as mock_queue, ): mock_bucket_exists.return_value = False @@ -84,8 +84,8 @@ async def test_workspace_doctor_missing_bucket_without_autofix( workspace = await WorkspaceFactory.create() with ( - patch("extralit_server.contexts.files.bucket_exists") as mock_bucket_exists, - patch("extralit_server.contexts.files.create_bucket") as mock_create_bucket, + patch("extralit_server.contexts.buckets.exists") as mock_bucket_exists, + patch("extralit_server.contexts.buckets.create") as mock_create_bucket, patch("extralit_server.jobs.queues.DEFAULT_QUEUE") as mock_queue, ): mock_bucket_exists.return_value = False diff --git a/extralit-server/tests/unit/contexts/test_files_store.py b/extralit-server/tests/unit/contexts/test_files_store.py new file mode 100644 index 000000000..2009c72ef --- /dev/null +++ b/extralit-server/tests/unit/contexts/test_files_store.py @@ -0,0 +1,185 @@ +"""End-to-end coverage of the object-storage port against a real `LocalStore`. + +Every other test in the suite patches `contexts.files.*`, so until this module nothing +exercised a storage backend at all. `LocalStore` is real code rather than a mock, which is +the point of the obstore port: these assertions hold for `S3Store` too. +""" + +import pytest +from fastapi import HTTPException + +from extralit_server.contexts import files +from extralit_server.contexts.files import ObjectStorage + +BUCKET = "test-workspace" + + +@pytest.fixture +def storage(monkeypatch, tmp_path) -> ObjectStorage: + monkeypatch.setattr(files.settings, "s3_endpoint", None) + monkeypatch.setattr(files.settings, "home_path", str(tmp_path)) + return ObjectStorage() + + +@pytest.mark.asyncio +class TestRoundtrip: + async def test_put_then_get_returns_the_body(self, storage): + await files.put_object(storage, BUCKET, "a/b.txt", b"hello world", content_type="text/plain") + + file = await files.get_object(storage, BUCKET, "a/b.txt") + + assert bytes(await file.response.bytes_async()) == b"hello world" + assert file.metadata.size == 11 + assert file.metadata.object_name == "a/b.txt" + + async def test_put_accepts_str_and_encodes_utf8(self, storage): + await files.put_object(storage, BUCKET, "s.json", '{"k": "é"}', content_type="application/json") + + file = await files.get_object(storage, BUCKET, "s.json") + + assert bytes(await file.response.bytes_async()) == '{"k": "é"}'.encode() + + async def test_a_key_containing_dots_roundtrips(self, storage): + # The client this replaced derived version paths with `Path.with_suffix`, which + # silently truncated at the first dot: `pdf/a.b.pdf` came back as `pdf/a`. + await files.put_object(storage, BUCKET, "pdf/a.b.pdf", b"%PDF-1.4") + + file = await files.get_object(storage, BUCKET, "pdf/a.b.pdf") + + assert file.metadata.object_name == "pdf/a.b.pdf" + assert bytes(await file.response.bytes_async()) == b"%PDF-1.4" + + async def test_put_overwrites_in_place(self, storage): + await files.put_object(storage, BUCKET, "k.txt", b"first") + await files.put_object(storage, BUCKET, "k.txt", b"second") + + file = await files.get_object(storage, BUCKET, "k.txt") + + assert bytes(await file.response.bytes_async()) == b"second" + + +@pytest.mark.asyncio +class TestRanges: + async def test_a_single_byte_range_returns_one_byte(self, storage): + # An HTTP range is inclusive of its end, obstore's is exclusive. Off by one here + # and `Range: bytes=0-0` returns nothing. + await files.put_object(storage, BUCKET, "k.bin", b"0123456789") + store = storage.store_for(BUCKET) + + result = await store.get_async("k.bin", options={"range": (0, 0 + 1)}) + + assert bytes(result.bytes()) == b"0" + + async def test_a_range_reaching_the_last_byte(self, storage): + await files.put_object(storage, BUCKET, "k.bin", b"0123456789") + store = storage.store_for(BUCKET) + + result = await store.get_async("k.bin", options={"range": (5, 9 + 1)}) + + assert bytes(result.bytes()) == b"56789" + + +@pytest.mark.asyncio +class TestListing: + async def test_list_is_recursive_by_default(self, storage): + for key in ("p/one.txt", "p/nested/two.txt", "other.txt"): + await files.put_object(storage, BUCKET, key, b"x") + + listing = await files.list_objects(storage, BUCKET, prefix="p") + + assert {object.object_name for object in listing.objects} == {"p/one.txt", "p/nested/two.txt"} + + async def test_non_recursive_list_stops_at_the_delimiter(self, storage): + for key in ("p/one.txt", "p/nested/two.txt"): + await files.put_object(storage, BUCKET, key, b"x") + + listing = await files.list_objects(storage, BUCKET, prefix="p", recursive=False) + + assert {object.object_name for object in listing.objects} == {"p/one.txt"} + + async def test_listing_an_empty_prefix_is_not_an_error(self, storage): + await files.put_object(storage, BUCKET, "k.txt", b"x") + + listing = await files.list_objects(storage, BUCKET, prefix="absent") + + assert list(listing.objects) == [] + + +@pytest.mark.asyncio +class TestDeletion: + async def test_delete_then_get_raises_404(self, storage): + await files.put_object(storage, BUCKET, "k.txt", b"x") + await files.delete_object(storage, BUCKET, "k.txt") + + with pytest.raises(HTTPException) as error: + await files.get_object(storage, BUCKET, "k.txt") + + assert error.value.status_code == 404 + + async def test_deleting_an_absent_key_is_a_no_op(self, storage): + await files.delete_object(storage, BUCKET, "never/existed.txt") + + async def test_a_deleted_object_leaves_the_listing(self, storage): + # The delete-marker model this replaced kept returning deleted keys from listings. + await files.put_object(storage, BUCKET, "p/k.txt", b"x") + await files.delete_object(storage, BUCKET, "p/k.txt") + + listing = await files.list_objects(storage, BUCKET, prefix="p") + + assert list(listing.objects) == [] + + +@pytest.mark.asyncio +class TestContentType: + async def test_it_is_derived_from_the_key_extension(self, storage): + await files.put_object(storage, BUCKET, "a/b.txt", b"x", content_type="text/plain") + + file = await files.get_object(storage, BUCKET, "a/b.txt") + + assert file.metadata.content_type == "text/plain" + + async def test_extensionless_minted_keys_fall_back_to_their_prefix(self, storage): + # `pdf/{id}` and `thumbnails/{id}` carry no extension, and LocalStore cannot persist + # attributes, so the prefix map is the only thing standing between a local-dev PDF + # and being served as application/octet-stream. + await files.put_object(storage, BUCKET, "pdf/some-uuid", b"%PDF-1.4", content_type="application/pdf") + + file = await files.get_object(storage, BUCKET, "pdf/some-uuid") + + assert file.metadata.content_type == "application/pdf" + + async def test_an_unknown_key_defaults_to_octet_stream(self, storage): + await files.put_object(storage, BUCKET, "mystery", b"x") + + file = await files.get_object(storage, BUCKET, "mystery") + + assert file.metadata.content_type == "application/octet-stream" + + +@pytest.mark.asyncio +class TestDocumentUrls: + async def test_download_file_content_reads_the_whole_body(self, storage): + await files.put_object(storage, BUCKET, "pdf/doc-1", b"%PDF-1.4 body") + url = files.get_proxy_document_url(BUCKET, "pdf/doc-1") + + assert await files.download_file_content(storage, url) == b"%PDF-1.4 body" + + async def test_a_malformed_url_is_rejected(self, storage): + with pytest.raises(ValueError): + await files.download_file_content(storage, "https://example.com/not-a-proxy-url") + + async def test_local_storage_cannot_presign_so_it_keeps_the_proxy_url(self, storage): + url = files.get_proxy_document_url(BUCKET, "pdf/doc-1") + + assert await files.get_presigned_url_from_document_url(storage, url, expires=60) == url + + +@pytest.mark.asyncio +class TestPutDocumentFile: + async def test_the_first_upload_returns_a_proxy_url(self, storage): + from uuid import uuid4 + + url = await files.put_document_file(storage, BUCKET, uuid4(), b"%PDF-1.4", "a.pdf") + + assert url is not None + assert url.startswith(f"/api/v1/file/{BUCKET}/pdf/") diff --git a/extralit-server/tests/unit/contexts/test_schema_versions.py b/extralit-server/tests/unit/contexts/test_schema_versions.py index 1916bb9ee..fc4d8c4e4 100644 --- a/extralit-server/tests/unit/contexts/test_schema_versions.py +++ b/extralit-server/tests/unit/contexts/test_schema_versions.py @@ -1,4 +1,5 @@ -from unittest.mock import AsyncMock, patch +import tempfile +from unittest.mock import patch import pandas as pd import pandera.pandas as pa @@ -7,9 +8,11 @@ from sqlalchemy.dialects import postgresql, sqlite from extralit_server.contexts import schema_versions +from extralit_server.contexts.files import ObjectStorage from extralit_server.enums import DatasetStatus, FieldType from extralit_server.errors.future import UnprocessableEntityError from extralit_server.models.database import Dataset, Field +from extralit_server.settings import settings from extralit_server.webhooks.v1.enums import DatasetEvent from tests.factories import DatasetFactory, TextFieldFactory @@ -32,22 +35,15 @@ async def _fields_for(db, dataset_id) -> list[Field]: return list((await db.execute(stmt)).scalars().all()) -def _s3_client() -> AsyncMock: - """A stand-in S3 client good enough for `files_ctx.put_object`'s head_object round-trip. +def _storage() -> ObjectStorage: + """A real `LocalStore`-backed storage rooted in a temp dir. - A bare `AsyncMock()` doesn't work here: every attribute of an unspecced AsyncMock is - itself an AsyncMock, so `head_response.get(...)` inside `put_object` returns an - un-awaited coroutine instead of a value. Stub `head_object` to return a plain dict. + `publish_version` writes the body through `files_ctx.put_object`, and the assertions here + are about what lands in Postgres, not about mocking the write away. """ - client = AsyncMock() - client.head_object.return_value = { - "ETag": '"etag"', - "ContentLength": 0, - "LastModified": None, - "ContentType": "application/json", - "Metadata": {}, - } - return client + settings.s3_endpoint = None + settings.home_path = tempfile.mkdtemp() + return ObjectStorage() class TestDeriveColumnFields: @@ -117,7 +113,7 @@ async def test_publish_creates_version_one_and_leaves_the_dataset_a_draft(self, # stays the sole draft -> ready transition, so a schema-backed dataset gets the same # DatasetPublishValidator checks as an annotation one and stays configurable until then. dataset = await DatasetFactory.create(status=DatasetStatus.draft) - version = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + version = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert version.version == 1 assert version.dataset_id == dataset.id assert dataset.current_schema_version_id == version.id @@ -125,41 +121,41 @@ async def test_publish_creates_version_one_and_leaves_the_dataset_a_draft(self, async def test_publish_materializes_column_fields(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") fields = await _fields_for(db, dataset.id) assert {f.name for f in fields} == {"population", "n_arms"} assert all(f.settings["type"] == FieldType.column for f in fields) async def test_republishing_is_idempotent_for_unchanged_columns(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") - v2 = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") + v2 = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert v2.version == 2 fields = await _fields_for(db, dataset.id) assert len(fields) == 2 # upserted, not duplicated async def test_republishing_adds_newly_declared_columns_while_still_a_draft(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") - await schema_versions.publish_version(db, _s3_client(), dataset, body=_wider_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_wider_body(), bucket="ws") fields = await _fields_for(db, dataset.id) assert {f.name for f in fields} == {"population", "n_arms", "outcome"} async def test_second_version_links_the_first_as_parent(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - v1 = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") - v2 = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + v1 = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") + v2 = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert v2.parent_version_id == v1.id async def test_publish_uploads_the_body_under_a_versioned_key(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - version = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + version = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert version.object_key == f"schemas/{dataset.id}/v1.json" async def test_invalid_body_is_rejected_before_anything_is_written(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) with pytest.raises(UnprocessableEntityError, match="not a valid Pandera DataFrameSchema"): - await schema_versions.publish_version(db, _s3_client(), dataset, body="{not pandera}", bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body="{not pandera}", bucket="ws") assert dataset.current_schema_version_id is None assert await _fields_for(db, dataset.id) == [] @@ -171,13 +167,13 @@ async def test_a_first_schema_version_on_an_already_published_dataset_is_rejecte dataset = await DatasetFactory.create(status=DatasetStatus.ready) with pytest.raises(UnprocessableEntityError, match="cannot be added to a published dataset"): - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert dataset.current_schema_version_id is None async def test_republishing_with_a_changed_column_dtype_is_rejected(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") changed_body = pa.DataFrameSchema( { @@ -187,7 +183,7 @@ async def test_republishing_with_a_changed_column_dtype_is_rejected(self, db): ).to_json() with pytest.raises(UnprocessableEntityError, match="cannot change dtype"): - await schema_versions.publish_version(db, _s3_client(), dataset, body=changed_body, bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=changed_body, bucket="ws") # Rejected before any write: no second version, no dtype mutation on the existing field. assert [v.version for v in await schema_versions.list_versions(db, dataset)] == [1] @@ -197,8 +193,8 @@ async def test_republishing_with_a_changed_column_dtype_is_rejected(self, db): async def test_republishing_with_an_unchanged_column_dtype_is_allowed(self, db): # Unchanged existing columns are always legal. Only a *changed* dtype is rejected. dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") - v2 = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") + v2 = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert v2.version == 2 async def test_adding_a_column_to_a_published_dataset_is_rejected(self, db): @@ -207,11 +203,11 @@ async def test_adding_a_column_to_a_published_dataset_is_rejected(self, db): # leave the dataset unwritable at `PUT /datasets/{id}/records/bulk`. Reject at publish # so the failure stays at the call that caused it. dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") await dataset.update(db, status=DatasetStatus.ready) with pytest.raises(UnprocessableEntityError, match="cannot be added to a published dataset"): - await schema_versions.publish_version(db, _s3_client(), dataset, body=_wider_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_wider_body(), bucket="ws") assert [v.version for v in await schema_versions.list_versions(db, dataset)] == [1] assert {f.name for f in await _fields_for(db, dataset.id)} == {"population", "n_arms"} @@ -224,7 +220,7 @@ async def test_a_column_colliding_with_an_annotation_field_is_rejected(self, db) await TextFieldFactory.create(name="population", dataset=dataset) with pytest.raises(UnprocessableEntityError, match="collides with an existing text field"): - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert await schema_versions.list_versions(db, dataset) == [] fields_by_name = {f.name: f for f in await _fields_for(db, dataset.id)} @@ -238,7 +234,7 @@ async def test_publish_notifies_the_dataset_updated_webhook_event(self, db): with patch("extralit_server.contexts.schema_versions.notify_dataset_event_v1") as mock_notify: mock_notify.return_value = [] - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") mock_notify.assert_awaited_once() awaited_args = mock_notify.await_args.args @@ -252,8 +248,8 @@ async def test_every_republish_notifies_updated_too(self, db): with patch("extralit_server.contexts.schema_versions.notify_dataset_event_v1") as mock_notify: mock_notify.return_value = [] - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") - v2 = await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") + v2 = await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert v2.version == 2 assert [call.args[1] for call in mock_notify.await_args_list] == [ @@ -267,7 +263,7 @@ async def test_publish_of_a_column_less_body_still_creates_a_version(self, db): # `Field.upsert_many` raises on an empty `objects` list, so publish_version must # skip the upsert call rather than blow up on a degenerate-but-valid schema. dataset = await DatasetFactory.create(status=DatasetStatus.draft) - version = await schema_versions.publish_version(db, _s3_client(), dataset, body=_empty_body(), bucket="ws") + version = await schema_versions.publish_version(db, _storage(), dataset, body=_empty_body(), bucket="ws") assert version.version == 1 assert await _fields_for(db, dataset.id) == [] @@ -305,7 +301,7 @@ async def _spy(statement, *args, **kwargs): return await original_execute(statement, *args, **kwargs) with patch.object(db, "execute", _spy): - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert any(getattr(statement, "_for_update_arg", None) is not None for statement in executed), ( "publish_version must take a row lock before allocating a version number" @@ -317,11 +313,11 @@ class TestReadVersions: async def test_list_versions_is_ordered_by_version_number(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) for _ in range(3): - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert [v.version for v in await schema_versions.list_versions(db, dataset)] == [1, 2, 3] async def test_get_version_by_number(self, db): dataset = await DatasetFactory.create(status=DatasetStatus.draft) - await schema_versions.publish_version(db, _s3_client(), dataset, body=_body(), bucket="ws") + await schema_versions.publish_version(db, _storage(), dataset, body=_body(), bucket="ws") assert (await schema_versions.get_version_by_number(db, dataset.id, 1)).version == 1 assert await schema_versions.get_version_by_number(db, dataset.id, 99) is None diff --git a/extralit-server/tests/unit/jobs/test_document_jobs.py b/extralit-server/tests/unit/jobs/test_document_jobs.py index 5c06c46bd..dfef4ffb3 100644 --- a/extralit-server/tests/unit/jobs/test_document_jobs.py +++ b/extralit-server/tests/unit/jobs/test_document_jobs.py @@ -61,7 +61,7 @@ async def put_object(client, workspace, key, data, **kwargs): patch(f"{MODULE}.get_current_job", return_value=current_job), patch(f"{MODULE}.AsyncSessionLocal") as session, ): - files.get_s3_client = AsyncMock(return_value=MagicMock()) + files.get_storage = AsyncMock(return_value=MagicMock()) files.download_file_content = AsyncMock(return_value=b"%PDF original") files.get_thumbnail_s3_object_path.return_value = f"thumbnails/{DOCUMENT_ID}" files.put_object = AsyncMock(side_effect=put_object) diff --git a/extralit-server/tests/unit/jobs/test_ocr_jobs.py b/extralit-server/tests/unit/jobs/test_ocr_jobs.py index f4df568dd..aadf95854 100644 --- a/extralit-server/tests/unit/jobs/test_ocr_jobs.py +++ b/extralit-server/tests/unit/jobs/test_ocr_jobs.py @@ -31,12 +31,12 @@ async def update_metadata(db, document_id, mutate): db = MagicMock() with ( - patch(f"{MODULE}.files.get_s3_client", AsyncMock(return_value=AsyncMock())), + patch(f"{MODULE}.files.get_storage", AsyncMock(return_value=AsyncMock())), patch(f"{MODULE}.files.download_file_content", AsyncMock(return_value=b"%PDF-1.4")), patch(f"{MODULE}.route_parser", return_value=("pdf_inspector", {"pages_needing_ocr": [2]})), patch(f"{MODULE}.get_parser", return_value=lambda *args, **kwargs: doc), patch(f"{MODULE}.LayoutStore.for_workspace", return_value=store), - patch(f"{MODULE}.storage.store_layout", store_layout), + patch(f"{MODULE}.layout_storage.store_layout", store_layout), patch(f"{MODULE}.update_processing_metadata", update_metadata), patch(f"{MODULE}.get_current_job", return_value=MagicMock(meta={"workflow_id": "wf-1"})), patch(f"{MODULE}.writer_skip_reason", AsyncMock(return_value=None)) as skip, diff --git a/extralit-server/uv.lock b/extralit-server/uv.lock index 82fff85e5..8b3160720 100644 --- a/extralit-server/uv.lock +++ b/extralit-server/uv.lock @@ -1105,7 +1105,6 @@ name = "extralit-server" source = { editable = "." } dependencies = [ { name = "aioboto3" }, - { name = "aiofiles" }, { name = "aiosqlite" }, { name = "alembic" }, { name = "authlib" }, @@ -1152,7 +1151,6 @@ dependencies = [ { name = "standardwebhooks" }, { name = "tenacity" }, { name = "typer" }, - { name = "types-aiobotocore-s3" }, { name = "uvicorn", extra = ["standard"] }, { name = "xxhash" }, ] @@ -1184,7 +1182,6 @@ dev = [ [package.metadata] requires-dist = [ { name = "aioboto3", specifier = ">=13.1.1" }, - { name = "aiofiles", specifier = ">=24.1.0" }, { name = "aiosqlite", specifier = "==0.20.0" }, { name = "alembic", specifier = ">=1.13.0" }, { name = "asyncpg", marker = "extra == 'postgresql'", specifier = ">=0.30.0" }, @@ -1233,7 +1230,6 @@ requires-dist = [ { name = "standardwebhooks", specifier = ">=1.0.0" }, { name = "tenacity", specifier = ">=9.1.2" }, { name = "typer", specifier = ">=0.19.1" }, - { name = "types-aiobotocore-s3", specifier = "==2.24.2" }, { name = "uvicorn", extras = ["standard"], specifier = ">=0.32.0" }, { name = "xxhash", specifier = ">=3.6.0" }, ] @@ -4369,18 +4365,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" }, ] -[[package]] -name = "types-aiobotocore-s3" -version = "2.24.2" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.12'" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/16/f7/278bd63254f296a8fef102a9e3b6281ad02b1e8a4bb702e21c5e201d9f50/types_aiobotocore_s3-2.24.2.tar.gz", hash = "sha256:7274bcca558385ef5f98037f5474f4583663634141d46d9d978ab7d37ca96324", size = 76388, upload-time = "2025-09-06T01:36:47.406Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/6d/10/4e09f98ba4d7c422ffd1b7291c8556fa66af0cabba4694e91fb6be077c47/types_aiobotocore_s3-2.24.2-py3-none-any.whl", hash = "sha256:b5541e994f9c5f41d551570be8788d4189fb55759f8f7d15a2893b29eae3114f", size = 83875, upload-time = "2025-09-06T01:36:45.804Z" }, -] - [[package]] name = "types-deprecated" version = "1.3.1.20260130" From 370bab02a19748702d11fa07e303027685272480 Mon Sep 17 00:00:00 2001 From: JonnyTran Date: Wed, 19 Aug 2026 14:11:26 -0700 Subject: [PATCH 3/3] fix(server): resolve storage lazily and drain versioned buckets Two HIGH findings from review of the obstore port (roborev #354, #355). `contexts/imports.py` read `shared_resources["storage"]` directly. Removing `create_s3_client()` from the lifespan left nothing to populate it eagerly, so a bulk upload arriving before any other storage-dependent request raised "S3 client not available" -- and an RQ worker, which never runs the lifespan at all, would have failed every time. It now awaits `get_storage()`, which resolves and caches on first use. `contexts/buckets.delete` drained with obstore's `list()`, which returns current objects only. Removing `put_bucket_versioning` does not disable versioning on buckets that already have it, so a workspace created before this branch keeps its noncurrent versions and delete markers and `DeleteBucket` fails with `BucketNotEmpty`. Worse, a key hidden behind a delete marker does not appear in the current-object listing at all, so the drain removed nothing. Deletion now pages `list_object_versions` and batch-deletes every version and delete marker, falling back to the old drain on a backend without version listing. The previous commit message claimed this case was already fixed; it was only ever true of buckets created after the change. For the buckets that already exist, `suspend_versioning` stops them minting new versions. It runs when `create` lands on a pre-existing bucket and is exposed as a `bucket_versioning` workspace-doctor check, which is the operator-facing migration path. This also closes the leak opened by dropping the `expire-noncurrent-layout-versions` lifecycle rule: layout JSON is rewritten on every commit, so a still-versioned bucket grew without bound and nothing reaped it. Verified against MinIO on a bucket seeded with 3 versions and 1 delete marker: the old path fails with BucketNotEmpty, the new path drains and deletes it, and suspension leaves the bucket `Suspended`. --- .../api/handlers/v1/workspaces.py | 42 +++- .../src/extralit_server/contexts/buckets.py | 83 +++++++- .../src/extralit_server/contexts/imports.py | 7 +- .../tests/unit/contexts/test_buckets.py | 183 ++++++++++++++++++ 4 files changed, 303 insertions(+), 12 deletions(-) create mode 100644 extralit-server/tests/unit/contexts/test_buckets.py diff --git a/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py b/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py index 207737987..c22270ebe 100644 --- a/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py +++ b/extralit-server/src/extralit_server/api/handlers/v1/workspaces.py @@ -183,6 +183,7 @@ async def workspace_doctor( Checks: - S3 bucket exists (can auto-fix) + - Bucket is not object-versioned (can auto-fix) - RQ worker pool connectivity (informational) """ await authorize(current_user, WorkspacePolicy.get(workspace_id)) @@ -232,7 +233,42 @@ async def workspace_doctor( ) ) - # Check 2: RQ worker pool connectivity + # Check 2: legacy object versioning. Buckets created before versioning was removed still + # mint a new version on every write, and the lifecycle rule that expired them is gone. + if bucket_exists: + if autofix: + try: + suspended = await buckets.normalize_versioning(storage, workspace.name) + checks.append( + WorkspaceDoctorCheckResult( + check_name="bucket_versioning", + status="ok", + message=f"Object versioning on '{workspace.name}' has been suspended" + if suspended + else f"Object versioning is not enabled on '{workspace.name}'", + fixed=suspended, + ) + ) + except Exception as e: + checks.append( + WorkspaceDoctorCheckResult( + check_name="bucket_versioning", + status="warning", + message=f"Could not check object versioning on '{workspace.name}': {e!s}", + fixed=False, + ) + ) + else: + checks.append( + WorkspaceDoctorCheckResult( + check_name="bucket_versioning", + status="ok", + message="Object versioning not checked (autofix disabled)", + fixed=False, + ) + ) + + # Check 3: RQ worker pool connectivity try: from extralit_server.jobs.queues import DEFAULT_QUEUE @@ -258,7 +294,7 @@ async def workspace_doctor( ) ) - # Check 3: Elasticsearch indexes for datasets (informational only) + # Check 4: Elasticsearch indexes for datasets (informational only) try: # Get datasets for this workspace from sqlalchemy import select @@ -312,7 +348,7 @@ async def workspace_doctor( ) ) - # Check 4: Database connections health with autofix + # Check 5: Database connections health with autofix try: import asyncio diff --git a/extralit-server/src/extralit_server/contexts/buckets.py b/extralit-server/src/extralit_server/contexts/buckets.py index 4c47cf51f..ed295fe57 100644 --- a/extralit-server/src/extralit_server/contexts/buckets.py +++ b/extralit-server/src/extralit_server/contexts/buckets.py @@ -18,6 +18,9 @@ _LOGGER = logging.getLogger(__name__) +# DeleteObjects caps one request at 1000 keys. +_DELETE_BATCH = 1000 + def _remote() -> bool: return all([settings.s3_endpoint, settings.s3_access_key, settings.s3_secret_key]) @@ -74,6 +77,8 @@ async def create(storage: ObjectStorage, workspace_name: str) -> None: except ClientError as e: if e.response["Error"]["Code"] not in ["BucketAlreadyOwnedByYou", "BucketAlreadyExists"]: raise + # Pre-dates the removal of object versioning, so it may still be Enabled. + await suspend_versioning(client, bucket) except ClientError as e: _LOGGER.error(f"Error creating bucket {bucket}: {e}") raise HTTPException(status_code=500, detail=f"Error creating bucket: {e!s}") @@ -82,19 +87,68 @@ async def create(storage: ObjectStorage, workspace_name: str) -> None: raise HTTPException(status_code=500, detail=f"Internal server error: {e!s}") -async def delete(storage: ObjectStorage, workspace_name: str) -> None: - """Drain the workspace's bucket and delete it.""" +async def suspend_versioning(client, bucket: str) -> bool: + """Stop a legacy bucket from minting new object versions. True if it was Enabled. + + Buckets created before versioning was removed still carry `Status: Enabled`, and the + lifecycle rule that expired their noncurrent layout versions went away with it -- so every + layout rewrite would grow the bucket with nothing reaping it. Suspending is the migration; + existing noncurrent versions stay until the bucket is drained or deleted. + """ + try: + current = await client.get_bucket_versioning(Bucket=bucket) + if current.get("Status") != "Enabled": + return False + + await client.put_bucket_versioning(Bucket=bucket, VersioningConfiguration={"Status": "Suspended"}) + _LOGGER.info(f"Suspended object versioning on legacy bucket {bucket}") + return True + except ClientError as e: + # A backend without versioning support never had the problem in the first place. + _LOGGER.warning(f"Could not suspend versioning on bucket {bucket}: {e}") + return False + + +async def normalize_versioning(storage: ObjectStorage, workspace_name: str) -> bool: + """Suspend versioning on the workspace's bucket if it is still Enabled. True if it changed.""" bucket, _prefix = workspace_root(workspace_name) if not _remote(): - shutil.rmtree(_local_root(bucket), ignore_errors=True) - return + return False + async with _client() as client: + return await suspend_versioning(client, bucket) + + +async def _drain_every_version(client, bucket: str) -> None: + """Delete every object version and delete marker in the bucket. + + A current-objects-only drain leaves the noncurrent versions and delete markers of a legacy + versioned bucket behind, and `DeleteBucket` then fails with `BucketNotEmpty`. Objects in an + unversioned bucket come back with `VersionId: "null"`, which is a valid delete target, so + this one path covers both. + """ + paginator = client.get_paginator("list_object_versions") + async for page in paginator.paginate(Bucket=bucket): + targets = [ + {"Key": entry["Key"], "VersionId": entry["VersionId"]} + for section in ("Versions", "DeleteMarkers") + for entry in page.get(section, []) + ] + for start in range(0, len(targets), _DELETE_BATCH): + await client.delete_objects( + Bucket=bucket, + Delete={"Objects": targets[start : start + _DELETE_BATCH], "Quiet": True}, + ) + + +async def _drain_current_only(storage: ObjectStorage, bucket: str) -> None: + """Fallback drain for a backend that does not implement ListObjectVersions.""" store = storage.store_for(bucket) try: listing = await list_objects(storage, bucket) except HTTPException: - return # No such bucket; nothing to drain or delete. + return for object in listing.objects: try: @@ -102,8 +156,27 @@ async def delete(storage: ObjectStorage, workspace_name: str) -> None: except Exception as remove_err: _LOGGER.warning(f"Error removing object {object.object_name} during bucket delete: {remove_err}") + +async def delete(storage: ObjectStorage, workspace_name: str) -> None: + """Drain the workspace's bucket and delete it.""" + bucket, _prefix = workspace_root(workspace_name) + + if not _remote(): + shutil.rmtree(_local_root(bucket), ignore_errors=True) + return + try: async with _client() as client: + try: + await _drain_every_version(client, bucket) + except ClientError as e: + code = e.response["Error"]["Code"] + if code == "NoSuchBucket": + return # Nothing to drain or delete. + if code != "NotImplemented": + raise + await _drain_current_only(storage, bucket) + await client.delete_bucket(Bucket=bucket) _LOGGER.info(f"Successfully deleted bucket: {bucket}") except ClientError as e: diff --git a/extralit-server/src/extralit_server/contexts/imports.py b/extralit-server/src/extralit_server/contexts/imports.py index 695d0c9cd..c8c1e1752 100644 --- a/extralit-server/src/extralit_server/contexts/imports.py +++ b/extralit-server/src/extralit_server/contexts/imports.py @@ -23,7 +23,6 @@ ) from extralit_server.contexts import files as file_context from extralit_server.database import AsyncSessionLocal -from extralit_server.helpers import shared_resources from extralit_server.models.database import Document, ImportHistory, Workspace from extralit_server.workflows.documents import create_document_workflow @@ -374,9 +373,9 @@ async def process_bulk_upload( ) reference_to_doc[doc.reference] = doc - storage = shared_resources.get("storage") - if storage is None: - raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="S3 client not available") + # Not `shared_resources["storage"]`: nothing populates it eagerly, and an RQ worker never + # runs the app lifespan at all, so the cache is empty until some other caller fills it. + storage = await file_context.get_storage() # Process each reference: upload files to S3, create documents, start workflows job_ids: dict[str, str] = {} diff --git a/extralit-server/tests/unit/contexts/test_buckets.py b/extralit-server/tests/unit/contexts/test_buckets.py new file mode 100644 index 000000000..dd6e4cfaa --- /dev/null +++ b/extralit-server/tests/unit/contexts/test_buckets.py @@ -0,0 +1,183 @@ +"""Bucket lifecycle against a faked S3 admin client. + +The drain is the part worth pinning: a legacy workspace bucket still carries +`Status: Enabled`, so a current-objects-only drain leaves noncurrent versions and delete +markers behind and `DeleteBucket` fails with `BucketNotEmpty`. +""" + +from contextlib import asynccontextmanager +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from botocore.exceptions import ClientError + +from extralit_server.contexts import buckets + +BUCKET = "test-workspace" + + +def _client_error(code: str) -> ClientError: + return ClientError({"Error": {"Code": code, "Message": code}}, "op") + + +class FakeClient: + """Just enough S3 admin surface for buckets.py, recording what it was asked to delete.""" + + def __init__(self, pages=None, versioning="Enabled"): + self._pages = pages if pages is not None else [{}] + self._versioning = versioning + self.deleted: list[dict] = [] + self.delete_calls = 0 + self.deleted_bucket: str | None = None + self.versioning_set: str | None = None + self.create_bucket = AsyncMock() + self.head_bucket = AsyncMock() + + def get_paginator(self, name): + assert name == "list_object_versions" + pages = self._pages + + class _Paginator: + def paginate(self, **kwargs): + async def _gen(): + for page in pages: + yield page + + return _gen() + + return _Paginator() + + async def delete_objects(self, Bucket, Delete): + self.delete_calls += 1 + self.deleted.extend(Delete["Objects"]) + + async def delete_bucket(self, Bucket): + self.deleted_bucket = Bucket + + async def get_bucket_versioning(self, Bucket): + return {"Status": self._versioning} if self._versioning else {} + + async def put_bucket_versioning(self, Bucket, VersioningConfiguration): + self.versioning_set = VersioningConfiguration["Status"] + + +@asynccontextmanager +async def _yield(client): + yield client + + +@pytest.fixture +def remote(monkeypatch): + monkeypatch.setattr(buckets, "_remote", lambda: True) + + +@pytest.mark.asyncio +class TestDelete: + async def test_it_deletes_noncurrent_versions_and_delete_markers(self, remote): + client = FakeClient( + pages=[ + { + "Versions": [ + {"Key": "pdf/a", "VersionId": "v2"}, + {"Key": "pdf/a", "VersionId": "v1"}, + ], + "DeleteMarkers": [{"Key": "pdf/gone", "VersionId": "v3"}], + } + ] + ) + + with patch.object(buckets, "_client", lambda: _yield(client)): + await buckets.delete(MagicMock(), BUCKET) + + assert client.deleted == [ + {"Key": "pdf/a", "VersionId": "v2"}, + {"Key": "pdf/a", "VersionId": "v1"}, + {"Key": "pdf/gone", "VersionId": "v3"}, + ] + assert client.deleted_bucket == BUCKET + + async def test_an_unversioned_bucket_drains_through_the_same_path(self, remote): + # Objects in a bucket that was never versioned come back as VersionId "null". + client = FakeClient(pages=[{"Versions": [{"Key": "k.txt", "VersionId": "null"}]}]) + + with patch.object(buckets, "_client", lambda: _yield(client)): + await buckets.delete(MagicMock(), BUCKET) + + assert client.deleted == [{"Key": "k.txt", "VersionId": "null"}] + assert client.deleted_bucket == BUCKET + + async def test_deletes_are_batched_under_the_thousand_key_cap(self, remote): + page = {"Versions": [{"Key": f"k{i}", "VersionId": "null"} for i in range(2500)]} + client = FakeClient(pages=[page]) + + with patch.object(buckets, "_client", lambda: _yield(client)): + await buckets.delete(MagicMock(), BUCKET) + + assert client.delete_calls == 3 + assert len(client.deleted) == 2500 + + async def test_an_empty_bucket_is_still_deleted(self, remote): + client = FakeClient(pages=[{}]) + + with patch.object(buckets, "_client", lambda: _yield(client)): + await buckets.delete(MagicMock(), BUCKET) + + assert client.deleted == [] + assert client.deleted_bucket == BUCKET + + async def test_a_missing_bucket_is_not_an_error(self, remote): + client = FakeClient() + client.get_paginator = MagicMock(side_effect=_client_error("NoSuchBucket")) + + with patch.object(buckets, "_client", lambda: _yield(client)): + await buckets.delete(MagicMock(), BUCKET) + + assert client.deleted_bucket is None + + async def test_it_falls_back_when_the_backend_lacks_version_listing(self, remote): + client = FakeClient() + client.get_paginator = MagicMock(side_effect=_client_error("NotImplemented")) + + with ( + patch.object(buckets, "_client", lambda: _yield(client)), + patch.object(buckets, "_drain_current_only", AsyncMock()) as fallback, + ): + await buckets.delete(MagicMock(), BUCKET) + + assert fallback.await_count == 1 + assert client.deleted_bucket == BUCKET + + +@pytest.mark.asyncio +class TestSuspendVersioning: + async def test_an_enabled_bucket_is_suspended(self): + client = FakeClient(versioning="Enabled") + + assert await buckets.suspend_versioning(client, BUCKET) is True + assert client.versioning_set == "Suspended" + + async def test_an_unversioned_bucket_is_left_alone(self): + client = FakeClient(versioning=None) + + assert await buckets.suspend_versioning(client, BUCKET) is False + assert client.versioning_set is None + + async def test_a_backend_without_versioning_support_is_not_fatal(self): + client = FakeClient() + client.get_bucket_versioning = AsyncMock(side_effect=_client_error("NotImplemented")) + + assert await buckets.suspend_versioning(client, BUCKET) is False + + async def test_creating_over_a_legacy_bucket_suspends_it(self, remote): + client = FakeClient(versioning="Enabled") + client.create_bucket = AsyncMock(side_effect=_client_error("BucketAlreadyOwnedByYou")) + + with patch.object(buckets, "_client", lambda: _yield(client)): + await buckets.create(MagicMock(), BUCKET) + + assert client.versioning_set == "Suspended" + + async def test_local_storage_has_nothing_to_normalize(self, monkeypatch): + monkeypatch.setattr(buckets, "_remote", lambda: False) + + assert await buckets.normalize_versioning(MagicMock(), BUCKET) is False