skytools.dbstruct

1 """Find table structure and allow CREATE/DROP elements from it. 2 """ 3 4 import sys, re 5 6 from sqltools import fq_name_parts, get_table_oid 7 8 __all__ = ['TableStruct', 9 'T_TABLE', 'T_CONSTRAINT', 'T_INDEX', 'T_TRIGGER', 10 'T_RULE', 'T_GRANT', 'T_OWNER', 'T_PKEY', 'T_ALL'] 11 12 T_TABLE = 1 << 0 13 T_CONSTRAINT = 1 << 1 14 T_INDEX = 1 << 2 15 T_TRIGGER = 1 << 3 16 T_RULE = 1 << 4 17 T_GRANT = 1 << 5 18 T_OWNER = 1 << 6 19 T_PKEY = 1 << 20 # special, one of constraints 20 T_ALL = ( T_TABLE | T_CONSTRAINT | T_INDEX 21 | T_TRIGGER | T_RULE | T_GRANT | T_OWNER ) 22 23 # 24 # Utility functions 25 # 26

27 -def find_new_name(curs, name):

28 """Create new object name for case the old exists. 29 30 Needed when creating a new table besides old one. 31 """ 32 # cut off previous numbers 33 m = re.search('_[0-9]+$', name) 34 if m: 35 name = name[:m.start()] 36 37 # now loop 38 for i in range(1, 1000): 39 tname = "%s_%d" % (name, i) 40 q = "select count(1) from pg_class where relname = %s" 41 curs.execute(q, [tname]) 42 if curs.fetchone()[0] == 0: 43 return tname 44 45 # failed 46 raise Exception('find_new_name failed')

47

48 -def rx_replace(rx, sql, new_part):

49 """Find a regex match and replace that part with new_part.""" 50 m = re.search(rx, sql, re.I) 51 if not m: 52 raise Exception('rx_replace failed') 53 p1 = sql[:m.start()] 54 p2 = sql[m.end():] 55 return p1 + new_part + p2

56 57 # 58 # Schema objects 59 # 60

61 -class TElem(object):

62 """Keeps info about one metadata object.""" 63 SQL = "" 64 type = 0

65 - def get_create_sql(self, curs):

66 """Return SQL statement for creating or None if not supported.""" 67 return None

68 - def get_drop_sql(self, curs):

69 """Return SQL statement for dropping or None of not supported.""" 70 return None

71

72 -class TConstraint(TElem):

73 """Info about constraint.""" 74 type = T_CONSTRAINT 75 SQL = """ 76 SELECT conname as name, pg_get_constraintdef(oid) as def, contype 77 FROM pg_constraint WHERE conrelid = %(oid)s 78 """

79 - def __init__(self, table_name, row):

80 self.table_name = table_name 81 self.name = row['name'] 82 self.defn = row['def'] 83 self.contype = row['contype'] 84 85 # tag pkeys 86 if self.contype == 'p': 87 self.type += T_PKEY

88

89 - def get_create_sql(self, curs, new_table_name=None):

90 fmt = "ALTER TABLE ONLY %s ADD CONSTRAINT %s %s;" 91 if new_table_name: 92 name = self.name 93 if self.contype in ('p', 'u'): 94 name = find_new_name(curs, self.name) 95 sql = fmt % (new_table_name, name, self.defn) 96 else: 97 sql = fmt % (self.table_name, self.name, self.defn) 98 return sql

99

100 - def get_drop_sql(self, curs):

101 fmt = "ALTER TABLE ONLY %s DROP CONSTRAINT %s;" 102 sql = fmt % (self.table_name, self.name) 103 return sql

104

105 -class TIndex(TElem):

106 """Info about index.""" 107 type = T_INDEX 108 SQL = """ 109 SELECT n.nspname || '.' || c.relname as name, 110 pg_get_indexdef(i.indexrelid) as defn 111 FROM pg_index i, pg_class c, pg_namespace n 112 WHERE c.oid = i.indexrelid AND i.indrelid = %(oid)s 113 AND n.oid = c.relnamespace 114 AND NOT EXISTS 115 (select objid from pg_depend 116 where classid = %(pg_class_oid)s 117 and objid = c.oid 118 and deptype = 'i') 119 """

120 - def __init__(self, table_name, row):

121 self.name = row['name'] 122 self.defn = row['defn'] + ';'

123

124 - def get_create_sql(self, curs, new_table_name = None):

