⚡️ Speed up method S3DataSource.put_bucket_versioning by 48%

[codeflash-ai]

📄 48% (0.48x) speedup for S3DataSource.put_bucket_versioning in backend/python/app/sources/external/s3/s3.py

⏱️ Runtime : 1.34 milliseconds → 909 microseconds (best of 230 runs)

📝 Explanation and details

The optimized code achieves a 47% runtime speedup by eliminating redundant S3 client creation overhead through client connection pooling.

Key Optimization: Cached S3 Client Context

• Original: Creates a new async with session.client('s3') context for every put_bucket_versioning call, requiring connection setup/teardown each time
• Optimized: Introduces _get_s3_client() method that caches the S3 client instance (_s3_client_instance) and reuses the same connection across multiple calls

Performance Impact Analysis:
The line profiler shows the critical improvement in put_bucket_versioning:

• Original: 15.6% of time spent on async with session.client('s3') (1.48762e+06 ns across 1274 hits)
• Optimized: This overhead is eliminated after the first call since the client is cached

Throughput Benefits:

• 1.3% throughput improvement (145,053 → 146,970 ops/sec) demonstrates that the optimization maintains concurrent execution efficiency while reducing per-operation overhead
• The modest throughput gain reflects that async operations are primarily I/O bound, but the reduced connection setup still provides measurable benefit

Test Case Performance:
The optimization particularly benefits workloads with:

• Concurrent execution tests: Multiple simultaneous calls reuse the same cached client
• High-volume tests: 100+ concurrent calls see reduced connection overhead
• Repeated operations: Subsequent calls after the first avoid client initialization entirely

Resource Management:
The optimization includes proper cleanup via _finalize() and __aexit__() methods to ensure the cached client context is properly closed, preventing resource leaks while maintaining the performance benefits of connection reuse.

✅ Correctness verification report:

Test	Status
⚙️ Existing Unit Tests	🔘 None Found
🌀 Generated Regression Tests	✅ 913 Passed
⏪ Replay Tests	🔘 None Found
🔎 Concolic Coverage Tests	🔘 None Found
📊 Tests Coverage	83.3%

🌀 Generated Regression Tests and Runtime

import asyncio # used to run async functions
from typing import Any, Dict, Optional

import pytest # used for our unit tests
from app.sources.external.s3.s3 import S3DataSource

Minimal stub for aioboto3.Session and s3_client

class DummyS3Client:
def init(self, response=None, raise_client_error=False, raise_exception=False):
self.response = response
self.raise_client_error = raise_client_error
self.raise_exception = raise_exception

async def put_bucket_versioning(self, **kwargs):
    # Simulate error scenarios
    if self.raise_client_error:
        # Simulate botocore.exceptions.ClientError
        class DummyClientError(Exception):
            def __init__(self):
                self.response = {'Error': {'Code': 'AccessDenied', 'Message': 'Access Denied'}}
        raise DummyClientError()
    if self.raise_exception:
        raise Exception("Unexpected error!")
    # Simulate normal response
    return self.response

async def __aenter__(self):
    return self

async def __aexit__(self, exc_type, exc_val, exc_tb):
    return False

class DummySession:
def init(self, s3_client):
self._s3_client = s3_client

def client(self, service_name):
    return self._s3_client

class DummyS3ClientWrapper:
def init(self, session):
self._session = session

def get_session(self):
    return self._session

--- Function under test (EXACT COPY) ---

try:
import aioboto3 # type: ignore
from botocore.exceptions import ClientError # type: ignore
except ImportError:
# For test environment, we don't actually need aioboto3
aioboto3 = None
ClientError = Exception

