Initial Commit

cursor.py

@@ -0,0 +1,42 @@
+import logging
+from typing import Optional
+
+from psycopg import AsyncCursor, sql
+from psycopg.abc import Query, Params
+from psycopg._encodings import conn_encoding
+
+logger = logging.getLogger(__name__)
+
+
+class AsyncLoggingCursor(AsyncCursor):
+    def mogrify_query(self, query: Query):
+        if isinstance(query, str):
+            msg = query
+        elif isinstance(query, (sql.SQL, sql.Composed)):
+            msg = query.as_string(self)
+        elif isinstance(query, bytes):
+            msg = query.decode(conn_encoding(self._conn.pgconn), 'replace')
+        else:
+            msg = repr(query)
+        return msg
+
+    async def execute(self, query: Query, params: Optional[Params] = None, **kwargs):
+        if logging.DEBUG >= logger.getEffectiveLevel():
+            msg = self.mogrify_query(query)
+            logger.debug(
+                "Executing query (%s) with values %s", msg, params,
+                extra={'action': "Query Execute"}
+            )
+        try:
+            return await super().execute(query, params=params, **kwargs)
+        except Exception:
+            msg = self.mogrify_query(query)
+            logger.exception(
+                "Exception during query execution. Query (%s) with parameters %s.",
+                msg, params,
+                extra={'action': "Query Execute"},
+                stack_info=True
+            )
+        else:
+            # TODO: Possibly log execution time
+            pass