"""Find table structure and allow CREATE/DROP elements from it. """ import re import skytools from skytools import quote_ident, quote_fqident __all__ = ['TableStruct', 'SeqStruct', 'T_TABLE', 'T_CONSTRAINT', 'T_INDEX', 'T_TRIGGER', 'T_RULE', 'T_GRANT', 'T_OWNER', 'T_PKEY', 'T_ALL', 'T_SEQUENCE', 'T_PARENT', 'T_DEFAULT'] T_TABLE = 1 << 0 T_CONSTRAINT = 1 << 1 T_INDEX = 1 << 2 T_TRIGGER = 1 << 3 T_RULE = 1 << 4 T_GRANT = 1 << 5 T_OWNER = 1 << 6 T_SEQUENCE = 1 << 7 T_PARENT = 1 << 8 T_DEFAULT = 1 << 9 T_PKEY = 1 << 20 # special, one of constraints T_ALL = ( T_TABLE | T_CONSTRAINT | T_INDEX | T_SEQUENCE | T_TRIGGER | T_RULE | T_GRANT | T_OWNER | T_DEFAULT ) # # Utility functions # def find_new_name(curs, name): """Create new object name for case the old exists. Needed when creating a new table besides old one. """ # cut off previous numbers m = re.search('_[0-9]+$', name) if m: name = name[:m.start()] # now loop for i in range(1, 1000): tname = "%s_%d" % (name, i) q = "select count(1) from pg_class where relname = %s" curs.execute(q, [tname]) if curs.fetchone()[0] == 0: return tname # failed raise Exception('find_new_name failed') def rx_replace(rx, sql, new_part): """Find a regex match and replace that part with new_part.""" m = re.search(rx, sql, re.I) if not m: raise Exception('rx_replace failed: rx=%r sql=%r new=%r' % (rx, sql, new_part)) p1 = sql[:m.start()] p2 = sql[m.end():] return p1 + new_part + p2 # # Schema objects # class TElem(object): """Keeps info about one metadata object.""" SQL = "" type = 0 def get_create_sql(self, curs, new_name = None): """Return SQL statement for creating or None if not supported.""" return None def get_drop_sql(self, curs): """Return SQL statement for dropping or None of not supported.""" return None @classmethod def get_load_sql(cls, pgver): """Return SQL statement for finding objects.""" return cls.SQL class TConstraint(TElem): """Info about constraint.""" type = T_CONSTRAINT SQL = """ SELECT c.conname as name, pg_get_constraintdef(c.oid) as def, c.contype, i.indisclustered as is_clustered FROM pg_constraint c LEFT JOIN pg_index i ON c.conrelid = i.indrelid AND c.conname = (SELECT r.relname FROM pg_class r WHERE r.oid = i.indexrelid) WHERE c.conrelid = %(oid)s AND c.contype != 'f' """ def __init__(self, table_name, row): """Init constraint.""" self.table_name = table_name self.name = row['name'] self.defn = row['def'] self.contype = row['contype'] self.is_clustered = row['is_clustered'] # tag pkeys if self.contype == 'p': self.type += T_PKEY def get_create_sql(self, curs, new_table_name=None): """Generate creation SQL.""" # no ONLY here as table with childs (only case that matters) # cannot have contraints that childs do not have fmt = "ALTER TABLE %s ADD CONSTRAINT %s\n %s;" if new_table_name: name = self.name if self.contype in ('p', 'u'): name = find_new_name(curs, self.name) qtbl = quote_fqident(new_table_name) qname = quote_ident(name) else: qtbl = quote_fqident(self.table_name) qname = quote_ident(self.name) sql = fmt % (qtbl, qname, self.defn) if self.is_clustered: sql +=' ALTER TABLE ONLY %s\n CLUSTER ON %s;' % (qtbl, qname) return sql def get_drop_sql(self, curs): """Generate removal sql.""" fmt = "ALTER TABLE ONLY %s\n DROP CONSTRAINT %s;" sql = fmt % (quote_fqident(self.table_name), quote_ident(self.name)) return sql class TIndex(TElem): """Info about index.""" type = T_INDEX SQL = """ SELECT n.nspname || '.' || c.relname as name, pg_get_indexdef(i.indexrelid) as defn, c.relname as local_name, i.indisclustered as is_clustered FROM pg_index i, pg_class c, pg_namespace n WHERE c.oid = i.indexrelid AND i.indrelid = %(oid)s AND n.oid = c.relnamespace AND NOT EXISTS (select objid from pg_depend where classid = %(pg_class_oid)s and objid = c.oid and deptype = 'i') """ def __init__(self, table_name, row): self.name = row['name'] self.defn = row['defn'].replace(' USING ', '\n USING ', 1) + ';' self.is_clustered = row['is_clustered'] self.table_name = table_name self.local_name = row['local_name'] def get_create_sql(self, curs, new_table_name = None): """Generate creation SQL.""" if new_table_name: # fixme: seems broken iname = find_new_name(curs, self.name) tname = new_table_name pnew = "INDEX %s ON %s " % (quote_ident(iname), quote_fqident(tname)) rx = r"\bINDEX[ ][a-z0-9._]+[ ]ON[ ][a-z0-9._]+[ ]" sql = rx_replace(rx, self.defn, pnew) else: sql = self.defn iname = self.local_name tname = self.table_name if self.is_clustered: sql += ' ALTER TABLE ONLY %s\n CLUSTER ON %s;' % ( quote_fqident(tname), quote_ident(iname)) return sql def get_drop_sql(self, curs): return 'DROP INDEX %s;' % quote_fqident(self.name) class TRule(TElem): """Info about rule.""" type = T_RULE SQL = """SELECT rw.*, pg_get_ruledef(rw.oid) as def FROM pg_rewrite rw WHERE rw.ev_class = %(oid)s AND rw.rulename <> '_RETURN'::name """ def __init__(self, table_name, row, new_name = None): self.table_name = table_name self.name = row['rulename'] self.defn = row['def'] self.enabled = row.get('ev_enabled', 'O') def get_create_sql(self, curs, new_table_name = None): """Generate creation SQL.""" if not new_table_name: sql = self.defn table = self.table_name else: idrx = r'''([a-z0-9._]+|"([^"]+|"")+")+''' # fixme: broken / quoting rx = r"\bTO[ ]" + idrx rc = re.compile(rx, re.X) m = rc.search(self.defn) if not m: raise Exception('Cannot find table name in rule') old_tbl = m.group(1) new_tbl = quote_fqident(new_table_name) sql = self.defn.replace(old_tbl, new_tbl) table = new_table_name if self.enabled != 'O': # O - rule fires in origin and local modes # D - rule is disabled # R - rule fires in replica mode # A - rule fires always action = {'R': 'ENABLE REPLICA', 'A': 'ENABLE ALWAYS', 'D': 'DISABLE'} [self.enabled] sql += ('\nALTER TABLE %s %s RULE %s;' % (table, action, self.name)) return sql def get_drop_sql(self, curs): return 'DROP RULE %s ON %s' % (quote_ident(self.name), quote_fqident(self.table_name)) class TTrigger(TElem): """Info about trigger.""" type = T_TRIGGER def __init__(self, table_name, row): self.table_name = table_name self.name = row['name'] self.defn = row['def'] + ';' self.defn = self.defn.replace('FOR EACH', '\n FOR EACH', 1) def get_create_sql(self, curs, new_table_name = None): """Generate creation SQL.""" if not new_table_name: return self.defn # fixme: broken / quoting rx = r"\bON[ ][a-z0-9._]+[ ]" pnew = "ON %s " % new_table_name return rx_replace(rx, self.defn, pnew) def get_drop_sql(self, curs): return 'DROP TRIGGER %s ON %s' % (quote_ident(self.name), quote_fqident(self.table_name)) @classmethod def get_load_sql(cls, pg_vers): """Return SQL statement for finding objects.""" sql = "SELECT tgname as name, pg_get_triggerdef(oid) as def "\ " FROM pg_trigger "\ " WHERE tgrelid = %(oid)s AND " if pg_vers >= 90000: sql += "NOT tgisinternal" else: sql += "NOT tgisconstraint" return sql class TParent(TElem): """Info about trigger.""" type = T_PARENT SQL = """ SELECT n.nspname||'.'||c.relname AS name FROM pg_inherits i JOIN pg_class c ON i.inhparent = c.oid JOIN pg_namespace n ON c.relnamespace = n.oid WHERE i.inhrelid = %(oid)s """ def __init__(self, table_name, row): self.name = table_name self.parent_name = row['name'] def get_create_sql(self, curs, new_table_name = None): return 'ALTER TABLE ONLY %s\n INHERIT %s' % (quote_fqident(self.name), quote_fqident(self.parent_name)) def get_drop_sql(self, curs): return 'ALTER TABLE ONLY %s\n NO INHERIT %s' % (quote_fqident(self.name), quote_fqident(self.parent_name)) class TOwner(TElem): """Info about table owner.""" type = T_OWNER SQL = """ SELECT pg_get_userbyid(relowner) as owner FROM pg_class WHERE oid = %(oid)s """ def __init__(self, table_name, row, new_name = None): self.table_name = table_name self.name = 'Owner' self.owner = row['owner'] def get_create_sql(self, curs, new_name = None): """Generate creation SQL.""" if not new_name: new_name = self.table_name return 'ALTER TABLE %s\n OWNER TO %s;' % (quote_fqident(new_name), quote_ident(self.owner)) class TGrant(TElem): """Info about permissions.""" type = T_GRANT SQL = "SELECT relacl FROM pg_class where oid = %(oid)s" # Sync with: src/include/utils/acl.h acl_map = { 'a': 'INSERT', 'r': 'SELECT', 'w': 'UPDATE', 'd': 'DELETE', 'D': 'TRUNCATE', 'x': 'REFERENCES', 't': 'TRIGGER', 'X': 'EXECUTE', 'U': 'USAGE', 'C': 'CREATE', 'T': 'TEMPORARY', 'c': 'CONNECT', # old 'R': 'RULE', } def acl_to_grants(self, acl): if acl == "arwdRxt": # ALL for tables return "ALL" i = 0 lst1 = [] lst2 = [] while i < len(acl): a = self.acl_map[acl[i]] if i+1 < len(acl) and acl[i+1] == '*': lst2.append(a) i += 2 else: lst1.append(a) i += 1 return ", ".join(lst1), ", ".join(lst2) def parse_relacl(self, relacl): """Parse ACL to tuple of (user, acl, who)""" if relacl is None: return [] tup_list = [] for sacl in skytools.parse_pgarray(relacl): acl = skytools.parse_acl(sacl) if not acl: continue tup_list.append(acl) return tup_list def __init__(self, table_name, row, new_name = None): self.name = table_name self.acl_list = self.parse_relacl(row['relacl']) def get_create_sql(self, curs, new_name = None): """Generate creation SQL.""" if not new_name: new_name = self.name qtarget = quote_fqident(new_name) sql_list = [] for role, acl, who in self.acl_list: qrole = quote_ident(role) astr1, astr2 = self.acl_to_grants(acl) if astr1: sql = "GRANT %s ON %s\n TO %s;" % (astr1, qtarget, qrole) sql_list.append(sql) if astr2: sql = "GRANT %s ON %s\n TO %s WITH GRANT OPTION;" % (astr2, qtarget, qrole) sql_list.append(sql) return "\n".join(sql_list) def get_drop_sql(self, curs): sql_list = [] for user, acl, who in self.acl_list: sql = "REVOKE ALL FROM %s ON %s;" % (quote_ident(user), quote_fqident(self.name)) sql_list.append(sql) return "\n".join(sql_list) class TColumnDefault(TElem): """Info about table column default value.""" type = T_DEFAULT SQL = """ select a.attname as name, pg_get_expr(d.adbin, d.adrelid) as expr from pg_attribute a left join pg_attrdef d on (d.adrelid = a.attrelid and d.adnum = a.attnum) where a.attrelid = %(oid)s and not a.attisdropped and a.atthasdef and a.attnum > 0 order by a.attnum; """ def __init__(self, table_name, row): self.table_name = table_name self.name = row['name'] self.expr = row['expr'] def get_create_sql(self, curs, new_name = None): """Generate creation SQL.""" tbl = new_name or self.table_name sql = "ALTER TABLE ONLY %s ALTER COLUMN %s\n SET DEFAULT %s;" % ( quote_fqident(tbl), quote_ident(self.name), self.expr) return sql def get_drop_sql(self, curs): return "ALTER TABLE %s ALTER COLUMN %s\n DROP DEFAULT;" % ( quote_fqident(self.table_name), quote_ident(self.name)) class TColumn(TElem): """Info about table column.""" SQL = """ select a.attname as name, quote_ident(a.attname) as qname, format_type(a.atttypid, a.atttypmod) as dtype, a.attnotnull, (select max(char_length(aa.attname)) from pg_attribute aa where aa.attrelid = %(oid)s) as maxcol, pg_get_serial_sequence(%(fq2name)s, a.attname) as seqname from pg_attribute a left join pg_attrdef d on (d.adrelid = a.attrelid and d.adnum = a.attnum) where a.attrelid = %(oid)s and not a.attisdropped and a.attnum > 0 order by a.attnum; """ seqname = None def __init__(self, table_name, row): self.name = row['name'] fname = row['qname'].ljust(row['maxcol'] + 3) self.column_def = fname + ' ' + row['dtype'] if row['attnotnull']: self.column_def += ' not null' self.sequence = None if row['seqname']: self.seqname = skytools.unquote_fqident(row['seqname']) class TGPDistKey(TElem): """Info about GreenPlum table distribution keys""" SQL = """ select a.attname as name from pg_attribute a, gp_distribution_policy p where p.localoid = %(oid)s and a.attrelid = %(oid)s and a.attnum = any(p.attrnums) order by a.attnum; """ def __init__(self, table_name, row): self.name = row['name'] class TTable(TElem): """Info about table only (columns).""" type = T_TABLE def __init__(self, table_name, col_list, dist_key_list = None): self.name = table_name self.col_list = col_list self.dist_key_list = dist_key_list def get_create_sql(self, curs, new_name = None): """Generate creation SQL.""" if not new_name: new_name = self.name sql = "CREATE TABLE %s (" % quote_fqident(new_name) sep = "\n " for c in self.col_list: sql += sep + c.column_def sep = ",\n " sql += "\n)" if self.dist_key_list is not None: if self.dist_key_list != []: sql += "\ndistributed by(%s)" % ','.join(c.name for c in self.dist_key_list) else: sql += '\ndistributed randomly' sql += ";" return sql def get_drop_sql(self, curs): return "DROP TABLE %s;" % quote_fqident(self.name) class TSeq(TElem): """Info about sequence.""" type = T_SEQUENCE SQL = """SELECT *, %(owner)s as "owner" from %(fqname)s """ def __init__(self, seq_name, row): self.name = seq_name defn = '' self.owner = row['owner'] if row['increment_by'] != 1: defn += ' INCREMENT BY %d' % row['increment_by'] if row['min_value'] != 1: defn += ' MINVALUE %d' % row['min_value'] if row['max_value'] != 9223372036854775807: defn += ' MAXVALUE %d' % row['max_value'] last_value = row['last_value'] if row['is_called']: last_value += row['increment_by'] if last_value >= row['max_value']: raise Exception('duh, seq passed max_value') if last_value != 1: defn += ' START %d' % last_value if row['cache_value'] != 1: defn += ' CACHE %d' % row['cache_value'] if row['is_cycled']: defn += ' CYCLE ' if self.owner: defn += ' OWNED BY %s' % self.owner self.defn = defn def get_create_sql(self, curs, new_seq_name = None): """Generate creation SQL.""" # we are in table def, forget full def if self.owner: sql = "ALTER SEQUENCE %s\n OWNED BY %s;" % ( quote_fqident(self.name), self.owner ) return sql name = self.name if new_seq_name: name = new_seq_name sql = 'CREATE SEQUENCE %s %s;' % (quote_fqident(name), self.defn) return sql def get_drop_sql(self, curs): if self.owner: return '' return 'DROP SEQUENCE %s;' % quote_fqident(self.name) # # Main table object, loads all the others # class BaseStruct(object): """Collects and manages all info about a higher-level db object. Allow to issue CREATE/DROP statements about any group of elements. """ object_list = [] def __init__(self, curs, name): """Initializes class by loading info about table_name from database.""" self.name = name self.fqname = quote_fqident(name) def _load_elem(self, curs, name, args, eclass): """Fetch element(s) from db.""" elem_list = [] #print "Loading %s, name=%s, args=%s" % (repr(eclass), repr(name), repr(args)) sql = eclass.get_load_sql(curs.connection.server_version) curs.execute(sql % args) for row in curs.fetchall(): elem_list.append(eclass(name, row)) return elem_list def create(self, curs, objs, new_table_name = None, log = None): """Issues CREATE statements for requested set of objects. If new_table_name is giver, creates table under that name and also tries to rename all indexes/constraints that conflict with existing table. """ for o in self.object_list: if o.type & objs: sql = o.get_create_sql(curs, new_table_name) if not sql: continue if log: log.info('Creating %s' % o.name) log.debug(sql) curs.execute(sql) def drop(self, curs, objs, log = None): """Issues DROP statements for requested set of objects.""" # make sure the creating & dropping happen in reverse order olist = self.object_list[:] olist.reverse() for o in olist: if o.type & objs: sql = o.get_drop_sql(curs) if not sql: continue if log: log.info('Dropping %s' % o.name) log.debug(sql) curs.execute(sql) def get_create_sql(self, objs): res = [] for o in self.object_list: if o.type & objs: sql = o.get_create_sql(None, None) if sql: res.append(sql) return "".join(res) class TableStruct(BaseStruct): """Collects and manages all info about table. Allow to issue CREATE/DROP statements about any group of elements. """ def __init__(self, curs, table_name): """Initializes class by loading info about table_name from database.""" BaseStruct.__init__(self, curs, table_name) self.table_name = table_name # fill args schema, name = skytools.fq_name_parts(table_name) args = { 'schema': schema, 'table': name, 'fqname': self.fqname, 'fq2name': skytools.quote_literal(self.fqname), 'oid': skytools.get_table_oid(curs, table_name), 'pg_class_oid': skytools.get_table_oid(curs, 'pg_catalog.pg_class'), } # load table struct self.col_list = self._load_elem(curs, self.name, args, TColumn) # if db is GP then read also table distribution keys if skytools.exists_table(curs, "pg_catalog.gp_distribution_policy"): self.dist_key_list = self._load_elem(curs, self.name, args, TGPDistKey) else: self.dist_key_list = None self.object_list = [ TTable(table_name, self.col_list, self.dist_key_list) ] self.seq_list = [] # load seqs for col in self.col_list: if col.seqname: fqname = quote_fqident(col.seqname) owner = self.fqname + '.' + quote_ident(col.name) seq_args = { 'fqname': fqname, 'owner': skytools.quote_literal(owner) } self.seq_list += self._load_elem(curs, col.seqname, seq_args, TSeq) self.object_list += self.seq_list # load additional objects to_load = [TColumnDefault, TConstraint, TIndex, TTrigger, TRule, TGrant, TOwner, TParent] for eclass in to_load: self.object_list += self._load_elem(curs, self.name, args, eclass) def get_column_list(self): """Returns list of column names the table has.""" res = [] for c in self.col_list: res.append(c.name) return res class SeqStruct(BaseStruct): """Collects and manages all info about sequence. Allow to issue CREATE/DROP statements about any group of elements. """ def __init__(self, curs, seq_name): """Initializes class by loading info about table_name from database.""" BaseStruct.__init__(self, curs, seq_name) # fill args args = { 'fqname': self.fqname, 'owner': 'null' } # load table struct self.object_list = self._load_elem(curs, seq_name, args, TSeq) def test(): from skytools import connect_database db = connect_database("dbname=fooz") curs = db.cursor() s = TableStruct(curs, "public.data1") s.drop(curs, T_ALL) s.create(curs, T_ALL) s.create(curs, T_ALL, "data1_new") s.create(curs, T_PKEY) if __name__ == '__main__': test()