feat: support JSON types, expression language

sqlalchemy_bigquery/__init__.py

@@ -37,6 +37,7 @@
     FLOAT64,
     INT64,
     INTEGER,
+    JSON,
     NUMERIC,
     RECORD,
     STRING,
@@ -74,6 +75,7 @@
     "FLOAT64",
     "INT64",
     "INTEGER",
+    "JSON",
     "NUMERIC",
     "RECORD",
     "STRING",

sqlalchemy_bigquery/_json.py

@@ -0,0 +1,106 @@
+from enum import auto, Enum
+import sqlalchemy
+from sqlalchemy.sql import sqltypes
+
+
+class _FormatTypeMixin:
+    def _format_value(self, value):
+        raise NotImplementedError()
+
+    def bind_processor(self, dialect):
+        super_proc = self.string_bind_processor(dialect)
+
+        def process(value):
+            value = self._format_value(value)
+            if super_proc:
+                value = super_proc(value)
+            return value
+
+        return process
+
+    def literal_processor(self, dialect):
+        super_proc = self.string_literal_processor(dialect)
+
+        def process(value):
+            value = self._format_value(value)
+            if super_proc:
+                value = super_proc(value)
+            return value
+
+        return process
+
+
+class JSON(sqltypes.JSON):
+    def bind_expression(self, bindvalue):
+        # JSON query parameters are STRINGs
+        return sqlalchemy.func.PARSE_JSON(bindvalue, type_=self)
+
+    def literal_processor(self, dialect):
+        super_proc = self.bind_processor(dialect)
+
+        def process(value):
+            value = super_proc(value)
+            return repr(value)
+
+        return process
+
+    class Comparator(sqltypes.JSON.Comparator):
+        def _generate_converter(self, name, lax):
+            prefix = "LAX_" if lax else ""
+            func_ = getattr(sqlalchemy.func, f"{prefix}{name}")
+            return func_
+
+        def as_boolean(self, lax=False):
+            func_ = self._generate_converter("BOOL", lax)
+            return func_(self.expr, type_=sqltypes.Boolean)
+
+        def as_string(self, lax=False):
+            func_ = self._generate_converter("STRING", lax)
+            return func_(self.expr, type_=sqltypes.String)
+
+        def as_integer(self, lax=False):
+            func_ = self._generate_converter("INT64", lax)
+            return func_(self.expr, type_=sqltypes.Integer)
+
+        def as_float(self, lax=False):
+            func_ = self._generate_converter("FLOAT64", lax)
+            return func_(self.expr, type_=sqltypes.Float)
+
+        def as_numeric(self, precision, scale, asdecimal=True):
+            # No converter available in BigQuery
+            raise NotImplementedError()
+
+    comparator_factory = Comparator
+
+    class JSONPathMode(Enum):
+        LAX = auto()
+        LAX_RECURSIVE = auto()
+
+
+class JSONPathType(_FormatTypeMixin, sqltypes.JSON.JSONPathType):
+    def _mode_prefix(self, mode):
+        if mode == JSON.JSONPathMode.LAX:
+            mode_prefix = "lax"
+        elif mode == JSON.JSONPathMode.LAX_RECURSIVE:
+            mode_prefix = "lax recursive"
+        else:
+            raise NotImplementedError(f"Unhandled JSONPathMode: {mode}")
+        return mode_prefix
+
+    def _format_value(self, value):
+        if isinstance(value[0], JSON.JSONPathMode):
+            mode = value[0]
+            mode_prefix = self._mode_prefix(mode)
+            value = value[1:]
+        else:
+            mode_prefix = ""
+
+        return "%s$%s" % (
+            mode_prefix + " " if mode_prefix else "",
+            "".join(
+                [
+                    "[%s]" % elem if isinstance(elem, int) else '."%s"' % elem
+                    for elem in value
+                ]
+            ),
+        )

sqlalchemy_bigquery/_types.py

@@ -27,6 +27,7 @@
 except ImportError:  # pragma: NO COVER
     pass
 
