fix(sqlalchemy-bigquery): escape JSON subfields

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/_struct.py

@@ -91,6 +91,11 @@ def _setup_getitem(self, name):
             subtype = self.expr.type._STRUCT_byname.get(name.lower())
             if subtype is None:
                 raise KeyError(name)
+            # Reuse `operators.json_getitem_op` to represent STRUCT subfield
+            # access. `BigQueryCompiler.visit_json_getitem_op_binary` in
+            # `base.py` checks whether the left operand is a `STRUCT` (vs. a
+            # `sqlalchemy.JSON` column) and compiles STRUCT subfields as
+            # backtick-quoted SQL identifiers (`struct_col.`field``).
             operator = operators.json_getitem_op
             index = _field_index(self, name, operator)
             return operator, index, subtype
@@ -112,9 +117,3 @@ def _field_index(self, name, operator):
         operator=operator,
         bindparam_type=sqlalchemy.types.String(),
     )
-
-
-# class SQLCompiler:
-#     def visit_json_getitem_op_binary(self, binary, operator_, **kw):
-#         left = self.process(binary.left, **kw)
-#         return f"{left}.{binary.right.value}"

packages/sqlalchemy-bigquery/sqlalchemy_bigquery/base.py

@@ -59,7 +59,7 @@
 from sqlalchemy.sql.selectable import CTE
 from sqlalchemy.sql.sqltypes import Integer, NullType, Numeric, String
 
-from . import _helpers, _types
+from . import _helpers, _struct, _types
 from .parse_url import parse_url
 
 # Illegal characters is intended to be all characters that are not explicitly
@@ -85,8 +85,22 @@ def __init__(self, dialect):
         super(BigQueryIdentifierPreparer, self).__init__(
             dialect,
             initial_quote="`",
+            escape_quote="`",
         )
 
