bybit_btcusdt_ingest/storage.go

package main

import (
	"database/sql"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	_ "modernc.org/sqlite"
)

// StorageManager handles all SQLite database operations: initialization,
// raw tick writing, feature writing, hourly migration, and pruning.
type StorageManager struct {
	cfg        Config
	hotDBPath  string
	featDBPath string
	archiveDir string
}

// NewStorageManager creates directories, initializes databases, and returns a ready manager.
func NewStorageManager(cfg Config) (*StorageManager, error) {
	archiveDir := filepath.Join(cfg.DataDir, "archive")
	if err := os.MkdirAll(archiveDir, 0755); err != nil {
		return nil, fmt.Errorf("create archive dir: %w", err)
	}

	sm := &StorageManager{
		cfg:        cfg,
		hotDBPath:  filepath.Join(cfg.DataDir, "hot_ticks.db"),
		featDBPath: filepath.Join(cfg.DataDir, "features.db"),
		archiveDir: archiveDir,
	}

	if err := sm.initHotDB(); err != nil {
		return nil, fmt.Errorf("init hot db: %w", err)
	}
	if err := sm.initFeaturesDB(); err != nil {
		return nil, fmt.Errorf("init features db: %w", err)
	}

	return sm, nil
}

// OpenHotDB returns a new connection to the hot ticks database.
func (sm *StorageManager) OpenHotDB() (*sql.DB, error) {
	db, err := sql.Open("sqlite", sm.hotDBPath)
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
		db.Close()
		return nil, err
	}
	if _, err := db.Exec("PRAGMA synchronous=NORMAL"); err != nil {
		db.Close()
		return nil, err
	}
	return db, nil
}

// OpenFeaturesDB returns a new connection to the features database.
func (sm *StorageManager) OpenFeaturesDB() (*sql.DB, error) {
	db, err := sql.Open("sqlite", sm.featDBPath)
	if err != nil {
		return nil, err
	}
	if _, err := db.Exec("PRAGMA journal_mode=WAL"); err != nil {
		db.Close()
		return nil, err
	}
	if _, err := db.Exec("PRAGMA synchronous=NORMAL"); err != nil {
		db.Close()
		return nil, err
	}
	return db, nil
}

func (sm *StorageManager) initHotDB() error {
	db, err := sm.OpenHotDB()
	if err != nil {
		return err
	}
	defer db.Close()

	_, err = db.Exec(`
		PRAGMA auto_vacuum=INCREMENTAL;

		CREATE TABLE IF NOT EXISTS btc_ticks (
			timestamp INTEGER NOT NULL,
			price REAL NOT NULL,
			volume REAL NOT NULL,
			side TEXT NOT NULL
		);

		CREATE INDEX IF NOT EXISTS idx_timestamp ON btc_ticks(timestamp);
	`)
	return err
}

func (sm *StorageManager) initFeaturesDB() error {
	db, err := sm.OpenFeaturesDB()
	if err != nil {
		return err
	}
	defer db.Close()

	_, err = db.Exec(`
		PRAGMA auto_vacuum=INCREMENTAL;

		CREATE TABLE IF NOT EXISTS five_second_features (
			timestamp INTEGER PRIMARY KEY,
			log_return REAL NOT NULL,
			realized_vol REAL NOT NULL,
			ofi REAL NOT NULL,
			volume_sum REAL NOT NULL,
			close_price REAL NOT NULL,
			vwap REAL NOT NULL
		);
	`)
	return err
}

// weeklyArchivePath resolves the archive DB file path for a given timestamp.
func (sm *StorageManager) weeklyArchivePath(timestampMs int64) string {
	t := time.UnixMilli(timestampMs)
	year, week := t.ISOWeek()
	filename := fmt.Sprintf("btc_ticks_%d_W%02d.db", year, week)
	return filepath.Join(sm.archiveDir, filename)
}

