Accept Merge Request #185: (script -> master)

Merge Request: Support to clean all projects

Created By: @王琼钰
Accepted By: @王琼钰
URL: https://tjwater.coding.net/p/tjwatercloud/d/TJWaterServer/git/merge/185?initial=true

api/__init__.py

@@ -1,4 +1,4 @@
-from .project import list_project, have_project, create_project, delete_project
+from .project import list_project, have_project, create_project, delete_project, clean_project
 from .project import is_project_open, get_project_open_count, open_project, close_project
 from .project import copy_project
 

api/project.py

@@ -6,6 +6,8 @@ from .connection import g_conn_dict as conn
 
 _project_open_count: dict[str, int] = {}
 
+_server_databases = ['template0', 'template1', 'postgres', 'project']
+
 def list_project() -> list[str]:
     ps = []
     with pg.connect(conninfo="dbname=postgres host=127.0.0.1", autocommit=True) as conn:
@@ -33,6 +35,21 @@ def delete_project(name: str) -> None:
         with conn.cursor() as cur:
             cur.execute(f"drop database {name}")
 
+def clean_project(excluded: list[str] = []) -> None:
+    projects = list_project()
+    with pg.connect(conninfo="dbname=postgres host=127.0.0.1", autocommit=True) as conn:
+        with conn.cursor(row_factory=dict_row) as cur:
+            row = cur.execute(f"select current_database()").fetchone()
+            if row != None:
+                current_db = row['current_database']
+                if current_db in projects:
+                    projects.remove(current_db)
+            for project in projects:
+                if project in _server_databases or project in excluded:
+                    continue
+                cur.execute(f"select pg_terminate_backend(pid) from pg_stat_activity where datname = '{project}'")
+                cur.execute(f"drop database {project}")
+
 def open_project(name: str) -> None:
     if name not in _project_open_count:
         conn[name] = pg.connect(conninfo=f"dbname={name} host=127.0.0.1", autocommit=True)

tjnetwork.py

@@ -170,6 +170,9 @@ def create_project(name: str) -> None:
 def delete_project(name: str) -> None:
     return api.delete_project(name)
 
+def clean_project(excluded: list[str] = []) -> None:
+    return api.clean_project(excluded)
+
 def is_project_open(name: str) -> bool:
     return api.is_project_open(name)