project/tbsg/db.py

"""
数据库模块 — 与 main1.py 分离，独立保存采集数据
与美团共用同一 MySQL（drug_retrieve 库）和相同的两张表：
  - retrieve_scrape_data       商品数据主表（INSERT IGNORE，字段与美团 save_to_database 对齐）
  - retrieve_scrape_shop_info  店铺信息表（存在则补全缺失字段，不存在则插入）
饿了么平台标识：主表 platform_id='12'，店铺表 platform=12
（美团平台注释：0=全网 1=淘宝 2=京东 3=拼多多 4=美团 5=药师帮 6=1药城 7=药易购 8=药房网 9=药便宜 10=天猫药药 11=药易购）
"""
import time

import pymysql

from area import AreaService

DB_HOST = "120.24.26.108"
DB_PORT = 3307
DB_USER = "collect_user"
DB_PASSWORD = "collect123456"
DB_NAME = "drug_retrieve_test"

PLATFORM_ID = "12"   # 主表 platform_id（varchar）
PLATFORM = 12        # 店铺表 platform（int）

# 省市推断服务（与美团同款：area.py + city.json + addr_prefix.json）
_area_service = None


def _get_area_service():
    global _area_service
    if _area_service is None:
        _area_service = AreaService("city.json", "addr_prefix.json")
    return _area_service


def _infer_region(address: str):
    """从地址推断省市（美团同款逻辑），返回 (province, city, province_id, city_id)"""
    if not address:
        return "", "", 0, 0
    try:
        m = _get_area_service().search_area(address)
        if m:
            prov, city = (m.province or ""), (m.city or "")
            # 直辖市(省名==市名, 如"重庆市"): 省份只填简称不带"市"
            if prov and prov == city and prov.endswith("市"):
                prov = prov[:-1]
            return prov, city, (m.province_id or 0), (m.city_id or 0)
    except Exception as e:
        print(f"[db] 省市推断失败: {e}")
    return "", "", 0, 0


def get_conn():
    return pymysql.connect(
        host=DB_HOST, port=DB_PORT, user=DB_USER, password=DB_PASSWORD,
        database=DB_NAME, charset="utf8mb4", autocommit=False,
    )


def _to_price(v) -> float:
    """价格字符串 → decimal(10,2) 兼容的浮点，非数字返回 0.0"""
    try:
        return float(str(v).replace("￥", "").replace("¥", "").strip())
    except (ValueError, TypeError):
        return 0.0


def _today() -> str:
    return time.strftime("%Y-%m-%d")


def _now() -> str:
    return time.strftime("%Y-%m-%d %H:%M:%S")


def _license(record: dict) -> dict:
    """从采集记录中取营业执照 dict（容错）"""
    lic = record.get("license") or {}
    return lic if isinstance(lic, dict) else {}


def save_to_database(record: dict) -> bool:
    """保存商品数据到主表 retrieve_scrape_data（与美团 save_to_database 完全一致的字段，未采集的留空/默认）"""
    shop = str(record.get("shop") or "").strip()
    lic = _license(record)
    company = str(lic.get("单位名称") or "").strip()
    address = str(lic.get("地址") or "").strip()
    qualification = str(record.get("license_no") or lic.get("社会信用代码") or "").strip()
    price = _to_price(record.get("price"))
    province, city, province_id, city_id = _infer_region(address)
    sql = """
        INSERT IGNORE INTO retrieve_scrape_data (
            enterprise_id, platform_id, platform_item_id, province_id, city_id,
            province_name, city_name, area_info, product_name, product_specs,
            one_box_price, manufacture_date, expiry_date, manufacturer, approval_number,
            is_sold_out, online_posting_count, continuous_listing_count, link_url,
            store_name, store_url, shipment_province_id, shipment_province_name,
            shipment_city_id, shipment_city_name, company_name, qualification_number,
            scrape_date, min_price, number, sales, inventory, snapshot_url,
            product_brand, search_name, insert_time, update_time, collect_config_info,
            collect_equipment_account_id, collect_region_id, collect_round,
            shop_id, company_id, task_id
        ) VALUES (
            %s, %s, %s, %s, %s,
            %s, %s, %s, %s, %s,
            %s, %s, %s, %s, %s,
            %s, %s, %s, %s,
            %s, %s, %s, %s,
            %s, %s, %s, %s,
            %s, %s, %s, %s, %s, %s,
            %s, %s, %s, %s, %s, %s, %s,
            %s, %s, %s, %s
        )
    """
    params = (
        int(record.get("enterprise_id") or 0),      # enterprise_id（调度任务传入，手动模式为0）
        PLATFORM_ID,                # platform_id
        # platform_item_id：主表唯一索引 uniq_product_key 含此列，留空会与同采集日其他商品冲突（只存1条），
        # 用「店铺名_商品标题」保证唯一，重复采集同商品自动去重
        f"{shop}_{record.get('title') or ''}"[:120],
        province_id, city_id,       # province_id, city_id（执照地址推断）
        province, city,             # province_name, city_name
        address,                    # area_info（详细地址=执照地址）
        str(record.get("title") or "")[:100],   # product_name
        str(record.get("product_specs") or ""), # product_specs（任务给的规格代码）
        price,                      # one_box_price
        "",                         # manufacture_date（未采集）
        str(record.get("validity") or ""),      # expiry_date（说明书有效期）
        "",                         # manufacturer（未采集）
        str(record.get("approval_no") or ""),   # approval_number
        0,                          # is_sold_out（未采集）
        1, 1,                       # online_posting_count, continuous_listing_count（未采集）
        str(record.get("link") or ""),          # link_url
        shop,                       # store_name
        "",                         # store_url（未采集）
        0, "", 0, "",               # shipment_*（未采集）
        company,                    # company_name（执照单位名称）
        qualification,              # qualification_number
        _today(),                   # scrape_date
        price,                      # min_price
        1,                          # number（未采集）
        str(record.get("sales") or ""),         # sales（月售）
        "",                         # inventory（未采集）
        str(record.get("snapshot_url") or ""),   # snapshot_url（网页快照OSS链接）
        str(record.get("brand") or ""),         # product_brand（任务品牌）
        str(record.get("search_name") or ""),   # search_name
        _now(), _now(),             # insert_time, update_time
        str(record.get("collect_config_info") or ""),        # collect_config_info（调度任务传入）
        record.get("collect_equipment_account_id") or 0,     # collect_equipment_account_id
        record.get("collect_region_id") or 0,                # collect_region_id
        record.get("collect_round") or "",                   # collect_round
        shop, shop,                 # shop_id, company_id
        record.get("task_id") or 0, # task_id（调度任务传入）
    )
    return _exec(sql, params, "保存商品数据")