// initArchiveDB ensures the archive database has the correct schema.
func (sm *StorageManager) initArchiveDB(path string) error {
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return err
	}
	defer db.Close()

	_, err = db.Exec(`
		PRAGMA journal_mode=WAL;
		PRAGMA synchronous=NORMAL;

		CREATE TABLE IF NOT EXISTS btc_ticks (
			timestamp INTEGER NOT NULL,
			price REAL NOT NULL,
			volume REAL NOT NULL,
			side TEXT NOT NULL
		);

		CREATE INDEX IF NOT EXISTS idx_timestamp ON btc_ticks(timestamp);
	`)
	return err
}

// RunHourlyMaintenance performs tick migration and feature pruning.
func (sm *StorageManager) RunHourlyMaintenance() {
	nowMs := time.Now().UnixMilli()
	log.Println("[maintenance] Starting hourly database maintenance...")

	// Part A: Migrate raw ticks older than retention window
	cutoff12h := nowMs - int64(sm.cfg.HotRetentionHours)*60*60*1000
	if err := sm.migrateRawTicks(cutoff12h); err != nil {
		log.Printf("[maintenance] Raw tick migration failed: %v", err)
	}

	// Part B: Prune features older than retention window
	cutoffFeatures := nowMs - int64(sm.cfg.FeatureRetentionDays)*24*60*60*1000
	if err := sm.pruneFeatures(cutoffFeatures); err != nil {
		log.Printf("[maintenance] Feature pruning failed: %v", err)
	}

	log.Println("[maintenance] Hourly maintenance complete.")
}

// StartupRecovery checks for stale ticks and migrates them before normal operation.
func (sm *StorageManager) StartupRecovery() error {
	log.Println("[startup] Checking for stale ticks in hot database...")

	db, err := sm.OpenHotDB()
	if err != nil {
		return err
	}
	defer db.Close()

	cutoffMs := time.Now().UnixMilli() - int64(sm.cfg.HotRetentionHours)*60*60*1000
	var count int64
	err = db.QueryRow("SELECT COUNT(*) FROM btc_ticks WHERE timestamp < ?", cutoffMs).Scan(&count)
	if err != nil {
		return fmt.Errorf("count stale ticks: %w", err)
	}

	if count > 0 {
		log.Printf("[startup] Found %d stale ticks, triggering immediate migration...", count)
		if err := sm.migrateRawTicks(cutoffMs); err != nil {
			return fmt.Errorf("startup migration: %w", err)
		}
	} else {
		log.Println("[startup] No stale ticks found, hot database is clean.")
	}

	return nil
}

func (sm *StorageManager) migrateRawTicks(cutoffMs int64) error {
	db, err := sm.OpenHotDB()
	if err != nil {
		return err
	}
	defer db.Close()

	// Find the range of ticks that need migration
	var minTS, maxTS sql.NullInt64
	err = db.QueryRow(`
		SELECT MIN(timestamp), MAX(timestamp)
		FROM btc_ticks
		WHERE timestamp < ?
	`, cutoffMs).Scan(&minTS, &maxTS)
	if err != nil {
		return fmt.Errorf("query migration range: %w", err)
	}

	if !minTS.Valid || !maxTS.Valid {
		log.Println("[maintenance] No raw ticks older than cutoff to migrate.")
		return nil
	}

	// Process ticks week by week to handle cross-week boundaries
	currentTs := minTS.Int64
	for currentTs <= maxTS.Int64 {
		t := time.UnixMilli(currentTs)
		year, week := t.ISOWeek()
		archivePath := sm.weeklyArchivePath(currentTs)

		log.Printf("[maintenance] Migrating ticks for %d W%02d -> %s",
			year, week, filepath.Base(archivePath))

		// Calculate the start of the next ISO week (Monday 00:00:00 UTC)
		nextWeekStart := sm.nextISOWeekStartMs(currentTs)

		// The upper bound for this chunk: either next week start or cutoff, whichever is smaller
		chunkEnd := nextWeekStart
		if cutoffMs < chunkEnd {
			chunkEnd = cutoffMs
		}

		// Ensure archive DB exists and has schema
		if err := sm.initArchiveDB(archivePath); err != nil {
			return fmt.Errorf("init archive db: %w", err)
		}

		// Attach and atomically migrate
		if err := sm.atomicMigrate(db, archivePath, currentTs, chunkEnd); err != nil {
			return fmt.Errorf("atomic migrate: %w", err)
		}

		currentTs = nextWeekStart
	}

	// Reclaim space
	if _, err := db.Exec("PRAGMA incremental_vacuum(500)"); err != nil {
		log.Printf("[maintenance] incremental_vacuum warning: %v", err)
	}

	log.Println("[maintenance] Raw tick migration completed successfully.")
	return nil
}