+from ._json import JSON
 from ._struct import STRUCT
 
 _type_map = {
@@ -41,6 +42,7 @@
     "FLOAT": sqlalchemy.types.Float,
     "INT64": sqlalchemy.types.Integer,
     "INTEGER": sqlalchemy.types.Integer,
+    "JSON": JSON,
     "NUMERIC": sqlalchemy.types.Numeric,
     "RECORD": STRUCT,
     "STRING": sqlalchemy.types.String,
@@ -61,6 +63,7 @@
 FLOAT = _type_map["FLOAT"]
 INT64 = _type_map["INT64"]
 INTEGER = _type_map["INTEGER"]
+JSON = _type_map["JSON"]
 NUMERIC = _type_map["NUMERIC"]
 RECORD = _type_map["RECORD"]
 STRING = _type_map["STRING"]

sqlalchemy_bigquery/base.py

@@ -59,7 +59,7 @@
 import re
 
 from .parse_url import parse_url
-from . import _helpers, _struct, _types
+from . import _helpers, _json, _struct, _types
 import sqlalchemy_bigquery_vendored.sqlalchemy.postgresql.base as vendored_postgresql
 
 # Illegal characters is intended to be all characters that are not explicitly
@@ -547,6 +547,13 @@ def visit_bindparam(
         bq_type = self.dialect.type_compiler.process(type_)
         bq_type = self.__remove_type_parameter(bq_type)
 
+        if bq_type == "JSON":
+            # FIXME: JSON is not a member of `SqlParameterScalarTypes` in the DBAPI
+            # For now, we hack around this by:
+            # - Rewriting the bindparam type to STRING
+            # - Applying a bind expression that converts the parameter back to JSON
+            bq_type = "STRING"
+
         assert_(param != "%s", f"Unexpected param: {param}")
 
         if bindparam.expanding:  # pragma: NO COVER
@@ -571,6 +578,12 @@ def visit_getitem_binary(self, binary, operator_, **kw):
         right = self.process(binary.right, **kw)
         return f"{left}[OFFSET({right})]"
 
+    def visit_json_path_getitem_op_binary(self, binary, operator, **kw):
+        return "JSON_QUERY(%s, %s)" % (
+            self.process(binary.left, **kw),
+            self.process(binary.right, **kw),
+        )
+
     def _get_regexp_args(self, binary, kw):
         string = self.process(binary.left, **kw)
         pattern = self.process(binary.right, **kw)
@@ -641,6 +654,12 @@ def visit_NUMERIC(self, type_, **kw):
 
     visit_DECIMAL = visit_NUMERIC
 
+    def visit_JSON(self, type_, **kw):
+        return "JSON"
+
+    def visit_json_path(self, type_, **kw):
+        return "STRING"
+
 
 class BigQueryDDLCompiler(DDLCompiler):
     option_datatype_mapping = {
@@ -1076,6 +1095,8 @@ class BigQueryDialect(DefaultDialect):
         sqlalchemy.sql.sqltypes.TIMESTAMP: BQTimestamp,
         sqlalchemy.sql.sqltypes.ARRAY: BQArray,
         sqlalchemy.sql.sqltypes.Enum: sqlalchemy.sql.sqltypes.Enum,
+        sqlalchemy.sql.sqltypes.JSON: _json.JSON,
+        sqlalchemy.sql.sqltypes.JSON.JSONPathType: _json.JSONPathType,
     }
 
     def __init__(
@@ -1086,6 +1107,8 @@ def __init__(
         credentials_info=None,
         credentials_base64=None,
         list_tables_page_size=1000,
+        json_serializer=None,
+        json_deserializer=None,
         *args,
         **kwargs,
     ):
@@ -1098,6 +1121,8 @@ def __init__(
         self.identifier_preparer = self.preparer(self)
         self.dataset_id = None
         self.list_tables_page_size = list_tables_page_size
+        self._json_serializer = json_serializer
+        self._json_deserializer = json_deserializer
 
     @classmethod
     def dbapi(cls):