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
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
|
#!/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, connstring):
"""Connects to the database.
:param connstring: Connection string containing user and database.
:type connstring: str
"""
self._pool = await asyncpg.create_pool(connstring)
async def upsert_type(self, shard, obj, objtype, many=False):
"""Upserts an object of given `objtype` into the corresponding database.
:param shard: Region in which an object's id is unique.
:type shard: str
: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 = [[
shard + 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 json)")
await conn.executemany("INSERT INTO " + objtype + " " +
"VALUES ($1, $2) " +
"ON CONFLICT (id) DO " +
"UPDATE SET data=$2",
arr)
async def upsert(self, shard, obj, many=False):
"""Upserts an object into the corresponding database.
:param shard: Region in which an object's id is unique.
:type shard: str
: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(shard, objects, objtype, True))
tasks.append(task)
await asyncio.gather(*tasks)
async def meta(self, key, value=None):
"""Sets or gets data into a dictionary-like database object.
:param key: ID of the value.
:type key: str
:param value: (optional) Value to set.
:type value: object
:return: If `value` is None, the value the database entry.
:rtype: object
"""
async with self._pool.acquire() as conn:
async with conn.transaction():
await conn.execute("CREATE TABLE IF NOT EXISTS meta " +
"(key TEXT PRIMARY KEY, value json)")
if value is not None:
value = json.dumps(value)
await conn.execute("INSERT INTO meta(key, value) " +
"VALUES ($1, $2) " +
"ON CONFLICT (key) DO " +
"UPDATE SET value=$2",
key, value)
else:
result = await conn.fetch("SELECT value FROM meta " +
"WHERE key='" + key + "'")
if len(result) == 0:
raise KeyError
return json.loads(result[0]["value"])
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):
"""Returns the result of an SQL query.
:param query: SQL query to execute.
:type query: str
:return: List of results.
:rtype: list of dict
"""
async with self._pool.acquire() as conn:
async with conn.transaction():
return await conn.fetch(query)
|
