Skip to content
24 changes: 23 additions & 1 deletion ami/main/admin.py
Original file line number Diff line number Diff line change
Expand Up @@ -566,11 +566,31 @@ class SiteAdmin(admin.ModelAdmin[Site]):
class S3StorageSourceAdmin(admin.ModelAdmin[S3StorageSource]):
"""Admin panel example for ``S3StorageSource`` model."""

list_display = ("name", "bucket", "prefix", "size", "total_files", "last_checked")
list_display = (
"name",
"uri",
"size",
"total_files",
"is_private",
"last_checked",
"project",
"updated_at",
)

def size(self, obj) -> str:
return filesizeformat(obj.total_size)

@admin.display(description="S3 URI", ordering="bucket")
def uri(self, obj) -> str:
return obj.uri()

@admin.display(boolean=True)
def is_private(self, obj) -> bool:
"""
If a public base URL is set, the source is considered public.
"""
return not bool(obj.public_base_url)

@admin.action()
def calculate_size_async(self, request: HttpRequest, queryset: QuerySet[S3StorageSource]) -> None:
queued_tasks = [tasks.calculate_storage_size.apply_async([source.pk]) for source in queryset]
Expand All @@ -586,6 +606,8 @@ def count_files(self, request: HttpRequest, queryset: QuerySet[S3StorageSource])
source.count_files()
self.message_user(request, f"File count calculated for {queryset.count()} source(s).")

list_filter = ("project", "bucket")

actions = [calculate_size_async, count_files]


Expand Down
22 changes: 22 additions & 0 deletions ami/main/migrations/0081_s3storagesource_region.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
# Generated by Django 4.2.10 on 2026-01-27 18:00

from django.db import migrations, models


class Migration(migrations.Migration):
dependencies = [
("main", "0080_userprojectmembership"),
]

operations = [
migrations.AddField(
model_name="s3storagesource",
name="region",
field=models.CharField(
blank=True,
help_text="AWS region (e.g., 'us-east-1', 'eu-west-1'). Leave blank for Swift/MinIO storage.",
max_length=255,
null=True,
),
),
]
7 changes: 7 additions & 0 deletions ami/main/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -1456,6 +1456,12 @@ class S3StorageSource(BaseModel):

name = models.CharField(max_length=255)
bucket = models.CharField(max_length=255)
region = models.CharField(
max_length=255,
null=True,
blank=True,
help_text="AWS region (e.g., 'us-east-1', 'eu-west-1'). Leave blank for Swift/MinIO storage.",
)
prefix = models.CharField(max_length=255, blank=True)
access_key = models.TextField()
secret_key = models.TextField()
Expand All @@ -1475,6 +1481,7 @@ class S3StorageSource(BaseModel):
def config(self) -> ami.utils.s3.S3Config:
return ami.utils.s3.S3Config(
bucket_name=self.bucket,
region=self.region,
prefix=self.prefix,
access_key_id=self.access_key,
secret_access_key=self.secret_key,
Expand Down
2 changes: 2 additions & 0 deletions ami/tests/fixtures/storage.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
access_key_id=settings.S3_TEST_KEY,
secret_access_key=settings.S3_TEST_SECRET,
bucket_name=settings.S3_TEST_BUCKET,
region=settings.S3_TEST_REGION,
prefix="test_prefix",
public_base_url=f"http://minio:9000/{settings.S3_TEST_BUCKET}/test_prefix",
# public_base_url="http://minio:9001",
Expand All @@ -33,6 +34,7 @@ def create_storage_source(project: Project, name: str, prefix: str = S3_TEST_CON
access_key=S3_TEST_CONFIG.access_key_id,
secret_key=S3_TEST_CONFIG.secret_access_key,
public_base_url=S3_TEST_CONFIG.public_base_url,
region=S3_TEST_CONFIG.region,
),
)
return data_source
Expand Down
51 changes: 45 additions & 6 deletions ami/utils/s3.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from dataclasses import dataclass

import boto3
import boto3.resources.base
import boto3.session
Comment thread
mihow marked this conversation as resolved.
import botocore
import botocore.config
import botocore.exceptions
Expand All @@ -37,6 +37,7 @@ class S3Config:
secret_access_key: str
bucket_name: str
prefix: str
region: str | None = None
public_base_url: str | None = None

