Skip to content

Commit

Permalink
Fix broken network connection
Browse files Browse the repository at this point in the history
  • Loading branch information
rubickcz committed Aug 16, 2023
1 parent fcad8de commit 935772f
Show file tree
Hide file tree
Showing 3 changed files with 31 additions and 7 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/django.yml
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ jobs:

services:
dynamodb:
image: amazon/dynamodb-local
image: amazon/dynamodb-local:1.22.0
ports:
- 8000:8000

Expand Down
26 changes: 22 additions & 4 deletions pydjamodb/connection.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import logging
import multiprocessing
import time

from django.conf import settings
Expand All @@ -11,6 +12,7 @@
from botocore.client import ClientError

logger = logging.getLogger('pydjamodb.units')
connection_lock = multiprocessing.RLock()


class TableConnection(BaseTableConnection):
Expand Down Expand Up @@ -183,15 +185,31 @@ class TestTableConnection:

def __init__(self, wrapped_connection, prefix=None):
self._wrapped_connection = wrapped_connection
self._wrapped_table_name = wrapped_connection.table_name
self._is_test_clean_required = False
if prefix:
self._wrapped_connection.table_name = 'test_{}_{}'.format(prefix, self._wrapped_connection.table_name)
else:
self._wrapped_connection.table_name = 'test_{}'.format(self._wrapped_connection.table_name)
self._patch_connection()
self.set_table_name(prefix)

def __getattr__(self, attr):
return getattr(self._wrapped_connection, attr)

def _patch_connection(self):
original_method = self._wrapped_connection.connection.client._endpoint.http_session.send

def patched_method(*args, **kwargs):
with connection_lock:
# Botocore's underlying network implementation does not work well when accessed from multiple
# subprocessses. We must ensure that only one subprocess at a time uses the network connection.
return original_method(*args, **kwargs)

self._wrapped_connection.connection.client._endpoint.http_session.send = patched_method

def set_table_name(self, prefix=None):
if prefix:
self._wrapped_connection.table_name = 'test_{}_{}'.format(prefix, self._wrapped_table_name)
else:
self._wrapped_connection.table_name = 'test_{}'.format(self._wrapped_table_name)

def update_item(self, *args, **kwargs):
self._is_test_clean_required = True
return self._wrapped_connection.update_item(*args, **kwargs)
Expand Down
10 changes: 8 additions & 2 deletions pydjamodb/test_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,14 @@ def set_dynamodb_test_autoclean():

def init_pynamodb_test_prefix(prefix=None):
for model_class in dynamodb_model_classes:
model_class._connection = None
model_class._connection = TestTableConnection(model_class._get_connection(), prefix)
if isinstance(model_class._connection, TestTableConnection):
# Botocore's underlying network implementation does not work well when accessed from multiple subprocessses.
# We must ensure, that no new connection is created in a subprocess. At this point, in a subprocess, we
# always have existing connection from the parent process, which we can use.
model_class._connection.set_table_name(prefix)
else:
model_class._connection = None
model_class._connection = TestTableConnection(model_class._get_connection(), prefix)


def remove_pynamodb_table(model_class):
Expand Down

0 comments on commit 935772f

Please sign in to comment.