uzdb/internal/database/manager.go

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
}