Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 34 additions & 0 deletions label_studio/io_storages/s3/tests/test_user_agent.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,34 @@
"""Tests for the S3 client user agent.

Verifies get_client_and_resource attaches a label-studio user agent to both the
client and resource while preserving signature_version and a custom endpoint_url
(the path used for Amazon S3 and any S3-compatible object store such as
Backblaze B2, Cloudflare R2, or MinIO).
"""

import unittest

from io_storages.s3.utils import _get_user_agent_extra, get_client_and_resource


class TestS3ClientUserAgent(unittest.TestCase):
def test_user_agent_extra_format(self):
ua = _get_user_agent_extra()
assert ua.startswith('label-studio/')

def test_client_and_resource_carry_user_agent(self):
client, resource = get_client_and_resource()
expected = _get_user_agent_extra()
assert expected in (client.meta.config.user_agent_extra or '')
assert expected in (resource.meta.client.meta.config.user_agent_extra or '')

def test_signature_version_preserved(self):
client, _ = get_client_and_resource()
assert client.meta.config.signature_version == 's3v4'

def test_custom_endpoint_preserved(self):
endpoint = 'https://your-s3-endpoint.example.com'
client, resource = get_client_and_resource(s3_endpoint=endpoint)
assert client.meta.endpoint_url == endpoint
assert resource.meta.client.meta.endpoint_url == endpoint
assert _get_user_agent_extra() in (client.meta.config.user_agent_extra or '')
14 changes: 12 additions & 2 deletions label_studio/io_storages/s3/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import base64
import fnmatch
import importlib.metadata
import logging
import mimetypes
import re
Expand All @@ -16,6 +17,14 @@
logger = logging.getLogger(__name__)


def _get_user_agent_extra():
try:
version = importlib.metadata.version('label-studio')
except importlib.metadata.PackageNotFoundError:
version = 'dev'
return f'label-studio/{version}'


def get_client_and_resource(
aws_access_key_id=None, aws_secret_access_key=None, aws_session_token=None, region_name=None, s3_endpoint=None
):
Expand All @@ -37,8 +46,9 @@ def get_client_and_resource(
s3_endpoint = s3_endpoint or get_env('S3_ENDPOINT')
if s3_endpoint:
settings['endpoint_url'] = s3_endpoint
client = session.client('s3', config=boto3.session.Config(signature_version='s3v4'), **settings)
resource = session.resource('s3', config=boto3.session.Config(signature_version='s3v4'), **settings)
config = boto3.session.Config(signature_version='s3v4', user_agent_extra=_get_user_agent_extra())
client = session.client('s3', config=config, **settings)
resource = session.resource('s3', config=config, **settings)
return client, resource


Expand Down
Loading