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