fetch_ml/internal/api/factory.go
Jeremie Fraeys 23e5f3d1dc
refactor(api): internal refactoring for TUI and worker modules
- Refactor internal/worker and internal/queue packages
- Update cmd/tui for monitoring interface
- Update test configurations
2026-02-20 15:51:23 -05:00

294 lines
7.9 KiB
Go

package api
import (
"context"
"os"
"strings"
apimiddleware "github.com/jfraeys/fetch_ml/internal/api/middleware"
"github.com/jfraeys/fetch_ml/internal/audit"
"github.com/jfraeys/fetch_ml/internal/config"
"github.com/jfraeys/fetch_ml/internal/experiment"
"github.com/jfraeys/fetch_ml/internal/jupyter"
"github.com/jfraeys/fetch_ml/internal/logging"
"github.com/jfraeys/fetch_ml/internal/middleware"
"github.com/jfraeys/fetch_ml/internal/queue"
"github.com/jfraeys/fetch_ml/internal/storage"
)
// initializeComponents initializes all server components
func (s *Server) initializeComponents() error {
// Setup logging
if err := s.config.EnsureLogDirectory(); err != nil {
return err
}
s.logger = s.setupLogger()
// Initialize experiment manager
if err := s.initExperimentManager(); err != nil {
return err
}
// Initialize task queue
if err := s.initTaskQueue(); err != nil {
return err
}
// Initialize database
if err := s.initDatabase(); err != nil {
return err
}
// Initialize database schema (if DB enabled)
if err := s.initDatabaseSchema(); err != nil {
return err
}
// Initialize security
s.initSecurity()
// Initialize Jupyter service manager
s.initJupyterServiceManager()
// Initialize OpenAPI validation middleware (optional, doesn't block startup)
s.initValidationMiddleware()
return nil
}
// setupLogger creates and configures the logger
func (s *Server) setupLogger() *logging.Logger {
logger := logging.NewLoggerFromConfig(s.config.Logging)
ctx := logging.EnsureTrace(context.Background())
return logger.Component(ctx, "api-server")
}
// initExperimentManager initializes the experiment manager
func (s *Server) initExperimentManager() error {
s.expManager = experiment.NewManager(s.config.BasePath)
if err := s.expManager.Initialize(); err != nil {
return err
}
s.logger.Info("experiment manager initialized", "base_path", s.config.BasePath)
return nil
}
// initTaskQueue initializes the task queue
func (s *Server) initTaskQueue() error {
backend := strings.ToLower(strings.TrimSpace(s.config.Queue.Backend))
if backend == "" {
backend = "redis"
}
redisAddr := strings.TrimSpace(s.config.Redis.Addr)
if redisAddr == "" {
redisAddr = "localhost:6379"
}
if strings.TrimSpace(s.config.Redis.URL) != "" {
redisAddr = strings.TrimSpace(s.config.Redis.URL)
}
backendCfg := queue.BackendConfig{
Backend: queue.QueueBackend(backend),
RedisAddr: redisAddr,
RedisPassword: s.config.Redis.Password,
RedisDB: s.config.Redis.DB,
SQLitePath: s.config.Queue.SQLitePath,
FilesystemPath: s.config.Queue.FilesystemPath,
FallbackToFilesystem: s.config.Queue.FallbackToFilesystem,
MetricsFlushInterval: 0,
}
taskQueue, err := queue.NewBackend(backendCfg)
if err != nil {
return err
}
s.taskQueue = taskQueue
if backend == "sqlite" {
s.logger.Info("task queue initialized", "backend", backend, "sqlite_path", s.config.Queue.SQLitePath)
} else {
s.logger.Info("task queue initialized", "backend", backend, "redis_addr", redisAddr)
}
// Add cleanup function
s.cleanupFuncs = append(s.cleanupFuncs, func() {
s.logger.Info("stopping task queue...")
if err := s.taskQueue.Close(); err != nil {
s.logger.Error("failed to stop task queue", "error", err)
} else {
s.logger.Info("task queue stopped")
}
})
return nil
}
// initDatabase initializes the database connection
func (s *Server) initDatabase() error {
if s.config.Database.Type == "" {
return nil
}
dbConfig := storage.DBConfig{
Type: s.config.Database.Type,
Connection: s.config.Database.Connection,
Host: s.config.Database.Host,
Port: s.config.Database.Port,
Username: s.config.Database.Username,
Password: s.config.Database.Password,
Database: s.config.Database.Database,
}
db, err := storage.NewDB(dbConfig)
if err != nil {
return err
}
s.db = db
s.logger.Info("database initialized", "type", s.config.Database.Type)
// Clear database password from memory after connection
if s.config.Database.Password != "" {
s.config.Database.Password = "[CLEARED]"
}
// Add cleanup function
s.cleanupFuncs = append(s.cleanupFuncs, func() {
s.logger.Info("closing database connection...")
if err := s.db.Close(); err != nil {
s.logger.Error("failed to close database", "error", err)
} else {
s.logger.Info("database connection closed")
}
})
return nil
}
// initDatabaseSchema initializes the database schema
func (s *Server) initDatabaseSchema() error {
if s.db == nil {
return nil
}
schema, err := storage.SchemaForDBType(s.config.Database.Type)
if err != nil {
return err
}
if err := s.db.Initialize(schema); err != nil {
return err
}
s.logger.Info("database schema initialized", "type", s.config.Database.Type)
return nil
}
// initSecurity initializes security middleware
func (s *Server) initSecurity() {
authConfig := s.config.BuildAuthConfig()
rlOpts := s.buildRateLimitOptions()
// Read JWT secret securely and clear from environment
jwtSecret := SecureEnvVar("JWT_SECRET")
s.sec = middleware.NewSecurityMiddleware(authConfig, jwtSecret, rlOpts)
}
// buildRateLimitOptions builds rate limit options from configuration
func (s *Server) buildRateLimitOptions() *middleware.RateLimitOptions {
if !s.config.Security.RateLimit.Enabled || s.config.Security.RateLimit.RequestsPerMinute <= 0 {
return nil
}
return &middleware.RateLimitOptions{
RequestsPerMinute: s.config.Security.RateLimit.RequestsPerMinute,
BurstSize: s.config.Security.RateLimit.BurstSize,
}
}
// initJupyterServiceManager initializes the Jupyter service manager
func (s *Server) initJupyterServiceManager() {
serviceConfig := jupyter.GetDefaultServiceConfig()
sm, err := jupyter.NewServiceManager(s.logger, serviceConfig)
if err != nil {
s.logger.Error("failed to initialize Jupyter service manager", "error", err)
return
}
s.jupyterServiceMgr = sm
s.logger.Info("jupyter service manager initialized")
}
// initAuditLogger initializes the audit logger
func (s *Server) initAuditLogger() *audit.Logger {
if !s.config.Security.AuditLogging.Enabled || s.config.Security.AuditLogging.LogPath == "" {
return nil
}
al, err := audit.NewLogger(
s.config.Security.AuditLogging.Enabled,
s.config.Security.AuditLogging.LogPath,
s.logger,
)
if err != nil {
s.logger.Warn("failed to initialize audit logger", "error", err)
return nil
}
s.auditLogger = al
// Add cleanup function
s.cleanupFuncs = append(s.cleanupFuncs, func() {
s.logger.Info("closing audit logger...")
if err := al.Close(); err != nil {
s.logger.Error("failed to close audit logger", "error", err)
}
})
return al
}
// initValidationMiddleware initializes the OpenAPI validation middleware
func (s *Server) initValidationMiddleware() {
// Only initialize if OpenAPI spec exists
if _, err := os.Stat("api/openapi.yaml"); err != nil {
s.logger.Debug("OpenAPI spec not found, skipping validation middleware")
return
}
vm, err := apimiddleware.NewValidationMiddleware("api/openapi.yaml")
if err != nil {
s.logger.Warn("failed to initialize OpenAPI validation middleware", "error", err)
return
}
s.validationMiddleware = vm
s.logger.Info("OpenAPI validation middleware initialized")
}
// getSecurityConfig extracts security config from server config
func getSecurityConfig(cfg *ServerConfig) *config.SecurityConfig {
return &config.SecurityConfig{
AllowedOrigins: cfg.Security.AllowedOrigins,
ProductionMode: cfg.Security.ProductionMode,
APIKeyRotationDays: cfg.Security.APIKeyRotationDays,
AuditLogging: config.AuditLoggingConfig{
Enabled: cfg.Security.AuditLogging.Enabled,
LogPath: cfg.Security.AuditLogging.LogPath,
},
IPWhitelist: cfg.Security.IPWhitelist,
}
}
// SecureEnvVar reads an environment variable and clears it from the environment
// to prevent exposure via /proc/*/environ. Returns empty string if not set.
func SecureEnvVar(name string) string {
value := os.Getenv(name)
if value != "" {
os.Unsetenv(name)
}
return value
}