Function extract_from_part

mycli/packages/sql_utils.py:128–170

(parsed: TokenList, stop_at_punctuation: bool = True)

def extract_from_part(parsed: TokenList, stop_at_punctuation: bool = True) -> Generator[Any, None, None]:
    tbl_prefix_seen = False
    for item in parsed.tokens:
        if tbl_prefix_seen:
            if is_subselect(item):
                yield from extract_from_part(item, stop_at_punctuation)
            elif stop_at_punctuation and item.ttype is Punctuation:
                return None
            # Multiple JOINs in the same query won't work properly since
            # "ON" is a keyword and will trigger the next elif condition.
            # So instead of stooping the loop when finding an "ON" skip it
            # eg: 'SELECT * FROM abc JOIN def ON abc.id = def.abc_id JOIN ghi'
            elif item.ttype is Keyword and item.value.upper() == "ON":
                tbl_prefix_seen = False
                continue
            # An incomplete nested select won't be recognized correctly as a
            # sub-select. eg: 'SELECT * FROM (SELECT id FROM user'. This causes
            # the second FROM to trigger this elif condition resulting in a
            # StopIteration. So we need to ignore the keyword if the keyword
            # FROM.
            # Also 'SELECT * FROM abc JOIN def' will trigger this elif
            # condition. So we need to ignore the keyword JOIN and its variants
            # INNER JOIN, FULL OUTER JOIN, etc.
            elif item.ttype is Keyword and (not item.value.upper() == "FROM") and (not item.value.upper().endswith("JOIN")):
                return None
            else:
                yield item
        elif (item.ttype is Keyword or item.ttype is Keyword.DML) and item.value.upper() in (
            "COPY",
            "FROM",
            "INTO",
            "UPDATE",
            "TABLE",
            "JOIN",
        ):
            tbl_prefix_seen = True
        # 'SELECT a, FROM abc' will detect FROM as part of the column list.
        # So this check here is necessary.
        elif isinstance(item, IdentifierList):
            for identifier in item.get_identifiers():
                if identifier.ttype is Keyword and identifier.value.upper() == "FROM":
                    tbl_prefix_seen = True
                    break


def extract_table_identifiers(token_stream: Generator[Any, None, None]) -> Generator[tuple[str | None, str, str], None, None]: