// Copyright 2026 Arctel.net // SPDX-License-Identifier: Apache-2.0 package repository import ( "Wavelet/plugins/domain/admin/errs" "Wavelet/plugins/domain/admin/model" "context" "database/sql" "errors" "fmt" "os" "os/exec" "strings" "sync" "time" ) const ( defaultSQLiteDBPath = "./data/wavelet.db" logDBNameSQLite = "sqlite" ) var ( dbConfigMu sync.RWMutex dbConfig = model.DatabaseConfig{ SQLitePath: defaultSQLiteDBPath, } ) // SetDBConfig sets the database configuration. func SetDBConfig(cfg model.DatabaseConfig) { dbConfigMu.Lock() defer dbConfigMu.Unlock() dbConfig = cfg } // GetDBConfig gets the database configuration. func GetDBConfig() model.DatabaseConfig { dbConfigMu.RLock() defer dbConfigMu.RUnlock() return dbConfig } // sqliteDatabasePath resolves the effective SQLite file path from configuration. func sqliteDatabasePath() string { name := GetDBConfig().SQLitePath if name == "" { name = defaultSQLiteDBPath } return name } // QuoteTableName escapes a raw identifier for use inside a quoted SQL fragment. func QuoteTableName(table string) string { return `"` + strings.ReplaceAll(table, `"`, `""`) + `"` } // GetSQLiteOverview collects the SQLite runtime overview. func GetSQLiteOverview(ctx context.Context) (model.DBOverviewResponse, error) { gormDB := GetDB(ctx) if gormDB == nil { return model.DBOverviewResponse{}, errs.ErrDatabaseUninitialized } name := sqliteDatabasePath() var version string var ver string if err := gormDB.Raw("SELECT sqlite_version()").Scan(&ver).Error; err == nil { version = "SQLite " + ver } else { version = "SQLite" } var sizeStr string if fi, err := os.Stat(name); err == nil { size := fi.Size() if size < 0 { size = 0 } sizeStr = model.FormatBytes(uint64(size)) } else { sizeStr = "0 B" } var tableCount int64 if err := gormDB.Raw("SELECT count(*) FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'").Scan(&tableCount).Error; err != nil { tableCount = 0 } var connCount int64 if sqlDB, err := gormDB.DB(); err == nil { connCount = int64(sqlDB.Stats().OpenConnections) } else { connCount = 1 } return model.DBOverviewResponse{ Type: logDBNameSQLite, Version: version, Name: name, Size: sizeStr, TableCount: tableCount, Connections: connCount, }, nil } // GetPostgresOverview collects the PostgreSQL runtime overview. func GetPostgresOverview(ctx context.Context) (model.DBOverviewResponse, error) { gormDB := GetDB(ctx) if gormDB == nil { return model.DBOverviewResponse{}, errs.ErrDatabaseUninitialized } name := GetDBConfig().Database var version string var ver string if err := gormDB.Raw("SELECT version()").Scan(&ver).Error; err == nil { version = ver } else { version = "PostgreSQL" } var sizeStr string var sizeBytes sql.NullInt64 if err := gormDB.Raw("SELECT pg_database_size(current_database())").Scan(&sizeBytes).Error; err == nil && sizeBytes.Valid { size := sizeBytes.Int64 if size < 0 { size = 0 } sizeStr = model.FormatBytes(uint64(size)) } else { sizeStr = "0 B" } var tableCount int64 if err := gormDB.Raw("SELECT count(*) FROM information_schema.tables WHERE table_schema = current_schema()").Scan(&tableCount).Error; err != nil { tableCount = 0 } var connCount int64 var pgc sql.NullInt64 if err := gormDB.Raw("SELECT count(*) FROM pg_stat_activity WHERE datname = current_database()").Scan(&pgc).Error; err == nil && pgc.Valid { connCount = pgc.Int64 } else { if sqlDB, err := gormDB.DB(); err == nil { connCount = int64(sqlDB.Stats().OpenConnections) } else { connCount = 1 } } return model.DBOverviewResponse{ Type: "postgres", Version: version, Name: name, Size: sizeStr, TableCount: tableCount, Connections: connCount, }, nil } // ListDatabaseTableNames returns every user table of the active database. func ListDatabaseTableNames(ctx context.Context) ([]string, error) { gormDB := GetDB(ctx) if gormDB == nil { return nil, errs.ErrDatabaseUninitialized } var tables []string var err error if !GetDBConfig().Enabled { err = gormDB.Raw("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' ORDER BY name").Scan(&tables).Error } else { err = gormDB.Raw("SELECT table_name FROM information_schema.tables WHERE table_schema = current_schema() ORDER BY table_name").Scan(&tables).Error } if err != nil { return nil, err } return tables, nil } // CountDatabaseTableRows counts the rows of the quoted table. func CountDatabaseTableRows(ctx context.Context, quotedTable string) (int64, error) { gormDB := GetDB(ctx) if gormDB == nil { return 0, errs.ErrDatabaseUninitialized } var total int64 if err := gormDB.Raw("SELECT count(*) FROM " + quotedTable).Scan(&total).Error; err != nil { return 0, errs.NewInvalidInputError(err.Error()) } return total, nil } // QueryDatabaseTableRows loads one page of raw rows from the quoted table. func QueryDatabaseTableRows( ctx context.Context, quotedTable string, limit int, offset int, ) ([]string, []map[string]any, error) { gormDB := GetDB(ctx) if gormDB == nil { return nil, nil, errs.ErrDatabaseUninitialized } rows, err := gormDB.Raw("SELECT * FROM "+quotedTable+" LIMIT ? OFFSET ?", limit, offset).Rows() if err != nil { return nil, nil, errs.NewInvalidInputError(err.Error()) } defer func() { _ = rows.Close() }() cols, err := rows.Columns() if err != nil { return nil, nil, err } results, err := scanTableRows(rows, cols) if err != nil { return nil, nil, err } return cols, results, nil } // RunSelectSQL executes an arbitrary select-like statement. func RunSelectSQL(ctx context.Context, sqlStr string) ([]string, []map[string]any, error) { gormDB := GetDB(ctx) if gormDB == nil { return nil, nil, errs.ErrDatabaseUninitialized } rows, err := gormDB.Raw(sqlStr).Rows() if err != nil { return nil, nil, errs.NewInvalidInputError(err.Error()) } defer func() { _ = rows.Close() }() cols, err := rows.Columns() if err != nil { return nil, nil, err } results, err := scanTableRows(rows, cols) if err != nil { return nil, nil, err } return cols, results, nil } // RunMutationSQL executes a non-query statement and reports affected rows. func RunMutationSQL(ctx context.Context, sqlStr string) (int64, error) { gormDB := GetDB(ctx) if gormDB == nil { return 0, errs.ErrDatabaseUninitialized } tx := gormDB.Exec(sqlStr) if tx.Error != nil { return 0, errs.NewInvalidInputError(tx.Error.Error()) } return tx.RowsAffected, nil } // scanTableRows decodes every row of the result set into a column keyed map. func scanTableRows(rows *sql.Rows, cols []string) ([]map[string]any, error) { results := make([]map[string]any, 0) for rows.Next() { row, err := scanRowAsMap(rows, cols) if err != nil { return nil, err } results = append(results, row) } return results, nil } // scanRowAsMap decodes a single row, normalising driver byte slices to strings. func scanRowAsMap(rows *sql.Rows, cols []string) (map[string]any, error) { columns := make([]any, len(cols)) columnPointers := make([]any, len(cols)) for i := range columns { columnPointers[i] = &columns[i] } if err := rows.Scan(columnPointers...); err != nil { return nil, err } rowMap := make(map[string]any) for i, colName := range cols { val := columns[i] if b, ok := val.([]byte); ok { rowMap[colName] = string(b) continue } rowMap[colName] = val } return rowMap, nil } // GetSQLiteInfo collects the SQLite type/name/version triple. func GetSQLiteInfo(ctx context.Context) model.DatabaseInfoResponse { cfg := GetDBConfig() info := model.DatabaseInfoResponse{ Type: logDBNameSQLite, Name: cfg.SQLitePath, Version: "SQLite", } if info.Name == "" { info.Name = defaultSQLiteDBPath } gormDB := GetDB(ctx) if gormDB == nil { return info } var ver string if err := gormDB.Raw("SELECT sqlite_version()").Scan(&ver).Error; err == nil && ver != "" { info.Version = "SQLite " + ver } return info } // GetPostgresInfo collects the PostgreSQL type/name/version triple. func GetPostgresInfo(ctx context.Context) model.DatabaseInfoResponse { cfg := GetDBConfig() info := model.DatabaseInfoResponse{ Type: "postgres", Name: cfg.Database, Version: "PostgreSQL", } gormDB := GetDB(ctx) if gormDB == nil { return info } var ver string if err := gormDB.Raw("SELECT version()").Scan(&ver).Error; err == nil && ver != "" { info.Version = ver } return info } // OpenSQLiteExportFile opens the active SQLite database file together with its stat info. func OpenSQLiteExportFile() (*os.File, os.FileInfo, error) { // export db file path is trusted f, err := os.Open(sqliteDatabasePath()) if err != nil { return nil, nil, fmt.Errorf("%s: %w", errs.ErrOpenDatabaseFileFailed, err) } fi, err := f.Stat() if err != nil { _ = f.Close() return nil, nil, fmt.Errorf("%s: %w", errs.ErrReadDatabaseFileInfoFailed, err) } return f, fi, nil } // NewPgDumpCommand builds the streaming pg_dump command for the active database. func NewPgDumpCommand(ctx context.Context) (*exec.Cmd, string, error) { dbCfg := GetDBConfig() pgDumpPath, err := exec.LookPath("pg_dump") if err != nil { return nil, "", errors.New(errs.ErrPgDumpUnavailable) } args := []string{ "--no-password", "-h", dbCfg.Host, "-p", fmt.Sprintf("%d", dbCfg.Port), "-U", dbCfg.Username, dbCfg.Database, } //nolint:gosec // pg_dump args are constructed from validated db config cmd := exec.CommandContext(ctx, pgDumpPath, args...) if dbCfg.Password != "" { cmd.Env = append(os.Environ(), "PGPASSWORD="+dbCfg.Password) } else { cmd.Env = os.Environ() } fileName := fmt.Sprintf("wavelet_%s.sql", time.Now().Format("20060102_150405")) return cmd, fileName, nil }