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