diff --git a/databend_sqlalchemy/connector.py b/databend_sqlalchemy/connector.py index e9de5fb..28ad167 100644 --- a/databend_sqlalchemy/connector.py +++ b/databend_sqlalchemy/connector.py @@ -41,6 +41,9 @@ def escape_string(self, item): item.replace("\\", "\\\\").replace("'", "\\'").replace("%", "%%") ) + def escape_binary(self, item): + return f"unhex('{bytes(item).hex()}')" + def escape_item(self, item): if item is None: return "NULL" @@ -59,6 +62,8 @@ def escape_item(self, item): return self.escape_string(item.strftime("%Y-%m-%d")) + "::date" elif isinstance(item, dict): return self.escape_string(f'parse_json({json.dumps(item)})') + elif isinstance(item, (bytes, bytearray, memoryview)): + return self.escape_binary(item) else: return self.escape_string(item) diff --git a/tests/unit/test_param_escaper.py b/tests/unit/test_param_escaper.py new file mode 100644 index 0000000..6ed6a8e --- /dev/null +++ b/tests/unit/test_param_escaper.py @@ -0,0 +1,34 @@ +from unittest import mock + +import pytest + +from databend_sqlalchemy.connector import Cursor, ParamEscaper + + +@pytest.mark.parametrize( + "raw,expected", + [ + (b"", "unhex('')"), + (b"\x00\x01", "unhex('0001')"), + (b"abc", "unhex('616263')"), + (b"\xc3\xa9", "unhex('c3a9')"), + (b"\x00\xff", "unhex('00ff')"), + (b"\xff", "unhex('ff')"), + (b"\x80", "unhex('80')"), + (b"7\xe7\x9f", "unhex('37e79f')"), + (bytearray(b"\x00\xff"), "unhex('00ff')"), + (memoryview(b"\x00\xff"), "unhex('00ff')"), + ], +) +def test_escape_binary_values(raw, expected): + escaper = ParamEscaper() + assert escaper.escape_item(raw) == expected + + +def test_cursor_mogrify_binary_params(): + cur = Cursor(mock.Mock()) + sql = cur.mogrify( + "INSERT INTO probe (id, payload) VALUES (%(id)s, %(payload)s)", + {"id": 1, "payload": b"\x00\xff"}, + ) + assert sql == "INSERT INTO probe (id, payload) VALUES (1, unhex('00ff'))"