Add missing reset

Author: fangpenlin

e2e/packages/client/reset_db.go

@@ -0,0 +1,162 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"strings"
+
+	"github.com/jackc/pgx/v5"
+)
+
+// ServicePortProvider provides a way to get the mapped port for a service
+type ServicePortProvider interface {
+	GetServicePort(ctx context.Context, serviceName string, internalPort string) (string, error)
+}
+
+// DatabaseConfig holds database connection configuration
+type DatabaseConfig struct {
+	User     string
+	Password string
+	Database string
+	Host     string
+	Port     string
+}
+
+// DefaultDatabaseConfig returns the default database configuration
+func DefaultDatabaseConfig() DatabaseConfig {
+	return DatabaseConfig{
+		User:     "infisical",
+		Password: "infisical",
+		Database: "infisical",
+		Host:     "localhost",
+		Port:     "5432",
+	}
+}
+
+// ResetDBOptions holds options for resetting the database
+type ResetDBOptions struct {
+	SkipTables map[string]struct{} // Tables to skip when truncating (e.g., migrations)
+	DBConfig   DatabaseConfig
+}
+
+// DefaultResetDBOptions returns default options for resetting the database
+func DefaultResetDBOptions() ResetDBOptions {
+	return ResetDBOptions{
+		SkipTables: map[string]struct{}{
+			"public.infisical_migrations":      {},
+			"public.infisical_migrations_lock": {},
+		},
+		DBConfig: DefaultDatabaseConfig(),
+	}
+}
+
+// ResetDB resets the PostgreSQL database.
+// It accepts a port provider to get service ports, and options to configure the reset behavior.
+func ResetDB(ctx context.Context, opts ...func(*ResetDBOptions)) error {
+	options := DefaultResetDBOptions()
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	// Reset PostgreSQL database
+	if err := resetPostgresDB(ctx, options); err != nil {
+		return fmt.Errorf("failed to reset PostgreSQL database: %w", err)
+	}
+
+	return nil
+}
+
+// resetPostgresDB resets the PostgreSQL database by truncating all tables (except skipped ones)
+// and inserting a default super_admin record.
+func resetPostgresDB(ctx context.Context, opts ResetDBOptions) error {
+	// Build connection string using config
+	connStr := fmt.Sprintf("postgresql://%s:%s@%s:%s/%s",
+		opts.DBConfig.User,
+		opts.DBConfig.Password,
+		opts.DBConfig.Host,
+		opts.DBConfig.Port,
+		opts.DBConfig.Database,
+	)
+
+	conn, err := pgx.Connect(ctx, connStr)
+	if err != nil {
+		slog.Error("Unable to connect to database", "err", err)
+		return err
+	}
+	defer conn.Close(ctx)
+
+	// Get all tables
+	query := `
+		SELECT table_schema, table_name
+		FROM information_schema.tables
+		WHERE table_type = 'BASE TABLE'
+		  AND table_schema NOT IN ('pg_catalog', 'information_schema')
+		ORDER BY table_schema, table_name;
+	`
+
+	rows, err := conn.Query(ctx, query)
+	if err != nil {
+		slog.Error("Unable to execute query", "query", query, "err", err)
+		return err
+	}
+	defer rows.Close()
+
+	tables := make([]string, 0)
+	for rows.Next() {
+		var schema, table string
+		if err := rows.Scan(&schema, &table); err != nil {
+			slog.Error("Scan failed", "error", err)
+			return err
+		}
+		tables = append(tables, fmt.Sprintf("%s.%s", schema, table))
+	}
+	if err := rows.Err(); err != nil {
+		slog.Error("Row iteration error", "error", err)
+		return err
+	}
+
+	// Build truncate statements
+	var builder strings.Builder
+	for _, table := range tables {
+		if _, ok := opts.SkipTables[table]; ok {
+			continue
+		}
+		builder.WriteString(fmt.Sprintf("TRUNCATE TABLE %s RESTART IDENTITY CASCADE;\n", table))
+	}
+
+	truncateQuery := builder.String()
+	if truncateQuery != "" {
+		_, err = conn.Exec(ctx, truncateQuery)
+		if err != nil {
+			slog.Error("Truncate failed", "error", err)
+			return err
+		}
+		slog.Info("Truncate all tables successfully")
+	}
+
+	// Insert default super_admin record
+	_, err = conn.Exec(ctx,
+		`INSERT INTO public.super_admin ("id", "fipsEnabled", "initialized", "allowSignUp") VALUES ($1, $2, $3, $4)`,
+		"00000000-0000-0000-0000-000000000000", true, false, true)
+	if err != nil {
+		slog.Error("Failed to insert super_admin", "error", err)
+		return err
+	}
+
+	return nil
+}
+
+// WithSkipTables sets the tables to skip when truncating
+func WithSkipTables(tables map[string]struct{}) func(*ResetDBOptions) {
+	return func(opts *ResetDBOptions) {
+		opts.SkipTables = tables
+	}
+}
+
+// WithDatabaseConfig sets the database configuration
+func WithDatabaseConfig(config DatabaseConfig) func(*ResetDBOptions) {
+	return func(opts *ResetDBOptions) {
+		opts.DBConfig = config
+	}
+}

