1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
|
#!/usr/bin/python
import asyncio
import json
import asyncpg
class Database(object):
"""Database wrapper class"""
def __init__(self):
self._pool = None
async def connect(self, host, port, user, password, database):
"""Connects to the database.
:param connstring: Connection string containing user and database.
:type connstring: str
"""
while True: # retry until connection succeeds
try:
print("attempting to connect to db…")
self._pool = await asyncpg.create_pool(
host=host, port=port, user=user,
password=password, database=database)
break
except asyncpg.exceptions.CannotConnectNowError:
await self._pool.close()
print("Database is not ready yet. Retrying…")
await asyncio.sleep(5)
async def upsert_type(self, obj, objtype, many=False):
"""Upserts an object of given `objtype` into the corresponding database.
:param obj: Object to upsert.
:type obj: dict or list
:param objtype: Object type and table name.
:type objtype: str
:param many: (optional) Whether `obj` is a list
of objects of the same objtype.
:type many: bool
"""
if not many:
obj = [obj]
arr = [[
(j["attributes"].get("shardId") or "") + j["id"],
json.dumps(j)
] for j in obj]
async with self._pool.acquire() as conn:
async with conn.transaction():
await conn.execute("CREATE TABLE IF NOT EXISTS " + objtype +
" (id TEXT PRIMARY KEY, data jsonb)")
await conn.executemany("INSERT INTO " + objtype + " " +
"VALUES ($1, $2) " +
"ON CONFLICT (id) DO " +
"UPDATE SET data=$2",
arr)
async def upsert(self, obj, many=False):
"""Upserts an object into the corresponding database.
:param obj: Object to upsert.
:type obj: dict
:param many: (optional) Whether `obj` is a list
of objects.
:type many: bool
"""
if not many:
obj = [obj]
objectmap = dict()
# figure out the type of each object and sort into map
for j in obj:
objtype = j["type"]
if objtype not in objectmap:
objectmap[objtype] = []
objectmap[objtype].append(j)
# execute bulk upsert for each type
tasks = []
for objtype, objects in objectmap.items():
task = asyncio.ensure_future(
self.upsert_type(objects, objtype, True))
tasks.append(task)
await asyncio.gather(*tasks)
async def execute(self, query, args):
"""Runs an SQL statement.
:param query: SQL query to execute.
:type query: str
:param args: Query arguments.
:type args: list of objects
"""
async with self._pool.acquire() as conn:
async with conn.transaction():
await conn.execute(query, args)
async def select(self, query, *args):
"""Returns the result of an SQL query.
:param query: SQL query to execute.
:type query: str
:param args: Query arguments.
:return: List of results.
:rtype: list of dict
"""
async with self._pool.acquire() as conn:
async with conn.transaction():
return await conn.fetch(query, *args)
|