def get_existing_license(shop: str) -> dict:
    """
    查询店铺已采集的资质信息（platform+shop 唯一键），已采集过的直接从数据库获取。
    返回 {"license_no": "", "company": "", "address": ""}，无记录返回空 dict。
    """
    shop = str(shop or "").strip()
    if not shop:
        return {}
    conn = get_conn()
    try:
        cur = conn.cursor()
        cur.execute(
            "SELECT qualification_number, business_license_company, business_license_address "
            "FROM retrieve_scrape_shop_info WHERE platform=%s AND shop_id=%s LIMIT 1",
            (PLATFORM, shop),
        )
        row = cur.fetchone()
        if row:
            return {
                "license_no": str(row[0] or "").strip(),
                "company": str(row[1] or "").strip(),
                "address": str(row[2] or "").strip(),
            }
    except Exception as e:
        print(f"[db] 查询已有资质失败: {e}")
    finally:
        try:
            conn.close()
        except Exception:
            pass
    return {}


def save_shop_info_to_database(record: dict) -> bool:
    """保存店铺信息到 retrieve_scrape_shop_info（存在则补全缺失字段，不存在则插入）"""
    shop = str(record.get("shop") or "").strip()
    lic = _license(record)
    company = str(lic.get("单位名称") or "").strip()
    address = str(lic.get("地址") or "").strip()
    qualification = str(record.get("license_no") or lic.get("社会信用代码") or "").strip()
    province, city, _, _ = _infer_region(address)
    if not shop:
        return False

    conn = get_conn()
    try:
        cur = conn.cursor()
        cur.execute(
            "SELECT qualification_number, business_license_company, business_license_address, contact_address, province, city "
            "FROM retrieve_scrape_shop_info WHERE platform=%s AND shop_id=%s AND company_id=%s",
            (PLATFORM, shop, shop),
        )
        row = cur.fetchone()
        if row:
            # 已存在：用新采集到的值补全缺失字段
            update_sql = """
                UPDATE retrieve_scrape_shop_info SET
                    qualification_number = %s, business_license_company = %s,
                    business_license_address = %s, contact_address = %s,
                    province = %s, city = %s, update_time = %s
                WHERE platform=%s AND shop_id=%s AND company_id=%s
            """
            cur.execute(update_sql, (
                qualification or row[0], company or row[1],
                address or row[2], address or row[3],
                province or row[4], city or row[5], _now(),
                PLATFORM, shop, shop,
            ))
        else:
            insert_sql = """
                INSERT INTO retrieve_scrape_shop_info (
                    platform, shop_id, company_id, shop, shop_url, business_license_url,
                    province, city, contact_address,
                    qualification_number, business_license_company, business_license_address,
                    scrape_date, create_time, update_time, ai_generate
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            cur.execute(insert_sql, (
                PLATFORM, shop, shop, shop, "", None,      # platform, shop_id, company_id, shop, shop_url, business_license_url（未采集）
                province, city, address,                   # province, city（推断）, contact_address（执照地址）
                qualification, company, address,           # 资质编号, 执照公司名, 执照地址
                _today(), _now(), _now(), 0,               # scrape_date, create_time, update_time, ai_generate（未采集）
            ))
        conn.commit()
        cur.close()
        return True
    except Exception as e:
        print(f"[db] 保存店铺数据失败: {e}")
        try:
            conn.rollback()
        except Exception:
            pass
        return False
    finally:
        try:
            conn.close()
        except Exception:
            pass


def save_record(record: dict) -> None:
    """保存一条采集记录（商品数据 + 店铺信息），任一步失败不抛异常"""
    try:
        ok1 = save_to_database(record)
    except Exception as e:
        print(f"[db] 保存商品数据异常: {e}")
        ok1 = False
    try:
        ok2 = save_shop_info_to_database(record)
    except Exception as e:
        print(f"[db] 保存店铺信息异常: {e}")
        ok2 = False
    print(f"[db] 入库完成: 商品{'成功' if ok1 else '失败'} 店铺{'成功' if ok2 else '失败'} <- {record.get('shop', '')}")


def _exec(sql: str, params: tuple, tag: str) -> bool:
    conn = get_conn()
    try:
        cur = conn.cursor()
        cur.execute(sql, params)
        conn.commit()
        cur.close()
        return True
    except Exception as e:
        print(f"[db] {tag}失败: {e}")
        try:
            conn.rollback()
        except Exception:
            pass
        return False
    finally:
        try:
            conn.close()
        except Exception:
            pass