diff --git a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py index 4a13d27f7b28..70666366fe80 100644 --- a/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py +++ b/packages/sqlalchemy-spanner/google/cloud/sqlalchemy_spanner/sqlalchemy_spanner.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import base64 +import collections.abc import re import sqlalchemy @@ -37,7 +38,7 @@ selectable, ) from sqlalchemy.sql.default_comparator import operator_lookup -from sqlalchemy.sql.operators import json_getitem_op +from sqlalchemy.sql.operators import json_getitem_op, json_path_getitem_op from google.cloud import spanner_dbapi from google.cloud.sqlalchemy_spanner import version as sqlalchemy_spanner_version @@ -88,6 +89,96 @@ def reset_connection(dbapi_conn, connection_record, reset_state=None): OPERATORS[json_getitem_op] = operator_lookup["json_getitem_op"] +_SPANNER_STRING_ESCAPE_TABLE = str.maketrans( + { + "\\": "\\\\", + '"': '\\"', + "'": "\\'", + "\n": "\\n", + "\r": "\\r", + "\t": "\\t", + **{ + chr(c): f"\\u{c:04x}" + for c in (*range(0x00, 0x20), 0x7F) + if chr(c) not in ("\n", "\r", "\t") + }, + } +) + + +def _escape_json_path_key(key): + """Escape a JSON member key for use inside double-quoted GoogleSQL JSONPath. + + GoogleSQL's JSONPath lexer (`kEscKeyLexerStandard`) consumes `\\[^"]` as a + two-character pair during tokenization, while `RemoveBackSlashFollowedByChar` + only unescapes `\\"`. Doubling backslashes only when immediately preceding + `"` or end-of-string preserves interior backslashes while preventing any + trailing or quote-adjacent backslash from unescaping a `"` delimiter. + """ + s = str(key) + s = re.sub(r'(\\+)(?="|\Z)', r"\g<1>\g<1>", s) + return s.replace('"', '\\"') + + +def _escape_spanner_string_literal(value, dialect=None): + """Escape a string value as a double-quoted GoogleSQL string literal.""" + escaped = str(value).translate(_SPANNER_STRING_ESCAPE_TABLE) + if dialect is not None and dialect.identifier_preparer._double_percents: + escaped = escaped.replace("%", "%%") + return f'"{escaped}"' + + +def _is_json_path_sequence(value): + """Return True if value is a sequence of JSON path steps (not str/bytes).""" + return not isinstance(value, (str, bytes, bytearray)) and isinstance( + value, collections.abc.Sequence + ) + + +def _format_json_path_step(elem): + """Format a single JSONPath member key or integer array index step.""" + if isinstance(elem, int) and not isinstance(elem, bool): + return "[%s]" % elem + return '."%s"' % _escape_json_path_key(elem) + + +class _FormatTypeMixin: + def _format_value(self, value): + if _is_json_path_sequence(value): + return "$%s" % "".join(_format_json_path_step(elem) for elem in value) + return "$%s" % _format_json_path_step(value) + + def bind_processor(self, dialect): + super_proc = self.string_bind_processor(dialect) + + def process(value): + if value is None: + return None + formatted = self._format_value(value) + if super_proc: + formatted = super_proc(formatted) + return formatted + + return process + + def literal_processor(self, dialect): + def process(value): + if value is None: + return "NULL" + formatted = self._format_value(value) + return _escape_spanner_string_literal(formatted, dialect) + + return process + + +class SpannerJSONIndexType(_FormatTypeMixin, types.JSON.JSONIndexType): + """Format a JSON index for Spanner JSON_VALUE path.""" + + +class SpannerJSONPathType(_FormatTypeMixin, types.JSON.JSONPathType): + """Format a JSON path sequence for Spanner JSON_VALUE path.""" + + # PickleType that can be used with Spanner. # Binary values are automatically encoded/decoded to/from base64. # Usage: @@ -341,52 +432,47 @@ def visit_like_op_binary(self, binary, operator, **kw): binary.right._compiler_dispatch(self, **kw), ) - def _generate_generic_binary(self, binary, opstring, eager_grouping=False, **kw): - """The method is overriden to process JSON data type cases.""" - _in_binary = kw.get("_in_binary", False) - - kw["_in_binary"] = True + def _json_value_binary(self, binary, eager_grouping=False, **kw): + """Build a parameterized or safely escaped JSON_VALUE() function call.""" + left_text = self.process(binary.left, eager_grouping=eager_grouping, **kw) + right = binary.right - if isinstance(opstring, str): - text = ( - binary.left._compiler_dispatch( - self, eager_grouping=eager_grouping, **kw - ) - + opstring - + binary.right._compiler_dispatch( - self, eager_grouping=eager_grouping, **kw + if isinstance(right, elements.BindParameter): + if not isinstance( + right.type, (types.JSON.JSONIndexType, types.JSON.JSONPathType) + ): + right = right._clone(maintain_key=True) + right.type = ( + SpannerJSONPathType() + if getattr(binary, "operator", None) is json_path_getitem_op + else SpannerJSONIndexType() ) - ) - if _in_binary and eager_grouping: - text = "(%s)" % text + right_text = self.process(right, eager_grouping=eager_grouping, **kw) + elif isinstance(right, elements.ClauseElement): + right_text = self.process(right, eager_grouping=eager_grouping, **kw) else: - # got JSON data - right_value = getattr( - binary.right, "value", None - ) or binary.right._compiler_dispatch( - self, eager_grouping=eager_grouping, **kw + raw_value = getattr(right, "value", right) + right_text = SpannerJSONIndexType().literal_processor(self.dialect)( + raw_value ) - text = ( - binary.left._compiler_dispatch( - self, eager_grouping=eager_grouping, **kw - ) - + """, "$.""" - + str(right_value) - + '"' + 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 ) - text = "JSON_VALUE(%s)" % text + return self._json_value_binary(binary, eager_grouping=eager_grouping, **kw) - return text + 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) def visit_json_path_getitem_op_binary(self, binary, operator, **kw): """Build a JSON_VALUE() function call.""" - expr = """JSON_VALUE(%s, "$.%s")""" - - return expr % ( - self.process(binary.left, **kw), - self.process(binary.right, **kw), - ) + return self._json_value_binary(binary, **kw) def render_literal_value(self, value, type_): """Render the value of a bind parameter as a quoted literal. @@ -407,8 +493,28 @@ def render_literal_value(self, value, type_): # (e.g. JSON) return self.process(elements.Null._instance()) - raw = ["\\", "'", '"', "\n", "\t", "\r"] - if isinstance(value, str) and any(single in value for single in raw): + if isinstance(type_, (types.JSON.JSONIndexType, types.JSON.JSONPathType)): + processor = type_._cached_literal_processor(self.dialect) + if processor: + return processor(value) + + raw = ("\\", "'", '"', "\n", "\t", "\r") + has_ctrl = isinstance(value, str) and any( + (ord(ch) < 0x20 and ch not in ("\n", "\t")) or ord(ch) == 0x7F + for ch in value + ) + if isinstance(value, str) and ( + has_ctrl or any(single in value for single in raw) + ): + if ( + has_ctrl + or '"""' in value + or value.endswith('"') + or value.endswith("\\") + ): + return _escape_spanner_string_literal(value, self.dialect) + if self.dialect.identifier_preparer._double_percents: + value = value.replace("%", "%%") value = 'r"""{}"""'.format(value) return value else: @@ -905,6 +1011,10 @@ class SpannerDialect(DefaultDialect): execution_ctx_cls = SpannerExecutionContext _json_serializer = JsonObject _json_deserializer = JsonObject + colspecs = { + types.JSON.JSONIndexType: SpannerJSONIndexType, + types.JSON.JSONPathType: SpannerJSONPathType, + } @classmethod def dbapi(cls): diff --git a/packages/sqlalchemy-spanner/tests/mockserver_tests/test_json.py b/packages/sqlalchemy-spanner/tests/mockserver_tests/test_json.py index 2cec82a2ebc3..1cfbb1e96ecc 100644 --- a/packages/sqlalchemy-spanner/tests/mockserver_tests/test_json.py +++ b/packages/sqlalchemy-spanner/tests/mockserver_tests/test_json.py @@ -30,6 +30,7 @@ from tests.mockserver_tests.mock_server_test_base import ( MockServerTestBase, add_result, + add_single_result, add_update_count, ) @@ -143,6 +144,57 @@ def _test_select_json(self, description, expected): venue = session.execute(select(Venue)).first()[0] eq_(venue.description, expected) + def _test_select_json_value(self, expr_fn, expected_param, expected_result): + from tests.mockserver_tests.json_model import Venue + + sql = "SELECT JSON_VALUE(venues.description, @a0) AS anon_1 \nFROM venues" + add_single_result(sql, "anon_1", TypeCode.STRING, [(expected_result,)]) + engine = self.create_engine() + + with Session(engine.execution_options(read_only=True)) as session: + val = session.scalar(select(expr_fn(Venue.description))) + eq_(val, expected_result) + + requests = [ + r for r in self.spanner_service.requests if isinstance(r, ExecuteSqlRequest) + ] + eq_(1, len(requests)) + eq_(sql, requests[0].sql) + eq_(expected_param, requests[0].params["a0"]) + eq_(TypeCode.STRING, requests[0].param_types["a0"].code) + return requests[0] + + def test_select_json_key(self): + self._test_select_json_value( + lambda col: col["type"].as_string(), '$."type"', "Stadium" + ) + + def test_select_json_int_index_zero(self): + self._test_select_json_value( + lambda col: col[0].as_string(), "$[0]", "first_item" + ) + + def test_select_json_path(self): + self._test_select_json_value( + lambda col: col[("info", 0, "size")].as_string(), + '$."info"[0]."size"', + "Great", + ) + + def test_select_json_key_sql_injection_prevented(self): + malicious_key = ( + 'display") FROM venues WHERE FALSE UNION ALL ' + "SELECT (SELECT secret FROM secrets) -- injected" + ) + req = self._test_select_json_value( + lambda col: col[malicious_key].as_string(), + '$."display\\") FROM venues WHERE FALSE UNION ALL ' + 'SELECT (SELECT secret FROM secrets) -- injected"', + None, + ) + assert "UNION ALL" not in req.sql + assert "secrets" not in req.sql + def add_venue_query_result(sql: str, description: str): result = result_set.ResultSet( diff --git a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py index f44599bb4fcb..d073c53b090a 100644 --- a/packages/sqlalchemy-spanner/tests/unit/test_dialect.py +++ b/packages/sqlalchemy-spanner/tests/unit/test_dialect.py @@ -12,8 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from types import SimpleNamespace from unittest.mock import MagicMock +import sqlalchemy as sa +from sqlalchemy.sql.compiler import OPERATORS +from sqlalchemy.sql.operators import json_getitem_op from sqlalchemy.testing import eq_ from sqlalchemy.testing.plugin.plugin_base import fixtures @@ -100,3 +104,253 @@ def test_max_size_exported(self): eq_(SpannerDialect.max_size, MAX_SIZE) eq_(int_from_size("MAX"), 2621440) eq_(int_from_size("100"), 100) + + def test_json_getitem_sql_injection_prevented(self): + """Test that malicious JSON keys cannot inject SQL.""" + dialect = SpannerDialect() + records = sa.Table("records", sa.MetaData(), sa.Column("data", sa.JSON)) + malicious_key = ( + 'display") FROM records WHERE FALSE UNION ALL ' + "SELECT (SELECT secret FROM secrets) -- injected" + ) + + stmt = sa.select(records.c.data[malicious_key].as_string()) + + # 1. Default parameterized compilation + compiled = stmt.compile(dialect=dialect) + eq_( + compiled.string, + "SELECT JSON_VALUE(records.data, :data_1) AS anon_1 \nFROM records", + ) + assert "UNION ALL" not in compiled.string + assert "secrets" not in compiled.string + raw_params = compiled.construct_params() + eq_(len(raw_params), 1) + param_name, raw_val = next(iter(raw_params.items())) + processed_val = compiled._bind_processors[param_name](raw_val) + eq_( + processed_val, + '$."display\\") FROM records WHERE FALSE UNION ALL ' + 'SELECT (SELECT secret FROM secrets) -- injected"', + ) + + # 2. Literal binds compilation (e.g., Computed columns or literal_binds=True) + compiled_literal = stmt.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + eq_( + compiled_literal.string, + 'SELECT JSON_VALUE(records.data, "$.\\"display\\\\\\") ' + "FROM records WHERE FALSE UNION ALL " + 'SELECT (SELECT secret FROM secrets) -- injected\\"") AS anon_1 \n' + "FROM records", + ) + + # 3. Direct call to _generate_generic_binary + 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] + ) + eq_(direct_sql, "JSON_VALUE(records.data, :data_1)") + + def test_json_getitem_adversarial_keys_and_types(self): + """Test JSON indexing with 0, empty string, backslashes, and quotes.""" + dialect = SpannerDialect(paramstyle="format") + records = sa.Table("records", sa.MetaData(), sa.Column("data", sa.JSON)) + + cases = [ + ("display", '$."display"', '"$.\\"display\\""'), + (0, "$[0]", '"$[0]"'), + (3, "$[3]", '"$[3]"'), + (True, '$."True"', '"$.\\"True\\""'), + (False, '$."False"', '"$.\\"False\\""'), + ("", '$.""', '"$.\\"\\""'), + ('a"b', '$."a\\"b"', '"$.\\"a\\\\\\"b\\""'), + ("a'b", '$."a\'b"', '"$.\\"a\\\'b\\""'), + ("a\\b", '$."a\\b"', '"$.\\"a\\\\b\\""'), + ("a\\", '$."a\\\\"', '"$.\\"a\\\\\\\\\\""'), + ("foo\\\n", '$."foo\\\n"', '"$.\\"foo\\\\\\n\\""'), + ('a\\"b', '$."a\\\\\\"b"', '"$.\\"a\\\\\\\\\\\\\\"b\\""'), + ('"""', '$."\\"\\"\\""', '"$.\\"\\\\\\"\\\\\\"\\\\\\"\\""'), + ( + "a\nb\tc\rd\0e\x7f", + '$."a\nb\tc\rd\0e\x7f"', + '"$.\\"a\\nb\\tc\\rd\\u0000e\\u007f\\""', + ), + ("%s", '$."%s"', '"$.\\"%%s\\""'), + ] + + for key, expected_bind, expected_literal in cases: + stmt = sa.select(records.c.data[key].as_string()) + compiled = stmt.compile(dialect=dialect) + eq_( + compiled.string, + "SELECT JSON_VALUE(records.data, %s) AS anon_1 \nFROM records", + ) + param_name, raw_val = next(iter(compiled.construct_params().items())) + processed = compiled._bind_processors[param_name](raw_val) + eq_(processed, expected_bind) + + compiled_lit = stmt.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + eq_( + compiled_lit.string, + f"SELECT JSON_VALUE(records.data, {expected_literal}) AS anon_1 \n" + "FROM records", + ) + + # Also test non-ClauseElement fallback in _generate_generic_binary + compiler = dialect.statement_compiler(dialect, sa.select(records.c.data)) + for right_obj in ( + SimpleNamespace(value='display") --'), + MagicMock(value='display") --'), + ): + fake_binary = SimpleNamespace( + left=records.c.data, + right=right_obj, + operator=json_getitem_op, + ) + eq_( + compiler._generate_generic_binary( + fake_binary, OPERATORS[json_getitem_op] + ), + 'JSON_VALUE(records.data, "$.\\"display\\\\\\") --\\"")', + ) + + def test_json_path_getitem_compilation_and_escaping(self): + """Test JSON path tuple indexing in parameterized and literal modes.""" + dialect = SpannerDialect(paramstyle="format") + records = sa.Table("records", sa.MetaData(), sa.Column("data", sa.JSON)) + path = ("users", 0, True, 'first"name', "back\\") + + stmt = sa.select(records.c.data[path].as_string()) + compiled = stmt.compile(dialect=dialect) + eq_( + compiled.string, + "SELECT JSON_VALUE(records.data, %s) AS anon_1 \nFROM records", + ) + param_name, raw_val = next(iter(compiled.construct_params().items())) + processed = compiled._bind_processors[param_name](raw_val) + eq_(processed, '$."users"[0]."True"."first\\"name"."back\\\\"') + + compiled_lit = stmt.compile( + dialect=dialect, compile_kwargs={"literal_binds": True} + ) + eq_( + compiled_lit.string, + "SELECT JSON_VALUE(records.data, " + '"$.\\"users\\"[0].\\"True\\".\\"first\\\\\\"name\\".\\"back\\\\\\\\\\"") ' + "AS anon_1 \nFROM records", + ) + + def test_json_getitem_statement_cache_and_bindparam_reuse(self): + """Test statement caching and explicit bindparam reuse without mutation.""" + dialect = SpannerDialect() + records = sa.Table( + "records", + sa.MetaData(), + sa.Column("data", sa.JSON), + sa.Column("name", sa.String), + ) + cache = {} + + stmt1 = sa.select(records.c.data["key1"].as_string()) + stmt2 = sa.select(records.c.data["key2"].as_string()) + + c1, params1, *_ = stmt1._compile_w_cache( + dialect, compiled_cache=cache, column_keys=[] + ) + c2, params2, *_ = stmt2._compile_w_cache( + dialect, compiled_cache=cache, column_keys=[] + ) + eq_(len(cache), 1) + + constructed1 = c1.construct_params(extracted_parameters=params1) + constructed2 = c2.construct_params(extracted_parameters=params2) + bname1, val1 = next(iter(constructed1.items())) + bname2, val2 = next(iter(constructed2.items())) + + eq_(c1._bind_processors[bname1](val1), '$."key1"') + eq_(c2._bind_processors[bname2](val2), '$."key2"') + + # Explicit valued bindparam reused across JSON and non-JSON expressions + bp = sa.bindparam("k", "shared_key") + bp_cache = {} + stmt_bp1 = sa.select(records.c.data[bp].as_string()) + stmt_bp2 = sa.select(records.c.data[bp].as_string()) + c_bp1, p_bp1, *_ = stmt_bp1._compile_w_cache( + dialect, compiled_cache=bp_cache, column_keys=[] + ) + c_bp2, p_bp2, *_ = stmt_bp2._compile_w_cache( + dialect, compiled_cache=bp_cache, column_keys=[] + ) + eq_(len(bp_cache), 1) + eq_( + c_bp2._bind_processors["k"]( + c_bp2.construct_params(extracted_parameters=p_bp2)["k"] + ), + '$."shared_key"', + ) + + # Ensure bp.type was not mutated in-place for non-JSON comparisons + stmt_name = sa.select(records.c.name).where(records.c.name == bp) + c_name = stmt_name.compile(dialect=dialect) + name_val = c_name.construct_params()["k"] + if "k" in c_name._bind_processors: + name_val = c_name._bind_processors["k"](name_val) + eq_(name_val, "shared_key") + + # Explicit bindparam with tuple path (compile-time, cache hit, and runtime) + path_cache = {} + stmt_path = sa.select( + records.c.data[sa.bindparam("p", value=("a", 0, "b"))].as_string() + ) + stmt_untyped = sa.select(records.c.data[sa.bindparam("p")].as_string()) + c_path, p_path, *_ = stmt_path._compile_w_cache( + dialect, compiled_cache=path_cache, column_keys=[] + ) + c_untyped, _, *_ = stmt_untyped._compile_w_cache( + dialect, compiled_cache=path_cache, column_keys=[] + ) + eq_(len(path_cache), 1) + eq_( + c_path._bind_processors["p"]( + c_path.construct_params(extracted_parameters=p_path)["p"] + ), + '$."a"[0]."b"', + ) + eq_(c_untyped._bind_processors["p"]("my_key"), '$."my_key"') + eq_(c_untyped._bind_processors["p"](0), "$[0]") + eq_(c_untyped._bind_processors["p"](("x", 1, "y")), '$."x"[1]."y"') + + def test_render_literal_value_triple_quote_and_trailing_escape(self): + """Test render_literal_value escapes strings breaking raw triple quotes.""" + dialect = SpannerDialect(paramstyle="format") + compiler = dialect.statement_compiler(dialect, sa.select(sa.literal("x"))) + + eq_( + compiler.render_literal_value('a"""b', sa.String()), + '"a\\"\\"\\"b"', + ) + eq_( + compiler.render_literal_value('ends_with_quote"', sa.String()), + '"ends_with_quote\\""', + ) + eq_( + compiler.render_literal_value("ends_with_backslash\\", sa.String()), + '"ends_with_backslash\\\\"', + ) + eq_( + compiler.render_literal_value("has\rcarriage_return", sa.String()), + '"has\\rcarriage_return"', + ) + eq_( + compiler.render_literal_value("has\x00null", sa.String()), + '"has\\u0000null"', + ) + eq_( + compiler.render_literal_value('100% "quoted" mid', sa.String()), + 'r"""100%% "quoted" mid"""', + )