Introducing count_distinct() function

[borsitsb]

Hey,

I work in Rollbar, and we created a modified SQL language called RQL, so our users can access their data easily. I analyze the queries to find possible ways to improve the system. I found the sqlglot library an elegant way to complete this task, but I need some help with the followings:

1. We introduced a function called count_distinct(exp) in RQL, but sqlglot doesn't know it, obviously. It identifies the function as ANONYMOUS. How can I "tell" sqlglot that it's a valid function, basically count(distinct foo)? Is there an easy way to do that?
2. Do I need to introduce a new dialect to correctly parse RQL queries? I don't want to transpile, it wouldn't make much sense. I do parsing only.
   Thanks for your help in advance! :)

Cheers,
Botond

[barakalon]

2 options:

(1) No custom dialect

Anonymous expressions are "valid". You can read their name:

import sqlglot
from sqlglot import expressions as exp

sqlglot.parse_one("SELECT count_distinct(x)").find(exp.Anonymous).name
# 'count_distinct'

(2) Custom dialect

If you want to parse it as a normal "COUNT(DISTINCT ...)" expression:

import sqlglot
from sqlglot import expressions as exp
from sqlglot.dialects.postgres import Postgres

def _count_distinct(args):
    return exp.Count.from_arg_list([exp.Distinct(this=args[0])])

class RQL(Postgres):  # If RQL is based on another dialect, subclass it
    class Parser(Postgres.Parser):
        FUNCTIONS = {
            **Postgres.Parser.FUNCTIONS,
            "COUNT_DISTINCT": _count_distinct
        }

sqlglot.parse_one("SELECT count_distinct(x)", read="rql")
# (SELECT expressions: 
#   (COUNT this: 
#     (DISTINCT this: 
#       (COLUMN this: 
#         (IDENTIFIER this: x, quoted: False)))))

You can also parse it into a custom function class (e.g. class CountDistinct), if you prefer that. But that might be similar to just using Anonymous.name, depending on what you're doing.