sensitive_fields = ["access_key_id", "secret_access_key"]
Expand Down Expand Up @@ -94,44 +95,79 @@ def get_session(config: S3Config) -> boto3.session.Session:
session = boto3.Session(
aws_access_key_id=config.access_key_id,
aws_secret_access_key=config.secret_access_key,
region_name=config.region,
)
return session


def get_s3_client(config: S3Config) -> S3Client:
session = get_session(config)

# Always use signature version 4
boto_config = botocore.config.Config(signature_version="s3v4")

if config.endpoint_url:
client = session.client(
service_name="s3",
endpoint_url=config.endpoint_url,
aws_access_key_id=config.access_key_id,
aws_secret_access_key=config.secret_access_key,
config=botocore.config.Config(signature_version="s3v4"),
region_name=config.region,
config=boto_config,
)
else:
client = session.client(
service_name="s3",
aws_access_key_id=config.access_key_id,
aws_secret_access_key=config.secret_access_key,
region_name=config.region,
config=boto_config,
)

client = typing.cast(S3Client, client)
Comment thread
mihow marked this conversation as resolved.
return client


def get_resource(config: S3Config) -> S3ServiceResource:
session = get_session(config)
boto_config = botocore.config.Config(signature_version="s3v4")
s3 = session.resource(
"s3",
endpoint_url=config.endpoint_url,
# api_version="s3v4",
region_name=config.region,
config=boto_config,
)
s3 = typing.cast(S3ServiceResource, s3)
Comment thread
mihow marked this conversation as resolved.
return s3


def create_bucket(config: S3Config, bucket_name: str, exists_ok: bool = True) -> CreateBucketOutputTypeDef | None:
"""
Create an S3 bucket.

Note: This is primarily used for testing. In production, users are expected to
create their own buckets and provide credentials to Antenna.

Args:
config: S3 configuration including region
bucket_name: Name of the bucket to create
exists_ok: If True, don't raise an error if bucket already exists

Returns:
CreateBucketOutputTypeDef or None if bucket already exists and exists_ok=True
"""
client = get_s3_client(config)
try:
# Create bucket if it doesn't exist
return client.create_bucket(Bucket=bucket_name)
# AWS requires CreateBucketConfiguration for non-us-east-1 regions
# See: https://docs.aws.amazon.com/AmazonS3/latest/API/API_CreateBucket.html
if config.region and config.region != "us-east-1":
return client.create_bucket(
Bucket=bucket_name,
CreateBucketConfiguration={"LocationConstraint": config.region},
)
else:
# us-east-1 or no region (Swift/MinIO) - don't specify CreateBucketConfiguration
return client.create_bucket(Bucket=bucket_name)
except botocore.exceptions.ClientError as e:
error_code = e.response.get("Error", {}).get("Code", "UnknownBotoError")
if error_code == "BucketAlreadyOwnedByYou" and exists_ok:
Expand Down Expand Up @@ -584,7 +620,9 @@ def read_image(config: S3Config, key: str) -> PIL.Image.Image:
obj = bucket.Object(key)
logger.info(f"Fetching image {key} from S3")
try:
img = PIL.Image.open(obj.get()["Body"])
# StreamingBody inherits from io.IOBase, but type checkers don't see that
fp = obj.get()["Body"]
img = PIL.Image.open(fp) # type: ignore[arg-type]
except PIL.UnidentifiedImageError:
logger.error(f"Could not read image {key}")
raise
Expand Down Expand Up @@ -677,6 +715,7 @@ def test():
bucket_name="test",
prefix="",
public_base_url="http://minio:9000/test",
region=None,
)

projects = list_projects(config)
Expand Down
1 change: 1 addition & 0 deletions config/settings/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -449,6 +449,7 @@
S3_TEST_KEY = env("MINIO_ROOT_USER", default=None) # type: ignore[no-untyped-call]
S3_TEST_SECRET = env("MINIO_ROOT_PASSWORD", default=None) # type: ignore[no-untyped-call]
S3_TEST_BUCKET = env("MINIO_TEST_BUCKET", default="ami-test") # type: ignore[no-untyped-call]
S3_TEST_REGION = env("MINIO_REGION", default=None) # type: ignore[no-untyped-call]


# Default processing service settings
Expand Down