xhs_crawler/db.py

from typing import List

from tortoise import Tortoise, run_async

from config.db_config import *
from tools import utils


def get_platform_models() -> List[str]:
    models = ["aerich.models", "store.xhs", "store.douyin", "store.bilibili", "store.kuaishou", "store.weibo"]
    return models


async def init_db(create_db: bool = False) -> None:
    await Tortoise.init(
        db_url=RELATION_DB_URL,
        modules={'models': get_platform_models()},
        _create_db=create_db
    )

async def close() -> None:
    await Tortoise.close_connections()

async def init():
    await init_db(create_db=True)
    await Tortoise.generate_schemas()
    await init_aerich()
    utils.logger.info("[db.init] Init DB Success!")

TORTOISE_ORM_CONFIG = {
    "connections": {"default": RELATION_DB_URL},
    "apps": {
        "models": {
            "models": get_platform_models(),
            "default_connection": "default",
        },
    },
}

async def init_aerich():
    Tortoise.get_connection('default').execute_insert("""INSERT INTO aerich VALUES(1,'0_20240302150109_init.py','models','{"models.Aerich": {"name": "models.Aerich", "app": "models", "table": "aerich", "abstract": false, "description": null, "docstring": null, "unique_together": [], "indexes": [], "pk_field": {"name": "id", "field_type": "IntField", "db_column": "id", "python_type": "int", "generated": true, "nullable": false, "unique": true, "indexed": true, "default": null, "description": null, "docstring": null, "constraints": {"ge": 1, "le": 2147483647}, "db_field_types": {"": "INT"}}, "data_fields": [{"name": "version", "field_type": "CharField", "db_column": "version", "python_type": "str", "generated": false, "nullable": false, "unique": false, "indexed": false, "default": null, "description": null, "docstring": null, "constraints": {"max_length": 255}, "db_field_types": {"": "VARCHAR(255)"}}, {"name": "app", "field_type": "CharField", "db_column": "app", "python_type": "str", "generated": false, "nullable": false, "unique": false, "indexed": false, "default": null, "description": null, "docstring": null, "constraints": {"max_length": 100}, "db_field_types": {"": "VARCHAR(100)"}}, {"name": "content", "field_type": "JSONField", "db_column": "content", "python_type": "Union[dict, list]", "generated": false, "nullable": false, "unique": false, "indexed": false, "default": null, "description": null, "docstring": null, "constraints": {}, "db_field_types": {"": "JSON", "mssql": "NVARCHAR(MAX)", "oracle": "NCLOB", "postgres": "JSONB"}}], "fk_fields": [], "backward_fk_fields": [], "o2o_fields": [], "backward_o2o_fields": [], "m2m_fields": []}, "models.XHSNote": {"name": "models.XHSNote", "app": "models", "table": "xhs_note", "abstract": false, "description": "\u5c0f\u7ea2\u4e66\u7b14\u8bb0", "docstring": null, "unique_together": [], "indexes": [], "pk_field": {"name": "id", "field_type": "IntField", "db_column": "id", "python_type": "int", "generated": true, "nullable": false, "unique": true, "indexed": true, "default": null, "description": "\u81ea\u589eID", "docstring": null, "constraints": {"ge": 1, "le": 2147483647}, "db_field_types": {"": "INT"}}, "data_fields": [{"name": "user_id", "field_type": "CharField", "db_column": "user_id", "python_type": "str", "generated": false, "nullable": false, "unique": false, "indexed": false, "default": null, "description": "\u7528\u6237ID", "docstring": null, "constraints": {"max_length": 64}, "db_field_types": {"": "VARCHAR(64)"}}, {"name": "nickname", "field_type": "CharField", "db_column": "nickname", "python_type": "str", "generated": false, "nullable": true, "unique": false, "indexed": false, "default": null, "description": "\u7528\u6237\u6635\u79f0", "docstring": null, "constraints": {"max_length": 64}, "db_field_types": {"": "VARCHAR(64)"}}, {"name": "avatar", "field_type": "CharField", "db_column": "avatar", "python_type": "str", "generated": false, "nullable": true, "unique": false, "indexed": false, "default": null, "description": "\u7528\u6237\u5934\u50cf\u5730\u5740", "docstring": null, "constraints": {"max_length": 255}, "db_field_types": {"": "VARCHAR(255)"}}, {"name": "ip_location", "field_type": "CharField", "db_column": "ip_location", "python_type": "str", "generated": false, "nullable": true, "unique": false, "indexed": false, "default": null, "description": "\u8bc4\u8bba\u65f6\u7684IP\u5730\u5740", "docstring": null, "constraints": {"max_length": 255}, "db_field_types": {"": "VARCHAR(255)"}}, {"name": "add_ts", "field_type": "BigIntField", "db_column": "add_ts", "python_type": "int", "generated": false, "nullable": false, "unique": false, "indexed": false, "default": null, "description": "\u8bb0\u5f55\u6dfb\u52a0\u65f6\u95f4\u6233", "docstring": null, "constraints": {"ge": -9223372036854775808, "le": 9223372036854775807}, "db_field_types": {"": "BIGINT", "oracle": "INT"}}, {"name": "last_modify_ts", "field_type": "BigIntField", "db_column": "last_modify_ts", "python_type": "int", "generated": false, "nullable": false, "unique": false, "indexed": false, "default":


if __name__ == '__main__':
    run_async(init())