|
31 | 31 | class MemoryStateClient(StateClient): |
32 | 32 | """A state client whose backend is a dict.""" |
33 | 33 |
|
34 | | - def __init__(self, checkpoints=None, fail_writes=False): |
35 | | - self.backend: Dict[str, str] = {} |
| 34 | + def __init__(self, checkpoints=None, fail_writes=False, backend=None): |
| 35 | + self.backend: Dict[str, str] = dict(backend or {}) |
36 | 36 | self.writes: List[List[Tuple[str, str]]] = [] |
37 | 37 | self.fail_writes = fail_writes |
38 | 38 | super().__init__( |
@@ -423,6 +423,106 @@ def test_resolvers_of_an_ended_sync_do_not_checkpoint_into_the_next_one(): |
423 | 423 | assert all(cursor.startswith("next-") for cursor in cursors) |
424 | 424 |
|
425 | 425 |
|
| 426 | +def test_resolvers_of_an_ended_sync_cannot_change_the_state_of_the_next_one(): |
| 427 | + plugin_attributes: Dict[str, StateClient] = {} |
| 428 | + next_started, leftover_done = threading.Event(), threading.Event() |
| 429 | + leftover_errors: List[Exception] = [] |
| 430 | + |
| 431 | + class LeftoverResolver(TableResolver): |
| 432 | + def __init__(self): |
| 433 | + super().__init__(Table("leftover", [Column("id", pa.int64())])) |
| 434 | + |
| 435 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 436 | + yield {"id": 1} |
| 437 | + try: |
| 438 | + assert next_started.wait(5), "the next sync did not start" |
| 439 | + plugin_attributes["state"].set_key("leftover_cursor", "ended") |
| 440 | + except Exception as e: |
| 441 | + leftover_errors.append(e) |
| 442 | + finally: |
| 443 | + leftover_done.set() |
| 444 | + |
| 445 | + class NextResolver(CursorResolver): |
| 446 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 447 | + next_started.set() |
| 448 | + assert leftover_done.wait(5), "the leftover resolver did not finish" |
| 449 | + yield from super().resolve(client, parent_resource) |
| 450 | + |
| 451 | + scheduler = Scheduler(4) |
| 452 | + try: |
| 453 | + ended = scheduler.sync(TestClient(), [LeftoverResolver()]) |
| 454 | + while not isinstance(next(ended), SyncInsertMessage): |
| 455 | + pass |
| 456 | + ended.close() |
| 457 | + |
| 458 | + plugin_attributes["state"] = MemoryStateClient( |
| 459 | + checkpoints=scheduler.checkpoints |
| 460 | + ) |
| 461 | + scheduler.set_post_sync_hook(plugin_attributes["state"].flush) |
| 462 | + messages = list( |
| 463 | + scheduler.sync( |
| 464 | + TestClient(), [NextResolver(plugin_attributes["state"], "next")] |
| 465 | + ) |
| 466 | + ) |
| 467 | + finally: |
| 468 | + scheduler.shutdown() |
| 469 | + assert leftover_errors == [] |
| 470 | + assert plugin_attributes["state"].get_key("leftover_cursor") is None |
| 471 | + checkpoints = [ |
| 472 | + checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage) |
| 473 | + ] |
| 474 | + assert [c["cursor"] for c in checkpoints] == [f"next-{i}" for i in range(1, 30)] |
| 475 | + |
| 476 | + |
| 477 | +class SetKeyThenWait(TableResolver): |
| 478 | + def __init__(self, state: StateClient, release: threading.Event): |
| 479 | + super().__init__(Table("set_then_wait", [Column("id", pa.int64())])) |
| 480 | + self._state = state |
| 481 | + self._release = release |
| 482 | + |
| 483 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 484 | + self._state.set_key("cursor", "ended") |
| 485 | + yield {"id": 1} |
| 486 | + self._release.wait(5) |
| 487 | + |
| 488 | + |
| 489 | +def test_a_state_client_reused_across_syncs_drops_unflushed_changes_of_an_ended_sync(): |
| 490 | + release = threading.Event() |
| 491 | + scheduler = Scheduler(4) |
| 492 | + try: |
| 493 | + state = MemoryStateClient( |
| 494 | + checkpoints=scheduler.checkpoints, backend={"cursor": "flushed"} |
| 495 | + ) |
| 496 | + ended = scheduler.sync(TestClient(), [SetKeyThenWait(state, release)]) |
| 497 | + while not isinstance(next(ended), SyncInsertMessage): |
| 498 | + pass |
| 499 | + ended.close() |
| 500 | + release.set() |
| 501 | + |
| 502 | + scheduler.set_post_sync_hook(state.flush) |
| 503 | + cursors_seen: List[Optional[str]] = [] |
| 504 | + |
| 505 | + class ReadCursor(TableResolver): |
| 506 | + def __init__(self): |
| 507 | + super().__init__(Table("read_cursor", [Column("id", pa.int64())])) |
| 508 | + |
| 509 | + def resolve(self, client, parent_resource) -> Generator[Any, None, None]: |
| 510 | + cursors_seen.append(state.get_key("cursor")) |
| 511 | + yield {"id": 1} |
| 512 | + state.set_key("other", "next") |
| 513 | + |
| 514 | + messages = list(scheduler.sync(TestClient(), [ReadCursor()])) |
| 515 | + list(scheduler.sync(TestClient(), [ReadCursor()])) |
| 516 | + finally: |
| 517 | + release.set() |
| 518 | + scheduler.shutdown() |
| 519 | + assert cursors_seen == ["flushed", "flushed"] |
| 520 | + checkpoints = [ |
| 521 | + checkpoint_rows(m) for m in messages if isinstance(m, SyncCheckpointMessage) |
| 522 | + ] |
| 523 | + assert checkpoints == [{"other": "next"}] |
| 524 | + |
| 525 | + |
426 | 526 | class HookCursorResolver(TableResolver): |
427 | 527 | """Checkpoints its cursor from a resource hook, before the resource is on the queue.""" |
428 | 528 |
|
|
0 commit comments