"""This file and its contents are licensed under the Apache License 2.0. Please see the included NOTICE for copyright information and LICENSE for a copy of the license. """ import json import logging import re from typing import Union from urllib.parse import urlparse import boto3 from core.feature_flags import flag_set from core.redis import start_job_async_or_sync from django.conf import settings from django.db import models from django.db.models.signals import post_save, pre_delete from django.dispatch import receiver from django.utils.translation import gettext_lazy as _ from io_storages.base_models import ( ExportStorage, ExportStorageLink, ImportStorage, ImportStorageLink, ProjectStorageMixin, ) from io_storages.s3.utils import ( catch_and_reraise_from_none, get_client_and_resource, resolve_s3_url, ) from io_storages.utils import StorageObject, load_tasks_json, storage_can_resolve_bucket_url from tasks.models import Annotation from label_studio.io_storages.s3.utils import AWS logger = logging.getLogger(__name__) logging.getLogger('botocore').setLevel(logging.CRITICAL) boto3.set_stream_logger(level=logging.INFO) clients_cache = {} class S3StorageMixin(models.Model): bucket = models.TextField(_('bucket'), null=True, blank=True, help_text='S3 bucket name') prefix = models.TextField(_('prefix'), null=True, blank=True, help_text='S3 bucket prefix') regex_filter = models.TextField( _('regex_filter'), null=True, blank=True, help_text='Cloud storage regex for filtering objects', ) use_blob_urls = models.BooleanField( _('use_blob_urls'), default=False, help_text='Interpret objects as BLOBs and generate URLs', ) aws_access_key_id = models.TextField(_('aws_access_key_id'), null=True, blank=True, help_text='AWS_ACCESS_KEY_ID') aws_secret_access_key = models.TextField( _('aws_secret_access_key'), null=True, blank=True, help_text='AWS_SECRET_ACCESS_KEY', ) aws_session_token = models.TextField(_('aws_session_token'), null=True, blank=True, help_text='AWS_SESSION_TOKEN') aws_sse_kms_key_id = models.TextField( _('aws_sse_kms_key_id'), null=True, blank=True, help_text='AWS SSE KMS Key ID' ) region_name = models.TextField(_('region_name'), null=True, blank=True, help_text='AWS Region') s3_endpoint = models.TextField(_('s3_endpoint'), null=True, blank=True, help_text='S3 Endpoint') @catch_and_reraise_from_none def get_client_and_resource(self): # s3 client initialization ~ 100 ms, for 30 tasks it's a 3 seconds, so we need to cache it cache_key = f'{self.aws_access_key_id}:{self.aws_secret_access_key}:{self.aws_session_token}:{self.region_name}:{self.s3_endpoint}' if cache_key in clients_cache: return clients_cache[cache_key] result = get_client_and_resource( self.aws_access_key_id, self.aws_secret_access_key, self.aws_session_token, self.region_name, self.s3_endpoint, ) clients_cache[cache_key] = result return result def get_client(self): client, _ = self.get_client_and_resource() return client def get_client_and_bucket(self, validate_connection=True): client, s3 = self.get_client_and_resource() if validate_connection: self.validate_connection(client) return client, s3.Bucket(self.bucket) @catch_and_reraise_from_none def validate_connection(self, client=None): logger.debug('validate_connection') if client is None: client = self.get_client() # TODO(jo): add check for write access for .*Export.* classes is_export = 'Export' in self.__class__.__name__ if self.prefix: logger.debug( f'[Class {self.__class__.__name__}]: Test connection to bucket {self.bucket} with prefix {self.prefix} using ListObjectsV2 operation' ) result = client.list_objects_v2(Bucket=self.bucket, Prefix=self.prefix, MaxKeys=1) # We expect 1 key with the prefix for imports. For exports it's okay if there are 0 with the prefix. expected_keycount = 0 if is_export else 1 if (keycount := result.get('KeyCount')) is None or keycount < expected_keycount: raise KeyError(f'{self.url_scheme}://{self.bucket}/{self.prefix} not found.') else: logger.debug( f'[Class {self.__class__.__name__}]: Test connection to bucket {self.bucket} using HeadBucket operation' ) client.head_bucket(Bucket=self.bucket) @property def path_full(self): prefix = self.prefix or '' return f'{self.url_scheme}://{self.bucket}/{prefix}' @property def type_full(self): return 'Amazon AWS S3' @catch_and_reraise_from_none def get_bytes_stream(self, uri, range_header=None): """Get file directly from S3 using iter_chunks without wrapper. This method forwards Range headers directly to S3 and returns the raw stream. Note: The returned stream is NOT seekable and will break if seeking backwards. Args: uri: The S3 URI of the file to retrieve range_header: Optional HTTP Range header to forward to S3 Returns: Tuple of (stream, content_type, metadata) where metadata contains important S3 headers like ETag, ContentLength, etc. """ # Parse URI to get bucket and key parsed_uri = urlparse(uri, allow_fragments=False) bucket_name = parsed_uri.netloc key = parsed_uri.path.lstrip('/') # Get S3 client client = self.get_client() try: # Forward Range header to S3 if provided request_params = {'Bucket': bucket_name, 'Key': key} if range_header: request_params['Range'] = range_header # Get the object from S3 response = client.get_object(**request_params) # Extract metadata to return metadata = { 'ETag': response.get('ETag'), 'ContentLength': response.get('ContentLength'), 'ContentRange': response.get('ContentRange'), 'LastModified': response.get('LastModified'), 'StatusCode': response['ResponseMetadata']['HTTPStatusCode'], } # Return the streaming body directly return response['Body'], response.get('ContentType'), metadata except Exception as e: logger.error(f'Error getting direct stream from S3 for uri {uri}: {e}', exc_info=True) return None, None, {} class Meta: abstract = True class S3ImportStorageBase(S3StorageMixin, ImportStorage): url_scheme = 's3' presign = models.BooleanField(_('presign'), default=True, help_text='Generate presigned URLs') presign_ttl = models.PositiveSmallIntegerField( _('presign_ttl'), default=1, help_text='Presigned URLs TTL (in minutes)' ) recursive_scan = models.BooleanField( _('recursive scan'), default=False, help_text=_('Perform recursive scan over the bucket content'), ) @catch_and_reraise_from_none def iter_objects(self): _, bucket = self.get_client_and_bucket() list_kwargs = {} if self.prefix: list_kwargs['Prefix'] = self.prefix.rstrip('/') + '/' if not self.recursive_scan: list_kwargs['Delimiter'] = '/' bucket_iter = bucket.objects.filter(**list_kwargs).all() regex = re.compile(str(self.regex_filter)) if self.regex_filter else None for obj in bucket_iter: key = obj.key if key.endswith('/'): logger.debug(key + ' is skipped because it is a folder') continue if regex and not regex.match(key): logger.debug(key + ' is skipped by regex filter') continue logger.debug(f's3 {key} has passed the regex filter') yield obj @catch_and_reraise_from_none def iter_keys(self): for obj in self.iter_objects(): yield obj.key def get_unified_metadata(self, obj): return { 'key': obj.key, 'last_modified': obj.last_modified, 'size': obj.size, } @catch_and_reraise_from_none def scan_and_create_links(self): return self._scan_and_create_links(S3ImportStorageLink) @catch_and_reraise_from_none def get_data(self, key) -> list[StorageObject]: uri = f'{self.url_scheme}://{self.bucket}/{key}' if self.use_blob_urls: data_key = settings.DATA_UNDEFINED_NAME task = {data_key: uri} return [StorageObject(key=key, task_data=task)] # read task json from bucket and validate it _, s3 = self.get_client_and_resource() bucket = s3.Bucket(self.bucket) obj = s3.Object(bucket.name, key).get()['Body'].read() return load_tasks_json(obj, key) @catch_and_reraise_from_none def generate_http_url(self, url): return resolve_s3_url(url, self.get_client(), self.presign, expires_in=self.presign_ttl * 60) @catch_and_reraise_from_none def can_resolve_url(self, url: Union[str, None]) -> bool: return storage_can_resolve_bucket_url(self, url) @catch_and_reraise_from_none def get_blob_metadata(self, key): return AWS.get_blob_metadata( key, self.bucket, aws_access_key_id=self.aws_access_key_id, aws_secret_access_key=self.aws_secret_access_key, aws_session_token=self.aws_session_token, region_name=self.region_name, s3_endpoint=self.s3_endpoint, ) class Meta: abstract = True class S3ImportStorage(ProjectStorageMixin, S3ImportStorageBase): class Meta: abstract = False class S3ExportStorage(S3StorageMixin, ExportStorage): @catch_and_reraise_from_none def save_annotation(self, annotation): client, s3 = self.get_client_and_resource() logger.debug(f'Creating new object on {self.__class__.__name__} Storage {self} for annotation {annotation}') ser_annotation = self._get_serialized_data(annotation) # get key that identifies this object in storage key = S3ExportStorageLink.get_key(annotation) key = str(self.prefix) + '/' + key if self.prefix else key # put object into storage additional_params = {} self.cached_user = getattr(self, 'cached_user', self.project.organization.created_by) if flag_set( 'fflag_feat_back_lsdv_3958_server_side_encryption_for_target_storage_short', user=self.cached_user, ): if self.aws_sse_kms_key_id: additional_params['SSEKMSKeyId'] = self.aws_sse_kms_key_id additional_params['ServerSideEncryption'] = 'aws:kms' else: additional_params['ServerSideEncryption'] = 'AES256' s3.Object(self.bucket, key).put(Body=json.dumps(ser_annotation), **additional_params) # create link if everything ok S3ExportStorageLink.create(annotation, self) @catch_and_reraise_from_none def delete_annotation(self, annotation): client, s3 = self.get_client_and_resource() logger.debug(f'Deleting object on {self.__class__.__name__} Storage {self} for annotation {annotation}') # get key that identifies this object in storage key = S3ExportStorageLink.get_key(annotation) key = str(self.prefix) + '/' + key if self.prefix else key # delete object from storage s3.Object(self.bucket, key).delete() # delete link if everything ok S3ExportStorageLink.objects.filter(storage=self, annotation=annotation).delete() def async_export_annotation_to_s3_storages(annotation: 'Annotation | int'): if isinstance(annotation, int): try: annotation = Annotation.objects.get(pk=annotation) except Annotation.DoesNotExist: logger.info(f'Annotation {annotation} no longer exists, skipping S3 export') return project = annotation.project if hasattr(project, 'io_storages_s3exportstorages'): for storage in project.io_storages_s3exportstorages.all(): logger.debug(f'Export {annotation} to S3 storage {storage}') storage.save_annotation(annotation) @receiver(post_save, sender=Annotation) def export_annotation_to_s3_storages(sender, instance, **kwargs): storages = getattr(instance.project, 'io_storages_s3exportstorages', None) if storages and storages.exists(): # avoid excess jobs in rq start_job_async_or_sync(async_export_annotation_to_s3_storages, instance.pk) @receiver(pre_delete, sender=Annotation) def delete_annotation_from_s3_storages(sender, instance, **kwargs): links = S3ExportStorageLink.objects.filter(annotation=instance) for link in links: storage = link.storage if storage.can_delete_objects: logger.debug(f'Delete {instance} from S3 storage {storage}') # nosec storage.delete_annotation(instance) class S3ImportStorageLink(ImportStorageLink): storage = models.ForeignKey(S3ImportStorage, on_delete=models.CASCADE, related_name='links') @classmethod def exists(cls, keys, storage) -> set[str]: super_exists = super(S3ImportStorageLink, cls).exists # TODO: this is a workaround to be compatible with old keys version - remove it later prefix = str(storage.prefix) or '' if prefix: return ( super_exists(keys, storage) | super_exists([prefix + key for key in keys], storage) | super_exists([prefix + '/' + key for key in keys], storage) ) else: return super_exists(keys, storage) class S3ExportStorageLink(ExportStorageLink): storage = models.ForeignKey(S3ExportStorage, on_delete=models.CASCADE, related_name='links')