1+ import os
2+ import subprocess
3+ from flask import Flask , render_template , request , flash , redirect , url_for
4+
5+ app = Flask (__name__ )
6+ app .secret_key = os .urandom (24 ) # Needed for flashing messages
7+
8+ BACKUP_DIR = "/backups"
9+ # --- Database Connection Details (Read from Environment Variables) ---
10+ # These will be needed for the restore operation
11+ TARGET_DB_HOST = os .environ .get ("TARGET_DB_HOST" , "db" ) # Default to 'db' service name
12+ TARGET_DB_PORT = os .environ .get ("TARGET_DB_PORT" , "5432" )
13+ TARGET_DB_USER = os .environ .get ("TARGET_DB_USER" , "postgres" )
14+ TARGET_DB_PASSWORD = os .environ .get ("TARGET_DB_PASSWORD" ) # MUST be provided
15+
16+ DEFAULT_RESTORE_SUFFIX = "_restore"
17+
18+ def get_backup_files ():
19+ """Scans the backup directory for .sql.gz files."""
20+ backups = []
21+ if not os .path .isdir (BACKUP_DIR ):
22+ flash (f"Error: Backup directory '{ BACKUP_DIR } , "danger" )
23+ return backups
24+ try :
25+ for filename in sorted (os .listdir (BACKUP_DIR ), reverse = True ):
26+ if filename .endswith (".sql.gz" ):
27+ # Extract original DB name (assuming format DB_NAME_backup_TIMESTAMP.sql.gz)
28+ parts = filename .split ('_backup_' )
29+ original_db = parts [0 ] if len (parts ) > 0 else "unknown_db"
30+ backups .append ({"filename" : filename , "original_db" : original_db })
31+ except OSError as e :
32+ flash (f"Error reading backup directory '{ BACKUP_DIR } { e } , "danger" )
33+ return backups
34+
35+ @app .route ('/' )
36+ def index ():
37+ """Displays the list of backups and the restore form."""
38+ backup_files = get_backup_files ()
39+ return render_template ('index.html' , backups = backup_files , default_suffix = DEFAULT_RESTORE_SUFFIX )
40+
41+ @app .route ('/restore' , methods = ['POST' ])
42+ def restore_backup ():
43+ """Handles the restore request."""
44+ selected_backup = request .form .get ('backup_file' )
45+ restore_suffix = request .form .get ('restore_suffix' , DEFAULT_RESTORE_SUFFIX ).strip ()
46+
47+ if not selected_backup :
48+ flash ("Error: No backup file selected." , "danger" )
49+ return redirect (url_for ('index' ))
50+
51+ if not restore_suffix :
52+ restore_suffix = DEFAULT_RESTORE_SUFFIX # Ensure default if user enters empty string
53+
54+ # Basic validation for suffix (prevent potentially harmful characters)
55+ if not restore_suffix .replace ('_' , '' ).isalnum ():
56+ flash (f"Error: Invalid suffix '{ restore_suffix } , "danger" )
57+ return redirect (url_for ('index' ))
58+
59+ backup_path = os .path .join (BACKUP_DIR , selected_backup )
60+
61+ if not os .path .exists (backup_path ):
62+ flash (f"Error: Selected backup file '{ selected_backup } , "danger" )
63+ return redirect (url_for ('index' ))
64+
65+ # Extract original DB name again for safety
66+ parts = selected_backup .split ('_backup_' )
67+ original_db_name = parts [0 ] if len (parts ) > 0 else None
68+
69+ if not original_db_name :
70+ flash (f"Error: Could not determine original database name from '{ selected_backup } , "danger" )
71+ return redirect (url_for ('index' ))
72+
73+ new_db_name = f"{ original_db_name } { restore_suffix }
74+
75+ # --- Input Validation ---
76+ if not TARGET_DB_PASSWORD :
77+ flash ("Error: TARGET_DB_PASSWORD environment variable is not set. Cannot connect to target database." , "danger" )
78+ return redirect (url_for ('index' ))
79+
80+ # --- Restore Process ---
81+ try :
82+ # 1. Create the new database
83+ flash (f"Attempting to create new database: '{ new_db_name } { TARGET_DB_HOST } , "info" )
84+ create_db_cmd = [
85+ "createdb" ,
86+ "-h" , TARGET_DB_HOST ,
87+ "-p" , TARGET_DB_PORT ,
88+ "-U" , TARGET_DB_USER ,
89+ new_db_name
90+ ]
91+ env = os .environ .copy ()
92+ env ['PGPASSWORD' ] = TARGET_DB_PASSWORD
93+ # Use check_output to capture stderr on failure
94+ subprocess .check_output (create_db_cmd , stderr = subprocess .STDOUT , env = env )
95+ flash (f"Successfully created database '{ new_db_name } , "success" )
96+
97+ # 2. Decompress and restore using psql
98+ flash (f"Attempting to restore '{ selected_backup } { new_db_name } , "info" )
99+ # Command: gunzip < backup_path | psql -h host -p port -U user -d new_db_name
100+ # We use shell=True because of the pipe. Be cautious with user input (already validated suffix).
101+ restore_cmd_str = f"gunzip < \" { backup_path } \" | psql -h \" { TARGET_DB_HOST } \" -p \" { TARGET_DB_PORT } \" -U \" { TARGET_DB_USER } \" -d \" { new_db_name } \" --quiet"
102+
103+ # Run the command
104+ # Use check_output to capture stderr on failure
105+ process = subprocess .run (restore_cmd_str , shell = True , capture_output = True , text = True , env = env )
106+
107+ if process .returncode == 0 :
108+ flash (f"Successfully restored backup '{ selected_backup } { new_db_name } , "success" )
109+ else :
110+ # Attempt to drop the partially created/restored DB on failure
111+ flash (f"Error during restore process (Exit Code: { process .returncode } , "danger" )
112+ flash (f"Stderr: { process .stderr } , "warning" )
113+ flash (f"Stdout: { process .stdout } , "warning" )
114+ try :
115+ drop_db_cmd = [
116+ "dropdb" , "--if-exists" ,
117+ "-h" , TARGET_DB_HOST ,
118+ "-p" , TARGET_DB_PORT ,
119+ "-U" , TARGET_DB_USER ,
120+ new_db_name
121+ ]
122+ subprocess .check_output (drop_db_cmd , stderr = subprocess .STDOUT , env = env )
123+ flash (f"Cleaned up (dropped) database '{ new_db_name } , "info" )
124+ except subprocess .CalledProcessError as drop_e :
125+ flash (f"Error during cleanup (dropping '{ new_db_name } { drop_e .output .decode ()} , "danger" )
126+
127+ except subprocess .CalledProcessError as e :
128+ flash (f"Error executing command: { e .cmd } , "danger" )
129+ flash (f"Output: { e .output .decode ()} , "warning" )
130+ # If createdb failed, no need to drop
131+ except Exception as e :
132+ flash (f"An unexpected error occurred: { e } , "danger" )
133+
134+ return redirect (url_for ('index' ))
135+
136+
137+ if __name__ == '__main__' :
138+ # Use 0.0.0.0 to be accessible within Docker network
139+ app .run (host = '0.0.0.0' , port = 5001 , debug = True ) # Use a different port like 5001
0 commit comments