Fix RecursionError with psycopg when DEBUG=True

dts_test_project/dts_test_app/migrations/0003_test_add_db_index.py

@@ -2,7 +2,6 @@
 from __future__ import unicode_literals
 
 from django.db import models, migrations
-from django.conf import settings
 
 
 class Migration(migrations.Migration):

dts_test_project/dts_test_app/migrations/0004_test_alter_unique.py

@@ -2,7 +2,6 @@
 from __future__ import unicode_literals
 
 from django.db import models, migrations
-from django.conf import settings
 
 
 class Migration(migrations.Migration):

src/tenant_schemas/apps.py

@@ -97,7 +97,7 @@ def best_practice(app_configs, **kwargs):
     if not isinstance(default_storage, TenantStorageMixin):
         errors.append(
             Warning(
-                f"Your default storage engine is not tenant aware.",
+                "Your default storage engine is not tenant aware.",
                 hint="Set settings.STORAGES default backend to "
                 "'tenant_schemas.storage.TenantFileSystemStorage'",
                 id="tenant_schemas.W003",

src/tenant_schemas/postgresql_backend/base.py

@@ -1,5 +1,6 @@
 import re
 import warnings
+from contextvars import ContextVar
 
 from django.conf import settings
 from django.contrib.contenttypes.models import ContentType
@@ -18,14 +19,19 @@
     raise ImproperlyConfigured("Error loading psycopg2 or psycopg module")
 
 
-ORIGINAL_BACKEND = getattr(settings, 'ORIGINAL_BACKEND', 'django.db.backends.postgresql')
+ORIGINAL_BACKEND = getattr(
+    settings, "ORIGINAL_BACKEND", "django.db.backends.postgresql"
+)
 original_backend = django.db.utils.load_backend(ORIGINAL_BACKEND)
 
-EXTRA_SEARCH_PATHS = getattr(settings, 'PG_EXTRA_SEARCH_PATHS', [])
+EXTRA_SEARCH_PATHS = getattr(settings, "PG_EXTRA_SEARCH_PATHS", [])
+
+# ContextVar to prevent recursion when setting search_path under DEBUG=True with psycopg3
+_SETTING_SEARCH_PATH = ContextVar("ts_setting_search_path", default=False)
 
 # from the postgresql doc
-SQL_IDENTIFIER_RE = re.compile(r'^[_a-zA-Z][_a-zA-Z0-9]{,62}$')
-SQL_SCHEMA_NAME_RESERVED_RE = re.compile(r'^pg_', re.IGNORECASE)
+SQL_IDENTIFIER_RE = re.compile(r"^[_a-zA-Z][_a-zA-Z0-9]{,62}$")
+SQL_SCHEMA_NAME_RESERVED_RE = re.compile(r"^pg_", re.IGNORECASE)
 
 
 def _is_valid_identifier(identifier):
@@ -50,6 +56,7 @@ class DatabaseWrapper(original_backend.DatabaseWrapper):
     """
     Adds the capability to manipulate the search_path using set_tenant and set_schema_name
     """
+
     include_public_schema = True
 
     def __init__(self, *args, **kwargs):
@@ -58,16 +65,19 @@ def __init__(self, *args, **kwargs):
         # Use a patched version of the DatabaseIntrospection that only returns the table list for the
         # currently selected schema.
         self.introspection = DatabaseSchemaIntrospection(self)
+        self._ts_last_path_sig = None  # Cache for last applied search path signature
         self.set_schema_to_public()
 
     def close(self):
         self.search_path_set = False
+        self._ts_last_path_sig = None  # Clear cache on close
         super().close()
 
     def rollback(self):
         super().rollback()
         # Django's rollback clears the search path so we have to set it again the next time.
         self.search_path_set = False
+        self._ts_last_path_sig = None  # Clear cache on rollback
 
     def set_tenant(self, tenant, include_public=True):
         """
@@ -87,6 +97,7 @@ def set_schema(self, schema_name, include_public=True):
         self.include_public_schema = include_public
         self.set_settings_schema(schema_name)
         self.search_path_set = False
+        self._ts_last_path_sig = None  # Clear cache when schema changes
         # Content type can no longer be cached as public and tenant schemas
         # have different models. If someone wants to change this, the cache
         # needs to be separated between public and shared schemas. If this
@@ -103,18 +114,47 @@ def set_schema_to_public(self):
         self.set_schema(get_public_schema_name())
 
     def set_settings_schema(self, schema_name):
-        self.settings_dict['SCHEMA'] = schema_name
+        self.settings_dict["SCHEMA"] = schema_name
 
     def get_schema(self):
-        warnings.warn("connection.get_schema() is deprecated, use connection.schema_name instead.",
-                      category=DeprecationWarning)
+        warnings.warn(
+            "connection.get_schema() is deprecated, use connection.schema_name instead.",
+            category=DeprecationWarning,
+        )
         return self.schema_name
 
     def get_tenant(self):
-        warnings.warn("connection.get_tenant() is deprecated, use connection.tenant instead.",
-                      category=DeprecationWarning)
+        warnings.warn(
+            "connection.get_tenant() is deprecated, use connection.tenant instead.",
+            category=DeprecationWarning,
+        )
         return self.tenant
 
+    def _should_set_search_path(self, path_sig):
+        """
+        Determine if search_path needs to be set based on current configuration.
+
+        Returns True if:
+        - Limit set calls is disabled OR search_path is not set
+        - AND the path signature has changed
+        """
+        return (
+            not get_limit_set_calls() or not self.search_path_set
+        ) and self._ts_last_path_sig != path_sig
+
+    def _get_raw_cursor(self, cursor_for_search_path):
+        """
+        Get the raw DB-API cursor for psycopg2/psycopg3 compatibility.
+
+        In psycopg2, cursor_for_search_path may have a 'cursor' attribute 
+        pointing to the raw DB-API cursor.
+        In psycopg3, the cursor object itself is the raw DB-API cursor.
+        """
+        if hasattr(cursor_for_search_path, "cursor"):
+            return cursor_for_search_path.cursor
+        else:
+            return cursor_for_search_path
+
     def _cursor(self, name=None):
         """
         Here it happens. We hope every Django db operation using PostgreSQL
@@ -126,56 +166,92 @@ def _cursor(self, name=None):
         else:
             cursor = super()._cursor()
 
-        # optionally limit the number of executions - under load, the execution
-        # of `set search_path` can be quite time consuming
-        if (not get_limit_set_calls()) or not self.search_path_set:
-            # Actual search_path modification for the cursor. Database will
-            # search schemata from left to right when looking for the object
-            # (table, index, sequence, etc.).
-            if not self.schema_name:
-                raise ImproperlyConfigured("Database schema not set. Did you forget "
-                                           "to call set_schema() or set_tenant()?")
-            _check_schema_name(self.schema_name)
-            public_schema_name = get_public_schema_name()
-            search_paths = []
-
-            if self.schema_name == public_schema_name:
-                search_paths = [public_schema_name]
-            elif self.include_public_schema:
-                search_paths = [self.schema_name, public_schema_name]
-            else:
-                search_paths = [self.schema_name]
-
-            search_paths.extend(EXTRA_SEARCH_PATHS)
-
-            if name:
-                # Named cursor can only be used once
-                cursor_for_search_path = self.connection.cursor()
-            else:
-                # Reuse
-                cursor_for_search_path = cursor
-
-            # In the event that an error already happened in this transaction and we are going
-            # to rollback we should just ignore database error when setting the search_path
-            # if the next instruction is not a rollback it will just fail also, so
-            # we do not have to worry that it's not the good one
-            try:
-                cursor_for_search_path.execute('SET search_path = {0}'.format(','.join(search_paths)))
-            except (django.db.utils.DatabaseError, InternalError):
-                self.search_path_set = False
-            else:
-                self.search_path_set = True
+        # Calculate search paths for current tenant configuration
+        if not self.schema_name:
+            raise ImproperlyConfigured(
+                "Database schema not set. Did you forget "
+                "to call set_schema() or set_tenant()?"
+            )
+
+        _check_schema_name(self.schema_name)
+        public_schema_name = get_public_schema_name()
+        search_paths = []
+
+        if self.schema_name == public_schema_name:
+            search_paths = [public_schema_name]
+        elif self.include_public_schema:
+            search_paths = [self.schema_name, public_schema_name]
+        else:
+            search_paths = [self.schema_name]
+
+        search_paths.extend(EXTRA_SEARCH_PATHS)
+        path_sig = tuple(search_paths)
 
-            if name:
-                cursor_for_search_path.close()
+        # Check if we need to set the search path
+        if self._should_set_search_path(path_sig):
+            # Prevent recursion during debug/mogrify operations with psycopg3
+            if _SETTING_SEARCH_PATH.get():
+                return cursor
+
+            token = _SETTING_SEARCH_PATH.set(True)
+            try:
+                if name:
+                    # Named cursor can only be used once
+                    cursor_for_search_path = self.connection.cursor()
+                else:
+                    # Reuse - get raw cursor to avoid Django's debug wrapper
+                    cursor_for_search_path = cursor
+                    raw_cursor = self._get_raw_cursor(cursor_for_search_path)
+
+                # In the event that an error already happened in this transaction and we are going
+                # to rollback we should just ignore database error when setting the search_path
+                # if the next instruction is not a rollback it will just fail also, so
+                # we do not have to worry that it's not the good one
+                try:
+                    # Use set_config with parameters instead of raw SQL formatting to avoid
+                    # triggering Django's debug SQL logging that causes psycopg3 recursion
+                    if name:
+                        cursor_for_search_path.execute(
+                            "SELECT set_config('search_path', %s, false)",
+                            [",".join(search_paths)],
+                        )
+                    else:
+                        raw_cursor.execute(
+                            "SELECT set_config('search_path', %s, false)",
+                            [",".join(search_paths)],
+                        )
+                except (django.db.utils.DatabaseError, InternalError):
+                    self.search_path_set = False
+                    self._ts_last_path_sig = None
+                else:
+                    self.search_path_set = True
+                    self._ts_last_path_sig = path_sig
+
+                if name:
+                    cursor_for_search_path.close()
+            finally:
+                _SETTING_SEARCH_PATH.reset(token)
 
         return cursor
 
+    def last_executed_query(self, cursor, sql, params):
+        """
+        Override to avoid opening a fresh cursor during mogrify when there are no params.
+        This helps prevent recursion issues with psycopg3 when DEBUG=True.
+        """
+        if (
+            params is None or len(params) == 0
+        ):  # no need to mogrify, avoids opening a fresh cursor
+            return sql
+        # Delegate to the operations class
+        return self.ops.last_executed_query(cursor, sql, params)
+
 
 class FakeTenant:
     """
     We can't import any db model in a backend (apparently?), so this class is used
     for wrapping schema names in a tenant-like structure.
     """
+
     def __init__(self, schema_name):
         self.schema_name = schema_name

src/tenant_schemas/tests/test_apps.py

@@ -72,7 +72,7 @@ def test_storage_engines(self):
         self.assertBestPractice(
             [
                 Warning(
-                    f"Your default storage engine is not tenant aware.",
+                    "Your default storage engine is not tenant aware.",
                     hint="Set settings.STORAGES default backend to "
                     "'tenant_schemas.storage.TenantFileSystemStorage'",
                     id="tenant_schemas.W003",