#!/usr/bin/env python3 """ Setup script for unified data storage system. Initializes TimescaleDB schema and verifies setup. """ import asyncio import asyncpg import logging import sys from pathlib import Path # Add parent directory to path sys.path.insert(0, str(Path(__file__).parent.parent)) from core.unified_storage_schema import UnifiedStorageSchemaManager from core.config import get_config # Setup logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' ) logger = logging.getLogger(__name__) async def create_connection_pool(config): """Create database connection pool.""" try: # Build connection string db_config = config.get('database', {}) # Default values host = db_config.get('host', 'localhost') port = db_config.get('port', 5432) database = db_config.get('name', 'trading_data') user = db_config.get('user', 'postgres') password = db_config.get('password', 'postgres') logger.info(f"Connecting to database: {host}:{port}/{database}") pool = await asyncpg.create_pool( host=host, port=port, database=database, user=user, password=password, min_size=2, max_size=10, command_timeout=60 ) logger.info("Database connection pool created") return pool except Exception as e: logger.error(f"Failed to create connection pool: {e}") raise async def verify_timescaledb(pool): """Verify TimescaleDB extension is available.""" try: async with pool.acquire() as conn: # Check if TimescaleDB extension exists result = await conn.fetchval(""" SELECT EXISTS ( SELECT FROM pg_extension WHERE extname = 'timescaledb' ) """) if result: # Get TimescaleDB version version = await conn.fetchval("SELECT extversion FROM pg_extension WHERE extname = 'timescaledb'") logger.info(f"TimescaleDB extension found (version {version})") return True else: logger.warning("TimescaleDB extension not found, attempting to create...") # Try to create extension await conn.execute("CREATE EXTENSION IF NOT EXISTS timescaledb CASCADE") logger.info("TimescaleDB extension created successfully") return True except Exception as e: logger.error(f"Failed to verify TimescaleDB: {e}") logger.error("Please ensure TimescaleDB is installed: https://docs.timescale.com/install/latest/") return False async def setup_schema(pool): """Set up the complete unified storage schema.""" try: schema_manager = UnifiedStorageSchemaManager(pool) logger.info("Starting schema setup...") success = await schema_manager.setup_complete_schema() if success: logger.info("Schema setup completed successfully") # Verify schema logger.info("Verifying schema...") verified = await schema_manager.verify_schema() if verified: logger.info("Schema verification passed") # Get schema info info = await schema_manager.get_schema_info() logger.info("\n=== Schema Information ===") logger.info(f"Migrations applied: {len(info.get('migrations', []))}") logger.info(f"Tables created: {len(info.get('tables', []))}") logger.info(f"Hypertables: {len(info.get('hypertables', []))}") logger.info(f"Continuous aggregates: {len(info.get('continuous_aggregates', []))}") # Display table sizes logger.info("\n=== Table Sizes ===") for table in info.get('tables', []): logger.info(f" {table['tablename']}: {table['size']}") # Display hypertables logger.info("\n=== Hypertables ===") for ht in info.get('hypertables', []): logger.info(f" {ht['hypertable_name']}: {ht['num_chunks']} chunks, " f"compression={'enabled' if ht['compression_enabled'] else 'disabled'}") # Display continuous aggregates if info.get('continuous_aggregates'): logger.info("\n=== Continuous Aggregates ===") for ca in info.get('continuous_aggregates', []): logger.info(f" {ca['view_name']}: {ca.get('size', 'N/A')}") return True else: logger.error("Schema verification failed") return False else: logger.error("Schema setup failed") return False except Exception as e: logger.error(f"Error during schema setup: {e}") return False async def test_basic_operations(pool): """Test basic database operations.""" try: logger.info("\n=== Testing Basic Operations ===") async with pool.acquire() as conn: # Test insert into ohlcv_data logger.info("Testing OHLCV insert...") await conn.execute(""" INSERT INTO ohlcv_data (timestamp, symbol, timeframe, open_price, high_price, low_price, close_price, volume) VALUES (NOW(), 'ETH/USDT', '1s', 2000.0, 2001.0, 1999.0, 2000.5, 100.0) ON CONFLICT (timestamp, symbol, timeframe) DO NOTHING """) logger.info("✓ OHLCV insert successful") # Test query logger.info("Testing OHLCV query...") result = await conn.fetchrow(""" SELECT * FROM ohlcv_data WHERE symbol = 'ETH/USDT' ORDER BY timestamp DESC LIMIT 1 """) if result: logger.info(f"✓ OHLCV query successful: {dict(result)}") # Test order book insert logger.info("Testing order book insert...") await conn.execute(""" INSERT INTO order_book_snapshots (timestamp, symbol, exchange, bids, asks, mid_price, spread) VALUES (NOW(), 'ETH/USDT', 'binance', '[]'::jsonb, '[]'::jsonb, 2000.0, 0.1) ON CONFLICT (timestamp, symbol, exchange) DO NOTHING """) logger.info("✓ Order book insert successful") # Test imbalances insert logger.info("Testing imbalances insert...") await conn.execute(""" INSERT INTO order_book_imbalances (timestamp, symbol, imbalance_1s, imbalance_5s, imbalance_15s, imbalance_60s) VALUES (NOW(), 'ETH/USDT', 0.5, 0.4, 0.3, 0.2) ON CONFLICT (timestamp, symbol) DO NOTHING """) logger.info("✓ Imbalances insert successful") logger.info("\n✓ All basic operations successful") return True except Exception as e: logger.error(f"Basic operations test failed: {e}") return False async def main(): """Main setup function.""" logger.info("=== Unified Data Storage Setup ===\n") pool = None try: # Load configuration config = get_config() # Create connection pool pool = await create_connection_pool(config) # Verify TimescaleDB if not await verify_timescaledb(pool): logger.error("TimescaleDB verification failed") return 1 # Setup schema if not await setup_schema(pool): logger.error("Schema setup failed") return 1 # Test basic operations if not await test_basic_operations(pool): logger.error("Basic operations test failed") return 1 logger.info("\n=== Setup Complete ===") logger.info("Unified data storage system is ready to use!") return 0 except Exception as e: logger.error(f"Setup failed: {e}") return 1 finally: if pool: await pool.close() logger.info("Database connection pool closed") if __name__ == "__main__": exit_code = asyncio.run(main()) sys.exit(exit_code)