@@ -7,166 +7,44 @@ import (
77 "fmt"
88 "time"
99
10- "github.com/cybertec-postgresql/pgcov/pkg/types"
1110 "github.com/jackc/pgx/v5/pgxpool"
1211)
1312
14- // CreateTempDatabase creates a temporary database with a unique name
15- func CreateTempDatabase ( ctx context. Context , pool * Pool ) ( * types. TempDatabase , error ) {
16- // Generate unique database name
13+ // CreateTempDatabase creates a temporary database and returns a pool connected to it.
14+ // The database name is accessible via pool.Config().ConnConfig.Database.
15+ func CreateTempDatabase ( ctx context. Context , adminPool * Pool ) ( * pgxpool. Pool , error ) {
1716 timestamp := time .Now ().Format ("20060102_150405" )
1817 randomBytes := make ([]byte , 4 )
1918 if _ , err := rand .Read (randomBytes ); err != nil {
2019 return nil , fmt .Errorf ("failed to generate random suffix: %w" , err )
2120 }
2221 randomSuffix := hex .EncodeToString (randomBytes )
23-
2422 dbName := fmt .Sprintf ("pgcov_test_%s_%s" , timestamp , randomSuffix )
2523
26- // Create database (must use a connection to template database)
27- conn , err := pool .Acquire (ctx )
28- if err != nil {
29- return nil , fmt .Errorf ("failed to acquire connection: %w" , err )
30- }
31- defer conn .Release ()
32-
33- // CREATE DATABASE cannot run in a transaction
34- createSQL := fmt .Sprintf ("CREATE DATABASE %s" , dbName )
35- _ , err = conn .Exec (ctx , createSQL )
24+ _ , err := adminPool .Exec (ctx , fmt .Sprintf ("CREATE DATABASE %s" , dbName ))
3625 if err != nil {
3726 return nil , fmt .Errorf ("failed to create temporary database: %w" , err )
3827 }
3928
40- // Build connection string for the new database by replacing the dbname in the original connection string
41- // Parse the original connection config and change the database
42- baseConfig , err := pgxpool .ParseConfig (pool .config .ConnectionString )
43- if err != nil {
44- // Fallback: if parsing fails, drop the database and return error
45- _ , _ = conn .Exec (ctx , fmt .Sprintf ("DROP DATABASE IF EXISTS %s" , dbName ))
46- return nil , fmt .Errorf ("failed to parse base connection string: %w" , err )
47- }
29+ // Build connection string for the new database, preserving all original options (sslmode, etc.)
30+ config := adminPool .Pool .Config ()
31+ config .ConnConfig .Database = dbName
4832
49- // Build new connection string with the new database name
50- connString := fmt .Sprintf ("host=%s port=%d user=%s password=%s dbname=%s sslmode=%s" ,
51- baseConfig .ConnConfig .Host ,
52- baseConfig .ConnConfig .Port ,
53- baseConfig .ConnConfig .User ,
54- baseConfig .ConnConfig .Password ,
55- dbName ,
56- "prefer" ) // Default to prefer for sslmode
57-
58- return & types.TempDatabase {
59- Name : dbName ,
60- CreatedAt : time .Now (),
61- ConnectionString : connString ,
62- }, nil
63- }
64-
65- // DestroyTempDatabase drops a temporary database
66- func DestroyTempDatabase (ctx context.Context , pool * Pool , tempDB * types.TempDatabase ) error {
67- if tempDB == nil {
68- return nil
69- }
70-
71- conn , err := pool .Acquire (ctx )
72- if err != nil {
73- return fmt .Errorf ("failed to acquire connection: %w" , err )
74- }
75- defer conn .Release ()
76-
77- // Terminate all connections to the database first (PostgreSQL 13+)
78- terminateSQL := `SELECT pg_terminate_backend(pid)
79- FROM pg_stat_activity
80- WHERE datname = $1 AND pid <> pg_backend_pid()`
81-
82- _ , _ = conn .Exec (ctx , terminateSQL , tempDB .Name )
83-
84- // Drop database with FORCE option (PostgreSQL 13+)
85- dropSQL := fmt .Sprintf ("DROP DATABASE IF EXISTS %s WITH (FORCE)" , tempDB .Name )
86- _ , err = conn .Exec (ctx , dropSQL )
33+ tempPool , err := pgxpool .NewWithConfig (ctx , config )
8734 if err != nil {
88- return fmt .Errorf ("failed to drop temporary database %s: %w" , tempDB .Name , err )
89- }
90-
91- // Verify database was actually dropped
92- if err := verifyDatabaseDropped (ctx , conn , tempDB .Name ); err != nil {
93- return fmt .Errorf ("cleanup verification failed for database %s: %w" , tempDB .Name , err )
35+ _ , _ = adminPool .Exec (ctx , fmt .Sprintf ("DROP DATABASE IF EXISTS %s" , dbName ))
36+ return nil , fmt .Errorf ("failed to connect to temp database: %w" , err )
9437 }
9538
96- return nil
39+ return tempPool , nil
9740}
9841
99- // verifyDatabaseDropped checks that a database no longer exists in PostgreSQL catalog
100- func verifyDatabaseDropped (ctx context.Context , conn * pgxpool.Conn , dbName string ) error {
101- verifySQL := `
102- SELECT EXISTS(
103- SELECT 1
104- FROM pg_database
105- WHERE datname = $1
106- )
107- `
108-
109- var exists bool
110- err := conn .QueryRow (ctx , verifySQL , dbName ).Scan (& exists )
111- if err != nil {
112- return fmt .Errorf ("failed to verify database deletion: %w" , err )
113- }
114-
115- if exists {
116- return fmt .Errorf ("database %s still exists after DROP command" , dbName )
117- }
118-
119- return nil
120- }
121-
122- // CleanupStaleTempDatabases removes old pgcov temporary databases
123- // This is useful for cleanup after crashes or interrupted test runs
124- // Returns list of cleaned databases and any errors encountered
125- func CleanupStaleTempDatabases (ctx context.Context , pool * Pool , _ time.Duration ) ([]string , error ) {
126- conn , err := pool .Acquire (ctx )
127- if err != nil {
128- return nil , fmt .Errorf ("failed to acquire connection: %w" , err )
129- }
130- defer conn .Release ()
131-
132- // Find pgcov temp databases
133- query := `
134- SELECT datname
135- FROM pg_database
136- WHERE datname LIKE 'pgcov_test_%'
137- `
138-
139- rows , err := conn .Query (ctx , query )
140- if err != nil {
141- return nil , fmt .Errorf ("failed to query temp databases: %w" , err )
142- }
143- defer rows .Close ()
144-
145- var cleaned []string
146- var failedCleanup []string
147-
148- for rows .Next () {
149- var dbName string
150- if err := rows .Scan (& dbName ); err != nil {
151- continue
152- }
153-
154- // Extract timestamp from database name
155- // Format: pgcov_test_YYYYMMDD_HHMMSS_randomhex
156- tempDB := & types.TempDatabase {Name : dbName }
157-
158- // Attempt to drop (will fail if database is in use)
159- if err := DestroyTempDatabase (ctx , pool , tempDB ); err == nil {
160- cleaned = append (cleaned , dbName )
161- } else {
162- failedCleanup = append (failedCleanup , dbName )
163- }
164- }
165-
166- // Report cleanup failures as non-fatal warning
167- if len (failedCleanup ) > 0 {
168- return cleaned , fmt .Errorf ("failed to cleanup %d databases: %v (may be in use)" , len (failedCleanup ), failedCleanup )
42+ // DestroyTempDatabase closes the temp pool and drops its underlying database.
43+ func DestroyTempDatabase (ctx context.Context , adminPool * Pool , tempPool * pgxpool.Pool ) error {
44+ if tempPool == nil {
45+ return nil
16946 }
170-
171- return cleaned , nil
47+ tempPool .Close ()
48+ _ , err := adminPool .Exec (ctx , fmt .Sprintf ("DROP DATABASE IF EXISTS %s WITH (FORCE)" , tempPool .Config ().ConnConfig .Database ))
49+ return err
17250}
0 commit comments