diff --git a/functions/backup_restore/database/base.py b/functions/backup_restore/database/base.py index b93e1e5..50694f0 100644 --- a/functions/backup_restore/database/base.py +++ b/functions/backup_restore/database/base.py @@ -33,10 +33,10 @@ def __init__(self, app_name: str): self.dump_command = None self.error = None - # Fetch database name and user if needed - if self.chart_info.chart_name != "immich": - self.database_name = self.fetch_database_name() - self.database_user = self.fetch_database_user() or self.database_name + # # Fetch database name and user if needed + # if self.chart_info.chart_name != "immich": + self.database_name = self.fetch_database_name() + self.database_user = self.fetch_database_user() or self.database_name def fetch_primary_pod(self, timeout=600, interval=5) -> str: """ diff --git a/functions/backup_restore/database/restore.py b/functions/backup_restore/database/restore.py index 77d7a4e..2a753e4 100644 --- a/functions/backup_restore/database/restore.py +++ b/functions/backup_restore/database/restore.py @@ -82,10 +82,10 @@ def restore(self, timeout=300, interval=5) -> Dict[str, str]: self.command, self.open_mode = self._get_restore_command() try: - # Drop and recreate the database before restoring - drop_create_result = self._drop_and_create_database() - if not drop_create_result["success"]: - result["message"] = drop_create_result["message"] + # Drop all objects in the database before restoring + drop_all_objects_result = self._drop_all_objects() + if not drop_all_objects_result["success"]: + result["message"] = drop_all_objects_result["message"] self.logger.error(result["message"]) return result @@ -153,9 +153,9 @@ def _get_restore_command(self) -> Tuple[str, str]: self.logger.debug(f"Restore command for app {self.app_name}: {command}") return command, open_mode - def _drop_and_create_database(self) -> Dict[str, str]: + def _drop_all_objects(self) -> Dict[str, str]: """ - Drop and recreate the database. + Drop all objects in the database. Returns: dict: Result containing status and message. @@ -165,7 +165,7 @@ def _drop_and_create_database(self) -> Dict[str, str]: "message": "" } - drop_command = [ + drop_all_objects_command = [ "k3s", "kubectl", "exec", "--namespace", self.namespace, "--stdin", @@ -173,41 +173,45 @@ def _drop_and_create_database(self) -> Dict[str, str]: self.primary_pod, "--", "psql", + "--dbname", self.database_name, "--command", - f"DROP DATABASE IF EXISTS {self.database_name};" - ] - - create_command = [ - "k3s", "kubectl", "exec", - "--namespace", self.namespace, - "--stdin", - "--container", "postgres", - self.primary_pod, - "--", - "psql", - "--command", - f"CREATE DATABASE {self.database_name} OWNER {self.database_user};" + """ + DO $$ DECLARE + r RECORD; + BEGIN + -- drop all tables + FOR r IN (SELECT tablename FROM pg_tables WHERE schemaname = current_schema()) LOOP + EXECUTE 'DROP TABLE IF EXISTS ' || quote_ident(r.tablename) || ' CASCADE'; + END LOOP; + -- drop all sequences + FOR r IN (SELECT sequencename FROM pg_sequences WHERE schemaname = current_schema()) LOOP + EXECUTE 'DROP SEQUENCE IF EXISTS ' || quote_ident(r.sequencename) || ' CASCADE'; + END LOOP; + -- drop all views + FOR r IN (SELECT viewname FROM pg_views WHERE schemaname = current_schema()) LOOP + EXECUTE 'DROP VIEW IF EXISTS ' || quote_ident(r.viewname) || ' CASCADE'; + END LOOP; + -- drop all functions + FOR r IN (SELECT proname FROM pg_proc p JOIN pg_namespace n ON p.pronamespace = n.oid WHERE n.nspname = current_schema()) LOOP + EXECUTE 'DROP FUNCTION IF EXISTS ' || quote_ident(r.proname) || ' CASCADE'; + END LOOP; + END $$; + """ ] try: - drop_process = subprocess.run(drop_command, capture_output=True, text=True) + drop_process = subprocess.run(drop_all_objects_command, capture_output=True, text=True) if drop_process.returncode != 0: - result["message"] = f"Failed to drop database: {drop_process.stderr}" - self.logger.error(result["message"]) - return result - - create_process = subprocess.run(create_command, capture_output=True, text=True) - if create_process.returncode != 0: - result["message"] = f"Failed to create database: {create_process.stderr}" + result["message"] = f"Failed to drop all objects in database: {drop_process.stderr}" self.logger.error(result["message"]) return result result["success"] = True - result["message"] = "Database dropped and recreated successfully." + result["message"] = "All objects in database dropped successfully." self.logger.debug(result["message"]) except Exception as e: - message = f"Failed to drop and create database: {e}" + message = f"Failed to drop all objects in database: {e}" self.logger.error(message, exc_info=True) result["message"] = message @@ -278,4 +282,4 @@ def _execute_restore_command(self, retries=3, wait=5) -> Dict[str, str]: result["message"] = f"{result['message']} Restore failed after retrying." self.logger.error(result["message"]) - return result \ No newline at end of file + return result