database_loader.py 6.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  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. """
  12. logger.info("--- 开始初始化数据库 ---")
  13. try:
  14. cnx = mysql.connector.connect(**settings.DATABASE_CONFIG)
  15. cursor = cnx.cursor()
  16. cursor.execute(f"CREATE DATABASE IF NOT EXISTS {settings.DB_NAME} DEFAULT CHARACTER SET 'utf8mb4'")
  17. logger.info(f"数据库 '{settings.DB_NAME}' 已准备就绪。")
  18. cnx.database = settings.DB_NAME
  19. # 1. 创建 cards 表
  20. cards_table = (
  21. f"CREATE TABLE IF NOT EXISTS `{settings.DB_CARD_TABLE_NAME}` ("
  22. " `id` INT AUTO_INCREMENT PRIMARY KEY,"
  23. " `card_name` VARCHAR(255) NULL COMMENT '卡牌的通用名称',"
  24. " `cardNo` VARCHAR(100) NULL COMMENT '卡牌编号',"
  25. " `card_type` VARCHAR(50) NOT NULL DEFAULT 'pokemon',"
  26. " `detection_score` DECIMAL(4, 2) NULL,"
  27. " `modified_score` DECIMAL(4, 2) NULL,"
  28. " `is_edited` BOOLEAN NOT NULL DEFAULT FALSE,"
  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. # 临时代码
  36. # ----------------- 新增无损增加字段的逻辑 -----------------
  37. try:
  38. cursor.execute(f"SHOW COLUMNS FROM `{settings.DB_CARD_TABLE_NAME}` LIKE 'cardNo'")
  39. if not cursor.fetchone():
  40. logger.info(f"表 '{settings.DB_CARD_TABLE_NAME}' 中不存在 'cardNo' 字段,正在安全添加以兼容旧数据...")
  41. # ALTER TABLE 添加列到 card_name 之后
  42. cursor.execute(
  43. f"ALTER TABLE `{settings.DB_CARD_TABLE_NAME}` ADD COLUMN `cardNo` VARCHAR(100) NULL COMMENT '卡牌编号' AFTER `card_name`")
  44. logger.info("字段 'cardNo' 添加成功。")
  45. except mysql.connector.Error as err:
  46. logger.error(f"检查/添加 'cardNo' 字段失败: {err}")
  47. # --------------------------------------------------------
  48. # 2. 创建 card_images 表 (主要计算图)
  49. card_images_table = (
  50. f"CREATE TABLE IF NOT EXISTS `{settings.DB_IMAGE_TABLE_NAME}` ("
  51. " `id` INT AUTO_INCREMENT PRIMARY KEY,"
  52. " `card_id` INT NOT NULL COMMENT '关联的卡牌ID',"
  53. " `image_type` VARCHAR(50) NOT NULL COMMENT '图片类型 (front_coaxial, back_ring, etc)',"
  54. " `image_name` VARCHAR(255) NULL,"
  55. " `image_path` VARCHAR(512) NOT NULL,"
  56. " `detection_image_path` VARCHAR(512) NULL,"
  57. " `modified_image_path` VARCHAR(512) NULL,"
  58. " `detection_json` JSON NOT NULL,"
  59. " `modified_json` JSON NULL,"
  60. " `is_edited` BOOLEAN NOT NULL DEFAULT FALSE,"
  61. " `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
  62. " `updated_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
  63. f" FOREIGN KEY (`card_id`) REFERENCES `{settings.DB_CARD_TABLE_NAME}`(`id`) ON DELETE CASCADE,"
  64. " UNIQUE KEY `uk_card_image_type` (`card_id`, `image_type`)"
  65. ") ENGINE=InnoDB COMMENT='存储主要卡牌图片及检测数据'"
  66. )
  67. cursor.execute(card_images_table)
  68. logger.info(f"数据表 '{settings.DB_IMAGE_TABLE_NAME}' 已准备就绪。")
  69. # 3. 创建 card_gray_images 表 (辅助灰度图)
  70. card_gray_images_table = (
  71. f"CREATE TABLE IF NOT EXISTS `{settings.DB_GRAY_IMAGE_TABLE_NAME}` ("
  72. " `id` INT AUTO_INCREMENT PRIMARY KEY,"
  73. " `card_id` INT NOT NULL COMMENT '关联的卡牌ID',"
  74. " `image_type` VARCHAR(50) NOT NULL COMMENT '类型 (front_gray, back_gray)',"
  75. " `image_path` VARCHAR(512) NOT NULL,"
  76. " `created_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP,"
  77. " `updated_at` TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,"
  78. f" FOREIGN KEY (`card_id`) REFERENCES `{settings.DB_CARD_TABLE_NAME}`(`id`) ON DELETE CASCADE,"
  79. " UNIQUE KEY `uk_card_gray_type` (`card_id`, `image_type`)"
  80. ") ENGINE=InnoDB COMMENT='存储辅助灰度图'"
  81. )
  82. cursor.execute(card_gray_images_table)
  83. logger.info(f"数据表 '{settings.DB_GRAY_IMAGE_TABLE_NAME}' 已准备就绪。")
  84. except mysql.connector.Error as err:
  85. logger.error(f"数据库初始化失败: {err}")
  86. exit(1)
  87. finally:
  88. if 'cursor' in locals() and cursor:
  89. cursor.close()
  90. if 'cnx' in locals() and cnx.is_connected():
  91. cnx.close()
  92. logger.info("--- 数据库初始化完成 ---")
  93. def load_database_pool():
  94. global db_connection_pool
  95. if db_connection_pool is None:
  96. logger.info("--- 创建数据库连接池 ---")
  97. try:
  98. db_connection_pool = mysql.connector.pooling.MySQLConnectionPool(
  99. pool_name="mypool",
  100. pool_size=5,
  101. **settings.DATABASE_CONFIG_WITH_DB
  102. )
  103. logger.info("--- 数据库连接池创建成功 ---")
  104. except mysql.connector.Error as err:
  105. logger.error(f"创建数据库连接池失败: {err}")
  106. exit(1)
  107. def close_database_pool():
  108. logger.info("--- 数据库连接池将自动关闭 ---")
  109. # --- FastAPI 依赖注入 ---
  110. def get_db_connection():
  111. if db_connection_pool is None:
  112. raise RuntimeError("数据库连接池未初始化")
  113. db_conn = None
  114. try:
  115. db_conn = db_connection_pool.get_connection()
  116. yield db_conn
  117. except mysql.connector.Error as err:
  118. logger.error(f"获取数据库连接失败: {err}")
  119. finally:
  120. if db_conn and db_conn.is_connected():
  121. db_conn.close()