sqlleaf 0.0.1__py3-none-any.whl

sqlleaf/processors/collector/collector.py

@@ -0,0 +1,827 @@
+import logging
+import typing as t
+from collections import Counter
+from dataclasses import dataclass, field
+
+import sqlglot
+from sqlglot import TokenType, exp
+from sqlglot.expressions import ColumnDef
+from sqlleaf.processors.transformer import udf
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+
+from sqlleaf import exception, mappings, settings, util, typing
+from sqlleaf.models.query import (
+    CallQuery,
+    CopyQuery,
+    CTASQuery,
+    DatabaseQuery,
+    DeleteQuery,
+    ExecuteQuery,
+    InsertQuery,
+    MergeQuery,
+    MultitableInsertQuery,
+    PrepareQuery,
+    ProcedureQuery,
+    PutQuery,
+    Q,
+    QueryHolder,
+    ReplaceQuery,
+    SchemaQuery,
+    SelectQuery,
+    SequenceQuery,
+    SetQuery,
+    StageQuery,
+    TableQuery,
+    TriggerQuery,
+    TypeQuery,
+    UnloadQuery,
+    UpdateQuery,
+    UserDefinedFunctionQuery,
+    ValuesQuery,
+    ViewQuery,
+)
+from sqlleaf.processors.collector import substitute
+from sqlleaf.processors.transformer.expressions.row import simplify_row_in_values
+
+logger = logging.getLogger("sqlleaf")
+
+"""
+Parses text for SQL statements and collects them into Query models.
+"""
+
+
+@dataclass(frozen=True)
+class CollectQueryResult:
+    queries: t.List[QueryHolder] = field(default_factory=list)  # Successfully collected queries
+    unknown: t.Dict[str, int] = field(default_factory=dict)  # Unsupported by sqlleaf (no handler)
+    unsupported: t.List[t.Tuple[int, exp.Expr]] = field(default_factory=list)  # Unsupported by sqlglot (no grammar)
+    object_mapping: mappings.ObjectMapping = None
+
+
+def _is_prepare_supported(stmt: exp.Command) -> bool:
+    """
+    Check if a PREPARE statement for Postgres is supported.
+    Syntax: PREPARE name AS statement
+    The variant where parameters are provided is not yet supported.
+    """
+    expression_name = stmt.expression.name
+    tokens = sqlglot.tokenize(expression_name, dialect="postgres")
+
+    if len(tokens) < 3:
+        raise exception.InvalidQueryError(f"Invalid syntax for PREPARE expression: {stmt.sql(dialect='postgres')}")
+
+    # We cannot process arguments yet
+    if tokens[1].token_type != TokenType.ALIAS or tokens[1].text.upper() != "AS":
+        if tokens[1].token_type == TokenType.L_PAREN:
+            logger.warning("PREPARE with arguments is not currently supported.")
+        return False
+
+    # Ensure there's an 'AS' token
+    if not any(t.token_type == TokenType.ALIAS and t.text.upper() == "AS" for t in tokens):
+        raise exception.InvalidQueryError(f"Could not find 'AS' in PREPARE expression: {expression_name}")
+
+    return True
+
+
+def _is_execute_supported(stmt: exp.Command) -> bool:
+    """
+    Check if an EXECUTE statement for Postgres is supported.
+    Syntax: EXECUTE name [ ( parameter [, ...] ) ]
+
+    Only the bare-name form (exactly one token, no argument parentheses) is supported.
+    If the expression contains more than one token it means arguments were supplied
+    (e.g. EXECUTE my_plan(arg1, arg2)), which is not yet handled.
+    """
+    expression_name = stmt.expression.name
+    tokens = sqlglot.tokenize(expression_name, dialect="postgres")
+
+    if not tokens or len(tokens) > 1:
+        raise exception.InvalidQueryError(f"Invalid syntax for EXECUTE expression: {stmt.sql(dialect='postgres')}")
+
+    return True
+
+
+def _is_replace_supported(stmt: exp.Command, dialect: str) -> bool:
+    """
+    Check if a REPLACE statement is supported.
+    Transform to INSERT INTO to verify syntax.
+    """
+    expression = stmt.args.get("expression")
+    new_sql = f"INSERT {expression.this}" if expression else "INSERT"
+    try:
+        sqlglot.parse_one(new_sql, dialect=dialect)
+        return True
+    except Exception as e:
+        logger.warning(f"Invalid REPLACE statement syntax: {e}")
+        return False
+
+
+def _classify_command(stmt: exp.Command, dialect: str) -> tuple[str, bool]:
+    """
+    Classify an exp.Command node into a (kind, is_supported) pair.
+    kind='' means the command is not recognised and should be treated as unsupported.
+    """
+    if dialect in ["athena", "redshift"] and stmt.name == "UNLOAD":
+        return "unload", True
+    if dialect == "postgres" and stmt.name == "PREPARE":
+        return "prepare", _is_prepare_supported(stmt)
+    if dialect == "postgres" and stmt.name == "EXECUTE":
+        return "execute", _is_execute_supported(stmt)
+    if dialect == "mysql" and stmt.name == "REPLACE":
+        return "replace", _is_replace_supported(stmt, dialect)
+    if stmt.this.upper() == "CALL":
+        return "call", True
+    return "", False
+
+
+def collect_queries(text: str, dialect: str, object_mapping: mappings.ObjectMapping | None = None) -> CollectQueryResult:
+    """
+    Parse a series of SQL statements provided as text.
+    This includes tables, views, procedures, functions, sequences, etc.
+
+    Each query may contain multiple child queries. For example, a stored procedure often
+    has multiple individual queries. Each of these individual queries may also have
+    subqueries. For example, a MERGE query often has INSERTs or UPDATEs in its WHEN clauses.
+
+    The statements must be provided in the order in which they depend on each other.
+    If B depends on A, A must be created before B.
+    """
+    queries = {}
+    unknown = {}
+    unsupported = []
+    counts: Counter[str] = Counter()
+
+    if not object_mapping:
+        object_mapping = mappings.ObjectMapping(dialect=dialect)
+
+    # Parse the statements
+    parsed = sqlglot.parse(text, dialect=dialect)
+    parsed = _split_combined_statements(parsed)
+
+    for index, stmt in enumerate(parsed):
+        if not stmt:
+            continue
+
+        kind = ""
+        if isinstance(stmt, exp.Command):
+            kind, supported = _classify_command(stmt, dialect)
+            if not kind:
+                logger.warning(f"Unsupported statement: {stmt.sql(dialect=dialect)}")
+                unsupported.append((index, stmt))
+                continue
+            if not supported:
+                unsupported.append((index, stmt))
+                continue
+
+        # Remove duplicate queries
+        sql_text = stmt.sql(dialect=dialect)
+        _id = util.short_sha256_hash(sql_text)
+        if _id in queries:
+            logger.debug(f"Skipping duplicate query: {sql_text}")
+            continue
+
+        if not kind:
+            stmt, kind = _determine_query_kind(stmt, dialect)
+
+        if kind not in _QUERY_PROCESSORS:
+            unknown[kind] = unknown[kind] + 1 if kind in unknown else 1
+            continue
+
+        # Convert the statement to uppercase if the dialect supports it
+        stmt = normalize_identifiers(stmt, dialect=dialect, store_original_column_identifiers=True)
+        _unnest_values_inside_select(stmt, dialect=dialect)
+
+        # Substitute any Snowflake session variables ($var) in non-SET statements
+        if kind != "set" and object_mapping.session_variables:
+            stmt = substitute.substitute_session_variables(stmt, object_mapping.session_variables)
+
+        query: t.Optional[Q] = _QUERY_PROCESSORS[kind](
+            statement=stmt, dialect=dialect, object_mapping=object_mapping, statement_index=index
+        )
+        if query:
+            holder = QueryHolder(original=query)
+            _collect_query_children(query, holder, dialect, object_mapping)
+            queries[_id] = holder
+            counts[kind] += 1
+
+    logger.debug("Found statements: %s", dict(+counts))
+    if unknown:
+        logger.warning("Unknown statements: %s", dict(unknown.items()))
+    if unsupported:
+        logger.warning("Unsupported statements: %s", len(unsupported))
+
+    return CollectQueryResult(queries=list(queries.values()), unknown=unknown, unsupported=unsupported, object_mapping=object_mapping)
+
+
+def _determine_query_kind(statement: exp.Expr, dialect: str) -> t.Tuple[exp.Expr, str]:
+    """
+    Determine a query's "kind" from the expression, which maps to how it will be processed.
+    """
+    if statement.key == "create" and isinstance(statement, exp.Create):
+        if statement.kind == "TABLE":
+            if isinstance(statement.expression, (exp.Select, exp.Values)) or statement.find(exp.ExecuteAsProperty):
+                kind = "ctas"
+            else:
+                kind = "table"
+        else:
+            kind = (statement.kind or "").lower()
+    elif statement.key == "select" and "into" in statement.args:
+        # sqlglot rewrites 'SELECT INTO' to 'CREATE TABLE AS' during parse()
+        # but it's not shown until we produce it with sql(), so we re-parse it
+        if dialect in ["redshift", "postgres"]:
+            statement = sqlglot.parse_one(statement.sql(dialect=""), dialect=dialect)
+            kind = "ctas"
+        else:
+            message = f"Expression 'SELECT INTO' has not been implemented yet for dialect: {dialect}"
+            raise exception.UnsupportedFeatureError(message=message)
+    else:
+        kind = statement.key.lower()
+
+    return statement, kind
+
+
+def _collect_writable_cte_queries(
+    parent_query: Q, parent_holder: QueryHolder, dialect: str, object_mapping: mappings.ObjectMapping
+):
+    """
+    Collect any writable (DML) CTEs and attach them as downstream holders to the parent query.
+
+    If this query is of the form:
+        WITH cte AS (
+            INSERT ... RETURNING ...
+        )
+        INSERT INTO ...
+
+    then the outer and inner queries form a parent-child relationship.
+    The inner query is left as-is and copied, while the outer query transformer its
+    inner query's SELECT columns with the RETURNING columns. This is so that
+    the lineage functions collect the right columns during expression traversal.
+    The two queries are processed independently later.
+    """
+    for i, cte in enumerate(parent_query.get_ctes()):
+        cte_expr = cte.this
+
+        if isinstance(cte_expr, exp.Select):
+            continue
+
+        query = _process_unnamed(cte_expr, dialect, object_mapping, i)
+        if not query:
+            logger.warning(f"Skipping unsupported query type in CTE: {type(cte_expr)}")
+            continue
+
+        # Detach the query in the AST so that certain transformations work later
+        downstream_holder = QueryHolder(original=query)
+        parent_holder.add_downstream_holder(downstream_holder)
+        _collect_query_children(query, downstream_holder, dialect, object_mapping)
+
+
+def _collect_call_substitutions(
+    query: Q,
+    holder: QueryHolder,
+    dialect: str,
+    object_mapping: mappings.ObjectMapping,
+) -> None:
+    """
+    Performs argument substitution for CallQuery, ExecuteQuery, CTASQuery (EXECUTE AS),
+    and any query containing UDF call sites.
+    The resulting inner statements are classified as Queries and attached as downstream holders.
+    """
+    subst_statements: t.List[exp.Expr] = []
+
+    if isinstance(query, CallQuery):
+        subst_statements = substitute.substitute_call(query=query)
+    elif isinstance(query, ExecuteQuery):
+        subst_statements = substitute.substitute_execute(query=query)
+    elif isinstance(query, CTASQuery) and query.source_info.type == typing.SqlObjectType.PREPARED_STATEMENT:
+        subst_statements = [substitute.substitute_create_execute(query=query)]
+    else:
+        expression = util.copy_expression(query.statement)
+        while True:
+            node, matched_udf = udf.find_next_udf_call(expression, query.object_mapping)
+            if not node:
+                break
+
+            inner_exprs = udf.substitute_udf(node, matched_udf)
+            subst_statements.extend(inner_exprs)
+            # Replace the processed call site with a placeholder so find_next_udf_call()
+            # no longer matches it on the next iteration, allowing the loop to advance
+            # to the next UDF call. The statement is a working copy and is not used
+            # after this function; only the extracted subst_statements matter.
+            node.replace(exp.Null())
+
+    parent_index = query.get_statement_index()
+    for i, stmt in enumerate(subst_statements):
+        composite_index = f"{parent_index}:{i}"
+        substituted_query = _process_unnamed(
+            statement=stmt,
+            dialect=dialect,
+            object_mapping=object_mapping,
+            statement_index=composite_index,
+        )
+        if substituted_query:
+            downstream_holder = QueryHolder(original=substituted_query)
+            holder.add_downstream_holder(downstream_holder)
+            # Recurse: the substituted child may itself contain calls/UDFs
+            _collect_query_children(substituted_query, downstream_holder, dialect, object_mapping)
+
+
+def _collect_query_children(query: Q, parent_holder: QueryHolder, dialect: str, object_mapping: mappings.ObjectMapping):
+    """
+    Collect any nested child queries for a given query and attach them to the holder.
+    """
+    if isinstance(query, InsertQuery):
+        _collect_insert_children(query, parent_holder, object_mapping)
+    elif isinstance(query, MergeQuery):
+        _collect_merge_children(query, parent_holder, object_mapping)
+    elif isinstance(query, MultitableInsertQuery):
+        _collect_multitable_insert_children(query, parent_holder, object_mapping)
+    elif isinstance(query, UserDefinedFunctionQuery):
+        _collect_udf_children(query, parent_holder, dialect, object_mapping)
+
+    # Always check for writable CTEs regardless of query type above
+    if not isinstance(query, (CopyQuery, PutQuery)):
+        _collect_writable_cte_queries(query, parent_holder, dialect, object_mapping)
+
+    # Perform call-site substitution immediately at collection time
+    _collect_call_substitutions(query, parent_holder, dialect, object_mapping)
+
+
+def _collect_udf_children(
+    query: UserDefinedFunctionQuery,
+    parent_holder: QueryHolder,
+    dialect: str,
+    object_mapping: mappings.ObjectMapping,
+):
+    """
+    Extract the function body from the UDF and parse it into individual statements.
+    Circular references in UDF/procedure definitions could cause infinite recursion.
+    """
+    body_expr = query.statement.args.get("expression")
+    if not body_expr:
+        return
+
+    for i, stmt in enumerate(util.iter_inner_statements(body_expr, dialect)):
+        child_query = _process_unnamed(stmt, dialect, object_mapping, i)
+
+        # If it's a simple select and _process_unnamed skipped it, create a SelectQuery
+        # TODO: VALUES and SELECT should be collected
+        if not child_query and isinstance(stmt, exp.Select):
+            child_query = SelectQuery(expr=stmt, dialect=dialect, object_mapping=object_mapping, statement_index=i)
+
+        if child_query:
+            downstream_holder = QueryHolder(original=child_query)
+            parent_holder.add_downstream_holder(downstream_holder)
+            _collect_query_children(child_query, downstream_holder, dialect, object_mapping)
+
+
+def _collect_insert_children(query: InsertQuery, parent_holder: QueryHolder, object_mapping: mappings.ObjectMapping):
+    """
+    Collect any additional queries inside an INSERT. For Postgres, this is 'INSERT .. ON CONFLICT DO UPDATE'.
+    """
+    if not isinstance(query.statement, exp.Insert):
+        return
+
+    on_conflict = query.statement.args.get("conflict")
+
+    if not isinstance(on_conflict, exp.OnConflict) or on_conflict.args["action"].name == "DO NOTHING":
+        return
+
+    update_query = UpdateQuery(
+        expr=on_conflict,
+        dialect=query.dialect,
+        object_mapping=object_mapping,
+        statement_index=0,
+        table=query.get_target_as_table(),
+    )
+    downstream_holder = QueryHolder(original=update_query)
+    parent_holder.add_downstream_holder(downstream_holder)
+
+
+def _collect_merge_children(
+    parent_query: MergeQuery, parent_holder: QueryHolder, object_mapping: mappings.ObjectMapping
+):
+    """
+    Transform any nested statements (INSERT or UPDATE) into fully qualified queries.
+
+    This is to allow the statements to be processed independently of the parent MERGE query.
+
+    For example, the merge query:
+
+        MERGE INTO fruit.processed AS t
+        USING fruit.raw AS s
+        ON t.kind = s.kind
+        WHEN MATCHED THEN
+            UPDATE SET name = s.name
+        WHEN NOT MATCHED THEN
+            INSERT (label) VALUES (s.kind);
+
+    has 2 nested queries that get transformed into:
+
+        UPDATE fruit.processed AS t
+        SET name = s.name
+        FROM fruit.raw AS t
+        WHERE t.kind = s.kind
+
+        INSERT INTO fruit.processed t
+        SELECT s.kind as label
+        FROM fruit.raw s;
+    """
+    merge = parent_query
+    parent_expr = parent_query.statement
+    whens = [when.args["then"] for when in parent_expr.args["whens"].expressions]
+
+    for i, when in enumerate(whens):
+        # Ensure the full expression tree is kept
+        when_expr = util.copy_expression(when)
+
+        if isinstance(when_expr, exp.Update):
+            update_query = UpdateQuery(
+                expr=when_expr,
+                dialect=parent_query.dialect,
+                object_mapping=object_mapping,
+                statement_index=i,
+                table=merge.get_target_as_table(),
+            )
+            downstream_holder = QueryHolder(original=update_query)
+            parent_holder.add_downstream_holder(downstream_holder)
+
+        elif isinstance(when_expr, exp.Insert):
+            insert_query = InsertQuery(
+                expr=when_expr,
+                dialect=parent_query.dialect,
+                object_mapping=object_mapping,
+                statement_index=i,
+                table=merge.get_target_as_table(),
+            )
+            insert_query.target_info = merge.target_info
+            downstream_holder = QueryHolder(original=insert_query)
+            parent_holder.add_downstream_holder(downstream_holder)
+
+
+def _collect_multitable_insert_children(
+    parent_query: MultitableInsertQuery, parent_holder: QueryHolder, object_mapping: mappings.ObjectMapping
+):
+    """
+    Extract the ConditionalInsert (exp.Insert) branches from the Multitable statement.
+    """
+    statement = parent_query.statement
+    for i, branch in enumerate(statement.expressions):
+        if not isinstance(branch, exp.ConditionalInsert):
+            continue
+
+        insert_expr = branch.this
+        if not isinstance(insert_expr, exp.Insert):
+            continue
+
+        copied_expr = util.copy_expression(insert_expr)
+
+        if not copied_expr.expression and not copied_expr.args.get("source"):
+            # If the branch doesn't have its own source, use the parent's shared source.
+            # This happens in Snowflake multi-table inserts without explicit VALUES.
+            copied_expr.set("expression", parent_query.source_info.expression)
+
+        child_query = InsertQuery(
+            expr=copied_expr,
+            dialect=parent_query.dialect,
+            object_mapping=object_mapping,
+            statement_index=i,
+        )
+        downstream_holder = QueryHolder(original=child_query)
+        parent_holder.add_downstream_holder(downstream_holder)
+
+
+_UNNAMED_TYPE_MAP: dict[type, type] = {
+    exp.Insert: InsertQuery,
+    exp.Update: UpdateQuery,
+    exp.Merge: MergeQuery,
+    exp.MultitableInserts: MultitableInsertQuery,
+    exp.Delete: DeleteQuery,
+    exp.Select: SelectQuery,
+    exp.Copy: CopyQuery,
+    exp.Values: ValuesQuery,
+    exp.Put: PutQuery,
+    exp.Set: SetQuery,
+}
+
+
+def _process_unnamed(
+    statement: exp.Expr, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int | str
+) -> Q | None:
+    """
+    Process an unnamed statement - one not inside a 'CREATE <name>' statement.
+    """
+    query_class = _UNNAMED_TYPE_MAP.get(type(statement))
+    if query_class is not None:
+        query = query_class(
+            expr=statement, dialect=dialect, object_mapping=object_mapping, statement_index=statement_index
+        )
+        if isinstance(query, DeleteQuery) and not statement.find(exp.Insert, exp.Update, exp.Merge):
+            logger.warning(
+                "Skipping statement: A DELETE query must have a data-modifying statement, "
+                "such as an INSERT, to contain lineage."
+            )
+        return query
+
+    if isinstance(statement, exp.Command) and statement.this.upper() == "CALL":
+        return CallQuery(
+            expr=statement, dialect=dialect, object_mapping=object_mapping, statement_index=statement_index
+        )
+    if isinstance(statement, exp.Command) and statement.this.upper() == "EXECUTE":
+        return ExecuteQuery(
+            expr=statement, dialect=dialect, object_mapping=object_mapping, statement_index=statement_index
+        )
+    if isinstance(statement, exp.Create):
+        if statement.kind == "TABLE":
+            if isinstance(statement.expression, (exp.Select, exp.Values)) or statement.find(exp.ExecuteAsProperty):
+                return _process_views_and_ctas(statement, dialect, object_mapping, statement_index)
+            else:
+                return _process_tables(statement, dialect, object_mapping, statement_index)
+        elif statement.kind == "VIEW":
+            return _process_views_and_ctas(statement, dialect, object_mapping, statement_index)
+
+    return None
+
+
+def _process_tables(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int | str
+) -> Q | None:
+    """
+    Process a 'CREATE TABLE' statement.
+    """
+    query = None
+
+    if statement.kind == "TABLE":
+        # CREATE TABLE ...
+        query = TableQuery(
+            expr=statement, dialect=dialect, object_mapping=object_mapping, statement_index=statement_index
+        )
+        query.set_column_defs()
+        object_mapping.add_table_query(
+            query=query,
+            column_mapping=query.get_column_names_with_types(include_system=True),
+        )
+    elif statement.kind == "SEQUENCE":
+        query = SequenceQuery(
+            expr=statement, dialect=dialect, object_mapping=object_mapping, statement_index=statement_index
+        )
+        object_mapping.add_sequence_query(query=query)
+
+    return query
+
+
+def _process_views_and_ctas(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int | str
+) -> Q:
+    """
+    Convert a series of `CREATE VIEW/TABLE AS ...` SQL DDL statements into sqlglot's MappingSchema
+    to extract the table and column details.
+    """
+    util.qualify_and_annotate(statement, dialect, object_mapping, remove_added_aliases=True)
+    col_defs = _determine_column_defs(statement, dialect, object_mapping)
+
+    if statement.kind == "VIEW":
+        # CREATE VIEW ...
+        query = ViewQuery(
+            expr=statement,
+            dialect=dialect,
+            object_mapping=object_mapping,
+            columns=col_defs,
+            statement_index=statement_index,
+        )
+    elif statement.kind == "TABLE":
+        # CREATE TABLE AS ...
+        query = CTASQuery(
+            expr=statement,
+            dialect=dialect,
+            object_mapping=object_mapping,
+            columns=col_defs,
+            statement_index=statement_index,
+        )
+        query.system_column_defs = settings.system_columns(dialect)
+    else:
+        raise exception.UnsupportedFeatureError(message=f"Unhandled situation for query: {statement.kind}")
+
+    object_mapping.add_table_query(
+        query=query,
+        column_mapping=query.get_column_names_with_types(include_system=True),
+    )
+    return query
+
+
+def _unnest_values_inside_select(statement: exp.Create, dialect: str):
+    """
+    Replace SELECT * FROM (VALUES ()) with VALUES ().
+    This prevents sqlglot from assigning its own aliases.
+    """
+    # TODO: this should be done in a transform, checked against self.statement
+    for values_expr in statement.find_all(exp.Values):
+        simplify_row_in_values(values_expr, dialect)
+        parent = values_expr.parent_select
+        while isinstance(parent, exp.Select) and parent.is_star and parent.parent_select:
+            parent = parent.parent_select
+
+        if parent and parent.parent:
+            # sqlglot converts "SELECT UNION VALUES" into "SELECT UNION SELECT * FROM (VALUES ()) AS _values"
+            # We undo that transformation (risky, as actual query may have that alias)
+            if alias := values_expr.args["alias"]:
+                if alias.name != "_values":
+                    return statement
+                alias.pop()
+            parent.replace(values_expr)
+
+    return statement
+
+
+def _determine_column_defs(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping
+) -> t.List[ColumnDef]:
+    """
+    Look up the columns for 'y' in 'INSERT INTO x TABLE y'
+    """
+    if isinstance(statement.expression, exp.Values):
+        # CREATE TABLE AS VALUES
+        columns = [stmt.name for stmt in statement.this.expressions]
+        types = [val.type for val in statement.expression.expressions[0].expressions]
+
+        if not columns:
+            columns = util.default_column_index_iterator(dialect, types)
+
+        col_defs = [exp.ColumnDef(this=col_name, kind=col_type) for col_name, col_type in zip(columns, types)]
+    elif exec_prop := statement.find(exp.ExecuteAsProperty):
+        # CREATE TABLE AS EXECUTE
+        plan_name = exec_prop.this.name
+        plan_table = exp.to_table(plan_name)
+        matched_prepare = object_mapping.lookup_prepare_query(plan_table, raise_on_missing=True)
+
+        source = matched_prepare.statement
+        if isinstance(source, exp.Select):
+            col_defs = [
+                exp.ColumnDef(this=exp.to_identifier(col.alias_or_name), kind=col.unalias().type)
+                for col in source.expressions
+            ]
+        else:
+            col_defs = []
+    else:
+        col_defs = [exp.ColumnDef(this=exp.to_identifier(s.alias), kind=s.type) for s in statement.selects]
+
+    return col_defs
+
+
+def _process_functions(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    """
+    Process a "CREATE FUNCTION" statement.
+    """
+    query = UserDefinedFunctionQuery(
+        expr=statement,
+        dialect=dialect,
+        object_mapping=object_mapping,
+        statement_index=statement_index,
+    )
+    object_mapping.add_udf_query(query, column_mapping=query.get_column_names_with_types(include_system=True))
+
+    return query
+
+
+def _process_triggers(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    """
+    Process a "CREATE TRIGGER" statement.
+    """
+    query = TriggerQuery(statement, dialect, object_mapping, statement_index)
+    object_mapping.add_trigger_query(query)
+    return query
+
+
+def _process_stored_procedures(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    """
+    Process a "CREATE PROCEDURE" statement.
+    """
+    query = ProcedureQuery(
+        expr=statement, dialect=dialect, object_mapping=object_mapping, statement_index=statement_index
+    )
+    object_mapping.add_procedure_query(query)
+    return query
+
+
+def _process_stage(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = StageQuery(statement, dialect, object_mapping=object_mapping, statement_index=statement_index)
+    object_mapping.add_stage_query(query)
+    return query
+
+
+def _process_unload(
+    statement: exp.Command, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = UnloadQuery(statement, dialect, object_mapping, statement_index)
+    return query
+
+
+def _process_replace(
+    statement: exp.Command, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = ReplaceQuery(statement, dialect, object_mapping, statement_index)
+    return query
+
+
+def _process_prepare(
+    statement: exp.Command, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = PrepareQuery(statement, dialect, object_mapping, statement_index)
+    object_mapping.add_prepare_query(query)
+    return query
+
+
+def _process_execute(
+    statement: exp.Command, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = ExecuteQuery(statement, dialect, object_mapping, statement_index)
+    return query
+
+
+def _process_type(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = TypeQuery(statement, dialect, object_mapping, statement_index)
+    object_mapping.add_type_query(query)
+    return query
+
+
+def _process_schema(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = SchemaQuery(statement, dialect, object_mapping, statement_index)
+    object_mapping.add_schema_query(query)
+    return query
+
+
+def _process_database(
+    statement: exp.Create, dialect: str, object_mapping: mappings.ObjectMapping, statement_index: int
+) -> Q:
+    query = DatabaseQuery(statement, dialect, object_mapping, statement_index)
+    object_mapping.add_database_query(query)
+    return query
+
+
+_QUERY_PROCESSORS: dict[str, t.Callable] = {
+    "table": _process_tables,
+    "ctas": _process_views_and_ctas,
+    "view": _process_views_and_ctas,
+    "sequence": _process_tables,
+    "procedure": _process_stored_procedures,
+    "multitableinserts": _process_unnamed,
+    "function": _process_functions,
+    "database": _process_database,
+    "trigger": _process_triggers,
+    "select": _process_unnamed,
+    "insert": _process_unnamed,
+    "update": _process_unnamed,
+    "merge": _process_unnamed,
+    "delete": _process_unnamed,
+    "schema": _process_schema,
+    "unload": _process_unload,
+    "replace": _process_replace,
+    "prepare": _process_prepare,
+    "execute": _process_execute,
+    "stage": _process_stage,
+    "call": _process_unnamed,
+    "copy": _process_unnamed,
+    "put": _process_unnamed,
+    "set": _process_unnamed,
+    "values": _process_unnamed,
+    "type": _process_type,
+}
+
+
+def _split_combined_statements(parsed: t.List[exp.Expr]) -> t.List[exp.Expr]:
+    """
+    Workaround for sqlglot incorrectly bundling statements following a block-based CREATE.
+    For example, with the statement "CREATE PROCEDURE .. BEGIN .. END; CALL();" the CALL()
+    is included with the procedure definition.
+    """
+    new_parsed = []
+    for stmt in parsed:
+        if isinstance(stmt, exp.Create) and isinstance(stmt.args.get("expression"), exp.Block):
+            block = stmt.args["expression"]
+            end_index = -1
+            for i, expr in enumerate(block.expressions):
+                if isinstance(expr, exp.EndStatement):
+                    end_index = i
+                    break
+
+            # Separate the statements
+            if end_index != -1 and end_index < len(block.expressions) - 1:
+                extra = block.expressions[end_index + 1 :]
+                block.set("expressions", block.expressions[: end_index + 1])
+                new_parsed.append(stmt)
+                new_parsed.extend(extra)
+                continue
+
+        new_parsed.append(stmt)
+    return new_parsed