|
@@ -0,0 +1,95 @@
|
|
|
|
|
+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()
|
|
|
|
|
+}
|