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
|
#!/usr/bin/python
import psycopg2
import psycopg2.extras
class Database(object):
"""Database wrapper class"""
def __init__(self, connstring):
"""Connects to the database.
:param connstring: Connection string containing user and database.
:type connstring: str
"""
self._connection = psycopg2.connect(connstring)
self._c = self._connection.cursor()
def upsert_type(self, shard, json, 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 json: Object to upsert.
:type json: dict
:param objtype: Object type and table name.
:type objtype: str
:param many: (optional) Whether `json` is an array
of objects of the same objtype.
:type many: bool
"""
if not many:
json = [json]
arr = [{
"id": shard + j["id"],
"data": psycopg2.extras.Json(j)
} for j in json]
self._c.execute("CREATE TABLE IF NOT EXISTS " + objtype +
" (id TEXT PRIMARY KEY, data json)")
self._c.executemany("INSERT INTO " + objtype + " " +
"VALUES (%(id)s, %(data)s) " +
"ON CONFLICT (id) DO " +
"UPDATE SET data=%(data)s",
arr)
self._connection.commit()
def upsert(self, shard, json, many=False):
"""Upserts an object into the corresponding database.
:param shard: Region in which an object's id is unique.
:type shard: str
:param json: Object to upsert.
:type json: dict
:param many: (optional) Whether `json` is an array
of objects.
:type many: bool
"""
if not many:
json = [json]
objectmap = dict()
# figure out the type of each object and sort into map
for j in json:
objtype = j["type"]
if objtype not in objectmap:
objectmap[objtype] = []
objectmap[objtype].append(j)
# execute bulk upsert for each type
for objtype, objects in objectmap.items():
self.upsert_type(shard, objects, objtype, True)
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
"""
self._c.execute(query)
return self._c.fetchall()
|
