codekingpro/portable-devtools
115k
1##########################################################################2#3# pgAdmin 4 - PostgreSQL Tools4#5# Copyright (C) 2013 - 2024, The pgAdmin Development Team6# This software is released under the PostgreSQL Licence7#8##########################################################################9 10"""A blueprint module implementing the sql auto complete feature."""11 12import re13import operator14from itertools import count15from .completion import Completion16from collections import namedtuple, defaultdict, OrderedDict17 18from .sqlcompletion import (19 FromClauseItem,20 suggest_type,21 Database,22 Schema,23 Table,24 Function,25 Column,26 View,27 Keyword,28 Datatype,29 Alias,30 JoinCondition,31 Join32)33from .parseutils.meta import FunctionMetadata, ColumnMetadata, ForeignKey34from .parseutils.utils import last_word35from .parseutils.tables import TableReference36from .prioritization import PrevalenceCounter37from flask import render_template38from pgadmin.utils.driver import get_driver39from config import PG_DEFAULT_DRIVER40from pgadmin.utils.preferences import Preferences41 42Match = namedtuple("Match", ["completion", "priority"])43 44_SchemaObject = namedtuple("SchemaObject", "name schema meta")45 46 47def SchemaObject(name, schema=None, meta=None):48 return _SchemaObject(name, schema, meta)49 50 51# Regex for finding "words" in documents.52_FIND_WORD_RE = re.compile(r'([\w]+|[^a-zA-Z0-9_\s]+)')53_FIND_BIG_WORD_RE = re.compile(r'([^\s]+)')54 55_Candidate = namedtuple("Candidate",56 "completion prio meta synonyms prio2 display")57 58 59def Candidate(60 completion, prio=None, meta=None, synonyms=None, prio2=None, display=None61):62 return _Candidate(63 completion, prio, meta, synonyms or [completion], prio2,64 display or completion65 )66 67 68# Used to strip trailing '::some_type' from default-value expressions69arg_default_type_strip_regex = re.compile(r"::[\w\.]+(\[\])?$")70 71 72def normalize_ref(ref):73 return ref if ref[0] == '"' else '"' + ref.lower() + '"'74 75 76def generate_alias(tbl):77 """Generate a table alias, consisting of all upper-case letters in78 the table name, or, if there are no upper-case letters, the first letter +79 all letters preceded by _80 param tbl - unescaped name of the table to alias81 """82 return "".join(83 [letter for letter in tbl if letter.isupper()] or84 [letter for letter, prev in zip(tbl, "_" + tbl)85 if prev == "_" and letter != "_"]86 )87 88 89class SQLAutoComplete():90 """91 class SQLAutoComplete92 93 This class is used to provide the postgresql's autocomplete feature.94 This class used sqlparse to parse the given sql and psycopg to make95 the connection and get the tables, schemas, functions etc. based on96 the query.97 """98 def __init__(self, **kwargs):99 """100 This method is used to initialize the class.101 102 Args:103 **kwargs : N number of parameters104 """105 106 self.sid = kwargs['sid'] if 'sid' in kwargs else None107 self.conn = kwargs['conn'] if 'conn' in kwargs else None108 self.keywords = []109 self.name_pattern = re.compile(r"^[_a-z][_a-z0-9\$]*$")110 111 self.databases = []112 self.functions = []113 self.datatypes = []114 self.dbmetadata = \115 {"tables": {}, "views": {}, "functions": {}, "datatypes": {}}116 self.text_before_cursor = None117 118 manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(self.sid)119 120 # we will set template path for sql scripts121 self.sql_path = 'sqlautocomplete/sql/#{0}#'.format(manager.version)122 123 self.search_path = []124 schema_names = []125 if self.conn.connected():126 # Fetch the search path127 self._set_search_path()128 129 # Fetch the schema names130 self._fetch_schema_name(schema_names)131 132 pref = Preferences.module('sqleditor')133 keywords_in_uppercase = \134 pref.preference('keywords_in_uppercase').get()135 136 # Fetch the keywords137 query = render_template("/".join([self.sql_path, 'keywords.sql']))138 # If setting 'Keywords in uppercase' is set to True in139 # Preferences then fetch the keywords in upper case.140 if keywords_in_uppercase:141 query = render_template(142 "/".join([self.sql_path, 'keywords.sql']), upper_case=True)143 status, res = self.conn.execute_dict(query)144 if status:145 for record in res['rows']:146 # 'public' is a keyword in EPAS database server. Don't add147 # this into the list of keywords.148 # This is a hack to fix the issue in autocomplete.149 if record['word'].lower() == 'public':150 continue151 self.keywords.append(record['word'])152 153 self.prioritizer = PrevalenceCounter(self.keywords)154 155 self.reserved_words = set()156 for x in self.keywords:157 self.reserved_words.update(x.split())158 159 self.all_completions = set(self.keywords)160 self.extend_schemata(schema_names)161 162 # Below are the configurable options in pgcli which we don't have163 # in pgAdmin4 at the moment. Setting the default value from the pgcli's164 # config file.165 self.signature_arg_style = '{arg_name} {arg_type}'166 self.call_arg_style = '{arg_name: <{max_arg_len}} := {arg_default}'167 self.call_arg_display_style = '{arg_name}'168 self.call_arg_oneliner_max = 2169 self.search_path_filter = True170 self.generate_aliases = False171 self.insert_col_skip_patterns = [172 re.compile(r'^now\(\)$'),173 re.compile(r'^nextval\(')]174 self.qualify_columns = 'if_more_than_one_table'175 self.asterisk_column_order = 'table_order'176 177 def _set_search_path(self):178 query = render_template(179 "/".join([self.sql_path, 'schema.sql']), search_path=True)180 status, res = self.conn.execute_dict(query)181 if status:182 for record in res['rows']:183 self.search_path.append(record['schema'])184 185 def _fetch_schema_name(self, schema_names):186 query = render_template("/".join([self.sql_path, 'schema.sql']))187 status, res = self.conn.execute_dict(query)188 if status:189 for record in res['rows']:190 schema_names.append(record['schema'])191 192 def escape_name(self, name):193 if name and (194 (not self.name_pattern.match(name)) or195 (name.upper() in self.reserved_words) or196 (name.upper() in self.functions)197 ):198 name = '"%s"' % name199 200 return name201 202 def escape_schema(self, name):203 return "'{}'".format(self.unescape_name(name))204 205 def unescape_name(self, name):206 """ Unquote a string."""207 if name and name[0] == '"' and name[-1] == '"':208 name = name[1:-1]209 210 return name211 212 def escaped_names(self, names):213 return [self.escape_name(name) for name in names]214 215 def extend_database_names(self, databases):216 self.databases.extend(databases)217 218 def extend_keywords(self, additional_keywords):219 self.keywords.extend(additional_keywords)220 self.all_completions.update(additional_keywords)221 222 def extend_schemata(self, schemata):223 224 # schemata is a list of schema names225 schemata = self.escaped_names(schemata)226 metadata = self.dbmetadata["tables"]227 for schema in schemata:228 metadata[schema] = {}229 230 # dbmetadata.values() are the 'tables' and 'functions' dicts231 for metadata in self.dbmetadata.values():232 for schema in schemata:233 metadata[schema] = {}234 235 self.all_completions.update(schemata)236 237 def extend_casing(self, words):238 """extend casing data239 240 :return:241 """242 # casing should be a dict {lowercasename:PreferredCasingName}243 self.casing = {word.lower(): word for word in words}244 245 def extend_relations(self, data, kind):246 """extend metadata for tables or views.247 248 :param data: list of (schema_name, rel_name) tuples249 :param kind: either 'tables' or 'views'250 251 :return:252 253 """254 255 data = [self.escaped_names(d) for d in data]256 257 # dbmetadata['tables']['schema_name']['table_name'] should be an258 # OrderedDict {column_name:ColumnMetaData}.259 metadata = self.dbmetadata[kind]260 for schema, relname in data:261 try:262 metadata[schema][relname] = OrderedDict()263 except KeyError:264 print('%r %r listed in unrecognized schema %r',265 kind, relname, schema)266 267 self.all_completions.add(relname)268 269 def extend_columns(self, column_data, kind):270 """extend column metadata.271 272 :param column_data: list of (schema_name, rel_name, column_name,273 column_type, has_default, default) tuples274 :param kind: either 'tables' or 'views'275 276 :return:277 278 """279 metadata = self.dbmetadata[kind]280 for schema, relname, colname, datatype, \281 has_default, default in column_data:282 (schema, relname, colname) = self.escaped_names(283 [schema, relname, colname])284 column = ColumnMetadata(285 name=colname,286 datatype=datatype,287 has_default=has_default,288 default=default,289 )290 metadata[schema][relname][colname] = column291 self.all_completions.add(colname)292 293 def extend_functions(self, func_data):294 295 # func_data is a list of function metadata namedtuples296 297 # dbmetadata['schema_name']['functions']['function_name'] should return298 # the function metadata namedtuple for the corresponding function299 metadata = self.dbmetadata["functions"]300 301 for f in func_data:302 schema, func = self.escaped_names([f.schema_name, f.func_name])303 304 if func in metadata[schema]:305 metadata[schema][func].append(f)306 else:307 metadata[schema][func] = [f]308 309 self.all_completions.add(func)310 311 self._refresh_arg_list_cache()312 313 def _refresh_arg_list_cache(self):314 # We keep a cache of315 # {function_usage:{function_metadata: function_arg_list_string}}316 # This is used when suggesting functions, to avoid the latency that317 # would result if we'd recalculate the arg lists each time we suggest318 # functions (in large DBs)319 320 self._arg_list_cache = {321 usage: {322 meta: self._arg_list(meta, usage)323 for sch, funcs in self.dbmetadata["functions"].items()324 for func, metas in funcs.items()325 for meta in metas326 }327 for usage in ("call", "call_display", "signature")328 }329 330 def extend_foreignkeys(self, fk_data):331 332 # fk_data is a list of ForeignKey namedtuples, with fields333 # parentschema, childschema, parenttable, childtable,334 # parentcolumns, childcolumns335 336 # These are added as a list of ForeignKey namedtuples to the337 # ColumnMetadata namedtuple for both the child and parent338 meta = self.dbmetadata["tables"]339 340 for fk in fk_data:341 e = self.escaped_names342 parentschema, childschema = e([fk.parentschema, fk.childschema])343 parenttable, childtable = e([fk.parenttable, fk.childtable])344 childcol, parcol = e([fk.childcolumn, fk.parentcolumn])345 346 if childtable not in meta[childschema] or \347 parenttable not in meta[parentschema] or \348 childcol not in meta[childschema][childtable] or \349 parcol not in meta[parentschema][parenttable]:350 continue351 352 childcolmeta = meta[childschema][childtable][childcol]353 parcolmeta = meta[parentschema][parenttable][parcol]354 fk = ForeignKey(355 parentschema, parenttable, parcol,356 childschema, childtable, childcol357 )358 childcolmeta.foreignkeys.append(fk)359 parcolmeta.foreignkeys.append(fk)360 361 def extend_datatypes(self, type_data):362 363 # dbmetadata['datatypes'][schema_name][type_name] should store type364 # metadata, such as composite type field names. Currently, we're not365 # storing any metadata beyond typename, so just store None366 meta = self.dbmetadata["datatypes"]367 368 for t in type_data:369 schema, type_name = self.escaped_names(t)370 meta[schema][type_name] = None371 self.all_completions.add(type_name)372 373 def set_search_path(self, search_path):374 self.search_path = self.escaped_names(search_path)375 376 def reset_completions(self):377 self.databases = []378 self.special_commands = []379 self.search_path = []380 self.dbmetadata = \381 {"tables": {}, "views": {}, "functions": {}, "datatypes": {}}382 self.all_completions = set(self.keywords + self.functions)383 384 def find_matches(self, text, collection, mode="strict", meta=None):385 """Find completion matches for the given text.386 387 Given the user's input text and a collection of available388 completions, find completions matching the last word of the389 text.390 391 `collection` can be either a list of strings or a list of Candidate392 namedtuples.393 `mode` can be either 'fuzzy', or 'strict'394 'fuzzy': fuzzy matching, ties broken by name prevalance395 `keyword`: start only matching, ties broken by keyword prevalance396 397 yields prompt_toolkit Completion instances for any matches found398 in the collection of available completions.399 400 """401 if not collection:402 return []403 prio_order = [404 "keyword",405 "function",406 "view",407 "table",408 "datatype",409 "database",410 "schema",411 "column",412 "table alias",413 "join",414 "name join",415 "fk join",416 "table format",417 ]418 type_priority = prio_order.index(meta) if meta in prio_order else -1419 text = last_word(text, include="most_punctuations").lower()420 text_len = len(text)421 422 if text and text[0] == '"':423 # text starts with double quote; user is manually escaping a name424 # Match on everything that follows the double-quote. Note that425 # text_len is calculated before removing the quote, so the426 # Completion.position value is correct427 text = text[1:]428 429 if mode == "fuzzy":430 fuzzy = True431 priority_func = self.prioritizer.name_count432 else:433 fuzzy = False434 priority_func = self.prioritizer.keyword_count435 436 # Construct a `_match` function for either fuzzy or non-fuzzy matching437 # The match function returns a 2-tuple used for sorting the matches,438 # or None if the item doesn't match439 # Note: higher priority values mean more important, so use negative440 # signs to flip the direction of the tuple441 if fuzzy:442 regex = ".*?".join(map(re.escape, text))443 pat = re.compile("(%s)" % regex)444 445 def _match(item):446 if item.lower()[: len(text) + 1] in (text, text + " "):447 # Exact match of first word in suggestion448 # This is to get exact alias matches to the top449 # E.g. for input `e`, 'Entries E' should be on top450 # (before e.g. `EndUsers EU`)451 return float("Infinity"), -1452 r = pat.search(self.unescape_name(item.lower()))453 if r:454 return -len(r.group()), -r.start()455 456 else:457 match_end_limit = len(text)458 459 def _match(item):460 # Text starts with double quote; Remove quoting and461 # match on everything that follows the double-quote.462 item = self.unescape_name(item.lower())463 match_point = item.find(text, 0, match_end_limit)464 if match_point >= 0:465 # Use negative infinity to force keywords to sort after all466 # fuzzy matches467 return -float("Infinity"), -match_point468 469 matches = []470 for cand in collection:471 if isinstance(cand, _Candidate):472 item, prio, display_meta, synonyms, prio2, display = cand473 if display_meta is None:474 display_meta = meta475 syn_matches = (_match(x) for x in synonyms)476 # Nones need to be removed to avoid max() crashing in Python 3477 syn_matches = [m for m in syn_matches if m]478 sort_key = max(syn_matches) if syn_matches else None479 else:480 item, display_meta, prio, prio2, display = \481 cand, meta, 0, 0, cand482 sort_key = _match(cand)483 484 if sort_key:485 if display_meta and len(display_meta) > 50:486 # Truncate meta-text to 50 characters, if necessary487 display_meta = display_meta[:47] + "..."488 489 # Lexical order of items in the collection, used for490 # tiebreaking items with the same match group length and start491 # position. Since we use *higher* priority to mean "more492 # important," we use -ord(c) to prioritize "aa" > "ab" and end493 # with 1 to prioritize shorter strings (ie "user" > "users").494 # We first do a case-insensitive sort and then a495 # case-sensitive one as a tie breaker.496 # We also use the unescape_name to make sure quoted names have497 # the same priority as unquoted names.498 lexical_priority = (499 tuple(500 0 if c in " _" else -ord(c)501 for c in self.unescape_name(item.lower())502 ) +503 (1,) +504 tuple(c for c in item)505 )506 507 priority = (508 sort_key,509 type_priority,510 prio,511 priority_func(item),512 prio2,513 lexical_priority,514 )515 matches.append(516 Match(517 completion=Completion(518 text=item,519 start_position=-text_len,520 display_meta=display_meta,521 display=display,522 ),523 priority=priority,524 )525 )526 return matches527 528 def get_completions(self, text, text_before_cursor):529 self.text_before_cursor = text_before_cursor530 531 word_before_cursor = self.get_word_before_cursor(word=True)532 matches = []533 suggestions = suggest_type(text, text_before_cursor)534 535 for suggestion in suggestions:536 suggestion_type = type(suggestion)537 538 # Map suggestion type to method539 # e.g. 'table' -> self.get_table_matches540 matcher = self.suggestion_matchers[suggestion_type]541 matches.extend(matcher(self, suggestion, word_before_cursor))542 543 # Sort matches so highest priorities are first544 matches = \545 sorted(matches, key=operator.attrgetter("priority"), reverse=True)546 547 result = dict()548 for m in matches:549 name = m.completion.display550 result[name] = {'object_type': m.completion.display_meta}551 552 return result553 554 def get_column_matches(self, suggestion, word_before_cursor):555 schema = None556 if len(suggestion.table_refs) > 0 and \557 hasattr(suggestion.table_refs[0], 'schema') and \558 suggestion.table_refs[0].schema != '':559 schema = suggestion.table_refs[0].schema560 561 # Tables and Views should be populated first.562 self.fetch_schema_objects(schema, 'tables')563 self.fetch_schema_objects(schema, 'views')564 565 tables = suggestion.table_refs566 do_qualify = (567 suggestion.qualifiable and568 {569 "always": True,570 "never": False,571 "if_more_than_one_table": len(tables) > 1,572 }[self.qualify_columns]573 )574 575 def qualify(col, tbl):576 return (tbl + '.' + col) if do_qualify else col577 578 scoped_cols = \579 self.populate_scoped_cols(tables, suggestion.local_tables)580 581 def make_cand(name, ref):582 synonyms = (name, generate_alias(name))583 return Candidate(qualify(name, ref), 0, "column", synonyms)584 585 def flat_cols():586 return [587 make_cand(c.name, t.ref)588 for t, cols in scoped_cols.items()589 for c in cols590 ]591 592 if suggestion.require_last_table:593 # require_last_table is used for 'tb11 JOIN tbl2 USING594 # (...' which should595 # suggest only columns that appear in the last table and one more596 ltbl = tables[-1].ref597 other_tbl_cols = {598 c.name for t, cs in scoped_cols.items()599 if t.ref != ltbl for c in cs600 }601 scoped_cols = {602 t: [col for col in cols if col.name in other_tbl_cols]603 for t, cols in scoped_cols.items()604 if t.ref == ltbl605 }606 607 lastword = last_word(word_before_cursor, include="most_punctuations")608 if lastword == "*":609 if suggestion.context == "insert":610 611 def filter(col):612 if not col.has_default:613 return True614 return not any(615 p.match(col.default)616 for p in self.insert_col_skip_patterns617 )618 619 scoped_cols = {620 t: [col for col in cols if filter(col)]621 for t, cols in scoped_cols.items()622 }623 if self.asterisk_column_order == "alphabetic":624 for cols in scoped_cols.values():625 cols.sort(key=operator.attrgetter("name"))626 if (627 lastword != word_before_cursor and628 len(tables) == 1 and629 word_before_cursor[-len(lastword) - 1] == "."630 ):631 # User typed x.*; replicate "x." for all columns except the632 # first, which gets the original (as we only replace the "*"")633 sep = ", " + word_before_cursor[:-1]634 collist = sep.join(c.completion for c in flat_cols())635 else:636 collist = ", ".join(qualify(c.name, t.ref)637 for t, cs in scoped_cols.items()638 for c in cs)639 640 return [641 Match(642 completion=Completion(643 collist, -1, display_meta="columns", display="*"644 ),645 priority=(1, 1, 1),646 )647 ]648 649 return self.find_matches(word_before_cursor, flat_cols(),650 meta="column")651 652 def alias(self, tbl, tbls):653 """Generate a unique table alias654 tbl - name of the table to alias, quoted if it needs to be655 tbls - TableReference iterable of tables already in query656 """657 tbls = set(normalize_ref(t.ref) for t in tbls)658 if self.generate_aliases:659 tbl = generate_alias(self.unescape_name(tbl))660 if normalize_ref(tbl) not in tbls:661 return tbl662 elif tbl[0] == '"':663 aliases = ('"' + tbl[1:-1] + str(i) + '"' for i in count(2))664 else:665 aliases = (tbl + str(i) for i in count(2))666 return next(a for a in aliases if normalize_ref(a) not in tbls)667 668 def get_join_matches(self, suggestion, word_before_cursor):669 tbls = suggestion.table_refs670 cols = self.populate_scoped_cols(tbls)671 # Set up some data structures for efficient access672 qualified = {normalize_ref(t.ref): t.schema for t in tbls}673 ref_prio = {normalize_ref(t.ref): n for n, t in enumerate(tbls)}674 refs = {normalize_ref(t.ref) for t in tbls}675 other_tbls = {(t.schema, t.name) for t in list(cols)[:-1]}676 joins = []677 # Iterate over FKs in existing tables to find potential joins678 fks = (679 (fk, rtbl, rcol)680 for rtbl, rcols in cols.items()681 for rcol in rcols682 for fk in rcol.foreignkeys683 )684 col = namedtuple("col", "schema tbl col")685 for fk, rtbl, rcol in fks:686 right = col(rtbl.schema, rtbl.name, rcol.name)687 child = col(fk.childschema, fk.childtable, fk.childcolumn)688 parent = col(fk.parentschema, fk.parenttable, fk.parentcolumn)689 left = child if parent == right else parent690 if suggestion.schema and left.schema != suggestion.schema:691 continue692 693 if self.generate_aliases or normalize_ref(left.tbl) in refs:694 lref = self.alias(left.tbl, suggestion.table_refs)695 join = "{0} {4} ON {4}.{1} = {2}.{3}".format(696 left.tbl, left.col, rtbl.ref, right.col, lref697 )698 else:699 join = "{0} ON {0}.{1} = {2}.{3}".format(700 left.tbl, left.col, rtbl.ref, right.col701 )702 alias = generate_alias(left.tbl)703 synonyms = [704 join,705 "{0} ON {0}.{1} = {2}.{3}".format(706 alias, left.col, rtbl.ref, right.col707 ),708 ]709 # Schema-qualify if (1) new table in same schema as old, and old710 # is schema-qualified, or (2) new in other schema, except public711 if not suggestion.schema and (712 qualified[normalize_ref(rtbl.ref)] and713 left.schema == right.schema or714 left.schema not in (right.schema, "public")715 ):716 join = left.schema + "." + join717 prio = ref_prio[normalize_ref(rtbl.ref)] * 2 + (718 0 if (left.schema, left.tbl) in other_tbls else 1719 )720 joins.append(Candidate(join, prio, "join", synonyms=synonyms))721 722 return self.find_matches(word_before_cursor, joins, meta="join")723 724 def get_join_condition_matches(self, suggestion, word_before_cursor):725 col = namedtuple("col", "schema tbl col")726 tbls = self.populate_scoped_cols(suggestion.table_refs).items727 cols = [(t, c) for t, cs in tbls() for c in cs]728 try:729 lref = (suggestion.parent or suggestion.table_refs[-1]).ref730 ltbl, lcols = [(t, cs) for (t, cs) in tbls() if t.ref == lref][-1]731 except IndexError: # The user typed an incorrect table qualifier732 return []733 conds, found_conds = [], set()734 735 def add_cond(lcol, rcol, rref, prio, meta):736 prefix = "" if suggestion.parent else ltbl.ref + "."737 cond = prefix + lcol + " = " + rref + "." + rcol738 if cond not in found_conds:739 found_conds.add(cond)740 conds.append(Candidate(cond, prio + ref_prio[rref], meta))741 742 def list_dict(pairs): # Turns [(a, b), (a, c)] into {a: [b, c]}743 d = defaultdict(list)744 for pair in pairs:745 d[pair[0]].append(pair[1])746 return d747 748 # Tables that are closer to the cursor get higher prio749 ref_prio = \750 {tbl.ref: num for num, tbl in enumerate(suggestion.table_refs)}751 # Map (schema, table, col) to tables752 coldict = list_dict(753 ((t.schema, t.name, c.name), t) for t, c in cols if t.ref != lref754 )755 # For each fk from the left table, generate a join condition if756 # the other table is also in the scope757 fks = ((fk, lcol.name) for lcol in lcols for fk in lcol.foreignkeys)758 for fk, lcol in fks:759 left = col(ltbl.schema, ltbl.name, lcol)760 child = col(fk.childschema, fk.childtable, fk.childcolumn)761 par = col(fk.parentschema, fk.parenttable, fk.parentcolumn)762 left, right = (child, par) if left == child else (par, child)763 for rtbl in coldict[right]:764 add_cond(left.col, right.col, rtbl.ref, 2000, "fk join")765 # For name matching, use a {(colname, coltype): TableReference} dict766 coltyp = namedtuple("coltyp", "name datatype")767 col_table = list_dict((coltyp(c.name, c.datatype), t) for t, c in cols)768 # Find all name-match join conditions769 for c in (coltyp(c.name, c.datatype) for c in lcols):770 for rtbl in (t for t in col_table[c] if t.ref != ltbl.ref):771 prio = 1000 if c.datatype in (772 "integer", "bigint", "smallint") else 0773 add_cond(c.name, c.name, rtbl.ref, prio, "name join")774 775 return self.find_matches(word_before_cursor, conds, meta="join")776 777 def get_function_matches(self, suggestion, word_before_cursor,778 alias=False):779 if suggestion.usage == "from":780 # Only suggest functions allowed in FROM clause781 782 def filt(f):783 return (784 not f.is_aggregate and785 not f.is_window and786 not f.is_extension and787 (788 f.is_public or789 f.schema_name in self.search_path or790 f.schema_name == suggestion.schema791 )792 )793 794 else:795 alias = False796 797 def filt(f):798 return not f.is_extension and (799 f.is_public or f.schema_name == suggestion.schema800 )801 802 arg_mode = {"signature": "signature", "special": None}.get(803 suggestion.usage, "call"804 )805 806 # Function overloading means we way have multiple functions of the same807 # name at this point, so keep unique names only808 all_functions = self.populate_functions(suggestion.schema, filt)809 funcs = {self._make_cand(f, alias, suggestion, arg_mode)810 for f in all_functions}811 812 matches = self.find_matches(word_before_cursor, funcs, meta="function")813 814 if not suggestion.schema and not suggestion.usage:815 # also suggest hardcoded functions using startswith matching816 predefined_funcs = self.find_matches(817 word_before_cursor, self.functions, mode="strict",818 meta="function"819 )820 matches.extend(predefined_funcs)821 822 return matches823 824 def get_schema_matches(self, suggestion, word_before_cursor):825 schema_names = self.dbmetadata["tables"].keys()826 827 # Unless we're sure the user really wants them, hide schema names828 # starting with pg_, which are mostly temporary schemas829 if not word_before_cursor.startswith("pg_"):830 schema_names = [s for s in schema_names if not s.startswith("pg_")]831 832 if suggestion.quoted:833 schema_names = [self.escape_schema(s) for s in schema_names]834 835 return self.find_matches(word_before_cursor, schema_names,836 meta="schema")837 838 def get_from_clause_item_matches(self, suggestion, word_before_cursor):839 alias = self.generate_aliases840 s = suggestion841 t_sug = Table(s.schema, s.table_refs, s.local_tables)842 v_sug = View(s.schema, s.table_refs)843 f_sug = Function(s.schema, s.table_refs, usage="from")844 return (845 self.get_table_matches(t_sug, word_before_cursor, alias) +846 self.get_view_matches(v_sug, word_before_cursor, alias) +847 self.get_function_matches(f_sug, word_before_cursor, alias)848 )849 850 def _arg_list(self, func, usage):851 """Returns a an arg list string, e.g. `(_foo:=23)` for a func.852 853 :param func is a FunctionMetadata object854 :param usage is 'call', 'call_display' or 'signature'855 856 """857 template = {858 "call": self.call_arg_style,859 "call_display": self.call_arg_display_style,860 "signature": self.signature_arg_style,861 }[usage]862 args = func.args()863 if (not template or (usage == "call" and len(args) < 2) or864 (usage == "call" and func.has_variadic())):865 return "()"866 multiline = usage == "call" and len(args) > self.call_arg_oneliner_max867 max_arg_len = max(len(a.name) for a in args) if multiline else 0868 args = (869 self._format_arg(template, arg, arg_num + 1, max_arg_len)870 for arg_num, arg in enumerate(args)871 )872 if multiline:873 return "(" + ",".join("\n " + a for a in args if a) + "\n)"874 else:875 return "(" + ", ".join(a for a in args if a) + ")"876 877 def _format_arg(self, template, arg, arg_num, max_arg_len):878 if not template:879 return None880 if arg.has_default:881 arg_default = "NULL" if arg.default is None else arg.default882 # Remove trailing ::(schema.)type883 arg_default = arg_default_type_strip_regex.sub("", arg_default)884 else:885 arg_default = ""886 return template.format(887 max_arg_len=max_arg_len,888 arg_name=arg.name,889 arg_num=arg_num,890 arg_type=arg.datatype,891 arg_default=arg_default,892 )893 894 def _make_cand(self, tbl, do_alias, suggestion, arg_mode=None):895 """Returns a Candidate namedtuple.896 897 :param tbl is a SchemaObject898 :param arg_mode determines what type of arg list to suffix for899 functions.900 Possible values: call, signature901 902 """903 cased_tbl = tbl.name904 if do_alias:905 alias = self.alias(cased_tbl, suggestion.table_refs)906 synonyms = (cased_tbl, generate_alias(cased_tbl))907 maybe_alias = (" " + alias) if do_alias else ""908 maybe_schema = (tbl.schema + ".") if tbl.schema else ""909 suffix = self._arg_list_cache[arg_mode][tbl.meta] if arg_mode else ""910 if arg_mode == "call":911 display_suffix = self._arg_list_cache["call_display"][tbl.meta]912 elif arg_mode == "signature":913 display_suffix = self._arg_list_cache["signature"][tbl.meta]914 else:915 display_suffix = ""916 item = maybe_schema + cased_tbl + suffix + maybe_alias917 display = maybe_schema + cased_tbl + display_suffix + maybe_alias918 prio2 = 0 if tbl.schema else 1919 return Candidate(item, synonyms=synonyms, prio2=prio2, display=display)920 921 def get_table_matches(self, suggestion, word_before_cursor, alias=False):922 tables = self.populate_schema_objects(suggestion.schema, "tables")923 tables.extend(924 SchemaObject(tbl.name) for tbl in suggestion.local_tables)925 926 # Unless we're sure the user really wants them, don't suggest the927 # pg_catalog tables that are implicitly on the search path928 if not suggestion.schema and \929 (not word_before_cursor.startswith("pg_")):930 tables = [t for t in tables if not t.name.startswith("pg_")]931 tables = [self._make_cand(t, alias, suggestion) for t in tables]932 return self.find_matches(word_before_cursor, tables, meta="table")933 934 def get_view_matches(self, suggestion, word_before_cursor, alias=False):935 views = self.populate_schema_objects(suggestion.schema, "views")936 937 if not suggestion.schema and (938 not word_before_cursor.startswith("pg_")):939 views = [v for v in views if not v.name.startswith("pg_")]940 views = [self._make_cand(v, alias, suggestion) for v in views]941 return self.find_matches(word_before_cursor, views, meta="view")942 943 def get_alias_matches(self, suggestion, word_before_cursor):944 aliases = suggestion.aliases945 return self.find_matches(word_before_cursor, aliases,946 meta="table alias")947 948 def get_database_matches(self, _, word_before_cursor):949 return self.find_matches(word_before_cursor, self.databases,950 meta="database")951 952 def get_keyword_matches(self, suggestion, word_before_cursor):953 return self.find_matches(word_before_cursor, self.keywords,954 meta="keyword")955 956 def get_datatype_matches(self, suggestion, word_before_cursor):957 # suggest custom datatypes958 types = self.populate_schema_objects(suggestion.schema, "datatypes")959 types = [self._make_cand(t, False, suggestion) for t in types]960 matches = self.find_matches(word_before_cursor, types, meta="datatype")961 962 if not suggestion.schema:963 # Also suggest hardcoded types964 matches.extend(965 self.find_matches(966 word_before_cursor, self.datatypes, mode="strict",967 meta="datatype"968 )969 )970 971 return matches972 973 def get_word_before_cursor(self, word=False):974 """975 Give the word before the cursor.976 If we have whitespace before the cursor this returns an empty string.977 978 Args:979 word:980 """981 982 if self.text_before_cursor[-1:].isspace():983 return ''984 else:985 return self.text_before_cursor[self.find_start_of_previous_word(986 word=word987 ):]988 989 def find_start_of_previous_word(self, count=1, word=False):990 """991 Return an index relative to the cursor position pointing to the start992 of the previous word. Return `None` if nothing was found.993 994 Args:995 count:996 word:997 """998 999 # Reverse the text before the cursor, in order to do an efficient1000 # backwards search.1001 text_before_cursor = self.text_before_cursor[::-1]1002 1003 regex = _FIND_BIG_WORD_RE if word else _FIND_WORD_RE1004 iterator = regex.finditer(text_before_cursor)1005 1006 try:1007 for i, match in enumerate(iterator):1008 if i + 1 == count:1009 return - match.end(1)1010 except StopIteration:1011 pass1012 1013 suggestion_matchers = {1014 FromClauseItem: get_from_clause_item_matches,1015 JoinCondition: get_join_condition_matches,1016 Join: get_join_matches,1017 Column: get_column_matches,1018 Function: get_function_matches,1019 Schema: get_schema_matches,1020 Table: get_table_matches,1021 View: get_view_matches,1022 Alias: get_alias_matches,1023 Database: get_database_matches,1024 Keyword: get_keyword_matches,1025 Datatype: get_datatype_matches,1026 }1027 1028 def populate_scoped_cols(self, scoped_tbls, local_tbls=()):1029 """Find all columns in a set of scoped_tables.1030 1031 :param scoped_tbls: list of TableReference namedtuples1032 :param local_tbls: tuple(TableMetadata)1033 :return: {TableReference:{colname:ColumnMetaData}}1034 1035 """1036 ctes = {normalize_ref(t.name): t.columns for t in local_tbls}1037 columns = OrderedDict()1038 meta = self.dbmetadata1039 1040 def addcols(schema, rel, alias, reltype, cols):1041 tbl = TableReference(schema, rel, alias, reltype == "functions")1042 if tbl not in columns:1043 columns[tbl] = []1044 columns[tbl].extend(cols)1045 1046 for tbl in scoped_tbls:1047 # Local tables should shadow database tables1048 if tbl.schema is None and normalize_ref(tbl.name) in ctes:1049 cols = ctes[normalize_ref(tbl.name)]1050 addcols(None, tbl.name, "CTE", tbl.alias, cols)1051 continue1052 schemas = [tbl.schema] if tbl.schema else self.search_path1053 for schema in schemas:1054 relname = self.escape_name(tbl.name)1055 schema = self.escape_name(schema)1056 if tbl.is_function:1057 # Return column names from a set-returning function1058 # Get an array of FunctionMetadata objects1059 functions = meta["functions"].get(schema, {}).get(relname)1060 for func in functions or []:1061 # func is a FunctionMetadata object1062 cols = func.fields()1063 addcols(schema, relname, tbl.alias, "functions", cols)1064 else:1065 for reltype in ("tables", "views"):1066 cols = meta[reltype].get(schema, {}).get(relname)1067 if cols:1068 cols = cols.values()1069 addcols(schema, relname, tbl.alias, reltype, cols)1070 break1071 1072 return columns1073 1074 def _get_schemas(self, obj_typ, schema):1075 """Returns a list of schemas from which to suggest objects.1076 1077 :param schema is the schema qualification input by the user (if any)1078 1079 """1080 metadata = self.dbmetadata[obj_typ]1081 if schema:1082 schema = self.escape_name(schema)1083 return [schema] if schema in metadata else []1084 return self.search_path if self.search_path_filter else metadata.keys()1085 1086 def _maybe_schema(self, schema, parent):1087 return None if parent or schema in self.search_path else schema1088 1089 def populate_schema_objects(self, schema, obj_type):1090 """Returns a list of SchemaObjects representing tables or views.1091 1092 :param schema is the schema qualification input by the user (if any)1093 1094 """1095 # Fetch the schema objects first1096 self.fetch_schema_objects(schema, obj_type)1097 1098 return [1099 SchemaObject(1100 name=obj,1101 schema=(self._maybe_schema(schema=sch, parent=schema))1102 )1103 for sch in self._get_schemas(obj_type, schema)1104 for obj in self.dbmetadata[obj_type][sch].keys()1105 ]1106 1107 def populate_functions(self, schema, filter_func):1108 """Returns a list of function SchemaObjects.1109 1110 :param filter_func is a function that accepts a FunctionMetadata1111 namedtuple and returns a boolean indicating whether that1112 function should be kept or discarded1113 1114 """1115 1116 # Fetch the functions list1117 self.fetch_functions(schema)1118 1119 # Because of multiple dispatch, we can have multiple functions1120 # with the same name, which is why `for meta in metas` is necessary1121 # in the comprehensions below1122 return [1123 SchemaObject(1124 name=func,1125 schema=(self._maybe_schema(schema=sch, parent=schema)),1126 meta=meta,1127 )1128 for sch in self._get_schemas("functions", schema)1129 for (func, metas) in self.dbmetadata["functions"][sch].items()1130 for meta in metas1131 if filter_func(meta)1132 ]1133 1134 def _get_schema_obj_query(self, schema, obj_type):1135 """1136 Get query according object type like tables, views, etc...1137 :param schema: schema flag to include schema in clause.1138 :param obj_type: object type.1139 :return: query according to object type and in_clause1140 if schema flag in true.1141 """1142 in_clause = ''1143 query = ''1144 1145 if schema:1146 in_clause = '\'' + schema.replace("'", "''") + '\''1147 else:1148 for r in self.search_path:1149 in_clause += '\'' + r + '\','1150 # Remove extra comma1151 if len(in_clause) > 0:1152 in_clause = in_clause[:-1]1153 1154 if obj_type == 'tables':1155 query = render_template("/".join([self.sql_path, 'tableview.sql']),1156 schema_names=in_clause,1157 object_name='tables')1158 elif obj_type == 'views':1159 query = render_template("/".join([self.sql_path, 'tableview.sql']),1160 schema_names=in_clause,1161 object_name='views')1162 elif obj_type == 'datatypes':1163 query = render_template("/".join([self.sql_path, 'datatypes.sql']),1164 schema_names=in_clause)1165 1166 return query, in_clause1167 1168 def fetch_schema_objects(self, schema, obj_type):1169 """1170 This function is used to fetch schema objects like tables, views, etc..1171 :return:1172 """1173 data = []1174 query, in_clause = self._get_schema_obj_query(schema, obj_type)1175 1176 if self.conn.connected():1177 status, res = self.conn.execute_dict(query)1178 if status:1179 for record in res['rows']:1180 data.append(1181 (record['schema_name'], record['object_name'])1182 )1183 1184 if (obj_type == 'tables' or obj_type == 'views') and len(data) > 0:1185 self.extend_relations(data, obj_type)1186 self.extend_columns(1187 self.fetch_columns(in_clause, obj_type), obj_type1188 )1189 if obj_type == 'tables':1190 self.extend_foreignkeys(1191 self.fetch_foreign_keys(in_clause)1192 )1193 elif obj_type == 'datatypes' and len(data) > 0:1194 self.extend_datatypes(data)1195 1196 def _get_function_sql(self, schema):1197 """1198 Check for schema inclusion and fetch sql for functions.1199 :param schema: include schema flag.1200 :return: sql query for functions, and in_clause value.