mirror of
https://github.com/MHSanaei/3x-ui.git
synced 2026-09-28 04:36:40 +08:00
n0ctal
66df77665f
With XUI_DB_TYPE=postgres every test package shared one database and worked in public. Go runs package test binaries concurrently, so migrations raced and rows a previous run left behind leaked into the next. testpg.IsolatePackage creates a schema for the calling package, puts it first on search_path and drops it when the package finishes. It returns at once unless XUI_DB_TYPE is postgres. internal/web/service's TestMain adopts it.
96 lines
2.5 KiB
Go
96 lines
2.5 KiB
Go
package testpg
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"fmt"
|
|
"net/url"
|
|
"os"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
const (
|
|
dbTypeEnv = "XUI_DB_TYPE"
|
|
dbDSNEnv = "XUI_DB_DSN"
|
|
)
|
|
|
|
// IsolatePackage gives one test package its own PostgreSQL schema: package test
|
|
// binaries run concurrently, and sharing public lets their migrations race.
|
|
func IsolatePackage(packageName string) (func(), error) {
|
|
if os.Getenv(dbTypeEnv) != "postgres" {
|
|
return func() {}, nil
|
|
}
|
|
baseDSN := strings.TrimSpace(os.Getenv(dbDSNEnv))
|
|
if baseDSN == "" {
|
|
return func() {}, nil
|
|
}
|
|
|
|
suffix := make([]byte, 8)
|
|
if _, err := rand.Read(suffix); err != nil {
|
|
return nil, fmt.Errorf("generate PostgreSQL test schema suffix: %w", err)
|
|
}
|
|
schema := fmt.Sprintf("xui_%s_%d_%s", sanitize(packageName), os.Getpid(), hex.EncodeToString(suffix))
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
admin, err := pgxpool.New(ctx, baseDSN)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("open PostgreSQL test database: %w", err)
|
|
}
|
|
if _, err := admin.Exec(ctx, "CREATE SCHEMA "+pgx.Identifier{schema}.Sanitize()); err != nil {
|
|
admin.Close()
|
|
return nil, fmt.Errorf("create PostgreSQL test schema: %w", err)
|
|
}
|
|
isolatedDSN, err := withSearchPath(baseDSN, schema)
|
|
if err != nil {
|
|
admin.Close()
|
|
return nil, err
|
|
}
|
|
if err := os.Setenv(dbDSNEnv, isolatedDSN); err != nil {
|
|
admin.Close()
|
|
return nil, fmt.Errorf("set isolated PostgreSQL test DSN: %w", err)
|
|
}
|
|
|
|
return func() {
|
|
cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cleanupCancel()
|
|
_, _ = admin.Exec(cleanupCtx, "DROP SCHEMA "+pgx.Identifier{schema}.Sanitize()+" CASCADE")
|
|
admin.Close()
|
|
_ = os.Setenv(dbDSNEnv, baseDSN)
|
|
}, nil
|
|
}
|
|
|
|
func withSearchPath(dsn, schema string) (string, error) {
|
|
u, err := url.Parse(dsn)
|
|
if err == nil && (u.Scheme == "postgres" || u.Scheme == "postgresql") {
|
|
query := u.Query()
|
|
query.Set("search_path", schema)
|
|
u.RawQuery = query.Encode()
|
|
return u.String(), nil
|
|
}
|
|
if strings.ContainsAny(schema, " '[]=\\") {
|
|
return "", fmt.Errorf("unsafe PostgreSQL test schema name")
|
|
}
|
|
return strings.TrimSpace(dsn) + " search_path=" + schema, nil
|
|
}
|
|
|
|
func sanitize(value string) string {
|
|
var result strings.Builder
|
|
for _, r := range strings.ToLower(value) {
|
|
if r >= 'a' && r <= 'z' || r >= '0' && r <= '9' || r == '_' {
|
|
result.WriteRune(r)
|
|
} else {
|
|
result.WriteByte('_')
|
|
}
|
|
}
|
|
if result.Len() == 0 {
|
|
return "pkg"
|
|
}
|
|
return result.String()
|
|
}
|