1import psycopg2.extensions
2from django.db.backends.postgresql import base
3
4from django_prometheus.db.common import DatabaseWrapperMixin, ExportingCursorWrapper
5
6
7class DatabaseFeatures(base.DatabaseFeatures):
8    """Our database has the exact same features as the base one."""
9
10    pass
11
12
13class DatabaseWrapper(DatabaseWrapperMixin, base.DatabaseWrapper):
14    def get_connection_params(self):
15        conn_params = super().get_connection_params()
16        conn_params["cursor_factory"] = ExportingCursorWrapper(
17            psycopg2.extensions.cursor, self.alias, self.vendor
18        )
19        return conn_params
20
21    def create_cursor(self, name=None):
22        # cursor_factory is a kwarg to connect() so restore create_cursor()'s
23        # default behavior
24        return base.DatabaseWrapper.create_cursor(self, name=name)
25