Skip to content

Commit 4502520

Browse files
mykauldkropachev
authored andcommitted
Add unit test for Session.set_keyspace double-quote escaping
Test that Session.set_keyspace properly escapes double quotes in keyspace names (e.g. 'my"ks' -> USE "my""ks") to prevent CQL injection. Also verifies simple keyspace names are not unnecessarily quoted. Requested in review of PR #758.
1 parent e10bf39 commit 4502520

1 file changed

Lines changed: 27 additions & 0 deletions

File tree

tests/unit/test_cluster.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,33 @@ def test_default_serial_consistency_level_legacy(self, *_):
254254
assert f.message.serial_consistency_level == cl_override
255255

256256

257+
258+
@mock_session_pools
259+
def test_set_keyspace_escapes_quotes(self, *_):
260+
"""
261+
Test that Session.set_keyspace properly escapes double quotes in
262+
keyspace names to prevent CQL injection.
263+
Requested in review of PR #758.
264+
"""
265+
c = Cluster(protocol_version=4)
266+
s = Session(c, [Host("127.0.0.1", SimpleConvictionPolicy, host_id=uuid.uuid4())])
267+
c.connection_class.initialize_reactor()
268+
269+
s.execute = Mock()
270+
271+
s.set_keyspace('my"ks')
272+
query = s.execute.call_args[0][0]
273+
assert query == 'USE "my""ks"', (
274+
"Double quotes in keyspace name must be escaped as double-double quotes, "
275+
"got: %r" % query)
276+
277+
# Also verify a simple keyspace name doesn't get unnecessarily quoted
278+
s.execute.reset_mock()
279+
s.set_keyspace('simple_ks')
280+
query = s.execute.call_args[0][0]
281+
assert query == 'USE simple_ks', (
282+
"Simple keyspace names should not be quoted, got: %r" % query)
283+
257284
class ProtocolVersionTests(unittest.TestCase):
258285

259286
def test_protocol_downgrade_test(self):

0 commit comments

Comments
 (0)