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
|
#!/usr/bin/python3
import asyncio
import os
import glob
import logging
import json
import asyncpg
import joblib.joblib
queue_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"
}
db_config = {
"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 Worker(object):
def __init__(self):
self._queue = None
self._pool = None
self._queries = {}
async def connect(self, dbconf, queuedb):
"""Connect to database."""
logging.warning("connecting to database")
self._queue = joblib.joblib.JobQueue()
await self._queue.connect(**queuedb)
await self._queue.setup()
self._pool = await asyncpg.create_pool(**dbconf)
async def setup(self):
"""Initialize the database."""
scriptroot = os.path.realpath(
os.path.join(os.getcwd(), os.path.dirname(__file__)))
for path in glob.glob(scriptroot + "/queries/*/*.sql"):
# utf-8-sig is used by pgadmin, doesn't hurt to specify
# directory names: web target table
table = os.path.basename(os.path.dirname(path))
with open(path, "r", encoding="utf-8-sig") as file:
try:
self._queries[table].append(file.read())
except KeyError:
self._queries[table] = [file.read()]
logging.info("loaded query '%s'", table)
async def _execute_job(self, jobid, payload):
"""Finish a job."""
object_id = payload["id"]
table = payload["type"]
if table not in self._queries:
return
async with self._pool.acquire() as con:
logging.debug("%s: compiling '%s' from '%s'",
jobid, object_id, table)
for query in self._queries[table]:
async with con.transaction():
await con.execute(query, object_id)
async def _work(self):
"""Fetch a job and run it."""
jobid, payload, _ = await self._queue.acquire(jobtype="compile")
if jobid is None:
raise LookupError("no jobs available")
await self._execute_job(jobid, payload)
await self._queue.finish(jobid)
async def run(self):
"""Start jobs forever."""
while True:
try:
await self._work()
except LookupError:
await asyncio.sleep(1)
async def start(self, number=1):
"""Start jobs in background."""
for _ in range(number):
asyncio.ensure_future(self.run())
async def startup():
worker = Worker()
await worker.connect(db_config, queue_db)
await worker.setup()
await worker.start(4)
logging.basicConfig(
filename="logs/apigrabber.log",
filemode="a",
level=logging.DEBUG
)
console = logging.StreamHandler()
console.setLevel(logging.WARNING)
logging.getLogger("").addHandler(console)
loop = asyncio.get_event_loop()
loop.run_until_complete(startup())
loop.run_forever()
|
