def start_queue(
connection_string,
job_filepath,
schema_name="public",
workers=10,
cooperative=None,
schedule_frequency=None,
job_args=None,
children=None,
):
"""Creates, configures and runs a job queue.
:param connection_string: a Postgres connection string
:param job_filepath: the path to a python script that can configure the queue
:param schema_name: the name of the schema that contains the queue tables
:param workers: the number of concurrent workers to run
"""
model_access = ModelAccess(get_pg_core(connection_string), search_path=schema_name)
model_access.open(autocommit=True)
queue = Queue(model_access, worker_count=workers, cooperative=cooperative, schedule_frequency=schedule_frequency)
job_module = load_module(job_filepath, None)
job_module.setup_jobs(queue, job_args)
def stop(sig, _):
"""Stops the queue in the manner specified by the signal.
:param sig: the signal receieved
"""
queue.stop(stop_mode=STOP_SIGNALS[sig])
if children:
for pid in children:
os.kill(pid, sig)
try:
os.waitpid(pid, 0)
except ChildProcessError:
# Child already shut down before we started waiting on it.
pass
for sig in STOP_SIGNALS:
signal.signal(sig, stop)
signal.signal(signal.SIGINFO, lambda n, f: print(queue.status(), file=sys.stderr))
log_queue_info(job_filepath, workers)
queue.start()
class TransferTracker(object):
"""The TransferTracker class keeps track of data being transferred from one database or table to
another. Specifically, it makes it possible to record the old and new ids for a relation, and
query for the new ID.
"""
def __init__(self, data_access_core):
"""Initialize the TransferTracker class."""
self.data_access = ModelAccess(data_access_core)
def open(self, autocommit=False):
"""Call-through to data_access.open."""
self.data_access.open(autocommit=autocommit)
return self
def close(self, commit=True):
"""Call-through to data_access.close."""
self.data_access.close(commit=commit)
return self
def commit(self):
"""Call-through to data_access.commit."""
self.data_access.commit()
return self
def rollback(self):
"""Call-through to data_access.rollback."""
self.data_access.rollback()
return self
def setup(self):
"""Creates the default relations and transfers tables. The SQL used may not work on all
databases. (It was written for SQLite3)
"""
cmds = [
"""
create table if not exists relations (
id integer not null primary key,
name text not null unique,
completed_at datetime
);
""",
"""
create table if not exists transfers (
relation_id integer not null references relations (id) on delete cascade,
old_id text not null,
new_id text,
primary key (relation_id, old_id)
);
""",
"""
create index if not exists transfers_relation_id_idx on transfers (relation_id);
"""
]
for cmd in cmds:
self.data_access.execute(cmd)
self.data_access.commit()
return self
def cleanup(self):
"""Cleanup the database, calling vacuum and analyze."""
self.data_access.execute("vacuum")
self.data_access.execute("analyze")
def reset(self, relation_name=None):
"""Reset the transfer info for a particular relation, or if none is given, for all relations.
"""
if relation_name is not None:
self.data_access.delete("relations", dict(name=relation_name))
else:
self.data_access.delete("relations", "1=1")
return self
def start_transfer(self, relation_name):
"""Write records to the data source indicating that a transfer has been started for a particular
relation.
"""
self.reset(relation_name)
relation = Relation(name=relation_name)
self.data_access.insert_model(relation)
return relation
def register_transfer(self, relation, old_id, new_id):
"""Register the old and new ids for a particular record in a relation."""
transfer = Transfer(relation_id=relation.id, old_id=old_id, new_id=new_id)
self.data_access.insert_model(transfer)
return transfer
def complete_transfer(self, relation, cleanup=True):
"""Write records to the data source indicating that a transfer has been completed for a
particular relation.
"""
relation.completed_at = utc_now().isoformat()
self.data_access.update_model(relation)
if cleanup:
self.cleanup()
return relation
def is_transfer_complete(self, relation_name):
"""Checks to see if a tansfer has been completed."""
phold = self.data_access.sql_writer.to_placeholder()
return self.data_access.find_model(
Relation,
("name = {0} and completed_at is not null".format(phold), [relation_name])) is not None
def get_new_id(self, relation_name, old_id, strict=False):
"""Given a relation name and its old ID, get the new ID for a relation. If strict is true, an
error is thrown if no record is found for the relation and old ID.
"""
record = self.data_access.find(
"relations as r inner join transfers as t on r.id = t.relation_id",
(("r.name", relation_name), ("t.old_id", old_id)),
columns="new_id")
if record:
return record[0]
else:
if strict:
raise KeyError("{0} with id {1} not found".format(relation_name, old_id))
def id_getter(self, relation_name, strict=False):
"""Returns a function that accepts an old_id and returns the new ID for the enclosed relation
name."""
def get_id(old_id):
"""Get the new ID for the enclosed relation, given an old ID."""
return self.get_new_id(relation_name, old_id, strict)
return get_id
