From 6a507dca7b472cf559b423254869e9d87e9a01db Mon Sep 17 00:00:00 2001
From: DingZQ <ding@aecforge.com>
Date: Fri, 7 Feb 2025 23:12:28 +0800
Subject: [PATCH] Add more project methods

---
 api/project.py | 79 ++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 79 insertions(+)

diff --git a/api/project.py b/api/project.py
index 2f882c7..a24660f 100644
--- a/api/project.py
+++ b/api/project.py
@@ -1,3 +1,4 @@
+import os
 import psycopg as pg
 from psycopg.rows import dict_row
 from .connection import g_conn_dict as conn
@@ -28,6 +29,83 @@ def copy_project(source: str, new: str) -> None:
         with conn.cursor() as cur:
             cur.execute(f'create database "{new}" with template = {source}')
 
+# 2025-02-07, WMH
+class CopyProjectEx:
+    @ staticmethod
+    def create_database(connection, new_db):
+        with connection.cursor() as cursor:
+            cursor.execute(f'create database "{new_db}"')
+        connection.commit()
+
+    @staticmethod
+    def execute_pg_dump(hostname, source_db, exclude_table_list):
+        dump_command_structure = (
+            f'pg_dump -h {hostname} -F c -s -f source_db_structure.dump {source_db}'
+        )
+        os.system(dump_command_structure)
+
+        if exclude_table_list is not None:
+            exclude_table = ' '.join(['-T {}'.format(i) for i in exclude_table_list])
+            dump_command_db = (
+                f'pg_dump -h {hostname} -F c -a {exclude_table} -f source_db.dump {source_db}'
+            )
+        else:
+            dump_command_db = (
+                f'pg_dump -h {hostname} -F c -a -f source_db.dump {source_db}'
+            )
+        os.system(dump_command_db)
+
+    @staticmethod
+    def execute_pg_restore(hostname, new_db):
+        restore_command_structure = (
+            f'pg_restore -h {hostname} -d {new_db} source_db_structure.dump'
+        )
+        os.system(restore_command_structure)
+
+        restore_command_db = (
+            f'pg_restore -h {hostname} -d {new_db} source_db.dump'
+        )
+        os.system(restore_command_db)
+
+    @staticmethod
+    def init_operation_table(connection, excluded_table):
+        with connection.cursor() as cursor:
+            if 'operation' in excluded_table:
+                insert_query \
+                    = "insert into operation (id, redo, undo, redo_cs, undo_cs) values (0, '', '', '', '')"
+                cursor.execute(insert_query)
+
+            if 'current_operation' in excluded_table:
+                insert_query \
+                    = "insert into current_operation (id) values (0)"
+                cursor.execute(insert_query)
+
+            if 'restore_operation' in excluded_table:
+                insert_query \
+                    = "insert into restore_operation (id) values (0)"
+                cursor.execute(insert_query)
+
+            if 'batch_operation' in excluded_table:
+                insert_query \
+                    = "insert into batch_operation (id, redo, undo, redo_cs, undo_cs) values (0, '', '', '', '')"
+                cursor.execute(insert_query)
+
+            if 'operation_table' in excluded_table:
+                insert_query \
+                    = "insert into operation_table (option) values ('operation')"
+                cursor.execute(insert_query)
+        connection.commit()
+
+    def __call__(self, source: str, new: str, excluded_table: [str] = None) -> None:
+        connection = pg.connect(conninfo="dbname=postgres host=127.0.0.1", autocommit=True)
+
+        self.create_database(connection, new)
+        self.execute_pg_dump('127.0.0.1', source, excluded_table)
+        self.execute_pg_restore('127.0.0.1', new)
+
+        connection = pg.connect(conninfo=f"dbname='{new}' host=127.0.0.1", autocommit=True)
+        self.init_operation_table(connection, excluded_table)
+
 
 def create_project(name: str) -> None:
     return copy_project('project', name)
@@ -69,3 +147,4 @@ def close_project(name: str) -> None:
     if name in conn:
         conn[name].close()
         del conn[name]
+