piker/tests/test_storage_nativedb.py

365 lines
9.4 KiB
Python

'''
NativeDB durability and history-preservation regressions.
'''
from fcntl import (
flock,
LOCK_EX,
LOCK_UN,
)
import os
from pathlib import Path
import numpy as np
import polars as pl
import pytest
import trio
from piker import tsp
from piker.data._source import def_iohlcv_fields
from piker.storage.nativedb import NativeStorageClient
def mk_ohlcv(
times: tuple[float, ...],
closes: tuple[float, ...] | None = None,
) -> np.ndarray:
'''
Build a minimal structured OHLCV frame.
'''
array = np.zeros(
len(times),
dtype=np.dtype(def_iohlcv_fields),
)
array['index'] = np.arange(len(times))
array['time'] = times
array['close'] = closes or times
return array
def run(coro) -> None:
'''
Run a NativeDB operation with a bounded Trio clock.
'''
async def main() -> None:
with trio.fail_after(1):
await coro
trio.run(main)
@pytest.mark.parametrize('use_polars', [False, True])
def test_numpy_and_polars_round_trip(
tmp_path: Path,
use_polars: bool,
) -> None:
'''
Both supported frame types survive durable serialization.
'''
client = NativeStorageClient(tmp_path)
array = mk_ohlcv(
(60, 120, 180),
(1, 2, 3),
)
payload = tsp.np2pl(array) if use_polars else array
run(client.write_ohlcv('x.test', payload, 60))
loaded = trio.run(client.read_ohlcv, 'x.test', 60)
assert loaded['time'].tolist() == [60, 120, 180]
assert loaded['close'].tolist() == [1, 2, 3]
def test_update_preserves_history_and_resolves_conflicts(
tmp_path: Path,
) -> None:
'''
Incremental writes retain old rows and prefer incoming bars.
'''
client = NativeStorageClient(tmp_path)
old = mk_ohlcv(
(60, 120, 180),
(1, 2, 3),
)
incoming = mk_ohlcv(
(180, 240),
(30, 4),
)
run(client.write_ohlcv('x.test', old, 60))
run(client.update_ohlcv('x.test', incoming, 60))
stored = pl.read_parquet(client.mk_path('x.test', 60))
assert stored['time'].to_list() == [60, 120, 180, 240]
assert stored['close'].to_list() == [1, 2, 30, 4]
assert stored['index'].to_list() == [0, 1, 2, 3]
def test_write_ohlcv_remains_an_explicit_replacement(
tmp_path: Path,
) -> None:
'''
Full repair writes can intentionally remove persisted rows.
'''
client = NativeStorageClient(tmp_path)
old = mk_ohlcv((60, 120, 180))
replacement = mk_ohlcv((120, 180))
run(client.write_ohlcv('x.test', old, 60))
run(client.write_ohlcv('x.test', replacement, 60))
stored = pl.read_parquet(client.mk_path('x.test', 60))
assert stored['time'].to_list() == [120, 180]
def test_failed_write_preserves_file_and_cache(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
'''
A failed parquet write leaves durable and cached data intact.
'''
client = NativeStorageClient(tmp_path)
old = mk_ohlcv((60, 120))
incoming = mk_ohlcv((180,))
run(client.write_ohlcv('x.test', old, 60))
def fail_write(
df: pl.DataFrame,
file: Path,
*args,
**kwargs,
) -> None:
Path(file).write_bytes(b'partial parquet')
raise OSError('simulated write failure')
monkeypatch.setattr(
pl.DataFrame,
'write_parquet',
fail_write,
)
with pytest.raises(OSError, match='simulated write failure'):
run(client.update_ohlcv('x.test', incoming, 60))
stored = pl.read_parquet(client.mk_path('x.test', 60))
cached = trio.run(client.as_df, 'x.test', 60, False)
assert stored['time'].to_list() == [60, 120]
assert cached['time'].to_list() == [60, 120]
assert not list(tmp_path.glob('*.tmp'))
def test_invalid_temporary_parquet_is_not_committed(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
'''
The reopened candidate must validate before atomic replacement.
'''
client = NativeStorageClient(tmp_path)
old = mk_ohlcv((60, 120))
run(client.write_ohlcv('x.test', old, 60))
real_read = pl.read_parquet
def corrupt_candidate(source, *args, **kwargs):
if Path(source).suffix == '.tmp':
return pl.DataFrame({'time': [180]})
return real_read(source, *args, **kwargs)
monkeypatch.setattr(pl, 'read_parquet', corrupt_candidate)
with pytest.raises(ValueError, match='missing columns'):
run(client.update_ohlcv(
'x.test',
mk_ohlcv((180,)),
60,
))
path = client.mk_path('x.test', 60)
stored = real_read(path)
cached = trio.run(client.as_df, 'x.test', 60, False)
assert stored['time'].to_list() == [60, 120]
assert cached['time'].to_list() == [60, 120]
assert not list(tmp_path.glob('*.tmp'))
def test_directory_sync_failure_keeps_visible_cache_consistent(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
'''
A post-replace durability error leaves cache matching the file.
'''
client = NativeStorageClient(tmp_path)
run(client.write_ohlcv('x.test', mk_ohlcv((60,)), 60))
real_fsync = os.fsync
calls: int = 0
def fail_directory_sync(fd: int) -> None:
nonlocal calls
calls += 1
if calls == 2:
raise OSError('simulated directory sync failure')
real_fsync(fd)
monkeypatch.setattr(os, 'fsync', fail_directory_sync)
with pytest.raises(OSError, match='directory sync failure'):
run(client.update_ohlcv(
'x.test',
mk_ohlcv((120,)),
60,
))
stored = pl.read_parquet(client.mk_path('x.test', 60))
cached = trio.run(client.as_df, 'x.test', 60, False)
assert stored['time'].to_list() == [60, 120]
assert cached['time'].to_list() == [60, 120]
@pytest.mark.parametrize(
'times',
[
(0, 60),
(60, 60),
(120, 60),
],
)
def test_update_rejects_invalid_timestamps(
tmp_path: Path,
times: tuple[float, ...],
) -> None:
'''
Incremental input must have positive, increasing timestamps.
'''
client = NativeStorageClient(tmp_path)
incoming = mk_ohlcv(times)
with pytest.raises(ValueError):
run(client.update_ohlcv('x.test', incoming, 60))
assert not client.mk_path('x.test', 60).exists()
def test_replacement_rejects_invalid_timestamps(
tmp_path: Path,
) -> None:
'''
Explicit replacement enforces the same durable invariants.
'''
client = NativeStorageClient(tmp_path)
with pytest.raises(ValueError):
run(client.write_ohlcv(
'x.test',
mk_ohlcv((60, 60)),
60,
))
assert not client.mk_path('x.test', 60).exists()
def test_write_rejects_invalid_schema_and_values(
tmp_path: Path,
) -> None:
'''
Durable OHLCV columns must exist and contain finite numbers.
'''
client = NativeStorageClient(tmp_path)
with pytest.raises(ValueError, match='missing columns'):
run(client.write_ohlcv(
'x.test',
pl.DataFrame({'time': [60]}),
60,
))
invalid = mk_ohlcv((60,))
invalid['close'] = np.nan
with pytest.raises(ValueError, match='finite numeric'):
run(client.write_ohlcv('x.test', invalid, 60))
def test_series_do_not_interfere(
tmp_path: Path,
) -> None:
'''
FQME and timeframe keys isolate merge and cache state.
'''
client = NativeStorageClient(tmp_path)
run(client.write_ohlcv('x.test', mk_ohlcv((60,)), 60))
run(client.write_ohlcv('x.test', mk_ohlcv((1, 2)), 1))
run(client.write_ohlcv('y.test', mk_ohlcv((60, 120)), 60))
run(client.update_ohlcv('x.test', mk_ohlcv((120,)), 60))
x_60 = pl.read_parquet(client.mk_path('x.test', 60))
x_1 = pl.read_parquet(client.mk_path('x.test', 1))
y_60 = pl.read_parquet(client.mk_path('y.test', 60))
assert x_60['time'].to_list() == [60, 120]
assert x_1['time'].to_list() == [1, 2]
assert y_60['time'].to_list() == [60, 120]
def test_index_files_ignores_lock_and_stale_temp_files(
tmp_path: Path,
) -> None:
'''
Crash leftovers and writer locks are not durable series entries.
'''
client = NativeStorageClient(tmp_path)
run(client.write_ohlcv('x.test', mk_ohlcv((60,)), 60))
(tmp_path / 'x.test.ohlcv60s.parquet.crash.tmp').touch()
index = client.index_files()
assert list(index) == ['x.test']
assert index['x.test']['period'] == 60
def test_contended_file_lock_yields_to_trio(
tmp_path: Path,
) -> None:
'''
Cross-client lock contention does not block the actor loop.
'''
client = NativeStorageClient(tmp_path)
run(client.write_ohlcv('x.test', mk_ohlcv((60,)), 60))
path = client.mk_path('x.test', 60)
lock_path = path.with_name(f'.{path.name}.lock')
async def main() -> None:
done = trio.Event()
async def update() -> None:
await client.update_ohlcv(
'x.test',
mk_ohlcv((120,)),
60,
)
done.set()
with lock_path.open('a+b') as lock_file:
flock(lock_file.fileno(), LOCK_EX)
async with trio.open_nursery() as nursery:
nursery.start_soon(update)
await trio.sleep(0.03)
assert not done.is_set()
flock(lock_file.fileno(), LOCK_UN)
with trio.fail_after(0.5):
await done.wait()
trio.run(main)
stored = pl.read_parquet(path)
assert stored['time'].to_list() == [60, 120]