def __init__(self, rowhandler, table_name, table_mode, cf, log):
self.part_map = {}
self.rowhandler = rowhandler
self.table_name = table_name
self.quoted_name = quote_fqident(table_name)
self.log = log
if table_mode == 'direct':
self.split = False
elif table_mode == 'split':
self.split = True
smode = cf.get('split_mode', 'by-batch-time')
sfield = None
if smode.find(':') > 0:
smode, sfield = smode.split(':', 1)
self.split_field = sfield
self.split_part = cf.get(
'split_part', '%(table_name)s_%(year)s_%(month)s_%(day)s')
self.split_part_template = cf.get('split_part_template', '')
if smode == 'by-batch-time':
self.split_format = self.split_date_from_batch
elif smode == 'by-event-time':
self.split_format = self.split_date_from_event
elif smode == 'by-date-field':
self.split_format = self.split_date_from_field
else:
raise UsageError('Bad value for split_mode: ' + smode)
self.log.debug(
"%s: split_mode=%s, split_field=%s, split_part=%s" %
(self.table_name, smode, self.split_field, self.split_part))
elif table_mode == 'ignore':
pass
else:
raise UsageError('Bad value for table_mode: ' + table_mode)
def split_format(self, ev, data):
"""Generates part table name from template"""
if self.conf.part_mode == 'batch_time':
dtm = self.batch_info['batch_end']
elif self.conf.part_mode == 'event_time':
dtm = ev.ev_time
elif self.conf.part_mode == 'current_time':
dtm = datetime.datetime.now()
elif self.conf.part_mode == 'date_field':
dt_str = data[self.conf.part_field]
if dt_str is None:
raise Exception('part_field(%s) is NULL: %s' %
(self.conf.part_field, ev))
dtm = datetime.datetime.strptime(dt_str[:19], "%Y-%m-%d %H:%M:%S")
else:
raise UsageError('Bad value for part_mode: %s' %\
self.conf.part_mode)
vals = {
'parent': self.dest_table,
'year': "%04d" % dtm.year,
'month': "%02d" % dtm.month,
'day': "%02d" % dtm.day,
'hour': "%02d" % dtm.hour,
}
return (self.get_part_name() % vals, dtm)
def cmd_takeover(self, old_node_name):
"""Generic node switchover."""
self.log.info("old: %s" % old_node_name)
self.load_local_info()
new_node_name = self.options.node
if not new_node_name:
worker = self.options.consumer
if not worker:
raise UsageError('old node not given')
if self.queue_info.local_node.worker_name != worker:
raise UsageError('old node not given')
new_node_name = self.local_node
if not old_node_name:
raise UsageError('old node not given')
if old_node_name not in self.queue_info.member_map:
raise UsageError('Unknown node: %s' % old_node_name)
if self.options.dead_root:
otype = 'root'
failover = True
elif self.options.dead_branch:
otype = 'branch'
failover = True
else:
onode = self.get_node_info(old_node_name)
otype = onode.type
failover = False
if failover:
self.cmd_tag_dead(old_node_name)
new_node = self.get_node_info(new_node_name)
if old_node_name == new_node.name:
self.log.info("same node?")
return
if otype == 'root':
self.takeover_root(old_node_name, new_node_name, failover)
else:
self.takeover_nonroot(old_node_name, new_node_name, failover)
# switch subscribers around
if self.options.all or failover:
for n in self.find_subscribers_for(old_node_name):
self.node_change_provider(n, new_node_name)
def check_part(self, curs, dst, pkey_list):
if skytools.exists_table(curs, dst):
return
if not self.split_part_template:
raise UsageError('Partition %s does not exist and split_part_template not specified' % dst)
vals = {
'dest': quote_fqident(dst),
'part': quote_fqident(dst),
'parent': quote_fqident(self.table_name),
'pkey': ",".join(pkey_list), # quoting?
}
sql = self.split_part_template % vals
curs.execute(sql)
def init_state(self, tbl):
cf = self.cf
if tbl in cf.cf.sections():
cf = cf.clone(tbl)
table_mode = cf.get('table_mode', 'ignore')
row_mode = cf.get('row_mode', 'plain')
if table_mode == 'ignore':
tblhandler = IgnoreTable
else:
tblhandler = TableHandler
if row_mode == 'plain':
rowhandler = BasicLoader
elif row_mode == 'keep_latest':
rowhandler = KeepLatestLoader
elif row_mode == 'keep_all':
rowhandler = KeepAllLoader
elif row_mode == 'bulk':
rowhandler = BulkLoader
else:
raise UsageError('Bad row_mode: '+row_mode)
self.table_state[tbl] = tblhandler(rowhandler, tbl, table_mode, cf, self.log)
def cmd_drop_node(self, node_name):
"""Drop a node."""
self.load_local_info()
try:
node = self.load_node_info(node_name)
if node:
# see if we can safely drop
subscriber_list = self.get_node_subscriber_list(node_name)
if subscriber_list:
raise UsageError('node still has subscribers')
except skytools.DBError:
pass
try:
# unregister node location from root node (event will be added to queue)
root_db = self.find_root_db()
q = "select * from pgq_node.unregister_location(%s, %s)"
self.exec_cmd(root_db, q, [self.queue_name, node_name])
except skytools.DBError, d:
self.log.warning("Unregister from root failed: %s", str(d))
def create_node(self, node_type, node_name, node_location):
"""Generic node init."""
provider_loc = self.options.provider
if node_type not in ('root', 'branch', 'leaf'):
raise Exception('unknown node type')
# connect to database
db = self.get_database("new_node", connstr=node_location)
# check if code is installed
self.install_code(db)
# query current status
res = self.exec_query(db, "select * from pgq_node.get_node_info(%s)",
[self.queue_name])
info = res[0]
if info['node_type'] is not None:
self.log.info("Node is already initialized as %s" %
info['node_type'])
return
self.log.info("Initializing node")
node_attrs = {}
worker_name = self.options.worker
if not worker_name:
raise Exception('--worker required')
combined_queue = self.options.merge
if combined_queue and node_type != 'leaf':
raise Exception('--merge can be used only for leafs')
if self.options.sync_watermark:
if node_type != 'branch':
raise UsageError(
'--sync-watermark can be used only for branch nodes')
node_attrs['sync_watermark'] = self.options.sync_watermark
# register member
if node_type == 'root':
global_watermark = None
combined_queue = None
provider_name = None
self.exec_cmd(
db,
"select * from pgq_node.register_location(%s, %s, %s, false)",
[self.queue_name, node_name, node_location])
self.exec_cmd(
db,
"select * from pgq_node.create_node(%s, %s, %s, %s, %s, %s, %s)",
[
self.queue_name, node_type, node_name, worker_name,
provider_name, global_watermark, combined_queue
])
provider_db = None
else:
if not provider_loc:
raise Exception('Please specify --provider')
root_db = self.find_root_db(provider_loc)
queue_info = self.load_queue_info(root_db)
# check if member already exists
if queue_info.get_member(node_name) is not None:
self.log.error("Node '%s' already exists" % node_name)
sys.exit(1)
combined_set = None
provider_db = self.get_database('provider_db',
connstr=provider_loc)
q = "select node_type, node_name from pgq_node.get_node_info(%s)"
res = self.exec_query(provider_db, q, [self.queue_name])
row = res[0]
if not row['node_name']:
raise Exception("provider node not found")
provider_name = row['node_name']
# register member on root
self.exec_cmd(
root_db,
"select * from pgq_node.register_location(%s, %s, %s, false)",
[self.queue_name, node_name, node_location])
# lookup provider
provider = queue_info.get_member(provider_name)
if not provider:
self.log.error("Node %s does not exist" % provider_name)
sys.exit(1)
# register on provider
self.exec_cmd(
provider_db,
"select * from pgq_node.register_location(%s, %s, %s, false)",
[self.queue_name, node_name, node_location])
rows = self.exec_cmd(
provider_db,
"select * from pgq_node.register_subscriber(%s, %s, %s, null)",
[self.queue_name, node_name, worker_name])
global_watermark = rows[0]['global_watermark']
# initialize node itself
# insert members
self.exec_cmd(
db,
"select * from pgq_node.register_location(%s, %s, %s, false)",
[self.queue_name, node_name, node_location])
for m in queue_info.member_map.values():
self.exec_cmd(
db,
"select * from pgq_node.register_location(%s, %s, %s, %s)",
[self.queue_name, m.name, m.location, m.dead])
# real init
self.exec_cmd(
db,
"select * from pgq_node.create_node(%s, %s, %s, %s, %s, %s, %s)",
[
self.queue_name, node_type, node_name, worker_name,
provider_name, global_watermark, combined_queue
])
self.extra_init(node_type, db, provider_db)
if node_attrs:
s_attrs = skytools.db_urlencode(node_attrs)
self.exec_cmd(db, "select * from pgq_node.set_node_attrs(%s, %s)",
[self.queue_name, s_attrs])
self.log.info("Done")
