codekingpro/portable-devtools

1import re2import sqlparse3from sqlparse.sql import Identifier4from sqlparse.tokens import Token, Error5 6cleanup_regex = {7    # This matches only alphanumerics and underscores.8    "alphanum_underscore": re.compile(r"(\w+)$"),9    # This matches everything except spaces, parens, colon, and comma10    "many_punctuations": re.compile(r"([^():,\s]+)$"),11    # This matches everything except spaces, parens, colon, comma, and period12    "most_punctuations": re.compile(r"([^\.():,\s]+)$"),13    # This matches everything except a space.14    "all_punctuations": re.compile(r"([^\s]+)$"),15}16 17 18def last_word(text, include="alphanum_underscore"):19    r"""20    Find the last word in a sentence.21 22    >>> last_word('abc')23    'abc'24    >>> last_word(' abc')25    'abc'26    >>> last_word('')27    ''28    >>> last_word(' ')29    ''30    >>> last_word('abc ')31    ''32    >>> last_word('abc def')33    'def'34    >>> last_word('abc def ')35    ''36    >>> last_word('abc def;')37    ''38    >>> last_word('bac $def')39    'def'40    >>> last_word('bac $def', include='most_punctuations')41    '$def'42    >>> last_word('bac \def', include='most_punctuations')43    '\\\\def'44    >>> last_word('bac \def;', include='most_punctuations')45    '\\\\def;'46    >>> last_word('bac::def', include='most_punctuations')47    'def'48    >>> last_word('"foo*bar', include='most_punctuations')49    '"foo*bar'50    """51 52    if not text:  # Empty string53        return ""54 55    if text[-1].isspace():56        return ""57    else:58        regex = cleanup_regex[include]59        matches = regex.search(text)60        if matches:61            return matches.group(0)62        else:63            return ""64 65 66def find_prev_keyword(sql, n_skip=0):67    """Find the last sql keyword in an SQL statement68 69    Returns the value of the last keyword, and the text of the query with70    everything after the last keyword stripped71    """72    if not sql.strip():73        return None, ""74 75    parsed = sqlparse.parse(sql)[0]76    flattened = list(parsed.flatten())77    flattened = flattened[: len(flattened) - n_skip]78 79    logical_operators = ("AND", "OR", "NOT", "BETWEEN")80 81    for t in reversed(flattened):82        if t.value == "(" or (83            t.is_keyword and (t.value.upper() not in logical_operators)84        ):85            # Find the location of token t in the original parsed statement86            # We can't use parsed.token_index(t) because t may be a child token87            # inside a TokenList, in which case token_index throws an error88            # Minimal example:89            #   p = sqlparse.parse('select * from foo where bar')90            #   t = list(p.flatten())[-3]  # The "Where" token91            #   p.token_index(t)  # Throws ValueError: not in list92            idx = flattened.index(t)93 94            # Combine the string values of all tokens in the original list95            # up to and including the target keyword token t, to produce a96            # query string with everything after the keyword token removed97            text = "".join(tok.value for tok in flattened[: idx + 1])98            return t, text99 100    return None, ""101 102 103# Postgresql dollar quote signs look like `$$` or `$tag$`104dollar_quote_regex = re.compile(r"^\$[^$]*\$$")105 106 107def is_open_quote(sql):108    """Returns true if the query contains an unclosed quote"""109 110    # parsed can contain one or more semi-colon separated commands111    parsed = sqlparse.parse(sql)112    return any(_parsed_is_open_quote(p) for p in parsed)113 114 115def _parsed_is_open_quote(parsed):116    # Look for unmatched single quotes, or unmatched dollar sign quotes117    return any(tok.match(Token.Error, ("'", "$")) for tok in parsed.flatten())118 119 120def parse_partial_identifier(word):121    """Attempt to parse a (partially typed) word as an identifier122 123    word may include a schema qualification, like `schema_name.partial_name`124    or `schema_name.` There may also be unclosed quotation marks, like125    `"schema`, or `schema."partial_name`126 127    :param word: string representing a (partially complete) identifier128    :return: sqlparse.sql.Identifier, or None129    """130 131    p = sqlparse.parse(word)[0]132    n_tok = len(p.tokens)133    if n_tok == 1 and isinstance(p.tokens[0], Identifier):134        return p.tokens[0]135    elif p.token_next_by(m=(Error, '"'))[1]:136        # An unmatched double quote, e.g. '"foo', 'foo."', or 'foo."bar'137        # Close the double quote, then reparse138        return parse_partial_identifier(word + '"')139    else:140        return None141