Refactor how we provide cursor factory (#376)

Author: asherf

django_prometheus/db/backends/common.py

@@ -0,0 +1,9 @@
+from django import VERSION
+
+
+def get_postgres_cursor_class():
+    if VERSION < (4, 2):
+        from psycopg2.extensions import cursor as cursor_cls
+    else:
+        from django.db.backends.postgresql.base import Cursor as cursor_cls
+    return cursor_cls

django_prometheus/db/backends/postgis/base.py

@@ -1,19 +1,16 @@
-from django import VERSION
 from django.contrib.gis.db.backends.postgis import base
 
+from django_prometheus.db.backends.common import get_postgres_cursor_class
 from django_prometheus.db.common import DatabaseWrapperMixin, ExportingCursorWrapper
 
-if VERSION < (4, 2):
-    from psycopg2.extensions import cursor as cursor_cls
-else:
-    from django.db.backends.postgresql.base import Cursor as cursor_cls
-
 
 class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
-    def get_connection_params(self):
-        conn_params = super().get_connection_params()
-        conn_params["cursor_factory"] = ExportingCursorWrapper(cursor_cls, "postgis", self.vendor)
-        return conn_params
+    def get_new_connection(self, *args, **kwargs):
+        conn = super().get_new_connection(*args, **kwargs)
+        conn.cursor_factory = ExportingCursorWrapper(
+            conn.cursor_factory or get_postgres_cursor_class(), "postgis", self.vendor
+        )
+        return conn
 
     def create_cursor(self, name=None):
         # cursor_factory is a kwarg to connect() so restore create_cursor()'s

django_prometheus/db/backends/postgresql/base.py

@@ -1,19 +1,16 @@
-from django import VERSION
 from django.contrib.gis.db.backends.postgis import base
 
+from django_prometheus.db.backends.common import get_postgres_cursor_class
 from django_prometheus.db.common import DatabaseWrapperMixin, ExportingCursorWrapper
 
-if VERSION < (4, 2):
-    from psycopg2.extensions import cursor as cursor_cls
-else:
-    from django.db.backends.postgresql.base import Cursor as cursor_cls
-
 
 class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
-    def get_connection_params(self):
-        conn_params = super().get_connection_params()
-        conn_params["cursor_factory"] = ExportingCursorWrapper(cursor_cls, self.alias, self.vendor)
-        return conn_params
+    def get_new_connection(self, *args, **kwargs):
+        conn = super().get_new_connection(*args, **kwargs)
+        conn.cursor_factory = ExportingCursorWrapper(
+            conn.cursor_factory or get_postgres_cursor_class(), self.alias, self.vendor
+        )
+        return conn
 
     def create_cursor(self, name=None):
         # cursor_factory is a kwarg to connect() so restore create_cursor()'s