feat: Schema migration

Author: codefionn

msgtausch-srv/stats/postgres.go

@@ -45,142 +45,19 @@ func NewPostgreSQLCollector(connectionString string) (*PostgreSQLCollector, erro
 
 // initSchema creates the necessary tables if they don't exist
 func (p *PostgreSQLCollector) initSchema() error {
-	queries := []string{
-		`CREATE TABLE IF NOT EXISTS connections (
-			id SERIAL PRIMARY KEY,
-			connection_uuid TEXT,
-			client_ip INET NOT NULL,
-			target_host TEXT NOT NULL,
-			target_port INTEGER NOT NULL,
-			protocol TEXT NOT NULL,
-			started_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
-			ended_at TIMESTAMP WITH TIME ZONE,
-			bytes_sent BIGINT DEFAULT 0,
-			bytes_received BIGINT DEFAULT 0,
-			duration_ms INTEGER,
-			close_reason TEXT
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_connections_client_ip ON connections(client_ip)`,
-		`CREATE INDEX IF NOT EXISTS idx_connections_target_host ON connections(target_host)`,
-		`CREATE INDEX IF NOT EXISTS idx_connections_started_at ON connections(started_at)`,
-		`CREATE INDEX IF NOT EXISTS idx_connections_uuid ON connections(connection_uuid)`,
-
-		`CREATE TABLE IF NOT EXISTS http_requests (
-			id SERIAL PRIMARY KEY,
-			connection_id INTEGER NOT NULL REFERENCES connections(id) ON DELETE CASCADE,
-			method TEXT NOT NULL,
-			url TEXT NOT NULL,
-			host TEXT NOT NULL,
-			user_agent TEXT,
-			content_length BIGINT DEFAULT 0,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_http_requests_connection_id ON http_requests(connection_id)`,
-		`CREATE INDEX IF NOT EXISTS idx_http_requests_timestamp ON http_requests(timestamp)`,
-
-		`CREATE TABLE IF NOT EXISTS http_responses (
-			id SERIAL PRIMARY KEY,
-			connection_id INTEGER NOT NULL REFERENCES connections(id) ON DELETE CASCADE,
-			request_id INTEGER REFERENCES http_requests(id) ON DELETE CASCADE,
-			status_code INTEGER NOT NULL,
-			content_length BIGINT DEFAULT 0,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_http_responses_connection_id ON http_responses(connection_id)`,
-
-		`CREATE TABLE IF NOT EXISTS errors (
-			id SERIAL PRIMARY KEY,
-			connection_id INTEGER REFERENCES connections(id) ON DELETE CASCADE,
-			error_type TEXT NOT NULL,
-			error_message TEXT NOT NULL,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_errors_timestamp ON errors(timestamp)`,
-
-		`CREATE TABLE IF NOT EXISTS security_events (
-			id SERIAL PRIMARY KEY,
-			client_ip INET NOT NULL,
-			target_host TEXT NOT NULL,
-			event_type TEXT NOT NULL,
-			reason TEXT,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_security_events_client_ip ON security_events(client_ip)`,
-		`CREATE INDEX IF NOT EXISTS idx_security_events_timestamp ON security_events(timestamp)`,
-
-		`CREATE TABLE IF NOT EXISTS data_transfers (
-			id SERIAL PRIMARY KEY,
-			connection_id INTEGER NOT NULL REFERENCES connections(id) ON DELETE CASCADE,
-			bytes_sent BIGINT DEFAULT 0,
-			bytes_received BIGINT DEFAULT 0,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_data_transfers_connection_id ON data_transfers(connection_id)`,
-
-		`CREATE TABLE IF NOT EXISTS recorded_http_requests (
-			id SERIAL PRIMARY KEY,
-			connection_id INTEGER NOT NULL REFERENCES connections(id) ON DELETE CASCADE,
-			method TEXT NOT NULL,
-			url TEXT NOT NULL,
-			host TEXT NOT NULL,
-			user_agent TEXT,
-			request_headers JSONB, -- JSON encoded headers
-			request_body BYTEA,
-			request_body_size INTEGER DEFAULT 0,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_recorded_http_requests_connection_id ON recorded_http_requests(connection_id)`,
-		`CREATE INDEX IF NOT EXISTS idx_recorded_http_requests_timestamp ON recorded_http_requests(timestamp)`,
-
-		`CREATE TABLE IF NOT EXISTS recorded_http_responses (
-			id SERIAL PRIMARY KEY,
-			connection_id INTEGER NOT NULL REFERENCES connections(id) ON DELETE CASCADE,
-			status_code INTEGER NOT NULL,
-			response_headers JSONB, -- JSON encoded headers
-			response_body BYTEA,
-			response_body_size INTEGER DEFAULT 0,
-			timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-		)`,
-		`CREATE INDEX IF NOT EXISTS idx_recorded_http_responses_connection_id ON recorded_http_responses(connection_id)`,
-		`CREATE INDEX IF NOT EXISTS idx_recorded_http_responses_timestamp ON recorded_http_responses(timestamp)`,
-
-		// New: request body parts table for streaming
-		`CREATE TABLE IF NOT EXISTS recorded_http_request_body_parts (
-            id SERIAL PRIMARY KEY,
-            request_id INTEGER NOT NULL REFERENCES recorded_http_requests(id) ON DELETE CASCADE,
-            seq_no INTEGER NOT NULL,
-            data BYTEA,
-            part_size INTEGER DEFAULT 0,
-            timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-        )`,
-		`CREATE INDEX IF NOT EXISTS idx_req_body_parts_request_id ON recorded_http_request_body_parts(request_id)`,
-		`CREATE INDEX IF NOT EXISTS idx_req_body_parts_req_seq ON recorded_http_request_body_parts(request_id, seq_no)`,
-
-		// New: response body parts table for streaming
-		`CREATE TABLE IF NOT EXISTS recorded_http_response_body_parts (
-            id SERIAL PRIMARY KEY,
-            response_id INTEGER NOT NULL REFERENCES recorded_http_responses(id) ON DELETE CASCADE,
-            seq_no INTEGER NOT NULL,
-            data BYTEA,
-            part_size INTEGER DEFAULT 0,
-            timestamp TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW()
-        )`,
-		`CREATE INDEX IF NOT EXISTS idx_resp_body_parts_response_id ON recorded_http_response_body_parts(response_id)`,
-		`CREATE INDEX IF NOT EXISTS idx_resp_body_parts_resp_seq ON recorded_http_response_body_parts(response_id, seq_no)`,
-	}
-
-	for _, query := range queries {
-		if _, err := p.db.Exec(query); err != nil {
-			return fmt.Errorf("failed to execute schema query: %w", err)
-		}
-	}
+	logger.Debug("Initializing PostgreSQL schema using schema-driven approach")
+	initializer := NewSchemaInitializer(p.db, "postgres")
 
-	// Add migration for existing tables to add connection_uuid column
-	_, _ = p.db.Exec(`ALTER TABLE connections ADD COLUMN connection_uuid TEXT`)
+	if err := initializer.ValidateAndInitialize(); err != nil {
+		return fmt.Errorf("schema initialization failed: %w", err)
+	}
 
+	logger.Info("Schema initialization completed using new schema-driven approach")
 	return nil
 }
 
+
+
 // StartConnection records the start of a connection (legacy method for backward compatibility)
 func (p *PostgreSQLCollector) StartConnection(ctx context.Context, clientIP, targetHost string, targetPort int, protocol string) (int64, error) {
 	var id int64