package dbmanager import ( "context" "database/sql" "sync" "time" "github.com/uptrace/bun" "github.com/uptrace/bun/schema" "go.mongodb.org/mongo-driver/mongo" "gorm.io/gorm" "github.com/bitechdev/ResolveSpec/pkg/common" "github.com/bitechdev/ResolveSpec/pkg/common/adapters/database" ) // Connection represents a single named database connection type Connection interface { // Metadata Name() string Type() DatabaseType // ORM Access (SQL databases only) Bun() (*bun.DB, error) GORM() (*gorm.DB, error) Native() (*sql.DB, error) // Common Database interface (for SQL databases) Database() (common.Database, error) // MongoDB Access (MongoDB only) MongoDB() (*mongo.Client, error) // Lifecycle Connect(ctx context.Context) error Close() error HealthCheck(ctx context.Context) error Reconnect(ctx context.Context) error // Stats Stats() *ConnectionStats } // ConnectionStats contains statistics about a database connection type ConnectionStats struct { Name string Type DatabaseType Connected bool LastHealthCheck time.Time HealthCheckStatus string // SQL connection pool stats OpenConnections int InUse int Idle int WaitCount int64 WaitDuration time.Duration MaxIdleClosed int64 MaxLifetimeClosed int64 } // sqlConnection implements Connection for SQL databases (PostgreSQL, SQLite, MSSQL) type sqlConnection struct { name string dbType DatabaseType config ConnectionConfig provider Provider // Lazy-initialized ORM instances (all wrap the same sql.DB) nativeDB *sql.DB bunDB *bun.DB gormDB *gorm.DB // Adapters for common.Database interface bunAdapter *database.BunAdapter gormAdapter *database.GormAdapter nativeAdapter common.Database // State connected bool mu sync.RWMutex // Health check lastHealthCheck time.Time healthCheckStatus string } // newSQLConnection creates a new SQL connection func newSQLConnection(name string, dbType DatabaseType, config ConnectionConfig, provider Provider) *sqlConnection { return &sqlConnection{ name: name, dbType: dbType, config: config, provider: provider, } } // Name returns the connection name func (c *sqlConnection) Name() string { return c.name } // Type returns the database type func (c *sqlConnection) Type() DatabaseType { return c.dbType } // Connect establishes the database connection func (c *sqlConnection) Connect(ctx context.Context) error { c.mu.Lock() defer c.mu.Unlock() if c.connected { return ErrAlreadyConnected } if err := c.provider.Connect(ctx, &c.config); err != nil { return NewConnectionError(c.name, "connect", err) } c.connected = true return nil } // Close closes the database connection and all ORM instances func (c *sqlConnection) Close() error { c.mu.Lock() defer c.mu.Unlock() if !c.connected { return nil } // Close Bun if initialized if c.bunDB != nil { if err := c.bunDB.Close(); err != nil { return NewConnectionError(c.name, "close bun", err) } } // GORM doesn't have a separate close - it uses the underlying sql.DB // Close the provider (which closes the underlying sql.DB) if err := c.provider.Close(); err != nil { return NewConnectionError(c.name, "close", err) } c.connected = false c.nativeDB = nil c.bunDB = nil c.gormDB = nil c.bunAdapter = nil c.gormAdapter = nil c.nativeAdapter = nil return nil } // HealthCheck verifies the connection is alive func (c *sqlConnection) HealthCheck(ctx context.Context) error { c.mu.Lock() defer c.mu.Unlock() c.lastHealthCheck = time.Now() if !c.connected { c.healthCheckStatus = "disconnected" return ErrConnectionClosed } if err := c.provider.HealthCheck(ctx); err != nil { c.healthCheckStatus = "unhealthy: " + err.Error() return NewConnectionError(c.name, "health check", err) } c.healthCheckStatus = "healthy" return nil } // Reconnect closes and re-establishes the connection func (c *sqlConnection) Reconnect(ctx context.Context) error { if err := c.Close(); err != nil { return err } return c.Connect(ctx) } // Native returns the native *sql.DB connection func (c *sqlConnection) Native() (*sql.DB, error) { c.mu.RLock() if c.nativeDB != nil { defer c.mu.RUnlock() return c.nativeDB, nil } c.mu.RUnlock() c.mu.Lock() defer c.mu.Unlock() // Double-check after acquiring write lock if c.nativeDB != nil { return c.nativeDB, nil } if !c.connected { return nil, ErrConnectionClosed } // Get native connection from provider db, err := c.provider.GetNative() if err != nil { return nil, NewConnectionError(c.name, "get native", err) } c.nativeDB = db return c.nativeDB, nil } // Bun returns a Bun ORM instance wrapping the native connection func (c *sqlConnection) Bun() (*bun.DB, error) { c.mu.RLock() if c.bunDB != nil { defer c.mu.RUnlock() return c.bunDB, nil } c.mu.RUnlock() c.mu.Lock() defer c.mu.Unlock() // Double-check after acquiring write lock if c.bunDB != nil { return c.bunDB, nil } // Get native connection first native, err := c.provider.GetNative() if err != nil { return nil, NewConnectionError(c.name, "get bun", err) } // Create Bun DB wrapping the same sql.DB dialect := c.getBunDialect() c.bunDB = bun.NewDB(native, dialect) return c.bunDB, nil } // GORM returns a GORM instance wrapping the native connection func (c *sqlConnection) GORM() (*gorm.DB, error) { c.mu.RLock() if c.gormDB != nil { defer c.mu.RUnlock() return c.gormDB, nil } c.mu.RUnlock() c.mu.Lock() defer c.mu.Unlock() // Double-check after acquiring write lock if c.gormDB != nil { return c.gormDB, nil } // Get native connection first native, err := c.provider.GetNative() if err != nil { return nil, NewConnectionError(c.name, "get gorm", err) } // Create GORM DB wrapping the same sql.DB dialector := c.getGORMDialector(native) db, err := gorm.Open(dialector, &gorm.Config{}) if err != nil { return nil, NewConnectionError(c.name, "initialize gorm", err) } c.gormDB = db return c.gormDB, nil } // Database returns the common.Database interface using the configured default ORM func (c *sqlConnection) Database() (common.Database, error) { c.mu.RLock() defaultORM := c.config.DefaultORM c.mu.RUnlock() switch ORMType(defaultORM) { case ORMTypeBun: return c.getBunAdapter() case ORMTypeGORM: return c.getGORMAdapter() case ORMTypeNative: return c.getNativeAdapter() default: // Default to Bun return c.getBunAdapter() } } // MongoDB returns an error for SQL connections func (c *sqlConnection) MongoDB() (*mongo.Client, error) { return nil, ErrNotMongoDB } // Stats returns connection statistics func (c *sqlConnection) Stats() *ConnectionStats { c.mu.RLock() defer c.mu.RUnlock() stats := &ConnectionStats{ Name: c.name, Type: c.dbType, Connected: c.connected, LastHealthCheck: c.lastHealthCheck, HealthCheckStatus: c.healthCheckStatus, } // Get SQL stats if connected if c.connected && c.provider != nil { if providerStats := c.provider.Stats(); providerStats != nil { stats.OpenConnections = providerStats.OpenConnections stats.InUse = providerStats.InUse stats.Idle = providerStats.Idle stats.WaitCount = providerStats.WaitCount stats.WaitDuration = providerStats.WaitDuration stats.MaxIdleClosed = providerStats.MaxIdleClosed stats.MaxLifetimeClosed = providerStats.MaxLifetimeClosed } } return stats } // getBunAdapter returns or creates the Bun adapter func (c *sqlConnection) getBunAdapter() (common.Database, error) { c.mu.RLock() if c.bunAdapter != nil { defer c.mu.RUnlock() return c.bunAdapter, nil } c.mu.RUnlock() c.mu.Lock() defer c.mu.Unlock() if c.bunAdapter != nil { return c.bunAdapter, nil } bunDB, err := c.Bun() if err != nil { return nil, err } c.bunAdapter = database.NewBunAdapter(bunDB) return c.bunAdapter, nil } // getGORMAdapter returns or creates the GORM adapter func (c *sqlConnection) getGORMAdapter() (common.Database, error) { c.mu.RLock() if c.gormAdapter != nil { defer c.mu.RUnlock() return c.gormAdapter, nil } c.mu.RUnlock() c.mu.Lock() defer c.mu.Unlock() if c.gormAdapter != nil { return c.gormAdapter, nil } gormDB, err := c.GORM() if err != nil { return nil, err } c.gormAdapter = database.NewGormAdapter(gormDB) return c.gormAdapter, nil } // getNativeAdapter returns or creates the native adapter func (c *sqlConnection) getNativeAdapter() (common.Database, error) { c.mu.RLock() if c.nativeAdapter != nil { defer c.mu.RUnlock() return c.nativeAdapter, nil } c.mu.RUnlock() c.mu.Lock() defer c.mu.Unlock() if c.nativeAdapter != nil { return c.nativeAdapter, nil } native, err := c.Native() if err != nil { return nil, err } // Create a native adapter based on database type switch c.dbType { case DatabaseTypePostgreSQL: c.nativeAdapter = database.NewPgSQLAdapter(native) case DatabaseTypeSQLite: // For SQLite, we'll use the PgSQL adapter as it works with standard sql.DB c.nativeAdapter = database.NewPgSQLAdapter(native) case DatabaseTypeMSSQL: // For MSSQL, we'll use the PgSQL adapter as it works with standard sql.DB c.nativeAdapter = database.NewPgSQLAdapter(native) default: return nil, ErrUnsupportedDatabase } return c.nativeAdapter, nil } // getBunDialect returns the appropriate Bun dialect for the database type func (c *sqlConnection) getBunDialect() schema.Dialect { switch c.dbType { case DatabaseTypePostgreSQL: return database.GetPostgresDialect() case DatabaseTypeSQLite: return database.GetSQLiteDialect() case DatabaseTypeMSSQL: return database.GetMSSQLDialect() default: // Default to PostgreSQL return database.GetPostgresDialect() } } // getGORMDialector returns the appropriate GORM dialector for the database type func (c *sqlConnection) getGORMDialector(db *sql.DB) gorm.Dialector { switch c.dbType { case DatabaseTypePostgreSQL: return database.GetPostgresDialector(db) case DatabaseTypeSQLite: return database.GetSQLiteDialector(db) case DatabaseTypeMSSQL: return database.GetMSSQLDialector(db) default: // Default to PostgreSQL return database.GetPostgresDialector(db) } } // mongoConnection implements Connection for MongoDB type mongoConnection struct { name string config ConnectionConfig provider Provider // MongoDB client client *mongo.Client // State connected bool mu sync.RWMutex // Health check lastHealthCheck time.Time healthCheckStatus string } // newMongoConnection creates a new MongoDB connection func newMongoConnection(name string, config ConnectionConfig, provider Provider) *mongoConnection { return &mongoConnection{ name: name, config: config, provider: provider, } } // Name returns the connection name func (c *mongoConnection) Name() string { return c.name } // Type returns the database type (MongoDB) func (c *mongoConnection) Type() DatabaseType { return DatabaseTypeMongoDB } // Connect establishes the MongoDB connection func (c *mongoConnection) Connect(ctx context.Context) error { c.mu.Lock() defer c.mu.Unlock() if c.connected { return ErrAlreadyConnected } if err := c.provider.Connect(ctx, &c.config); err != nil { return NewConnectionError(c.name, "connect", err) } // Get the mongo client client, err := c.provider.GetMongo() if err != nil { return NewConnectionError(c.name, "get mongo client", err) } c.client = client c.connected = true return nil } // Close closes the MongoDB connection func (c *mongoConnection) Close() error { c.mu.Lock() defer c.mu.Unlock() if !c.connected { return nil } if err := c.provider.Close(); err != nil { return NewConnectionError(c.name, "close", err) } c.connected = false c.client = nil return nil } // HealthCheck verifies the MongoDB connection is alive func (c *mongoConnection) HealthCheck(ctx context.Context) error { c.mu.Lock() defer c.mu.Unlock() c.lastHealthCheck = time.Now() if !c.connected { c.healthCheckStatus = "disconnected" return ErrConnectionClosed } if err := c.provider.HealthCheck(ctx); err != nil { c.healthCheckStatus = "unhealthy: " + err.Error() return NewConnectionError(c.name, "health check", err) } c.healthCheckStatus = "healthy" return nil } // Reconnect closes and re-establishes the MongoDB connection func (c *mongoConnection) Reconnect(ctx context.Context) error { if err := c.Close(); err != nil { return err } return c.Connect(ctx) } // MongoDB returns the MongoDB client func (c *mongoConnection) MongoDB() (*mongo.Client, error) { c.mu.RLock() defer c.mu.RUnlock() if !c.connected || c.client == nil { return nil, ErrConnectionClosed } return c.client, nil } // Bun returns an error for MongoDB connections func (c *mongoConnection) Bun() (*bun.DB, error) { return nil, ErrNotSQLDatabase } // GORM returns an error for MongoDB connections func (c *mongoConnection) GORM() (*gorm.DB, error) { return nil, ErrNotSQLDatabase } // Native returns an error for MongoDB connections func (c *mongoConnection) Native() (*sql.DB, error) { return nil, ErrNotSQLDatabase } // Database returns an error for MongoDB connections func (c *mongoConnection) Database() (common.Database, error) { return nil, ErrNotSQLDatabase } // Stats returns connection statistics for MongoDB func (c *mongoConnection) Stats() *ConnectionStats { c.mu.RLock() defer c.mu.RUnlock() return &ConnectionStats{ Name: c.name, Type: DatabaseTypeMongoDB, Connected: c.connected, LastHealthCheck: c.lastHealthCheck, HealthCheckStatus: c.healthCheckStatus, } }