Find an expression at the top level

[avin-kavish]

WITH
  sales_by_person AS (
    SELECT
      p.firstname || ' ' || p.lastname AS person_name,
      SUM(sod.orderqty) AS total_sales
    FROM
      sales.salesorderdetail sod
      JOIN sales.salesorderheader soh ON sod.salesorderid = soh.salesorderid
      JOIN sales.customer c ON soh.customerid = c.customerid
      JOIN person.person p ON c.personid = p.businessentityid
    GROUP BY
      p.firstname,
      p.lastname
    ORDER BY
      total_sales DESC
    LIMIT
      10
  )
SELECT
  person_name,
  total_sales
FROM
  sales_by_person;

Basically, given a query like this, I want to check if there's a top level limit statement (not in subqueries/ctes).

I tried,

parsed_query = sqlglot.parse_one(query)
print(parsed_query.find(exp.Limit))

But this is descending the tree and searching at depths > 1, so it returns limit 10 as a match.

How do I check for only top level constructs?

[georgesittas]

If you know that query will always be a SELECT expression (i.e. no DDLs), you can do what you want in a simpler way:

parsed_query = sqlglot.parse_one(query)
print(parsed_query.args.get("limit") or "The query has no LIMIT clause")

[avin-kavish]

Thanks. This is actually part of a larger validation routine,

def _validate_query(sql_query):
    try:
        parsed_query, *extra = sqlglot.parse(sql_query)
    except sqlglot.errors.ParseError as e:
        raise Exception from e
    if not parsed_query:
        raise Exception
    if len(extra):
        raise Exception

    if (
        parsed_query.find(exp.Insert)
        or parsed_query.find(exp.Update)
        or parsed_query.find(exp.Delete)
    ):
        raise Exception("Contains prohibited ops")

    if parsed_query.find(exp.Into):
        raise Exception("Selecting INTO new table")

I'm trying to ensure that the passed sql query is a single statement that selects some rows. So, no DDL, mutations, grant/revoke etc.

So with the initial approach I took above I'd have to check for all invalid statements one by one. So now I'm thinking of changing it to just check for a select statement. It also has to not be a select that's creating a new view/table or not part of a update etc.

Do you have any suggestions on how to implement this?

[georgesittas]

Yeah, you can do an instance check to see if the parsed expression is an exp.Select.