diff --git a/apps/application/migrations/0015_re_database_index.py b/apps/application/migrations/0015_re_database_index.py new file mode 100644 index 000000000..740a2a2d2 --- /dev/null +++ b/apps/application/migrations/0015_re_database_index.py @@ -0,0 +1,63 @@ +# Generated by Django 4.2.15 on 2024-09-18 16:14 +import logging + +import psycopg2 +from django.db import migrations +from psycopg2 import extensions + +from smartdoc.const import CONFIG + + +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 = psycopg2.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.set_isolation_level(extensions.ISOLATION_LEVEL_AUTOCOMMIT) + 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): + dependencies = [ + ('application', '0014_application_problem_optimization_prompt'), + ] + + operations = [ + migrations.RunPython(re_index, atomic=False) + ]