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")