|
30 | 30 | class MemoryStateClient(StateClient): |
31 | 31 | """A state client whose backend is a dict.""" |
32 | 32 |
|
33 | | - def __init__(self, checkpoints=None, fail_writes=False): |
34 | | - self.backend: Dict[str, str] = {} |
| 33 | + def __init__(self, checkpoints=None, fail_writes=False, backend=None): |
| 34 | + self.backend: Dict[str, str] = dict(backend or {}) |
35 | 35 | self.writes: List[List[Tuple[str, str]]] = [] |
36 | 36 | self.fail_writes = fail_writes |
37 | 37 | super().__init__( |
@@ -395,3 +395,103 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one(): |
395 | 395 | ] |
396 | 396 | assert cursors |
397 | 397 | assert all(cursor.startswith("next-") for cursor in cursors) |
| 398 | + |
| 399 | + |
| 400 | +def test_resolvers_of_an_ended_sync_cannot_change_the_state_of_the_next_one(): |
| 401 | + plugin_attributes: Dict[str, StateClient] = {} |
| 402 | + next_started, leftover_done = threading.Event(), threading.Event() |
| 403 | + leftover_errors: List[Exception] = [] |
| 404 | + |
| 405 | + class LeftoverResolver(TableResolver): |
| 406 | + def __init__(self): |
| 407 | + super().__init__(Table("leftover", [Column("id", pa.int64())])) |
| 408 | + |
| 409 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 410 | + yield {"id": 1} |
| 411 | + try: |
| 412 | + assert next_started.wait(5), "the next sync did not start" |
| 413 | + plugin_attributes["state"].set_key("leftover_cursor", "ended") |
| 414 | + except Exception as e: |
| 415 | + leftover_errors.append(e) |
| 416 | + finally: |
| 417 | + leftover_done.set() |
| 418 | + |
| 419 | + class NextResolver(CursorResolver): |
| 420 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 421 | + next_started.set() |
| 422 | + assert leftover_done.wait(5), "the leftover resolver did not finish" |
| 423 | + yield from super().resolve(client, parent_resource) |
| 424 | + |
| 425 | + scheduler = Scheduler(4) |
| 426 | + try: |
| 427 | + ended = scheduler.sync(TestClient(), [LeftoverResolver()]) |
| 428 | + while not isinstance(next(ended), SyncInsertMessage): |
| 429 | + pass |
| 430 | + ended.close() |
| 431 | + |
| 432 | + plugin_attributes["state"] = MemoryStateClient( |
| 433 | + checkpoints=scheduler.checkpoints |
| 434 | + ) |
| 435 | + scheduler.set_post_sync_hook(plugin_attributes["state"].flush) |
| 436 | + messages = list( |
| 437 | + scheduler.sync( |
| 438 | + TestClient(), [NextResolver(plugin_attributes["state"], "next")] |
| 439 | + ) |
| 440 | + ) |
| 441 | + finally: |
| 442 | + scheduler.shutdown() |
| 443 | + assert leftover_errors == [] |
| 444 | + assert plugin_attributes["state"].get_key("leftover_cursor") is None |
| 445 | + checkpoints = [ |
| 446 | + checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage) |
| 447 | + ] |
| 448 | + assert [c["cursor"] for c in checkpoints] == [f"next-{i}" for i in range(1, 30)] |
| 449 | + |
| 450 | + |
| 451 | +class SetKeyThenWait(TableResolver): |
| 452 | + def __init__(self, state: StateClient, release: threading.Event): |
| 453 | + super().__init__(Table("set_then_wait", [Column("id", pa.int64())])) |
| 454 | + self._state = state |
| 455 | + self._release = release |
| 456 | + |
| 457 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 458 | + self._state.set_key("cursor", "ended") |
| 459 | + yield {"id": 1} |
| 460 | + self._release.wait(5) |
| 461 | + |
| 462 | + |
| 463 | +def test_a_state_client_reused_across_syncs_drops_unflushed_changes_of_an_ended_sync(): |
| 464 | + release = threading.Event() |
| 465 | + scheduler = Scheduler(4) |
| 466 | + try: |
| 467 | + state = MemoryStateClient( |
| 468 | + checkpoints=scheduler.checkpoints, backend={"cursor": "flushed"} |
| 469 | + ) |
| 470 | + ended = scheduler.sync(TestClient(), [SetKeyThenWait(state, release)]) |
| 471 | + while not isinstance(next(ended), SyncInsertMessage): |
| 472 | + pass |
| 473 | + ended.close() |
| 474 | + release.set() |
| 475 | + |
| 476 | + scheduler.set_post_sync_hook(state.flush) |
| 477 | + cursors_seen: List[Optional[str]] = [] |
| 478 | + |
| 479 | + class ReadCursor(TableResolver): |
| 480 | + def __init__(self): |
| 481 | + super().__init__(Table("read_cursor", [Column("id", pa.int64())])) |
| 482 | + |
| 483 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 484 | + cursors_seen.append(state.get_key("cursor")) |
| 485 | + yield {"id": 1} |
| 486 | + state.set_key("other", "next") |
| 487 | + |
| 488 | + messages = list(scheduler.sync(TestClient(), [ReadCursor()])) |
| 489 | + list(scheduler.sync(TestClient(), [ReadCursor()])) |
| 490 | + finally: |
| 491 | + release.set() |
| 492 | + scheduler.shutdown() |
| 493 | + assert cursors_seen == ["flushed", "flushed"] |
| 494 | + checkpoints = [ |
| 495 | + checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage) |
| 496 | + ] |
| 497 | + assert checkpoints == [{"other": "next"}] |
0 commit comments