Skip to content

Commit 4a84e26

Browse files
Add code refactored from seed-python-docker-rabbitmq repo
1 parent ac46e77 commit 4a84e26

File tree

3 files changed

+11
-15
lines changed

3 files changed

+11
-15
lines changed

db_wrapper/client.py

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -105,9 +105,5 @@ async def execute_and_return(
105105
async with self._connection.cursor() as cursor:
106106
await self._execute_query(cursor, query, params)
107107

108-
result = []
109-
110-
async for row in cursor:
111-
result.append(row)
112-
108+
result: List[T] = await cursor.fetchall()
113109
return result

db_wrapper/connection.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import logging
66
from typing import Optional
77

8+
from psycopg2.extras import RealDictCursor # type: ignore
89
from psycopg2 import OperationalError # type: ignore
910

1011
# no stubs available, starting out by just determining correct return
@@ -48,7 +49,9 @@ async def _try_connect(
4849

4950
while connection is None:
5051
try:
51-
connection = await aiopg.connect(dsn)
52+
connection = await aiopg.connect(
53+
dsn,
54+
cursor_factory=RealDictCursor)
5255
except OperationalError as err:
5356
print(type(err))
5457
if retries > 12:

db_wrapper/model.py

Lines changed: 6 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -63,18 +63,15 @@ def __init__(self, client: Client, table: sql.Composable) -> None:
6363
async def one(self, item: T) -> T:
6464
"""Create one new record with a given item."""
6565
columns: List[sql.Identifier] = []
66-
values: List[sql.Identifier] = []
66+
values: List[sql.Literal] = []
6767

6868
for column, value in item.items():
69-
if column == '_id':
70-
values.append(sql.Identifier(str(value)))
71-
else:
72-
values.append(sql.Identifier(value))
69+
values.append(sql.Literal(value))
7370

7471
columns.append(sql.Identifier(column))
7572

7673
query = sql.SQL(
77-
'INSERT INTO {table}({columns}) '
74+
'INSERT INTO {table} ({columns}) '
7875
'VALUES ({values}) '
7976
'RETURNING *;'
8077
).format(
@@ -103,10 +100,10 @@ async def one_by_id(self, id_value: str) -> T:
103100
query = sql.SQL(
104101
'SELECT * '
105102
'FROM {table} '
106-
'WHERE id = {id_value};'
103+
'WHERE _id = {id_value};'
107104
).format(
108105
table=self._table,
109-
id_value=sql.Identifier(id_value)
106+
id_value=sql.Literal(id_value)
110107
)
111108

112109
result: List[T] = await self._client.execute_and_return(query)
@@ -178,7 +175,7 @@ def __init__(self, client: Client, table: sql.Composable) -> None:
178175
self._client = client
179176
self._table = table
180177

181-
async def one(self, id_value: str) -> T:
178+
async def one_by_id(self, id_value: str) -> T:
182179
"""Delete one record with matching ID."""
183180
query = sql.SQL(
184181
'DELETE FROM {table} '

0 commit comments

Comments
 (0)