path: root/pgcli/packages/sqlcompletion.py

import sys
import re
import sqlparse
from collections import namedtuple
from sqlparse.sql import Comparison, Identifier, Where
from .parseutils.utils import last_word, find_prev_keyword, parse_partial_identifier
from .parseutils.tables import extract_tables
from .parseutils.ctes import isolate_query_ctes
from pgspecial.main import parse_special_command


Special = namedtuple("Special", [])
Database = namedtuple("Database", [])
Schema = namedtuple("Schema", ["quoted"])
Schema.__new__.__defaults__ = (False,)
# FromClauseItem is a table/view/function used in the FROM clause
# `table_refs` contains the list of tables/... already in the statement,
# used to ensure that the alias we suggest is unique
FromClauseItem = namedtuple("FromClauseItem", "schema table_refs local_tables")
Table = namedtuple("Table", ["schema", "table_refs", "local_tables"])
TableFormat = namedtuple("TableFormat", [])
View = namedtuple("View", ["schema", "table_refs"])
# JoinConditions are suggested after ON, e.g. 'foo.barid = bar.barid'
JoinCondition = namedtuple("JoinCondition", ["table_refs", "parent"])
# Joins are suggested after JOIN, e.g. 'foo ON foo.barid = bar.barid'
Join = namedtuple("Join", ["table_refs", "schema"])

Function = namedtuple("Function", ["schema", "table_refs", "usage"])
# For convenience, don't require the `usage` argument in Function constructor
Function.__new__.__defaults__ = (None, tuple(), None)
Table.__new__.__defaults__ = (None, tuple(), tuple())
View.__new__.__defaults__ = (None, tuple())
FromClauseItem.__new__.__defaults__ = (None, tuple(), tuple())

Column = namedtuple(
    "Column",
    ["table_refs", "require_last_table", "local_tables", "qualifiable", "context"],
)
Column.__new__.__defaults__ = (None, None, tuple(), False, None)

Keyword = namedtuple("Keyword", ["last_token"])
Keyword.__new__.__defaults__ = (None,)
NamedQuery = namedtuple("NamedQuery", [])
Datatype = namedtuple("Datatype", ["schema"])
Alias = namedtuple("Alias", ["aliases"])

Path = namedtuple("Path", [])


class SqlStatement(object):
    def __init__(self, full_text, text_before_cursor):
        self.identifier = None
        self.word_before_cursor = word_before_cursor = last_word(
            text_before_cursor, include="many_punctuations"
        )
        full_text = _strip_named_query(full_text)
        text_before_cursor = _strip_named_query(text_before_cursor)

        full_text, text_before_cursor, self.local_tables = isolate_query_ctes(
            full_text, text_before_cursor
        )

        self.text_before_cursor_including_last_word = text_before_cursor

        # If we've partially typed a word then word_before_cursor won't be an
        # empty string. In that case we want to remove the partially typed
        # string before sending it to the sqlparser. Otherwise the last token
        # will always be the partially typed string which renders the smart
        # completion useless because it will always return the list of
        # keywords as completion.
        if self.word_before_cursor:
            if word_before_cursor[-1] == "(" or word_before_cursor[0] == "\\":
                parsed = sqlparse.parse(text_before_cursor)
            else:
                text_before_cursor = text_before_cursor[: -len(word_before_cursor)]
                parsed = sqlparse.parse(text_before_cursor)
                self.identifier = parse_partial_identifier(word_before_cursor)
        else:
            parsed = sqlparse.parse(text_before_cursor)

        full_text, text_before_cursor, parsed = _split_multiple_statements(
            full_text, text_before_cursor, parsed
        )

        self.full_text = full_text
        self.text_before_cursor = text_before_cursor
        self.parsed = parsed

        self.last_token = parsed and parsed.token_prev(len(parsed.tokens))[1] or ""

    def is_insert(self):
        return self.parsed.token_first().value.lower() == "insert"

    def get_tables(self, scope="full"):
        """ Gets the tables available in the statement.
        param `scope:` possible values: 'full', 'insert', 'before'
        If 'insert', only the first table is returned.
        If 'before', only tables before the cursor are returned.
        If not 'insert' and the stmt is an insert, the first table is skipped.
        """
        tables = extract_tables(
            self.full_text if scope == "full" else self.text_before_cursor
        )
        if scope == "insert":
            tables = tables[:1]
        elif self.is_insert():
            tables = tables[1:]
        return tables

    def get_previous_token(self, token):
        return self.parsed.token_prev(self.parsed.token_index(token))[1]

    def get_identifier_schema(self):
        schema = (self.identifier and self.identifier.get_parent_name()) or None
        # If schema name is unquoted, lower-case it
        if schema and self.identifier.value[0] != '"':
            schema = schema.lower()

        return schema

    def reduce_to_prev_keyword(self, n_skip=0):
        prev_keyword, self.text_before_cursor = find_prev_keyword(
            self.text_before_cursor, n_skip=n_skip
        )
        return prev_keyword


