File: test_cluster_transaction.py

package info (click to toggle)
python-redis 6.4.0-1
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 9,432 kB
  • sloc: python: 60,318; sh: 179; makefile: 128
file content (399 lines) | stat: -rw-r--r-- 15,589 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
from typing import Tuple
from unittest.mock import patch, Mock

import pytest

import redis
from redis import CrossSlotTransactionError, RedisClusterException
from redis.asyncio import RedisCluster, Connection
from redis.asyncio.cluster import ClusterNode, NodesManager
from redis.asyncio.retry import Retry
from redis.backoff import NoBackoff
from redis.cluster import PRIMARY
from tests.conftest import skip_if_server_version_lt


def _find_source_and_target_node_for_slot(
    r: RedisCluster, slot: int
) -> Tuple[ClusterNode, ClusterNode]:
    """Returns a pair of ClusterNodes, where the first node is the
    one that owns the slot and the second is a possible target
    for that slot, i.e. a primary node different from the first
    one.
    """
    node_migrating = r.nodes_manager.get_node_from_slot(slot)
    assert node_migrating, f"No node could be found that owns slot #{slot}"

    available_targets = [
        n
        for n in r.nodes_manager.startup_nodes.values()
        if node_migrating.name != n.name and n.server_type == PRIMARY
    ]

    assert available_targets, f"No possible target nodes for slot #{slot}"
    return node_migrating, available_targets[0]