125 if not new_table_name: 126 return self.defn 127 name = find_new_name(curs, self.name) 128 pnew = "INDEX %s ON %s " % (name, new_table_name) 129 rx = r"\bINDEX[ ][a-z0-9._]+[ ]ON[ ][a-z0-9._]+[ ]" 130 sql = rx_replace(rx, self.defn, pnew) 131 return sql

132 - def get_drop_sql(self, curs):

133 return 'DROP INDEX %s;' % self.name

134

135 -class TRule(TElem):

136 """Info about rule.""" 137 type = T_RULE 138 SQL = """ 139 SELECT rulename as name, pg_get_ruledef(oid) as def 140 FROM pg_rewrite 141 WHERE ev_class = %(oid)s AND rulename <> '_RETURN'::name 142 """

143 - def __init__(self, table_name, row, new_name = None):

144 self.table_name = table_name 145 self.name = row['name'] 146 self.defn = row['def']

147

148 - def get_create_sql(self, curs, new_table_name = None):

149 if not new_table_name: 150 return self.defn 151 rx = r"\bTO[ ][a-z0-9._]+[ ]DO[ ]" 152 pnew = "TO %s DO " % new_table_name 153 return rx_replace(rx, self.defn, pnew)

154

155 - def get_drop_sql(self, curs):

156 return 'DROP RULE %s ON %s' % (self.name, self.table_name)

157

158 -class TTrigger(TElem):

159 """Info about trigger.""" 160 type = T_TRIGGER 161 SQL = """ 162 SELECT tgname as name, pg_get_triggerdef(oid) as def 163 FROM pg_trigger 164 WHERE tgrelid = %(oid)s AND NOT tgisconstraint 165 """

166 - def __init__(self, table_name, row):

167 self.table_name = table_name 168 self.name = row['name'] 169 self.defn = row['def'] + ';'

170

171 - def get_create_sql(self, curs, new_table_name = None):

172 if not new_table_name: 173 return self.defn 174 175 rx = r"\bON[ ][a-z0-9._]+[ ]" 176 pnew = "ON %s " % new_table_name 177 return rx_replace(rx, self.defn, pnew)

178

179 - def get_drop_sql(self, curs):

180 return 'DROP TRIGGER %s ON %s' % (self.name, self.table_name)

181

182 -class TOwner(TElem):

183 """Info about table owner.""" 184 type = T_OWNER 185 SQL = """ 186 SELECT pg_get_userbyid(relowner) as owner FROM pg_class 187 WHERE oid = %(oid)s 188 """

189 - def __init__(self, table_name, row, new_name = None):

190 self.table_name = table_name 191 self.name = 'Owner' 192 self.owner = row['owner']

193

194 - def get_create_sql(self, curs, new_name = None):

195 if not new_name: 196 new_name = self.table_name 197 return 'ALTER TABLE %s OWNER TO %s;' % (new_name, self.owner)

198

199 -class TGrant(TElem):

200 """Info about permissions.""" 201 type = T_GRANT 202 SQL = "SELECT relacl FROM pg_class where oid = %(oid)s" 203 acl_map = { 204 'r': 'SELECT', 'w': 'UPDATE', 'a': 'INSERT', 'd': 'DELETE', 205 'R': 'RULE', 'x': 'REFERENCES', 't': 'TRIGGER', 'X': 'EXECUTE', 206 'U': 'USAGE', 'C': 'CREATE', 'T': 'TEMPORARY' 207 }

208 - def acl_to_grants(self, acl):

209 if acl == "arwdRxt": # ALL for tables 210 return "ALL" 211 return ", ".join([ self.acl_map[c] for c in acl ])

212

213 - def parse_relacl(self, relacl):

214 if relacl is None: 215 return [] 216 if len(relacl) > 0 and relacl[0] == '{' and relacl[-1] == '}': 217 relacl = relacl[1:-1] 218 list = [] 219 for f in relacl.split(','): 220 user, tmp = f.strip('"').split('=') 221 acl, who = tmp.split('/') 222 list.append((user, acl, who)) 223 return list

224

225 - def __init__(self, table_name, row, new_name = None):

226 self.name = table_name 227 self.acl_list = self.parse_relacl(row['relacl'])

228

229 - def get_create_sql(self, curs, new_name = None):

230 if not new_name: 231 new_name = self.name 232 233 list = [] 234 for user, acl, who in self.acl_list: 235 astr = self.acl_to_grants(acl) 236 sql = "GRANT %s ON %s TO %s;" % (astr, new_name, user) 237 list.append(sql) 238 return "\n".join(list)

