codekingpro/portable-devtools

1import sqlparse2from collections import namedtuple3from sqlparse.sql import IdentifierList, Identifier, Function4from sqlparse.tokens import Keyword, DML, Punctuation5 6TableReference = namedtuple(7    "TableReference", ["schema", "name", "alias", "is_function"]8)9TableReference.ref = property(10    lambda self: self.alias or (11        self.name12        if self.name.islower() or self.name[0] == '"'13        else '"' + self.name + '"'14    )15)16 17 18# This code is borrowed from sqlparse example script.19# <url>20def is_subselect(parsed):21    if not parsed.is_group:22        return False23    for item in parsed.tokens:24        if item.ttype is DML and item.value.upper() in (25            "SELECT",26            "INSERT",27            "UPDATE",28            "CREATE",29            "DELETE",30        ):31            return True32    return False33 34 35def _identifier_is_function(identifier):36    return any(isinstance(t, Function) for t in identifier.tokens)37 38 39def extract_from_part(parsed, stop_at_punctuation=True):40    tbl_prefix_seen = False41    for item in parsed.tokens:42        if tbl_prefix_seen:43            if is_subselect(item):44                yield from extract_from_part(item, stop_at_punctuation)45            elif stop_at_punctuation and item.ttype is Punctuation:46                return47            # An incomplete nested select won't be recognized correctly as a48            # sub-select. eg: 'SELECT * FROM (SELECT id FROM user'. This causes49            # the second FROM to trigger this elif condition resulting in a50            # `return`. So we need to ignore the keyword if the keyword51            # FROM.52            # Also 'SELECT * FROM abc JOIN def' will trigger this elif53            # condition. So we need to ignore the keyword JOIN and its variants54            # INNER JOIN, FULL OUTER JOIN, etc.55            elif (56                item.ttype is Keyword and57                item.value.upper() != "FROM" and58                (not item.value.upper().endswith("JOIN"))59            ):60                tbl_prefix_seen = False61            else:62                yield item63        elif item.ttype is Keyword or item.ttype is Keyword.DML:64            item_val = item.value.upper()65            if item_val in (66                "COPY",67                "FROM",68                "INTO",69                "UPDATE",70                "TABLE",71            ) or item_val.endswith("JOIN"):72                tbl_prefix_seen = True73        # 'SELECT a, FROM abc' will detect FROM as part of the column list.74        # So this check here is necessary.75        elif isinstance(item, IdentifierList):76            for identifier in item.get_identifiers():77                if identifier.ttype is Keyword and \78                        identifier.value.upper() == "FROM":79                    tbl_prefix_seen = True80                    break81 82 83def extract_table_identifiers(token_stream, allow_functions=True):84    """yields tuples of TableReference namedtuples"""85 86    # We need to do some massaging of the names because postgres is case-87    # insensitive and '"Foo"' is not the same table as 'Foo' (while 'foo' is)88    def parse_identifier(item):89        name = item.get_real_name()90        schema_name = item.get_parent_name()91        alias = item.get_alias()92        if not name:93            schema_name = None94            name = item.get_name()95            alias = alias or name96        schema_quoted = schema_name and item.value[0] == '"'97        if schema_name and not schema_quoted:98            schema_name = schema_name.lower()99        quote_count = item.value.count('"')100        name_quoted = quote_count > 2 or (quote_count and not schema_quoted)101        alias_quoted = alias and item.value[-1] == '"'102        if alias_quoted or name_quoted and not alias and name.islower():103            alias = '"' + (alias or name) + '"'104        if name and not name_quoted and not name.islower():105            if not alias:106                alias = name107            name = name.lower()108        return schema_name, name, alias109 110    try:111        for item in token_stream:112            if isinstance(item, IdentifierList):113                for identifier in item.get_identifiers():114                    # Sometimes Keywords (such as FROM ) are classified as115                    # identifiers which don't have the get_real_name() method.116                    try:117                        schema_name = identifier.get_parent_name()118                        real_name = identifier.get_real_name()119                        is_function = allow_functions and \120                            _identifier_is_function(identifier)121                    except AttributeError:122                        continue123                    if real_name:124                        yield TableReference(125                            schema_name, real_name, identifier.get_alias(),126                            is_function127                        )128            elif isinstance(item, Identifier):129                schema_name, real_name, alias = parse_identifier(item)130                is_function = allow_functions and _identifier_is_function(item)131 132                yield TableReference(schema_name, real_name, alias,133                                     is_function)134            elif isinstance(item, Function):135                schema_name, real_name, alias = parse_identifier(item)136                yield TableReference(None, real_name, alias, allow_functions)137    except StopIteration:138        return139 140 141# extract_tables is inspired from examples in the sqlparse lib.142def extract_tables(sql):143    """Extract the table names from an SQL statment.144 145    Returns a list of TableReference namedtuples146 147    """148    parsed = sqlparse.parse(sql)149    if not parsed:150        return ()151 152    # INSERT statements must stop looking for tables at the sign of first153    # Punctuation. eg: INSERT INTO abc (col1, col2) VALUES (1, 2)154    # abc is the table name, but if we don't stop at the first lparen, then155    # we'll identify abc, col1 and col2 as table names.156    insert_stmt = parsed[0].token_first().value.lower() == "insert"157    stream = extract_from_part(parsed[0], stop_at_punctuation=insert_stmt)158 159    # Kludge: sqlparse mistakenly identifies insert statements as160    # function calls due to the parenthesized column list, e.g. interprets161    # "insert into foo (bar, baz)" as a function call to foo with arguments162    # (bar, baz). So don't allow any identifiers in insert statements163    # to have is_function=True164    identifiers = extract_table_identifiers(stream,165                                            allow_functions=not insert_stmt)166    # In the case 'sche.<cursor>', we get an empty TableReference; remove that167    return tuple(i for i in identifiers if i.name)168