|
33 | 33 | pytestmark = pytest.mark.asyncio |
34 | 34 |
|
35 | 35 |
|
| 36 | +def _claim_structure_tables_in_process( |
| 37 | + db_path: str, |
| 38 | + sessions_table: str, |
| 39 | + messages_table: str, |
| 40 | + ready: Any, |
| 41 | + start: Any, |
| 42 | + results: Any, |
| 43 | +) -> None: |
| 44 | + """Construct a create_tables session in a child process and report the outcome.""" |
| 45 | + ready.set() |
| 46 | + start.wait(timeout=30) |
| 47 | + try: |
| 48 | + session = AdvancedSQLiteSession( |
| 49 | + session_id="concurrent", |
| 50 | + db_path=db_path, |
| 51 | + create_tables=True, |
| 52 | + sessions_table=sessions_table, |
| 53 | + messages_table=messages_table, |
| 54 | + ) |
| 55 | + session.close() |
| 56 | + except ValueError: |
| 57 | + results.put(("rejected", sessions_table)) |
| 58 | + except BaseException as exc: # pragma: no cover - surfaced in the assertion below |
| 59 | + results.put((f"error:{type(exc).__name__}", sessions_table)) |
| 60 | + else: |
| 61 | + results.put(("claimed", sessions_table)) |
| 62 | + |
| 63 | + |
36 | 64 | def _multiprocessing_context() -> Any: |
37 | 65 | method = "spawn" if sys.platform == "win32" else "forkserver" |
38 | 66 | return multiprocessing.get_context(method) |
@@ -3844,3 +3872,117 @@ async def test_structure_tables_accept_equivalent_identifier_casing(tmp_path: Pa |
3844 | 3872 | assert await recased.get_items() == [{"role": "user", "content": "first"}] |
3845 | 3873 | finally: |
3846 | 3874 | recased.close() |
| 3875 | + |
| 3876 | + |
| 3877 | +async def test_structure_tables_reject_distinct_non_ascii_identifiers(tmp_path: Path) -> None: |
| 3878 | + """SQLite folds identifiers with ASCII rules, so these are two different pairs. |
| 3879 | +
|
| 3880 | + Python's `casefold()` equates `ßsessions` and `sssessions`, which would let the second pair |
| 3881 | + through and restore the cross-table mixing this change prevents. |
| 3882 | + """ |
| 3883 | + assert "ßsessions".casefold() == "sssessions".casefold() |
| 3884 | + |
| 3885 | + db_path = tmp_path / "advanced_non_ascii_structure.db" |
| 3886 | + first = AdvancedSQLiteSession( |
| 3887 | + session_id="shared", |
| 3888 | + db_path=db_path, |
| 3889 | + create_tables=True, |
| 3890 | + sessions_table="ßsessions", |
| 3891 | + messages_table="ßmessages", |
| 3892 | + ) |
| 3893 | + try: |
| 3894 | + await first.add_items([{"role": "user", "content": "first"}]) |
| 3895 | + |
| 3896 | + with pytest.raises(ValueError, match="ßsessions"): |
| 3897 | + AdvancedSQLiteSession( |
| 3898 | + session_id="shared", |
| 3899 | + db_path=db_path, |
| 3900 | + create_tables=True, |
| 3901 | + sessions_table="sssessions", |
| 3902 | + messages_table="ssmessages", |
| 3903 | + ) |
| 3904 | + |
| 3905 | + assert await first.get_items() == [{"role": "user", "content": "first"}] |
| 3906 | + finally: |
| 3907 | + first.close() |
| 3908 | + |
| 3909 | + |
| 3910 | +async def test_no_create_session_rejects_a_database_without_an_owner(tmp_path: Path) -> None: |
| 3911 | + """A no-create session must not open a file before a pair has claimed the structure tables. |
| 3912 | +
|
| 3913 | + Accepting it would let another pair claim the tables afterwards, leaving this session writing |
| 3914 | + and reading structure rows owned by that other pair. |
| 3915 | + """ |
| 3916 | + db_path = tmp_path / "advanced_unclaimed_structure.db" |
| 3917 | + |
| 3918 | + with pytest.raises(ValueError, match="create_tables=True"): |
| 3919 | + AdvancedSQLiteSession(session_id="shared", db_path=db_path, create_tables=False) |
| 3920 | + |
| 3921 | + owner = AdvancedSQLiteSession(session_id="shared", db_path=db_path, create_tables=True) |
| 3922 | + try: |
| 3923 | + await owner.add_items([{"role": "user", "content": "first"}]) |
| 3924 | + finally: |
| 3925 | + owner.close() |
| 3926 | + |
| 3927 | + # Once a pair owns the layout, the same pair may open it without creating anything. |
| 3928 | + reader = AdvancedSQLiteSession(session_id="shared", db_path=db_path, create_tables=False) |
| 3929 | + try: |
| 3930 | + assert await reader.get_items() == [{"role": "user", "content": "first"}] |
| 3931 | + finally: |
| 3932 | + reader.close() |
| 3933 | + |
| 3934 | + |
| 3935 | +@pytest.mark.review_optional |
| 3936 | +async def test_concurrent_structure_table_claims_leave_one_coherent_owner(tmp_path: Path) -> None: |
| 3937 | + """Two processes claiming a fresh file with different pairs must not split the layout.""" |
| 3938 | + db_path = tmp_path / "advanced_concurrent_claim.db" |
| 3939 | + pairs = [("a_sessions", "a_messages"), ("b_sessions", "b_messages")] |
| 3940 | + |
| 3941 | + # Create the file with WAL already enabled. Switching journal mode needs an exclusive lock, |
| 3942 | + # so two processes opening a brand-new file race there first, in SQLiteSession.__init__, |
| 3943 | + # before either reaches the structure tables. Settling it up front keeps this test about the |
| 3944 | + # claim. |
| 3945 | + with contextlib.closing(sqlite3.connect(db_path)) as setup_conn: |
| 3946 | + setup_conn.execute("PRAGMA journal_mode=WAL") |
| 3947 | + |
| 3948 | + context = _multiprocessing_context() |
| 3949 | + start = context.Event() |
| 3950 | + results = context.Queue() |
| 3951 | + ready_events = [context.Event(), context.Event()] |
| 3952 | + processes = [ |
| 3953 | + context.Process( |
| 3954 | + target=_claim_structure_tables_in_process, |
| 3955 | + args=(str(db_path), sessions_table, messages_table, ready, start, results), |
| 3956 | + ) |
| 3957 | + for (sessions_table, messages_table), ready in zip(pairs, ready_events, strict=False) |
| 3958 | + ] |
| 3959 | + |
| 3960 | + try: |
| 3961 | + for process in processes: |
| 3962 | + process.start() |
| 3963 | + for ready in ready_events: |
| 3964 | + assert ready.wait(timeout=30) |
| 3965 | + start.set() |
| 3966 | + for process in processes: |
| 3967 | + process.join(timeout=30) |
| 3968 | + assert process.exitcode == 0 |
| 3969 | + |
| 3970 | + outcomes = [results.get(timeout=5), results.get(timeout=5)] |
| 3971 | + claimed = [pair for status, pair in outcomes if status == "claimed"] |
| 3972 | + assert len(claimed) == 1, outcomes |
| 3973 | + assert all(status in {"claimed", "rejected"} for status, _ in outcomes), outcomes |
| 3974 | + |
| 3975 | + # Every owner-bearing structure table must name the one pair that won. |
| 3976 | + winner = claimed[0] |
| 3977 | + with contextlib.closing(sqlite3.connect(db_path)) as conn: |
| 3978 | + for table in ("message_structure", "turn_usage"): |
| 3979 | + owners = { |
| 3980 | + row[3]: row[2] for row in conn.execute(f"PRAGMA foreign_key_list({table})") |
| 3981 | + } |
| 3982 | + assert owners["session_id"] == winner, (table, owners) |
| 3983 | + finally: |
| 3984 | + start.set() |
| 3985 | + for process in processes: |
| 3986 | + if process.is_alive(): |
| 3987 | + process.terminate() |
| 3988 | + process.join(timeout=5) |
0 commit comments