def alias(self, tbl, tbls):
""" Generate a unique table alias
tbl - name of the table to alias, quoted if it needs to be
tbls - TableReference iterable of tables already in query
"""
tbl = self.case(tbl)
tbls = set(normalize_ref(t.ref) for t in tbls)
if self.generate_aliases:
tbl = generate_alias(self.unescape_name(tbl))
if normalize_ref(tbl) not in tbls:
return tbl
elif tbl[0] == '"':
aliases = ('"' + tbl[1:-1] + str(i) + '"' for i in count(2))
else:
aliases = (tbl + str(i) for i in count(2))
return next(a for a in aliases if normalize_ref(a) not in tbls)
def populate_scoped_cols(self, scoped_tbls, local_tbls=()):
""" Find all columns in a set of scoped_tables
:param scoped_tbls: list of TableReference namedtuples
:param local_tbls: tuple(TableMetadata)
:return: {TableReference:{colname:ColumnMetaData}}
"""
ctes = dict((normalize_ref(t.name), t.columns) for t in local_tbls)
columns = OrderedDict()
meta = self.metadata
def addcols(schema, rel, alias, reltype, cols):
tbl = TableReference(schema, rel, alias, reltype == 'functions')
if tbl not in columns:
columns[tbl] = []
columns[tbl].extend(cols)
for tbl in scoped_tbls:
# Local tables should shadow database tables
if tbl.schema is None and normalize_ref(tbl.name) in ctes:
cols = ctes[normalize_ref(tbl.name)]
addcols(None, tbl.name, 'CTE', tbl.alias, cols)
continue
schemas = [tbl.schema] if tbl.schema else self.search_path
for schema in schemas:
relname = self.escape_name(tbl.name)
schema = self.escape_name(schema)
if tbl.is_function:
# Return column names from a set-returning function
# Get an array of FunctionMetadata objects
functions = meta['functions'].get(schema, {}).get(relname)
for func in (functions or []):
# func is a FunctionMetadata object
cols = func.fields()
addcols(schema, relname, tbl.alias, 'functions', cols)
else:
for reltype in ('tables', 'views'):
cols = meta[reltype].get(schema, {}).get(relname)
if cols:
cols = cols.values()
addcols(schema, relname, tbl.alias, reltype, cols)
break
return columns
def get_join_matches(self, suggestion, word_before_cursor):
tbls = suggestion.table_refs
cols = self.populate_scoped_cols(tbls)
# Set up some data structures for efficient access
qualified = dict((normalize_ref(t.ref), t.schema) for t in tbls)
ref_prio = dict((normalize_ref(t.ref), n) for n, t in enumerate(tbls))
refs = set(normalize_ref(t.ref) for t in tbls)
other_tbls = set((t.schema, t.name) for t in list(cols)[:-1])
joins = []
# Iterate over FKs in existing tables to find potential joins
fks = ((fk, rtbl, rcol) for rtbl, rcols in cols.items()
for rcol in rcols for fk in rcol.foreignkeys)
col = namedtuple('col', 'schema tbl col')
for fk, rtbl, rcol in fks:
right = col(rtbl.schema, rtbl.name, rcol.name)
child = col(fk.childschema, fk.childtable, fk.childcolumn)
parent = col(fk.parentschema, fk.parenttable, fk.parentcolumn)
left = child if parent == right else parent
if suggestion.schema and left.schema != suggestion.schema:
continue
c = self.case
if self.generate_aliases or normalize_ref(left.tbl) in refs:
lref = self.alias(left.tbl, suggestion.table_refs)
join = '{0} {4} ON {4}.{1} = {2}.{3}'.format(
c(left.tbl), c(left.col), rtbl.ref, c(right.col), lref)
else:
join = '{0} ON {0}.{1} = {2}.{3}'.format(
c(left.tbl), c(left.col), rtbl.ref, c(right.col))
alias = generate_alias(self.case(left.tbl))
synonyms = [
join,
'{0} ON {0}.{1} = {2}.{3}'.format(alias, c(left.col), rtbl.ref,
c(right.col))
]
# Schema-qualify if (1) new table in same schema as old, and old
# is schema-qualified, or (2) new in other schema, except public
if not suggestion.schema and (qualified[normalize_ref(
rtbl.ref)] and left.schema == right.schema or left.schema
not in (right.schema, 'default')):
join = left.schema + '.' + join
prio = ref_prio[normalize_ref(rtbl.ref)] * 2 + (0 if (
left.schema, left.tbl) in other_tbls else 1)
joins.append(Candidate(join, prio, 'join', synonyms=synonyms))
return self.find_matches(word_before_cursor, joins, meta='join')
