Map JSON to Tokens to SQL

[chriskathumbi]

Hello,

Does sqlglot offer a way to map JSON to Tokens and the SQL?
I've been trying to work on something and I'd like to support SQL to JSON and back.
Is there a better way to do this?
Sorry for the long code block ...

from enum import StrEnum
from datetime import date, datetime
from typing import Annotated, Literal, Union, Optional

from pydantic import (
    BaseModel,
    BeforeValidator,
    field_validator,
    Field,
    StrictBool,
    StrictFloat,
    StrictInt,
)


class LogicalOperator(StrEnum):
    AND = "and"
    OR = "or"


class Operator(StrEnum):
    EQ = "="
    NEQ = "!="
    GT = ">"
    GTE = ">="
    LT = "<"
    LTE = "<="


class IsOperator(StrEnum):
    IS = "is"
    IS_NOT = "is not"


class LikeOperator(StrEnum):
    CONTAINS = "contains"
    BEGINS_WITH = "begins with"
    ENDS_WITH = "ends with"

    def __str__(self):
        return "ilike"


class BetweenOperator(StrEnum):
    BETWEEN = "between"
    NOT_BETWEEN = "not between"


class InOperator(StrEnum):
    IN = "in"
    NOT_IN = "not in"


class FilterType(StrEnum):
    RULE = "rule"
    GROUP = "group"


# DateType = Union[date, datetime]
DateType = Annotated[
    Union[date, datetime],
    BeforeValidator(lambda x: datetime.strptime(str(x), "%Y-%m-%d").date()),
]
NumericType = StrictFloat | StrictInt
BetweenValueType = DateType | NumericType | str
InValueType = StrictBool | BetweenValueType


def check_list_type(arr: list):
    arr_types = {type(i) for i in arr}
    if len(arr_types) > 1:
        raise ValueError("elements must be of the same type")
    return arr


class BaseQuery(BaseModel):
    table: str
    column: str
    negate: Optional[StrictBool] = False
    include_nulls: Optional[StrictBool] = False

    def __str__(self) -> str:
        column_ref = f"{self.table}.{self.column}"
        query = f"{column_ref}" + " {operator} {value}"

        if self.include_nulls:
            query = f"({query} or {column_ref} is null)"
        if self.negate:
            query = f"not {query}"

        return query


class BetweenQuery(BaseQuery):
    operator: BetweenOperator
    value: list[BetweenValueType] = Field(..., min_length=2, max_length=2)

    _value = field_validator("value")(check_list_type)

    def __str__(self) -> str:
        if isinstance(self.value[0], (int, float)):
            value = f"{self.value[0]} and {self.value[1]}"
        elif isinstance(self.value[0], (date, str)):
            value = f"'{self.value[0]}' and '{self.value[1]}'"

        return super().__str__().format(operator=self.operator, value=value)


class InQuery(BaseQuery):
    operator: InOperator
    value: list[InValueType] = Field(..., max_length=20)

    _value = field_validator("value")(check_list_type)

    def __str__(self) -> str:
        value = ""
        if isinstance(self.value[0], (int, float, bool)):
            value += ", ".join(str(i) for i in self.value)
        elif isinstance(self.value[0], (date, str)):
            value += ", ".join(f"'{i}'" for i in self.value)

        value = f"({value})"
        return super().__str__().format(operator=self.operator.value, value=value)


class Query(BaseQuery):
    operator: Operator
    value: StrictBool | DateType | NumericType | str

    def __str__(self) -> str:
        if isinstance(self.value, (int, float, bool)):
            value = self.value
        elif isinstance(self.value, (date, str)):
            value = f"'{self.value}'"

        return super().__str__().format(operator=self.operator, value=value)


class LikeQuery(BaseQuery):
    operator: LikeOperator
    value: str

    def __str__(self) -> str:
        if self.operator == LikeOperator.BEGINS_WITH:
            value = f"'{self.value}%'"
        elif self.operator == LikeOperator.ENDS_WITH:
            value = f"'%{self.value}'"
        else:
            value = f"'%{self.value}%'"

        return super().__str__().format(operator=self.operator, value=value)


class IntervalUnit(StrEnum):
    HOUR = "hour"
    DAY = "day"
    WEEK = "week"
    MONTH = "month"
    YEAR = "year"


class IntervalValue(BaseModel):
    value: NumericType
    unit: IntervalUnit
    before: bool = True

    def __str__(self) -> str:
        return f"interval '{self.value} {self.unit}'"


class IntervalQuery(BaseQuery):
    operator: Operator
    value: IntervalValue

    def __str__(self) -> str:
        if self.value.before:
            value = f"current_date - {self.value}"
        else:
            value = f"current_date + {self.value}"

        return super().__str__().format(operator=self.operator, value=value)


class NullQuery(BaseQuery):
    table: str
    column: str
    operator: IsOperator
    value: Literal["NULL"]

    def __str__(self):
        return super().__str__().format(operator=self.operator, value=self.value)


QueryType = Union[InQuery, BetweenQuery, Query, LikeQuery, IntervalQuery, NullQuery]


class Rule(BaseModel):
    type: FilterType
    query: QueryType

    def __str__(self) -> str:
        return f"{self.query}"


class RuleGroup(BaseModel):
    type: FilterType
    logicalOperator: LogicalOperator
    group: list[Union[Rule, "RuleGroup"]]

    def __str__(self) -> str:
        group = [f"{g}" for g in self.group]
        return "(" + f" {self.logicalOperator} ".join(group) + ")"


class OrderBy(BaseModel):
    table: str
    column: str
    order: Optional[Literal["asc", "desc"]] = "asc"
    nulls: Optional[Literal["first", "last"]] = None

    def __str__(self):
        order_by = f"{self.table}.{self.column} {self.order}"
        if self.nulls is not None:
            order_by += f" nulls {self.nulls}"
        return order_by


class OrderByGroup(BaseModel):
    order_by: list[OrderBy]

    def __str__(self) -> str:
        order_by = [str(i) for i in self.order_by]
        return "order by " + ", ".join(order_by)


class LimitOffset(BaseModel):
    limit: int
    offset: Optional[int] = None

    def __str__(self) -> str:
        limit = f"limit {self.limit}"
        if self.offset is not None:
            limit += f" offset {self.offset}"
        return limit


if __name__ == "__main__":
    ...

[georgesittas]

I've been trying to work on something and I'd like to support SQL to JSON and back.

Check out the dump method on the AST. If this isn't enough, please share a smaller snippet (few lines) that demonstrates what you're after.