e2e/packages/client/reset_redis.go

@@ -0,0 +1,83 @@
+package client
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// RedisConfig holds Redis connection configuration
+type RedisConfig struct {
+	Host     string
+	Port     int
+	Password string
+}
+
+// DefaultRedisConfig returns the default Redis configuration
+func DefaultRedisConfig() RedisConfig {
+	return RedisConfig{
+		Host:     "localhost",
+		Port:     6379,
+		Password: "",
+	}
+}
+
+// ResetRedisOptions holds options for resetting Redis
+type ResetRedisOptions struct {
+	RedisConfig RedisConfig
+}
+
+// DefaultResetRedisOptions returns default options for resetting Redis
+func DefaultResetRedisOptions() ResetRedisOptions {
+	return ResetRedisOptions{
+		RedisConfig: DefaultRedisConfig(),
+	}
+}
+
+// ResetRedis resets the Redis database by flushing all keys.
+// It accepts a port provider to get service ports, and options to configure the reset behavior.
+func ResetRedis(ctx context.Context, opts ...func(*ResetRedisOptions)) error {
+	options := DefaultResetRedisOptions()
+	for _, opt := range opts {
+		opt(&options)
+	}
+
+	return resetRedisDB(ctx, options)
+}
+
+// resetRedisDB resets the Redis database by flushing all keys.
+func resetRedisDB(ctx context.Context, opts ResetRedisOptions) error {
+	addr := fmt.Sprintf("%s:%d", opts.RedisConfig.Host, opts.RedisConfig.Port)
+	rdb := redis.NewClient(&redis.Options{
+		Addr:     addr,
+		Password: opts.RedisConfig.Password,
+	})
+	defer func() {
+		_ = rdb.Close()
+	}()
+
+	// Test the connection
+	pong, err := rdb.Ping(ctx).Result()
+	if err != nil {
+		return fmt.Errorf("failed to connect to Redis: %w", err)
+	}
+	slog.Info("Connected to Redis", "pong", pong)
+
+	// Clear all keys in the current database
+	err = rdb.FlushAll(ctx).Err()
+	if err != nil {
+		return fmt.Errorf("failed to flush Redis database: %w", err)
+	}
+	slog.Info("All keys cleared successfully from Redis database")
+
+	return nil
+}
+
+// WithRedisConfig sets the Redis configuration
+func WithRedisConfig(config RedisConfig) func(*ResetRedisOptions) {
+	return func(opts *ResetRedisOptions) {
+		opts.RedisConfig = config
+	}
+}