+    def _escape_identifier(self, value):
+        # SQLAlchemy's base IdentifierPreparer defaults to escaping quotes by
+        # doubling them (e.g. `""`), whereas BigQuery GoogleSQL escapes
+        # backslashes and backticks inside backtick-quoted identifiers using a
+        # leading backslash (`\\` and ``\` ``).
+        value = value.replace("\\", "\\\\").replace("`", "\\`")
+        if self._double_percents:
+            value = value.replace("%", "%%")
+        return value
+
+    def _unescape_identifier(self, value):
+        return value.replace("\\`", "`").replace("\\\\", "\\")
+
     def quote_column(self, value):
         """
         Quote a column.
@@ -216,9 +230,114 @@ def visit_insert(self, insert_stmt, asfrom=False, **kw):
             insert_stmt, asfrom=False, **kw
         )
 
+    @staticmethod
+    def _get_struct_type(type_):
+        # Unwrap TypeDecorator.impl so custom types wrapping STRUCT are still
+        # recognized as STRUCT columns during compilation.
+        while type_ is not None:
+            if isinstance(type_, _struct.STRUCT):
+                return type_
+            type_ = getattr(type_, "impl", None)
+        return None
+
     def visit_json_getitem_op_binary(self, binary, operator_, **kw):
+        # Both `STRUCT` subfield access (see `_struct.py`) and `sqlalchemy.JSON`
+        # subscripting (`json_col[key]`) use `operators.json_getitem_op`, but
+        # they compile to completely different BigQuery constructs:
+        #
+        # 1. For `STRUCT`: subfields are SQL identifiers (`struct_col.`field``),
+        #    which cannot be query parameters in BigQuery. We verify that the
+        #    requested field exists in the static STRUCT schema (in case
+        #    `json_getitem_op` was constructed directly, bypassing
+        #    `STRUCT.Comparator._setup_getitem`) and backtick-quote the
+        #    identifier via `quote_identifier`.
+        struct_type = self._get_struct_type(binary.left.type)
+        if struct_type is not None:
+            field_name = getattr(binary.right, "value", None)
+            if (
+                not isinstance(field_name, str)
+                or field_name.lower() not in struct_type._STRUCT_byname
+            ):
+                raise sqlalchemy.exc.CompileError(
+                    f"Invalid STRUCT field name: {field_name!r}"
+                )
+            left = self.process(binary.left, **kw)
+            return f"{left}.{self.preparer.quote_identifier(field_name)}"
+
+        # 2. For `JSON`: if a sequence (tuple/list) or `JSONPathType` parameter
+        #    was passed to `json_getitem_op`, delegate to
+        #    `visit_json_path_getitem_op_binary` so it compiles as a multi-step
+        #    `JSON_QUERY` path lookup.
+        if isinstance(binary.right, sqlalchemy.sql.elements.BindParameter) and (
+            _is_json_path_sequence(binary.right.value)
+            or isinstance(
+                sqlalchemy.sql.type_api.to_instance(binary.right.type),
+                sqlalchemy.sql.sqltypes.JSON.JSONPathType,
+            )
+        ):
+            return self.visit_json_path_getitem_op_binary(binary, operator_, **kw)
+
+        # Otherwise, compile single-step `JSON` subscripting using BigQuery's
+        # native JSON subscript operator `{left}[{right}]` rather than
+        # `JSON_VALUE` or unquoted dot access:
+        # - Unlike raw dot interpolation (`{left}.{right}`), `{left}[{right}]`
+        #   accepts query parameters (`%(param:STRING)s` / `%(param:INT64)s`),
+        #   preventing SQL injection and supporting integer array indices.
+        # - Unlike `JSON_VALUE` (which extracts a scalar `STRING` and returns
+        #   `NULL` for nested objects/arrays), `{left}[{right}]` returns `JSON`,
+        #   preserving the return type and chained subscripts (`col["a"][0]["b"]`).
+        left = self.process(binary.left, **kw)
+        right = binary.right
+        if isinstance(right, sqlalchemy.sql.elements.BindParameter):
+            if isinstance(right.value, bool):
+                # In Python `isinstance(True, int)` is True, so SQLAlchemy's
+                # `JSON.Comparator` assigns `JSONIntIndexType` to `col[True]`.
+                # BigQuery rejects `BOOL` as a JSON subscript, and JSON object
+                # keys are strings, so coerce boolean subscripts to string keys.
+                right = right._clone(maintain_key=True)
+                right.value = str(right.value)
+                right.type = BQJSONStrIndexType()
+            elif isinstance(right.type, type):
+                # When subscripting with an untyped `sa.bindparam("k")`,
+                # SQLAlchemy's `coercions.expect` sets `bindparam.type` to the
+                # `JSONStrIndexType` class rather than an instance. Instantiate
+                # it on a clone so `type_compiler` can dispatch it.
+                right = right._clone(maintain_key=True)
+                right.type = sqlalchemy.sql.type_api.to_instance(right.type)
+            elif isinstance(right.type, NullType):
+                # When `json_getitem_op` is constructed directly with an untyped
+                # BindParameter, infer the index type so `visit_bindparam` can
+                # emit the required `:INT64` or `:STRING` BigQuery type suffix.
+                right = right._clone(maintain_key=True)
+                right.type = (
+                    BQJSONIntIndexType()
+                    if isinstance(right.value, int)
+                    else BQJSONStrIndexType()
+                )
+        right_text = self.process(right, **kw)
+        return f"{left}[{right_text}]"
+
+    def visit_json_path_getitem_op_binary(self, binary, operator_, **kw):
+        # BigQuery's bracket subscript `json_col[...]` only steps one key or
+        # array offset at a time. Multi-step path tuples (`col[("a", 0, "b")]`)
+        # use `JSON_QUERY(json_col, json_path)`, which also returns a `JSON`
+        # value (consistent with single-step subscripting) and accepts a
+        # parameterized JSONPath string formatted by `BQJSONPathType`.
         left = self.process(binary.left, **kw)
-        return f"{left}.{binary.right.value}"
+        right = binary.right
+        if isinstance(right, sqlalchemy.sql.elements.BindParameter):
+            right_type = sqlalchemy.sql.type_api.to_instance(right.type)
+            if not isinstance(right_type, sqlalchemy.sql.sqltypes.JSON.JSONPathType):
+                right_type = (
+                    BQJSONPathType()
+                    if _is_json_path_sequence(right.value)
+                    else BQString()
+                )
+            if right_type is not right.type:
+                right = right._clone(maintain_key=True)
+                right.type = right_type
+        right_text = self.process(right, **kw)
+        return f"JSON_QUERY({left}, {right_text})"
 
     def visit_table_valued_alias(self, element, **kw):
         # When using table-valued functions, like UNNEST, BigQuery requires a
@@ -533,6 +652,16 @@ def visit_bindparam(
         if literal_binds or isinstance(type_, NullType):
             return param
 
+        # SQLAlchemy's generic `JSONElementType` subclasses (`JSONIndexType`,
+        # `JSONIntIndexType`, `JSONStrIndexType`, `JSONPathType`) do not define
+        # `__visit_name__`. Because `bindparam.type` is not automatically
+        # adapted to the dialect before `visit_bindparam`, we resolve its
+        # dialect implementation via `BigQueryDialect.colspecs` (`BQJSON*Type`)
+        # so `self.dialect.type_compiler.process(type_)` can emit `:STRING` or
+        # `:INT64`.
+        if isinstance(type_, sqlalchemy.sql.sqltypes.JSON.JSONElementType):
+            type_ = type_.dialect_impl(self.dialect)
+
         if (
             isinstance(type_, Numeric)
             and (type_.precision is None or type_.scale is None)
@@ -637,6 +766,15 @@ def visit_BINARY(self, type_, **kw):
     def visit_JSON(self, type_, **kw):
         return "JSON"
 
+    def visit_json_int_index(self, type_, **kw):
+        return "INT64"
+
+    def visit_json_str_index(self, type_, **kw):
+        return "STRING"
+
+    visit_json_index = visit_json_str_index
+    visit_json_path = visit_json_str_index
+
     def visit_NUMERIC(self, type_, **kw):
         if (type_.precision is not None) and isinstance(
             kw.get("type_expression"), Column
@@ -1007,11 +1145,78 @@ def process_string_literal(value):
     return repr(value.replace("%", "%%"))
 
 
+def _escape_json_path_key(key: str) -> str:
+    # Always quote member keys in BigQuery JSONPath (`$."key"`) and escape
+    # backslashes and double quotes so keys containing dots, brackets, quotes,
+    # spaces, or `$` cannot alter the JSONPath structure.
+    escaped = key.replace("\\", "\\\\").replace('"', '\\"')
+    return f'."{escaped}"'
+
+
+def _format_json_path_step(step: Any) -> str:
+    # Check `not isinstance(step, bool)` because `bool` is a subclass of `int`
+    # in Python, and boolean keys in a path tuple should be treated as string
+    # member names rather than `[1]` / `[0]` array offsets.
+    if isinstance(step, int) and not isinstance(step, bool):
+        return f"[{step}]"
+    return _escape_json_path_key(str(step))
+
+
+def _is_json_path_sequence(value: Any) -> bool:
+    return isinstance(value, (tuple, list))
+
+
 class BQString(String):
     def literal_processor(self, dialect):
         return process_string_literal
 
 
+# SQLAlchemy's built-in `JSON.JSONIndexType`, `JSON.JSONIntIndexType`, and
+# `JSON.JSONStrIndexType` do not define `__visit_name__`, which causes
+# `BigQueryTypeCompiler.process()` to fail when `visit_bindparam` formats the
+# `%(name:TYPE)s` parameter placeholder. These dialect subclasses supply
+# `__visit_name__` and `process_string_literal` (to escape `%` as `%%` when
+# `literal_binds=True`) and are registered in `BigQueryDialect.colspecs`.
+class BQJSONIndexType(sqlalchemy.sql.sqltypes.JSON.JSONIndexType):
+    __visit_name__ = "json_index"
+
+    def string_literal_processor(self, dialect):
+        return process_string_literal
+
+
+class BQJSONIntIndexType(sqlalchemy.sql.sqltypes.JSON.JSONIntIndexType):
+    __visit_name__ = "json_int_index"
+
+
+class BQJSONStrIndexType(sqlalchemy.sql.sqltypes.JSON.JSONStrIndexType):
+    __visit_name__ = "json_str_index"
+
+    def string_literal_processor(self, dialect):
+        return process_string_literal
+
+
+class BQJSONPathType(sqlalchemy.sql.sqltypes.JSON.JSONPathType):
+    __visit_name__ = "json_path"
+
+    def _format_value(self, value):
+        # Convert a sequence of path steps like `("a", 0, 'b"c')` into a
+        # BigQuery JSONPath string `'$."a"[0]."b\\"c"'` for `JSON_QUERY`.
+        path = "".join(_format_json_path_step(elem) for elem in value)
+        return f"${path}"
+
+    def bind_processor(self, dialect):
+        def process(value):
+            return self._format_value(value)
+
+        return process
+
+    def literal_processor(self, dialect):
+        def process(value):
+            return process_string_literal(self._format_value(value))
+
+        return process
+
+
 class BQBinary(sqlalchemy.sql.sqltypes._Binary):
     @staticmethod
     def __process_bytes_literal(value):
@@ -1079,6 +1284,10 @@ class BigQueryDialect(DefaultDialect):
     supports_native_decimal = True
     description_encoding = None
     supports_native_boolean = True
+    # google-cloud-bigquery's DB-API cursor already parses JSON fields into
+    # Python objects via CELL_DATA_PARSER.json_to_py, so SQLAlchemy's default
+    # JSON.result_processor must not call json.loads() a second time.
+    supports_native_json_deserialization = True
     supports_simple_order_by_label = True
     postfetch_lastrowid = False
     preexecute_autoincrement_sequences = False
@@ -1092,6 +1301,10 @@ 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.JSONIndexType: BQJSONIndexType,
+        sqlalchemy.sql.sqltypes.JSON.JSONIntIndexType: BQJSONIntIndexType,
+        sqlalchemy.sql.sqltypes.JSON.JSONStrIndexType: BQJSONStrIndexType,
+        sqlalchemy.sql.sqltypes.JSON.JSONPathType: BQJSONPathType,
     }
 
     def __init__(

packages/sqlalchemy-bigquery/tests/system/test__struct.py

@@ -117,7 +117,7 @@ def test_complex_literals_pr_67(engine, bigquery_dataset, metadata):
         sqlalchemy.select((table.c.dimensions.x__count + 5).label("c")).compile(engine)
     )
     want = (
-        f"SELECT `{table_name}`.`dimensions`.x__count + %(param_1:INT64)s AS `c` \n"
+        f"SELECT `{table_name}`.`dimensions`.`x__count` + %(param_1:INT64)s AS `c` \n"
         f"FROM `{table_name}`"
     )
 
@@ -185,7 +185,7 @@ def test_unnest_and_struct_access_233(engine, bigquery_dataset, metadata):
         f" AS `another_mock_objects`"
         f") AS `anon_1` "
         f"ON "
-        f"`anon_1`.`another_mock_objects`.object_id = "
+        f"`anon_1`.`another_mock_objects`.`object_id` = "
         f"`{bigquery_dataset}.Mock`.`mock_id`"
     )
     assert got == want

packages/sqlalchemy-bigquery/tests/system/test_sqlalchemy_bigquery.py

@@ -930,3 +930,62 @@ def test_not_regexp_match(session, table):
 
     number_of_non_barrowst_records = 993
     assert len(results) == number_of_non_barrowst_records
+
+
+def test_json_subscripts(engine, bigquery_dataset, metadata):
+    table_name = f"{bigquery_dataset}.test_json_subscripts"
+    table = Table(
+        table_name,
+        metadata,
+        Column("id", sqlalchemy.Integer),
+        Column("payload", sqlalchemy.JSON),
+    )
+    metadata.create_all(engine)
+
+    sample_json = (
+        '{"name": "alice", "items": [{"id": 10}, {"id": 20}], '
+        '"a`b": "tricky", "True": "bool_val"}'
+    )
+    with engine.connect() as conn:
+        conn.execute(
+            table.insert().from_select(
+                ["id", "payload"],
+                select(
+                    sqlalchemy.literal(1),
+                    func.parse_json(sample_json),
+                ),
+            )
+        )
+
+        stmt = select(
+            table.c.payload,
+            table.c.payload["name"],
+            table.c.payload["items"][1],
+            table.c.payload["items"][1]["id"],
+            table.c.payload[("items", 0, "id")],
+            table.c.payload["a`b"],
+            table.c.payload[True],
+            table.c.payload["a` FROM `events`; DROP TABLE `x`; --"],
+        ).where(func.json_value(table.c.payload["name"]) == "alice")
+
+        row = conn.execute(stmt).fetchone()
+        assert row == (
+            {
+                "name": "alice",
+                "items": [{"id": 10}, {"id": 20}],
+                "a`b": "tricky",
+                "True": "bool_val",
+            },
+            "alice",
+            {"id": 20},
+            20,
+            10,
+            "tricky",
+            "bool_val",
+            None,
+        )
+
+        # Also verify literal_binds=True execution against BigQuery
+        compiled_literal = stmt.compile(engine, compile_kwargs={"literal_binds": True})
+        literal_row = conn.exec_driver_sql(str(compiled_literal)).fetchone()
+        assert literal_row == row