def mk_update_sql(row, tbl, pkey_list, field_map = None):
r"""Generate UPDATE statement from dict data.
>>> mk_update_sql({'id': 0, 'id2': '2', 'data': 'str\\'}, 'Table', ['id', 'id2'])
'update only public."Table" set data = E\'str\\\\\' where id = \'0\' and id2 = \'2\';'
"""
if len(pkey_list) < 1:
raise Exception("update needs pkeys")
set_list = []
whe_list = []
pkmap = {}
for k in pkey_list:
pkmap[k] = 1
new_k = field_map and field_map[k] or k
col = quote_ident(new_k)
val = quote_literal(row[k])
whe_list.append("%s = %s" % (col, val))
if field_map:
for src, dst in field_map.iteritems():
if src not in pkmap:
col = quote_ident(dst)
val = quote_literal(row[src])
set_list.append("%s = %s" % (col, val))
else:
for col, val in row.iteritems():
if col not in pkmap:
col = quote_ident(col)
val = quote_literal(val)
set_list.append("%s = %s" % (col, val))
return "update only %s set %s where %s;" % (quote_fqident(tbl),
", ".join(set_list), " and ".join(whe_list))
def mk_update_sql(row, tbl, pkey_list, field_map=None):
r"""Generate UPDATE statement from dict data.
>>> mk_update_sql({'id': 0, 'id2': '2', 'data': 'str\\'}, 'Table', ['id', 'id2'])
'update only public."Table" set data = E\'str\\\\\' where id = \'0\' and id2 = \'2\';'
"""
if len(pkey_list) < 1:
raise Exception("update needs pkeys")
set_list = []
whe_list = []
pkmap = {}
for k in pkey_list:
pkmap[k] = 1
new_k = field_map and field_map[k] or k
col = quote_ident(new_k)
val = quote_literal(row[k])
whe_list.append("%s = %s" % (col, val))
if field_map:
for src, dst in field_map.iteritems():
if src not in pkmap:
col = quote_ident(dst)
val = quote_literal(row[src])
set_list.append("%s = %s" % (col, val))
else:
for col, val in row.iteritems():
if col not in pkmap:
col = quote_ident(col)
val = quote_literal(val)
set_list.append("%s = %s" % (col, val))
return "update only %s set %s where %s;" % (
quote_fqident(tbl), ", ".join(set_list), " and ".join(whe_list))
def mk_update_sql(row, tbl, pkey_list, field_map = None):
"""Generate UPDATE statement from dict data."""
if len(pkey_list) < 1:
raise Exception("update needs pkeys")
set_list = []
whe_list = []
pkmap = {}
for k in pkey_list:
pkmap[k] = 1
new_k = field_map and field_map[k] or k
col = quote_ident(new_k)
val = quote_literal(row[k])
whe_list.append("%s = %s" % (col, val))
if field_map:
for src, dst in field_map.iteritems():
if src not in pkmap:
col = quote_ident(dst)
val = quote_literal(row[src])
set_list.append("%s = %s" % (col, val))
else:
for col, val in row.iteritems():
if col not in pkmap:
col = quote_ident(col)
val = quote_literal(val)
set_list.append("%s = %s" % (col, val))
return "update %s set %s where %s;" % (quote_fqident(tbl),
", ".join(set_list), " and ".join(whe_list))
def get_create_sql(self, curs, new_table_name=None):
fmt = "ALTER TABLE ONLY %s ADD CONSTRAINT %s %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)
return sql
def get_create_sql(self, curs, new_table_name=None):
fmt = "ALTER TABLE ONLY %s ADD CONSTRAINT %s %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)
return sql
def magic_insert(curs, tablename, data, fields = None, use_insert = 0, quoted_table = False):
r"""Copy/insert a list of dict/list data to database.
If curs == None, then the copy or insert statements are returned
as string. For list of dict the field list is optional, as its
possible to guess them from dict keys.
Example:
>>> magic_insert(None, 'tbl', [[1, '1'], [2, '2']], ['col1', 'col2'])
'COPY public.tbl (col1,col2) FROM STDIN;\n1\t1\n2\t2\n\\.\n'
"""
if len(data) == 0:
return
# decide how to process
if hasattr(data[0], 'keys'):
if fields == None:
fields = data[0].keys()
if use_insert:
row_func = _gen_dict_insert
else:
row_func = _gen_dict_copy
else:
if fields == None:
raise Exception("Non-dict data needs field list")
if use_insert:
row_func = _gen_list_insert
else:
row_func = _gen_list_copy
qfields = [quote_ident(f) for f in fields]
if quoted_table:
qtablename = tablename
else:
qtablename = quote_fqident(tablename)
# init processing
buf = StringIO()
if curs == None and use_insert == 0:
fmt = "COPY %s (%s) FROM STDIN;\n"
buf.write(fmt % (qtablename, ",".join(qfields)))
# process data
for row in data:
buf.write(row_func(qtablename, row, fields, qfields))
buf.write("\n")
# if user needs only string, return it
if curs == None:
if use_insert == 0:
buf.write("\\.\n")
return buf.getvalue()
# do the actual copy/inserts
if use_insert:
curs.execute(buf.getvalue())
else:
buf.seek(0)
hdr = "%s (%s)" % (qtablename, ",".join(qfields))
curs.copy_from(buf, hdr)
def get_drop_sql(self, curs):
list = []
for user, acl, who in self.acl_list:
sql = "REVOKE ALL FROM %s ON %s;" % (quote_ident(user),
quote_fqident(self.name))
list.append(sql)
return "\n".join(list)
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 = fq_name_parts(table_name)
args = {
'schema': schema,
'table': name,
'fqname': self.fqname,
'fq2name': quote_literal(self.fqname),
'oid': get_table_oid(curs, table_name),
'pg_class_oid': get_table_oid(curs, 'pg_catalog.pg_class'),
}
# load table struct
self.col_list = self._load_elem(curs, self.name, args, TColumn)
self.object_list = [ TTable(table_name, self.col_list) ]
self.seq_list = []
# load seqs
for col in self.col_list:
if col.seqname:
owner = self.fqname + '.' + quote_ident(col.name)
seq_args = { 'fqname': col.seqname, 'owner': 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 = [TConstraint, TIndex, TTrigger, TRule, TGrant, TOwner, TParent]
for eclass in to_load:
self.object_list += self._load_elem(curs, self.name, args, eclass)
def full_copy(tablename, src_curs, dst_curs, column_list=[], condition=None):
"""COPY table from one db to another."""
qtable = quote_fqident(tablename)
if column_list:
qfields = ",".join([quote_ident(f) for f in column_list])
src = dst = "%s (%s)" % (qtable, qfields)
else:
qfields = '*'
src = dst = qtable
if condition:
src = "(SELECT %s FROM %s WHERE %s)" % (qfields, qtable, condition)
if hasattr(src_curs, 'copy_expert'):
sql_to = "COPY %s TO stdout" % src
sql_from = "COPY %s FROM stdin" % dst
buf = CopyPipe(dst_curs, sql_from=sql_from)
src_curs.copy_expert(sql_to, buf)
else:
if condition:
# regular psycopg copy_to generates invalid sql for subselect copy
raise Exception('copy_expert() is needed for conditional copy')
buf = CopyPipe(dst_curs, dst)
src_curs.copy_to(buf, src)
buf.flush()
return (buf.total_bytes, buf.total_rows)
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 = fq_name_parts(table_name)
args = {
'schema': schema,
'table': name,
'fqname': self.fqname,
'fq2name': quote_literal(self.fqname),
'oid': get_table_oid(curs, table_name),
'pg_class_oid': get_table_oid(curs, 'pg_catalog.pg_class'),
}
# load table struct
self.col_list = self._load_elem(curs, self.name, args, TColumn)
self.object_list = [ TTable(table_name, self.col_list) ]
self.seq_list = []
# load seqs
for col in self.col_list:
if col.seqname:
owner = self.fqname + '.' + quote_ident(col.name)
seq_args = { 'fqname': col.seqname, 'owner': 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 = [TConstraint, TIndex, TTrigger, TRule, TGrant, TOwner, TParent]
for eclass in to_load:
self.object_list += self._load_elem(curs, self.name, args, eclass)
def full_copy(tablename, src_curs, dst_curs, column_list = [], condition = None):
"""COPY table from one db to another."""
qtable = quote_fqident(tablename)
if column_list:
qfields = ",".join([quote_ident(f) for f in column_list])
src = dst = "%s (%s)" % (qtable, qfields)
else:
qfields = '*'
src = dst = qtable
if condition:
src = "(SELECT %s FROM %s WHERE %s)" % (qfields, qtable, condition)
if hasattr(src_curs, 'copy_expert'):
sql_to = "COPY %s TO stdout" % src
sql_from = "COPY %s FROM stdin" % dst
buf = CopyPipe(dst_curs, sql_from = sql_from)
src_curs.copy_expert(sql_to, buf)
else:
if condition:
# regular psycopg copy_to generates invalid sql for subselect copy
raise Exception('copy_expert() is needed for conditional copy')
buf = CopyPipe(dst_curs, dst)
src_curs.copy_to(buf, src)
buf.flush()
return (buf.total_bytes, buf.total_rows)
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 %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)
return sql
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 CLUSTER ON %s;' % (
quote_fqident(tname), quote_ident(iname))
return sql
def mk_insert_sql(row, tbl, pkey_list = None, field_map = None):
"""Generate INSERT statement from dict data."""
col_list = []
val_list = []
if field_map:
for src, dst in field_map.iteritems():
col_list.append(quote_ident(dst))
val_list.append(quote_literal(row[src]))
else:
for c, v in row.iteritems():
col_list.append(quote_ident(c))
val_list.append(quote_literal(v))
col_str = ", ".join(col_list)
val_str = ", ".join(val_list)
return "insert into %s (%s) values (%s);" % (
quote_fqident(tbl), col_str, val_str)
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_create_sql(self, curs, new_table_name = None):
if not new_table_name:
return self.defn
# fixme: seems broken
name = find_new_name(curs, self.name)
pnew = "INDEX %s ON %s " % (quote_ident(name), quote_fqident(new_table_name))
rx = r"\bINDEX[ ][a-z0-9._]+[ ]ON[ ][a-z0-9._]+[ ]"
sql = rx_replace(rx, self.defn, pnew)
return sql
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 %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 CLUSTER ON %s;' % (qtbl, qname)
return sql
def magic_insert(curs, tablename, data, fields = None, use_insert = 0):
r"""Copy/insert a list of dict/list data to database.
If curs == None, then the copy or insert statements are returned
as string. For list of dict the field list is optional, as its
possible to guess them from dict keys.
Example:
>>> magic_insert(None, 'tbl', [[1, '1'], [2, '2']], ['col1', 'col2'])
'COPY public.tbl (col1,col2) FROM STDIN;\n1\t1\n2\t2\n\\.\n'
"""
if len(data) == 0:
return
# decide how to process
if hasattr(data[0], 'keys'):
if fields == None:
fields = data[0].keys()
if use_insert:
row_func = _gen_dict_insert
else:
row_func = _gen_dict_copy
else:
if fields == None:
raise Exception("Non-dict data needs field list")
if use_insert:
row_func = _gen_list_insert
else:
row_func = _gen_list_copy
qfields = [quote_ident(f) for f in fields]
qtablename = quote_fqident(tablename)
# init processing
buf = StringIO()
if curs == None and use_insert == 0:
fmt = "COPY %s (%s) FROM STDIN;\n"
buf.write(fmt % (qtablename, ",".join(qfields)))
# process data
for row in data:
buf.write(row_func(qtablename, row, fields, qfields))
buf.write("\n")
# if user needs only string, return it
if curs == None:
if use_insert == 0:
buf.write("\\.\n")
return buf.getvalue()
# do the actual copy/inserts
if use_insert:
curs.execute(buf.getvalue())
else:
buf.seek(0)
hdr = "%s (%s)" % (qtablename, ",".join(qfields))
curs.copy_from(buf, hdr)
def get_create_sql(self, curs, new_name = None):
if not new_name:
new_name = self.name
list = []
for user, acl, who in self.acl_list:
astr = self.acl_to_grants(acl)
sql = "GRANT %s ON %s TO %s;" % (astr, quote_fqident(new_name), quote_ident(user))
list.append(sql)
return "\n".join(list)
def get_create_sql(self, curs, new_table_name=None):
if not new_table_name:
return self.defn
# fixme: seems broken
name = find_new_name(curs, self.name)
pnew = "INDEX %s ON %s " % (quote_ident(name),
quote_fqident(new_table_name))
rx = r"\bINDEX[ ][a-z0-9._]+[ ]ON[ ][a-z0-9._]+[ ]"
sql = rx_replace(rx, self.defn, pnew)
return sql
def get_create_sql(self, curs, new_name=None):
if not new_name:
new_name = self.name
list = []
for user, acl, who in self.acl_list:
astr = self.acl_to_grants(acl)
sql = "GRANT %s ON %s TO %s;" % (astr, quote_fqident(new_name),
quote_ident(user))
list.append(sql)
return "\n".join(list)
def mk_insert_sql(row, tbl, pkey_list=None, field_map=None):
"""Generate INSERT statement from dict data.
>>> mk_insert_sql({'id': '1', 'data': None}, 'tbl')
"insert into public.tbl (data, id) values (null, '1');"
"""
col_list = []
val_list = []
if field_map:
for src, dst in field_map.iteritems():
col_list.append(quote_ident(dst))
val_list.append(quote_literal(row[src]))
else:
for c, v in row.iteritems():
col_list.append(quote_ident(c))
val_list.append(quote_literal(v))
col_str = ", ".join(col_list)
val_str = ", ".join(val_list)
return "insert into %s (%s) values (%s);" % (quote_fqident(tbl), col_str,
val_str)
def mk_delete_sql(row, tbl, pkey_list, field_map = None):
"""Generate DELETE statement from dict data."""
if len(pkey_list) < 1:
raise Exception("delete needs pkeys")
whe_list = []
for k in pkey_list:
new_k = field_map and field_map[k] or k
col = quote_ident(new_k)
val = quote_literal(row[k])
whe_list.append("%s = %s" % (col, val))
whe_str = " and ".join(whe_list)
return "delete from %s where %s;" % (quote_fqident(tbl), whe_str)
def mk_delete_sql(row, tbl, pkey_list, field_map=None):
"""Generate DELETE statement from dict data."""
if len(pkey_list) < 1:
raise Exception("delete needs pkeys")
whe_list = []
for k in pkey_list:
new_k = field_map and field_map[k] or k
col = quote_ident(new_k)
val = quote_literal(row[k])
whe_list.append("%s = %s" % (col, val))
whe_str = " and ".join(whe_list)
return "delete from only %s where %s;" % (quote_fqident(tbl), whe_str)
def full_copy(tablename, src_curs, dst_curs, column_list=[]):
"""COPY table from one db to another."""
qtable = quote_fqident(tablename)
if column_list:
qfields = [quote_ident(f) for f in column_list]
hdr = "%s (%s)" % (qtable, ",".join(qfields))
else:
hdr = qtable
if hasattr(src_curs, "copy_expert"):
sql_to = "COPY %s TO stdout" % hdr
sql_from = "COPY %s FROM stdout" % hdr
buf = CopyPipe(dst_curs, sql_from=sql_from)
src_curs.copy_expert(sql_to, buf)
else:
buf = CopyPipe(dst_curs, hdr)
src_curs.copy_to(buf, hdr)
buf.flush()
return (buf.total_bytes, buf.total_rows)
def full_copy(tablename, src_curs, dst_curs, column_list = []):
"""COPY table from one db to another."""
qtable = quote_fqident(tablename)
if column_list:
qfields = [quote_ident(f) for f in column_list]
hdr = "%s (%s)" % (qtable, ",".join(qfields))
else:
hdr = qtable
if hasattr(src_curs, 'copy_expert'):
sql_to = "COPY %s TO stdout" % hdr
sql_from = "COPY %s FROM stdout" % hdr
buf = CopyPipe(dst_curs, sql_from = sql_from)
src_curs.copy_expert(sql_to, buf)
else:
buf = CopyPipe(dst_curs, hdr)
src_curs.copy_to(buf, hdr)
buf.flush()
return (buf.total_bytes, buf.total_rows)
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))
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):
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)
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 OWNER TO %s;' % (quote_fqident(new_name), quote_ident(self.owner))
def get_drop_sql(self, curs):
return 'DROP TRIGGER %s ON %s' % (quote_ident(self.name), quote_fqident(self.table_name))
def get_drop_sql(self, curs):
"""Generate removal sql."""
fmt = "ALTER TABLE ONLY %s DROP CONSTRAINT %s;"
sql = fmt % (quote_fqident(self.table_name), quote_ident(self.name))
return sql
def get_drop_sql(self, curs):
fmt = "ALTER TABLE ONLY %s DROP CONSTRAINT %s;"
sql = fmt % (quote_fqident(self.table_name), quote_ident(self.name))
return sql
def get_create_sql(self, curs, new_name=None):
if not new_name:
new_name = self.table_name
return 'ALTER TABLE %s OWNER TO %s;' % (quote_fqident(new_name),
quote_ident(self.owner))
def get_drop_sql(self, curs):
return 'DROP TRIGGER %s ON %s' % (quote_ident(
self.name), quote_fqident(self.table_name))