--- TESTS ---

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_success():
# Basic test: normal successful response
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
result = await s3ds.put_bucket_versioning(
Bucket="my-bucket",
VersioningConfiguration={"Status": "Enabled"}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_error_dict():
# Basic test: S3 returns error in dict
dummy_response = {'Error': {'Code': 'InvalidRequest', 'Message': 'Invalid bucket'}}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
result = await s3ds.put_bucket_versioning(
Bucket="bad-bucket",
VersioningConfiguration={"Status": "Suspended"}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_none_response():
# Edge case: None response from S3
s3_client = DummyS3Client(response=None)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
result = await s3ds.put_bucket_versioning(
Bucket="my-bucket",
VersioningConfiguration={"Status": "Enabled"}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_with_all_optional_parameters():
# Basic test: All optional parameters provided
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
result = await s3ds.put_bucket_versioning(
Bucket="my-bucket",
VersioningConfiguration={"Status": "Enabled"},
ChecksumAlgorithm="SHA256",
MFA="123456",
ExpectedBucketOwner="ownerid"
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_clienterror_handling():
# Edge case: Simulate ClientError exception
s3_client = DummyS3Client(raise_client_error=True)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
result = await s3ds.put_bucket_versioning(
Bucket="forbidden-bucket",
VersioningConfiguration={"Status": "Enabled"}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_unexpected_exception():
# Edge case: Simulate generic exception
s3_client = DummyS3Client(raise_exception=True)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
result = await s3ds.put_bucket_versioning(
Bucket="my-bucket",
VersioningConfiguration={"Status": "Enabled"}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_concurrent_execution():
# Edge case: Run multiple concurrent requests
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)

async def run_one(i):
    return await s3ds.put_bucket_versioning(
        Bucket=f"bucket-{i}",
        VersioningConfiguration={"Status": "Enabled"}
    )

# Run 10 concurrent requests
results = await asyncio.gather(*(run_one(i) for i in range(10)))
for res in results:
    pass

@pytest.mark.asyncio
async def test_put_bucket_versioning_large_scale_concurrent():
# Large scale: Run 100 concurrent requests
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)

async def run_one(i):
    return await s3ds.put_bucket_versioning(
        Bucket=f"bucket-{i}",
        VersioningConfiguration={"Status": "Enabled"}
    )

results = await asyncio.gather(*(run_one(i) for i in range(100)))
for res in results:
    pass

@pytest.mark.asyncio
async def test_put_bucket_versioning_edge_missing_bucket():
# Edge case: Missing required Bucket parameter (should raise TypeError)
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
with pytest.raises(TypeError):
await s3ds.put_bucket_versioning(
VersioningConfiguration={"Status": "Enabled"}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_edge_missing_versioning_configuration():
# Edge case: Missing required VersioningConfiguration parameter (should raise TypeError)
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
with pytest.raises(TypeError):
await s3ds.put_bucket_versioning(
Bucket="my-bucket"
)

--- Throughput tests ---

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_small_load():
# Throughput test: 10 requests
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
tasks = [
s3ds.put_bucket_versioning(
Bucket=f"bucket-{i}",
VersioningConfiguration={"Status": "Enabled"}
) for i in range(10)
]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_medium_load():
# Throughput test: 50 requests
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
tasks = [
s3ds.put_bucket_versioning(
Bucket=f"bucket-{i}",
VersioningConfiguration={"Status": "Enabled"}
) for i in range(50)
]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_high_load():
# Throughput test: 200 requests (stress test, but <1000)
dummy_response = {'Versioning': 'Enabled'}
s3_client = DummyS3Client(response=dummy_response)
session = DummySession(s3_client)
client_wrapper = DummyS3ClientWrapper(session)
s3ds = S3DataSource(client_wrapper)
tasks = [
s3ds.put_bucket_versioning(
Bucket=f"bucket-{i}",
VersioningConfiguration={"Status": "Enabled"}
) for i in range(200)
]
results = await asyncio.gather(*tasks)

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_mixed_load():
# Throughput test: Mix of successful and error responses
dummy_response = {'Versioning': 'Enabled'}
error_response = {'Error': {'Code': 'InvalidRequest', 'Message': 'Invalid bucket'}}
clients = [
DummyS3Client(response=dummy_response) if i % 2 == 0 else DummyS3Client(response=error_response)
for i in range(20)
]
sessions = [DummySession(c) for c in clients]
wrappers = [DummyS3ClientWrapper(s) for s in sessions]
sources = [S3DataSource(w) for w in wrappers]

async def run_one(i):
    return await sources[i].put_bucket_versioning(
        Bucket=f"bucket-{i}",
        VersioningConfiguration={"Status": "Enabled"}
    )

results = await asyncio.gather(*(run_one(i) for i in range(20)))
for i, res in enumerate(results):
    if i % 2 == 0:
        pass
    else:
        pass

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

#------------------------------------------------
import asyncio # used to run async functions
from typing import Any, Dict, Optional

import pytest # used for our unit tests
from app.sources.external.s3.s3 import S3DataSource

class ClientError(Exception):
def init(self, error_response, operation_name):
self.response = error_response
self.operation_name = operation_name
super().init(str(error_response))

class DummyS3Client:
"""Mock S3 client with put_bucket_versioning method."""
def init(self, response=None, raise_client_error=False, raise_exception=False):
self.response = response
self.raise_client_error = raise_client_error
self.raise_exception = raise_exception

async def put_bucket_versioning(self, **kwargs):
    if self.raise_client_error:
        raise ClientError({'Error': {'Code': 'AccessDenied', 'Message': 'You do not have permission'}}, 'PutBucketVersioning')
    if self.raise_exception:
        raise RuntimeError("Unexpected failure")
    return self.response

class DummySessionClient:
"""Mock aioboto3 session.client context manager."""
def init(self, s3_client):
self.s3_client = s3_client

async def __aenter__(self):
    return self.s3_client

async def __aexit__(self, exc_type, exc, tb):
    pass

class DummySession:
"""Mock aioboto3.Session with client() returning async context manager."""
def init(self, s3_client):
self.s3_client = s3_client

def client(self, service_name):
    return DummySessionClient(self.s3_client)

class DummyS3ClientBuilder:
"""Mock S3Client.get_session() returning DummySession."""
def init(self, s3_client):
self._session = DummySession(s3_client)

def get_session(self):
    return self._session

------------------- UNIT TESTS -------------------

1. Basic Test Cases

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_success():
"""Test basic successful call with required parameters."""
# Simulate S3 API returns a dict with no error
response = {'Versioning': 'Enabled'}
s3_client = DummyS3ClientBuilder(DummyS3Client(response=response))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='mybucket',
VersioningConfiguration={'Status': 'Enabled'}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_error_response_dict():
"""Test S3 returns error in response dict."""
error_response = {'Error': {'Code': 'AccessDenied', 'Message': 'You do not have permission'}}
s3_client = DummyS3ClientBuilder(DummyS3Client(response=error_response))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='mybucket',
VersioningConfiguration={'Status': 'Suspended'}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_none_response():
"""Test S3 returns None as response."""
s3_client = DummyS3ClientBuilder(DummyS3Client(response=None))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='mybucket',
VersioningConfiguration={'Status': 'Enabled'}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_non_dict_response():
"""Test S3 returns non-dict response (e.g., string)."""
s3_client = DummyS3ClientBuilder(DummyS3Client(response="OK"))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='mybucket',
VersioningConfiguration={'Status': 'Suspended'}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_basic_optional_parameters():
"""Test all optional parameters are passed and handled."""
response = {'Versioning': 'Enabled'}
s3_client = DummyS3ClientBuilder(DummyS3Client(response=response))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='bucket123',
VersioningConfiguration={'Status': 'Enabled'},
ChecksumAlgorithm='SHA256',
MFA='123456 789012',
ExpectedBucketOwner='ownerid'
)

2. Edge Test Cases

@pytest.mark.asyncio
async def test_put_bucket_versioning_clienterror():
"""Test S3 raises ClientError exception."""
s3_client = DummyS3ClientBuilder(DummyS3Client(raise_client_error=True))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='bucket',
VersioningConfiguration={'Status': 'Enabled'}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_unexpected_exception():
"""Test S3 raises unexpected exception."""
s3_client = DummyS3ClientBuilder(DummyS3Client(raise_exception=True))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='bucket',
VersioningConfiguration={'Status': 'Enabled'}
)

@pytest.mark.asyncio
async def test_put_bucket_versioning_concurrent_execution():
"""Test concurrent execution of multiple put_bucket_versioning calls."""
# Each call returns a different response
responses = [
{'Versioning': 'Enabled'},
{'Versioning': 'Suspended'},
{'Error': {'Code': 'AccessDenied', 'Message': 'No permission'}}
]
s3_clients = [
DummyS3ClientBuilder(DummyS3Client(response=responses[0])),
DummyS3ClientBuilder(DummyS3Client(response=responses[1])),
DummyS3ClientBuilder(DummyS3Client(response=responses[2]))
]
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Enabled'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)

@pytest.mark.asyncio
async def test_put_bucket_versioning_invalid_versioning_configuration():
"""Test invalid VersioningConfiguration value."""
# S3 returns error in response dict
error_response = {'Error': {'Code': 'MalformedXML', 'Message': 'The XML you provided was not well-formed'}}
s3_client = DummyS3ClientBuilder(DummyS3Client(response=error_response))
ds = S3DataSource(s3_client)
result = await ds.put_bucket_versioning(
Bucket='bucket',
VersioningConfiguration={'InvalidKey': 'InvalidValue'}
)

3. Large Scale Test Cases

@pytest.mark.asyncio
async def test_put_bucket_versioning_many_concurrent_calls():
"""Test large scale concurrent execution (up to 50 calls)."""
N = 50
response = {'Versioning': 'Enabled'}
s3_clients = [DummyS3ClientBuilder(DummyS3Client(response=response)) for _ in range(N)]
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Enabled'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)
for r in results:
pass

@pytest.mark.asyncio
async def test_put_bucket_versioning_many_concurrent_errors():
"""Test large scale concurrent errors (up to 20 calls)."""
N = 20
error_response = {'Error': {'Code': 'AccessDenied', 'Message': 'Denied'}}
s3_clients = [DummyS3ClientBuilder(DummyS3Client(response=error_response)) for _ in range(N)]
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Suspended'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)
for r in results:
pass

4. Throughput Test Cases

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_small_load():
"""Throughput test: small load (5 concurrent calls)."""
N = 5
response = {'Versioning': 'Enabled'}
s3_clients = [DummyS3ClientBuilder(DummyS3Client(response=response)) for _ in range(N)]
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Enabled'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)
for r in results:
pass

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_medium_load():
"""Throughput test: medium load (25 concurrent calls)."""
N = 25
response = {'Versioning': 'Suspended'}
s3_clients = [DummyS3ClientBuilder(DummyS3Client(response=response)) for _ in range(N)]
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Suspended'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)
for r in results:
pass

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_high_volume():
"""Throughput test: high volume (100 concurrent calls)."""
N = 100
response = {'Versioning': 'Enabled'}
s3_clients = [DummyS3ClientBuilder(DummyS3Client(response=response)) for _ in range(N)]
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Enabled'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)
for r in results:
pass

@pytest.mark.asyncio
async def test_put_bucket_versioning_throughput_mixed_success_and_error():
"""Throughput test: mix of success and error responses."""
N = 30
success_response = {'Versioning': 'Enabled'}
error_response = {'Error': {'Code': 'AccessDenied', 'Message': 'Denied'}}
s3_clients = []
for i in range(N):
if i % 2 == 0:
s3_clients.append(DummyS3ClientBuilder(DummyS3Client(response=success_response)))
else:
s3_clients.append(DummyS3ClientBuilder(DummyS3Client(response=error_response)))
data_sources = [S3DataSource(c) for c in s3_clients]
coros = [
ds.put_bucket_versioning(Bucket=f'bucket{i}', VersioningConfiguration={'Status': 'Enabled' if i % 2 == 0 else 'Suspended'})
for i, ds in enumerate(data_sources)
]
results = await asyncio.gather(*coros)
for i, r in enumerate(results):
if i % 2 == 0:
pass
else:
pass

codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

To edit these changes git checkout codeflash/optimize-S3DataSource.put_bucket_versioning-mhx9to8v and push.