package database import ( "database/sql" "fmt" "sync" "go.uber.org/zap" "uzdb/internal/config" "uzdb/internal/models" ) // ConnectionManager manages database connections for different database types type ConnectionManager struct { connections map[string]DatabaseConnector mu sync.RWMutex } // DatabaseConnector is the interface for all database connections type DatabaseConnector interface { GetDB() *sql.DB Close() error IsConnected() bool ExecuteQuery(sql string, args ...interface{}) (*models.QueryResult, error) ExecuteStatement(sql string, args ...interface{}) (*models.QueryResult, error) GetTables(schema string) ([]models.Table, error) GetTableStructure(tableName string) (*models.TableStructure, error) GetMetadata() (*models.DBMetadata, error) } // NewConnectionManager creates a new connection manager func NewConnectionManager() *ConnectionManager { return &ConnectionManager{ connections: make(map[string]DatabaseConnector), } } // GetConnection gets or creates a connection for a user connection config func (m *ConnectionManager) GetConnection(conn *models.UserConnection, password string) (DatabaseConnector, error) { m.mu.Lock() defer m.mu.Unlock() // Check if connection already exists if existing, ok := m.connections[conn.ID]; ok { if existing.IsConnected() { return existing, nil } // Connection is dead, remove it existing.Close() delete(m.connections, conn.ID) } // Create new connection based on type var connector DatabaseConnector var err error switch conn.Type { case models.ConnectionTypeMySQL: connector, err = NewMySQLConnection(conn, password) case models.ConnectionTypePostgreSQL: connector, err = NewPostgreSQLConnection(conn, password) case models.ConnectionTypeSQLite: connector, err = NewSQLiteConnection(conn) default: return nil, fmt.Errorf("unsupported database type: %s", conn.Type) } if err != nil { return nil, err } m.connections[conn.ID] = connector config.GetLogger().Info("connection created in manager", zap.String("connection_id", conn.ID), zap.String("type", string(conn.Type)), ) return connector, nil } // RemoveConnection removes a connection from the manager func (m *ConnectionManager) RemoveConnection(connectionID string) error { m.mu.Lock() defer m.mu.Unlock() if conn, ok := m.connections[connectionID]; ok { if err := conn.Close(); err != nil { return err } delete(m.connections, connectionID) config.GetLogger().Info("connection removed from manager", zap.String("connection_id", connectionID), ) } return nil } // GetAllConnections returns all managed connections func (m *ConnectionManager) GetAllConnections() map[string]DatabaseConnector { m.mu.RLock() defer m.mu.RUnlock() result := make(map[string]DatabaseConnector) for k, v := range m.connections { result[k] = v } return result } // CloseAll closes all managed connections func (m *ConnectionManager) CloseAll() error { m.mu.Lock() defer m.mu.Unlock() var lastErr error for id, conn := range m.connections { if err := conn.Close(); err != nil { lastErr = err config.GetLogger().Error("failed to close connection", zap.String("connection_id", id), zap.Error(err), ) } } m.connections = make(map[string]DatabaseConnector) if lastErr != nil { return lastErr } config.GetLogger().Info("all connections closed") return nil }