From 935772f0b31184eacf5d682ff8d7523f0bcbe191 Mon Sep 17 00:00:00 2001
From: Ondrej Kulaty <kulaty.o@gmail.com>
Date: Sun, 6 Aug 2023 22:59:51 +0200
Subject: [PATCH] Fix broken network connection

---
 .github/workflows/django.yml |  2 +-
 pydjamodb/connection.py      | 26 ++++++++++++++++++++++----
 pydjamodb/test_runner.py     | 10 ++++++++--
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/.github/workflows/django.yml b/.github/workflows/django.yml
index 319b11d..96ee7c3 100644
--- a/.github/workflows/django.yml
+++ b/.github/workflows/django.yml
@@ -16,7 +16,7 @@ jobs:
 
     services:
       dynamodb:
-        image: amazon/dynamodb-local
+        image: amazon/dynamodb-local:1.22.0
         ports:
           - 8000:8000
 
diff --git a/pydjamodb/connection.py b/pydjamodb/connection.py
index 8887331..0d921d1 100644
--- a/pydjamodb/connection.py
+++ b/pydjamodb/connection.py
@@ -1,4 +1,5 @@
 import logging
+import multiprocessing
 import time
 
 from django.conf import settings
@@ -11,6 +12,7 @@
 from botocore.client import ClientError
 
 logger = logging.getLogger('pydjamodb.units')
+connection_lock = multiprocessing.RLock()
 
 
 class TableConnection(BaseTableConnection):
@@ -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)
diff --git a/pydjamodb/test_runner.py b/pydjamodb/test_runner.py
index 90dd436..783c42b 100644
--- a/pydjamodb/test_runner.py
+++ b/pydjamodb/test_runner.py
@@ -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):