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
|
#!/usr/bin/python
import os
import glob
import logging
import asyncio
import asyncpg
source_db = {
"host": os.environ.get("POSTGRESQL_SOURCE_HOST") or "vaindock_postgres_raw",
"port": os.environ.get("POSTGRESQL_SOURCE_PORT") or 5532,
"user": os.environ.get("POSTGRESQL_SOURCE_USER") or "vainraw",
"password": os.environ.get("POSTGRESQL_SOURCE_PASSWORD") or "vainraw",
"database": os.environ.get("POSTGRESQL_SOURCE_DB") or "vainsocial-raw"
}
dest_db = {
"host": os.environ.get("POSTGRESQL_DEST_HOST") or "vaindock_postgres_web",
"port": os.environ.get("POSTGRESQL_DEST_PORT") or 5432,
"user": os.environ.get("POSTGRESQL_DEST_USER") or "vainweb",
"password": os.environ.get("POSTGRESQL_DEST_PASSWORD") or "vainweb",
"database": os.environ.get("POSTGRESQL_DEST_DB") or "vainsocial-web"
}
class Processor(object):
def __init__(self):
self._queries = {}
self._srcpool = self._destpool = None
def setup(self):
"""Load .sql files from queries/ directory.
File name is the table to insert into."""
self._queries = {}
scriptroot = os.path.realpath(
os.path.join(os.getcwd(), os.path.dirname(__file__)))
# glob *.sql
for fp in glob.glob(scriptroot + "/queries/*.sql"):
# utf-8-sig is used by pgadmin, doesn't hurt to specify
with open(fp, "r", encoding="utf-8-sig") as file:
table = os.path.splitext(os.path.basename(fp))[0]
logging.info("loaded query for '%s'", table)
self._queries[table] = file.read()
async def connect(self, source, dest):
"""Connect to database by arguments."""
self._srcpool = await asyncpg.create_pool(**source)
self._destpool = await asyncpg.create_pool(**dest)
async def run(self, stop_after=-1):
"""Execute a function for each row that is fetched with query."""
# TODO unnest
async with self._srcpool.acquire() as srccon:
async with self._destpool.acquire() as destcon:
for table, query in self._queries.items():
logging.info("running query for '%s'", table)
stop_after -= 1 # TODO for debugging, don't process whole table
async with srccon.transaction():
async with destcon.transaction():
async for rec in srccon.cursor(query):
await self.into(destcon, rec, table)
async def into(self, conn, data, table):
"""Insert a named tuple into a table."""
items = list(data.items())
keys, values = [x[0] for x in items], [x[1] for x in items]
placeholders = ["${}".format(i) for i, _ in enumerate(values, 1)]
query = "INSERT INTO {} (\"{}\") VALUES ({})".format(
table, "\", \"".join(keys), ", ".join(placeholders))
await conn.execute(query, (*data))
async def main():
pr = Processor()
await pr.connect(source_db, dest_db)
pr.setup()
await pr.run(stop_after=1000)
logging.basicConfig(level=logging.DEBUG)
loop = asyncio.get_event_loop()
loop.run_until_complete(main())
|
