1+ #!/usr/bin/env python3
2+ """
3+ Script to wait for PostgreSQL database to be ready and create it if it doesn't exist.
4+ """
5+
6+ import sys
7+ import psycopg2
8+ import os
9+ import re
10+ from urllib .parse import urlparse
11+ import time
12+
13+
14+ def create_database_if_not_exists ():
15+ """Attempt to connect to database and create it if it doesn't exist."""
16+ database_url = os .getenv ("DATABASE_URL" , "postgres://postgres:postgres@postgres:5432/postgres_service" )
17+
18+ # Parse the database URL to extract components
19+ parsed = urlparse (database_url )
20+
21+ # Create connection URL without database name (to connect to postgres default database)
22+ base_url = f"postgresql://{ parsed .username } { parsed .password } { parsed .hostname } { parsed .port or 5432 }
23+
24+ try :
25+ # First try to connect to the target database
26+ psycopg2 .connect (database_url )
27+ print ("Database connection successful" )
28+ return True
29+ except psycopg2 .OperationalError as e :
30+ # If database doesn't exist, try to create it
31+ if "does not exist" in str (e ) or "database" in str (e ).lower ():
32+ print (f"Database does not exist, attempting to create it..." )
33+ try :
34+ # Connect to postgres database to create the target database
35+ conn = psycopg2 .connect (base_url )
36+ conn .autocommit = True
37+ cursor = conn .cursor ()
38+
39+ # Extract database name from the original URL
40+ db_name = parsed .path .lstrip ('/' )
41+
42+ # Create the database
43+ cursor .execute (f"CREATE DATABASE { db_name } )
44+ cursor .close ()
45+ conn .close ()
46+
47+ print (f"Database '{ db_name } )
48+
49+ # Try connecting again to verify
50+ psycopg2 .connect (database_url )
51+ print ("Database connection verified after creation" )
52+ return True
53+ except Exception as create_error :
54+ print (f"Failed to create database: { create_error } )
55+ return False
56+ else :
57+ # Other connection error (host unreachable, etc.)
58+ print (f"Database connection error: { e } )
59+ return False
60+
61+
62+ def wait_for_database (max_retries = 60 , retry_interval = 1 ):
63+ """Wait for database to be available with retries."""
64+ for attempt in range (max_retries ):
65+ if create_database_if_not_exists ():
66+ return True
67+ print (f"Database not ready, retrying in { retry_interval } { attempt + 1 } { max_retries } )
68+ time .sleep (retry_interval )
69+
70+ print ("Failed to connect to database after maximum retries" )
71+ return False
72+
73+
74+ if __name__ == "__main__" :
75+ if wait_for_database ():
76+ sys .exit (0 )
77+ else :
78+ sys .exit (1 )
0 commit comments