Build 060.27: implement Persistent Market Data Storage

This commit is contained in:
2026-08-01 03:22:25 +03:00
parent cb8acfe5fe
commit 58e5a12a4d
54 changed files with 10243 additions and 101 deletions

View File

@@ -10,6 +10,11 @@ DB_NAME=dzentra_bot
DB_USER=dzentra_bot
DB_PASSWORD=change_me
MARKET_DATA_STORAGE_ENABLED=false
MARKET_DATA_STORAGE_POOL_MIN_SIZE=1
MARKET_DATA_STORAGE_POOL_MAX_SIZE=4
MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS=10
EXCHANGE_ENABLED=true
EXCHANGE_NAME=dzengi
EXCHANGE_BASE_URL=https://demo-api-adapter.dzengi.com

View File

@@ -3,5 +3,6 @@
aiogram==3.13.1
python-dotenv==1.0.1
psycopg[binary]==3.2.9
psycopg-pool==3.3.1
openpyxl==3.1.5
websockets==13.1
websockets==13.1

View File

@@ -7,6 +7,9 @@ from aiogram.client.default import DefaultBotProperties
from src.bootstrap.application import ApplicationComposition
from src.bootstrap.logging import setup_logging
from src.bootstrap.market_data_storage import (
build_market_data_storage,
)
from src.bootstrap.trade_stream_runtime import (
build_trade_stream_production_runtime,
)
@@ -41,8 +44,14 @@ def create_app() -> ApplicationComposition:
pass
raise
trade_stream_runtime = (
build_trade_stream_production_runtime(settings)
market_data_storage = build_market_data_storage(settings)
trade_stream_runtime = build_trade_stream_production_runtime(
settings,
trade_observation_sink=(
market_data_storage.trade_observation_sink
if market_data_storage is not None
else None
),
)
bot = Bot(
@@ -67,6 +76,9 @@ def create_app() -> ApplicationComposition:
"trade_stream_enabled": (
settings.trade_stream.enabled
),
"market_data_storage_enabled": (
settings.market_data_storage.enabled
),
},
)
except Exception:
@@ -76,4 +88,9 @@ def create_app() -> ApplicationComposition:
bot=bot,
dispatcher=dispatcher,
trade_stream_runtime=trade_stream_runtime,
market_data_storage_lifecycle=(
market_data_storage.lifecycle
if market_data_storage is not None
else None
),
)

View File

@@ -1,10 +1,14 @@
from __future__ import annotations
import asyncio
from collections.abc import Callable
from dataclasses import dataclass
from aiogram import Bot, Dispatcher
from src.bootstrap.market_data_storage import (
MarketDataStorageLifecycleProtocol,
)
from src.market_data.acquisition.runtime.trade_stream_production_runtime import (
TradeStreamProductionRuntimeProtocol,
)
@@ -19,6 +23,9 @@ class ApplicationComposition:
trade_stream_runtime: (
TradeStreamProductionRuntimeProtocol | None
)
market_data_storage_lifecycle: (
MarketDataStorageLifecycleProtocol | None
) = None
async def run_application(
@@ -31,25 +38,37 @@ async def run_application(
приложения. Остановка Telegram, Trade Stream и bot session
выполняется одним владельцем и не оставляет фоновых root tasks.
"""
polling_task = asyncio.create_task(
application.dispatcher.start_polling(
application.bot,
close_bot_session=False,
),
name="telegram-polling",
)
runtime_task = (
asyncio.create_task(
application.trade_stream_runtime.run(),
name="trade-stream-runtime",
)
if application.trade_stream_runtime is not None
else None
)
polling_task: asyncio.Task[None] | None = None
runtime_task: asyncio.Task[None] | None = None
primary_error: BaseException | None = None
try:
storage_lifecycle = (
application.market_data_storage_lifecycle
)
if storage_lifecycle is not None:
await _run_blocking_lifecycle_operation(
storage_lifecycle.start,
task_name="market-data-storage-startup",
)
polling_task = asyncio.create_task(
application.dispatcher.start_polling(
application.bot,
close_bot_session=False,
),
name="telegram-polling",
)
runtime_task = (
asyncio.create_task(
application.trade_stream_runtime.run(),
name="trade-stream-runtime",
)
if application.trade_stream_runtime is not None
else None
)
await _wait_for_root_tasks(
polling_task=polling_task,
runtime_task=runtime_task,
@@ -128,7 +147,7 @@ async def _wait_for_root_tasks(
async def _shutdown_application(
*,
application: ApplicationComposition,
polling_task: asyncio.Task[None],
polling_task: asyncio.Task[None] | None,
runtime_task: asyncio.Task[None] | None,
primary_error: BaseException | None,
) -> BaseException | None:
@@ -136,7 +155,7 @@ async def _shutdown_application(
polling_cancelled = False
runtime_cancelled = False
if not polling_task.done():
if polling_task is not None and not polling_task.done():
polling_cancelled = True
polling_task.cancel()
@@ -155,12 +174,13 @@ async def _shutdown_application(
runtime_cancelled = True
runtime_task.cancel()
cleanup_error = await _observe_root_task(
polling_task,
expected_cancellation=polling_cancelled,
primary_error=primary_error,
cleanup_error=cleanup_error,
)
if polling_task is not None:
cleanup_error = await _observe_root_task(
polling_task,
expected_cancellation=polling_cancelled,
primary_error=primary_error,
cleanup_error=cleanup_error,
)
if runtime_task is not None:
cleanup_error = await _observe_root_task(
@@ -170,6 +190,20 @@ async def _shutdown_application(
cleanup_error=cleanup_error,
)
storage_lifecycle = application.market_data_storage_lifecycle
if storage_lifecycle is not None:
try:
await _run_blocking_lifecycle_operation(
storage_lifecycle.stop,
task_name="market-data-storage-shutdown",
)
except BaseException as error:
cleanup_error = _merge_cleanup_error(
cleanup_error,
error,
)
try:
await application.bot.session.close()
except BaseException as error:
@@ -181,6 +215,44 @@ async def _shutdown_application(
return cleanup_error
async def _run_blocking_lifecycle_operation(
operation: Callable[[], None],
*,
task_name: str,
) -> None:
operation_task = asyncio.create_task(
asyncio.to_thread(operation),
name=task_name,
)
try:
await asyncio.shield(operation_task)
except asyncio.CancelledError as cancellation:
try:
await _wait_for_blocking_operation(operation_task)
except BaseException as operation_error:
cancellation.add_note(
"Application blocking lifecycle operation also failed: "
f"{type(operation_error).__name__}."
)
raise
async def _wait_for_blocking_operation(
operation_task: asyncio.Task[None],
) -> None:
while not operation_task.done():
try:
await asyncio.shield(operation_task)
except asyncio.CancelledError:
continue
except BaseException:
break
operation_task.result()
async def _observe_root_task(
task: asyncio.Task[None],
*,

View File

@@ -0,0 +1,156 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Protocol, runtime_checkable
from psycopg.conninfo import make_conninfo
from src.core.config import Settings
from src.market_data.storage.postgres_trade_repository import (
PostgresTradeRepository,
)
from src.market_data.storage.trade_storage_observation_sink import (
TradeStorageObservationSink,
)
from src.storage.migrations import StorageMigrationRunner
from src.storage.postgres_pool import PostgresConnectionPool
@runtime_checkable
class MarketDataStorageLifecycleProtocol(Protocol):
"""Синхронный жизненный цикл под управлением сборки приложения."""
def start(self) -> None:
"""Открыть пул и применить миграции до запуска Runtime."""
...
def stop(self) -> None:
"""Закрыть пул после завершения всех записей Runtime."""
...
class _ConnectionPoolLifecycleProtocol(Protocol):
def open(self) -> None:
...
def close(self) -> None:
...
class _MigrationRunnerProtocol(Protocol):
def run(self) -> tuple[int, ...]:
...
class MarketDataStorageLifecycle:
"""Управляет запуском пула, миграциями и очисткой частичного запуска."""
__slots__ = (
"_connection_pool",
"_migration_runner",
"_started",
)
def __init__(
self,
*,
connection_pool: _ConnectionPoolLifecycleProtocol,
migration_runner: _MigrationRunnerProtocol,
) -> None:
self._connection_pool = connection_pool
self._migration_runner = migration_runner
self._started = False
@property
def started(self) -> bool:
"""Показать, завершились ли запуск пула и миграции."""
return self._started
def start(self) -> None:
if self._started:
return
self._connection_pool.open()
try:
self._migration_runner.run()
except BaseException as error:
try:
self._connection_pool.close()
except BaseException as cleanup_error:
error.add_note(
"Market Data Storage startup cleanup also failed: "
f"{type(cleanup_error).__name__}."
)
raise
self._started = True
def stop(self) -> None:
self._started = False
self._connection_pool.close()
@dataclass(frozen=True, slots=True)
class MarketDataStorageBootstrapComposition:
"""Неактивный граф зависимостей постоянного хранилища сделок."""
connection_pool: PostgresConnectionPool
migration_runner: StorageMigrationRunner
trade_repository: PostgresTradeRepository
trade_observation_sink: TradeStorageObservationSink
lifecycle: MarketDataStorageLifecycle
def build_market_data_storage(
settings: Settings,
) -> MarketDataStorageBootstrapComposition | None:
"""Собрать зависимости хранилища без открытия PostgreSQL."""
storage_settings = settings.market_data_storage
if not storage_settings.enabled:
return None
if not settings.trade_stream.enabled:
raise RuntimeError(
"Trade Stream must be enabled when Market Data Storage "
"is enabled."
)
conninfo = make_conninfo(
host=settings.db_host,
port=settings.db_port,
dbname=settings.db_name,
user=settings.db_user,
password=settings.db_password,
)
connection_pool = PostgresConnectionPool(
conninfo=conninfo,
min_size=storage_settings.pool_min_size,
max_size=storage_settings.pool_max_size,
timeout_seconds=storage_settings.pool_timeout_seconds,
name="market-data-storage",
)
migration_runner = StorageMigrationRunner(
connection_provider=connection_pool.connection,
)
trade_repository = PostgresTradeRepository(
connection_provider=connection_pool.connection,
)
trade_observation_sink = TradeStorageObservationSink(
trade_storage=trade_repository,
venue=settings.exchange_name,
)
lifecycle = MarketDataStorageLifecycle(
connection_pool=connection_pool,
migration_runner=migration_runner,
)
return MarketDataStorageBootstrapComposition(
connection_pool=connection_pool,
migration_runner=migration_runner,
trade_repository=trade_repository,
trade_observation_sink=trade_observation_sink,
lifecycle=lifecycle,
)

View File

@@ -17,6 +17,9 @@ from src.market_data.acquisition.adapters.dzengi.websocket_inbound_message_class
from src.market_data.acquisition.adapters.dzengi.websocket_transport import (
DzengiWebSocketTransport,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.runtime.acquisition_runtime_event_logging_consumer import (
AcquisitionRuntimeEventLoggingConsumer,
)
@@ -39,6 +42,8 @@ from src.market_data.acquisition.trade_stream_runtime_composition import (
def build_trade_stream_production_runtime(
settings: Settings,
*,
trade_observation_sink: TradeObservationSinkProtocol | None = None,
) -> TradeStreamProductionRuntime | None:
"""
Собрать Production Trade Stream Runtime без запуска lifecycle.
@@ -98,6 +103,7 @@ def build_trade_stream_production_runtime(
scheduler_interval_seconds=(
trade_stream.scheduler_interval_seconds
),
trade_observation_sink=trade_observation_sink,
max_recovery_window_ms=trade_stream.recovery_window_ms,
)

View File

@@ -35,18 +35,28 @@ class TradeStreamSettings:
recovery_window_ms: int
@dataclass(frozen=True, slots=True)
class MarketDataStorageSettings:
"""Настройки persistent Canonical Market Data Storage."""
enabled: bool
pool_min_size: int
pool_max_size: int
pool_timeout_seconds: float
@dataclass(slots=True)
class Settings:
# Telegram
bot_token: str
bot_parse_mode: str
# App
# Приложение
app_env: str
log_level: str
tz: str
# Exchange
# Биржа
exchange_enabled: bool
exchange_name: str
exchange_base_url: str
@@ -58,23 +68,24 @@ class Settings:
default_symbol: str
trade_stream: TradeStreamSettings
# Database
# База данных
db_host: str
db_port: int
db_name: str
db_user: str
db_password: str
market_data_storage: MarketDataStorageSettings
# Debug helpers
# Отладочные параметры
debug_enabled: bool
journal_debug_enabled: bool
# helper: demo/live mode
# Вспомогательное свойство режима demo/live
def is_demo_mode(self) -> bool:
return "demo" in self.exchange_base_url.lower()
# parse bool
# Разбор булева значения
def _parse_bool(raw_value: str, default: bool = False) -> bool:
value = (raw_value or "").strip().lower()
if not value:
@@ -83,7 +94,7 @@ def _parse_bool(raw_value: str, default: bool = False) -> bool:
return value in {"1", "true", "yes", "on"}
# parse int
# Разбор целого числа
def _parse_int(raw_value: str, default: int) -> int:
value = (raw_value or "").strip()
if not value:
@@ -292,7 +303,69 @@ def _load_trade_stream_settings(
)
# load all settings
def _load_market_data_storage_settings(
*,
trade_stream_enabled: bool,
) -> MarketDataStorageSettings:
enabled = _parse_strict_bool(
os.getenv("MARKET_DATA_STORAGE_ENABLED", "false"),
name="MARKET_DATA_STORAGE_ENABLED",
default=False,
)
if not enabled:
return MarketDataStorageSettings(
enabled=False,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=10.0,
)
if not trade_stream_enabled:
raise RuntimeError(
"TRADE_STREAM_ENABLED must be true when "
"MARKET_DATA_STORAGE_ENABLED is true"
)
pool_min_size = _parse_positive_int(
os.getenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"1",
),
name="MARKET_DATA_STORAGE_POOL_MIN_SIZE",
default=1,
)
pool_max_size = _parse_positive_int(
os.getenv(
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"4",
),
name="MARKET_DATA_STORAGE_POOL_MAX_SIZE",
default=4,
)
if pool_max_size < pool_min_size:
raise ValueError(
"MARKET_DATA_STORAGE_POOL_MAX_SIZE must not be smaller "
"than MARKET_DATA_STORAGE_POOL_MIN_SIZE"
)
return MarketDataStorageSettings(
enabled=True,
pool_min_size=pool_min_size,
pool_max_size=pool_max_size,
pool_timeout_seconds=_parse_positive_float(
os.getenv(
"MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS",
"10",
),
name="MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS",
default=10.0,
),
)
# Загрузка всех настроек
def load_settings() -> Settings:
bot_token = os.getenv("BOT_TOKEN", "").strip()
@@ -304,12 +377,19 @@ def load_settings() -> Settings:
"",
).strip()
trade_stream = _load_trade_stream_settings(
exchange_base_url=exchange_base_url,
)
market_data_storage = _load_market_data_storage_settings(
trade_stream_enabled=trade_stream.enabled,
)
return Settings(
# Telegram
bot_token=bot_token,
bot_parse_mode=os.getenv("BOT_PARSE_MODE", "HTML").strip() or "HTML",
# App
# Приложение
app_env=os.getenv("APP_ENV", "dev").strip() or "dev",
log_level=os.getenv("LOG_LEVEL", "INFO").strip().upper() or "INFO",
tz=os.getenv("TZ", "Europe/Minsk").strip() or "Europe/Minsk",
@@ -318,7 +398,7 @@ def load_settings() -> Settings:
os.getenv("JOURNAL_DEBUG_ENABLED", "false")
),
# Exchange
# Биржа
exchange_enabled=_parse_bool(os.getenv("EXCHANGE_ENABLED", "false")),
exchange_name=os.getenv("EXCHANGE_NAME", "dzengi").strip() or "dzengi",
exchange_base_url=exchange_base_url,
@@ -329,14 +409,13 @@ def load_settings() -> Settings:
exchange_testnet=_parse_bool(os.getenv("EXCHANGE_TESTNET", "false")),
default_symbol=os.getenv("DEFAULT_SYMBOL", "ETH/USD_LEVERAGE").strip()
or "ETH/USD_LEVERAGE",
trade_stream=_load_trade_stream_settings(
exchange_base_url=exchange_base_url,
),
trade_stream=trade_stream,
# Database
# База данных
db_host=os.getenv("DB_HOST", "localhost").strip() or "localhost",
db_port=_parse_int(os.getenv("DB_PORT", "5432"), 5432),
db_name=os.getenv("DB_NAME", "dzentra_bot").strip() or "dzentra_bot",
db_user=os.getenv("DB_USER", "dzentra_bot").strip() or "dzentra_bot",
db_password=os.getenv("DB_PASSWORD", "").strip(),
market_data_storage=market_data_storage,
)

View File

@@ -0,0 +1,17 @@
from __future__ import annotations
from typing import Protocol, runtime_checkable
from src.market_data.acquisition.models.trade import Trade
@runtime_checkable
class TradeObservationSinkProtocol(Protocol):
"""Надёжный приёмник одного наблюдения канонической сделки."""
def persist(
self,
trade: Trade,
) -> None:
"""Надёжно сохранить наблюдение до продвижения контрольной точки."""
...

View File