def suggest_type(full_text, text_before_cursor):
    """Takes the full_text that is typed so far and also the text before the
    cursor to suggest completion type and scope.

    Returns a tuple with a type of entity ('table', 'column' etc) and a scope.
    A scope for a column category will be a list of tables.
    """

    if full_text.startswith("\\i "):
        return (Path(),)

    # This is a temporary hack; the exception handling
    # here should be removed once sqlparse has been fixed
    try:
        stmt = SqlStatement(full_text, text_before_cursor)
    except (TypeError, AttributeError):
        return []

    # Check for special commands and handle those separately
    if stmt.parsed:
        # Be careful here because trivial whitespace is parsed as a
        # statement, but the statement won't have a first token
        tok1 = stmt.parsed.token_first()
        if tok1 and tok1.value.startswith("\\"):
            text = stmt.text_before_cursor + stmt.word_before_cursor
            return suggest_special(text)

    return suggest_based_on_last_token(stmt.last_token, stmt)


named_query_regex = re.compile(r"^\s*\\ns\s+[A-z0-9\-_]+\s+")


def _strip_named_query(txt):
    """
    This will strip "save named query" command in the beginning of the line:
    '\ns zzz SELECT * FROM abc'   -> 'SELECT * FROM abc'
    '  \ns zzz SELECT * FROM abc' -> 'SELECT * FROM abc'
    """

    if named_query_regex.match(txt):
        txt = named_query_regex.sub("", txt)
    return txt


function_body_pattern = re.compile(r"(\$.*?\$)([\s\S]*?)\1", re.M)


def _find_function_body(text):
    split = function_body_pattern.search(text)
    return (split.start(2), split.end(2)) if split else (None, None)


def _statement_from_function(full_text, text_before_cursor, statement):
    current_pos = len(text_before_cursor)
    body_start, body_end = _find_function_body(full_text)
    if body_start is None:
        return full_text, text_before_cursor, statement
    if not body_start <= current_pos < body_end:
        return full_text, text_before_cursor, statement
    full_text = full_text[body_start:body_end]
    text_before_cursor = text_before_cursor[body_start:]
    parsed = sqlparse.parse(text_before_cursor)
    return _split_multiple_statements(full_text, text_before_cursor, parsed)


def _split_multiple_statements(full_text, text_before_cursor, parsed):
    if len(parsed) > 1:
        # Multiple statements being edited -- isolate the current one by
        # cumulatively summing statement lengths to find the one that bounds
        # the current position
        current_pos = len(text_before_cursor)
        stmt_start, stmt_end = 0, 0

        for statement in parsed:
            stmt_len = len(str(statement))
            stmt_start, stmt_end = stmt_end, stmt_end + stmt_len

            if stmt_end >= current_pos:
                text_before_cursor = full_text[stmt_start:current_pos]
                full_text = full_text[stmt_start:]
                break

    elif parsed:
        # A single statement
        statement = parsed[0]
    else:
        # The empty string
        return full_text, text_before_cursor, None

    token2 = None
    if statement.get_type() in ("CREATE", "CREATE OR REPLACE"):
        token1 = statement.token_first()
        if token1:
            token1_idx = statement.token_index(token1)
            token2 = statement.token_next(token1_idx)[1]
    if token2 and token2.value.upper() == "FUNCTION":
        full_text, text_before_cursor, statement = _statement_from_function(
            full_text, text_before_cursor, statement
        )
    return full_text, text_before_cursor, statement


SPECIALS_SUGGESTION = {
    "dT": Datatype,
    "df": Function,
    "dt": Table,
    "dv": View,
    "sf": Function,
}


def suggest_special(text):
    text = text.lstrip()
    cmd, _, arg = parse_special_command(text)

    if cmd == text:
        # Trying to complete the special command itself</