239

240 - def get_drop_sql(self, curs):

241 list = [] 242 for user, acl, who in self.acl_list: 243 sql = "REVOKE ALL FROM %s ON %s;" % (user, self.name) 244 list.append(sql) 245 return "\n".join(list)

246

247 -class TColumn(TElem):

248 """Info about table column.""" 249 SQL = """ 250 select a.attname as name, 251 a.attname || ' ' 252 || format_type(a.atttypid, a.atttypmod) 253 || case when a.attnotnull then ' not null' else '' end 254 || case when a.atthasdef then ' ' || d.adsrc else '' end 255 as def 256 from pg_attribute a left join pg_attrdef d 257 on (d.adrelid = a.attrelid and d.adnum = a.attnum) 258 where a.attrelid = %(oid)s 259 and not a.attisdropped 260 and a.attnum > 0 261 order by a.attnum; 262 """

263 - def __init__(self, table_name, row):

264 self.name = row['name'] 265 self.column_def = row['def']

266

267 -class TTable(TElem):

268 """Info about table only (columns).""" 269 type = T_TABLE

270 - def __init__(self, table_name, col_list):

271 self.name = table_name 272 self.col_list = col_list

273

274 - def get_create_sql(self, curs, new_name = None):

275 if not new_name: 276 new_name = self.name 277 sql = "create table %s (" % new_name 278 sep = "\n\t" 279 for c in self.col_list: 280 sql += sep + c.column_def 281 sep = ",\n\t" 282 sql += "\n);" 283 return sql

284

285 - def get_drop_sql(self, curs):

286 return "DROP TABLE %s;" % self.name

287 288 # 289 # Main table object, loads all the others 290 # 291

292 -class TableStruct(object):

293 """Collects and manages all info about table. 294 295 Allow to issue CREATE/DROP statements about any 296 group of elements. 297 """

298 - def __init__(self, curs, table_name):

299 """Initializes class by loading info about table_name from database.""" 300 301 self.table_name = table_name 302 303 # fill args 304 schema, name = fq_name_parts(table_name) 305 args = { 306 'schema': schema, 307 'table': name, 308 'oid': get_table_oid(curs, table_name), 309 'pg_class_oid': get_table_oid(curs, 'pg_catalog.pg_class'), 310 } 311 312 # load table struct 313 self.col_list = self._load_elem(curs, args, TColumn) 314 self.object_list = [ TTable(table_name, self.col_list) ] 315 316 # load additional objects 317 to_load = [TConstraint, TIndex, TTrigger, TRule, TGrant, TOwner] 318 for eclass in to_load: 319 self.object_list += self._load_elem(curs, args, eclass)

320

321 - def _load_elem(self, curs, args, eclass):

322 list = [] 323 curs.execute(eclass.SQL % args) 324 for row in curs.dictfetchall(): 325 list.append(eclass(self.table_name, row)) 326 return list

327

328 - def create(self, curs, objs, new_table_name = None, log = None):

329 """Issues CREATE statements for requested set of objects. 330 331 If new_table_name is giver, creates table under that name 332 and also tries to rename all indexes/constraints that conflict 333 with existing table. 334 """ 335 336 for o in self.object_list: 337 if o.type & objs: 338 sql = o.get_create_sql(curs, new_table_name) 339 if not sql: 340 continue 341 if log: 342 log.info('Creating %s' % o.name) 343 log.debug(sql) 344 curs.execute(sql)

345

346 - def drop(self, curs, objs, log = None):

347 """Issues DROP statements for requested set of objects.""" 348 for o in self.object_list: 349 if o.type & objs: 350 sql = o.get_drop_sql(curs) 351 if not sql: 352 continue 353 if log: 354 log.info('Dropping %s' % o.name) 355 log.debug(sql) 356 curs.execute(sql)

357

358 - def get_column_list(self):

359 """Returns list of column names the table has.""" 360 361 res = [] 362 for c in self.col_list: 363 res.append(c.name) 364 return res

365

366 -def test():

367 import psycopg 368 db = psycopg.connect("dbname=fooz") 369 curs = db.cursor() 370 371 s = TableStruct(curs, "public.data1") 372 373 s.drop(curs, T_ALL) 374 s.create(curs, T_ALL) 375 s.create(curs, T_ALL, "data1_new") 376 s.create(curs, T_PKEY)

377 378 if __name__ == '__main__': 379 test() 380

Source Code for Module skytools.dbstruct