Function find_sql_refs

marimo/_ast/sql_visitor.py:480–608

Find table and schema references in a SQL statement. Args: sql_statement: The SQL statement to parse. Returns: A set of unique SQLRefs, one for each table reference in the statement. Eg. SELECT * FROM schema1.test_table INNER JOIN schema2.test_table2 wo

(sql_statement: str)

def find_sql_refs(sql_statement: str) -> set[SQLRef]:
    """
    Find table and schema references in a SQL statement.

    Args:
        sql_statement: The SQL statement to parse.

    Returns:
        A set of unique SQLRefs, one for each table reference in the statement.
        Eg. SELECT * FROM schema1.test_table INNER JOIN schema2.test_table2
        would return two SQLRefs, one for the first table and one for the second.

    Note:
        When providing only a single qualification,
        DuckDB will interpret as either a catalog or a schema, as long as there are no conflicts.

        Eg. SELECT * FROM my_db.my_table, my_db can be a catalog or schema. If a catalog exists,
        then it would resolve to my_db.main.my_table.

        At the moment, we don't know this, so my_db is treated as a schema.
    """

    # Use sqlglot to parse ast (https://github.com/tobymao/sqlglot/blob/main/posts/ast_primer.md)

    DependencyManager.sqlglot.require(why="SQL parsing")

    from sqlglot import exp, parse
    from sqlglot.errors import OptimizeError
    from sqlglot.optimizer.scope import build_scope

    def get_ref_from_table(table: exp.Table) -> SQLRef | None:
        # The variables might be empty strings, if they are, we set them to None
        try:
            table_name = table.name or None
        except AttributeError:
            # sqlglot may return Table nodes with this=None (e.g. DROP SCHEMA)
            return None
        schema_name = table.db or None
        catalog_name = table.catalog or None

        if table_name is None:
            return None

        # Check if the table name looks like a URL or has a file extension.
        # These are often not actual table references, so we skip them.
        # Note that they can be valid table names, but we skip them to avoid circular deps
        if "://" in table_name or table_name.endswith(COMMON_FILE_EXTENSIONS):
            return None

        return SQLRef(
            table=table_name, schema=schema_name, catalog=catalog_name
        )

    from sqlglot.errors import ParseError

    try:
        with _loggers.suppress_warnings_logs("sqlglot"):
            expression_list = parse(sql_statement, dialect="duckdb")