Skip to content

Commit

Permalink
Switch back to botocore _make_api_call under the hood (#1079)
Browse files Browse the repository at this point in the history
  • Loading branch information
garrettheel authored Nov 29, 2023
1 parent c5e91ca commit 369f461
Show file tree
Hide file tree
Showing 12 changed files with 1,775 additions and 1,883 deletions.
222 changes: 222 additions & 0 deletions bench/benchmark.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,222 @@
import timeit
import io
import logging
import zlib
from datetime import datetime

import urllib3

COUNT = 1000

benchmark_results = []
benchmark_registry = {}


def register_benchmark(testname):
def _wrap(func):
benchmark_registry[testname] = func
return func
return _wrap


def results_new_benchmark(name: str) -> None:
benchmark_results.append((name, {}))
print(name)


def results_record_result(callback, count):
callback_name = callback.__name__
bench_name = callback_name.split('_', 1)[-1]
try:
results = timeit.repeat(
f"{callback_name}()",
setup=f"from __main__ import patch_urllib3, {callback_name}; patch_urllib3()",
repeat=10,
number=count,
)
except Exception:
logging.exception(f"error running {bench_name}")
return
result = count / min(results)
benchmark_results.append((bench_name, str(result)))

print(f"{bench_name}: {result:,.02f} calls/sec")


# =============================================================================
# Monkeypatching
# =============================================================================

def mock_urlopen(self, method, url, body, headers, **kwargs):
target = headers.get('X-Amz-Target')
if target.endswith(b'DescribeTable'):
body = """{
"Table": {
"TableName": "users",
"TableArn": "arn",
"CreationDateTime": "1421866952.062",
"ItemCount": 0,
"TableSizeBytes": 0,
"TableStatus": "ACTIVE",
"ProvisionedThroughput": {
"NumberOfDecreasesToday": 0,
"ReadCapacityUnits": 1,
"WriteCapacityUnits": 25
},
"AttributeDefinitions": [{"AttributeName": "user_name", "AttributeType": "S"}],
"KeySchema": [{"AttributeName": "user_name", "KeyType": "HASH"}],
"LocalSecondaryIndexes": [],
"GlobalSecondaryIndexes": []
}
}
"""
elif target.endswith(b'GetItem'):
# TODO: sometimes raise exc
body = """{
"Item": {
"user_name": {"S": "some_user"},
"email": {"S": "some_user@gmail.com"},
"first_name": {"S": "John"},
"last_name": {"S": "Doe"},
"phone_number": {"S": "4155551111"},
"country": {"S": "USA"},
"preferences": {
"M": {
"timezone": {"S": "America/New_York"},
"allows_notifications": {"BOOL": 1},
"date_of_birth": {"S": "2022-10-26T20:00:00.000000+0000"}
}
},
"last_login": {"S": "2022-10-27T20:00:00.000000+0000"}
}
}
"""
elif target.endswith(b'PutItem'):
body = """{
"Attributes": {
"user_name": {"S": "some_user"},
"email": {"S": "some_user@gmail.com"},
"first_name": {"S": "John"},
"last_name": {"S": "Doe"},
"phone_number": {"S": "4155551111"},
"country": {"S": "USA"},
"preferences": {
"M": {
"timezone": {"S": "America/New_York"},
"allows_notifications": {"BOOL": 1},
"date_of_birth": {"S": "2022-10-26T20:44:49.207740+0000"}
}
},
"last_login": {"S": "2022-10-27T20:00:00.000000+0000"}
}
}
"""
else:
body = ""

body_bytes = body.encode('utf-8')
headers = {
"content-type": "application/x-amz-json-1.0",
"content-length": str(len(body_bytes)),
"x-amz-crc32": str(zlib.crc32(body_bytes)),
"x-amz-requestid": "YB5DURFL1EQ6ULM39GSEEHFTYTPBBUXDJSYPFZPR4EL7M3AYV0RS",
}

# TODO: consumed capacity?

body = io.BytesIO(body_bytes)
resp = urllib3.HTTPResponse(
body,
preload_content=False,
headers=headers,
status=200,
)
resp.chunked = False
return resp


def patch_urllib3():
urllib3.connectionpool.HTTPConnectionPool.urlopen = mock_urlopen


# =============================================================================
# Setup
# =============================================================================

import os
from pynamodb.models import Model
from pynamodb.attributes import UnicodeAttribute, BooleanAttribute, MapAttribute, UTCDateTimeAttribute


os.environ["AWS_ACCESS_KEY_ID"] = "1"
os.environ["AWS_SECRET_ACCESS_KEY"] = "1"
os.environ["AWS_DEFAULT_REGION"] = "us-east-1"


class UserPreferences(MapAttribute):
timezone = UnicodeAttribute()
allows_notifications = BooleanAttribute()
date_of_birth = UTCDateTimeAttribute()


class UserModel(Model):
class Meta:
table_name = 'User'
max_retry_attempts = 0 # TODO: do this conditionally. need to replace the connection object
user_name = UnicodeAttribute(hash_key=True)
first_name = UnicodeAttribute()
last_name = UnicodeAttribute()
phone_number = UnicodeAttribute()
country = UnicodeAttribute()
email = UnicodeAttribute()
preferences = UserPreferences(null=True)
last_login = UTCDateTimeAttribute()


# =============================================================================
# GetItem
# =============================================================================

@register_benchmark("get_item")
def bench_get_item():
UserModel.get("username")


# =============================================================================
# PutItem
# =============================================================================

@register_benchmark("put_item")
def bench_put_item():
UserModel(
"username",
email="some_user@gmail.com",
first_name="John",
last_name="Doe",
phone_number="4155551111",
country="USA",
preferences=UserPreferences(
timezone="America/New_York",
allows_notifications=True,
date_of_birth=datetime.utcnow(),
),
last_login=datetime.utcnow(),
).save()


# =============================================================================
# Benchmarks.
# =============================================================================

def main():
results_new_benchmark("Basic operations")

results_record_result(benchmark_registry["get_item"], COUNT)
results_record_result(benchmark_registry["put_item"], COUNT)

print()
print("Above metrics are in call/sec, larger is better.")


if __name__ == "__main__":
main()
9 changes: 0 additions & 9 deletions docs/settings.rst
Original file line number Diff line number Diff line change
Expand Up @@ -34,15 +34,6 @@ The number of times to retry certain failed DynamoDB API calls. The most common
retries include ``ProvisionedThroughputExceededException`` and ``5xx`` errors.


base_backoff_ms
---------------

Default: ``25``

The base number of milliseconds used for `exponential backoff and jitter
<https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/>`_ on retries.


region
------

Expand Down
3 changes: 3 additions & 0 deletions mypy.ini
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,6 @@ ignore_errors = True
# TODO: burn these down
[mypy-tests.*]
ignore_errors = True

[mypy-benchmark]
ignore_errors = True
25 changes: 5 additions & 20 deletions pynamodb/connection/_botocore_private.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""
Type-annotates the private botocore APIs that we're currently relying on.
"""
from typing import Any, Dict, Optional
from typing import Dict

import botocore.client
import botocore.credentials
Expand All @@ -22,25 +22,10 @@ class BotocoreRequestSignerPrivate(botocore.signers.RequestSigner):
class BotocoreBaseClientPrivate(botocore.client.BaseClient):
_endpoint: BotocoreEndpointPrivate
_request_signer: BotocoreRequestSignerPrivate
_service_model: botocore.model.ServiceModel

def _resolve_endpoint_ruleset(
def _make_api_call(
self,
operation_model: botocore.model.OperationModel,
params: Dict[str, Any],
request_context: Dict[str, Any],
ignore_signing_region: bool = ...,
):
raise NotImplementedError

def _convert_to_request_dict(
self,
api_params: Dict[str, Any],
operation_model: botocore.model.OperationModel,
*,
endpoint_url: str = ..., # added in botocore 1.28
context: Optional[Dict[str, Any]] = ...,
headers: Optional[Dict[str, Any]] = ...,
set_user_agent_header: bool = ...,
) -> Dict[str, Any]:
operation_name: str,
operation_kwargs: Dict,
) -> Dict:
raise NotImplementedError
Loading

0 comments on commit 369f461

Please sign in to comment.