database_loader.py 4.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121
  1. import mysql.connector
  2. from mysql.connector import errorcode
  3. from .config import settings
  4. from app.core.logger import get_logger
  5. logger = get_logger(__name__)
  6. # 全局的连接池
  7. db_connection_pool = None
  8. def init_database():
  9. """
  10. 初始化数据库:如果数据库或表不存在,则创建它们。
  11. 主键统一为 'id'。
  12. """
  13. logger.info("--- 开始初始化数据库 ---")
  14. try:
  15. cnx = mysql.connector.connect(**settings.DATABASE_CONFIG)
  16. cursor = cnx.cursor()
  17. cursor.execute(f"CREATE DATABASE IF NOT EXISTS {settings.DB_NAME} DEFAULT CHARACTER SET 'utf8mb4'")
  18. logger.info(f"数据库 '{settings.DB_NAME}' 已准备就绪。")
  19. cnx.database = settings.DB_NAME
  20. # 1. 创建 cards 表 (主键为 id)
  21. cards_table = (
  22. f"CREATE TABLE IF NOT EXISTS `{settings.DB_CARD_TABLE_NAME}` ("
  23. " `id` INT AUTO_INCREMENT PRIMARY KEY,"
  24. " `card_name` VARCHAR(255) NULL COMMENT '卡牌的通用名称',"
  25. " `card_type` VARCHAR(50) NOT NULL DEFAULT 'pokemon' COMMENT '卡牌类型 (pokemon, basketball, etc)',"
  26. " `detection_score` DECIMAL(4, 2) NULL COMMENT '原始检测总分',"
  27. " `modified_score` DECIMAL(4, 2) NULL COMMENT '修改后总分',"
  28. " `is_edited` BOOLEAN NOT NULL DEFAULT FALSE COMMENT '是否被编辑过',"
  29. " `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
  30. " `updated_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP"
  31. ") ENGINE=InnoDB COMMENT='存储实体卡牌的核心信息'"
  32. )
  33. cursor.execute(cards_table)
  34. logger.info(f"数据表 '{settings.DB_CARD_TABLE_NAME}' 已准备就绪。")
  35. # 2. 创建 card_images 表 (主键为 id)
  36. card_images_table = (
  37. f"CREATE TABLE IF NOT EXISTS `{settings.DB_IMAGE_TABLE_NAME}` ("
  38. " `id` INT AUTO_INCREMENT PRIMARY KEY,"
  39. " `card_id` INT NOT NULL COMMENT '关联的卡牌ID',"
  40. " `image_type` VARCHAR(50) NOT NULL COMMENT '图片类型 (front_face, back_face, etc)',"
  41. " `image_name` VARCHAR(255) NULL,"
  42. " `image_path` VARCHAR(512) NOT NULL,"
  43. " `detection_image_path` VARCHAR(512) NULL COMMENT '检测结果图的路径',"
  44. " `modified_image_path` VARCHAR(512) NULL COMMENT '修改后结果图的路径',"
  45. " `detection_json` JSON NOT NULL COMMENT '原始检测JSON数据',"
  46. " `modified_json` JSON NULL COMMENT '修改后的JSON数据',"
  47. " `is_edited` BOOLEAN NOT NULL DEFAULT FALSE COMMENT '该图片记录是否被修改过',"
  48. " `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
  49. " `updated_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
  50. f" FOREIGN KEY (`card_id`) REFERENCES `{settings.DB_CARD_TABLE_NAME}`(`id`) ON DELETE CASCADE,"
  51. " UNIQUE KEY `uk_card_image_type` (`card_id`, `image_type`)"
  52. ") ENGINE=InnoDB COMMENT='存储所有卡牌的图片及检测数据'"
  53. )
  54. cursor.execute(card_images_table)
  55. logger.info(f"数据表 '{settings.DB_IMAGE_TABLE_NAME}' 已准备就绪。")
  56. except mysql.connector.Error as err:
  57. logger.error(f"数据库初始化失败: {err}")
  58. exit(1)
  59. finally:
  60. if 'cursor' in locals() and cursor:
  61. cursor.close()
  62. if 'cnx' in locals() and cnx.is_connected():
  63. cnx.close()
  64. logger.info("--- 数据库初始化完成 ---")
  65. def load_database_pool():
  66. """
  67. 在应用启动时创建数据库连接池。
  68. """
  69. global db_connection_pool
  70. if db_connection_pool is None:
  71. logger.info("--- 创建数据库连接池 ---")
  72. try:
  73. db_connection_pool = mysql.connector.pooling.MySQLConnectionPool(
  74. pool_name="mypool",
  75. pool_size=5,
  76. **settings.DATABASE_CONFIG_WITH_DB
  77. )
  78. logger.info("--- 数据库连接池创建成功 ---")
  79. except mysql.connector.Error as err:
  80. logger.error(f"创建数据库连接池失败: {err}")
  81. exit(1)
  82. def close_database_pool():
  83. """
  84. 在应用关闭时,不需要手动关闭连接池,连接器会自动处理。
  85. 这个函数留作备用。
  86. """
  87. logger.info("--- 数据库连接池将自动关闭 ---")
  88. # --- FastAPI 依赖注入 ---
  89. def get_db_connection():
  90. """
  91. 一个FastAPI依赖项,用于从池中获取数据库连接。
  92. 它确保连接在使用后返回到池中。
  93. """
  94. if db_connection_pool is None:
  95. raise RuntimeError("数据库连接池未初始化")
  96. db_conn = None
  97. try:
  98. db_conn = db_connection_pool.get_connection()
  99. yield db_conn
  100. except mysql.connector.Error as err:
  101. logger.error(f"获取数据库连接失败: {err}")
  102. finally:
  103. if db_conn and db_conn.is_connected():
  104. db_conn.close()