You can not select more than 25 topics
			Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
		
		
		
		
		
			
		
			
				
					
					
						
							302 lines
						
					
					
						
							8.8 KiB
						
					
					
				
			
		
		
		
			
			
			
		
		
	
	
							302 lines
						
					
					
						
							8.8 KiB
						
					
					
				
								package offset
							 | 
						|
								
							 | 
						|
								import (
							 | 
						|
									"database/sql"
							 | 
						|
									"fmt"
							 | 
						|
									"time"
							 | 
						|
								)
							 | 
						|
								
							 | 
						|
								// MigrationVersion represents a database migration version
							 | 
						|
								type MigrationVersion struct {
							 | 
						|
									Version     int
							 | 
						|
									Description string
							 | 
						|
									SQL         string
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// GetMigrations returns all available migrations for offset storage
							 | 
						|
								func GetMigrations() []MigrationVersion {
							 | 
						|
									return []MigrationVersion{
							 | 
						|
										{
							 | 
						|
											Version:     1,
							 | 
						|
											Description: "Create initial offset storage tables",
							 | 
						|
											SQL: `
							 | 
						|
												-- Partition offset checkpoints table
							 | 
						|
												-- TODO: Add _index as computed column when supported by database
							 | 
						|
												CREATE TABLE IF NOT EXISTS partition_offset_checkpoints (
							 | 
						|
													partition_key TEXT PRIMARY KEY,
							 | 
						|
													ring_size INTEGER NOT NULL,
							 | 
						|
													range_start INTEGER NOT NULL,
							 | 
						|
													range_stop INTEGER NOT NULL,
							 | 
						|
													unix_time_ns INTEGER NOT NULL,
							 | 
						|
													checkpoint_offset INTEGER NOT NULL,
							 | 
						|
													updated_at INTEGER NOT NULL
							 | 
						|
												);
							 | 
						|
												
							 | 
						|
												-- Offset mappings table for detailed tracking
							 | 
						|
												-- TODO: Add _index as computed column when supported by database
							 | 
						|
												CREATE TABLE IF NOT EXISTS offset_mappings (
							 | 
						|
													id INTEGER PRIMARY KEY AUTOINCREMENT,
							 | 
						|
													partition_key TEXT NOT NULL,
							 | 
						|
													kafka_offset INTEGER NOT NULL,
							 | 
						|
													smq_timestamp INTEGER NOT NULL,
							 | 
						|
													message_size INTEGER NOT NULL,
							 | 
						|
													created_at INTEGER NOT NULL,
							 | 
						|
													UNIQUE(partition_key, kafka_offset)
							 | 
						|
												);
							 | 
						|
												
							 | 
						|
												-- Schema migrations tracking table
							 | 
						|
												CREATE TABLE IF NOT EXISTS schema_migrations (
							 | 
						|
													version INTEGER PRIMARY KEY,
							 | 
						|
													description TEXT NOT NULL,
							 | 
						|
													applied_at INTEGER NOT NULL
							 | 
						|
												);
							 | 
						|
											`,
							 | 
						|
										},
							 | 
						|
										{
							 | 
						|
											Version:     2,
							 | 
						|
											Description: "Add indexes for performance optimization",
							 | 
						|
											SQL: `
							 | 
						|
												-- Indexes for performance
							 | 
						|
												CREATE INDEX IF NOT EXISTS idx_partition_offset_checkpoints_partition 
							 | 
						|
												ON partition_offset_checkpoints(partition_key);
							 | 
						|
												
							 | 
						|
												CREATE INDEX IF NOT EXISTS idx_offset_mappings_partition_offset 
							 | 
						|
												ON offset_mappings(partition_key, kafka_offset);
							 | 
						|
												
							 | 
						|
												CREATE INDEX IF NOT EXISTS idx_offset_mappings_timestamp 
							 | 
						|
												ON offset_mappings(partition_key, smq_timestamp);
							 | 
						|
												
							 | 
						|
												CREATE INDEX IF NOT EXISTS idx_offset_mappings_created_at 
							 | 
						|
												ON offset_mappings(created_at);
							 | 
						|
											`,
							 | 
						|
										},
							 | 
						|
										{
							 | 
						|
											Version:     3,
							 | 
						|
											Description: "Add partition metadata table for enhanced tracking",
							 | 
						|
											SQL: `
							 | 
						|
												-- Partition metadata table
							 | 
						|
												CREATE TABLE IF NOT EXISTS partition_metadata (
							 | 
						|
													partition_key TEXT PRIMARY KEY,
							 | 
						|
													ring_size INTEGER NOT NULL,
							 | 
						|
													range_start INTEGER NOT NULL,
							 | 
						|
													range_stop INTEGER NOT NULL,
							 | 
						|
													unix_time_ns INTEGER NOT NULL,
							 | 
						|
													created_at INTEGER NOT NULL,
							 | 
						|
													last_activity_at INTEGER NOT NULL,
							 | 
						|
													record_count INTEGER DEFAULT 0,
							 | 
						|
													total_size INTEGER DEFAULT 0
							 | 
						|
												);
							 | 
						|
												
							 | 
						|
												-- Index for partition metadata
							 | 
						|
												CREATE INDEX IF NOT EXISTS idx_partition_metadata_activity 
							 | 
						|
												ON partition_metadata(last_activity_at);
							 | 
						|
											`,
							 | 
						|
										},
							 | 
						|
									}
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// MigrationManager handles database schema migrations
							 | 
						|
								type MigrationManager struct {
							 | 
						|
									db *sql.DB
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// NewMigrationManager creates a new migration manager
							 | 
						|
								func NewMigrationManager(db *sql.DB) *MigrationManager {
							 | 
						|
									return &MigrationManager{db: db}
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// GetCurrentVersion returns the current schema version
							 | 
						|
								func (m *MigrationManager) GetCurrentVersion() (int, error) {
							 | 
						|
									// First, ensure the migrations table exists
							 | 
						|
									_, err := m.db.Exec(`
							 | 
						|
										CREATE TABLE IF NOT EXISTS schema_migrations (
							 | 
						|
											version INTEGER PRIMARY KEY,
							 | 
						|
											description TEXT NOT NULL,
							 | 
						|
											applied_at INTEGER NOT NULL
							 | 
						|
										)
							 | 
						|
									`)
							 | 
						|
									if err != nil {
							 | 
						|
										return 0, fmt.Errorf("failed to create migrations table: %w", err)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									var version sql.NullInt64
							 | 
						|
									err = m.db.QueryRow("SELECT MAX(version) FROM schema_migrations").Scan(&version)
							 | 
						|
									if err != nil {
							 | 
						|
										return 0, fmt.Errorf("failed to get current version: %w", err)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									if !version.Valid {
							 | 
						|
										return 0, nil // No migrations applied yet
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									return int(version.Int64), nil
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// ApplyMigrations applies all pending migrations
							 | 
						|
								func (m *MigrationManager) ApplyMigrations() error {
							 | 
						|
									currentVersion, err := m.GetCurrentVersion()
							 | 
						|
									if err != nil {
							 | 
						|
										return fmt.Errorf("failed to get current version: %w", err)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									migrations := GetMigrations()
							 | 
						|
									
							 | 
						|
									for _, migration := range migrations {
							 | 
						|
										if migration.Version <= currentVersion {
							 | 
						|
											continue // Already applied
							 | 
						|
										}
							 | 
						|
										
							 | 
						|
										fmt.Printf("Applying migration %d: %s\n", migration.Version, migration.Description)
							 | 
						|
										
							 | 
						|
										// Begin transaction
							 | 
						|
										tx, err := m.db.Begin()
							 | 
						|
										if err != nil {
							 | 
						|
											return fmt.Errorf("failed to begin transaction for migration %d: %w", migration.Version, err)
							 | 
						|
										}
							 | 
						|
										
							 | 
						|
										// Execute migration SQL
							 | 
						|
										_, err = tx.Exec(migration.SQL)
							 | 
						|
										if err != nil {
							 | 
						|
											tx.Rollback()
							 | 
						|
											return fmt.Errorf("failed to execute migration %d: %w", migration.Version, err)
							 | 
						|
										}
							 | 
						|
										
							 | 
						|
										// Record migration as applied
							 | 
						|
										_, err = tx.Exec(
							 | 
						|
											"INSERT INTO schema_migrations (version, description, applied_at) VALUES (?, ?, ?)",
							 | 
						|
											migration.Version,
							 | 
						|
											migration.Description,
							 | 
						|
											getCurrentTimestamp(),
							 | 
						|
										)
							 | 
						|
										if err != nil {
							 | 
						|
											tx.Rollback()
							 | 
						|
											return fmt.Errorf("failed to record migration %d: %w", migration.Version, err)
							 | 
						|
										}
							 | 
						|
										
							 | 
						|
										// Commit transaction
							 | 
						|
										err = tx.Commit()
							 | 
						|
										if err != nil {
							 | 
						|
											return fmt.Errorf("failed to commit migration %d: %w", migration.Version, err)
							 | 
						|
										}
							 | 
						|
										
							 | 
						|
										fmt.Printf("Successfully applied migration %d\n", migration.Version)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									return nil
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// RollbackMigration rolls back a specific migration (if supported)
							 | 
						|
								func (m *MigrationManager) RollbackMigration(version int) error {
							 | 
						|
									// TODO: Implement rollback functionality
							 | 
						|
									// ASSUMPTION: For now, rollbacks are not supported as they require careful planning
							 | 
						|
									return fmt.Errorf("migration rollbacks not implemented - manual intervention required")
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// GetAppliedMigrations returns a list of all applied migrations
							 | 
						|
								func (m *MigrationManager) GetAppliedMigrations() ([]AppliedMigration, error) {
							 | 
						|
									rows, err := m.db.Query(`
							 | 
						|
										SELECT version, description, applied_at 
							 | 
						|
										FROM schema_migrations 
							 | 
						|
										ORDER BY version
							 | 
						|
									`)
							 | 
						|
									if err != nil {
							 | 
						|
										return nil, fmt.Errorf("failed to query applied migrations: %w", err)
							 | 
						|
									}
							 | 
						|
									defer rows.Close()
							 | 
						|
									
							 | 
						|
									var migrations []AppliedMigration
							 | 
						|
									for rows.Next() {
							 | 
						|
										var migration AppliedMigration
							 | 
						|
										err := rows.Scan(&migration.Version, &migration.Description, &migration.AppliedAt)
							 | 
						|
										if err != nil {
							 | 
						|
											return nil, fmt.Errorf("failed to scan migration: %w", err)
							 | 
						|
										}
							 | 
						|
										migrations = append(migrations, migration)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									return migrations, nil
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// ValidateSchema validates that the database schema is up to date
							 | 
						|
								func (m *MigrationManager) ValidateSchema() error {
							 | 
						|
									currentVersion, err := m.GetCurrentVersion()
							 | 
						|
									if err != nil {
							 | 
						|
										return fmt.Errorf("failed to get current version: %w", err)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									migrations := GetMigrations()
							 | 
						|
									if len(migrations) == 0 {
							 | 
						|
										return nil
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									latestVersion := migrations[len(migrations)-1].Version
							 | 
						|
									if currentVersion < latestVersion {
							 | 
						|
										return fmt.Errorf("schema is outdated: current version %d, latest version %d", currentVersion, latestVersion)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									return nil
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// AppliedMigration represents a migration that has been applied
							 | 
						|
								type AppliedMigration struct {
							 | 
						|
									Version     int
							 | 
						|
									Description string
							 | 
						|
									AppliedAt   int64
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// getCurrentTimestamp returns the current timestamp in nanoseconds
							 | 
						|
								func getCurrentTimestamp() int64 {
							 | 
						|
									return time.Now().UnixNano()
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// CreateDatabase creates and initializes a new offset storage database
							 | 
						|
								func CreateDatabase(dbPath string) (*sql.DB, error) {
							 | 
						|
									// TODO: Support different database types (PostgreSQL, MySQL, etc.)
							 | 
						|
									// ASSUMPTION: Using SQLite for now, can be extended for other databases
							 | 
						|
									
							 | 
						|
									db, err := sql.Open("sqlite3", dbPath)
							 | 
						|
									if err != nil {
							 | 
						|
										return nil, fmt.Errorf("failed to open database: %w", err)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									// Configure SQLite for better performance
							 | 
						|
									pragmas := []string{
							 | 
						|
										"PRAGMA journal_mode=WAL",           // Write-Ahead Logging for better concurrency
							 | 
						|
										"PRAGMA synchronous=NORMAL",         // Balance between safety and performance
							 | 
						|
										"PRAGMA cache_size=10000",           // Increase cache size
							 | 
						|
										"PRAGMA foreign_keys=ON",            // Enable foreign key constraints
							 | 
						|
										"PRAGMA temp_store=MEMORY",          // Store temporary tables in memory
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									for _, pragma := range pragmas {
							 | 
						|
										_, err := db.Exec(pragma)
							 | 
						|
										if err != nil {
							 | 
						|
											db.Close()
							 | 
						|
											return nil, fmt.Errorf("failed to set pragma %s: %w", pragma, err)
							 | 
						|
										}
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									// Apply migrations
							 | 
						|
									migrationManager := NewMigrationManager(db)
							 | 
						|
									err = migrationManager.ApplyMigrations()
							 | 
						|
									if err != nil {
							 | 
						|
										db.Close()
							 | 
						|
										return nil, fmt.Errorf("failed to apply migrations: %w", err)
							 | 
						|
									}
							 | 
						|
									
							 | 
						|
									return db, nil
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// BackupDatabase creates a backup of the offset storage database
							 | 
						|
								func BackupDatabase(sourceDB *sql.DB, backupPath string) error {
							 | 
						|
									// TODO: Implement database backup functionality
							 | 
						|
									// ASSUMPTION: This would use database-specific backup mechanisms
							 | 
						|
									return fmt.Errorf("database backup not implemented yet")
							 | 
						|
								}
							 | 
						|
								
							 | 
						|
								// RestoreDatabase restores a database from a backup
							 | 
						|
								func RestoreDatabase(backupPath, targetPath string) error {
							 | 
						|
									// TODO: Implement database restore functionality
							 | 
						|
									// ASSUMPTION: This would use database-specific restore mechanisms
							 | 
						|
									return fmt.Errorf("database restore not implemented yet")
							 | 
						|
								}
							 |