195 lines
9.6 KiB
Python
195 lines
9.6 KiB
Python
from datetime import datetime, timezone, timedelta
|
|
from decimal import Decimal
|
|
from pathlib import Path
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
|
|
from psycopg import sql
|
|
import pytest
|
|
|
|
from btcdata_archive.archive import run
|
|
from btcdata_archive.postgres import Postgres
|
|
from btcdata_archive.schema import TABLES, encode_row, decode_row
|
|
from btcdata_archive.sqlite_archive import connect, validate
|
|
from btcdata_archive.weeks import boundaries, WEEK
|
|
|
|
pytestmark = pytest.mark.integration
|
|
CUTOFF = boundaries(datetime.now(timezone.utc))[1]
|
|
START = CUTOFF - 3 * WEEK
|
|
|
|
|
|
def seed(infra):
|
|
for offset, pk in [(0, 1), (14, 3), (21, 4)]:
|
|
instant = START + timedelta(days=offset, microseconds=123456)
|
|
infra.admin.execute('INSERT INTO binance_price VALUES (%s,%s,%s)', (pk, Decimal('1234567890.12345678'), instant))
|
|
infra.admin.execute("INSERT INTO strike_price VALUES (%s,'{\"a\":12345678901234567890.12345678901234567890,\"null\":null}'::jsonb,%s,NULL)", (pk, instant))
|
|
infra.admin.execute('INSERT INTO kraken_minute VALUES (%s,1.00000001,2,3,4,5,1234567890123456.12345678,7,%s)', (int(instant.timestamp()), instant))
|
|
|
|
|
|
def invoke(infra, mode='archive-and-prune', max_weeks=None):
|
|
return run(infra.pg, infra.store, infra.config, CUTOFF, mode, max_weeks)
|
|
|
|
|
|
def test_export_gap_precision_cutoff_restore(infra, tmp_path):
|
|
seed(infra)
|
|
original = {}
|
|
for table in TABLES:
|
|
projection = ','.join(n + '::text' if typ == 'jsonb' else n for n, typ, _ in table.columns)
|
|
original[table.name] = [encode_row(table, row) for row in infra.admin.execute(f'SELECT {projection} FROM public.{table.name} ORDER BY {table.pk}')]
|
|
# Exactly at cutoff must stay, regardless of current week.
|
|
infra.admin.execute('INSERT INTO binance_price VALUES (5,5,%s)', (CUTOFF,))
|
|
invoke(infra, 'verify-only')
|
|
assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 4
|
|
assert infra.admin.execute("SELECT count(*) FROM btcdata_archive.weeks WHERE status='verified'").fetchone()[0] == 3
|
|
invoke(infra)
|
|
assert infra.admin.execute('SELECT id FROM binance_price ORDER BY id').fetchall() == [(4,), (5,)]
|
|
assert infra.admin.execute('SELECT count(*) FROM kraken_minute').fetchone()[0] == 1
|
|
for offset in (0, 7, 14):
|
|
start = START + timedelta(days=offset)
|
|
key = infra.config.prefix + start.date().isoformat() + '.sqlite'
|
|
head = infra.store.head(key)
|
|
path = tmp_path / ('restored-' + str(offset) + '.sqlite')
|
|
infra.store.download(key, head['VersionId'], path)
|
|
meta = validate(path, start, infra.config.pg['dbname'], 2)
|
|
assert all(v['count'] == (0 if offset == 7 else 1) for v in meta['tables'].values())
|
|
# Restore every column through the documented type mapping into isolated PostgreSQL tables.
|
|
db = connect(path, readonly=True)
|
|
for table in TABLES:
|
|
restored = 'restored_' + str(offset) + '_' + table.name
|
|
infra.admin.execute(sql.SQL('CREATE TABLE {} (LIKE public.{})').format(sql.Identifier(restored), sql.Identifier(table.name)))
|
|
rows = db.execute(f'SELECT * FROM {table.name} ORDER BY {table.pk}').fetchall()
|
|
for row in rows:
|
|
placeholders = ','.join('%s::jsonb' if typ == 'jsonb' else '%s' for _, typ, _ in table.columns)
|
|
infra.admin.execute(sql.SQL(f'INSERT INTO {{}} VALUES ({placeholders})').format(sql.Identifier(restored)), decode_row(table, row))
|
|
projection = ','.join(n + '::text' if typ == 'jsonb' else n for n, typ, _ in table.columns)
|
|
roundtrip = [encode_row(table, r) for r in infra.admin.execute(sql.SQL(f'SELECT {projection} FROM {{}} ORDER BY {table.pk}').format(sql.Identifier(restored)))]
|
|
assert roundtrip == rows
|
|
assert all(r in original[table.name] for r in rows)
|
|
db.close()
|
|
assert infra.admin.execute("SELECT count(*) FROM btcdata_archive.weeks WHERE status='complete'").fetchone()[0] == 3
|
|
# Application latest-value reads still work with the retained rows.
|
|
assert infra.admin.execute('SELECT id FROM binance_price ORDER BY recorded_at DESC LIMIT 1').fetchone()[0] == 4
|
|
|
|
|
|
def test_schema_mismatch(infra):
|
|
seed(infra)
|
|
infra.admin.execute('ALTER TABLE strike_price ADD COLUMN unexpected text')
|
|
with pytest.raises(RuntimeError, match='migration required'):
|
|
invoke(infra)
|
|
assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3
|
|
|
|
|
|
def test_concurrent_lock_and_loss(infra, monkeypatch):
|
|
seed(infra)
|
|
second = Postgres(infra.config)
|
|
infra.pg.lock()
|
|
try:
|
|
with pytest.raises(RuntimeError, match='advisory lock'):
|
|
second.lock()
|
|
finally:
|
|
second.close()
|
|
# Closing the coordinating session after upload must prevent ANY deletion.
|
|
upload = infra.store.upload
|
|
def lose_lock(*args):
|
|
result = upload(*args)
|
|
infra.pg.close()
|
|
return result
|
|
monkeypatch.setattr(infra.store, 'upload', lose_lock)
|
|
# Don't acquire a second lock on the same session in run.
|
|
monkeypatch.setattr(infra.pg, 'lock', lambda: None)
|
|
with pytest.raises(Exception):
|
|
invoke(infra)
|
|
assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3
|
|
|
|
|
|
def test_changed_and_late_rows(infra, monkeypatch):
|
|
seed(infra)
|
|
upload = infra.store.upload
|
|
def change(*args):
|
|
version = upload(*args)
|
|
infra.admin.execute('UPDATE binance_price SET price=9 WHERE id=1')
|
|
infra.admin.execute('INSERT INTO binance_price VALUES (2,2,%s)', (START,))
|
|
return version
|
|
monkeypatch.setattr(infra.store, 'upload', change)
|
|
with pytest.raises(RuntimeError, match='reconciliation required'):
|
|
invoke(infra)
|
|
assert infra.admin.execute('SELECT id FROM binance_price ORDER BY id').fetchall() == [(1,), (2,), (3,), (4,)]
|
|
# Other unchanged archived rows were safely pruned; their data must survive the merge.
|
|
assert infra.admin.execute('SELECT count(*) FROM strike_price').fetchone()[0] == 2
|
|
monkeypatch.setattr(infra.store, 'upload', upload)
|
|
invoke(infra)
|
|
key = infra.config.prefix + START.date().isoformat() + '.sqlite'
|
|
path = infra.config.work_dir / 'late.sqlite'
|
|
infra.store.download(key, infra.store.head(key)['VersionId'], path)
|
|
db = connect(path, readonly=True)
|
|
assert db.execute('SELECT id,price FROM binance_price ORDER BY id').fetchall() == [(1,'9.00000000'), (2,'2.00000000')]
|
|
assert db.execute('SELECT count(*) FROM strike_price').fetchone()[0] == 1
|
|
db.close()
|
|
|
|
|
|
@pytest.mark.parametrize('boundary', ['uploaded', 'verified', 'pruning', 'deletion_batch', 'complete'])
|
|
def test_crash_restart_boundaries(infra, monkeypatch, boundary):
|
|
seed(infra)
|
|
ledger, upload = infra.pg.ledger, infra.store.upload
|
|
fired = False
|
|
def crash_ledger(start, status, *args, **kwargs):
|
|
nonlocal fired
|
|
result = ledger(start, status, *args, **kwargs)
|
|
target = 'pruning' if boundary == 'deletion_batch' else boundary
|
|
if not fired and status == target and (boundary != 'deletion_batch' or kwargs.get('deletions')):
|
|
fired = True
|
|
raise KeyboardInterrupt()
|
|
return result
|
|
def crash_upload(*args):
|
|
result = upload(*args)
|
|
raise KeyboardInterrupt()
|
|
monkeypatch.setattr(infra.pg, 'ledger', crash_ledger)
|
|
if boundary == 'uploaded':
|
|
monkeypatch.setattr(infra.store, 'upload', crash_upload)
|
|
with pytest.raises(KeyboardInterrupt):
|
|
invoke(infra)
|
|
infra.pg.close()
|
|
infra.pg = Postgres(infra.config)
|
|
monkeypatch.setattr(infra.store, 'upload', upload)
|
|
invoke(infra)
|
|
key = infra.config.prefix + START.date().isoformat() + '.sqlite'
|
|
path = infra.config.work_dir / 'resumed.sqlite'
|
|
infra.store.download(key, infra.store.head(key)['VersionId'], path)
|
|
assert all(s['count'] == 1 for s in validate(path, START, infra.config.pg['dbname'], 2)['tables'].values())
|
|
assert infra.admin.execute('SELECT id FROM binance_price').fetchall() == [(4,)]
|
|
|
|
|
|
@pytest.mark.parametrize('operation', ['upload', 'download'])
|
|
def test_network_failure_preserves_data(infra, monkeypatch, operation):
|
|
seed(infra)
|
|
def fail(*args):
|
|
raise OSError('network unavailable')
|
|
monkeypatch.setattr(infra.store, operation, fail)
|
|
with pytest.raises(OSError):
|
|
invoke(infra)
|
|
assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3
|
|
|
|
|
|
def test_missing_delete_privilege(infra):
|
|
seed(infra)
|
|
infra.admin.execute(sql.SQL('REVOKE DELETE ON binance_price FROM {}').format(sql.Identifier(infra.config.pg['user'])))
|
|
with pytest.raises(Exception):
|
|
invoke(infra)
|
|
assert infra.admin.execute('SELECT count(*) FROM binance_price').fetchone()[0] == 3
|
|
|
|
|
|
def test_wrapper_all_weeks_and_bound(infra):
|
|
seed(infra)
|
|
c = infra.config
|
|
env = {**os.environ, 'PGHOST': c.pg['host'], 'PGPORT': str(c.pg['port']), 'PGDATABASE': c.pg['dbname'],
|
|
'PGUSER': c.pg['user'], 'PGPASSWORD': c.pg['password'], 'S3_ENDPOINT_URL': c.endpoint,
|
|
'AWS_ACCESS_KEY_ID': c.access_key, 'AWS_SECRET_ACCESS_KEY': c.secret_key, 'ARCHIVE_BUCKET': c.bucket,
|
|
'ARCHIVE_WORK_DIR': str(c.work_dir), 'PATH': str(Path(sys.executable).parent) + ':' + os.environ['PATH']}
|
|
script = str(Path(__file__).resolve().parents[2] / 'scripts/archive-old-weeks.sh')
|
|
first = subprocess.run([script, '--mode', 'verify-only', '--max-weeks', '1'], env=env, capture_output=True, text=True)
|
|
assert first.returncode == 0, first.stdout
|
|
second = subprocess.run([script, '--mode', 'archive-and-prune'], env=env, capture_output=True, text=True)
|
|
assert second.returncode == 0, second.stdout
|
|
assert infra.admin.execute("SELECT count(*) FROM btcdata_archive.weeks WHERE status='complete'").fetchone()[0] == 3
|