diff options
Diffstat (limited to 'database.py')
| -rw-r--r-- | database.py | 113 |
1 files changed, 0 insertions, 113 deletions
diff --git a/database.py b/database.py deleted file mode 100644 index ccb00bc..0000000 --- a/database.py +++ /dev/null @@ -1,113 +0,0 @@ -#!/usr/bin/python - -import asyncio -import logging -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: - logging.warning("connecting to database") - 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() - logging.error("database is not ready, 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) |
