diff --git a/mylib/pgsql.py b/mylib/pgsql.py index c960318..fc7f98c 100644 --- a/mylib/pgsql.py +++ b/mylib/pgsql.py @@ -5,6 +5,7 @@ import logging import sys import psycopg2 +from psycopg2.extras import RealDictCursor from mylib.db import DB, DBFailToConnect @@ -114,20 +115,16 @@ class PgDB(DB): :return: List of selected rows as dict on success, False otherwise :rtype: list, bool """ - cursor = self._conn.cursor() + cursor = self._conn.cursor(cursor_factory=RealDictCursor) try: self._log_query(sql, params) cursor.execute(sql, params) results = cursor.fetchall() - return results + return list(map(dict, results)) except psycopg2.Error: self._log_query_exception(sql, params) return False - @staticmethod - def _map_row_fields_by_index(fields, row): - return {field: row[idx] for idx, field in enumerate(fields)} - # # Depreated helpers # diff --git a/tests/test_pgsql.py b/tests/test_pgsql.py index dd6112c..d31cd4f 100644 --- a/tests/test_pgsql.py +++ b/tests/test_pgsql.py @@ -3,6 +3,7 @@ import psycopg2 import pytest +from psycopg2.extras import RealDictCursor from mylib.pgsql import PgDB @@ -57,6 +58,7 @@ class FakePsycopg2: expected_sql = None expected_params = None + expected_cursor_factory = None expected_return = True expected_just_try = False expected_exception = False @@ -81,7 +83,8 @@ class FakePsycopg2: raise psycopg2.Error(f"set_client_encoding({arg[0]}): Expected exception") return self.expected_return - def cursor(self): + def cursor(self, cursor_factory=None): + assert cursor_factory is self.expected_cursor_factory return FakePsycopg2Cursor( self.expected_sql, self.expected_params, @@ -455,6 +458,7 @@ def test_doSQL_on_exception(fake_connected_pgdb): def test_doSelect(fake_connected_pgdb): fake_connected_pgdb._conn.expected_sql = "SELECT * FROM table WHERE test1 = %(test1)s" fake_connected_pgdb._conn.expected_params = {"test1": 1} + fake_connected_pgdb._conn.expected_cursor_factory = RealDictCursor fake_connected_pgdb._conn.expected_return = [{"test1": 1}] assert ( fake_connected_pgdb.doSelect( @@ -466,6 +470,7 @@ def test_doSelect(fake_connected_pgdb): def test_doSelect_without_params(fake_connected_pgdb): fake_connected_pgdb._conn.expected_sql = "SELECT * FROM table" + fake_connected_pgdb._conn.expected_cursor_factory = RealDictCursor fake_connected_pgdb._conn.expected_return = [{"test1": 1}] assert ( fake_connected_pgdb.doSelect(fake_connected_pgdb._conn.expected_sql) @@ -474,6 +479,7 @@ def test_doSelect_without_params(fake_connected_pgdb): def test_doSelect_on_exception(fake_connected_pgdb): + fake_connected_pgdb._conn.expected_cursor_factory = RealDictCursor fake_connected_pgdb._conn.expected_exception = True assert fake_connected_pgdb.doSelect("SELECT * FROM table") is False @@ -481,6 +487,7 @@ def test_doSelect_on_exception(fake_connected_pgdb): def test_doSelect_just_try(fake_connected_just_try_pgdb): fake_connected_just_try_pgdb._conn.expected_sql = "SELECT * FROM table WHERE test1 = %(test1)s" fake_connected_just_try_pgdb._conn.expected_params = {"test1": 1} + fake_connected_just_try_pgdb._conn.expected_cursor_factory = RealDictCursor fake_connected_just_try_pgdb._conn.expected_return = [{"test1": 1}] assert ( fake_connected_just_try_pgdb.doSelect(