piker/tests/test_storage_nativedb.py

'''
NativeDB durability and history-preservation regressions.

'''
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_pl2np_maps_fields_by_name() -> None:
    '''
    Polars conversion must not depend on DataFrame column positions.

    Provider and legacy Parquet frames can include extra columns or
    present canonical fields in a different order. The old ``zip()``
    conversion paired NumPy field names with DataFrame positions,
    silently assigning unrelated values. Build distinct canonical
    values, prepend provider-only ``count``, reverse canonical order,
    and prove every structured-array field is selected by its name.

    '''
    expected = mk_ohlcv(
        (60, 120),
        (1.15, 2.15),
    )
    expected['open'] = [1.1, 2.1]
    expected['high'] = [1.2, 2.2]
    expected['low'] = [1.0, 2.0]
    expected['volume'] = [10, 20]
    canonical = [name for name, _ in def_iohlcv_fields]
    reordered = (
        tsp.np2pl(expected)
        .with_columns(pl.Series('count', [3, 4]))
        .select(['count', *reversed(canonical)])
    )

    actual = tsp.pl2np(
        reordered,
        dtype=expected.dtype,
    )

    for field in canonical:
        assert actual[field].tolist() == expected[field].tolist()


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_fractional_timestamps_are_not_truncated(
    tmp_path: Path,
) -> None:
    '''
    Durable timestamp coercion must not silently alter provider data.

    NativeDB declares integer epoch seconds, while IB delivers its
    timestamps in a floating dtype. Casting before validation truncated
    fractional values and could collapse distinct rows at one second.
    Supply otherwise-valid half-second values and prove the write rejects
    them before publishing any Parquet path.

    '''
    client = NativeStorageClient(tmp_path)
    frame = tsp.np2pl(mk_ohlcv((60, 120))).with_columns(
        pl.Series('time', [60.5, 120.5])
    )

    with pytest.raises(ValueError, match='whole-second'):
        run(client.write_ohlcv('x.test', frame, 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_sidecar_files(
    tmp_path: Path,
) -> None:
    '''
    Legacy writer locks and crash leftovers are not series entries.

    Earlier deep-fix revisions created persistent
    ``.parquet.lock`` files beside each series. The upstream indexer
    mistakes those sidecars for Parquet data and crashes while parsing
    their period. Arrange both a legacy lock and stale temporary file,
    then prove exact-suffix indexing exposes only the durable series.

    '''
    client = NativeStorageClient(tmp_path)
    run(client.write_ohlcv('x.test', mk_ohlcv((60,)), 60))
    (tmp_path / '.x.test.ohlcv60s.parquet.lock').touch()
    (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_writes_create_no_lock_sidecars(
    tmp_path: Path,
) -> None:
    '''
    Actor-owned NativeDB writes must not create lock sidecars.

    ``datad`` already gives each persistent feed one parent history
    writer, with its child tasks writing distinct timeframe files. A
    redundant filesystem lock previously leaked ``.parquet.lock``
    files into NativeDB and made upstream ``flake_update`` crash during
    startup. Exercise replacement and incremental writes, then prove
    no lock artifact exists and the merged history remains intact.

    '''
    client = NativeStorageClient(tmp_path)
    run(client.write_ohlcv('x.test', mk_ohlcv((60,)), 60))
    run(client.update_ohlcv('x.test', mk_ohlcv((120,)), 60))
    path = client.mk_path('x.test', 60)
    stored = pl.read_parquet(path)

    assert not list(tmp_path.glob('*.lock'))
    assert stored['time'].to_list() == [60, 120]