From 7fc1d89659c1e86966b37128d27f98e3910fb73b Mon Sep 17 00:00:00 2001
From: CaptainB <bin@fit2cloud.com>
Date: Tue, 19 Aug 2025 15:48:38 +0800
Subject: [PATCH] fix: refactor database collation refresh and reindex logic
 with improved error handling

---
 .../0002_refresh_collation_reindex.py         | 61 +++++++++++++++----
 1 file changed, 48 insertions(+), 13 deletions(-)

diff --git a/apps/system_manage/migrations/0002_refresh_collation_reindex.py b/apps/system_manage/migrations/0002_refresh_collation_reindex.py
index b28f800d0..fe5fc4d12 100644
--- a/apps/system_manage/migrations/0002_refresh_collation_reindex.py
+++ b/apps/system_manage/migrations/0002_refresh_collation_reindex.py
@@ -1,26 +1,61 @@
+import logging
+
+import psycopg
 from django.db import migrations
 
-
-def refresh_collation_and_reindex(apps, schema_editor):
-    # 获取当前数据库名
-    db_name = schema_editor.connection.settings_dict["NAME"]
-    with schema_editor.connection.cursor() as cursor:
-        cursor.execute(f'ALTER DATABASE "{db_name}" REFRESH COLLATION VERSION;')
-        cursor.execute(f'REINDEX DATABASE "{db_name}";')
+from maxkb.const import CONFIG
 
 
-def noop(apps, schema_editor):
-    # 不可逆操作，留空
-    pass
+def get_connect(db_name):
+    conn_params = {
+        "dbname": db_name,
+        "user": CONFIG.get('DB_USER'),
+        "password": CONFIG.get('DB_PASSWORD'),
+        "host": CONFIG.get('DB_HOST'),
+        "port": CONFIG.get('DB_PORT')
+    }
+    # 建立连接
+    connect = psycopg.connect(**conn_params)
+    return connect
+
+
+def sql_execute(conn, reindex_sql: str, alter_database_sql: str):
+    """
+    执行一条sql
+    @param reindex_sql:
+    @param conn:
+    @param alter_database_sql:
+    """
+    conn.autocommit = True
+    with conn.cursor() as cursor:
+        cursor.execute(reindex_sql, [])
+        cursor.execute(alter_database_sql, [])
+        cursor.close()
+
+def re_index(apps, schema_editor):
+    app_db_name = CONFIG.get('DB_NAME')
+    try:
+        re_index_database(app_db_name)
+    except Exception as e:
+        logging.error(f'reindex database {app_db_name}发送错误:{str(e)}')
+    try:
+        re_index_database('root')
+    except Exception as e:
+        logging.error(f'reindex database root 发送错误:{str(e)}')
+
+
+def re_index_database(db_name):
+    db_conn = get_connect(db_name)
+    sql_execute(db_conn, f'REINDEX DATABASE "{db_name}";', f'ALTER DATABASE "{db_name}" REFRESH COLLATION VERSION;')
+    db_conn.close()
 
 
 class Migration(migrations.Migration):
-    atomic = False  # ALTER DATABASE/REINDEX 需在事务外执行
 
     dependencies = [
         ("system_manage", "0001_initial"),
     ]
 
     operations = [
-        migrations.RunPython(refresh_collation_and_reindex, reverse_code=noop),
-    ]
\ No newline at end of file
+        migrations.RunPython(re_index, atomic=False)
+    ]