AST Manipulation

If there are inaccuracies discovered with this documentation, please submit a GitHub issue.

The partiql-ast package provides a set of classes and interfaces to represent PartiQL statements. Performing some logic on a query or performing some tree-based rewrite can be performed using a couple partiql-ast APIs:

  • AstVisitor

  • AstRewriter which implements the AstVisitor interface and allows for easier rewriting of the syntax tree

These two APIs follow the visitor pattern.

A common use case could be to validate the nodes of the AST to allow (e.g. only permit SELECT-FROM-WHERE queries and reject non-SFW statements). This can be done by creating a class that implements the AstVisitor interface and checking that an AST Statement is a Query:

import org.partiql.ast.AstNode
import org.partiql.ast.AstVisitor
import org.partiql.ast.Query
import org.partiql.ast.Statement

/**
 * Simple AstVisitor with return type `R` of type `Boolean` and context type `C` of Unit.
 *
 * Visitor function returning `true` denotes a valid AST. Visitor function returning `false` denotes an invalid AST node.
 */
class DenyNonSFWVisitor : AstVisitor<Boolean, Unit>() {
    // `defaultReturn` must be specified for `AstVisitor` implementations
    override fun defaultReturn(node: AstNode, ctx: Unit): Boolean {
        // default of true means we allow all nodes by default
        return true
    }

    override fun visitStatement(node: Statement, ctx: Unit?): Boolean {
        return when (node) {
            is Query -> true
            else -> false // disallow non-Query statements
        }
    }
}

Alternatively, we could have written the same AST logic using per-function overrides:

import org.partiql.ast.AstNode
import org.partiql.ast.AstVisitor
import org.partiql.ast.ddl.CreateTable
import org.partiql.ast.dml.Delete
import org.partiql.ast.dml.Insert
import org.partiql.ast.dml.Replace
import org.partiql.ast.dml.Update
import org.partiql.ast.dml.Upsert

/**
 * Simple AstVisitor with return type `R` of type `Boolean` and context type `C` of Unit.
 *
 * Visitor function returning `true` denotes a valid AST. Visitor function returning `false` denotes an invalid AST node.
 */
public class DenyNonSFWVisitor : AstVisitor<Boolean, Unit>() {
    override fun defaultReturn(node: AstNode, ctx: Unit): Boolean {
        return true
    }

    override fun visitCreateTable(node: CreateTable, ctx: Unit): Boolean {
        return false
    }

    override fun visitDelete(node: Delete, ctx: Unit): Boolean {
        return false
    }

    override fun visitInsert(node: Insert, ctx: Unit): Boolean {
        return false
    }

    override fun visitReplace(node: Replace, ctx: Unit): Boolean {
        return false
    }

    override fun visitUpdate(node: Update, ctx: Unit): Boolean {
        return false
    }

    override fun visitUpsert(node: Upsert, ctx: Unit): Boolean {
        return false
    }
}

This visitor can be invoked by calling the .visit function with the node

// assume `sfw` is some `AstNode` that is an SFW query
assertTrue(DenyNonSFWVisitor().visit(sfw, Unit))

A similar use-case could be to deny all nodes unless explictly overridden. For this case, the defaultReturn would return false by default and nodes we allow will include explicit overrides:

class AllowVisitor : AstVisitor<Boolean, Unit>() {
    override fun defaultReturn(node: AstNode, ctx: Unit): Boolean {
        // default of false means we disallow all nodes by default
        return false
    }

    // Explicitly allow `Query`
    override fun visitQuery(node: Query?, ctx: Unit?): Boolean {
        return true
    }

    // Explicitly allow `Expr.QueryBodySFW`
    override fun visitQueryBodySFW(node: QueryBody.SFW?, ctx: Unit?): Boolean {
        return true
    }

    // Explicitly allow other expressions and clauses ...
}

AstRewriter provides a simpler interface to rewrite an AstNode. For example, the following rewrite to upper-case every identifier can be done by overriding the Identifier.Simple visit function:

/**
 * Rewriter extends AstVisitor with a default return type of [AstNode]. A context type can still be provided.
 */
public class BasicRewriter : AstRewriter<Unit>() {
    override fun visitIdentifierSimple(node: Identifier.Simple, ctx: Unit): AstNode {
        return Identifier.Simple(node.text.uppercase(), node.isRegular)
    }
}

The above examples did not make use of the context argument to the AstVisitor and AstRewriter. This argument can be used to pass state between the different visit* functions. We have some existing examples in our code base making use of the context: