Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -29,16 +29,14 @@
from sqlalchemy.pool import Pool
from sqlalchemy.sql import elements, expression
from sqlalchemy.sql.compiler import (
OPERATORS,
RESERVED_WORDS,
DDLCompiler,
GenericTypeCompiler,
IdentifierPreparer,
SQLCompiler,
selectable,
)
from sqlalchemy.sql.default_comparator import operator_lookup
from sqlalchemy.sql.operators import json_getitem_op, json_path_getitem_op
from sqlalchemy.sql.operators import json_path_getitem_op

from google.cloud import spanner_dbapi
from google.cloud.sqlalchemy_spanner import version as sqlalchemy_spanner_version
Expand Down Expand Up @@ -85,10 +83,6 @@ def reset_connection(dbapi_conn, connection_record, reset_state=None):
dbapi_conn.read_only = False


# register a method to get a single value of a JSON object
OPERATORS[json_getitem_op] = operator_lookup["json_getitem_op"]


_SPANNER_STRING_ESCAPE_TABLE = str.maketrans(
{
"\\": "\\\\",
Expand Down Expand Up @@ -486,14 +480,6 @@ def _json_value_binary(self, binary, eager_grouping=False, **kw):

return "JSON_VALUE(%s, %s)" % (left_text, right_text)

def _generate_generic_binary(self, binary, opstring, eager_grouping=False, **kw):
"""The method is overriden to process JSON data type cases."""
if isinstance(opstring, str):
return super()._generate_generic_binary(
binary, opstring, eager_grouping=eager_grouping, **kw
)
return self._json_value_binary(binary, eager_grouping=eager_grouping, **kw)

def visit_json_getitem_op_binary(self, binary, operator, **kw):
"""Build a JSON_VALUE() function call for single-key JSON indexing."""
return self._json_value_binary(binary, **kw)
Expand Down
31 changes: 31 additions & 0 deletions packages/sqlalchemy-spanner/tests/mockserver_tests/test_json.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,37 @@ def test_select_json_key_sql_injection_prevented(self):
assert "UNION ALL" not in req.sql
assert "secrets" not in req.sql

def test_select_json_where_filter(self):
from tests.mockserver_tests.json_model import Venue

sql = (
"SELECT venues.id, venues.name, venues.description \n"
"FROM venues \n"
"WHERE JSON_VALUE(venues.description, @a0) = @a1"
)
add_venue_query_result(sql, '{"type":"Stadium"}')
engine = self.create_engine()

with Session(engine.execution_options(read_only=True)) as session:
venue = (
session.execute(
select(Venue).where(Venue.description["type"] == "Stadium")
)
.scalars()
.first()
)
eq_(venue.name, "Test")

requests = [
r for r in self.spanner_service.requests if isinstance(r, ExecuteSqlRequest)
]
eq_(1, len(requests))
eq_(sql, requests[0].sql)
eq_('$."type"', requests[0].params["a0"])
eq_("Stadium", requests[0].params["a1"])
eq_(TypeCode.STRING, requests[0].param_types["a0"].code)
eq_(TypeCode.STRING, requests[0].param_types["a1"].code)


def add_venue_query_result(sql: str, description: str):
result = result_set.ResultSet(
Expand Down
61 changes: 53 additions & 8 deletions packages/sqlalchemy-spanner/tests/unit/test_dialect.py
Original file line number Diff line number Diff line change
Expand Up @@ -146,12 +146,11 @@ def test_json_getitem_sql_injection_prevented(self):
"FROM records",
)

# 3. Direct call to _generate_generic_binary
# 3. Verify OPERATORS is not polluted and direct binary dispatch works
assert json_getitem_op not in OPERATORS
binary_expr = records.c.data[malicious_key]
direct_compiler = dialect.statement_compiler(dialect, sa.select(records.c.data))
direct_sql = direct_compiler._generate_generic_binary(
binary_expr, OPERATORS[json_getitem_op]
)
direct_sql = direct_compiler.process(binary_expr)
eq_(direct_sql, "JSON_VALUE(records.data, :data_1)")

def test_json_getitem_adversarial_keys_and_types(self):
Expand Down Expand Up @@ -201,7 +200,7 @@ def test_json_getitem_adversarial_keys_and_types(self):
"FROM records",
)

# Also test non-ClauseElement fallback in _generate_generic_binary
# Also test non-ClauseElement fallback in visit_json_getitem_op_binary
compiler = dialect.statement_compiler(dialect, sa.select(records.c.data))
for right_obj in (
SimpleNamespace(value='display") --'),
Expand All @@ -213,9 +212,7 @@ def test_json_getitem_adversarial_keys_and_types(self):
operator=json_getitem_op,
)
eq_(
compiler._generate_generic_binary(
fake_binary, OPERATORS[json_getitem_op]
),
compiler.visit_json_getitem_op_binary(fake_binary, json_getitem_op),
'JSON_VALUE(records.data, "$.\\"display\\\\\\") --\\"")',
)

Expand Down Expand Up @@ -503,3 +500,51 @@ def test_explicit_literal_where_does_not_duplicate(self):
str(update_stmt.compile(dialect=dialect)),
"UPDATE my_table SET id=:id WHERE 1=1",
)

def test_json_binary_dispatch_and_operators_table_integrity(self):
"""Test that JSON binary dispatch uses visitor methods and OPERATORS is pure."""
dialect = SpannerDialect()
records = sa.Table(
"records",
sa.MetaData(),
sa.Column("id", sa.Integer),
sa.Column("data", sa.JSON),
)

# 1. Global OPERATORS table is not mutated and only contains str values
assert json_getitem_op not in OPERATORS
assert all(isinstance(value, str) for value in OPERATORS.values())

# 2. JSON indexing dispatches cleanly via visit_binary ->
# visit_json_getitem_op_binary
compiler = dialect.statement_compiler(dialect, sa.select(records))
json_binary = records.c.data["key"]
eq_(compiler.visit_binary(json_binary), "JSON_VALUE(records.data, :data_1)")

# 3. Standard binary operators that rely on base
# _generate_generic_binary still work
compiler = dialect.statement_compiler(dialect, sa.select(records))
generic_binary = records.c.id + 10
eq_(compiler.visit_binary(generic_binary), "records.id + :id_1")

compiler = dialect.statement_compiler(dialect, sa.select(records))
comparison_binary = records.c.id == 5
eq_(compiler.visit_binary(comparison_binary), "records.id = :id_1")

# 4. Comparison operators with JSON subscript (exercises both
# base _generate_generic_binary and visit_json_getitem_op_binary)
stmt_where = sa.select(records).where(records.c.data["type"] == "Stadium")
eq_(
str(stmt_where.compile(dialect=dialect)),
"SELECT records.id, records.data \n"
"FROM records \n"
"WHERE JSON_VALUE(records.data, :data_1) = :param_1",
)

stmt_is_null = sa.select(records).where(records.c.data["type"].is_(None))
eq_(
str(stmt_is_null.compile(dialect=dialect)),
"SELECT records.id, records.data \n"
"FROM records \n"
"WHERE JSON_VALUE(records.data, :data_1) IS NULL",
)
Loading