@pytest.mark.onlycluster
class TestClusterTransaction:
    @pytest.mark.onlycluster
    async def test_pipeline_is_true(self, r) -> None:
        "Ensure pipeline instances are not false-y"
        async with r.pipeline(transaction=True) as pipe:
            assert pipe

    @pytest.mark.onlycluster
    async def test_pipeline_empty_transaction(self, r):
        await r.set("a", 0)

        async with r.pipeline(transaction=True) as pipe:
            assert await pipe.execute() == []

    @pytest.mark.onlycluster
    async def test_executes_transaction_against_cluster(self, r) -> None:
        async with r.pipeline(transaction=True) as tx:
            tx.set("{foo}bar", "value1")
            tx.set("{foo}baz", "value2")
            tx.set("{foo}bad", "value3")
            tx.get("{foo}bar")
            tx.get("{foo}baz")
            tx.get("{foo}bad")
            assert await tx.execute() == [
                True,
                True,
                True,
                b"value1",
                b"value2",
                b"value3",
            ]

        await r.flushall()

        tx = r.pipeline(transaction=True)
        tx.set("{foo}bar", "value1")
        tx.set("{foo}baz", "value2")
        tx.set("{foo}bad", "value3")
        tx.get("{foo}bar")
        tx.get("{foo}baz")
        tx.get("{foo}bad")
        assert await tx.execute() == [True, True, True, b"value1", b"value2", b"value3"]

    @pytest.mark.onlycluster
    async def test_throws_exception_on_different_hash_slots(self, r):
        async with r.pipeline(transaction=True) as tx:
            tx.set("{foo}bar", "value1")
            tx.set("{foobar}baz", "value2")

            with pytest.raises(
                CrossSlotTransactionError,
                match="All keys involved in a cluster transaction must map to the same slot",
            ):
                await tx.execute()

    @pytest.mark.onlycluster
    async def test_throws_exception_with_watch_on_different_hash_slots(self, r):
        async with r.pipeline(transaction=True) as tx:
            with pytest.raises(
                RedisClusterException,
                match="WATCH - all keys must map to the same key slot",
            ):
                await tx.watch("key1", "key2")

    @pytest.mark.onlycluster
    async def test_transaction_with_watched_keys(self, r):
        await r.set("a", 0)

        async with r.pipeline(transaction=True) as pipe:
            await pipe.watch("a")
            a = await pipe.get("a")

            pipe.multi()
            pipe.set("a", int(a) + 1)
            assert await pipe.execute() == [True]

    @pytest.mark.onlycluster
    async def test_retry_transaction_during_unfinished_slot_migration(self, r):
        """
        When a transaction is triggered during a migration, MovedError
        or AskError may appear (depends on the key being already migrated
        or the key not existing already). The patch on parse_response
        simulates such an error, but the slot cache is not updated
        (meaning the migration is still ongoing) so the pipeline eventually
        fails as if it was retried but the migration is not yet complete.
        """
        key = "book"
        slot = r.keyslot(key)
        node_migrating, node_importing = _find_source_and_target_node_for_slot(r, slot)

        with (
            patch.object(ClusterNode, "parse_response") as parse_response,
            patch.object(
                NodesManager, "_update_moved_slots"
            ) as manager_update_moved_slots,
        ):

            def ask_redirect_effect(connection, *args, **options):
                if "MULTI" in args:
                    return
                elif "EXEC" in args:
                    raise redis.exceptions.ExecAbortError()

                raise redis.exceptions.AskError(f"{slot} {node_importing.name}")

            parse_response.side_effect = ask_redirect_effect

            async with r.pipeline(transaction=True) as pipe:
                pipe.set(key, "val")
                with pytest.raises(redis.exceptions.AskError) as ex:
                    await pipe.execute()

                assert str(ex.value).startswith(
                    "Command # 1 (SET book val) of pipeline caused error:"
                    f" {slot} {node_importing.name}"
                )

            manager_update_moved_slots.assert_called()

    @pytest.mark.onlycluster
    async def test_retry_transaction_during_slot_migration_successful(
        self, create_redis
    ):
        """
        If a MovedError or AskError appears when calling EXEC and no key is watched,
        the pipeline is retried after updating the node manager slot table. If the
        migration was completed, the transaction may then complete successfully.
        """
        r = await create_redis(flushdb=False)
        key = "book"
        slot = r.keyslot(key)
        node_migrating, node_importing = _find_source_and_target_node_for_slot(r, slot)

        with (
            patch.object(ClusterNode, "parse_response") as parse_response,
            patch.object(
                NodesManager, "_update_moved_slots"
            ) as manager_update_moved_slots,
        ):

            def ask_redirect_effect(conn, *args, **options):
                # first call should go here, we trigger an AskError
                if f"{conn.host}:{conn.port}" == node_migrating.name:
                    if "MULTI" in args:
                        return
                    elif "EXEC" in args:
                        raise redis.exceptions.ExecAbortError()

                    raise redis.exceptions.AskError(f"{slot} {node_importing.name}")
                # if the slot table is updated, the next call will go here
                elif f"{conn.host}:{conn.port}" == node_importing.name:
                    if "EXEC" in args:
                        return ["OK"]  # mock value to validate this section was called
                    return
                else:
                    assert False, f"unexpected node {conn.host}:{conn.port} was called"

            def update_moved_slot():  # simulate slot table update
                ask_error = r.nodes_manager._moved_exception
                assert ask_error is not None, "No AskError was previously triggered"
                assert f"{ask_error.host}:{ask_error.port}" == node_importing.name
                r.nodes_manager._moved_exception = None
                r.nodes_manager.slots_cache[slot] = [node_importing]

            parse_response.side_effect = ask_redirect_effect
            manager_update_moved_slots.side_effect = update_moved_slot

            result = None
            async with r.pipeline(transaction=True) as pipe:
                pipe.multi()
                pipe.set(key, "val")
                result = await pipe.execute()

            assert result and True in result, "Target node was not called"

    @pytest.mark.onlycluster
    async def test_retry_transaction_with_watch_after_slot_migration(self, r):
        """
        If a MovedError or AskError appears when calling WATCH, the client
        must attempt to recover itself before proceeding and no WatchError
        should appear.
        """
        key = "book"
        slot = r.keyslot(key)
        r.reinitialize_steps = 1

        # force a MovedError on the first call to pipe.watch()
        # by switching the node that owns the slot to another one
        _node_migrating, node_importing = _find_source_and_target_node_for_slot(r, slot)
        r.nodes_manager.slots_cache[slot] = [node_importing]

        async with r.pipeline(transaction=True) as pipe:
            await pipe.watch(key)
            pipe.multi()
            pipe.set(key, "val")
            assert await pipe.execute() == [True]

    @pytest.mark.onlycluster
    async def test_retry_transaction_with_watch_during_slot_migration(self, r):
        """
        If a MovedError or AskError appears when calling EXEC and keys were
        being watched before the migration started, a WatchError should appear.
        These errors imply resetting the connection and connecting to a new node,
        so watches are lost anyway and the client code must be notified.
        """
        key = "book"
        slot = r.keyslot(key)
        node_migrating, node_importing = _find_source_and_target_node_for_slot(r, slot)

        with patch.object(ClusterNode, "parse_response") as parse_response:

            def ask_redirect_effect(conn, *args, **options):
                if f"{conn.host}:{conn.port}" == node_migrating.name:
                    # we simulate the watch was sent before the migration started
                    if "WATCH" in args:
                        return b"OK"
                    # but the pipeline was triggered after the migration started
                    elif "MULTI" in args:
                        return
                    elif "EXEC" in args:
                        raise redis.exceptions.ExecAbortError()

                    raise redis.exceptions.AskError(f"{slot} {node_importing.name}")
                # we should not try to connect to any other node
                else:
                    assert False, f"unexpected node {conn.host}:{conn.port} was called"

            parse_response.side_effect = ask_redirect_effect

            async with r.pipeline(transaction=True) as pipe:
                await pipe.watch(key)

                pipe.multi()
                pipe.set(key, "val")
                with pytest.raises(redis.exceptions.WatchError) as ex:
                    await pipe.execute()

                assert str(ex.value).startswith(
                    "Slot rebalancing occurred while watching keys"
                )

    @pytest.mark.onlycluster
    async def test_retry_transaction_on_connection_error(self, r):
        key = "book"
        slot = r.keyslot(key)

        mock_connection = Mock(spec=Connection)
        mock_connection.read_response.side_effect = redis.exceptions.ConnectionError(
            "Conn error"
        )
        mock_connection.retry = Retry(NoBackoff(), 0)

        _node_migrating, node_importing = _find_source_and_target_node_for_slot(r, slot)
        node_importing._free.append(mock_connection)
        r.nodes_manager.slots_cache[slot] = [node_importing]
        r.reinitialize_steps = 1

        async with r.pipeline(transaction=True) as pipe:
            pipe.set(key, "val")
            assert await pipe.execute() == [True]

        assert mock_connection.read_response.call_count == 1

    @pytest.mark.onlycluster
    async def test_retry_transaction_on_connection_error_with_watched_keys(self, r):
        key = "book"
        slot = r.keyslot(key)

        mock_connection = Mock(spec=Connection)
        mock_connection.read_response.side_effect = redis.exceptions.ConnectionError(
            "Conn error"
        )
        mock_connection.retry = Retry(NoBackoff(), 0)

        _node_migrating, node_importing = _find_source_and_target_node_for_slot(r, slot)
        node_importing._free.append(mock_connection)
        r.nodes_manager.slots_cache[slot] = [node_importing]
        r.reinitialize_steps = 1

        async with r.pipeline(transaction=True) as pipe:
            await pipe.watch(key)

            pipe.multi()
            pipe.set(key, "val")
            assert await pipe.execute() == [True]

        assert mock_connection.read_response.call_count == 1

    @pytest.mark.onlycluster
    async def test_exec_error_raised(self, r):
        hashkey = "{key}"
        await r.set(f"{hashkey}:c", "a")

        async with r.pipeline(transaction=True) as pipe:
            pipe.set(f"{hashkey}:a", 1).set(f"{hashkey}:b", 2)
            pipe.lpush(f"{hashkey}:c", 3).set(f"{hashkey}:d", 4)
            with pytest.raises(redis.ResponseError) as ex:
                await pipe.execute()
            assert str(ex.value).startswith(
                "Command # 3 (LPUSH {key}:c 3) of pipeline caused error: "
            )

            # make sure the pipe was restored to a working state
            assert await pipe.set(f"{hashkey}:z", "zzz").execute() == [True]
            assert await r.get(f"{hashkey}:z") == b"zzz"

    @pytest.mark.onlycluster
    async def test_parse_error_raised(self, r):
        hashkey = "{key}"
        async with r.pipeline(transaction=True) as pipe:
            # the zrem is invalid because we don't pass any keys to it
            pipe.set(f"{hashkey}:a", 1).zrem(f"{hashkey}:b").set(f"{hashkey}:b", 2)
            with pytest.raises(redis.ResponseError) as ex:
                await pipe.execute()

            assert str(ex.value).startswith(
                "Command # 2 (ZREM {key}:b) of pipeline caused error: wrong number"
            )

            # make sure the pipe was restored to a working state
            assert await pipe.set(f"{hashkey}:z", "zzz").execute() == [True]
            assert await r.get(f"{hashkey}:z") == b"zzz"

    @pytest.mark.onlycluster
    async def test_transaction_callable(self, r):
        hashkey = "{key}"
        await r.set(f"{hashkey}:a", 1)
        await r.set(f"{hashkey}:b", 2)
        has_run = []

        async def my_transaction(pipe):
            a_value = await pipe.get(f"{hashkey}:a")
            assert a_value in (b"1", b"2")
            b_value = await pipe.get(f"{hashkey}:b")
            assert b_value == b"2"

            # silly run-once code... incr's "a" so WatchError should be raised
            # forcing this all to run again. this should incr "a" once to "2"
            if not has_run:
                await r.incr(f"{hashkey}:a")
                has_run.append("it has")

            pipe.multi()
            pipe.set(f"{hashkey}:c", int(a_value) + int(b_value))

        result = await r.transaction(my_transaction, f"{hashkey}:a", f"{hashkey}:b")
        assert result == [True]
        assert await r.get(f"{hashkey}:c") == b"4"

    @pytest.mark.onlycluster
    @skip_if_server_version_lt("2.0.0")
    async def test_transaction_discard(self, r):
        hashkey = "{key}"

        # pipelines enabled as transactions can be discarded at any point
        async with r.pipeline(transaction=True) as pipe:
            await pipe.watch(f"{hashkey}:key")
            await pipe.set(f"{hashkey}:key", "someval")
            await pipe.discard()

            assert not pipe._execution_strategy._watching
            assert not len(pipe)