From d8dff0118b5a8f7770b7b2ef275ffc01467ed177 Mon Sep 17 00:00:00 2001 From: "WQY\\qiong" Date: Wed, 22 Mar 2023 19:24:50 +0800 Subject: [PATCH] Support to clean all projects --- api/__init__.py | 2 +- api/project.py | 17 +++++++++++++++++ tjnetwork.py | 3 +++ 3 files changed, 21 insertions(+), 1 deletion(-) diff --git a/api/__init__.py b/api/__init__.py index ad08e19..00c1638 100644 --- a/api/__init__.py +++ b/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 diff --git a/api/project.py b/api/project.py index c9cdfb0..0db3c7b 100644 --- a/api/project.py +++ b/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) diff --git a/tjnetwork.py b/tjnetwork.py index b8cd257..12b9247 100644 --- a/tjnetwork.py +++ b/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)