@@ -2,6 +2,9 @@
from __future__ import annotations
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.consistency.trade_stream_protocol import (
TradeStreamConsistencyProtocol,
)
@@ -28,8 +31,23 @@ class TradeStreamConsistencyController(
def __init__(
self,
state_store: TradeStreamStateStoreProtocol,
*,
trade_observation_sink: TradeObservationSinkProtocol | None = None,
) -> None:
if (
trade_observation_sink is not None
and not isinstance(
trade_observation_sink,
TradeObservationSinkProtocol,
)
):
raise TypeError(
"trade_observation_sink must implement "
"TradeObservationSinkProtocol"
)
self._state_store = state_store
self._trade_observation_sink = trade_observation_sink
def accept(
self,
@@ -37,7 +55,16 @@ class TradeStreamConsistencyController(
) -> Trade | None:
state = self._get_state(trade.symbol)
return state.accept(trade)
sink = self._trade_observation_sink
return state.accept(
trade,
before_checkpoint=(
sink.persist
if sink is not None
else None
),
)
def _get_state(
self,
@@ -50,4 +77,4 @@ class TradeStreamConsistencyController(
При первом обращении состояние автоматически
создаётся специализированным хранилищем.
"""
return self._state_store.get_or_create(symbol)
return self._state_store.get_or_create(symbol)

View File

@@ -3,6 +3,7 @@
from __future__ import annotations
from collections import deque
from collections.abc import Callable
from dataclasses import dataclass, field
from src.market_data.acquisition.consistency.trade_stream_exceptions import (
@@ -51,6 +52,8 @@ class TradeStreamState:
def accept(
self,
trade: Trade,
*,
before_checkpoint: Callable[[Trade], None] | None = None,
) -> Trade | None:
"""
Проверить сделку на согласованность.
@@ -63,6 +66,11 @@ class TradeStreamState:
либо конфликтующем дубликате.
"""
if before_checkpoint is not None and not callable(
before_checkpoint
):
raise TypeError("before_checkpoint must be callable")
if trade.symbol != self.symbol:
raise ValueError(
f"Unexpected symbol: {trade.symbol!r}"
@@ -85,6 +93,10 @@ class TradeStreamState:
previous,
trade,
):
self._run_before_checkpoint(
trade,
callback=before_checkpoint,
)
return None
raise TradeConsistencyError()
@@ -96,10 +108,19 @@ class TradeStreamState:
previous,
trade,
):
self._run_before_checkpoint(
trade,
callback=before_checkpoint,
)
return None
raise TradeConsistencyError()
self._run_before_checkpoint(
trade,
callback=before_checkpoint,
)
self._append(trade)
self.last_trade_id = trade_id
@@ -107,6 +128,15 @@ class TradeStreamState:
return trade
@staticmethod
def _run_before_checkpoint(
trade: Trade,
*,
callback: Callable[[Trade], None] | None,
) -> None:
if callback is not None:
callback(trade)
@staticmethod
def _is_same_market_trade(
first: Trade,

View File

@@ -10,6 +10,7 @@ from src.market_data.acquisition.exceptions import (
WebSocketMessageDecodeError,
WebSocketTransportError,
)
from src.market_data.acquisition.models.trade import Trade
from src.market_data.acquisition.runtime.live_processing_gate import (
RuntimeLiveProcessingGateProtocol,
)
@@ -121,6 +122,7 @@ class TradeStreamProductionRuntime:
"_state",
"_startup_task",
"_receive_task",
"_market_processing_task",
"_scheduler_task",
"_scheduler_claimed",
"_supervisor_started",
@@ -195,6 +197,9 @@ class TradeStreamProductionRuntime:
self._state = TradeStreamProductionRuntimeState.STOPPED
self._startup_task: asyncio.Task[None] | None = None
self._receive_task: asyncio.Task[None] | None = None
self._market_processing_task: (
asyncio.Task[Trade | None] | None
) = None
self._scheduler_task: asyncio.Task[None] | None = None
self._scheduler_claimed = False
self._supervisor_started = False
@@ -322,6 +327,7 @@ class TradeStreamProductionRuntime:
self._state = TradeStreamProductionRuntimeState.STARTING
self._startup_task = None
self._receive_task = None
self._market_processing_task = None
self._scheduler_task = None
self._supervisor_started = False
self._session_started = False
@@ -485,10 +491,63 @@ class TradeStreamProductionRuntime:
"WebSocketInboundMessageKind"
)
self._trade_stream_service.handle_message(
await self._process_market_document(
document,
)
async def _process_market_document(
self,
document: object,
) -> Trade | None:
processing_task = asyncio.create_task(
asyncio.to_thread(
self._trade_stream_service.handle_message,
document,
),
name="trade-stream-market-processing",
)
self._market_processing_task = processing_task
try:
try:
return await asyncio.shield(
processing_task,
)
except asyncio.CancelledError as cancellation:
try:
await self._wait_for_market_processing_completion(
processing_task,
)
except BaseException as processing_error:
cancellation.add_note(
"Trade Stream market processing also failed: "
f"{type(processing_error).__name__}."
)
raise
except Exception as error:
self._live_processing_gate.fail(error)
raise
finally:
if self._market_processing_task is processing_task:
self._market_processing_task = None
@staticmethod
async def _wait_for_market_processing_completion(
processing_task: asyncio.Task[Trade | None],
) -> None:
while not processing_task.done():
try:
await asyncio.shield(
processing_task,
)
except asyncio.CancelledError:
continue
except BaseException:
break
processing_task.result()
async def _shutdown(
self,
*,
@@ -608,6 +667,7 @@ class TradeStreamProductionRuntime:
self._startup_task = None
self._receive_task = None
self._market_processing_task = None
self._scheduler_task = None
self._session_started = False
self._subscription_correlation_id = None

View File

@@ -15,6 +15,9 @@ from src.market_data.acquisition.adapters.dzengi.rest import (
from src.market_data.acquisition.consistency.trade_stream_consistency_controller import (
TradeStreamConsistencyController,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.consistency.trade_stream_state_store import (
TradeStreamStateStore,
)
@@ -92,6 +95,7 @@ class TradeStreamRuntimeComposition:
state_store: TradeStreamStateStore
consistency_controller: TradeStreamConsistencyController
trade_observation_sink: TradeObservationSinkProtocol | None
recovery_controller: TradeRecoveryController
recovery_window_planner: TradeRecoveryWindowPlanner
runtime_recovery_coordinator: RuntimeRecoveryCoordinator
@@ -121,6 +125,7 @@ def build_trade_stream_runtime_composition(
symbols: tuple[str, ...],
heartbeat_timeout_seconds: float,
scheduler_interval_seconds: float,
trade_observation_sink: TradeObservationSinkProtocol | None = None,
max_recovery_window_ms: int = DEFAULT_TRADE_RECOVERY_WINDOW_MS,
heartbeat_clock: Callable[[], float] = time.monotonic,
recovery_end_time_clock: RuntimeUnixTimeMillisecondsClock = (
@@ -144,6 +149,7 @@ def build_trade_stream_runtime_composition(
consistency_controller = TradeStreamConsistencyController(
state_store,
trade_observation_sink=trade_observation_sink,
)
recovery_controller = TradeRecoveryController(
@@ -216,6 +222,7 @@ def build_trade_stream_runtime_composition(
return TradeStreamRuntimeComposition(
state_store=state_store,
consistency_controller=consistency_controller,
trade_observation_sink=trade_observation_sink,
recovery_controller=recovery_controller,
recovery_window_planner=recovery_window_planner,
runtime_recovery_coordinator=runtime_recovery_coordinator,

View File

@@ -0,0 +1,75 @@
"""Контракты постоянного хранилища канонических рыночных данных."""
from src.market_data.storage.contracts import (
CandleStorageProtocol,
MarketDataBatchWriteResult,
MarketDataWriteResult,
MarketDataWriteStatus,
QuoteStorageProtocol,
TradeStorageProtocol,
)
from src.market_data.storage.exceptions import (
MarketDataStorageConfigurationError,
MarketDataStorageConflictError,
MarketDataStorageError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
)
from src.market_data.storage.market_data_storage import MarketDataStorage
from src.market_data.storage.postgres_candle_repository import (
PostgresCandleRepository,
)
from src.market_data.storage.postgres_quote_repository import (
PostgresQuoteRepository,
)
from src.market_data.storage.postgres_trade_repository import (
PostgresTradeRepository,
)
from src.market_data.storage.trade_storage_observation_sink import (
TradeObservationClock,
TradeStorageObservationSink,
system_utc_datetime,
)
from src.market_data.storage.postgres_partitions import (
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
MarketDataPartition,
MarketDataPartitionResult,
MarketDataPartitionType,
MarketDataRetentionEntryResult,
MarketDataRetentionPolicy,
MarketDataRetentionResult,
PostgresMarketDataPartitionManager,
PostgresMarketDataRetentionService,
build_monthly_partition,
)
__all__ = (
"CandleStorageProtocol",
"MarketDataBatchWriteResult",
"MarketDataPartition",
"MarketDataPartitionResult",
"MarketDataPartitionType",
"MarketDataRetentionEntryResult",
"MarketDataRetentionPolicy",
"MarketDataRetentionResult",
"MarketDataStorage",
"MarketDataStorageConfigurationError",
"MarketDataStorageConflictError",
"MarketDataStorageError",
"MarketDataStorageOperationError",
"MarketDataStorageValidationError",
"MarketDataWriteResult",
"MarketDataWriteStatus",
"PostgresCandleRepository",
"PostgresMarketDataPartitionManager",
"PostgresMarketDataRetentionService",
"PostgresQuoteRepository",
"PostgresTradeRepository",
"QuoteStorageProtocol",
"TradeStorageProtocol",
"TradeObservationClock",
"TradeStorageObservationSink",
"MARKET_DATA_PARTITION_ADVISORY_LOCK_ID",
"build_monthly_partition",
"system_utc_datetime",
)

View File

@@ -0,0 +1,112 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from enum import Enum
from typing import Protocol, runtime_checkable
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.acquisition.models.trade import Trade
class MarketDataWriteStatus(Enum):
"""Результат одной идемпотентной записи канонических рыночных данных."""
INSERTED = "inserted"
DUPLICATE = "duplicate"
PROVENANCE_UPDATED = "provenance_updated"
@dataclass(frozen=True, slots=True)
class MarketDataWriteResult:
"""Наблюдаемый результат одной записи в хранилище."""
status: MarketDataWriteStatus
def __post_init__(self) -> None:
if not isinstance(self.status, MarketDataWriteStatus):
raise TypeError("status must be MarketDataWriteStatus")
@dataclass(frozen=True, slots=True)
class MarketDataBatchWriteResult:
"""Счётчики одной атомарной пакетной записи."""
inserted_count: int
duplicate_count: int
provenance_updated_count: int
def __post_init__(self) -> None:
for field_name, value in (
("inserted_count", self.inserted_count),
("duplicate_count", self.duplicate_count),
(
"provenance_updated_count",
self.provenance_updated_count,
),
):
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"{field_name} must be an integer")
if value < 0:
raise ValueError(f"{field_name} must not be negative")
@property
def total_count(self) -> int:
return (
self.inserted_count
+ self.duplicate_count
+ self.provenance_updated_count
)
@runtime_checkable
class TradeStorageProtocol(Protocol):
"""Граница постоянного хранения канонических сделок только для записи."""
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
...
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
...
@runtime_checkable
class QuoteStorageProtocol(Protocol):
"""Граница постоянного хранения канонических котировок для записи."""
def store_quote(
self,
*,
venue: str,
quote: Quote,
) -> MarketDataWriteResult:
...
@runtime_checkable
class CandleStorageProtocol(Protocol):
"""Граница записи неизменяемых ревизий канонических свечей."""
def store_candle_revision(
self,
*,
venue: str,
candle: Candle,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
...

View File

@@ -0,0 +1,21 @@
from __future__ import annotations
class MarketDataStorageError(Exception):
"""Базовая ошибка постоянного хранилища рыночных данных."""
class MarketDataStorageConfigurationError(MarketDataStorageError):
"""Ошибка конфигурации или жизненного цикла хранилища."""
class MarketDataStorageValidationError(MarketDataStorageError):
"""Недопустимое значение на границе постоянного хранилища."""
class MarketDataStorageConflictError(MarketDataStorageError):
"""Один ключ хранения указывает на противоречащие канонические факты."""
class MarketDataStorageOperationError(MarketDataStorageError):
"""Ошибка операции постоянного хранилища с базой данных."""

View File

@@ -0,0 +1,98 @@
from __future__ import annotations
from datetime import datetime
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.acquisition.models.trade import Trade
from src.market_data.storage.contracts import (
CandleStorageProtocol,
MarketDataBatchWriteResult,
MarketDataWriteResult,
QuoteStorageProtocol,
TradeStorageProtocol,
)
class MarketDataStorage:
"""Фасад только для записи над тремя каноническими репозиториями."""
__slots__ = (
"_candle_storage",
"_quote_storage",
"_trade_storage",
)
def __init__(
self,
*,
trade_storage: TradeStorageProtocol,
quote_storage: QuoteStorageProtocol,
candle_storage: CandleStorageProtocol,
) -> None:
if not isinstance(trade_storage, TradeStorageProtocol):
raise TypeError("trade_storage must implement TradeStorageProtocol")
if not isinstance(quote_storage, QuoteStorageProtocol):
raise TypeError("quote_storage must implement QuoteStorageProtocol")
if not isinstance(candle_storage, CandleStorageProtocol):
raise TypeError(
"candle_storage must implement CandleStorageProtocol"
)
self._trade_storage = trade_storage
self._quote_storage = quote_storage
self._candle_storage = candle_storage
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
return self._trade_storage.store_trade(
venue=venue,
trade=trade,
observed_at=observed_at,
)
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
return self._trade_storage.store_trades(
venue=venue,
trades=trades,
observed_at=observed_at,
)
def store_quote(
self,
*,
venue: str,
quote: Quote,
) -> MarketDataWriteResult:
return self._quote_storage.store_quote(
venue=venue,
quote=quote,
)
def store_candle_revision(
self,
*,
venue: str,
candle: Candle,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
return self._candle_storage.store_candle_revision(
venue=venue,
candle=candle,
observed_at=observed_at,
is_final=is_final,
)

View File

@@ -0,0 +1,341 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
from typing import Any
from src.market_data.acquisition.models.candle import Candle
from src.market_data.storage.contracts import (
MarketDataWriteResult,
MarketDataWriteStatus,
)
from src.market_data.storage.exceptions import (
MarketDataStorageConflictError,
MarketDataStorageError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
)
from src.market_data.storage.postgres_repository_support import (
PostgresRepositoryConnectionProvider,
normalize_aware_datetime,
normalize_non_empty_text,
normalize_venue,
validate_decimal,
)
CANONICAL_CANDLE_SCHEMA_VERSION = 1
_INSERT_CANDLE_REVISION_SQL = """
INSERT INTO market_data.candle_revisions (
venue,
symbol,
interval,
open_time,
observed_at,
open_price,
high_price,
low_price,
close_price,
volume,
is_final,
source,
observation_sources,
canonical_schema_version
)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (venue, symbol, interval, open_time, observed_at) DO NOTHING
RETURNING 1
"""
_SELECT_CANDLE_REVISION_FOR_UPDATE_SQL = """
SELECT
open_price,
high_price,
low_price,
close_price,
volume,
is_final,
observation_sources,
canonical_schema_version
FROM market_data.candle_revisions
WHERE venue = %s
AND symbol = %s
AND interval = %s
AND open_time = %s
AND observed_at = %s
FOR UPDATE
"""
_UPDATE_CANDLE_PROVENANCE_SQL = """
UPDATE market_data.candle_revisions
SET observation_sources = %s
WHERE venue = %s
AND symbol = %s
AND interval = %s
AND open_time = %s
AND observed_at = %s
"""
@dataclass(frozen=True, slots=True)
class _PreparedCandleRevision:
symbol: str
interval: str
open_time: datetime
observed_at: datetime
open_price: Decimal
high_price: Decimal
low_price: Decimal
close_price: Decimal
volume: Decimal
is_final: bool
source: str
class PostgresCandleRepository:
"""Транзакционный модуль записи ревизий свечей в PostgreSQL."""
__slots__ = ("_connection_provider",)
def __init__(
self,
*,
connection_provider: PostgresRepositoryConnectionProvider,
) -> None:
if not callable(connection_provider):
raise TypeError("connection_provider must be callable")
self._connection_provider = connection_provider
def store_candle_revision(
self,
*,
venue: str,
candle: Candle,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
normalized_venue = normalize_venue(venue)
prepared = self._prepare_candle_revision(
candle=candle,
observed_at=observed_at,
is_final=is_final,
)
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
status = self._store_prepared_candle_revision(
cursor=cursor,
venue=normalized_venue,
candle=prepared,
)
except MarketDataStorageError:
raise
except Exception as error:
raise MarketDataStorageOperationError(
"Failed to store Canonical Candle revision."
) from error
return MarketDataWriteResult(status=status)
def _store_prepared_candle_revision(
self,
*,
cursor: Any,
venue: str,
candle: _PreparedCandleRevision,
) -> MarketDataWriteStatus:
cursor.execute(
_INSERT_CANDLE_REVISION_SQL,
(
venue,
candle.symbol,
candle.interval,
candle.open_time,
candle.observed_at,
candle.open_price,
candle.high_price,
candle.low_price,
candle.close_price,
candle.volume,
candle.is_final,
candle.source,
[candle.source],
CANONICAL_CANDLE_SCHEMA_VERSION,
),
)
if cursor.fetchone() is not None:
return MarketDataWriteStatus.INSERTED
cursor.execute(
_SELECT_CANDLE_REVISION_FOR_UPDATE_SQL,
self._identity_parameters(
venue=venue,
candle=candle,
),
)
existing = cursor.fetchone()
if existing is None:
raise MarketDataStorageOperationError(
"Conflicting Candle revision disappeared before locking."
)
(
existing_open_price,
existing_high_price,
existing_low_price,
existing_close_price,
existing_volume,
existing_is_final,
existing_sources,
existing_schema_version,
) = existing
if (
existing_schema_version != CANONICAL_CANDLE_SCHEMA_VERSION
or existing_open_price != candle.open_price
or existing_high_price != candle.high_price
or existing_low_price != candle.low_price
or existing_close_price != candle.close_price
or existing_volume != candle.volume
or existing_is_final is not candle.is_final
):
raise MarketDataStorageConflictError(
"Candle revision identity points to conflicting facts: "
f"venue={venue!r}, symbol={candle.symbol!r}, "
f"interval={candle.interval!r}, "
f"open_time={candle.open_time!r}, "
f"observed_at={candle.observed_at!r}."
)
observation_sources = tuple(existing_sources)
if candle.source in observation_sources:
return MarketDataWriteStatus.DUPLICATE
observation_sources += (candle.source,)
cursor.execute(
_UPDATE_CANDLE_PROVENANCE_SQL,
(
list(observation_sources),
*self._identity_parameters(
venue=venue,
candle=candle,
),
),
)
return MarketDataWriteStatus.PROVENANCE_UPDATED
@staticmethod
def _identity_parameters(
*,
venue: str,
candle: _PreparedCandleRevision,
) -> tuple[str, str, str, datetime, datetime]:
return (
venue,
candle.symbol,
candle.interval,
candle.open_time,
candle.observed_at,
)
@staticmethod
def _prepare_candle_revision(
*,
candle: Candle,
observed_at: datetime,
is_final: bool,
) -> _PreparedCandleRevision:
if not isinstance(candle, Candle):
raise MarketDataStorageValidationError(
"candle must be a Canonical Candle"
)
if not isinstance(is_final, bool):
raise MarketDataStorageValidationError(
"is_final must be a boolean"
)
symbol = normalize_non_empty_text(
candle.symbol,
field_name="candle.symbol",
).upper()
interval = normalize_non_empty_text(
candle.interval,
field_name="candle.interval",
)
source = normalize_non_empty_text(
candle.source,
field_name="candle.source",
)
open_time = normalize_aware_datetime(
candle.open_time,
field_name="candle.open_time",
)
normalized_observed_at = normalize_aware_datetime(
observed_at,
field_name="observed_at",
)
if normalized_observed_at < open_time:
raise MarketDataStorageValidationError(
"observed_at must not precede candle.open_time"
)
open_price = validate_decimal(
candle.open_price,
field_name="candle.open_price",
)
high_price = validate_decimal(
candle.high_price,
field_name="candle.high_price",
)
low_price = validate_decimal(
candle.low_price,
field_name="candle.low_price",
)
close_price = validate_decimal(
candle.close_price,
field_name="candle.close_price",
)
volume = validate_decimal(
candle.volume,
field_name="candle.volume",
allow_zero=True,
)
if low_price > high_price:
raise MarketDataStorageValidationError(
"candle.low_price must not exceed candle.high_price"
)
if not low_price <= open_price <= high_price:
raise MarketDataStorageValidationError(
"candle.open_price must be within low/high range"
)
if not low_price <= close_price <= high_price:
raise MarketDataStorageValidationError(
"candle.close_price must be within low/high range"
)
return _PreparedCandleRevision(
symbol=symbol,
interval=interval,
open_time=open_time,
observed_at=normalized_observed_at,
open_price=open_price,
high_price=high_price,
low_price=low_price,
close_price=close_price,
volume=volume,
is_final=is_final,
source=source,
)

View File

@@ -0,0 +1,745 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from enum import Enum
from typing import Any
from psycopg import sql
from src.market_data.storage.exceptions import (
MarketDataStorageConfigurationError,
MarketDataStorageError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
)
from src.market_data.storage.postgres_repository_support import (
PostgresRepositoryConnectionProvider,
normalize_aware_datetime,
)
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID = 0x445A504152544E
_SCHEMA_NAME = "market_data"
class MarketDataPartitionType(Enum):
"""Контролируемый набор секционированных таблиц рыночных данных."""
TRADES = "trades"
QUOTES = "quotes"
CANDLE_REVISIONS = "candle_revisions"
@dataclass(frozen=True, slots=True)
class _PartitionSpec:
parent_table: str
default_table: str
time_column: str
_PARTITION_SPECS = {
MarketDataPartitionType.TRADES: _PartitionSpec(
parent_table="trades",
default_table="trades_default",
time_column="executed_at",
),
MarketDataPartitionType.QUOTES: _PartitionSpec(
parent_table="quotes",
default_table="quotes_default",
time_column="received_at",
),
MarketDataPartitionType.CANDLE_REVISIONS: _PartitionSpec(
parent_table="candle_revisions",
default_table="candle_revisions_default",
time_column="open_time",
),
}
_SELECT_REGISTRY_ENTRY_SQL = """
SELECT partition_name, range_start, range_end, partition_bound
FROM market_data.partition_registry
WHERE data_type = %s
AND range_start = %s
FOR UPDATE
"""
_SELECT_EXPIRED_REGISTRY_ENTRIES_SQL = """
SELECT partition_name, range_start, range_end, partition_bound
FROM market_data.partition_registry
WHERE data_type = %s
AND range_end <= %s
ORDER BY range_start
FOR UPDATE
"""
_SELECT_RELATION_STATE_SQL = """
SELECT
EXISTS (
SELECT 1
FROM pg_catalog.pg_class AS child
JOIN pg_catalog.pg_namespace AS child_namespace
ON child_namespace.oid = child.relnamespace
WHERE child_namespace.nspname = %s
AND child.relname = %s
),
EXISTS (
SELECT 1
FROM pg_catalog.pg_inherits AS inheritance
JOIN pg_catalog.pg_class AS parent
ON parent.oid = inheritance.inhparent
JOIN pg_catalog.pg_namespace AS parent_namespace
ON parent_namespace.oid = parent.relnamespace
JOIN pg_catalog.pg_class AS child
ON child.oid = inheritance.inhrelid
JOIN pg_catalog.pg_namespace AS child_namespace
ON child_namespace.oid = child.relnamespace
WHERE parent_namespace.nspname = %s
AND parent.relname = %s
AND child_namespace.nspname = %s
AND child.relname = %s
),
(
SELECT pg_catalog.pg_get_expr(child.relpartbound, child.oid)
FROM pg_catalog.pg_inherits AS inheritance
JOIN pg_catalog.pg_class AS parent
ON parent.oid = inheritance.inhparent
JOIN pg_catalog.pg_namespace AS parent_namespace
ON parent_namespace.oid = parent.relnamespace
JOIN pg_catalog.pg_class AS child
ON child.oid = inheritance.inhrelid
JOIN pg_catalog.pg_namespace AS child_namespace
ON child_namespace.oid = child.relnamespace
WHERE parent_namespace.nspname = %s
AND parent.relname = %s
AND child_namespace.nspname = %s
AND child.relname = %s
)
"""
_INSERT_REGISTRY_ENTRY_SQL = """
INSERT INTO market_data.partition_registry (
data_type,
partition_name,
range_start,
range_end,
partition_bound
)
VALUES (%s, %s, %s, %s, %s)
"""
_DELETE_REGISTRY_ENTRY_SQL = """
DELETE FROM market_data.partition_registry
WHERE data_type = %s
AND partition_name = %s
AND range_start = %s
AND range_end = %s
"""
@dataclass(frozen=True, slots=True)
class MarketDataPartition:
"""Одна детерминированная граница месячной секции в UTC."""
data_type: MarketDataPartitionType
partition_name: str
range_start: datetime
range_end: datetime
@dataclass(frozen=True, slots=True)
class MarketDataPartitionResult:
"""Результат проверки существования одной месячной секции."""
partition: MarketDataPartition
created: bool
moved_row_count: int
def __post_init__(self) -> None:
if not isinstance(self.partition, MarketDataPartition):
raise TypeError("partition must be MarketDataPartition")
if not isinstance(self.created, bool):
raise TypeError("created must be a boolean")
if (
isinstance(self.moved_row_count, bool)
or not isinstance(self.moved_row_count, int)
):
raise TypeError("moved_row_count must be an integer")
if self.moved_row_count < 0:
raise ValueError("moved_row_count must not be negative")
@dataclass(frozen=True, slots=True)
class MarketDataRetentionPolicy:
"""Явные окна хранения; None означает отсутствие ограничения."""
enabled: bool = False
trade_days: int | None = None
quote_days: int | None = None
candle_days: int | None = None
def __post_init__(self) -> None:
if not isinstance(self.enabled, bool):
raise TypeError("enabled must be a boolean")
for field_name, value in (
("trade_days", self.trade_days),
("quote_days", self.quote_days),
("candle_days", self.candle_days),
):
if value is None:
continue
if isinstance(value, bool) or not isinstance(value, int):
raise TypeError(f"{field_name} must be an integer or None")
if value <= 0:
raise MarketDataStorageConfigurationError(
f"{field_name} must be positive when configured"
)
if self.enabled and all(
value is None
for value in (
self.trade_days,
self.quote_days,
self.candle_days,
)
):
raise MarketDataStorageConfigurationError(
"enabled retention requires at least one retention window"
)
def days_for(self, data_type: MarketDataPartitionType) -> int | None:
if not isinstance(data_type, MarketDataPartitionType):
raise TypeError("data_type must be MarketDataPartitionType")
return {
MarketDataPartitionType.TRADES: self.trade_days,
MarketDataPartitionType.QUOTES: self.quote_days,
MarketDataPartitionType.CANDLE_REVISIONS: self.candle_days,
}[data_type]
@dataclass(frozen=True, slots=True)
class MarketDataRetentionEntryResult:
"""Результат очистки одного типа канонических рыночных данных."""
data_type: MarketDataPartitionType
cutoff: datetime
dropped_partitions: tuple[str, ...]
dropped_row_count: int
deleted_row_count: int
@property
def total_removed_count(self) -> int:
return self.dropped_row_count + self.deleted_row_count
@dataclass(frozen=True, slots=True)
class MarketDataRetentionResult:
"""Атомарный результат одного явно запрошенного запуска очистки."""
entries: tuple[MarketDataRetentionEntryResult, ...]
@property
def total_removed_count(self) -> int:
return sum(entry.total_removed_count for entry in self.entries)
def build_monthly_partition(
*,
data_type: MarketDataPartitionType,
month: datetime,
) -> MarketDataPartition:
"""Создать детерминированное описание месячной секции в UTC."""
if not isinstance(data_type, MarketDataPartitionType):
raise MarketDataStorageValidationError(
"data_type must be MarketDataPartitionType"
)
normalized_month = normalize_aware_datetime(
month,
field_name="month",
)
range_start = normalized_month.replace(
day=1,
hour=0,
minute=0,
second=0,
microsecond=0,
)
try:
if range_start.month == 12:
range_end = range_start.replace(
year=range_start.year + 1,
month=1,
)
else:
range_end = range_start.replace(month=range_start.month + 1)
except ValueError as error:
raise MarketDataStorageValidationError(
"month cannot be represented as a complete UTC month"
) from error
return MarketDataPartition(
data_type=data_type,
partition_name=(
f"{data_type.value}_{range_start.year:04d}_{range_start.month:02d}"
),
range_start=range_start,
range_end=range_end,
)
class _PostgresPartitionOperations:
@staticmethod
def relation_state(
*,
cursor: Any,
spec: _PartitionSpec,
partition_name: str,
) -> tuple[bool, bool, str | None]:
cursor.execute(
_SELECT_RELATION_STATE_SQL,
(
_SCHEMA_NAME,
partition_name,
_SCHEMA_NAME,
spec.parent_table,
_SCHEMA_NAME,
partition_name,
_SCHEMA_NAME,
spec.parent_table,
_SCHEMA_NAME,
partition_name,
),
)
row = cursor.fetchone()
if not isinstance(row, tuple) or len(row) != 3:
raise MarketDataStorageOperationError(
"PostgreSQL returned invalid partition relation state."
)
partition_bound = row[2]
if partition_bound is not None and not isinstance(
partition_bound,
str,
):
raise MarketDataStorageOperationError(
"PostgreSQL returned invalid partition bound state."
)
return bool(row[0]), bool(row[1]), partition_bound
@classmethod
def assert_registered_partition(
cls,
*,
cursor: Any,
partition: MarketDataPartition,
registry_row: tuple[Any, ...],
) -> None:
if len(registry_row) != 4 or registry_row[:3] != (
partition.partition_name,
partition.range_start,
partition.range_end,
):
raise MarketDataStorageConfigurationError(
"Partition registry entry does not match the expected "
f"UTC month: {partition.partition_name!r}."
)
spec = _PARTITION_SPECS[partition.data_type]
relation_exists, is_attached, actual_bound = cls.relation_state(
cursor=cursor,
spec=spec,
partition_name=partition.partition_name,
)
if not relation_exists or not is_attached:
raise MarketDataStorageConfigurationError(
"Registered partition is missing or detached in PostgreSQL: "
f"{partition.partition_name!r}."
)
registered_bound = registry_row[3]
if (
not isinstance(registered_bound, str)
or not registered_bound.strip()
or actual_bound != registered_bound
):
raise MarketDataStorageConfigurationError(
"Registered partition bound differs from PostgreSQL: "
f"{partition.partition_name!r}."
)
class PostgresMarketDataPartitionManager:
"""Создаёт и подключает контролируемые месячные секции PostgreSQL."""
__slots__ = ("_connection_provider",)
def __init__(
self,
*,
connection_provider: PostgresRepositoryConnectionProvider,
) -> None:
if not callable(connection_provider):
raise TypeError("connection_provider must be callable")
self._connection_provider = connection_provider
def ensure_month_partition(
self,
*,
data_type: MarketDataPartitionType,
month: datetime,
) -> MarketDataPartitionResult:
partition = build_monthly_partition(
data_type=data_type,
month=month,
)
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_xact_lock(%s)",
(MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,),
)
cursor.execute(
_SELECT_REGISTRY_ENTRY_SQL,
(
partition.data_type.value,
partition.range_start,
),
)
registry_row = cursor.fetchone()
if registry_row is not None:
_PostgresPartitionOperations.assert_registered_partition(
cursor=cursor,
partition=partition,
registry_row=registry_row,
)
return MarketDataPartitionResult(
partition=partition,
created=False,
moved_row_count=0,
)
spec = _PARTITION_SPECS[partition.data_type]
relation_exists, is_attached, _ = (
_PostgresPartitionOperations.relation_state(
cursor=cursor,
spec=spec,
partition_name=partition.partition_name,
)
)
if relation_exists or is_attached:
raise MarketDataStorageConfigurationError(
"Unregistered PostgreSQL relation blocks managed "
f"partition {partition.partition_name!r}."
)
moved_row_count = self._create_partition(
cursor=cursor,
partition=partition,
spec=spec,
)
except MarketDataStorageError:
raise
except Exception as error:
raise MarketDataStorageOperationError(
"Failed to ensure monthly Market Data partition."
) from error
return MarketDataPartitionResult(
partition=partition,
created=True,
moved_row_count=moved_row_count,
)
@staticmethod
def _create_partition(
*,
cursor: Any,
partition: MarketDataPartition,
spec: _PartitionSpec,
) -> int:
parent = sql.Identifier(_SCHEMA_NAME, spec.parent_table)
default = sql.Identifier(_SCHEMA_NAME, spec.default_table)
child = sql.Identifier(_SCHEMA_NAME, partition.partition_name)
time_column = sql.Identifier(spec.time_column)
constraint_name = sql.Identifier(
f"{partition.partition_name}_{spec.time_column}_range"
)
range_start = sql.Literal(partition.range_start)
range_end = sql.Literal(partition.range_end)
cursor.execute(
sql.SQL("LOCK TABLE {} IN ACCESS EXCLUSIVE MODE").format(default)
)
cursor.execute(
sql.SQL("CREATE TABLE {} (LIKE {} INCLUDING ALL)").format(
child,
parent,
)
)
cursor.execute(
sql.SQL(
"ALTER TABLE {} ADD CONSTRAINT {} "
"CHECK ({} >= {} AND {} < {})"
).format(
child,
constraint_name,
time_column,
range_start,
time_column,
range_end,
),
)
cursor.execute(
sql.SQL(
"WITH moved_rows AS ("
"DELETE FROM {} WHERE {} >= %s AND {} < %s RETURNING *"
") INSERT INTO {} SELECT * FROM moved_rows"
).format(
default,
time_column,
time_column,
child,
),
(partition.range_start, partition.range_end),
)
moved_row_count = cursor.rowcount
if (
isinstance(moved_row_count, bool)
or not isinstance(moved_row_count, int)
or moved_row_count < 0
):
raise MarketDataStorageOperationError(
"PostgreSQL did not report the moved partition row count."
)
cursor.execute(
sql.SQL(
"ALTER TABLE {} ATTACH PARTITION {} "
"FOR VALUES FROM ({}) TO ({})"
).format(parent, child, range_start, range_end),
)
relation_exists, is_attached, partition_bound = (
_PostgresPartitionOperations.relation_state(
cursor=cursor,
spec=spec,
partition_name=partition.partition_name,
)
)
if (
not relation_exists
or not is_attached
or not isinstance(partition_bound, str)
or not partition_bound.strip()
):
raise MarketDataStorageOperationError(
"Attached partition has no canonical PostgreSQL bound."
)
cursor.execute(
_INSERT_REGISTRY_ENTRY_SQL,
(
partition.data_type.value,
partition.partition_name,
partition.range_start,
partition.range_end,
partition_bound,
),
)
return moved_row_count
class PostgresMarketDataRetentionService:
"""Применяет явную транзакционную политику хранения данных."""
__slots__ = ("_connection_provider",)
def __init__(
self,
*,
connection_provider: PostgresRepositoryConnectionProvider,
) -> None:
if not callable(connection_provider):
raise TypeError("connection_provider must be callable")
self._connection_provider = connection_provider
def apply(
self,
*,
policy: MarketDataRetentionPolicy,
now: datetime,
) -> MarketDataRetentionResult:
if not isinstance(policy, MarketDataRetentionPolicy):
raise TypeError("policy must be MarketDataRetentionPolicy")
if not policy.enabled:
return MarketDataRetentionResult(entries=())
normalized_now = normalize_aware_datetime(now, field_name="now")
configured = tuple(
(data_type, days)
for data_type in MarketDataPartitionType
if (days := policy.days_for(data_type)) is not None
)
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_xact_lock(%s)",
(MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,),
)
for data_type, _ in configured:
spec = _PARTITION_SPECS[data_type]
cursor.execute(
sql.SQL(
"LOCK TABLE {} "
"IN SHARE ROW EXCLUSIVE MODE"
).format(
sql.Identifier(
_SCHEMA_NAME,
spec.parent_table,
)
)
)
entries = tuple(
self._apply_one(
cursor=cursor,
data_type=data_type,
cutoff=normalized_now - timedelta(days=days),
)
for data_type, days in configured
)
except MarketDataStorageError:
raise
except Exception as error:
raise MarketDataStorageOperationError(
"Failed to apply Market Data retention policy."
) from error
return MarketDataRetentionResult(entries=entries)
@staticmethod
def _apply_one(
*,
cursor: Any,
data_type: MarketDataPartitionType,
cutoff: datetime,
) -> MarketDataRetentionEntryResult:
spec = _PARTITION_SPECS[data_type]
cursor.execute(
_SELECT_EXPIRED_REGISTRY_ENTRIES_SQL,
(data_type.value, cutoff),
)
registry_rows = tuple(cursor.fetchall())
dropped_partitions: list[str] = []
dropped_row_count = 0
for registry_row in registry_rows:
if not isinstance(registry_row, tuple) or len(registry_row) != 4:
raise MarketDataStorageOperationError(
"PostgreSQL returned an invalid partition registry row."
)
partition_name, range_start, range_end, _ = registry_row
partition = build_monthly_partition(
data_type=data_type,
month=range_start,
)
_PostgresPartitionOperations.assert_registered_partition(
cursor=cursor,
partition=partition,
registry_row=registry_row,
)
cursor.execute(
sql.SQL("SELECT COUNT(*) FROM {}").format(
sql.Identifier(_SCHEMA_NAME, partition_name)
)
)
count_row = cursor.fetchone()
if (
not isinstance(count_row, tuple)
or len(count_row) != 1
or isinstance(count_row[0], bool)
or not isinstance(count_row[0], int)
or count_row[0] < 0
):
raise MarketDataStorageOperationError(
"PostgreSQL returned an invalid partition row count."
)
dropped_row_count += count_row[0]
cursor.execute(
sql.SQL("DROP TABLE {}").format(
sql.Identifier(_SCHEMA_NAME, partition_name)
)
)
cursor.execute(
_DELETE_REGISTRY_ENTRY_SQL,
(
data_type.value,
partition_name,
range_start,
range_end,
),
)
if cursor.rowcount != 1:
raise MarketDataStorageOperationError(
"Partition registry entry disappeared during retention."
)
dropped_partitions.append(partition_name)
parent = sql.Identifier(_SCHEMA_NAME, spec.parent_table)
time_column = sql.Identifier(spec.time_column)
cursor.execute(
sql.SQL("DELETE FROM {} WHERE {} < %s").format(
parent,
time_column,
),
(cutoff,),
)
deleted_row_count = cursor.rowcount
if (
isinstance(deleted_row_count, bool)
or not isinstance(deleted_row_count, int)
or deleted_row_count < 0
):
raise MarketDataStorageOperationError(
"PostgreSQL did not report the retained row count."
)
return MarketDataRetentionEntryResult(
data_type=data_type,
cutoff=cutoff,
dropped_partitions=tuple(dropped_partitions),
dropped_row_count=dropped_row_count,
deleted_row_count=deleted_row_count,
)

View File

@@ -0,0 +1,270 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
from typing import Any
from src.market_data.acquisition.models.quote import Quote
from src.market_data.storage.contracts import (
MarketDataWriteResult,
MarketDataWriteStatus,
)
from src.market_data.storage.exceptions import (
MarketDataStorageConflictError,
MarketDataStorageError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
)
from src.market_data.storage.postgres_repository_support import (
PostgresRepositoryConnectionProvider,
normalize_aware_datetime,
normalize_non_empty_text,
normalize_venue,
validate_decimal,
)
CANONICAL_QUOTE_SCHEMA_VERSION = 1
_INSERT_QUOTE_SQL = """
INSERT INTO market_data.quotes (
venue,
symbol,
received_at,
exchange_timestamp,
last_price,
bid_price,
ask_price,
source,
observation_sources,
canonical_schema_version
)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (venue, symbol, received_at) DO NOTHING
RETURNING 1
"""
_SELECT_QUOTE_FOR_UPDATE_SQL = """
SELECT
exchange_timestamp,
last_price,
bid_price,
ask_price,
observation_sources,
canonical_schema_version
FROM market_data.quotes
WHERE venue = %s
AND symbol = %s
AND received_at = %s
FOR UPDATE
"""
_UPDATE_QUOTE_PROVENANCE_SQL = """
UPDATE market_data.quotes
SET observation_sources = %s
WHERE venue = %s
AND symbol = %s
AND received_at = %s
"""
@dataclass(frozen=True, slots=True)
class _PreparedQuote:
symbol: str
received_at: datetime
exchange_timestamp: datetime | None
last_price: Decimal
bid_price: Decimal
ask_price: Decimal
source: str
class PostgresQuoteRepository:
"""Транзакционный модуль записи снимков котировок в PostgreSQL."""
__slots__ = ("_connection_provider",)
def __init__(
self,
*,
connection_provider: PostgresRepositoryConnectionProvider,
) -> None:
if not callable(connection_provider):
raise TypeError("connection_provider must be callable")
self._connection_provider = connection_provider
def store_quote(
self,
*,
venue: str,
quote: Quote,
) -> MarketDataWriteResult:
normalized_venue = normalize_venue(venue)
prepared = self._prepare_quote(quote)
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
status = self._store_prepared_quote(
cursor=cursor,
venue=normalized_venue,
quote=prepared,
)
except MarketDataStorageError:
raise
except Exception as error:
raise MarketDataStorageOperationError(
"Failed to store Canonical Quote."
) from error
return MarketDataWriteResult(status=status)
def _store_prepared_quote(
self,
*,
cursor: Any,
venue: str,
quote: _PreparedQuote,
) -> MarketDataWriteStatus:
cursor.execute(
_INSERT_QUOTE_SQL,
(
venue,
quote.symbol,
quote.received_at,
quote.exchange_timestamp,
quote.last_price,
quote.bid_price,
quote.ask_price,
quote.source,
[quote.source],
CANONICAL_QUOTE_SCHEMA_VERSION,
),
)
if cursor.fetchone() is not None:
return MarketDataWriteStatus.INSERTED
cursor.execute(
_SELECT_QUOTE_FOR_UPDATE_SQL,
self._identity_parameters(
venue=venue,
quote=quote,
),
)
existing = cursor.fetchone()
if existing is None:
raise MarketDataStorageOperationError(
"Conflicting Quote disappeared before it could be locked."
)
(
existing_exchange_timestamp,
existing_last_price,
existing_bid_price,
existing_ask_price,
existing_sources,
existing_schema_version,
) = existing
if (
existing_schema_version != CANONICAL_QUOTE_SCHEMA_VERSION
or existing_exchange_timestamp != quote.exchange_timestamp
or existing_last_price != quote.last_price
or existing_bid_price != quote.bid_price
or existing_ask_price != quote.ask_price
):
raise MarketDataStorageConflictError(
"Quote identity points to conflicting canonical facts: "
f"venue={venue!r}, symbol={quote.symbol!r}, "
f"received_at={quote.received_at!r}."
)
observation_sources = tuple(existing_sources)
if quote.source in observation_sources:
return MarketDataWriteStatus.DUPLICATE
observation_sources += (quote.source,)
cursor.execute(
_UPDATE_QUOTE_PROVENANCE_SQL,
(
list(observation_sources),
*self._identity_parameters(
venue=venue,
quote=quote,
),
),
)
return MarketDataWriteStatus.PROVENANCE_UPDATED
@staticmethod
def _identity_parameters(
*,
venue: str,
quote: _PreparedQuote,
) -> tuple[str, str, datetime]:
return (
venue,
quote.symbol,
quote.received_at,
)
@staticmethod
def _prepare_quote(quote: Quote) -> _PreparedQuote:
if not isinstance(quote, Quote):
raise MarketDataStorageValidationError(
"quote must be a Canonical Quote"
)
symbol = normalize_non_empty_text(
quote.symbol,
field_name="quote.symbol",
).upper()
source = normalize_non_empty_text(
quote.source,
field_name="quote.source",
)
last_price = validate_decimal(
quote.last_price,
field_name="quote.last_price",
)
bid_price = validate_decimal(
quote.bid_price,
field_name="quote.bid_price",
)
ask_price = validate_decimal(
quote.ask_price,
field_name="quote.ask_price",
)
if bid_price > ask_price:
raise MarketDataStorageValidationError(
"quote.bid_price must not exceed quote.ask_price"
)
exchange_timestamp = (
normalize_aware_datetime(
quote.exchange_timestamp,
field_name="quote.exchange_timestamp",
)
if quote.exchange_timestamp is not None
else None
)
return _PreparedQuote(
symbol=symbol,
received_at=normalize_aware_datetime(
quote.received_at,
field_name="quote.received_at",
),
exchange_timestamp=exchange_timestamp,
last_price=last_price,
bid_price=bid_price,
ask_price=ask_price,
source=source,
)

View File

@@ -0,0 +1,85 @@
from __future__ import annotations
from collections.abc import Callable
from contextlib import AbstractContextManager
from datetime import datetime, timezone
from decimal import Decimal
from typing import Any
from src.market_data.storage.exceptions import (
MarketDataStorageValidationError,
)
PostgresRepositoryConnectionProvider = Callable[
[],
AbstractContextManager[Any],
]
def normalize_venue(venue: str) -> str:
return normalize_non_empty_text(
venue,
field_name="venue",
)
def normalize_non_empty_text(
value: str,
*,
field_name: str,
) -> str:
if not isinstance(value, str):
raise MarketDataStorageValidationError(
f"{field_name} must be a string"
)
normalized = value.strip()
if not normalized:
raise MarketDataStorageValidationError(
f"{field_name} must not be empty"
)
return normalized
def normalize_aware_datetime(
value: datetime,
*,
field_name: str,
) -> datetime:
if (
not isinstance(value, datetime)
or value.tzinfo is None
or value.utcoffset() is None
):
raise MarketDataStorageValidationError(
f"{field_name} must be timezone-aware datetime"
)
return value.astimezone(timezone.utc)
def validate_decimal(
value: Decimal,
*,
field_name: str,
allow_zero: bool = False,
) -> Decimal:
if not isinstance(value, Decimal) or not value.is_finite():
raise MarketDataStorageValidationError(
f"{field_name} must be a finite Decimal"
)
if allow_zero:
if value < 0:
raise MarketDataStorageValidationError(
f"{field_name} must not be negative"
)
elif value <= 0:
raise MarketDataStorageValidationError(
f"{field_name} must be positive"
)
return value

View File

@@ -0,0 +1,391 @@
from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime
from decimal import Decimal
from typing import Any
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.acquisition.trade_id_sequence import (
validate_signed_trade_id,
)
from src.market_data.storage.contracts import (
MarketDataBatchWriteResult,
MarketDataWriteResult,
MarketDataWriteStatus,
)
from src.market_data.storage.exceptions import (
MarketDataStorageConflictError,
MarketDataStorageError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
)
from src.market_data.storage.postgres_repository_support import (
PostgresRepositoryConnectionProvider,
normalize_aware_datetime,
normalize_non_empty_text,
normalize_venue,
validate_decimal,
)
CANONICAL_TRADE_SCHEMA_VERSION = 1
_INSERT_TRADE_SQL = """
INSERT INTO market_data.trades (
venue,
symbol,
trade_id,
executed_at,
price,
quantity,
aggressor_side,
source,
first_observed_at,
last_observed_at,
observation_sources,
canonical_schema_version
)
VALUES (
%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s
)
ON CONFLICT (venue, symbol, trade_id, executed_at) DO NOTHING
RETURNING 1
"""
_SELECT_TRADE_FOR_UPDATE_SQL = """
SELECT
price,
quantity,
aggressor_side,
first_observed_at,
last_observed_at,
observation_sources,
canonical_schema_version
FROM market_data.trades
WHERE venue = %s
AND symbol = %s
AND trade_id = %s
AND executed_at = %s
FOR UPDATE
"""
_UPDATE_TRADE_PROVENANCE_SQL = """
UPDATE market_data.trades
SET first_observed_at = %s,
last_observed_at = %s,
observation_sources = %s
WHERE venue = %s
AND symbol = %s
AND trade_id = %s
AND executed_at = %s
"""
@dataclass(frozen=True, slots=True)
class _PreparedTrade:
symbol: str
trade_id: int
executed_at: datetime
price: Decimal
quantity: Decimal
aggressor_side: str
source: str
observed_at: datetime
@property
def identity_order_key(self) -> tuple[str, datetime, int]:
return (
self.symbol,
self.executed_at,
self.trade_id,
)
class PostgresTradeRepository:
"""Транзакционный модуль записи канонических сделок в PostgreSQL."""
__slots__ = ("_connection_provider",)
def __init__(
self,
*,
connection_provider: PostgresRepositoryConnectionProvider,
) -> None:
if not callable(connection_provider):
raise TypeError("connection_provider must be callable")
self._connection_provider = connection_provider
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
normalized_venue = normalize_venue(venue)
prepared = self._prepare_trade(
trade=trade,
observed_at=observed_at,
)
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
status = self._store_prepared_trade(
cursor=cursor,
venue=normalized_venue,
trade=prepared,
)
except MarketDataStorageError:
raise
except Exception as error:
raise MarketDataStorageOperationError(
"Failed to store Canonical Trade."
) from error
return MarketDataWriteResult(status=status)
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
normalized_venue = normalize_venue(venue)
if not isinstance(trades, tuple):
raise MarketDataStorageValidationError(
"trades must be a tuple"
)
normalized_observed_at = normalize_aware_datetime(
observed_at,
field_name="observed_at",
)
prepared_trades = tuple(
self._prepare_trade(
trade=trade,
observed_at=normalized_observed_at,
)
for trade in trades
)
if not prepared_trades:
return MarketDataBatchWriteResult(
inserted_count=0,
duplicate_count=0,
provenance_updated_count=0,
)
inserted_count = 0
duplicate_count = 0
provenance_updated_count = 0
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
for prepared in sorted(
prepared_trades,
key=lambda item: item.identity_order_key,
):
status = self._store_prepared_trade(
cursor=cursor,
venue=normalized_venue,
trade=prepared,
)
if status is MarketDataWriteStatus.INSERTED:
inserted_count += 1
elif status is MarketDataWriteStatus.DUPLICATE:
duplicate_count += 1
else:
provenance_updated_count += 1
except MarketDataStorageError:
raise
except Exception as error:
raise MarketDataStorageOperationError(
"Failed to atomically store Canonical Trade batch."
) from error
return MarketDataBatchWriteResult(
inserted_count=inserted_count,
duplicate_count=duplicate_count,
provenance_updated_count=provenance_updated_count,
)
def _store_prepared_trade(
self,
*,
cursor: Any,
venue: str,
trade: _PreparedTrade,
) -> MarketDataWriteStatus:
cursor.execute(
_INSERT_TRADE_SQL,
(
venue,
trade.symbol,
trade.trade_id,
trade.executed_at,
trade.price,
trade.quantity,
trade.aggressor_side,
trade.source,
trade.observed_at,
trade.observed_at,
[trade.source],
CANONICAL_TRADE_SCHEMA_VERSION,
),
)
if cursor.fetchone() is not None:
return MarketDataWriteStatus.INSERTED
cursor.execute(
_SELECT_TRADE_FOR_UPDATE_SQL,
self._identity_parameters(
venue=venue,
trade=trade,
),
)
existing = cursor.fetchone()
if existing is None:
raise MarketDataStorageOperationError(
"Conflicting Trade disappeared before it could be locked."
)
(
existing_price,
existing_quantity,
existing_aggressor_side,
existing_first_observed_at,
existing_last_observed_at,
existing_sources,
existing_schema_version,
) = existing
if (
existing_schema_version != CANONICAL_TRADE_SCHEMA_VERSION
or existing_price != trade.price
or existing_quantity != trade.quantity
or existing_aggressor_side != trade.aggressor_side
):
raise MarketDataStorageConflictError(
"Trade identity points to conflicting canonical facts: "
f"venue={venue!r}, symbol={trade.symbol!r}, "
f"trade_id={trade.trade_id!r}, "
f"executed_at={trade.executed_at!r}."
)
first_observed_at = min(
existing_first_observed_at,
trade.observed_at,
)
last_observed_at = max(
existing_last_observed_at,
trade.observed_at,
)
observation_sources = tuple(existing_sources)
if trade.source not in observation_sources:
observation_sources += (trade.source,)
if (
first_observed_at == existing_first_observed_at
and last_observed_at == existing_last_observed_at
and observation_sources == tuple(existing_sources)
):
return MarketDataWriteStatus.DUPLICATE
cursor.execute(
_UPDATE_TRADE_PROVENANCE_SQL,
(
first_observed_at,
last_observed_at,
list(observation_sources),
*self._identity_parameters(
venue=venue,
trade=trade,
),
),
)
return MarketDataWriteStatus.PROVENANCE_UPDATED
@staticmethod
def _identity_parameters(
*,
venue: str,
trade: _PreparedTrade,
) -> tuple[str, str, int, datetime]:
return (
venue,
trade.symbol,
trade.trade_id,
trade.executed_at,
)
def _prepare_trade(
self,
*,
trade: Trade,
observed_at: datetime,
) -> _PreparedTrade:
if not isinstance(trade, Trade):
raise MarketDataStorageValidationError(
"trade must be a Canonical Trade"
)
symbol = normalize_non_empty_text(
trade.symbol,
field_name="trade.symbol",
).upper()
source = normalize_non_empty_text(
trade.source,
field_name="trade.source",
)
try:
validate_signed_trade_id(trade.trade_id)
except (TypeError, ValueError) as error:
raise MarketDataStorageValidationError(
"trade.trade_id must fit signed 32-bit range"
) from error
price = validate_decimal(
trade.price,
field_name="trade.price",
)
quantity = validate_decimal(
trade.quantity,
field_name="trade.quantity",
)
if not isinstance(trade.aggressor_side, TradeAggressorSide):
raise MarketDataStorageValidationError(
"trade.aggressor_side must be TradeAggressorSide"
)
return _PreparedTrade(
symbol=symbol,
trade_id=trade.trade_id,
executed_at=normalize_aware_datetime(
trade.executed_at,
field_name="trade.executed_at",
),
price=price,
quantity=quantity,
aggressor_side=trade.aggressor_side.value,
source=source,
observed_at=normalize_aware_datetime(
observed_at,
field_name="observed_at",
),
)

View File

@@ -0,0 +1,66 @@
from __future__ import annotations
from collections.abc import Callable
from datetime import datetime, timezone
from src.market_data.acquisition.models.trade import Trade
from src.market_data.storage.contracts import (
MarketDataWriteResult,
TradeStorageProtocol,
)
from src.market_data.storage.postgres_repository_support import (
normalize_venue,
)
TradeObservationClock = Callable[[], datetime]
def system_utc_datetime() -> datetime:
"""Вернуть текущее время UTC с часовым поясом."""
return datetime.now(timezone.utc)
class TradeStorageObservationSink:
"""Адаптер наблюдений канонических сделок к постоянному хранилищу."""
__slots__ = (
"_trade_storage",
"_venue",
"_clock",
)
def __init__(
self,
*,
trade_storage: TradeStorageProtocol,
venue: str,
clock: TradeObservationClock = system_utc_datetime,
) -> None:
if not isinstance(trade_storage, TradeStorageProtocol):
raise TypeError(
"trade_storage must implement TradeStorageProtocol"
)
if not callable(clock):
raise TypeError("clock must be callable")
self._trade_storage = trade_storage
self._venue = normalize_venue(venue)
self._clock = clock
def persist(
self,
trade: Trade,
) -> None:
result = self._trade_storage.store_trade(
venue=self._venue,
trade=trade,
observed_at=self._clock(),
)
if not isinstance(result, MarketDataWriteResult):
raise TypeError(
"trade_storage.store_trade() must return "
"MarketDataWriteResult"
)

View File

@@ -16,3 +16,11 @@ class InstrumentStoreError(StorageError):
# Ошибка хранилища канонических котировок.
class QuoteStoreError(StorageError):
"""Quote store contract or operation error."""
class PostgresConnectionPoolError(StorageError):
"""PostgreSQL connection pool lifecycle error."""
class StorageMigrationError(StorageError):
"""Versioned storage migration error."""

View File

@@ -0,0 +1,429 @@
from __future__ import annotations
from collections.abc import Callable, Iterable
from contextlib import AbstractContextManager
from dataclasses import dataclass
from typing import Any
from src.storage.exceptions import StorageMigrationError
STORAGE_MIGRATION_ADVISORY_LOCK_ID = 0x445A454E545241
_CREATE_HISTORY_TABLE_SQL = """
CREATE TABLE IF NOT EXISTS public.storage_schema_migrations (
version INTEGER PRIMARY KEY,
name TEXT UNIQUE NOT NULL,
applied_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
)
"""
_SELECT_APPLIED_MIGRATIONS_SQL = """
SELECT version, name
FROM public.storage_schema_migrations
ORDER BY version
"""
_INSERT_APPLIED_MIGRATION_SQL = """
INSERT INTO public.storage_schema_migrations (version, name)
VALUES (%s, %s)
"""
@dataclass(frozen=True, slots=True)
class StorageMigration:
"""Одна неизменяемая упорядоченная миграция схемы хранилища."""
version: int
name: str
statements: tuple[str, ...]
def __post_init__(self) -> None:
if (
isinstance(self.version, bool)
or not isinstance(self.version, int)
or self.version <= 0
):
raise ValueError("migration version must be a positive integer")
normalized_name = str(self.name or "").strip()
if not normalized_name:
raise ValueError("migration name must not be empty")
if not isinstance(self.statements, tuple) or not self.statements:
raise ValueError("migration statements must be a non-empty tuple")
if any(not str(statement or "").strip() for statement in self.statements):
raise ValueError("migration statements must not be empty")
object.__setattr__(self, "name", normalized_name)
STORAGE_MIGRATIONS = (
StorageMigration(
version=1,
name="create_market_data_schema",
statements=(
"CREATE SCHEMA IF NOT EXISTS market_data",
),
),
StorageMigration(
version=2,
name="create_canonical_trades",
statements=(
"""
CREATE TABLE market_data.trades (
venue TEXT NOT NULL,
symbol TEXT NOT NULL,
trade_id INTEGER NOT NULL,
executed_at TIMESTAMPTZ NOT NULL,
price NUMERIC NOT NULL,
quantity NUMERIC NOT NULL,
aggressor_side TEXT NOT NULL,
source TEXT NOT NULL,
first_observed_at TIMESTAMPTZ NOT NULL,
last_observed_at TIMESTAMPTZ NOT NULL,
canonical_schema_version INTEGER NOT NULL DEFAULT 1,
PRIMARY KEY (venue, symbol, trade_id, executed_at),
CHECK (trade_id BETWEEN -2147483648 AND 2147483647),
CHECK (BTRIM(venue) <> ''),
CHECK (BTRIM(symbol) <> ''),
CHECK (BTRIM(source) <> ''),
CHECK (price > 0),
CHECK (quantity > 0),
CHECK (aggressor_side IN ('buy', 'sell')),
CHECK (canonical_schema_version > 0),
CHECK (last_observed_at >= first_observed_at)
) PARTITION BY RANGE (executed_at)
""",
"""
CREATE TABLE market_data.trades_default
PARTITION OF market_data.trades DEFAULT
""",
"""
CREATE INDEX trades_event_order_idx
ON market_data.trades (
venue,
symbol,
executed_at,
trade_id
)
""",
),
),
StorageMigration(
version=3,
name="create_canonical_quotes",
statements=(
"""
CREATE TABLE market_data.quotes (
venue TEXT NOT NULL,
symbol TEXT NOT NULL,
received_at TIMESTAMPTZ NOT NULL,
exchange_timestamp TIMESTAMPTZ,
last_price NUMERIC NOT NULL,
bid_price NUMERIC NOT NULL,
ask_price NUMERIC NOT NULL,
source TEXT NOT NULL,
canonical_schema_version INTEGER NOT NULL DEFAULT 1,
PRIMARY KEY (venue, symbol, received_at),
CHECK (BTRIM(venue) <> ''),
CHECK (BTRIM(symbol) <> ''),
CHECK (BTRIM(source) <> ''),
CHECK (last_price > 0),
CHECK (bid_price > 0),
CHECK (ask_price > 0),
CHECK (bid_price <= ask_price),
CHECK (canonical_schema_version > 0)
) PARTITION BY RANGE (received_at)
""",
"""
CREATE TABLE market_data.quotes_default
PARTITION OF market_data.quotes DEFAULT
""",
),
),
StorageMigration(
version=4,
name="create_canonical_candle_revisions",
statements=(
"""
CREATE TABLE market_data.candle_revisions (
venue TEXT NOT NULL,
symbol TEXT NOT NULL,
interval TEXT NOT NULL,
open_time TIMESTAMPTZ NOT NULL,
observed_at TIMESTAMPTZ NOT NULL,
open_price NUMERIC NOT NULL,
high_price NUMERIC NOT NULL,
low_price NUMERIC NOT NULL,
close_price NUMERIC NOT NULL,
volume NUMERIC NOT NULL,
is_final BOOLEAN NOT NULL,
source TEXT NOT NULL,
canonical_schema_version INTEGER NOT NULL DEFAULT 1,
PRIMARY KEY (
venue,
symbol,
interval,
open_time,
observed_at
),
CHECK (BTRIM(venue) <> ''),
CHECK (BTRIM(symbol) <> ''),
CHECK (BTRIM(interval) <> ''),
CHECK (BTRIM(source) <> ''),
CHECK (open_price > 0),
CHECK (high_price > 0),
CHECK (low_price > 0),
CHECK (close_price > 0),
CHECK (volume >= 0),
CHECK (low_price <= high_price),
CHECK (open_price BETWEEN low_price AND high_price),
CHECK (close_price BETWEEN low_price AND high_price),
CHECK (canonical_schema_version > 0),
CHECK (observed_at >= open_time)
) PARTITION BY RANGE (open_time)
""",
"""
CREATE TABLE market_data.candle_revisions_default
PARTITION OF market_data.candle_revisions DEFAULT
""",
),
),
StorageMigration(
version=5,
name="add_trade_observation_sources",
statements=(
"""
ALTER TABLE market_data.trades
ADD COLUMN observation_sources TEXT[]
""",
"""
UPDATE market_data.trades
SET observation_sources = ARRAY[source]
""",
"""
ALTER TABLE market_data.trades
ALTER COLUMN observation_sources SET NOT NULL
""",
"""
ALTER TABLE market_data.trades
ADD CONSTRAINT trades_observation_sources_not_empty
CHECK (
CARDINALITY(observation_sources) > 0
AND ARRAY_POSITION(observation_sources, NULL) IS NULL
)
""",
),
),
StorageMigration(
version=6,
name="add_quote_and_candle_observation_sources",
statements=(
"""
ALTER TABLE market_data.quotes
ADD COLUMN observation_sources TEXT[]
""",
"""
UPDATE market_data.quotes
SET observation_sources = ARRAY[source]
""",
"""
ALTER TABLE market_data.quotes
ALTER COLUMN observation_sources SET NOT NULL
""",
"""
ALTER TABLE market_data.quotes
ADD CONSTRAINT quotes_observation_sources_not_empty
CHECK (
CARDINALITY(observation_sources) > 0
AND ARRAY_POSITION(observation_sources, NULL) IS NULL
)
""",
"""
ALTER TABLE market_data.candle_revisions
ADD COLUMN observation_sources TEXT[]
""",
"""
UPDATE market_data.candle_revisions
SET observation_sources = ARRAY[source]
""",
"""
ALTER TABLE market_data.candle_revisions
ALTER COLUMN observation_sources SET NOT NULL
""",
"""
ALTER TABLE market_data.candle_revisions
ADD CONSTRAINT candle_revisions_observation_sources_not_empty
CHECK (
CARDINALITY(observation_sources) > 0
AND ARRAY_POSITION(observation_sources, NULL) IS NULL
)
""",
),
),
StorageMigration(
version=7,
name="create_market_data_partition_registry",
statements=(
"""
CREATE TABLE market_data.partition_registry (
data_type TEXT NOT NULL,
partition_name TEXT NOT NULL,
range_start TIMESTAMPTZ NOT NULL,
range_end TIMESTAMPTZ NOT NULL,
partition_bound TEXT NOT NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
PRIMARY KEY (data_type, range_start),
UNIQUE (partition_name),
CHECK (
data_type IN (
'trades',
'quotes',
'candle_revisions'
)
),
CHECK (BTRIM(partition_name) <> ''),
CHECK (BTRIM(partition_bound) <> ''),
CHECK (range_end > range_start)
)
""",
),
),
)
StorageConnectionProvider = Callable[[], AbstractContextManager[Any]]
class StorageMigrationRunner:
"""Применяет миграции PostgreSQL в одной заблокированной транзакции."""
__slots__ = (
"_connection_provider",
"_migrations",
)
def __init__(
self,
*,
connection_provider: StorageConnectionProvider,
migrations: Iterable[StorageMigration] = STORAGE_MIGRATIONS,
) -> None:
if not callable(connection_provider):
raise TypeError("connection_provider must be callable")
normalized_migrations = tuple(migrations)
versions = tuple(
migration.version
for migration in normalized_migrations
)
if len(set(versions)) != len(versions):
raise ValueError("migration versions must be unique")
if versions != tuple(sorted(versions)):
raise ValueError("migrations must be ordered by version")
self._connection_provider = connection_provider
self._migrations = normalized_migrations
def run(self) -> tuple[int, ...]:
"""Применить ожидающие миграции и вернуть их версии."""
try:
with self._connection_provider() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_xact_lock(%s)",
(STORAGE_MIGRATION_ADVISORY_LOCK_ID,),
)
cursor.execute(_CREATE_HISTORY_TABLE_SQL)
cursor.execute(_SELECT_APPLIED_MIGRATIONS_SQL)
applied = {
int(version): str(name)
for version, name in cursor.fetchall()
}
configured = {
migration.version: migration
for migration in self._migrations
}
self._validate_applied(
applied=applied,
configured=configured,
)
applied_now: list[int] = []
for migration in self._migrations:
if migration.version in applied:
continue
for statement in migration.statements:
cursor.execute(statement)
cursor.execute(
_INSERT_APPLIED_MIGRATION_SQL,
(
migration.version,
migration.name,
),
)
applied_now.append(migration.version)
return tuple(applied_now)
except StorageMigrationError:
raise
except Exception as error:
raise StorageMigrationError(
"Failed to apply storage schema migrations."
) from error
@staticmethod
def _validate_applied(
*,
applied: dict[int, str],
configured: dict[int, StorageMigration],
) -> None:
for version, applied_name in applied.items():
migration = configured.get(version)
if migration is None:
raise StorageMigrationError(
"Database contains unknown storage migration "
f"version {version}."
)
if migration.name != applied_name:
raise StorageMigrationError(
"Storage migration name mismatch for version "
f"{version}: database={applied_name!r}, "
f"configured={migration.name!r}."
)
configured_versions = tuple(configured)
applied_versions = tuple(sorted(applied))
expected_prefix = configured_versions[: len(applied_versions)]
if applied_versions != expected_prefix:
raise StorageMigrationError(
"Applied storage migrations must form an ordered "
"prefix of configured migrations."
)
def run_storage_migrations(
connection_provider: StorageConnectionProvider | None = None,
) -> tuple[int, ...]:
"""Запустить миграции через явного поставщика соединений."""
if connection_provider is None:
from src.storage.session import get_connection
connection_provider = get_connection
return StorageMigrationRunner(
connection_provider=connection_provider,
).run()

View File

@@ -0,0 +1,156 @@
from __future__ import annotations
import math
from collections.abc import Callable
from contextlib import AbstractContextManager
from typing import Any
from src.storage.exceptions import PostgresConnectionPoolError
PostgresPoolFactory = Callable[..., Any]
def _default_pool_factory(**kwargs: Any) -> Any:
from psycopg_pool import ConnectionPool
return ConnectionPool(**kwargs)
class PostgresConnectionPool:
"""Явная обёртка управляемого жизненного цикла пула psycopg."""
__slots__ = (
"_conninfo",
"_min_size",
"_max_size",
"_name",
"_pool",
"_pool_factory",
"_timeout_seconds",
)
def __init__(
self,
*,
conninfo: str,
min_size: int = 1,
max_size: int = 4,
timeout_seconds: float = 10.0,
name: str = "dzentra-storage",
pool_factory: PostgresPoolFactory = _default_pool_factory,
) -> None:
normalized_conninfo = str(conninfo or "").strip()
normalized_name = str(name or "").strip()
if not normalized_conninfo:
raise ValueError("conninfo must not be empty")
if (
isinstance(min_size, bool)
or not isinstance(min_size, int)
or min_size <= 0
):
raise ValueError("min_size must be a positive integer")
if (
isinstance(max_size, bool)
or not isinstance(max_size, int)
or max_size < min_size
):
raise ValueError(
"max_size must be an integer not smaller than min_size"
)
if (
isinstance(timeout_seconds, bool)
or not isinstance(timeout_seconds, (int, float))
or not math.isfinite(float(timeout_seconds))
or timeout_seconds <= 0
):
raise ValueError("timeout_seconds must be positive and finite")
if not normalized_name:
raise ValueError("name must not be empty")
if not callable(pool_factory):
raise TypeError("pool_factory must be callable")
self._conninfo = normalized_conninfo
self._min_size = min_size
self._max_size = max_size
self._timeout_seconds = float(timeout_seconds)
self._name = normalized_name
self._pool_factory = pool_factory
self._pool: Any | None = None
@property
def is_open(self) -> bool:
return self._pool is not None
def open(self) -> None:
"""Открыть пул и проверить создание минимального числа соединений."""
if self._pool is not None:
return
pool: Any = None
try:
pool = self._pool_factory(
conninfo=self._conninfo,
min_size=self._min_size,
max_size=self._max_size,
timeout=self._timeout_seconds,
kwargs={"autocommit": False},
name=self._name,
open=False,
)
pool.open(
wait=True,
timeout=self._timeout_seconds,
)
except BaseException as error:
if pool is not None:
try:
pool.close(timeout=self._timeout_seconds)
except Exception as cleanup_error:
error.add_note(
"PostgreSQL pool cleanup also failed: "
f"{type(cleanup_error).__name__}."
)
if not isinstance(error, Exception):
raise
raise PostgresConnectionPoolError(
"Failed to open PostgreSQL connection pool."
) from error
self._pool = pool
def connection(self) -> AbstractContextManager[Any]:
"""Выдать одно транзакционное соединение из открытого пула."""
pool = self._pool
if pool is None:
raise PostgresConnectionPoolError(
"PostgreSQL connection pool is not open."
)
return pool.connection(timeout=self._timeout_seconds)
def close(self) -> None:
"""Закрыть пул; повторное закрытие ничего не делает."""
pool = self._pool
if pool is None:
return
self._pool = None
try:
pool.close(timeout=self._timeout_seconds)
except Exception as error:
raise PostgresConnectionPoolError(
"Failed to close PostgreSQL connection pool."
) from error

View File

@@ -0,0 +1,103 @@
from __future__ import annotations
import os
from collections.abc import Iterator
import psycopg
import pytest
from src.storage.migrations import StorageMigrationRunner
from src.storage.postgres_pool import PostgresConnectionPool
from tests.support.postgres_market_data import (
PostgresTestSettings,
acquire_postgres_test_lock,
connect_postgres_test_database,
load_postgres_test_settings,
release_postgres_test_lock,
reset_postgres_test_database,
)
@pytest.fixture(scope="session")
def postgres_test_settings() -> Iterator[PostgresTestSettings]:
try:
settings = load_postgres_test_settings(os.environ)
except ValueError as error:
pytest.fail(str(error), pytrace=False)
if settings is None:
pytest.skip(
"PostgreSQL integration is opt-in; set "
"DZENTRA_RUN_POSTGRES_TESTS=1 and "
"DZENTRA_TEST_POSTGRES_DSN explicitly."
)
try:
control = connect_postgres_test_database(settings)
except psycopg.Error as error:
pytest.fail(
"Could not connect to the explicit PostgreSQL test database: "
f"{type(error).__name__}.",
pytrace=False,
)
with control:
if not acquire_postgres_test_lock(control):
pytest.fail(
"Another integration session already owns this test database.",
pytrace=False,
)
try:
reset_postgres_test_database(
control,
expected_database_name=settings.database_name,
)
yield settings
finally:
reset_postgres_test_database(
control,
expected_database_name=settings.database_name,
)
release_postgres_test_lock(control)
@pytest.fixture(autouse=True)
def clean_postgres_test_database(
postgres_test_settings: PostgresTestSettings,
) -> Iterator[None]:
with connect_postgres_test_database(postgres_test_settings) as control:
reset_postgres_test_database(
control,
expected_database_name=postgres_test_settings.database_name,
)
yield
with connect_postgres_test_database(postgres_test_settings) as control:
reset_postgres_test_database(
control,
expected_database_name=postgres_test_settings.database_name,
)
@pytest.fixture
def migrated_postgres_pool(
postgres_test_settings: PostgresTestSettings,
) -> Iterator[PostgresConnectionPool]:
pool = PostgresConnectionPool(
conninfo=postgres_test_settings.dsn,
min_size=1,
max_size=4,
timeout_seconds=5.0,
name="market-data-integration",
)
pool.open()
StorageMigrationRunner(
connection_provider=pool.connection,
).run()
try:
yield pool
finally:
pool.close()

View File

@@ -0,0 +1,257 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from decimal import Decimal
import threading
import pytest
from src.market_data.acquisition.models.quote import Quote
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.storage import (
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
MarketDataPartitionType,
MarketDataRetentionPolicy,
MarketDataStorageOperationError,
PostgresMarketDataPartitionManager,
PostgresMarketDataRetentionService,
PostgresQuoteRepository,
PostgresTradeRepository,
)
from src.storage.postgres_pool import PostgresConnectionPool
from tests.support.postgres_market_data import (
PostgresTestSettings,
connect_postgres_test_database,
wait_for_postgres_advisory_lock_waiters,
)
pytestmark = pytest.mark.integration
VENUE = "dzengi"
SYMBOL = "BTC/USD_LEVERAGE"
def _trade(*, trade_id: int, executed_at: datetime) -> Trade:
return Trade(
symbol=SYMBOL,
trade_id=trade_id,
price=Decimal("100"),
quantity=Decimal("1"),
executed_at=executed_at,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi_websocket_trade",
)
def _quote(*, received_at: datetime) -> Quote:
return Quote(
symbol=SYMBOL,
last_price=Decimal("100"),
bid_price=Decimal("99"),
ask_price=Decimal("101"),
exchange_timestamp=received_at,
received_at=received_at,
source="dzengi",
)
def test_real_partition_creation_moves_default_row_and_is_idempotent(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
event_time = datetime(2026, 7, 15, tzinfo=timezone.utc)
repository = PostgresTradeRepository(
connection_provider=migrated_postgres_pool.connection,
)
manager = PostgresMarketDataPartitionManager(
connection_provider=migrated_postgres_pool.connection,
)
repository.store_trade(
venue=VENUE,
trade=_trade(trade_id=1, executed_at=event_time),
observed_at=event_time,
)
created = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=event_time,
)
repeated = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=event_time,
)
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT tableoid::regclass::text FROM market_data.trades"
)
relation = cursor.fetchone()
assert created.created is True
assert created.moved_row_count == 1
assert repeated.created is False
assert repeated.moved_row_count == 0
assert relation == ("market_data.trades_2026_07",)
def test_two_real_partition_callers_create_one_partition(
migrated_postgres_pool: PostgresConnectionPool,
postgres_test_settings: PostgresTestSettings,
) -> None:
month = datetime(2026, 8, 1, tzinfo=timezone.utc)
manager = PostgresMarketDataPartitionManager(
connection_provider=migrated_postgres_pool.connection,
)
start_barrier = threading.Barrier(3)
caller_ids: set[int] = set()
caller_ids_lock = threading.Lock()
def ensure_partition() -> tuple[bool, int]:
with caller_ids_lock:
caller_ids.add(threading.get_ident())
start_barrier.wait(timeout=5.0)
result = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=month,
)
return result.created, result.moved_row_count
with connect_postgres_test_database(
postgres_test_settings,
autocommit=False,
) as control:
with control.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_xact_lock(%s)",
(MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,),
)
with ThreadPoolExecutor(max_workers=2) as executor:
futures = tuple(executor.submit(ensure_partition) for _ in range(2))
start_barrier.wait(timeout=5.0)
try:
wait_for_postgres_advisory_lock_waiters(
control,
lock_id=MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
expected_count=2,
)
finally:
control.commit()
results = tuple(future.result(timeout=10.0) for future in futures)
assert len(caller_ids) == 2
assert sorted(results) == [(False, 0), (True, 0)]
def test_real_retention_uses_exact_cutoff(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
now = datetime(2026, 8, 15, 12, 0, tzinfo=timezone.utc)
cutoff = now - timedelta(days=10)
repository = PostgresTradeRepository(
connection_provider=migrated_postgres_pool.connection,
)
service = PostgresMarketDataRetentionService(
connection_provider=migrated_postgres_pool.connection,
)
for trade_id, executed_at in (
(1, cutoff - timedelta(microseconds=1)),
(2, cutoff),
(3, cutoff + timedelta(microseconds=1)),
):
repository.store_trade(
venue=VENUE,
trade=_trade(trade_id=trade_id, executed_at=executed_at),
observed_at=now,
)
result = service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=10),
now=now,
)
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT trade_id FROM market_data.trades ORDER BY trade_id"
)
remaining = tuple(row[0] for row in cursor.fetchall())
assert result.total_removed_count == 1
assert remaining == (2, 3)
def test_real_retention_failure_rolls_back_all_data_types(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
now = datetime(2026, 8, 15, 12, 0, tzinfo=timezone.utc)
old_time = now - timedelta(days=30)
trade_repository = PostgresTradeRepository(
connection_provider=migrated_postgres_pool.connection,
)
quote_repository = PostgresQuoteRepository(
connection_provider=migrated_postgres_pool.connection,
)
service = PostgresMarketDataRetentionService(
connection_provider=migrated_postgres_pool.connection,
)
trade_repository.store_trade(
venue=VENUE,
trade=_trade(trade_id=1, executed_at=old_time),
observed_at=now,
)
quote_repository.store_quote(
venue=VENUE,
quote=_quote(received_at=old_time),
)
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"""
CREATE FUNCTION market_data.reject_quote_delete()
RETURNS trigger
LANGUAGE plpgsql
AS $$
BEGIN
RAISE EXCEPTION 'injected quote retention failure';
END;
$$
"""
)
cursor.execute(
"""
CREATE TRIGGER reject_quote_delete
BEFORE DELETE ON market_data.quotes
FOR EACH ROW
EXECUTE FUNCTION market_data.reject_quote_delete()
"""
)
with pytest.raises(MarketDataStorageOperationError):
service.apply(
policy=MarketDataRetentionPolicy(
enabled=True,
trade_days=10,
quote_days=10,
),
now=now,
)
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute("SELECT COUNT(*) FROM market_data.trades")
trade_count = cursor.fetchone()
cursor.execute("SELECT COUNT(*) FROM market_data.quotes")
quote_count = cursor.fetchone()
assert trade_count == (1,)
assert quote_count == (1,)

View File

@@ -0,0 +1,272 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
from dataclasses import replace
from datetime import datetime, timedelta, timezone
from decimal import Decimal
import threading
import pytest
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.storage import (
MarketDataStorageConflictError,
MarketDataWriteStatus,
PostgresCandleRepository,
PostgresQuoteRepository,
PostgresTradeRepository,
)
from src.storage.postgres_pool import PostgresConnectionPool
pytestmark = pytest.mark.integration
VENUE = "dzengi"
SYMBOL = "BTC/USD_LEVERAGE"
EVENT_TIME = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
OBSERVED_AT = EVENT_TIME + timedelta(seconds=1)
def _trade(
*,
trade_id: int = 100,
executed_at: datetime = EVENT_TIME,
price: Decimal = Decimal("64159.45"),
source: str = "dzengi_websocket_trade",
symbol: str = SYMBOL,
) -> Trade:
return Trade(
symbol=symbol,
trade_id=trade_id,
price=price,
quantity=Decimal("0.125"),
executed_at=executed_at,
aggressor_side=TradeAggressorSide.BUY,
source=source,
)
def _quote(*, source: str = "dzengi") -> Quote:
return Quote(
symbol=SYMBOL,
last_price=Decimal("100"),
bid_price=Decimal("99"),
ask_price=Decimal("101"),
exchange_timestamp=EVENT_TIME,
received_at=OBSERVED_AT,
source=source,
)
def _candle(*, source: str = "rest_klines:bid") -> Candle:
return Candle(
symbol=SYMBOL,
interval="1m",
open_time=EVENT_TIME,
open_price=Decimal("100"),
high_price=Decimal("110"),
low_price=Decimal("90"),
close_price=Decimal("105"),
volume=Decimal("10"),
source=source,
)
def test_real_trade_insert_duplicate_provenance_and_conflict(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
repository = PostgresTradeRepository(
connection_provider=migrated_postgres_pool.connection,
)
websocket_trade = _trade()
inserted = repository.store_trade(
venue=VENUE,
trade=websocket_trade,
observed_at=OBSERVED_AT,
)
duplicate = repository.store_trade(
venue=VENUE,
trade=websocket_trade,
observed_at=OBSERVED_AT,
)
provenance = repository.store_trade(
venue=VENUE,
trade=replace(websocket_trade, source="dzengi"),
observed_at=OBSERVED_AT + timedelta(seconds=5),
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_trade(
venue=VENUE,
trade=replace(websocket_trade, price=Decimal("999")),
observed_at=OBSERVED_AT,
)
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"""
SELECT price, source, observation_sources,
first_observed_at, last_observed_at
FROM market_data.trades
"""
)
rows = tuple(cursor.fetchall())
assert inserted.status is MarketDataWriteStatus.INSERTED
assert duplicate.status is MarketDataWriteStatus.DUPLICATE
assert provenance.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert rows == (
(
Decimal("64159.45"),
"dzengi_websocket_trade",
["dzengi_websocket_trade", "dzengi"],
OBSERVED_AT,
OBSERVED_AT + timedelta(seconds=5),
),
)
def test_real_trade_batch_rolls_back_preceding_insert_on_conflict(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
repository = PostgresTradeRepository(
connection_provider=migrated_postgres_pool.connection,
)
existing = _trade(trade_id=2, symbol="B")
repository.store_trade(
venue=VENUE,
trade=existing,
observed_at=OBSERVED_AT,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_trades(
venue=VENUE,
trades=(
_trade(trade_id=1, symbol="A"),
replace(existing, price=Decimal("999")),
),
observed_at=OBSERVED_AT,
)
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT symbol, trade_id FROM market_data.trades"
)
rows = tuple(cursor.fetchall())
assert rows == (("B", 2),)
def test_two_real_trade_writers_converge_on_one_canonical_row(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
repository = PostgresTradeRepository(
connection_provider=migrated_postgres_pool.connection,
)
trades = (
_trade(source="dzengi_websocket_trade"),
_trade(source="dzengi"),
)
start_barrier = threading.Barrier(2)
caller_ids: set[int] = set()
caller_ids_lock = threading.Lock()
def store(trade: Trade) -> MarketDataWriteStatus:
with caller_ids_lock:
caller_ids.add(threading.get_ident())
start_barrier.wait(timeout=5.0)
return repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
).status
with ThreadPoolExecutor(max_workers=2) as executor:
statuses = tuple(executor.map(store, trades))
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT observation_sources FROM market_data.trades"
)
rows = tuple(cursor.fetchall())
assert set(statuses) == {
MarketDataWriteStatus.INSERTED,
MarketDataWriteStatus.PROVENANCE_UPDATED,
}
assert len(caller_ids) == 2
assert len(rows) == 1
assert set(rows[0][0]) == {"dzengi_websocket_trade", "dzengi"}
def test_real_quote_and_candle_revision_provenance_and_identity(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
quote_repository = PostgresQuoteRepository(
connection_provider=migrated_postgres_pool.connection,
)
candle_repository = PostgresCandleRepository(
connection_provider=migrated_postgres_pool.connection,
)
quote = _quote()
candle = _candle()
assert quote_repository.store_quote(
venue=VENUE,
quote=quote,
).status is MarketDataWriteStatus.INSERTED
assert quote_repository.store_quote(
venue=VENUE,
quote=replace(quote, source="dzengi_websocket_quote"),
).status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert candle_repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
).status is MarketDataWriteStatus.INSERTED
assert candle_repository.store_candle_revision(
venue=VENUE,
candle=replace(candle, source="secondary"),
observed_at=OBSERVED_AT,
is_final=False,
).status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert candle_repository.store_candle_revision(
venue=VENUE,
candle=replace(candle, close_price=Decimal("106")),
observed_at=OBSERVED_AT + timedelta(seconds=1),
is_final=True,
).status is MarketDataWriteStatus.INSERTED
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"SELECT source, observation_sources FROM market_data.quotes"
)
quote_rows = tuple(cursor.fetchall())
cursor.execute(
"""
SELECT is_final, source, observation_sources
FROM market_data.candle_revisions
ORDER BY observed_at
"""
)
candle_rows = tuple(cursor.fetchall())
assert quote_rows == (
("dzengi", ["dzengi", "dzengi_websocket_quote"]),
)
assert candle_rows == (
(False, "rest_klines:bid", ["rest_klines:bid", "secondary"]),
(True, "rest_klines:bid", ["rest_klines:bid"]),
)

View File

@@ -0,0 +1,204 @@
from __future__ import annotations
from concurrent.futures import ThreadPoolExecutor
import threading
import pytest
from src.storage.exceptions import StorageMigrationError
from src.storage.migrations import (
STORAGE_MIGRATIONS,
STORAGE_MIGRATION_ADVISORY_LOCK_ID,
StorageMigrationRunner,
)
from src.storage.postgres_pool import PostgresConnectionPool
from tests.support.postgres_market_data import (
PostgresTestSettings,
connect_postgres_test_database,
count_other_test_connections,
wait_for_postgres_advisory_lock_waiters,
)
pytestmark = pytest.mark.integration
def _open_pool(
settings: PostgresTestSettings,
*,
name: str,
) -> PostgresConnectionPool:
pool = PostgresConnectionPool(
conninfo=settings.dsn,
min_size=1,
max_size=2,
timeout_seconds=5.0,
name=name,
)
pool.open()
return pool
def test_real_migrations_create_expected_schema_and_are_idempotent(
postgres_test_settings: PostgresTestSettings,
) -> None:
pool = _open_pool(
postgres_test_settings,
name="schema-integration",
)
runner = StorageMigrationRunner(
connection_provider=pool.connection,
)
expected_versions = tuple(
migration.version for migration in STORAGE_MIGRATIONS
)
try:
assert runner.run() == expected_versions
assert runner.run() == ()
with pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"""
SELECT version, name
FROM public.storage_schema_migrations
ORDER BY version
"""
)
history = tuple(cursor.fetchall())
cursor.execute(
"""
SELECT child.relname
FROM pg_catalog.pg_inherits AS inheritance
JOIN pg_catalog.pg_class AS parent
ON parent.oid = inheritance.inhparent
JOIN pg_catalog.pg_namespace AS parent_namespace
ON parent_namespace.oid = parent.relnamespace
JOIN pg_catalog.pg_class AS child
ON child.oid = inheritance.inhrelid
WHERE parent_namespace.nspname = 'market_data'
AND parent.relname IN (
'trades',
'quotes',
'candle_revisions'
)
ORDER BY child.relname
"""
)
default_partitions = tuple(
row[0] for row in cursor.fetchall()
)
cursor.execute(
"""
SELECT parent.relname
FROM pg_catalog.pg_partitioned_table AS partitioned
JOIN pg_catalog.pg_class AS parent
ON parent.oid = partitioned.partrelid
JOIN pg_catalog.pg_namespace AS namespace
ON namespace.oid = parent.relnamespace
WHERE namespace.nspname = 'market_data'
ORDER BY parent.relname
"""
)
partitioned_tables = tuple(
row[0] for row in cursor.fetchall()
)
finally:
pool.close()
assert history == tuple(
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS
)
assert default_partitions == (
"candle_revisions_default",
"quotes_default",
"trades_default",
)
assert partitioned_tables == (
"candle_revisions",
"quotes",
"trades",
)
with connect_postgres_test_database(postgres_test_settings) as control:
assert count_other_test_connections(control) == 0
def test_two_real_migration_runners_apply_each_version_once(
postgres_test_settings: PostgresTestSettings,
) -> None:
pools = (
_open_pool(postgres_test_settings, name="migration-one"),
_open_pool(postgres_test_settings, name="migration-two"),
)
runners = tuple(
StorageMigrationRunner(connection_provider=pool.connection)
for pool in pools
)
expected = tuple(
migration.version for migration in STORAGE_MIGRATIONS
)
start_barrier = threading.Barrier(3)
caller_ids: set[int] = set()
caller_ids_lock = threading.Lock()
def run(runner: StorageMigrationRunner) -> tuple[int, ...]:
with caller_ids_lock:
caller_ids.add(threading.get_ident())
start_barrier.wait(timeout=5.0)
return runner.run()
try:
with connect_postgres_test_database(
postgres_test_settings,
autocommit=False,
) as control:
with control.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_xact_lock(%s)",
(STORAGE_MIGRATION_ADVISORY_LOCK_ID,),
)
with ThreadPoolExecutor(max_workers=2) as executor:
futures = tuple(executor.submit(run, runner) for runner in runners)
start_barrier.wait(timeout=5.0)
try:
wait_for_postgres_advisory_lock_waiters(
control,
lock_id=STORAGE_MIGRATION_ADVISORY_LOCK_ID,
expected_count=2,
)
finally:
control.commit()
results = tuple(future.result(timeout=10.0) for future in futures)
finally:
for pool in pools:
pool.close()
assert len(caller_ids) == 2
assert set(results) == {expected, ()}
def test_real_migration_history_mismatch_is_fatal(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
with migrated_postgres_pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"""
UPDATE public.storage_schema_migrations
SET name = 'tampered'
WHERE version = 1
"""
)
runner = StorageMigrationRunner(
connection_provider=migrated_postgres_pool.connection,
)
with pytest.raises(StorageMigrationError, match="name mismatch"):
runner.run()

View File

@@ -0,0 +1,445 @@
from __future__ import annotations
import asyncio
import threading
import time
from dataclasses import replace
from types import SimpleNamespace
from typing import Any, cast
import pytest
from aiogram import Bot, Dispatcher
from psycopg.conninfo import conninfo_to_dict
from src.bootstrap.application import (
ApplicationComposition,
run_application,
)
from src.bootstrap.market_data_storage import build_market_data_storage
from src.bootstrap.trade_stream_runtime import (
build_trade_stream_production_runtime,
)
from src.core.config import MarketDataStorageSettings, Settings
from src.market_data.storage import (
MarketDataStorageOperationError,
PostgresTradeRepository,
TradeStorageObservationSink,
)
from src.storage.postgres_pool import PostgresConnectionPool
from tests.integration.market_data.acquisition.runtime.loopback_trade_exchange import (
LoopbackHttpResponse,
LoopbackTradeEnvironment,
LoopbackTradeRestServer,
LoopbackTradeWebSocketServer,
wait_until,
)
from tests.support.postgres_market_data import (
PostgresTestSettings,
connect_postgres_test_database,
count_other_test_connections,
)
from tests.support.trade_stream_runtime import (
SYMBOL,
assert_no_owned_tasks,
build_runtime,
make_settings,
run_scenario,
start_runtime,
state_store_from,
stop_runtime,
)
pytestmark = pytest.mark.integration
class FakeBotSession:
def __init__(self) -> None:
self.close_calls = 0
async def close(self) -> None:
self.close_calls += 1
class ControlledDispatcher:
def __init__(self) -> None:
self.started = asyncio.Event()
self.release = asyncio.Event()
self.cancelled = asyncio.Event()
async def start_polling(
self,
bot: object,
*,
close_bot_session: bool,
) -> None:
del bot
assert close_bot_session is False
self.started.set()
try:
await self.release.wait()
except asyncio.CancelledError:
self.cancelled.set()
raise
def _sink(pool: PostgresConnectionPool) -> TradeStorageObservationSink:
return TradeStorageObservationSink(
trade_storage=PostgresTradeRepository(
connection_provider=pool.connection,
),
venue="dzengi",
)
def _fetch_trade_rows(
pool: PostgresConnectionPool,
) -> tuple[tuple[Any, ...], ...]:
with pool.connection() as connection:
with connection.cursor() as cursor:
cursor.execute(
"""
SELECT trade_id, source, observation_sources
FROM market_data.trades
ORDER BY executed_at, trade_id
"""
)
return tuple(cursor.fetchall())
def _recovered_trade(
*,
trade_id: int,
timestamp_ms: int,
price: str = "64555.56",
quantity: str = "0.003",
) -> dict[str, object]:
return {
"a": trade_id,
"p": price,
"q": quantity,
"T": timestamp_ms,
"m": False,
}
def _settings_for_database(
*,
postgres: PostgresTestSettings,
websocket_url: str,
rest_base_url: str,
) -> Settings:
parameters = conninfo_to_dict(postgres.dsn)
host_value = parameters.get("host") or parameters.get("hostaddr")
if host_value is None:
raise AssertionError("PostgreSQL integration DSN has no host")
port_value = parameters.get("port")
return replace(
make_settings(
websocket_url=websocket_url,
rest_base_url=rest_base_url,
),
db_host=str(host_value),
db_port=int(str(port_value if port_value is not None else 5432)),
db_name=str(parameters["dbname"]),
db_user=str(parameters.get("user", "")),
db_password=str(parameters.get("password", "")),
market_data_storage=MarketDataStorageSettings(
enabled=True,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=5.0,
),
)
def test_live_runtime_persists_before_checkpoint(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
async def scenario() -> None:
websocket = LoopbackTradeWebSocketServer()
rest = LoopbackTradeRestServer()
async with LoopbackTradeEnvironment(
websocket=websocket,
rest=rest,
) as environment:
runtime = build_runtime(
websocket_url=environment.websocket_url,
rest_base_url=rest.base_url,
trade_observation_sink=_sink(migrated_postgres_pool),
)
runtime_task: asyncio.Task[None] | None = None
try:
runtime_task = await start_runtime(runtime)
await websocket.wait_for_subscriptions(1)
timestamp_ms = time.time_ns() // 1_000_000
await websocket.send_trade(
0,
symbol=SYMBOL,
trade_id=600,
timestamp_ms=timestamp_ms,
)
state_store = state_store_from(runtime)
await wait_until(
lambda: (
state_store.contains(SYMBOL)
and state_store.get(SYMBOL).last_trade_id == 600
)
)
rows = await asyncio.to_thread(
_fetch_trade_rows,
migrated_postgres_pool,
)
assert rows == (
(
600,
"dzengi_websocket_trade",
["dzengi_websocket_trade"],
),
)
finally:
if runtime_task is not None:
await stop_runtime(runtime, runtime_task)
await assert_no_owned_tasks()
run_scenario(scenario())
def test_reconnect_recovery_and_live_share_real_repository(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
async def scenario() -> None:
base_time_ms = time.time_ns() // 1_000_000 - 1_000
websocket = LoopbackTradeWebSocketServer()
rest = LoopbackTradeRestServer(
responses=(
LoopbackHttpResponse(
body=[
_recovered_trade(
trade_id=700,
timestamp_ms=base_time_ms,
price="64555.55",
quantity="0.002",
),
_recovered_trade(
trade_id=701,
timestamp_ms=base_time_ms + 100,
)
]
),
)
)
async with LoopbackTradeEnvironment(
websocket=websocket,
rest=rest,
) as environment:
runtime = build_runtime(
websocket_url=environment.websocket_url,
rest_base_url=rest.base_url,
trade_observation_sink=_sink(migrated_postgres_pool),
)
runtime_task: asyncio.Task[None] | None = None
try:
runtime_task = await start_runtime(runtime)
await websocket.wait_for_subscriptions(1)
await websocket.send_trade(
0,
symbol=SYMBOL,
trade_id=700,
timestamp_ms=base_time_ms,
)
state_store = state_store_from(runtime)
await wait_until(
lambda: (
state_store.contains(SYMBOL)
and state_store.get(SYMBOL).last_trade_id == 700
)
)
await websocket.abort_connection(0)
await websocket.wait_for_subscriptions(2)
await rest.wait_for_requests(1)
await wait_until(
lambda: state_store.get(SYMBOL).last_trade_id == 701
)
await websocket.send_trade(
1,
symbol=SYMBOL,
trade_id=702,
timestamp_ms=base_time_ms + 200,
)
await wait_until(
lambda: state_store.get(SYMBOL).last_trade_id == 702
)
rows = await asyncio.to_thread(
_fetch_trade_rows,
migrated_postgres_pool,
)
assert tuple(row[0] for row in rows) == (700, 701, 702)
assert rows[0][1] == "dzengi_websocket_trade"
assert rows[0][2] == [
"dzengi_websocket_trade",
"dzengi",
]
assert rows[1][1] == "dzengi"
assert rows[1][2] == ["dzengi"]
finally:
if runtime_task is not None:
await stop_runtime(runtime, runtime_task)
await assert_no_owned_tasks()
run_scenario(scenario())
def test_real_persistence_failure_is_fatal_and_preserves_checkpoint(
migrated_postgres_pool: PostgresConnectionPool,
) -> None:
async def scenario() -> None:
base_time_ms = time.time_ns() // 1_000_000
websocket = LoopbackTradeWebSocketServer()
rest = LoopbackTradeRestServer()
async with LoopbackTradeEnvironment(
websocket=websocket,
rest=rest,
) as environment:
runtime = build_runtime(
websocket_url=environment.websocket_url,
rest_base_url=rest.base_url,
trade_observation_sink=_sink(migrated_postgres_pool),
)
runtime_task: asyncio.Task[None] | None = None
try:
runtime_task = await start_runtime(runtime)
await websocket.wait_for_subscriptions(1)
await websocket.send_trade(
0,
symbol=SYMBOL,
trade_id=800,
timestamp_ms=base_time_ms,
)
state_store = state_store_from(runtime)
await wait_until(
lambda: (
state_store.contains(SYMBOL)
and state_store.get(SYMBOL).last_trade_id == 800
)
)
await asyncio.to_thread(migrated_postgres_pool.close)
await websocket.send_trade(
0,
symbol=SYMBOL,
trade_id=801,
timestamp_ms=base_time_ms + 1,
)
with pytest.raises(MarketDataStorageOperationError):
await asyncio.wait_for(runtime_task, timeout=3.0)
assert state_store.get(SYMBOL).last_trade_id == 800
finally:
if runtime_task is not None and not runtime_task.done():
await stop_runtime(runtime, runtime_task)
await assert_no_owned_tasks()
run_scenario(scenario())
migrated_postgres_pool.open()
assert tuple(row[0] for row in _fetch_trade_rows(migrated_postgres_pool)) == (
800,
)
def test_application_owns_real_storage_before_and_after_runtime(
postgres_test_settings: PostgresTestSettings,
) -> None:
storage_holder: list[Any] = []
bot_session = FakeBotSession()
async def scenario() -> None:
websocket = LoopbackTradeWebSocketServer()
rest = LoopbackTradeRestServer()
async with LoopbackTradeEnvironment(
websocket=websocket,
rest=rest,
) as environment:
settings = _settings_for_database(
postgres=postgres_test_settings,
websocket_url=environment.websocket_url,
rest_base_url=rest.base_url,
)
storage = build_market_data_storage(settings)
assert storage is not None
storage_holder.append(storage)
runtime = build_trade_stream_production_runtime(
settings,
trade_observation_sink=storage.trade_observation_sink,
)
assert runtime is not None
dispatcher = ControlledDispatcher()
bot = SimpleNamespace(session=bot_session)
application_task = asyncio.create_task(
run_application(
ApplicationComposition(
bot=cast(Bot, cast(object, bot)),
dispatcher=cast(
Dispatcher,
cast(object, dispatcher),
),
trade_stream_runtime=runtime,
market_data_storage_lifecycle=storage.lifecycle,
)
)
)
await dispatcher.started.wait()
assert storage.connection_pool.is_open is True
await websocket.wait_for_subscriptions(1)
await websocket.send_trade(
0,
symbol=SYMBOL,
trade_id=900,
timestamp_ms=time.time_ns() // 1_000_000,
)
state_store = state_store_from(runtime)
await wait_until(
lambda: (
state_store.contains(SYMBOL)
and state_store.get(SYMBOL).last_trade_id == 900
)
)
dispatcher.release.set()
await asyncio.wait_for(application_task, timeout=5.0)
assert storage.connection_pool.is_open is False
assert bot_session.close_calls == 1
await assert_no_owned_tasks()
run_scenario(scenario())
assert storage_holder
with connect_postgres_test_database(postgres_test_settings) as control:
with control.cursor() as cursor:
cursor.execute(
"SELECT trade_id FROM market_data.trades ORDER BY trade_id"
)
assert tuple(row[0] for row in cursor.fetchall()) == (900,)
assert count_other_test_connections(control) == 0

View File

@@ -114,7 +114,11 @@ def load_live_trade_stream_test_config(
def build_live_trade_stream_settings(
config: LiveTradeStreamTestConfig,
) -> Settings:
from src.core.config import Settings, TradeStreamSettings
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
return Settings(
bot_token="live-verification-does-not-use-telegram",
@@ -147,6 +151,12 @@ def build_live_trade_stream_settings(
db_name="live-verification",
db_user="live-verification",
db_password="",
market_data_storage=MarketDataStorageSettings(
enabled=False,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=10.0,
),
debug_enabled=False,
journal_debug_enabled=False,
)

View File

@@ -0,0 +1,316 @@
from __future__ import annotations
import ipaddress
import re
import time
from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any
import psycopg
from psycopg.conninfo import conninfo_to_dict, make_conninfo
POSTGRES_TEST_FLAG = "DZENTRA_RUN_POSTGRES_TESTS"
POSTGRES_TEST_DSN = "DZENTRA_TEST_POSTGRES_DSN"
POSTGRES_TEST_APPLICATION_NAME = "dzentra-storage-integration"
POSTGRES_TEST_CONTROL_APPLICATION_NAME = (
"dzentra-storage-integration-control"
)
POSTGRES_TEST_DATABASE_PREFIX = "dzentra_test_"
POSTGRES_TEST_ADVISORY_LOCK_ID = 0x445A454E54524154
_SAFE_DATABASE_NAME = re.compile(r"^dzentra_test_[A-Za-z0-9_]+$")
_LOCAL_HOST_NAMES = frozenset({"localhost"})
@dataclass(frozen=True, slots=True)
class PostgresTestSettings:
"""Явные настройки подключения к одноразовой локальной базе."""
dsn: str
database_name: str
def load_postgres_test_settings(
environ: Mapping[str, str],
) -> PostgresTestSettings | None:
"""Загрузить явно включаемый DSN без настроек основной базы."""
enabled = str(environ.get(POSTGRES_TEST_FLAG, "")).strip()
if not enabled:
return None
if enabled != "1":
raise ValueError(f"{POSTGRES_TEST_FLAG} must be exactly '1'")
raw_dsn = str(environ.get(POSTGRES_TEST_DSN, "")).strip()
if not raw_dsn:
raise ValueError(
f"{POSTGRES_TEST_DSN} is required when {POSTGRES_TEST_FLAG}=1"
)
try:
parameters = conninfo_to_dict(raw_dsn)
except Exception as error:
raise ValueError(
f"{POSTGRES_TEST_DSN} is not a valid PostgreSQL DSN"
) from error
database_name = str(parameters.get("dbname", "")).strip()
if not _SAFE_DATABASE_NAME.fullmatch(database_name):
raise ValueError(
f"{POSTGRES_TEST_DSN} database name must start with "
f"{POSTGRES_TEST_DATABASE_PREFIX!r} and contain only safe characters"
)
_validate_local_endpoint(parameters)
try:
normalized_dsn = make_conninfo(
raw_dsn,
application_name=POSTGRES_TEST_APPLICATION_NAME,
connect_timeout=5,
)
except Exception as error:
raise ValueError(f"{POSTGRES_TEST_DSN} could not be normalized") from error
return PostgresTestSettings(
dsn=normalized_dsn,
database_name=database_name,
)
def connect_postgres_test_database(
settings: PostgresTestSettings,
*,
autocommit: bool = True,
) -> psycopg.Connection[Any]:
"""Подключиться после прохождения явных проверок безопасности."""
if not isinstance(settings, PostgresTestSettings):
raise TypeError("settings must be PostgresTestSettings")
control_dsn = make_conninfo(
settings.dsn,
application_name=POSTGRES_TEST_CONTROL_APPLICATION_NAME,
)
connection = psycopg.connect(
control_dsn,
autocommit=autocommit,
)
with connection.cursor() as cursor:
cursor.execute("SELECT current_database()")
row = cursor.fetchone()
if row != (settings.database_name,):
connection.close()
raise RuntimeError(
"Connected PostgreSQL database does not match the validated "
"disposable test database."
)
return connection
def reset_postgres_test_database(
connection: psycopg.Connection[Any],
*,
expected_database_name: str,
) -> None:
"""Удалить только объекты Build 060.27 из проверенной тестовой базы."""
if not connection.autocommit:
raise ValueError("test database reset requires autocommit")
if not _SAFE_DATABASE_NAME.fullmatch(expected_database_name):
raise ValueError("expected database name is not a safe test database")
with connection.cursor() as cursor:
cursor.execute(
"""
SELECT current_database(), current_setting('application_name')
"""
)
identity = cursor.fetchone()
if identity != (
expected_database_name,
POSTGRES_TEST_CONTROL_APPLICATION_NAME,
):
raise RuntimeError(
"Refusing destructive reset because the current PostgreSQL "
"connection is not the validated test control connection."
)
cursor.execute("DROP SCHEMA IF EXISTS market_data CASCADE")
cursor.execute(
"DROP TABLE IF EXISTS public.storage_schema_migrations"
)
def acquire_postgres_test_lock(
connection: psycopg.Connection[Any],
) -> bool:
"""Не допустить две разрушающие сессии в одной тестовой базе."""
with connection.cursor() as cursor:
cursor.execute(
"SELECT pg_try_advisory_lock(%s)",
(POSTGRES_TEST_ADVISORY_LOCK_ID,),
)
row = cursor.fetchone()
return row == (True,)
def release_postgres_test_lock(
connection: psycopg.Connection[Any],
) -> None:
with connection.cursor() as cursor:
cursor.execute(
"SELECT pg_advisory_unlock(%s)",
(POSTGRES_TEST_ADVISORY_LOCK_ID,),
)
def count_other_test_connections(
connection: psycopg.Connection[Any],
) -> int:
"""Посчитать оставшиеся соединения стенда и пула с тестовой базой."""
with connection.cursor() as cursor:
cursor.execute(
"""
SELECT COUNT(*)
FROM pg_stat_activity
WHERE datname = current_database()
AND application_name = %s
""",
(POSTGRES_TEST_APPLICATION_NAME,),
)
row = cursor.fetchone()
if (
not isinstance(row, tuple)
or len(row) != 1
or isinstance(row[0], bool)
or not isinstance(row[0], int)
):
raise RuntimeError("PostgreSQL returned an invalid connection count")
return row[0]
def wait_for_postgres_advisory_lock_waiters(
connection: psycopg.Connection[Any],
*,
lock_id: int,
expected_count: int,
timeout_seconds: float = 5.0,
) -> None:
"""Дождаться подтверждения конкуренции всех сессий за блокировку."""
if (
isinstance(lock_id, bool)
or not isinstance(lock_id, int)
or lock_id < 0
or lock_id > 0x7FFF_FFFF_FFFF_FFFF
):
raise ValueError("lock_id must be a non-negative signed BIGINT")
if (
isinstance(expected_count, bool)
or not isinstance(expected_count, int)
or expected_count <= 0
):
raise ValueError("expected_count must be a positive integer")
if timeout_seconds <= 0:
raise ValueError("timeout_seconds must be positive")
class_id = (lock_id >> 32) & 0xFFFF_FFFF
object_id = lock_id & 0xFFFF_FFFF
deadline = time.monotonic() + timeout_seconds
while True:
with connection.cursor() as cursor:
cursor.execute(
"""
SELECT COUNT(*)
FROM pg_catalog.pg_locks
WHERE locktype = 'advisory'
AND classid = %s
AND objid = %s
AND objsubid = 1
AND NOT granted
""",
(class_id, object_id),
)
row = cursor.fetchone()
if row == (expected_count,):
return
if time.monotonic() >= deadline:
observed_count = row[0] if isinstance(row, tuple) and row else row
raise TimeoutError(
"PostgreSQL did not observe all advisory-lock callers; "
f"expected {expected_count}, observed {observed_count!r}."
)
time.sleep(0.01)
def _validate_local_endpoint(parameters: Mapping[str, object]) -> None:
service = str(parameters.get("service", "")).strip()
host = str(parameters.get("host", "")).strip()
hostaddr = str(parameters.get("hostaddr", "")).strip()
if service:
raise ValueError(
f"{POSTGRES_TEST_DSN} must not use a PostgreSQL service"
)
if not host and not hostaddr:
raise ValueError(
f"{POSTGRES_TEST_DSN} must contain an explicit local host or "
"hostaddr"
)
for endpoint_host in _split_postgres_endpoints(host):
if not endpoint_host.startswith("/") and not _is_local_host(
endpoint_host
):
raise ValueError(
f"{POSTGRES_TEST_DSN} must target localhost or a local socket"
)
for endpoint_address in _split_postgres_endpoints(hostaddr):
if not _is_loopback_address(endpoint_address):
raise ValueError(
f"{POSTGRES_TEST_DSN} hostaddr must be a loopback address"
)
def _split_postgres_endpoints(value: str) -> tuple[str, ...]:
if not value:
return ()
endpoints = tuple(part.strip() for part in value.split(","))
if any(not endpoint for endpoint in endpoints):
raise ValueError(
f"{POSTGRES_TEST_DSN} must not contain an implicit endpoint"
)
return endpoints
def _is_local_host(value: str) -> bool:
return value.lower() in _LOCAL_HOST_NAMES or _is_loopback_address(value)
def _is_loopback_address(value: str) -> bool:
try:
return ipaddress.ip_address(value).is_loopback
except ValueError:
return False

View File

@@ -7,10 +7,17 @@ from typing import Any, Protocol
from src.bootstrap.trade_stream_runtime import (
build_trade_stream_production_runtime,
)
from src.core.config import Settings, TradeStreamSettings
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
from src.market_data.acquisition.consistency.trade_stream_state_store import (
TradeStreamStateStore,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.runtime.trade_stream_production_runtime import (
TradeStreamProductionRuntime,
TradeStreamProductionRuntimeState,
@@ -23,6 +30,8 @@ RUNTIME_CLEANUP_TIMEOUT_SECONDS = 5.0
OWNED_TASK_NAMES = frozenset(
{
"market-data-storage-shutdown",
"market-data-storage-startup",
"trade-stream-receive",
"trade-stream-runtime",
"trade-stream-runtime-recovery",
@@ -98,6 +107,12 @@ def make_settings(
db_name="integration",
db_user="integration",
db_password="",
market_data_storage=MarketDataStorageSettings(
enabled=False,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=10.0,
),
debug_enabled=False,
journal_debug_enabled=False,
)
@@ -111,6 +126,7 @@ def build_runtime(
close_timeout_seconds: float = 0.2,
heartbeat_timeout_seconds: float = 60.0,
scheduler_interval_seconds: float = 60.0,
trade_observation_sink: TradeObservationSinkProtocol | None = None,
) -> TradeStreamProductionRuntime:
runtime = build_trade_stream_production_runtime(
make_settings(
@@ -120,7 +136,8 @@ def build_runtime(
close_timeout_seconds=close_timeout_seconds,
heartbeat_timeout_seconds=heartbeat_timeout_seconds,
scheduler_interval_seconds=scheduler_interval_seconds,
)
),
trade_observation_sink=trade_observation_sink,
)
assert runtime is not None

View File

@@ -46,6 +46,7 @@ def make_settings() -> SimpleNamespace:
exchange_name="dzengi",
default_symbol="BTC/USD_LEVERAGE",
trade_stream=SimpleNamespace(enabled=True),
market_data_storage=SimpleNamespace(enabled=True),
)
@@ -56,6 +57,12 @@ def test_create_app_builds_one_application_composition(
bot = object()
dispatcher = object()
runtime = object()
storage_lifecycle = object()
storage_sink = object()
storage = SimpleNamespace(
lifecycle=storage_lifecycle,
trade_observation_sink=storage_sink,
)
journal = RecordingJournal()
observed_runtime_settings: list[object] = []
registered_bots: list[object] = []
@@ -82,8 +89,23 @@ def test_create_app_builds_one_application_composition(
lambda: journal,
)
def build_runtime(received_settings: object) -> object:
monkeypatch.setattr(
app_factory,
"build_market_data_storage",
lambda received_settings: (
storage
if received_settings is settings
else None
),
)
def build_runtime(
received_settings: object,
*,
trade_observation_sink: object,
) -> object:
observed_runtime_settings.append(received_settings)
assert trade_observation_sink is storage_sink
return runtime
monkeypatch.setattr(
@@ -118,10 +140,18 @@ def test_create_app_builds_one_application_composition(
assert application.bot is bot
assert application.dispatcher is dispatcher
assert application.trade_stream_runtime is runtime
assert (
application.market_data_storage_lifecycle
is storage_lifecycle
)
assert observed_runtime_settings == [settings]
assert registered_bots == [bot]
assert routed_dispatchers == [dispatcher]
assert journal.info_calls[0][2]["trade_stream_enabled"] is True
assert (
journal.info_calls[0][2]["market_data_storage_enabled"]
is True
)
def test_runtime_build_error_is_fatal(
@@ -149,9 +179,18 @@ def test_runtime_build_error_is_fatal(
"JournalService",
RecordingJournal,
)
monkeypatch.setattr(
app_factory,
"build_market_data_storage",
lambda settings: None,
)
def fail_runtime_build(settings: object) -> None:
del settings
def fail_runtime_build(
settings: object,
*,
trade_observation_sink: object,
) -> None:
del settings, trade_observation_sink
raise expected
monkeypatch.setattr(

View File

@@ -1,6 +1,7 @@
from __future__ import annotations
import asyncio
import threading
import pytest
@@ -46,12 +47,14 @@ class FakeDispatcher:
*,
return_immediately: bool = False,
error: BaseException | None = None,
calls: list[str] | None = None,
) -> None:
self.started = asyncio.Event()
self.release = asyncio.Event()
self.cancelled = asyncio.Event()
self.return_immediately = return_immediately
self.error = error
self.calls = calls
self.close_bot_session_values: list[bool] = []
async def start_polling(
@@ -61,6 +64,8 @@ class FakeDispatcher:
close_bot_session: bool,
) -> None:
del bot
if self.calls is not None:
self.calls.append("polling.start")
self.close_bot_session_values.append(close_bot_session)
self.started.set()
@@ -82,6 +87,7 @@ class FakeRuntime:
return_immediately: bool = False,
error: BaseException | None = None,
stop_error: BaseException | None = None,
calls: list[str] | None = None,
) -> None:
self.started = asyncio.Event()
self.release = asyncio.Event()
@@ -89,6 +95,7 @@ class FakeRuntime:
self.return_immediately = return_immediately
self.error = error
self.stop_error = stop_error
self.calls = calls
self.stop_calls = 0
@property
@@ -104,6 +111,8 @@ class FakeRuntime:
return self.state is TradeStreamProductionRuntimeState.RUNNING
async def run(self) -> None:
if self.calls is not None:
self.calls.append("runtime.run")
self.started.set()
if not self.return_immediately:
@@ -114,6 +123,8 @@ class FakeRuntime:
async def stop(self) -> None:
self.stop_calls += 1
if self.calls is not None:
self.calls.append("runtime.stop")
self.release.set()
self.stopped.set()
@@ -136,19 +147,289 @@ class BlockingStopRuntime(FakeRuntime):
self.stopped.set()
class FakeStorageLifecycle:
def __init__(
self,
*,
calls: list[str] | None = None,
start_error: BaseException | None = None,
stop_error: BaseException | None = None,
start_release: threading.Event | None = None,
stop_release: threading.Event | None = None,
) -> None:
self.calls = calls
self.start_error = start_error
self.stop_error = stop_error
self.start_release = start_release
self.stop_release = stop_release
self.start_entered = threading.Event()
self.stop_entered = threading.Event()
self.start_calls = 0
self.stop_calls = 0
self.started = False
def start(self) -> None:
self.start_calls += 1
if self.calls is not None:
self.calls.append("storage.start")
self.start_entered.set()
if self.start_release is not None:
if not self.start_release.wait(timeout=2.0):
raise TimeoutError("storage startup was not released")
if self.start_error is not None:
raise self.start_error
self.started = True
def stop(self) -> None:
self.stop_calls += 1
if self.calls is not None:
self.calls.append("storage.stop")
self.started = False
self.stop_entered.set()
if self.stop_release is not None:
if not self.stop_release.wait(timeout=2.0):
raise TimeoutError("storage shutdown was not released")
if self.stop_error is not None:
raise self.stop_error
def make_application(
*,
dispatcher: FakeDispatcher,
runtime: FakeRuntime | None,
bot: FakeBot | None = None,
storage_lifecycle: FakeStorageLifecycle | None = None,
) -> ApplicationComposition:
return ApplicationComposition(
bot=bot or FakeBot(), # type: ignore[arg-type]
dispatcher=dispatcher, # type: ignore[arg-type]
trade_stream_runtime=runtime,
market_data_storage_lifecycle=storage_lifecycle,
)
def test_storage_lifecycle_wraps_root_runtime_tasks() -> None:
async def scenario() -> None:
calls: list[str] = []
storage = FakeStorageLifecycle(calls=calls)
dispatcher = FakeDispatcher(
return_immediately=True,
calls=calls,
)
runtime = FakeRuntime(calls=calls)
await run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
storage_lifecycle=storage,
)
)
assert storage.start_calls == 1
assert storage.stop_calls == 1
assert calls.index("storage.start") < calls.index(
"polling.start"
)
assert calls.index("storage.start") < calls.index(
"runtime.run"
)
assert calls.index("runtime.stop") < calls.index(
"storage.stop"
)
asyncio.run(scenario())
def test_storage_startup_failure_prevents_root_task_start() -> None:
async def scenario() -> None:
expected = RuntimeError("storage startup failed")
storage = FakeStorageLifecycle(start_error=expected)
dispatcher = FakeDispatcher()
runtime = FakeRuntime()
bot = FakeBot()
with pytest.raises(RuntimeError) as error_info:
await run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
bot=bot,
storage_lifecycle=storage,
)
)
assert error_info.value is expected
assert dispatcher.started.is_set() is False
assert runtime.started.is_set() is False
assert runtime.stop_calls == 1
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
asyncio.run(scenario())
def test_cancellation_waits_for_storage_startup_before_cleanup() -> None:
async def scenario() -> None:
start_release = threading.Event()
storage = FakeStorageLifecycle(
start_release=start_release,
)
dispatcher = FakeDispatcher()
runtime = FakeRuntime()
bot = FakeBot()
task = asyncio.create_task(
run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
bot=bot,
storage_lifecycle=storage,
)
)
)
entered = await asyncio.to_thread(
storage.start_entered.wait,
1.0,
)
assert entered is True
task.cancel()
try:
await asyncio.sleep(0)
await asyncio.sleep(0)
assert task.done() is False
finally:
start_release.set()
with pytest.raises(asyncio.CancelledError):
await task
assert dispatcher.started.is_set() is False
assert runtime.started.is_set() is False
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
await asyncio.sleep(0)
assert not {
child.get_name()
for child in asyncio.all_tasks()
if child is not asyncio.current_task()
and not child.done()
and child.get_name()
in {
"application-shutdown",
"market-data-storage-shutdown",
"market-data-storage-startup",
"telegram-polling",
"trade-stream-runtime",
}
}
asyncio.run(scenario())
def test_storage_shutdown_error_is_reported_and_bot_still_closes() -> None:
async def scenario() -> None:
expected = RuntimeError("storage close failed")
storage = FakeStorageLifecycle(stop_error=expected)
dispatcher = FakeDispatcher(return_immediately=True)
bot = FakeBot()
with pytest.raises(RuntimeError) as error_info:
await run_application(
make_application(
dispatcher=dispatcher,
runtime=None,
bot=bot,
storage_lifecycle=storage,
)
)
assert error_info.value is expected
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
asyncio.run(scenario())
def test_storage_shutdown_error_does_not_replace_runtime_error() -> None:
async def scenario() -> None:
runtime_error = RuntimeError("runtime failed")
storage_error = RuntimeError("storage close failed")
storage = FakeStorageLifecycle(stop_error=storage_error)
dispatcher = FakeDispatcher()
runtime = FakeRuntime(
return_immediately=True,
error=runtime_error,
)
with pytest.raises(RuntimeError) as error_info:
await run_application(
make_application(
dispatcher=dispatcher,
runtime=runtime,
storage_lifecycle=storage,
)
)
assert error_info.value is runtime_error
assert any(
"cleanup also failed" in note
for note in getattr(runtime_error, "__notes__", ())
)
asyncio.run(scenario())
def test_repeated_cancellation_does_not_interrupt_storage_shutdown(
) -> None:
async def scenario() -> None:
stop_release = threading.Event()
storage = FakeStorageLifecycle(
stop_release=stop_release,
)
dispatcher = FakeDispatcher(return_immediately=True)
bot = FakeBot()
task = asyncio.create_task(
run_application(
make_application(
dispatcher=dispatcher,
runtime=None,
bot=bot,
storage_lifecycle=storage,
)
)
)
entered = await asyncio.to_thread(
storage.stop_entered.wait,
1.0,
)
assert entered is True
task.cancel()
task.cancel()
try:
await asyncio.sleep(0)
assert task.done() is False
finally:
stop_release.set()
with pytest.raises(asyncio.CancelledError):
await task
assert storage.stop_calls == 1
assert bot.session.close_calls == 1
asyncio.run(scenario())
def test_disabled_runtime_runs_only_polling_and_closes_bot() -> None:
async def scenario() -> None:
dispatcher = FakeDispatcher(return_immediately=True)

View File

@@ -0,0 +1,232 @@
from __future__ import annotations
from typing import Any, cast
import pytest
from psycopg.conninfo import conninfo_to_dict
from src.bootstrap.market_data_storage import (
MarketDataStorageLifecycle,
MarketDataStorageLifecycleProtocol,
build_market_data_storage,
)
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
def make_settings(
*,
storage_enabled: bool = True,
trade_stream_enabled: bool = True,
) -> Settings:
return Settings(
bot_token="test-token",
bot_parse_mode="HTML",
app_env="test",
log_level="INFO",
tz="UTC",
exchange_enabled=True,
exchange_name=" dzengi ",
exchange_base_url="https://rest.example.test",
exchange_ws_url="",
exchange_api_key="",
exchange_api_secret="",
exchange_timeout_sec=10,
exchange_testnet=True,
default_symbol="BTC/USD_LEVERAGE",
trade_stream=TradeStreamSettings(
enabled=trade_stream_enabled,
websocket_url="wss://stream.example.test",
symbols=("BTC/USD_LEVERAGE",),
open_timeout_seconds=10.0,
probe_timeout_seconds=20.0,
close_timeout_seconds=10.0,
heartbeat_timeout_seconds=30.0,
scheduler_interval_seconds=5.0,
recovery_window_ms=3_599_999,
),
db_host="db.example.test",
db_port=5544,
db_name="dzentra",
db_user="market-data",
db_password="p@ss word",
market_data_storage=MarketDataStorageSettings(
enabled=storage_enabled,
pool_min_size=2,
pool_max_size=6,
pool_timeout_seconds=7.5,
),
debug_enabled=False,
journal_debug_enabled=False,
)
class RecordingPool:
def __init__(
self,
*,
close_error: BaseException | None = None,
) -> None:
self.close_error = close_error
self.events: list[str] = []
def open(self) -> None:
self.events.append("pool.open")
def close(self) -> None:
self.events.append("pool.close")
if self.close_error is not None:
raise self.close_error
class RecordingMigrationRunner:
def __init__(
self,
*,
pool: RecordingPool,
error: BaseException | None = None,
) -> None:
self._pool = pool
self._error = error
self.calls = 0
def run(self) -> tuple[int, ...]:
self.calls += 1
self._pool.events.append("migrations.run")
if self._error is not None:
raise self._error
return (1,)
def make_lifecycle(
*,
migration_error: BaseException | None = None,
close_error: BaseException | None = None,
) -> tuple[
MarketDataStorageLifecycle,
RecordingPool,
RecordingMigrationRunner,
]:
pool = RecordingPool(close_error=close_error)
runner = RecordingMigrationRunner(
pool=pool,
error=migration_error,
)
lifecycle = MarketDataStorageLifecycle(
connection_pool=pool,
migration_runner=runner,
)
return lifecycle, pool, runner
def test_disabled_storage_builds_nothing() -> None:
assert build_market_data_storage(
make_settings(storage_enabled=False),
) is None
def test_enabled_storage_requires_enabled_trade_stream() -> None:
with pytest.raises(RuntimeError, match="Trade Stream"):
build_market_data_storage(
make_settings(trade_stream_enabled=False),
)
def test_builds_shared_graph_without_opening_pool() -> None:
composition = build_market_data_storage(make_settings())
assert composition is not None
assert composition.connection_pool.is_open is False
assert composition.lifecycle.started is False
migration_provider = cast(
Any,
composition.migration_runner._connection_provider,
)
repository_provider = cast(
Any,
composition.trade_repository._connection_provider,
)
assert migration_provider.__self__ is composition.connection_pool
assert repository_provider.__self__ is composition.connection_pool
assert (
composition.trade_observation_sink._trade_storage
is composition.trade_repository
)
assert composition.trade_observation_sink._venue == "dzengi"
conninfo = conninfo_to_dict(composition.connection_pool._conninfo)
assert conninfo == {
"host": "db.example.test",
"port": "5544",
"dbname": "dzentra",
"user": "market-data",
"password": "p@ss word",
}
def test_lifecycle_opens_pool_before_migrations_and_is_idempotent() -> None:
lifecycle, pool, runner = make_lifecycle()
lifecycle.start()
lifecycle.start()
assert isinstance(lifecycle, MarketDataStorageLifecycleProtocol)
assert lifecycle.started is True
assert pool.events == [
"pool.open",
"migrations.run",
]
assert runner.calls == 1
lifecycle.stop()
lifecycle.stop()
assert lifecycle.started is False
assert pool.events == [
"pool.open",
"migrations.run",
"pool.close",
"pool.close",
]
def test_migration_failure_closes_pool_and_preserves_error() -> None:
migration_error = RuntimeError("migration failed")
lifecycle, pool, _ = make_lifecycle(
migration_error=migration_error,
)
with pytest.raises(RuntimeError) as error_info:
lifecycle.start()
assert error_info.value is migration_error
assert lifecycle.started is False
assert pool.events == [
"pool.open",
"migrations.run",
"pool.close",
]
def test_migration_error_keeps_cleanup_failure_as_note() -> None:
migration_error = RuntimeError("migration failed")
lifecycle, _, _ = make_lifecycle(
migration_error=migration_error,
close_error=RuntimeError("close failed"),
)
with pytest.raises(RuntimeError) as error_info:
lifecycle.start()
assert error_info.value is migration_error
assert migration_error.__notes__ == [
"Market Data Storage startup cleanup also failed: RuntimeError."
]

View File

@@ -5,9 +5,10 @@ import json
import threading
import time
from collections.abc import Awaitable, Callable
from typing import Any
from typing import Any, cast
import pytest
from aiogram import Bot, Dispatcher
from websockets.protocol import State
import src.bootstrap.trade_stream_runtime as production_factory
@@ -15,7 +16,11 @@ from src.bootstrap.application import (
ApplicationComposition,
run_application,
)
from src.core.config import Settings, TradeStreamSettings
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
from src.market_data.acquisition.adapters.dzengi.websocket_transport import (
DzengiWebSocketTransport,
)
@@ -37,6 +42,8 @@ SYMBOL = "BTC/USD_LEVERAGE"
_OWNED_TASK_NAMES = frozenset(
{
"application-shutdown",
"market-data-storage-shutdown",
"market-data-storage-startup",
"telegram-polling",
"trade-stream-receive",
"trade-stream-runtime",
@@ -295,6 +302,12 @@ def make_settings(
db_name="test",
db_user="test",
db_password="test",
market_data_storage=MarketDataStorageSettings(
enabled=False,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=10.0,
),
debug_enabled=False,
journal_debug_enabled=False,
)
@@ -367,8 +380,8 @@ def make_application(
bot: FakeBot,
) -> ApplicationComposition:
return ApplicationComposition(
bot=bot, # type: ignore[arg-type]
dispatcher=dispatcher, # type: ignore[arg-type]
bot=cast(Bot, cast(object, bot)),
dispatcher=cast(Dispatcher, cast(object, dispatcher)),
trade_stream_runtime=runtime,
)
@@ -400,14 +413,23 @@ async def assert_no_owned_tasks() -> None:
def state_store_from(
runtime: TradeStreamProductionRuntime,
) -> TradeStreamStateStore:
runtime_graph: Any = runtime
return (
runtime
runtime_graph
._reconnect_recovery_coordinator
._recovery_coordinator
._state_store
)
def subscription_keys_from(
runtime: TradeStreamProductionRuntime,
) -> tuple[str, ...]:
runtime_graph: Any = runtime
return runtime_graph._subscription_manager.subscription_keys
def test_disabled_feature_runs_only_telegram_without_runtime_graph(
monkeypatch: pytest.MonkeyPatch,
) -> None:
@@ -529,7 +551,7 @@ def test_production_factory_processes_ack_trade_and_shutdown(
assert runtime.state is (
TradeStreamProductionRuntimeState.STOPPED
)
assert runtime._subscription_manager.subscription_keys == ()
assert subscription_keys_from(runtime) == ()
assert bot.session.close_calls == 1
assert len(transports) == 1
await assert_no_owned_tasks()
@@ -579,7 +601,7 @@ def test_connection_startup_failure_is_fatal_and_leaves_no_tasks(
TradeStreamProductionRuntimeState.FAILED
)
assert runtime._session.is_connected is False
assert runtime._subscription_manager.subscription_keys == ()
assert subscription_keys_from(runtime) == ()
assert bot.session.close_calls == 1
await assert_no_owned_tasks()
@@ -633,7 +655,7 @@ def test_subscription_startup_failure_rolls_back_concrete_graph(
TradeStreamProductionRuntimeState.FAILED
)
assert runtime._session.is_connected is False
assert runtime._subscription_manager.subscription_keys == ()
assert subscription_keys_from(runtime) == ()
assert bot.session.close_calls == 1
await assert_no_owned_tasks()

View File

@@ -1,13 +1,28 @@
from __future__ import annotations
from typing import Any
from src.bootstrap.trade_stream_runtime import (
build_trade_stream_production_runtime,
)
from src.core.config import Settings, TradeStreamSettings
from src.core.config import (
MarketDataStorageSettings,
Settings,
TradeStreamSettings,
)
from src.market_data.acquisition.runtime.trade_stream_production_runtime import (
TradeStreamProductionRuntime,
TradeStreamProductionRuntimeState,
)
from src.market_data.acquisition.models.trade import Trade
class RecordingTradeObservationSink:
def __init__(self) -> None:
self.observations: list[Trade] = []
def persist(self, trade: Trade) -> None:
self.observations.append(trade)
def make_settings(
@@ -49,6 +64,12 @@ def make_settings(
db_name="test",
db_user="test",
db_password="test",
market_data_storage=MarketDataStorageSettings(
enabled=False,
pool_min_size=1,
pool_max_size=4,
pool_timeout_seconds=10.0,
),
debug_enabled=False,
journal_debug_enabled=False,
)
@@ -80,33 +101,60 @@ def test_uses_one_shared_stateful_dependency_graph() -> None:
assert isinstance(runtime, TradeStreamProductionRuntime)
transport = runtime._transport
service = runtime._trade_stream_service
reconnect_recovery = runtime._reconnect_recovery_coordinator
runtime_graph: Any = runtime
transport = runtime_graph._transport
service = runtime_graph._trade_stream_service
reconnect_recovery = runtime_graph._reconnect_recovery_coordinator
recovery = reconnect_recovery._recovery_coordinator
assert runtime._session._transport is transport
assert runtime._subscription_manager._transport is transport
assert runtime._runtime_scheduler.liveness_probe is transport
assert runtime_graph._session._transport is transport
assert runtime_graph._subscription_manager._transport is transport
assert runtime_graph._runtime_scheduler.liveness_probe is transport
assert (
service._consistency_controller
is recovery._recovery_controller._consistency_controller
)
assert runtime._live_processing_gate is (
assert runtime_graph._live_processing_gate is (
reconnect_recovery.live_processing_gate
)
assert runtime._runtime_scheduler.runtime_supervisor is (
runtime._runtime_supervisor
assert runtime_graph._runtime_scheduler.runtime_supervisor is (
runtime_graph._runtime_supervisor
)
def test_optional_storage_sink_is_shared_by_live_and_recovery() -> None:
sink = RecordingTradeObservationSink()
runtime = build_trade_stream_production_runtime(
make_settings(),
trade_observation_sink=sink,
)
assert isinstance(runtime, TradeStreamProductionRuntime)
runtime_graph: Any = runtime
live_controller = (
runtime_graph._trade_stream_service._consistency_controller
)
recovery_controller = (
runtime_graph
._reconnect_recovery_coordinator
._recovery_coordinator
._recovery_controller
._consistency_controller
)
assert live_controller is recovery_controller
assert live_controller._trade_observation_sink is sink
def test_applies_explicit_transport_and_runtime_settings() -> None:
settings = make_settings()
runtime = build_trade_stream_production_runtime(settings)
assert isinstance(runtime, TradeStreamProductionRuntime)
transport = runtime._transport
runtime_graph: Any = runtime
transport = runtime_graph._transport
assert transport._url == "wss://stream.example.test/root/connect"
assert transport._headers == {
@@ -119,17 +167,17 @@ def test_applies_explicit_transport_and_runtime_settings() -> None:
assert transport._close_timeout == 9.0
assert transport._ping_interval is None
assert transport._ping_timeout is None
assert runtime._symbols == (
assert runtime_graph._symbols == (
"BTC/USD_LEVERAGE",
"ETH/USD_LEVERAGE",
)
assert runtime._runtime_scheduler.interval_seconds == 6.0
assert runtime_graph._runtime_scheduler.interval_seconds == 6.0
assert (
runtime._runtime_supervisor._heartbeat_monitor.timeout_seconds
runtime_graph._runtime_supervisor._heartbeat_monitor.timeout_seconds
== 31.0
)
assert (
runtime._reconnect_recovery_coordinator
runtime_graph._reconnect_recovery_coordinator
._recovery_coordinator
._window_planner
.max_window_ms
@@ -143,8 +191,9 @@ def test_recovery_rest_client_reuses_settings_snapshot() -> None:
runtime = build_trade_stream_production_runtime(settings)
assert isinstance(runtime, TradeStreamProductionRuntime)
runtime_graph: Any = runtime
document_source = (
runtime._reconnect_recovery_coordinator
runtime_graph._reconnect_recovery_coordinator
._recovery_coordinator
._recovery_controller
._document_source
@@ -154,4 +203,4 @@ def test_recovery_rest_client_reuses_settings_snapshot() -> None:
assert rest_client.settings is settings
assert rest_client.base_url == "https://rest.example.test"
assert rest_client.timeout == 17
assert "X-MBX-APIKEY" not in runtime._transport._headers
assert "X-MBX-APIKEY" not in runtime_graph._transport._headers

View File

@@ -17,6 +17,13 @@ _TRADE_STREAM_VARIABLES = (
"TRADE_STREAM_RECOVERY_WINDOW_MS",
)
_MARKET_DATA_STORAGE_VARIABLES = (
"MARKET_DATA_STORAGE_ENABLED",
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS",
)
def prepare_environment(
monkeypatch: pytest.MonkeyPatch,
@@ -28,6 +35,9 @@ def prepare_environment(
for variable in _TRADE_STREAM_VARIABLES:
monkeypatch.delenv(variable, raising=False)
for variable in _MARKET_DATA_STORAGE_VARIABLES:
monkeypatch.delenv(variable, raising=False)
def enable_trade_stream(
monkeypatch: pytest.MonkeyPatch,
@@ -200,3 +210,142 @@ def test_symbols_must_not_contain_empty_items(
match="empty symbols",
):
load_settings()
def test_market_data_storage_is_disabled_by_default(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
settings = load_settings()
assert settings.market_data_storage.enabled is False
assert settings.market_data_storage.pool_min_size == 1
assert settings.market_data_storage.pool_max_size == 4
assert settings.market_data_storage.pool_timeout_seconds == 10.0
def test_disabled_market_data_storage_ignores_dependent_values(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"invalid",
)
settings = load_settings()
assert settings.market_data_storage.enabled is False
def test_market_data_storage_flag_is_strict(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"sometimes",
)
with pytest.raises(
ValueError,
match="MARKET_DATA_STORAGE_ENABLED",
):
load_settings()
def test_enabled_market_data_storage_requires_trade_stream(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
with pytest.raises(
RuntimeError,
match="TRADE_STREAM_ENABLED",
):
load_settings()
def test_enabled_market_data_storage_parses_pool_settings(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
enable_trade_stream(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"2",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"7",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS",
"12.5",
)
storage = load_settings().market_data_storage
assert storage.enabled is True
assert storage.pool_min_size == 2
assert storage.pool_max_size == 7
assert storage.pool_timeout_seconds == 12.5
@pytest.mark.parametrize(
("variable", "value"),
(
("MARKET_DATA_STORAGE_POOL_MIN_SIZE", "0"),
("MARKET_DATA_STORAGE_POOL_MAX_SIZE", "invalid"),
("MARKET_DATA_STORAGE_POOL_TIMEOUT_SECONDS", "nan"),
),
)
def test_enabled_market_data_storage_rejects_invalid_pool_settings(
monkeypatch: pytest.MonkeyPatch,
variable: str,
value: str,
) -> None:
prepare_environment(monkeypatch)
enable_trade_stream(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
monkeypatch.setenv(variable, value)
with pytest.raises(ValueError, match=variable):
load_settings()
def test_pool_max_size_must_not_be_smaller_than_min_size(
monkeypatch: pytest.MonkeyPatch,
) -> None:
prepare_environment(monkeypatch)
enable_trade_stream(monkeypatch)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_ENABLED",
"true",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MIN_SIZE",
"5",
)
monkeypatch.setenv(
"MARKET_DATA_STORAGE_POOL_MAX_SIZE",
"4",
)
with pytest.raises(
ValueError,
match="POOL_MAX_SIZE",
):
load_settings()

View File

@@ -10,6 +10,9 @@ import pytest
from src.market_data.acquisition.consistency.trade_stream_consistency_controller import (
TradeStreamConsistencyController,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.consistency.trade_stream_exceptions import (
TradeConsistencyError,
TradeOrderingError,
@@ -47,6 +50,7 @@ def _trade(
price: Decimal = Decimal("50000.00"),
quantity: Decimal = Decimal("0.25"),
aggressor_side: TradeAggressorSide = TradeAggressorSide.BUY,
source: str = "dzengi",
) -> Trade:
return Trade(
symbol=symbol,
@@ -62,10 +66,29 @@ def _trade(
tzinfo=timezone.utc,
),
aggressor_side=aggressor_side,
source="dzengi",
source=source,
)
class RecordingTradeObservationSink:
def __init__(
self,
*,
error: Exception | None = None,
) -> None:
self.error = error
self.observations: list[Trade] = []
def persist(
self,
trade: Trade,
) -> None:
self.observations.append(trade)
if self.error is not None:
raise self.error
def test_creates_state_for_first_symbol(
controller: TradeStreamConsistencyController,
state_store: TradeStreamStateStore,
@@ -179,4 +202,128 @@ def test_propagates_consistency_error(
trade_id=100,
price=Decimal("50001.00"),
),
)
)
def test_persists_trade_before_advancing_checkpoint(
state_store: TradeStreamStateStore,
) -> None:
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
trade = _trade()
result = controller.accept(trade)
state = state_store.get(trade.symbol)
assert result is trade
assert sink.observations == [trade]
assert state.last_trade is trade
def test_persistence_failure_leaves_checkpoint_unchanged(
state_store: TradeStreamStateStore,
) -> None:
storage_error = RuntimeError("storage failed")
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
first_trade = _trade(trade_id=100)
failed_trade = _trade(trade_id=101)
controller.accept(first_trade)
sink.error = storage_error
with pytest.raises(
RuntimeError,
match="storage failed",
) as error_info:
controller.accept(failed_trade)
state = state_store.get(first_trade.symbol)
assert error_info.value is storage_error
assert state.last_trade is first_trade
assert state.last_trade_id == first_trade.trade_id
assert failed_trade.trade_id not in state._trades
sink.error = None
assert controller.accept(failed_trade) is failed_trade
def test_valid_duplicate_is_persisted_without_checkpoint_advance(
state_store: TradeStreamStateStore,
) -> None:
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
websocket_trade = _trade(
source="dzengi_websocket_trade",
)
rest_duplicate = _trade(
source="dzengi",
)
controller.accept(websocket_trade)
result = controller.accept(rest_duplicate)
state = state_store.get(websocket_trade.symbol)
assert result is None
assert sink.observations == [
websocket_trade,
rest_duplicate,
]
assert state.last_trade is websocket_trade
assert state.last_trade_id == websocket_trade.trade_id
def test_invalid_trades_do_not_reach_persistence_sink(
state_store: TradeStreamStateStore,
) -> None:
sink = RecordingTradeObservationSink()
controller = TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=sink,
)
first_trade = _trade(trade_id=100)
controller.accept(first_trade)
with pytest.raises(TradeOrderingError):
controller.accept(_trade(trade_id=99))
with pytest.raises(TradeConsistencyError):
controller.accept(
_trade(
trade_id=100,
price=Decimal("50001.00"),
)
)
assert sink.observations == [first_trade]
def test_rejects_invalid_trade_observation_sink(
state_store: TradeStreamStateStore,
) -> None:
with pytest.raises(
TypeError,
match="TradeObservationSinkProtocol",
):
TradeStreamConsistencyController(
state_store=state_store,
trade_observation_sink=object(), # type: ignore[arg-type]
)
def test_recording_sink_implements_public_protocol() -> None:
assert isinstance(
RecordingTradeObservationSink(),
TradeObservationSinkProtocol,
)

View File

@@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import json
import threading
from collections import deque
from collections.abc import Awaitable, Callable
from typing import cast
@@ -376,15 +377,20 @@ class FakeTradeStreamService:
subscribe_error: Exception | None = None,
handle_error: Exception | None = None,
subscribe_gate: asyncio.Event | None = None,
handle_entered: threading.Event | None = None,
handle_release: threading.Event | None = None,
) -> None:
self._calls = calls
self._subscribe_error = subscribe_error
self._handle_error = handle_error
self._subscribe_gate = subscribe_gate
self._handle_entered = handle_entered
self._handle_release = handle_release
self.subscribe_calls: list[tuple[str, ...]] = []
self.subscribe_correlation_ids: list[str | None] = []
self.documents: list[object] = []
self.handled = asyncio.Event()
self.handled = threading.Event()
self.handle_thread_ids: list[int] = []
self.subscribe_entered = asyncio.Event()
async def subscribe(
@@ -408,10 +414,17 @@ class FakeTradeStreamService:
self,
document: object,
) -> Trade | None:
self.handle_thread_ids.append(threading.get_ident())
self.documents.append(document)
self._calls.append("service.handle_message")
self.handled.set()
if self._handle_entered is not None:
self._handle_entered.set()
if self._handle_release is not None:
self._handle_release.wait()
if self._handle_error is not None:
raise self._handle_error
@@ -592,6 +605,8 @@ class RuntimeDependencies:
stop_error: Exception | None = None,
subscribe_error: Exception | None = None,
handle_error: Exception | None = None,
handle_entered: threading.Event | None = None,
handle_release: threading.Event | None = None,
clear_error: Exception | None = None,
start_gate: asyncio.Event | None = None,
connected_publish_gate: asyncio.Event | None = None,
@@ -633,6 +648,8 @@ class RuntimeDependencies:
subscribe_error=subscribe_error,
handle_error=handle_error,
subscribe_gate=subscribe_gate,
handle_entered=handle_entered,
handle_release=handle_release,
)
self.live_processing_gate = RuntimeLiveProcessingGate()
self.reconnect_recovery = FakeReconnectRecoveryCoordinator(
@@ -709,6 +726,16 @@ async def wait_until(
raise AssertionError("condition was not reached")
async def wait_for_thread_event(
event: threading.Event,
) -> None:
reached = await asyncio.to_thread(
event.wait,
1.0,
)
assert reached is True
def test_implements_public_protocol_and_uses_slots() -> None:
dependencies = RuntimeDependencies()
@@ -944,7 +971,7 @@ def test_decodes_and_forwards_market_messages(
dependencies.runtime.run(),
)
await dependencies.service.handled.wait()
await wait_for_thread_event(dependencies.service.handled)
await dependencies.runtime.stop()
await runtime_task
@@ -1295,7 +1322,10 @@ def test_transport_error_runs_recovery_and_resumes_receive_loop() -> None:
dependencies.runtime.run(),
)
await dependencies.service.handled.wait()
await wait_for_thread_event(dependencies.service.handled)
await wait_until(
lambda: dependencies.transport.receive_calls == 3
)
assert dependencies.runtime.running is True
await dependencies.runtime.stop()
@@ -1342,7 +1372,7 @@ def test_buffered_market_waits_for_reconnect_recovery() -> None:
assert dependencies.service.documents == []
reconnect_release.set()
await dependencies.service.handled.wait()
await wait_for_thread_event(dependencies.service.handled)
await dependencies.runtime.stop()
await runtime_task
@@ -1568,6 +1598,69 @@ def test_market_handler_error_is_terminal_and_not_wrapped() -> None:
assert dependencies.runtime.state is (
TradeStreamProductionRuntimeState.FAILED
)
assert dependencies.live_processing_gate.failed is True
def test_market_processing_runs_outside_event_loop_thread() -> None:
async def scenario() -> tuple[RuntimeDependencies, int]:
event_loop_thread_id = threading.get_ident()
dependencies = RuntimeDependencies(
incoming=(MARKET_MESSAGE,),
)
runtime_task = asyncio.create_task(
dependencies.runtime.run(),
)
await wait_for_thread_event(dependencies.service.handled)
await dependencies.runtime.stop()
await runtime_task
return dependencies, event_loop_thread_id
dependencies, event_loop_thread_id = asyncio.run(scenario())
assert dependencies.service.handle_thread_ids
assert dependencies.service.handle_thread_ids[0] != event_loop_thread_id
def test_stop_waits_for_inflight_market_processing() -> None:
async def scenario() -> RuntimeDependencies:
handle_entered = threading.Event()
handle_release = threading.Event()
dependencies = RuntimeDependencies(
incoming=(MARKET_MESSAGE,),
handle_entered=handle_entered,
handle_release=handle_release,
)
runtime_task = asyncio.create_task(
dependencies.runtime.run(),
)
await wait_for_thread_event(handle_entered)
stop_task = asyncio.create_task(
dependencies.runtime.stop(),
)
try:
await asyncio.sleep(0)
await asyncio.sleep(0)
assert stop_task.done() is False
assert dependencies.live_processing_gate.locked is True
assert dependencies.runtime._market_processing_task is not None
finally:
handle_release.set()
await stop_task
await runtime_task
return dependencies
dependencies = asyncio.run(scenario())
assert dependencies.runtime.state is (
TradeStreamProductionRuntimeState.STOPPED
)
assert dependencies.runtime._market_processing_task is None
def test_connect_failure_publishes_event_and_rolls_back() -> None:

View File

@@ -13,6 +13,9 @@ import pytest
from src.market_data.acquisition.adapters.dzengi.rest import (
DzengiTradesDocumentSource,
)
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
@@ -315,6 +318,25 @@ class RecordingSleep:
self.calls.append(seconds)
class RecordingTradeObservationSink:
def __init__(
self,
*,
fail_on_trade_id: int | None = None,
) -> None:
self._fail_on_trade_id = fail_on_trade_id
self.observations: list[Trade] = []
def persist(
self,
trade: Trade,
) -> None:
self.observations.append(trade)
if trade.trade_id == self._fail_on_trade_id:
raise RuntimeError("storage failed")
@dataclass(slots=True)
class CompositionDependencies:
session: FakeSession
@@ -336,6 +358,7 @@ def create_composition(
scheduler_interval_seconds: float = 1.0,
max_recovery_window_ms: int = 3_599_999,
probe_results: tuple[bool, ...] = (True,),
trade_observation_sink: TradeObservationSinkProtocol | None = None,
) -> tuple[
TradeStreamRuntimeComposition,
CompositionDependencies,
@@ -370,6 +393,7 @@ def create_composition(
symbols=(SYMBOL,),
heartbeat_timeout_seconds=heartbeat_timeout_seconds,
scheduler_interval_seconds=scheduler_interval_seconds,
trade_observation_sink=trade_observation_sink,
max_recovery_window_ms=max_recovery_window_ms,
heartbeat_clock=dependencies.heartbeat_clock,
recovery_end_time_clock=(
@@ -491,6 +515,120 @@ def test_live_stream_and_recovery_share_consistency_state() -> None:
)
def test_live_and_recovery_share_optional_persistence_sink() -> None:
sink = RecordingTradeObservationSink()
recovered_trade_id = 101
composition, _ = create_composition(
trade_observation_sink=sink,
recovery_document=[
make_raw_trade(
trade_id=recovered_trade_id,
timestamp=CHECKPOINT_TIME_MS + 1_000,
),
],
)
live_trade = (
composition.trade_stream_acquisition_service.handle_message(
{
"destination": "internal.trade",
}
)
)
recovery_result = composition.runtime_recovery_coordinator.recover(
symbol=SYMBOL,
recovery_end_time=RECOVERY_END_TIME_MS,
)
assert composition.trade_observation_sink is sink
assert isinstance(sink, TradeObservationSinkProtocol)
assert sink.observations[0] is live_trade
assert sink.observations[1] is recovery_result.last_trade
assert [
trade.trade_id
for trade in sink.observations
] == [100, recovered_trade_id]
def test_recovery_duplicate_updates_persistence_without_checkpoint_change(
) -> None:
live_trade = Trade(
symbol=SYMBOL,
trade_id=100,
price=Decimal("64556.00"),
quantity=Decimal("0.003"),
executed_at=CHECKPOINT_TIME,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi_websocket_trade",
)
sink = RecordingTradeObservationSink()
composition, _ = create_composition(
trade=live_trade,
trade_observation_sink=sink,
recovery_document=[
make_raw_trade(
trade_id=live_trade.trade_id,
timestamp=CHECKPOINT_TIME_MS,
),
],
)
composition.trade_stream_acquisition_service.handle_message(
{
"destination": "internal.trade",
}
)
result = composition.runtime_recovery_coordinator.recover(
symbol=SYMBOL,
recovery_end_time=RECOVERY_END_TIME_MS,
)
state = composition.state_store.get(SYMBOL)
assert result.is_empty is True
assert len(sink.observations) == 2
assert sink.observations[0] is live_trade
assert sink.observations[1].source == "dzengi"
assert state.last_trade is live_trade
def test_recovery_persistence_failure_does_not_advance_checkpoint() -> None:
sink = RecordingTradeObservationSink(
fail_on_trade_id=101,
)
composition, _ = create_composition(
trade_observation_sink=sink,
recovery_document=[
make_raw_trade(
trade_id=101,
timestamp=CHECKPOINT_TIME_MS + 1_000,
),
],
)
live_trade = (
composition.trade_stream_acquisition_service.handle_message(
{
"destination": "internal.trade",
}
)
)
with pytest.raises(
RuntimeError,
match="storage failed",
):
composition.runtime_recovery_coordinator.recover(
symbol=SYMBOL,
recovery_end_time=RECOVERY_END_TIME_MS,
)
state = composition.state_store.get(SYMBOL)
assert state.last_trade is live_trade
assert state.last_trade_id == 100
assert [trade.trade_id for trade in sink.observations] == [100, 101]
def test_runtime_components_share_lifecycle_dependencies() -> None:
composition, dependencies = create_composition()

View File

@@ -0,0 +1,121 @@
from __future__ import annotations
from datetime import datetime
from typing import Any
import pytest
from src.market_data.storage import (
CandleStorageProtocol,
MarketDataBatchWriteResult,
MarketDataWriteResult,
MarketDataWriteStatus,
QuoteStorageProtocol,
TradeStorageProtocol,
)
class RecordingMarketDataStorage:
def store_trade(
self,
*,
venue: str,
trade: Any,
observed_at: datetime,
) -> MarketDataWriteResult:
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def store_trades(
self,
*,
venue: str,
trades: tuple[Any, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
return MarketDataBatchWriteResult(
inserted_count=len(trades),
duplicate_count=0,
provenance_updated_count=0,
)
def store_quote(
self,
*,
venue: str,
quote: Any,
) -> MarketDataWriteResult:
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def store_candle_revision(
self,
*,
venue: str,
candle: Any,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def test_storage_protocols_are_runtime_checkable() -> None:
storage = RecordingMarketDataStorage()
assert isinstance(storage, TradeStorageProtocol)
assert isinstance(storage, QuoteStorageProtocol)
assert isinstance(storage, CandleStorageProtocol)
@pytest.mark.parametrize("status", tuple(MarketDataWriteStatus))
def test_write_result_accepts_every_defined_status(
status: MarketDataWriteStatus,
) -> None:
result = MarketDataWriteResult(status=status)
assert result.status is status
def test_write_result_rejects_unknown_status() -> None:
with pytest.raises(TypeError, match="MarketDataWriteStatus"):
MarketDataWriteResult(status="inserted") # type: ignore[arg-type]
def test_batch_result_reports_total_count() -> None:
result = MarketDataBatchWriteResult(
inserted_count=2,
duplicate_count=3,
provenance_updated_count=5,
)
assert result.total_count == 10
@pytest.mark.parametrize(
"field_name",
(
"inserted_count",
"duplicate_count",
"provenance_updated_count",
),
)
def test_batch_result_rejects_negative_count(field_name: str) -> None:
values = {
"inserted_count": 0,
"duplicate_count": 0,
"provenance_updated_count": 0,
}
values[field_name] = -1
with pytest.raises(ValueError, match=field_name):
MarketDataBatchWriteResult(**values)
@pytest.mark.parametrize("invalid_count", (True, 1.5, "1", None))
def test_batch_result_rejects_non_integer_count(
invalid_count: object,
) -> None:
with pytest.raises(TypeError, match="inserted_count"):
MarketDataBatchWriteResult(
inserted_count=invalid_count, # type: ignore[arg-type]
duplicate_count=0,
provenance_updated_count=0,
)

View File

@@ -0,0 +1,206 @@
from __future__ import annotations
from dataclasses import dataclass, field
from datetime import datetime, timezone
from decimal import Decimal
import pytest
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.storage import (
CandleStorageProtocol,
MarketDataBatchWriteResult,
MarketDataStorage,
MarketDataWriteResult,
MarketDataWriteStatus,
QuoteStorageProtocol,
TradeStorageProtocol,
)
NOW = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
@dataclass
class RecordingTradeStorage:
calls: list[tuple[str, object, datetime]] = field(default_factory=list)
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
self.calls.append((venue, trade, observed_at))
return MarketDataWriteResult(MarketDataWriteStatus.INSERTED)
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
self.calls.append((venue, trades, observed_at))
return MarketDataBatchWriteResult(len(trades), 0, 0)
@dataclass
class RecordingQuoteStorage:
calls: list[tuple[str, Quote]] = field(default_factory=list)
def store_quote(
self,
*,
venue: str,
quote: Quote,
) -> MarketDataWriteResult:
self.calls.append((venue, quote))
return MarketDataWriteResult(MarketDataWriteStatus.DUPLICATE)
@dataclass
class RecordingCandleStorage:
calls: list[tuple[str, Candle, datetime, bool]] = field(
default_factory=list
)
def store_candle_revision(
self,
*,
venue: str,
candle: Candle,
observed_at: datetime,
is_final: bool,
) -> MarketDataWriteResult:
self.calls.append((venue, candle, observed_at, is_final))
return MarketDataWriteResult(
MarketDataWriteStatus.PROVENANCE_UPDATED
)
def _trade() -> Trade:
return Trade(
symbol="BTC/USD_LEVERAGE",
trade_id=1,
price=Decimal("100"),
quantity=Decimal("1"),
executed_at=NOW,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi",
)
def _quote() -> Quote:
return Quote(
symbol="BTC/USD_LEVERAGE",
last_price=Decimal("100"),
bid_price=Decimal("99"),
ask_price=Decimal("101"),
exchange_timestamp=NOW,
received_at=NOW,
source="dzengi",
)
def _candle() -> Candle:
return Candle(
symbol="BTC/USD_LEVERAGE",
interval="1m",
open_time=NOW,
open_price=Decimal("100"),
high_price=Decimal("110"),
low_price=Decimal("90"),
close_price=Decimal("105"),
volume=Decimal("10"),
source="dzengi",
)
def _storage() -> tuple[
MarketDataStorage,
RecordingTradeStorage,
RecordingQuoteStorage,
RecordingCandleStorage,
]:
trades = RecordingTradeStorage()
quotes = RecordingQuoteStorage()
candles = RecordingCandleStorage()
return (
MarketDataStorage(
trade_storage=trades,
quote_storage=quotes,
candle_storage=candles,
),
trades,
quotes,
candles,
)
def test_facade_matches_all_write_protocols() -> None:
storage, *_ = _storage()
assert isinstance(storage, TradeStorageProtocol)
assert isinstance(storage, QuoteStorageProtocol)
assert isinstance(storage, CandleStorageProtocol)
def test_facade_delegates_without_replacing_models_or_results() -> None:
storage, trades, quotes, candles = _storage()
trade = _trade()
quote = _quote()
candle = _candle()
trade_result = storage.store_trade(
venue="venue",
trade=trade,
observed_at=NOW,
)
batch_result = storage.store_trades(
venue="venue",
trades=(trade,),
observed_at=NOW,
)
quote_result = storage.store_quote(venue="venue", quote=quote)
candle_result = storage.store_candle_revision(
venue="venue",
candle=candle,
observed_at=NOW,
is_final=True,
)
assert trades.calls == [
("venue", trade, NOW),
("venue", (trade,), NOW),
]
assert quotes.calls == [("venue", quote)]
assert candles.calls == [("venue", candle, NOW, True)]
assert trade_result.status is MarketDataWriteStatus.INSERTED
assert batch_result.inserted_count == 1
assert quote_result.status is MarketDataWriteStatus.DUPLICATE
assert candle_result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
@pytest.mark.parametrize(
"dependency_name",
("trade_storage", "quote_storage", "candle_storage"),
)
def test_facade_rejects_dependency_outside_protocol(
dependency_name: str,
) -> None:
dependencies: dict[str, object] = {
"trade_storage": RecordingTradeStorage(),
"quote_storage": RecordingQuoteStorage(),
"candle_storage": RecordingCandleStorage(),
}
dependencies[dependency_name] = object()
with pytest.raises(TypeError, match=dependency_name):
MarketDataStorage(**dependencies) # type: ignore[arg-type]

View File

@@ -0,0 +1,862 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
import re
import threading
from typing import Any
import pytest
from psycopg.sql import Composable
from src.market_data.storage import (
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
MarketDataPartitionResult,
MarketDataPartitionType,
MarketDataRetentionPolicy,
MarketDataStorageConfigurationError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
PostgresMarketDataPartitionManager,
PostgresMarketDataRetentionService,
build_monthly_partition,
)
UTC = timezone.utc
JUNE = datetime(2026, 6, 1, tzinfo=UTC)
JULY = datetime(2026, 7, 1, tzinfo=UTC)
AUGUST = datetime(2026, 8, 1, tzinfo=UTC)
RegistryKey = tuple[str, datetime]
RegistryRow = tuple[str, datetime, datetime, str]
RelationState = tuple[bool, bool, str | None]
@dataclass
class PartitionDatabase:
registry: dict[RegistryKey, RegistryRow] = field(default_factory=dict)
relations: dict[str, RelationState] = field(default_factory=dict)
def _partition_bound(range_start: datetime, range_end: datetime) -> str:
return f"FOR VALUES FROM ({range_start.isoformat()}) TO ({range_end.isoformat()})"
class PartitionCursor:
def __init__(self, connection: PartitionConnection) -> None:
self._connection = connection
self._fetchone: object = None
self._fetchall: list[tuple[Any, ...]] = []
self.rowcount = -1
def __enter__(self) -> PartitionCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str | Composable,
parameters: tuple[Any, ...] | None = None,
) -> None:
text = (
statement
if isinstance(statement, str)
else statement.as_string()
)
normalized = " ".join(text.split())
self._connection.calls.append((normalized, parameters))
self._fetchone = None
self._fetchall = []
self.rowcount = -1
if self._connection.interrupt_on in normalized:
raise KeyboardInterrupt
if self._connection.fail_on in normalized:
raise RuntimeError("database failed")
if normalized.startswith("SELECT pg_advisory_xact_lock"):
self._connection.acquire_advisory_lock()
return
if normalized.startswith("SELECT partition_name"):
assert parameters is not None
data_type, boundary = parameters
if "range_end <=" in normalized:
self._fetchall = sorted(
(
row
for (kind, _), row
in self._connection.working.registry.items()
if kind == data_type and row[2] <= boundary
),
key=lambda row: row[1],
)
else:
self._fetchone = self._connection.working.registry.get(
(data_type, boundary)
)
return
if "FROM pg_catalog.pg_class AS child" in normalized:
assert parameters is not None
partition_name = parameters[1]
self._fetchone = self._connection.working.relations.get(
partition_name,
(False, False, None),
)
return
if normalized.startswith("LOCK TABLE"):
if "SHARE ROW EXCLUSIVE" in normalized:
self._connection.acquire_write_lock()
return
if normalized.startswith("CREATE TABLE"):
partition_name = _second_qualified_identifier(normalized)
self._connection.working.relations[partition_name] = (
True,
False,
None,
)
return
if normalized.startswith("ALTER TABLE") and "ADD CONSTRAINT" in normalized:
return
if normalized.startswith("WITH moved_rows AS"):
self.rowcount = self._connection.moved_row_count
return
if normalized.startswith("ALTER TABLE") and "ATTACH PARTITION" in normalized:
partition_name = _attached_partition_identifier(normalized)
assert parameters is None
range_start, range_end = _timestamp_literals(normalized)
self._connection.working.relations[partition_name] = (
True,
True,
_partition_bound(range_start, range_end),
)
return
if normalized.startswith(
"INSERT INTO market_data.partition_registry"
):
assert parameters is not None
data_type, name, range_start, range_end, partition_bound = parameters
self._connection.working.registry[(data_type, range_start)] = (
name,
range_start,
range_end,
partition_bound,
)
self.rowcount = 1
return
if normalized.startswith("SELECT COUNT(*) FROM"):
partition_name = _first_qualified_identifier(normalized)
self._fetchone = (
self._connection.partition_row_counts.get(
partition_name,
0,
),
)
return
if normalized.startswith("DROP TABLE"):
partition_name = _first_qualified_identifier(normalized)
self._connection.working.relations.pop(partition_name, None)
return
if normalized.startswith(
"DELETE FROM market_data.partition_registry"
):
assert parameters is not None
data_type, name, range_start, range_end = parameters
key = (data_type, range_start)
row = self._connection.working.registry.get(key)
if row is not None and row[:3] == (
name,
range_start,
range_end,
):
del self._connection.working.registry[key]
self.rowcount = 1
else:
self.rowcount = 0
return
if normalized.startswith("DELETE FROM"):
parent_name = _first_qualified_identifier(normalized)
self.rowcount = self._connection.partial_delete_counts.get(
parent_name,
0,
)
return
raise AssertionError(f"Unexpected SQL: {normalized}")
def fetchone(self) -> object:
return self._fetchone
def fetchall(self) -> list[tuple[Any, ...]]:
return list(self._fetchall)
def _quoted_identifiers(statement: str) -> list[str]:
return statement.split('"')[1::2]
def _first_qualified_identifier(statement: str) -> str:
identifiers = _quoted_identifiers(statement)
return identifiers[1]
def _second_qualified_identifier(statement: str) -> str:
identifiers = _quoted_identifiers(statement)
return identifiers[1]
def _attached_partition_identifier(statement: str) -> str:
identifiers = _quoted_identifiers(statement)
return identifiers[3]
def _timestamp_literals(statement: str) -> tuple[datetime, datetime]:
values = tuple(
datetime.fromisoformat(value)
for value in re.findall(
r"'([^']+)'::timestamptz",
statement,
)
)
if len(values) != 2:
raise AssertionError(
f"Expected two timestamptz literals: {statement}"
)
return values
class PartitionConnection:
def __init__(self, database: PartitionDatabase) -> None:
self._database = database
self.working = deepcopy(database)
self.calls: list[tuple[str, tuple[Any, ...] | None]] = []
self.exit_exception_types: list[type[BaseException] | None] = []
self.fail_on = "__never_fail__"
self.interrupt_on = "__never_interrupt__"
self.moved_row_count = 0
self.partition_row_counts: dict[str, int] = {}
self.partial_delete_counts: dict[str, int] = {}
self.advisory_lock: threading.Lock | None = None
self.advisory_attempted: threading.Event | None = None
self.advisory_acquired: threading.Event | None = None
self.advisory_continue: threading.Event | None = None
self.write_lock: threading.Lock | None = None
self.write_lock_acquired: threading.Event | None = None
self.write_lock_continue: threading.Event | None = None
self._holds_advisory_lock = False
self._holds_write_lock = False
def __enter__(self) -> PartitionConnection:
self.working = deepcopy(self._database)
return self
def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: object,
) -> None:
self.exit_exception_types.append(exception_type)
try:
if exception_type is None:
self._database.registry = self.working.registry
self._database.relations = self.working.relations
finally:
if self._holds_write_lock:
assert self.write_lock is not None
self._holds_write_lock = False
self.write_lock.release()
if self._holds_advisory_lock:
assert self.advisory_lock is not None
self._holds_advisory_lock = False
self.advisory_lock.release()
return None
def cursor(self) -> PartitionCursor:
return PartitionCursor(self)
def acquire_advisory_lock(self) -> None:
lock = self.advisory_lock
if lock is None or self._holds_advisory_lock:
return
if self.advisory_attempted is not None:
self.advisory_attempted.set()
lock.acquire()
self._holds_advisory_lock = True
self.working = deepcopy(self._database)
if self.advisory_acquired is not None:
self.advisory_acquired.set()
if self.advisory_continue is not None:
if not self.advisory_continue.wait(timeout=3):
raise TimeoutError("advisory test gate timed out")
def acquire_write_lock(self) -> None:
lock = self.write_lock
if lock is None or self._holds_write_lock:
return
lock.acquire()
self._holds_write_lock = True
if self.write_lock_acquired is not None:
self.write_lock_acquired.set()
if self.write_lock_continue is not None:
if not self.write_lock_continue.wait(timeout=3):
raise TimeoutError("retention test gate timed out")
@dataclass
class PartitionProvider:
connection: PartitionConnection
calls: int = 0
def __call__(self) -> PartitionConnection:
self.calls += 1
return self.connection
class ConcurrentPartitionProvider:
def __init__(self, database: PartitionDatabase) -> None:
self.database = database
self.advisory_lock = threading.Lock()
self.first_acquired = threading.Event()
self.first_continue = threading.Event()
self.second_attempted = threading.Event()
self.connections: list[PartitionConnection] = []
self._provider_lock = threading.Lock()
def __call__(self) -> PartitionConnection:
with self._provider_lock:
caller_index = len(self.connections)
connection = PartitionConnection(self.database)
connection.advisory_lock = self.advisory_lock
if caller_index == 0:
connection.advisory_acquired = self.first_acquired
connection.advisory_continue = self.first_continue
elif caller_index == 1:
connection.advisory_attempted = self.second_attempted
self.connections.append(connection)
return connection
def _dependencies() -> tuple[
PartitionDatabase,
PartitionConnection,
PartitionProvider,
]:
database = PartitionDatabase()
connection = PartitionConnection(database)
provider = PartitionProvider(connection)
return database, connection, provider
def _register(
database: PartitionDatabase,
data_type: MarketDataPartitionType,
month: datetime,
) -> str:
partition = build_monthly_partition(data_type=data_type, month=month)
database.registry[(data_type.value, partition.range_start)] = (
partition.partition_name,
partition.range_start,
partition.range_end,
_partition_bound(partition.range_start, partition.range_end),
)
database.relations[partition.partition_name] = (
True,
True,
_partition_bound(partition.range_start, partition.range_end),
)
return partition.partition_name
def test_monthly_descriptor_uses_utc_and_handles_year_rollover() -> None:
local_time = datetime(
2027,
1,
1,
1,
tzinfo=timezone(timedelta(hours=3)),
)
partition = build_monthly_partition(
data_type=MarketDataPartitionType.TRADES,
month=local_time,
)
assert partition.partition_name == "trades_2026_12"
assert partition.range_start == datetime(2026, 12, 1, tzinfo=UTC)
assert partition.range_end == datetime(2027, 1, 1, tzinfo=UTC)
def test_monthly_descriptor_rejects_naive_datetime() -> None:
with pytest.raises(MarketDataStorageValidationError, match="month"):
build_monthly_partition(
data_type=MarketDataPartitionType.QUOTES,
month=JULY.replace(tzinfo=None),
)
def test_manager_creates_partition_moves_default_rows_and_registers_it() -> None:
database, connection, provider = _dependencies()
connection.moved_row_count = 3
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
result = manager.ensure_month_partition(
data_type=MarketDataPartitionType.QUOTES,
month=JULY + timedelta(days=20),
)
assert result.created is True
assert result.moved_row_count == 3
assert result.partition.partition_name == "quotes_2026_07"
assert database.relations["quotes_2026_07"] == (
True,
True,
_partition_bound(JULY, AUGUST),
)
assert database.registry[("quotes", JULY)] == (
"quotes_2026_07",
JULY,
AUGUST,
_partition_bound(JULY, AUGUST),
)
statements = [statement for statement, _ in connection.calls]
assert statements[0] == "SELECT pg_advisory_xact_lock(%s)"
assert connection.calls[0][1] == (
MARKET_DATA_PARTITION_ADVISORY_LOCK_ID,
)
assert statements.index(
'LOCK TABLE "market_data"."quotes_default" IN ACCESS EXCLUSIVE MODE'
) < next(
index
for index, statement in enumerate(statements)
if statement.startswith("WITH moved_rows AS")
)
assert any(
'CREATE TABLE "market_data"."quotes_2026_07"' in statement
for statement in statements
)
ddl_calls = tuple(
(statement, parameters)
for statement, parameters in connection.calls
if statement.startswith("ALTER TABLE")
)
assert ddl_calls
assert all(parameters is None for _, parameters in ddl_calls)
assert all("::timestamptz" in statement for statement, _ in ddl_calls)
def test_manager_is_idempotent_after_registered_partition_exists() -> None:
_, connection, provider = _dependencies()
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
first = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
call_count = len(connection.calls)
second = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
assert first.created is True
assert second.created is False
assert second.partition == first.partition
assert not any(
statement.startswith("CREATE TABLE")
for statement, _ in connection.calls[call_count:]
)
def test_two_manager_callers_participate_in_single_serialized_creation() -> None:
database = PartitionDatabase()
provider = ConcurrentPartitionProvider(database)
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
results: dict[str, MarketDataPartitionResult] = {}
errors: list[BaseException] = []
def run(caller: str) -> None:
try:
results[caller] = manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
except BaseException as error:
errors.append(error)
first = threading.Thread(target=run, args=("first",))
second = threading.Thread(target=run, args=("second",))
first.start()
try:
assert provider.first_acquired.wait(timeout=1)
second.start()
assert provider.second_attempted.wait(timeout=1)
assert second.is_alive()
finally:
provider.first_continue.set()
first.join(timeout=2)
second.join(timeout=2)
assert first.is_alive() is False
assert second.is_alive() is False
assert errors == []
assert {result.created for result in results.values()} == {
True,
False,
}
assert sum(
statement.startswith("CREATE TABLE")
for connection in provider.connections
for statement, _ in connection.calls
) == 1
assert ("trades", JULY) in database.registry
def test_unregistered_relation_is_not_silently_adopted() -> None:
database, _, provider = _dependencies()
database.relations["trades_2026_07"] = (True, False, None)
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(
MarketDataStorageConfigurationError,
match="Unregistered",
):
manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
def test_registered_but_detached_partition_is_rejected() -> None:
database, _, provider = _dependencies()
name = _register(database, MarketDataPartitionType.TRADES, JULY)
database.relations[name] = (True, False, None)
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(
MarketDataStorageConfigurationError,
match="missing or detached",
):
manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
def test_retention_rejects_actual_partition_bound_mismatch_before_drop() -> None:
database, connection, provider = _dependencies()
name = _register(database, MarketDataPartitionType.TRADES, JUNE)
database.relations[name] = (
True,
True,
_partition_bound(JUNE, AUGUST),
)
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
with pytest.raises(
MarketDataStorageConfigurationError,
match="bound differs",
):
service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=31),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert ("trades", JUNE) in database.registry
assert name in database.relations
assert not any(
statement.startswith("DROP TABLE")
for statement, _ in connection.calls
)
def test_partition_setup_error_rolls_back_created_relation_and_registry() -> None:
database, connection, provider = _dependencies()
connection.fail_on = "ATTACH PARTITION"
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(MarketDataStorageOperationError) as error_info:
manager.ensure_month_partition(
data_type=MarketDataPartitionType.CANDLE_REVISIONS,
month=JULY,
)
assert isinstance(error_info.value.__cause__, RuntimeError)
assert database.registry == {}
assert database.relations == {}
assert connection.exit_exception_types[-1] is RuntimeError
def test_partition_setup_interrupt_is_not_wrapped_and_rolls_back() -> None:
database, connection, provider = _dependencies()
connection.interrupt_on = "ATTACH PARTITION"
manager = PostgresMarketDataPartitionManager(
connection_provider=provider
)
with pytest.raises(KeyboardInterrupt):
manager.ensure_month_partition(
data_type=MarketDataPartitionType.TRADES,
month=JULY,
)
assert database.registry == {}
assert database.relations == {}
assert connection.exit_exception_types[-1] is KeyboardInterrupt
def test_retention_is_disabled_without_borrowing_connection() -> None:
_, _, provider = _dependencies()
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
result = service.apply(
policy=MarketDataRetentionPolicy(),
now=JULY.replace(tzinfo=None),
)
assert result.entries == ()
assert result.total_removed_count == 0
assert provider.calls == 0
@pytest.mark.parametrize("days", (0, -1))
def test_retention_rejects_non_positive_window(days: int) -> None:
with pytest.raises(MarketDataStorageConfigurationError):
MarketDataRetentionPolicy(enabled=True, trade_days=days)
def test_enabled_retention_requires_at_least_one_window() -> None:
with pytest.raises(MarketDataStorageConfigurationError, match="window"):
MarketDataRetentionPolicy(enabled=True)
def test_retention_drops_complete_month_and_deletes_partial_history() -> None:
database, connection, provider = _dependencies()
june_name = _register(database, MarketDataPartitionType.TRADES, JUNE)
july_name = _register(database, MarketDataPartitionType.TRADES, JULY)
connection.partition_row_counts[june_name] = 10
connection.partition_row_counts[july_name] = 20
connection.partial_delete_counts["trades"] = 4
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
result = service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=31),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert len(result.entries) == 1
entry = result.entries[0]
assert entry.data_type is MarketDataPartitionType.TRADES
assert entry.cutoff == datetime(2026, 7, 15, tzinfo=UTC)
assert entry.dropped_partitions == (june_name,)
assert entry.dropped_row_count == 10
assert entry.deleted_row_count == 4
assert entry.total_removed_count == 14
assert june_name not in database.relations
assert july_name in database.relations
assert ("trades", JUNE) not in database.registry
assert ("trades", JULY) in database.registry
assert any(
statement == (
'DELETE FROM "market_data"."trades" '
'WHERE "executed_at" < %s'
)
for statement, _ in connection.calls
)
def test_unconfigured_data_type_is_unlimited_and_not_touched() -> None:
database, connection, provider = _dependencies()
_register(database, MarketDataPartitionType.QUOTES, JUNE)
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
result = service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=30),
now=AUGUST,
)
assert tuple(entry.data_type for entry in result.entries) == (
MarketDataPartitionType.TRADES,
)
assert not any(
'"quotes"' in statement or parameters == ("quotes", JULY)
for statement, parameters in connection.calls
)
def test_parallel_writer_waits_until_retention_transaction_finishes() -> None:
_, connection, provider = _dependencies()
write_lock = threading.Lock()
connection.write_lock = write_lock
connection.write_lock_acquired = threading.Event()
connection.write_lock_continue = threading.Event()
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
retention_errors: list[BaseException] = []
writer_attempted = threading.Event()
writer_finished = threading.Event()
def retain() -> None:
try:
service.apply(
policy=MarketDataRetentionPolicy(
enabled=True,
trade_days=31,
),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
except BaseException as error:
retention_errors.append(error)
def write() -> None:
writer_attempted.set()
with write_lock:
writer_finished.set()
retention_thread = threading.Thread(target=retain)
writer_thread = threading.Thread(target=write)
retention_thread.start()
try:
assert connection.write_lock_acquired.wait(timeout=1)
writer_thread.start()
assert writer_attempted.wait(timeout=1)
assert writer_finished.wait(timeout=0.05) is False
finally:
assert connection.write_lock_continue is not None
connection.write_lock_continue.set()
retention_thread.join(timeout=2)
writer_thread.join(timeout=2)
assert retention_thread.is_alive() is False
assert writer_thread.is_alive() is False
assert retention_errors == []
assert writer_finished.is_set()
statements = [statement for statement, _ in connection.calls]
lock_index = statements.index(
'LOCK TABLE "market_data"."trades" '
"IN SHARE ROW EXCLUSIVE MODE"
)
registry_index = next(
index
for index, statement in enumerate(statements)
if statement.startswith("SELECT partition_name")
)
assert lock_index < registry_index
def test_retention_error_rolls_back_all_configured_data_types() -> None:
database, connection, provider = _dependencies()
trade_name = _register(
database,
MarketDataPartitionType.TRADES,
JUNE,
)
quote_name = _register(
database,
MarketDataPartitionType.QUOTES,
JUNE,
)
connection.partition_row_counts[trade_name] = 2
connection.partition_row_counts[quote_name] = 3
connection.fail_on = 'DROP TABLE "market_data"."quotes_2026_06"'
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
with pytest.raises(MarketDataStorageOperationError):
service.apply(
policy=MarketDataRetentionPolicy(
enabled=True,
trade_days=31,
quote_days=31,
),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert ("trades", JUNE) in database.registry
assert ("quotes", JUNE) in database.registry
assert trade_name in database.relations
assert quote_name in database.relations
assert connection.exit_exception_types[-1] is RuntimeError
def test_retention_interrupt_rolls_back_and_is_not_wrapped() -> None:
database, connection, provider = _dependencies()
trade_name = _register(
database,
MarketDataPartitionType.TRADES,
JUNE,
)
connection.interrupt_on = "DROP TABLE"
service = PostgresMarketDataRetentionService(
connection_provider=provider
)
with pytest.raises(KeyboardInterrupt):
service.apply(
policy=MarketDataRetentionPolicy(enabled=True, trade_days=31),
now=datetime(2026, 8, 15, tzinfo=UTC),
)
assert ("trades", JUNE) in database.registry
assert trade_name in database.relations
assert connection.exit_exception_types[-1] is KeyboardInterrupt

View File

@@ -0,0 +1,641 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from typing import Any
import pytest
from src.market_data.acquisition.models.candle import Candle
from src.market_data.acquisition.models.quote import Quote
from src.market_data.storage import (
CandleStorageProtocol,
MarketDataStorageConflictError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
MarketDataWriteStatus,
PostgresCandleRepository,
PostgresQuoteRepository,
QuoteStorageProtocol,
)
VENUE = "dzengi"
SYMBOL = "BTC/USD_LEVERAGE"
OPEN_TIME = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
RECEIVED_AT = OPEN_TIME + timedelta(seconds=1)
OBSERVED_AT = OPEN_TIME + timedelta(minutes=1)
QuoteKey = tuple[str, str, datetime]
CandleKey = tuple[str, str, str, datetime, datetime]
StorageRow = dict[str, Any]
@dataclass
class TransactionalMarketDataDatabase:
quote_rows: dict[QuoteKey, StorageRow] = field(default_factory=dict)
candle_rows: dict[CandleKey, StorageRow] = field(default_factory=dict)
class TransactionalMarketDataCursor:
def __init__(
self,
connection: TransactionalMarketDataConnection,
) -> None:
self._connection = connection
self._fetchone_result: object = None
def __enter__(self) -> TransactionalMarketDataCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str,
parameters: tuple[Any, ...],
) -> None:
normalized = " ".join(statement.split())
self._connection.calls.append((normalized, parameters))
if (
self._connection.interrupt_on is not None
and self._connection.interrupt_on in normalized
):
raise KeyboardInterrupt
if (
self._connection.fail_on is not None
and self._connection.fail_on in normalized
):
raise RuntimeError("database failed")
if normalized.startswith("INSERT INTO market_data.quotes"):
self._insert_quote(parameters)
return
if normalized.startswith("SELECT exchange_timestamp"):
self._select_quote(parameters)
return
if normalized.startswith("UPDATE market_data.quotes"):
self._update_quote(parameters)
return
if normalized.startswith(
"INSERT INTO market_data.candle_revisions"
):
self._insert_candle(parameters)
return
if normalized.startswith("SELECT open_price"):
self._select_candle(parameters)
return
if normalized.startswith(
"UPDATE market_data.candle_revisions"
):
self._update_candle(parameters)
return
raise AssertionError(f"Unexpected SQL: {normalized}")
def fetchone(self) -> object:
return self._fetchone_result
def _insert_quote(self, parameters: tuple[Any, ...]) -> None:
(
venue,
symbol,
received_at,
exchange_timestamp,
last_price,
bid_price,
ask_price,
source,
observation_sources,
canonical_schema_version,
) = parameters
key = (venue, symbol, received_at)
if key in self._connection.working_quote_rows:
self._fetchone_result = None
return
self._connection.working_quote_rows[key] = {
"exchange_timestamp": exchange_timestamp,
"last_price": last_price,
"bid_price": bid_price,
"ask_price": ask_price,
"source": source,
"observation_sources": list(observation_sources),
"canonical_schema_version": canonical_schema_version,
}
self._fetchone_result = (1,)
def _select_quote(self, parameters: tuple[Any, ...]) -> None:
row = self._connection.working_quote_rows.get(parameters)
if row is None:
self._fetchone_result = None
return
self._fetchone_result = (
row["exchange_timestamp"],
row["last_price"],
row["bid_price"],
row["ask_price"],
list(row["observation_sources"]),
row["canonical_schema_version"],
)
def _update_quote(self, parameters: tuple[Any, ...]) -> None:
sources, *identity = parameters
row = self._connection.working_quote_rows[tuple(identity)]
row["observation_sources"] = list(sources)
self._fetchone_result = None
def _insert_candle(self, parameters: tuple[Any, ...]) -> None:
(
venue,
symbol,
interval,
open_time,
observed_at,
open_price,
high_price,
low_price,
close_price,
volume,
is_final,
source,
observation_sources,
canonical_schema_version,
) = parameters
key = (venue, symbol, interval, open_time, observed_at)
if key in self._connection.working_candle_rows:
self._fetchone_result = None
return
self._connection.working_candle_rows[key] = {
"open_price": open_price,
"high_price": high_price,
"low_price": low_price,
"close_price": close_price,
"volume": volume,
"is_final": is_final,
"source": source,
"observation_sources": list(observation_sources),
"canonical_schema_version": canonical_schema_version,
}
self._fetchone_result = (1,)
def _select_candle(self, parameters: tuple[Any, ...]) -> None:
row = self._connection.working_candle_rows.get(parameters)
if row is None:
self._fetchone_result = None
return
self._fetchone_result = (
row["open_price"],
row["high_price"],
row["low_price"],
row["close_price"],
row["volume"],
row["is_final"],
list(row["observation_sources"]),
row["canonical_schema_version"],
)
def _update_candle(self, parameters: tuple[Any, ...]) -> None:
sources, *identity = parameters
row = self._connection.working_candle_rows[tuple(identity)]
row["observation_sources"] = list(sources)
self._fetchone_result = None
class TransactionalMarketDataConnection:
def __init__(self, database: TransactionalMarketDataDatabase) -> None:
self._database = database
self.calls: list[tuple[str, tuple[Any, ...]]] = []
self.exit_exception_types: list[type[BaseException] | None] = []
self.fail_on: str | None = None
self.interrupt_on: str | None = None
self.working_quote_rows: dict[QuoteKey, StorageRow] = {}
self.working_candle_rows: dict[CandleKey, StorageRow] = {}
def __enter__(self) -> TransactionalMarketDataConnection:
self.working_quote_rows = deepcopy(self._database.quote_rows)
self.working_candle_rows = deepcopy(self._database.candle_rows)
return self
def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: object,
) -> None:
self.exit_exception_types.append(exception_type)
if exception_type is None:
self._database.quote_rows = self.working_quote_rows
self._database.candle_rows = self.working_candle_rows
return None
def cursor(self) -> TransactionalMarketDataCursor:
return TransactionalMarketDataCursor(self)
@dataclass
class RecordingConnectionProvider:
connection: TransactionalMarketDataConnection
calls: int = 0
def __call__(self) -> TransactionalMarketDataConnection:
self.calls += 1
return self.connection
def _dependencies() -> tuple[
TransactionalMarketDataDatabase,
TransactionalMarketDataConnection,
RecordingConnectionProvider,
]:
database = TransactionalMarketDataDatabase()
connection = TransactionalMarketDataConnection(database)
provider = RecordingConnectionProvider(connection)
return database, connection, provider
def _quote(
*,
received_at: datetime = RECEIVED_AT,
exchange_timestamp: datetime | None = OPEN_TIME,
last_price: Decimal = Decimal("100"),
bid_price: Decimal = Decimal("99"),
ask_price: Decimal = Decimal("101"),
source: str = "dzengi",
) -> Quote:
return Quote(
symbol=SYMBOL,
last_price=last_price,
bid_price=bid_price,
ask_price=ask_price,
exchange_timestamp=exchange_timestamp,
received_at=received_at,
source=source,
)
def _candle(
*,
open_time: datetime = OPEN_TIME,
open_price: Decimal = Decimal("100"),
high_price: Decimal = Decimal("110"),
low_price: Decimal = Decimal("90"),
close_price: Decimal = Decimal("105"),
volume: Decimal = Decimal("10"),
source: str = "rest_klines:bid",
) -> Candle:
return Candle(
symbol=SYMBOL,
interval="1m",
open_time=open_time,
open_price=open_price,
high_price=high_price,
low_price=low_price,
close_price=close_price,
volume=volume,
source=source,
)
def test_quote_repository_matches_protocol() -> None:
_, _, provider = _dependencies()
assert isinstance(
PostgresQuoteRepository(connection_provider=provider),
QuoteStorageProtocol,
)
def test_quote_insert_normalizes_identity_and_preserves_payload() -> None:
database, _, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
quote = replace(_quote(), symbol=" btc/usd_leverage ")
result = repository.store_quote(venue=" dzengi ", quote=quote)
assert result.status is MarketDataWriteStatus.INSERTED
assert tuple(database.quote_rows) == ((VENUE, SYMBOL, RECEIVED_AT),)
row = next(iter(database.quote_rows.values()))
assert row["last_price"] == Decimal("100")
assert row["source"] == "dzengi"
assert row["observation_sources"] == ["dzengi"]
def test_quote_exact_duplicate_is_not_updated() -> None:
database, connection, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
quote = _quote()
repository.store_quote(venue=VENUE, quote=quote)
call_count = len(connection.calls)
result = repository.store_quote(venue=VENUE, quote=quote)
assert result.status is MarketDataWriteStatus.DUPLICATE
assert len(database.quote_rows) == 1
assert not any(
statement.startswith("UPDATE market_data.quotes")
for statement, _ in connection.calls[call_count:]
)
def test_quote_new_source_updates_only_provenance() -> None:
database, _, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
quote = _quote(source="dzengi")
repository.store_quote(venue=VENUE, quote=quote)
result = repository.store_quote(
venue=VENUE,
quote=replace(quote, source="dzengi_websocket_quote"),
)
row = next(iter(database.quote_rows.values()))
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["source"] == "dzengi"
assert row["observation_sources"] == [
"dzengi",
"dzengi_websocket_quote",
]
@pytest.mark.parametrize(
"conflicting_quote",
(
_quote(last_price=Decimal("102")),
_quote(bid_price=Decimal("98")),
_quote(ask_price=Decimal("102")),
_quote(exchange_timestamp=OPEN_TIME + timedelta(seconds=1)),
),
)
def test_quote_timestamp_collision_with_different_payload_is_conflict(
conflicting_quote: Quote,
) -> None:
database, connection, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
repository.store_quote(venue=VENUE, quote=_quote())
with pytest.raises(MarketDataStorageConflictError):
repository.store_quote(venue=VENUE, quote=conflicting_quote)
assert len(database.quote_rows) == 1
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
@pytest.mark.parametrize(
"invalid_quote",
(
_quote(received_at=RECEIVED_AT.replace(tzinfo=None)),
_quote(exchange_timestamp=OPEN_TIME.replace(tzinfo=None)),
_quote(bid_price=Decimal("102"), ask_price=Decimal("101")),
_quote(last_price=Decimal("NaN")),
),
)
def test_quote_validation_happens_before_connection(
invalid_quote: Quote,
) -> None:
_, _, provider = _dependencies()
repository = PostgresQuoteRepository(connection_provider=provider)
with pytest.raises(MarketDataStorageValidationError):
repository.store_quote(venue=VENUE, quote=invalid_quote)
assert provider.calls == 0
def test_quote_database_error_is_wrapped_and_rolled_back() -> None:
database, connection, provider = _dependencies()
connection.fail_on = "INSERT INTO market_data.quotes"
repository = PostgresQuoteRepository(connection_provider=provider)
with pytest.raises(MarketDataStorageOperationError) as error_info:
repository.store_quote(venue=VENUE, quote=_quote())
assert isinstance(error_info.value.__cause__, RuntimeError)
assert database.quote_rows == {}
assert connection.exit_exception_types[-1] is RuntimeError
def test_candle_repository_matches_protocol() -> None:
_, _, provider = _dependencies()
assert isinstance(
PostgresCandleRepository(connection_provider=provider),
CandleStorageProtocol,
)
def test_candle_revision_insert_normalizes_identity_and_payload() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = replace(
_candle(),
symbol=" btc/usd_leverage ",
interval=" 1m ",
)
result = repository.store_candle_revision(
venue=" dzengi ",
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
assert result.status is MarketDataWriteStatus.INSERTED
assert tuple(database.candle_rows) == (
(VENUE, SYMBOL, "1m", OPEN_TIME, OBSERVED_AT),
)
row = next(iter(database.candle_rows.values()))
assert row["is_final"] is False
assert row["observation_sources"] == ["rest_klines:bid"]
def test_candle_interval_case_is_preserved() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
repository.store_candle_revision(
venue=VENUE,
candle=replace(_candle(), interval="1M"),
observed_at=OBSERVED_AT,
is_final=True,
)
assert next(iter(database.candle_rows))[2] == "1M"
def test_candle_exact_duplicate_is_not_updated() -> None:
database, connection, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = _candle()
repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
call_count = len(connection.calls)
result = repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
assert result.status is MarketDataWriteStatus.DUPLICATE
assert len(database.candle_rows) == 1
assert not any(
statement.startswith("UPDATE market_data.candle_revisions")
for statement, _ in connection.calls[call_count:]
)
def test_candle_new_source_updates_only_provenance() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = _candle(source="rest_klines:bid")
repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=True,
)
result = repository.store_candle_revision(
venue=VENUE,
candle=replace(candle, source="secondary"),
observed_at=OBSERVED_AT,
is_final=True,
)
row = next(iter(database.candle_rows.values()))
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["source"] == "rest_klines:bid"
assert row["observation_sources"] == [
"rest_klines:bid",
"secondary",
]
def test_candle_final_state_at_later_observation_is_new_revision() -> None:
database, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
candle = _candle()
repository.store_candle_revision(
venue=VENUE,
candle=candle,
observed_at=OBSERVED_AT,
is_final=False,
)
result = repository.store_candle_revision(
venue=VENUE,
candle=replace(candle, close_price=Decimal("106")),
observed_at=OBSERVED_AT + timedelta(seconds=1),
is_final=True,
)
assert result.status is MarketDataWriteStatus.INSERTED
assert len(database.candle_rows) == 2
@pytest.mark.parametrize(
("conflicting_candle", "is_final"),
(
(replace(_candle(), close_price=Decimal("106")), False),
(replace(_candle(), volume=Decimal("11")), False),
(_candle(), True),
),
)
def test_candle_same_revision_identity_with_different_payload_is_conflict(
conflicting_candle: Candle,
is_final: bool,
) -> None:
database, connection, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
repository.store_candle_revision(
venue=VENUE,
candle=_candle(),
observed_at=OBSERVED_AT,
is_final=False,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_candle_revision(
venue=VENUE,
candle=conflicting_candle,
observed_at=OBSERVED_AT,
is_final=is_final,
)
assert len(database.candle_rows) == 1
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
@pytest.mark.parametrize(
("invalid_candle", "observed_at", "is_final"),
(
(_candle(), OPEN_TIME - timedelta(seconds=1), False),
(_candle(open_time=OPEN_TIME.replace(tzinfo=None)), OBSERVED_AT, False),
(_candle(volume=Decimal("-1")), OBSERVED_AT, False),
(_candle(high_price=Decimal("99")), OBSERVED_AT, False),
(_candle(), OBSERVED_AT, 1),
),
)
def test_candle_validation_happens_before_connection(
invalid_candle: Candle,
observed_at: datetime,
is_final: object,
) -> None:
_, _, provider = _dependencies()
repository = PostgresCandleRepository(connection_provider=provider)
with pytest.raises(MarketDataStorageValidationError):
repository.store_candle_revision(
venue=VENUE,
candle=invalid_candle,
observed_at=observed_at,
is_final=is_final, # type: ignore[arg-type]
)
assert provider.calls == 0
def test_candle_keyboard_interrupt_is_not_wrapped_and_rolls_back() -> None:
database, connection, provider = _dependencies()
connection.interrupt_on = "INSERT INTO market_data.candle_revisions"
repository = PostgresCandleRepository(connection_provider=provider)
with pytest.raises(KeyboardInterrupt):
repository.store_candle_revision(
venue=VENUE,
candle=_candle(),
observed_at=OBSERVED_AT,
is_final=False,
)
assert database.candle_rows == {}
assert connection.exit_exception_types[-1] is KeyboardInterrupt

View File

@@ -0,0 +1,638 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass, field, replace
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from typing import Any
import pytest
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.acquisition.trade_id_sequence import (
SIGNED_TRADE_ID_MAX,
SIGNED_TRADE_ID_MIN,
)
from src.market_data.storage import (
MarketDataStorageConflictError,
MarketDataStorageOperationError,
MarketDataStorageValidationError,
MarketDataWriteStatus,
PostgresTradeRepository,
TradeStorageProtocol,
)
VENUE = "dzengi"
SYMBOL = "BTC/USD_LEVERAGE"
EXECUTED_AT = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
OBSERVED_AT = EXECUTED_AT + timedelta(seconds=1)
TradeKey = tuple[str, str, int, datetime]
TradeRow = dict[str, Any]
@dataclass
class TransactionalTradeDatabase:
rows: dict[TradeKey, TradeRow] = field(default_factory=dict)
class TransactionalCursor:
def __init__(
self,
connection: TransactionalConnection,
) -> None:
self._connection = connection
self._fetchone_result: object = None
def __enter__(self) -> TransactionalCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str,
parameters: tuple[Any, ...],
) -> None:
normalized = " ".join(statement.split())
self._connection.calls.append((normalized, parameters))
if (
self._connection.interrupt_on is not None
and self._connection.interrupt_on in normalized
):
raise KeyboardInterrupt
if (
self._connection.fail_on is not None
and self._connection.fail_on in normalized
):
raise RuntimeError("database failed")
if normalized.startswith("INSERT INTO market_data.trades"):
self._insert(parameters)
return
if normalized.startswith("SELECT price, quantity"):
self._select(parameters)
return
if normalized.startswith("UPDATE market_data.trades"):
self._update(parameters)
return
raise AssertionError(f"Unexpected SQL: {normalized}")
def fetchone(self) -> object:
return self._fetchone_result
def _insert(self, parameters: tuple[Any, ...]) -> None:
(
venue,
symbol,
trade_id,
executed_at,
price,
quantity,
aggressor_side,
source,
first_observed_at,
last_observed_at,
observation_sources,
canonical_schema_version,
) = parameters
key = (venue, symbol, trade_id, executed_at)
working_rows = self._connection.working_rows
if key in working_rows:
self._fetchone_result = None
return
working_rows[key] = {
"price": price,
"quantity": quantity,
"aggressor_side": aggressor_side,
"source": source,
"first_observed_at": first_observed_at,
"last_observed_at": last_observed_at,
"observation_sources": list(observation_sources),
"canonical_schema_version": canonical_schema_version,
}
self._fetchone_result = (1,)
def _select(self, parameters: tuple[Any, ...]) -> None:
key = parameters
row = self._connection.working_rows.get(key)
if row is None:
self._fetchone_result = None
return
self._fetchone_result = (
row["price"],
row["quantity"],
row["aggressor_side"],
row["first_observed_at"],
row["last_observed_at"],
list(row["observation_sources"]),
row["canonical_schema_version"],
)
def _update(self, parameters: tuple[Any, ...]) -> None:
first_observed_at, last_observed_at, sources, *identity = parameters
row = self._connection.working_rows[tuple(identity)]
row["first_observed_at"] = first_observed_at
row["last_observed_at"] = last_observed_at
row["observation_sources"] = list(sources)
self._fetchone_result = None
class TransactionalConnection:
def __init__(self, database: TransactionalTradeDatabase) -> None:
self._database = database
self.calls: list[tuple[str, tuple[Any, ...]]] = []
self.exit_exception_types: list[type[BaseException] | None] = []
self.fail_on: str | None = None
self.interrupt_on: str | None = None
self.working_rows: dict[TradeKey, TradeRow] = {}
def __enter__(self) -> TransactionalConnection:
self.working_rows = deepcopy(self._database.rows)
return self
def __exit__(
self,
exception_type: type[BaseException] | None,
exception: BaseException | None,
traceback: object,
) -> None:
self.exit_exception_types.append(exception_type)
if exception_type is None:
self._database.rows = self.working_rows
return None
def cursor(self) -> TransactionalCursor:
return TransactionalCursor(self)
@dataclass
class RecordingConnectionProvider:
connection: TransactionalConnection
calls: int = 0
def __call__(self) -> TransactionalConnection:
self.calls += 1
return self.connection
def _trade(
*,
symbol: str = SYMBOL,
trade_id: int = 100,
executed_at: datetime = EXECUTED_AT,
price: Decimal = Decimal("64159.45"),
quantity: Decimal = Decimal("0.125"),
aggressor_side: TradeAggressorSide = TradeAggressorSide.BUY,
source: str = "dzengi_websocket_trade",
) -> Trade:
return Trade(
symbol=symbol,
trade_id=trade_id,
price=price,
quantity=quantity,
executed_at=executed_at,
aggressor_side=aggressor_side,
source=source,
)
def _repository() -> tuple[
PostgresTradeRepository,
TransactionalTradeDatabase,
TransactionalConnection,
RecordingConnectionProvider,
]:
database = TransactionalTradeDatabase()
connection = TransactionalConnection(database)
provider = RecordingConnectionProvider(connection)
repository = PostgresTradeRepository(
connection_provider=provider,
)
return repository, database, connection, provider
def _only_row(database: TransactionalTradeDatabase) -> TradeRow:
assert len(database.rows) == 1
return next(iter(database.rows.values()))
def test_repository_matches_trade_storage_protocol() -> None:
repository, _, _, _ = _repository()
assert isinstance(repository, TradeStorageProtocol)
def test_store_trade_inserts_canonical_payload_and_provenance() -> None:
repository, database, _, provider = _repository()
result = repository.store_trade(
venue=" dzengi ",
trade=_trade(symbol=" btc/usd_leverage "),
observed_at=OBSERVED_AT,
)
assert result.status is MarketDataWriteStatus.INSERTED
assert provider.calls == 1
assert tuple(database.rows) == (
(VENUE, SYMBOL, 100, EXECUTED_AT),
)
assert _only_row(database) == {
"price": Decimal("64159.45"),
"quantity": Decimal("0.125"),
"aggressor_side": "buy",
"source": "dzengi_websocket_trade",
"first_observed_at": OBSERVED_AT,
"last_observed_at": OBSERVED_AT,
"observation_sources": ["dzengi_websocket_trade"],
"canonical_schema_version": 1,
}
def test_exact_duplicate_does_not_update_row() -> None:
repository, database, connection, _ = _repository()
trade = _trade()
repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
)
calls_before_duplicate = len(connection.calls)
result = repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
)
assert result.status is MarketDataWriteStatus.DUPLICATE
assert _only_row(database)["observation_sources"] == [
"dzengi_websocket_trade"
]
assert not any(
statement.startswith("UPDATE market_data.trades")
for statement, _ in connection.calls[calls_before_duplicate:]
)
def test_second_transport_source_updates_provenance_not_market_fact() -> None:
repository, database, _, _ = _repository()
websocket_trade = _trade(source="dzengi_websocket_trade")
rest_trade = replace(websocket_trade, source="dzengi")
repository.store_trade(
venue=VENUE,
trade=websocket_trade,
observed_at=OBSERVED_AT,
)
result = repository.store_trade(
venue=VENUE,
trade=rest_trade,
observed_at=OBSERVED_AT + timedelta(seconds=5),
)
row = _only_row(database)
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["source"] == "dzengi_websocket_trade"
assert row["observation_sources"] == [
"dzengi_websocket_trade",
"dzengi",
]
assert row["first_observed_at"] == OBSERVED_AT
assert row["last_observed_at"] == OBSERVED_AT + timedelta(seconds=5)
def test_second_transport_source_updates_provenance_at_same_time() -> None:
repository, database, _, _ = _repository()
websocket_trade = _trade(source="dzengi_websocket_trade")
repository.store_trade(
venue=VENUE,
trade=websocket_trade,
observed_at=OBSERVED_AT,
)
result = repository.store_trade(
venue=VENUE,
trade=replace(websocket_trade, source="dzengi"),
observed_at=OBSERVED_AT,
)
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert _only_row(database)["observation_sources"] == [
"dzengi_websocket_trade",
"dzengi",
]
def test_earlier_observation_moves_only_first_observed_at() -> None:
repository, database, _, _ = _repository()
trade = _trade()
repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=OBSERVED_AT,
)
earlier = OBSERVED_AT - timedelta(seconds=5)
result = repository.store_trade(
venue=VENUE,
trade=trade,
observed_at=earlier,
)
row = _only_row(database)
assert result.status is MarketDataWriteStatus.PROVENANCE_UPDATED
assert row["first_observed_at"] == earlier
assert row["last_observed_at"] == OBSERVED_AT
@pytest.mark.parametrize(
"conflicting_trade",
(
_trade(price=Decimal("64160")),
_trade(quantity=Decimal("1")),
_trade(aggressor_side=TradeAggressorSide.SELL),
),
)
def test_same_identity_with_different_market_fact_is_conflict(
conflicting_trade: Trade,
) -> None:
repository, database, connection, _ = _repository()
original = _trade()
repository.store_trade(
venue=VENUE,
trade=original,
observed_at=OBSERVED_AT,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_trade(
venue=VENUE,
trade=conflicting_trade,
observed_at=OBSERVED_AT,
)
assert _only_row(database)["price"] == original.price
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
def test_empty_batch_returns_zero_without_borrowing_connection() -> None:
repository, _, _, provider = _repository()
result = repository.store_trades(
venue=VENUE,
trades=(),
observed_at=OBSERVED_AT,
)
assert result.total_count == 0
assert provider.calls == 0
def test_batch_returns_insert_duplicate_and_provenance_counts() -> None:
repository, _, _, _ = _repository()
duplicate = _trade(symbol="B", trade_id=2)
provenance = _trade(symbol="C", trade_id=3)
repository.store_trade(
venue=VENUE,
trade=duplicate,
observed_at=OBSERVED_AT + timedelta(seconds=1),
)
repository.store_trade(
venue=VENUE,
trade=provenance,
observed_at=OBSERVED_AT,
)
result = repository.store_trades(
venue=VENUE,
trades=(
replace(provenance, source="dzengi"),
duplicate,
_trade(symbol="A", trade_id=1),
),
observed_at=OBSERVED_AT + timedelta(seconds=1),
)
assert result.inserted_count == 1
assert result.duplicate_count == 1
assert result.provenance_updated_count == 1
assert result.total_count == 3
def test_batch_uses_stable_identity_order() -> None:
repository, _, connection, _ = _repository()
repository.store_trades(
venue=VENUE,
trades=(
_trade(symbol="C", trade_id=3),
_trade(symbol="A", trade_id=1),
_trade(symbol="B", trade_id=2),
),
observed_at=OBSERVED_AT,
)
inserted_symbols = tuple(
parameters[1]
for statement, parameters in connection.calls
if statement.startswith("INSERT INTO market_data.trades")
)
assert inserted_symbols == ("A", "B", "C")
def test_batch_conflict_rolls_back_preceding_insert() -> None:
repository, database, connection, _ = _repository()
existing = _trade(symbol="B", trade_id=2)
repository.store_trade(
venue=VENUE,
trade=existing,
observed_at=OBSERVED_AT,
)
with pytest.raises(MarketDataStorageConflictError):
repository.store_trades(
venue=VENUE,
trades=(
_trade(symbol="A", trade_id=1),
replace(existing, price=Decimal("999")),
),
observed_at=OBSERVED_AT,
)
assert tuple(database.rows) == (
(VENUE, "B", 2, EXECUTED_AT),
)
assert connection.exit_exception_types[-1] is (
MarketDataStorageConflictError
)
def test_batch_accepts_signed_rollover_boundary_as_distinct_trades() -> None:
repository, database, _, _ = _repository()
result = repository.store_trades(
venue=VENUE,
trades=(
_trade(
trade_id=SIGNED_TRADE_ID_MAX,
executed_at=EXECUTED_AT,
),
_trade(
trade_id=SIGNED_TRADE_ID_MIN,
executed_at=EXECUTED_AT + timedelta(milliseconds=1),
),
),
observed_at=OBSERVED_AT,
)
assert result.inserted_count == 2
assert len(database.rows) == 2
def test_database_error_is_wrapped_and_transaction_is_rolled_back() -> None:
repository, database, connection, _ = _repository()
connection.fail_on = "INSERT INTO market_data.trades"
with pytest.raises(MarketDataStorageOperationError) as error_info:
repository.store_trade(
venue=VENUE,
trade=_trade(),
observed_at=OBSERVED_AT,
)
assert isinstance(error_info.value.__cause__, RuntimeError)
assert database.rows == {}
assert connection.exit_exception_types[-1] is RuntimeError
def test_keyboard_interrupt_is_not_wrapped_and_rolls_back() -> None:
repository, database, connection, _ = _repository()
connection.interrupt_on = "INSERT INTO market_data.trades"
with pytest.raises(KeyboardInterrupt):
repository.store_trade(
venue=VENUE,
trade=_trade(),
observed_at=OBSERVED_AT,
)
assert database.rows == {}
assert connection.exit_exception_types[-1] is KeyboardInterrupt
@pytest.mark.parametrize("venue", ("", " ", None, 1))
def test_rejects_invalid_venue_without_connection(venue: object) -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError, match="venue"):
repository.store_trade(
venue=venue, # type: ignore[arg-type]
trade=_trade(),
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
def test_rejects_naive_observed_at_without_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(
MarketDataStorageValidationError,
match="observed_at",
):
repository.store_trade(
venue=VENUE,
trade=_trade(),
observed_at=OBSERVED_AT.replace(tzinfo=None),
)
assert provider.calls == 0
def test_rejects_naive_executed_at_without_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(
MarketDataStorageValidationError,
match="executed_at",
):
repository.store_trade(
venue=VENUE,
trade=_trade(executed_at=EXECUTED_AT.replace(tzinfo=None)),
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
@pytest.mark.parametrize(
"invalid_trade",
(
_trade(trade_id=SIGNED_TRADE_ID_MAX + 1),
_trade(price=Decimal("0")),
_trade(quantity=Decimal("NaN")),
replace(_trade(), aggressor_side="buy"), # type: ignore[arg-type]
),
)
def test_rejects_invalid_canonical_trade_before_connection(
invalid_trade: Trade,
) -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError):
repository.store_trade(
venue=VENUE,
trade=invalid_trade,
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
def test_batch_rejects_non_tuple_without_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError, match="tuple"):
repository.store_trades(
venue=VENUE,
trades=[_trade()], # type: ignore[arg-type]
observed_at=OBSERVED_AT,
)
assert provider.calls == 0
def test_batch_validates_every_trade_before_connection() -> None:
repository, _, _, provider = _repository()
with pytest.raises(MarketDataStorageValidationError):
repository.store_trades(
venue=VENUE,
trades=(
_trade(symbol="A"),
_trade(symbol=" "),
),
observed_at=OBSERVED_AT,
)
assert provider.calls == 0

View File

@@ -0,0 +1,170 @@
from __future__ import annotations
from datetime import datetime, timezone
from decimal import Decimal
import pytest
from src.market_data.acquisition.consistency.trade_observation_sink_protocol import (
TradeObservationSinkProtocol,
)
from src.market_data.acquisition.models.trade import (
Trade,
TradeAggressorSide,
)
from src.market_data.storage import (
MarketDataBatchWriteResult,
MarketDataStorageValidationError,
MarketDataWriteResult,
MarketDataWriteStatus,
TradeStorageObservationSink,
)
VENUE = "dzengi"
OBSERVED_AT = datetime(2026, 7, 31, 12, 0, tzinfo=timezone.utc)
def make_trade() -> Trade:
return Trade(
symbol="BTC/USD_LEVERAGE",
trade_id=100,
price=Decimal("64555.55"),
quantity=Decimal("0.002"),
executed_at=OBSERVED_AT,
aggressor_side=TradeAggressorSide.BUY,
source="dzengi_websocket_trade",
)
class RecordingTradeStorage:
def __init__(
self,
*,
result: object | None = None,
error: Exception | None = None,
) -> None:
self.result = result or MarketDataWriteResult(
status=MarketDataWriteStatus.INSERTED,
)
self.error = error
self.calls: list[tuple[str, Trade, datetime]] = []
def store_trade(
self,
*,
venue: str,
trade: Trade,
observed_at: datetime,
) -> MarketDataWriteResult:
self.calls.append((venue, trade, observed_at))
if self.error is not None:
raise self.error
return self.result # type: ignore[return-value]
def store_trades(
self,
*,
venue: str,
trades: tuple[Trade, ...],
observed_at: datetime,
) -> MarketDataBatchWriteResult:
raise AssertionError("Runtime persists observations one by one")
@pytest.mark.parametrize(
"status",
tuple(MarketDataWriteStatus),
)
def test_forwards_observation_and_accepts_all_success_statuses(
status: MarketDataWriteStatus,
) -> None:
storage = RecordingTradeStorage(
result=MarketDataWriteResult(status=status),
)
sink = TradeStorageObservationSink(
trade_storage=storage,
venue=" dzengi ",
clock=lambda: OBSERVED_AT,
)
trade = make_trade()
result = sink.persist(trade)
assert result is None
assert storage.calls == [
(
VENUE,
trade,
OBSERVED_AT,
)
]
def test_implements_acquisition_side_sink_protocol() -> None:
sink = TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(),
venue=VENUE,
)
assert isinstance(sink, TradeObservationSinkProtocol)
assert not hasattr(sink, "__dict__")
def test_storage_error_is_not_wrapped() -> None:
storage_error = RuntimeError("storage failed")
sink = TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(error=storage_error),
venue=VENUE,
clock=lambda: OBSERVED_AT,
)
with pytest.raises(
RuntimeError,
match="storage failed",
) as error_info:
sink.persist(make_trade())
assert error_info.value is storage_error
def test_rejects_invalid_storage_result() -> None:
sink = TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(result=object()),
venue=VENUE,
clock=lambda: OBSERVED_AT,
)
with pytest.raises(
TypeError,
match="MarketDataWriteResult",
):
sink.persist(make_trade())
def test_rejects_invalid_dependencies() -> None:
with pytest.raises(TypeError, match="TradeStorageProtocol"):
TradeStorageObservationSink(
trade_storage=object(), # type: ignore[arg-type]
venue=VENUE,
)
with pytest.raises(TypeError, match="clock must be callable"):
TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(),
venue=VENUE,
clock=object(), # type: ignore[arg-type]
)
def test_rejects_empty_venue_at_construction() -> None:
with pytest.raises(
MarketDataStorageValidationError,
match="venue must not be empty",
):
TradeStorageObservationSink(
trade_storage=RecordingTradeStorage(),
venue=" ",
)

View File

@@ -0,0 +1,295 @@
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any
import pytest
from src.storage.exceptions import StorageMigrationError
from src.storage.migrations import (
STORAGE_MIGRATION_ADVISORY_LOCK_ID,
STORAGE_MIGRATIONS,
StorageMigration,
StorageMigrationRunner,
)
@dataclass
class RecordingCursor:
applied_rows: list[tuple[int, str]] = field(default_factory=list)
fail_on: str | None = None
calls: list[tuple[str, object | None]] = field(default_factory=list)
def __enter__(self) -> RecordingCursor:
return self
def __exit__(self, *args: object) -> None:
return None
def execute(
self,
statement: str,
parameters: object | None = None,
) -> None:
normalized = " ".join(statement.split())
self.calls.append((normalized, parameters))
if self.fail_on is not None and self.fail_on in normalized:
raise RuntimeError("database failed")
def fetchall(self) -> list[tuple[int, str]]:
return list(self.applied_rows)
@dataclass
class RecordingConnection:
cursor_value: RecordingCursor
entered: int = 0
exited: int = 0
def __enter__(self) -> RecordingConnection:
self.entered += 1
return self
def __exit__(self, *args: object) -> None:
self.exited += 1
return None
def cursor(self) -> RecordingCursor:
return self.cursor_value
@dataclass
class RecordingProvider:
connection: RecordingConnection
calls: int = 0
def __call__(self) -> RecordingConnection:
self.calls += 1
return self.connection
def _runner(
*,
applied_rows: list[tuple[int, str]] | None = None,
fail_on: str | None = None,
migrations: tuple[StorageMigration, ...] = STORAGE_MIGRATIONS,
) -> tuple[
StorageMigrationRunner,
RecordingCursor,
RecordingConnection,
RecordingProvider,
]:
cursor = RecordingCursor(
applied_rows=applied_rows or [],
fail_on=fail_on,
)
connection = RecordingConnection(cursor)
provider = RecordingProvider(connection)
runner = StorageMigrationRunner(
connection_provider=provider,
migrations=migrations,
)
return runner, cursor, connection, provider
def test_default_migrations_have_stable_order_and_names() -> None:
assert tuple(
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS
) == (
(1, "create_market_data_schema"),
(2, "create_canonical_trades"),
(3, "create_canonical_quotes"),
(4, "create_canonical_candle_revisions"),
(5, "add_trade_observation_sources"),
(6, "add_quote_and_candle_observation_sources"),
(7, "create_market_data_partition_registry"),
)
def test_default_schema_defines_partitions_identities_and_constraints() -> None:
sql = "\n".join(
statement
for migration in STORAGE_MIGRATIONS
for statement in migration.statements
)
assert "CREATE SCHEMA IF NOT EXISTS market_data" in sql
assert "CREATE TABLE market_data.trades" in sql
assert "PRIMARY KEY (venue, symbol, trade_id, executed_at)" in sql
assert "trade_id BETWEEN -2147483648 AND 2147483647" in sql
assert sql.count("CHECK (BTRIM(venue) <> '')") == 3
assert sql.count("CHECK (BTRIM(symbol) <> '')") == 3
assert sql.count("CHECK (BTRIM(source) <> '')") == 3
assert "PARTITION BY RANGE (executed_at)" in sql
assert "CREATE TABLE market_data.quotes" in sql
assert "PRIMARY KEY (venue, symbol, received_at)" in sql
assert "PARTITION BY RANGE (received_at)" in sql
assert "CREATE TABLE market_data.candle_revisions" in sql
assert "open_time,\n observed_at" in sql
assert "PARTITION BY RANGE (open_time)" in sql
assert sql.count("PARTITION OF") == 3
assert "ADD COLUMN observation_sources TEXT[]" in sql
assert "SET observation_sources = ARRAY[source]" in sql
assert "CARDINALITY(observation_sources) > 0" in sql
assert "ALTER TABLE market_data.quotes" in sql
assert "ALTER TABLE market_data.candle_revisions" in sql
assert "quotes_observation_sources_not_empty" in sql
assert "candle_revisions_observation_sources_not_empty" in sql
assert "CREATE TABLE market_data.partition_registry" in sql
assert "PRIMARY KEY (data_type, range_start)" in sql
assert "UNIQUE (partition_name)" in sql
assert "partition_bound TEXT NOT NULL" in sql
assert "BTRIM(partition_bound) <> ''" in sql
assert "range_end > range_start" in sql
def test_run_locks_and_applies_every_pending_migration_in_order() -> None:
runner, cursor, connection, provider = _runner()
result = runner.run()
assert result == (1, 2, 3, 4, 5, 6, 7)
assert provider.calls == 1
assert connection.entered == 1
assert connection.exited == 1
assert cursor.calls[0] == (
"SELECT pg_advisory_xact_lock(%s)",
(STORAGE_MIGRATION_ADVISORY_LOCK_ID,),
)
assert "public.storage_schema_migrations" in cursor.calls[1][0]
inserted_versions = tuple(
parameters[0]
for statement, parameters in cursor.calls
if statement.startswith(
"INSERT INTO public.storage_schema_migrations"
)
and isinstance(parameters, tuple)
)
assert inserted_versions == (1, 2, 3, 4, 5, 6, 7)
def test_run_skips_already_applied_migrations() -> None:
applied = [
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS
]
runner, cursor, _, _ = _runner(applied_rows=applied)
result = runner.run()
assert result == ()
assert not any(
statement.startswith("CREATE SCHEMA")
or statement.startswith("CREATE TABLE market_data")
for statement, _ in cursor.calls
)
def test_run_applies_only_migrations_after_existing_prefix() -> None:
applied = [
(migration.version, migration.name)
for migration in STORAGE_MIGRATIONS[:2]
]
runner, cursor, _, _ = _runner(applied_rows=applied)
result = runner.run()
assert result == (3, 4, 5, 6, 7)
inserted_versions = tuple(
parameters[0]
for statement, parameters in cursor.calls
if statement.startswith(
"INSERT INTO public.storage_schema_migrations"
)
and isinstance(parameters, tuple)
)
assert inserted_versions == (3, 4, 5, 6, 7)
def test_run_rejects_unknown_applied_version() -> None:
runner, _, _, _ = _runner(applied_rows=[(99, "future")])
with pytest.raises(StorageMigrationError, match="unknown.*99"):
runner.run()
def test_run_rejects_changed_applied_name() -> None:
runner, _, _, _ = _runner(applied_rows=[(1, "renamed")])
with pytest.raises(StorageMigrationError, match="name mismatch"):
runner.run()
def test_run_rejects_non_prefix_applied_history() -> None:
second = STORAGE_MIGRATIONS[1]
runner, _, _, _ = _runner(
applied_rows=[(second.version, second.name)],
)
with pytest.raises(StorageMigrationError, match="ordered prefix"):
runner.run()
def test_database_error_is_wrapped() -> None:
runner, _, _, _ = _runner(fail_on="CREATE SCHEMA")
with pytest.raises(StorageMigrationError) as error_info:
runner.run()
assert isinstance(error_info.value.__cause__, RuntimeError)
def test_keyboard_interrupt_is_not_wrapped() -> None:
def interrupted_provider() -> Any:
raise KeyboardInterrupt
runner = StorageMigrationRunner(
connection_provider=interrupted_provider,
)
with pytest.raises(KeyboardInterrupt):
runner.run()
def test_runner_rejects_duplicate_versions() -> None:
migration = StorageMigration(
version=1,
name="one",
statements=("SELECT 1",),
)
with pytest.raises(ValueError, match="unique"):
StorageMigrationRunner(
connection_provider=lambda: None, # type: ignore[arg-type]
migrations=(migration, migration),
)
def test_runner_rejects_out_of_order_versions() -> None:
first = StorageMigration(1, "one", ("SELECT 1",))
second = StorageMigration(2, "two", ("SELECT 2",))
with pytest.raises(ValueError, match="ordered"):
StorageMigrationRunner(
connection_provider=lambda: None, # type: ignore[arg-type]
migrations=(second, first),
)
@pytest.mark.parametrize(
"arguments",
(
{"version": 0, "name": "zero", "statements": ("SELECT 0",)},
{"version": True, "name": "one", "statements": ("SELECT 1",)},
{"version": 1, "name": " ", "statements": ("SELECT 1",)},
{"version": 1, "name": "one", "statements": ()},
{"version": 1, "name": "one", "statements": (" ",)},
),
)
def test_migration_rejects_invalid_definition(
arguments: dict[str, object],
) -> None:
with pytest.raises(ValueError):
StorageMigration(**arguments) # type: ignore[arg-type]

View File

@@ -0,0 +1,236 @@
from __future__ import annotations
from contextlib import nullcontext
from typing import Any
import pytest
from src.storage.exceptions import PostgresConnectionPoolError
from src.storage.postgres_pool import PostgresConnectionPool
class RecordingPool:
def __init__(
self,
*,
open_error: BaseException | None = None,
close_error: Exception | None = None,
) -> None:
self.close_calls: list[float] = []
self.close_error = close_error
self.connection_calls: list[float] = []
self.open_calls: list[tuple[bool, float]] = []
self.open_error = open_error
self.connection_value = object()
def open(self, *, wait: bool, timeout: float) -> None:
self.open_calls.append((wait, timeout))
if self.open_error is not None:
raise self.open_error
def connection(self, *, timeout: float):
self.connection_calls.append(timeout)
return nullcontext(self.connection_value)
def close(self, *, timeout: float) -> None:
self.close_calls.append(timeout)
if self.close_error is not None:
raise self.close_error
class RecordingPoolFactory:
def __init__(self, pool: RecordingPool) -> None:
self.calls: list[dict[str, Any]] = []
self.pool = pool
def __call__(self, **kwargs: Any) -> RecordingPool:
self.calls.append(kwargs)
return self.pool
def _connection_pool(
*,
pool: RecordingPool | None = None,
) -> tuple[
PostgresConnectionPool,
RecordingPool,
RecordingPoolFactory,
]:
recording_pool = pool or RecordingPool()
factory = RecordingPoolFactory(recording_pool)
connection_pool = PostgresConnectionPool(
conninfo="postgresql://db/dzentra",
min_size=2,
max_size=6,
timeout_seconds=7.5,
name="market-data",
pool_factory=factory,
)
return connection_pool, recording_pool, factory
def test_constructor_does_not_create_or_open_pool() -> None:
connection_pool, pool, factory = _connection_pool()
assert connection_pool.is_open is False
assert factory.calls == []
assert pool.open_calls == []
def test_open_creates_pool_with_explicit_configuration_and_waits() -> None:
connection_pool, pool, factory = _connection_pool()
connection_pool.open()
assert connection_pool.is_open is True
assert factory.calls == [
{
"conninfo": "postgresql://db/dzentra",
"min_size": 2,
"max_size": 6,
"timeout": 7.5,
"kwargs": {"autocommit": False},
"name": "market-data",
"open": False,
}
]
assert pool.open_calls == [(True, 7.5)]
def test_repeated_open_is_no_op() -> None:
connection_pool, pool, factory = _connection_pool()
connection_pool.open()
connection_pool.open()
assert len(factory.calls) == 1
assert pool.open_calls == [(True, 7.5)]
def test_connection_requires_open_pool() -> None:
connection_pool, _, _ = _connection_pool()
with pytest.raises(PostgresConnectionPoolError, match="not open"):
connection_pool.connection()
def test_connection_borrows_from_pool_with_timeout() -> None:
connection_pool, pool, _ = _connection_pool()
connection_pool.open()
with connection_pool.connection() as connection:
assert connection is pool.connection_value
assert pool.connection_calls == [7.5]
def test_close_is_idempotent() -> None:
connection_pool, pool, _ = _connection_pool()
connection_pool.open()
connection_pool.close()
connection_pool.close()
assert connection_pool.is_open is False
assert pool.close_calls == [7.5]
def test_open_failure_closes_partial_pool_and_is_retryable() -> None:
pool = RecordingPool(open_error=RuntimeError("database unavailable"))
connection_pool, _, factory = _connection_pool(pool=pool)
with pytest.raises(PostgresConnectionPoolError) as error_info:
connection_pool.open()
assert isinstance(error_info.value.__cause__, RuntimeError)
assert connection_pool.is_open is False
assert pool.close_calls == [7.5]
pool.open_error = None
connection_pool.open()
assert len(factory.calls) == 2
assert connection_pool.is_open is True
def test_open_failure_preserves_cleanup_failure_as_note() -> None:
pool = RecordingPool(
open_error=RuntimeError("open failed"),
close_error=RuntimeError("close failed"),
)
connection_pool, _, _ = _connection_pool(pool=pool)
with pytest.raises(PostgresConnectionPoolError) as error_info:
connection_pool.open()
cause = error_info.value.__cause__
assert isinstance(cause, RuntimeError)
assert cause.__notes__ == [
"PostgreSQL pool cleanup also failed: RuntimeError."
]
def test_close_failure_leaves_wrapper_closed() -> None:
pool = RecordingPool(close_error=RuntimeError("close failed"))
connection_pool, _, _ = _connection_pool(pool=pool)
connection_pool.open()
with pytest.raises(PostgresConnectionPoolError) as error_info:
connection_pool.close()
assert isinstance(error_info.value.__cause__, RuntimeError)
assert connection_pool.is_open is False
connection_pool.close()
assert pool.close_calls == [7.5]
def test_keyboard_interrupt_is_not_wrapped() -> None:
def interrupted_factory(**kwargs: object) -> object:
raise KeyboardInterrupt
connection_pool = PostgresConnectionPool(
conninfo="postgresql://db/dzentra",
pool_factory=interrupted_factory,
)
with pytest.raises(KeyboardInterrupt):
connection_pool.open()
def test_keyboard_interrupt_during_pool_open_closes_partial_pool() -> None:
pool = RecordingPool(open_error=KeyboardInterrupt())
connection_pool, _, _ = _connection_pool(pool=pool)
with pytest.raises(KeyboardInterrupt):
connection_pool.open()
assert connection_pool.is_open is False
assert pool.open_calls == [(True, 7.5)]
assert pool.close_calls == [7.5]
@pytest.mark.parametrize(
("override", "message"),
(
({"conninfo": ""}, "conninfo"),
({"min_size": 0}, "min_size"),
({"min_size": True}, "min_size"),
({"min_size": 2, "max_size": 1}, "max_size"),
({"timeout_seconds": 0}, "timeout_seconds"),
({"timeout_seconds": float("inf")}, "timeout_seconds"),
({"name": " "}, "name"),
),
)
def test_constructor_rejects_invalid_configuration(
override: dict[str, object],
message: str,
) -> None:
arguments: dict[str, object] = {
"conninfo": "postgresql://db/dzentra",
}
arguments.update(override)
with pytest.raises(ValueError, match=message):
PostgresConnectionPool(**arguments) # type: ignore[arg-type]

View File

@@ -0,0 +1,190 @@
from __future__ import annotations
from typing import Any
import pytest
from psycopg.conninfo import conninfo_to_dict
from tests.support.postgres_market_data import (
POSTGRES_TEST_APPLICATION_NAME,
POSTGRES_TEST_CONTROL_APPLICATION_NAME,
load_postgres_test_settings,
reset_postgres_test_database,
)
class RecordingCursor:
def __init__(
self,
*,
identity: tuple[str, str],
statements: list[str],
) -> None:
self._identity = identity
self._statements = statements
def __enter__(self) -> RecordingCursor:
return self
def __exit__(self, *_: object) -> None:
return None
def execute(
self,
statement: object,
parameters: object = None,
) -> None:
del parameters
self._statements.append(" ".join(str(statement).split()))
def fetchone(self) -> tuple[str, str]:
return self._identity
class RecordingControlConnection:
def __init__(
self,
*,
identity: tuple[str, str],
autocommit: bool = True,
) -> None:
self.autocommit = autocommit
self.identity = identity
self.statements: list[str] = []
def cursor(self) -> RecordingCursor:
return RecordingCursor(
identity=self.identity,
statements=self.statements,
)
def test_postgres_harness_is_disabled_without_explicit_flag() -> None:
assert load_postgres_test_settings({}) is None
def test_postgres_harness_requires_exact_opt_in() -> None:
with pytest.raises(ValueError, match="exactly '1'"):
load_postgres_test_settings(
{"DZENTRA_RUN_POSTGRES_TESTS": "true"}
)
def test_postgres_harness_requires_explicit_dsn() -> None:
with pytest.raises(ValueError, match="DZENTRA_TEST_POSTGRES_DSN"):
load_postgres_test_settings(
{"DZENTRA_RUN_POSTGRES_TESTS": "1"}
)
@pytest.mark.parametrize(
"dsn",
(
"postgresql://localhost/production",
"postgresql://db.example.test/dzentra_test_060_27",
"hostaddr=192.0.2.10 dbname=dzentra_test_060_27",
"postgresql://localhost/dzentra_test_unsafe-name",
"service=production dbname=dzentra_test_shadow",
"dbname=dzentra_test_implicit_endpoint",
"host=/tmp,db.example.test dbname=dzentra_test_socket_fallback",
(
"hostaddr=127.0.0.1,192.0.2.10 "
"dbname=dzentra_test_address_fallback"
),
"host=localhost, dbname=dzentra_test_implicit_fallback",
),
)
def test_postgres_harness_rejects_unsafe_database_targets(
dsn: str,
) -> None:
with pytest.raises(ValueError):
load_postgres_test_settings(
{
"DZENTRA_RUN_POSTGRES_TESTS": "1",
"DZENTRA_TEST_POSTGRES_DSN": dsn,
}
)
@pytest.mark.parametrize(
"dsn",
(
"postgresql://localhost:5544/dzentra_test_060_27",
"host=/tmp dbname=dzentra_test_socket user=test",
"hostaddr=127.0.0.1 dbname=dzentra_test_loopback",
"hostaddr=::1 dbname=dzentra_test_ipv6",
),
)
def test_postgres_harness_accepts_only_explicit_local_test_database(
dsn: str,
) -> None:
settings = load_postgres_test_settings(
{
"DZENTRA_RUN_POSTGRES_TESTS": "1",
"DZENTRA_TEST_POSTGRES_DSN": dsn,
}
)
assert settings is not None
parsed = conninfo_to_dict(settings.dsn)
assert parsed["dbname"] == settings.database_name
assert parsed["application_name"] == POSTGRES_TEST_APPLICATION_NAME
assert parsed["connect_timeout"] == "5"
def test_postgres_harness_does_not_accept_environment_endpoint_fallback(
) -> None:
with pytest.raises(ValueError, match="explicit local host"):
load_postgres_test_settings(
{
"DZENTRA_RUN_POSTGRES_TESTS": "1",
"DZENTRA_TEST_POSTGRES_DSN": (
"dbname=dzentra_test_environment_fallback"
),
"PGHOST": "localhost",
}
)
def test_postgres_reset_revalidates_database_and_control_identity() -> None:
database_name = "dzentra_test_cleanup"
connection: Any = RecordingControlConnection(
identity=(
database_name,
POSTGRES_TEST_CONTROL_APPLICATION_NAME,
)
)
reset_postgres_test_database(
connection,
expected_database_name=database_name,
)
assert connection.statements == [
"SELECT current_database(), current_setting('application_name')",
"DROP SCHEMA IF EXISTS market_data CASCADE",
"DROP TABLE IF EXISTS public.storage_schema_migrations",
]
@pytest.mark.parametrize(
"identity",
(
("production", POSTGRES_TEST_CONTROL_APPLICATION_NAME),
("dzentra_test_cleanup", "unexpected-application"),
),
)
def test_postgres_reset_refuses_unvalidated_connection_before_drop(
identity: tuple[str, str],
) -> None:
connection: Any = RecordingControlConnection(identity=identity)
with pytest.raises(RuntimeError, match="Refusing destructive reset"):
reset_postgres_test_database(
connection,
expected_database_name="dzentra_test_cleanup",
)
assert connection.statements == [
"SELECT current_database(), current_setting('application_name')",
]