// nextISOWeekStartMs calculates the epoch ms of the next Monday 00:00:00 UTC
// relative to the given timestamp.
func (sm *StorageManager) nextISOWeekStartMs(timestampMs int64) int64 {
	t := time.UnixMilli(timestampMs).UTC()
	// Calculate days until next Monday
	daysUntilMonday := (8 - int(t.Weekday())) % 7
	if daysUntilMonday == 0 {
		daysUntilMonday = 7
	}
	nextMonday := time.Date(t.Year(), t.Month(), t.Day()+daysUntilMonday, 0, 0, 0, 0, time.UTC)
	return nextMonday.UnixMilli()
}

func (sm *StorageManager) atomicMigrate(hotDB *sql.DB, archivePath string, fromMs, toMs int64) error {
	// Attach the archive database
	_, err := hotDB.Exec("ATTACH DATABASE ? AS archive", archivePath)
	if err != nil {
		return fmt.Errorf("attach archive: %w", err)
	}
	defer hotDB.Exec("DETACH DATABASE archive")

	tx, err := hotDB.Begin()
	if err != nil {
		return fmt.Errorf("begin migration tx: %w", err)
	}

	// Insert into archive
	_, err = tx.Exec(`
		INSERT INTO archive.btc_ticks (timestamp, price, volume, side)
		SELECT timestamp, price, volume, side
		FROM main.btc_ticks
		WHERE timestamp >= ? AND timestamp < ?
	`, fromMs, toMs)
	if err != nil {
		tx.Rollback()
		return fmt.Errorf("insert into archive: %w", err)
	}

	// Delete from hot
	result, err := tx.Exec(`
		DELETE FROM main.btc_ticks
		WHERE timestamp >= ? AND timestamp < ?
	`, fromMs, toMs)
	if err != nil {
		tx.Rollback()
		return fmt.Errorf("delete from hot: %w", err)
	}

	migrated, _ := result.RowsAffected()
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit migration: %w", err)
	}

	log.Printf("[maintenance] Migrated %d ticks to archive.", migrated)
	return nil
}

func (sm *StorageManager) pruneFeatures(cutoffMs int64) error {
	db, err := sm.OpenFeaturesDB()
	if err != nil {
		return err
	}
	defer db.Close()

	cutoffTime := time.UnixMilli(cutoffMs)
	log.Printf("[maintenance] Pruning features older than %s", cutoffTime.Format(time.RFC3339))

	result, err := db.Exec("DELETE FROM five_second_features WHERE timestamp < ?", cutoffMs)
	if err != nil {
		return fmt.Errorf("delete old features: %w", err)
	}

	deleted, _ := result.RowsAffected()
	if deleted > 0 {
		log.Printf("[maintenance] Pruned %d old feature rows.", deleted)
		if _, err := db.Exec("PRAGMA incremental_vacuum(100)"); err != nil {
			log.Printf("[maintenance] feature vacuum warning: %v", err)
		}
	} else {
		log.Println("[maintenance] No feature rows needed pruning.")
	}

	return nil
}