lysync 1.0.0__py3-none-any.whl

lysync/export/utils.py

@@ -0,0 +1,330 @@
+"""SQL / 数据源工具函数（从 costexcel 移植并适配新模型）。
+
+主要提供：
+* 使用 SQLAlchemy 从外部数据库拉取数据（``fetch_data_from_temp_db``）。
+* 使用 sqlglot 在最外层 SELECT 上增删 WHERE / ORDER BY / LIMIT / OFFSET。
+* 计算 id 字段、外层 count SQL、数据指纹、SQL 性能评级等。
+
+参数命名由旧的 ``type`` 改为 ``db_type``，避免与内置关键字冲突。
+"""
+from __future__ import annotations
+
+import datetime
+import hashlib
+import json
+import socket
+from urllib.parse import quote
+
+from sqlglot import exp
+from sqlglot import parse_one
+
+
+def safe_quote(string):
+    if not string:
+        return string
+    return quote(str(string))
+
+
+# 统一使用项目共享的 Jinja2 环境渲染 SQL 模板（内置 now()/today()/sqlstr() 等，
+# 可通过 lysync.export.jinja_env.register_global / register_filter 扩展）。
+from lysync.export.jinja_env import get_jinja_env, render_sql_template  # noqa: E402,F401
+
+
+
+def fetch_data_from_temp_db(db_type, host, port, user, password, database, options, query, params=None, skip_rows=0):
+    """使用 SQLAlchemy 从外部临时数据库获取数据，返回 ``list[dict]``。"""
+    from sqlalchemy import create_engine, text
+
+    try:
+        port = int(port)
+    except Exception:
+        raise ValueError(f"端口应为整数：{host}")
+    try:
+        socket.gethostbyname(host)
+    except socket.gaierror:
+        raise ValueError(f"DNS解析失败，请确认配置了正确的数据库主机地址：{host}")
+
+    if db_type == "mysql":
+        dialect = "mysql+pymysql"
+    elif db_type == "postgresql":
+        dialect = "postgresql+psycopg2"
+    elif db_type == "tsql":
+        dialect = "mssql+pymssql"
+    else:
+        raise ValueError(f"不支持的数据库类型: {db_type}")
+
+    db_url = f"{dialect}://{safe_quote(user)}:{safe_quote(password)}@{safe_quote(host)}:{port}/{safe_quote(database)}"
+    connect_args = options if isinstance(options, dict) else {}
+
+    engine = create_engine(db_url, connect_args=connect_args)
+
+    with engine.connect() as connection:
+        import re
+
+        num_params = query.count("@p")
+        param_names = [f"p{i}" for i in range(num_params)]
+        query_named = re.sub(r"@p", lambda m, c=iter(param_names): f":{next(c)}", query)
+
+        params_dict = {}
+        if params:
+            if isinstance(params, (list, tuple)):
+                if len(params) != num_params:
+                    raise ValueError(f"查询有 {num_params} 个占位符，但传入了 {len(params)} 个参数")
+                params_dict = dict(zip(param_names, params))
+            else:
+                raise TypeError(f"不支持的 params 类型: {type(params)}")
+
+        stmt = text(query_named)
+        result = connection.execute(stmt, params_dict)
+
+        columns = list(result.keys())
+        return [dict(zip(columns, row)) for index, row in enumerate(result) if index >= skip_rows]
+
+
+def add_order_by_with_sqlglot(db_type, sql, order_fields):
+    """仅替换最外层 SELECT 的 ORDER BY。
+
+    Args:
+        order_fields: ``[(field, desc), ...]``，如 ``[("c2", True), ("c0", False)]``。
+    """
+    order_exprs = []
+    for field, desc in order_fields:
+        order_expr = exp.Ordered(this=parse_one(field, read=db_type), desc=desc)
+        order_exprs.append(order_expr)
+
+    new_order = exp.Order(expressions=order_exprs)
+
+    parsed_sql = parse_one(sql, dialect=db_type)
+    select_stmt = parsed_sql if isinstance(parsed_sql, exp.Select) else parsed_sql.find(exp.Select)
+    if select_stmt:
+        select_stmt.set("order", new_order)
+    return parsed_sql.sql(dialect=db_type)
+
+
+def add_where_condition_with_sqlglot(db_type, sql, new_condition_str):
+    """仅处理最外层 SELECT 的 WHERE 子句（已有则 AND 叠加）。"""
+    parsed_sql = parse_one(sql, dialect=db_type)
+    new_condition_parsed = parse_one(new_condition_str, dialect=db_type)
+
+    select_stmt = parsed_sql if isinstance(parsed_sql, exp.Select) else parsed_sql.find(exp.Select)
+    if select_stmt:
+        existing_where = select_stmt.args.get("where")
+        if existing_where:
+            combined = exp.And(this=existing_where.this, expression=new_condition_parsed)
+            select_stmt.set("where", exp.Where(this=combined))
+        else:
+            select_stmt.set("where", exp.Where(this=new_condition_parsed))
+    return parsed_sql.sql(dialect=db_type)
+
+
+def add_limit_with_sqlglot(db_type, sql, limit_value):
+    """仅替换或添加最外层 SELECT 的 LIMIT 子句。"""
+    parsed_sql = parse_one(sql, dialect=db_type)
+    new_limit = exp.Limit(expression=exp.Literal.number(limit_value))
+
+    select_stmt = parsed_sql if isinstance(parsed_sql, exp.Select) else parsed_sql.find(exp.Select)
+    if select_stmt:
+        select_stmt.set("limit", new_limit)
+    return parsed_sql.sql(dialect=db_type)
+
+
+def add_offset_with_sqlglot(db_type, sql, offset_value):
+    """仅替换或添加最外层 SELECT 的 OFFSET 子句。"""
+    if offset_value <= 0:
+        return sql
+    parsed_sql = parse_one(sql, dialect=db_type)
+    new_offset = exp.Offset(expression=exp.Literal.number(offset_value))
+
+    select_stmt = parsed_sql if isinstance(parsed_sql, exp.Select) else parsed_sql.find(exp.Select)
+    if select_stmt:
+        select_stmt.set("offset", new_offset)
+    return parsed_sql.sql(dialect=db_type)
+
+
+def get_id_field(sql, db_type):
+    """从 SELECT 语句中解析 ``id`` 字段对应的真实列名（默认 ``id``）。"""
+    parsed = parse_one(sql, dialect=db_type)
+    if not isinstance(parsed, exp.Select):
+        raise ValueError("不是SELECT查询语句")
+    id_field = "id"
+    for expr in parsed.expressions:
+        field_name = expr.this.sql() if getattr(expr, "this", None) else str(expr)
+        alias_name = expr.alias if getattr(expr, "alias", None) else field_name
+        if alias_name == "id":
+            id_field = field_name
+    return id_field
+
+
+def build_outer_count_sql(db_type: str, sql: str, alias_name: str = "total") -> str:
+    """生成 ``SELECT count(*) AS total FROM (...)`` 形式的外层计数 SQL。"""
+    expr = parse_one(sql, dialect=db_type).select(f"count(*) as {alias_name}", append=False, dialect=db_type)
+    return expr.sql(dialect=db_type)
+
+
+def build_outer_sql(db_type: str, sql: str) -> str:
+    """把原 SQL 包成子查询并 ``SELECT * FROM (<subquery>) AS t``。"""
+    original_expr = parse_one(sql, dialect=db_type)
+    subquery_expr = original_expr.copy()
+    subquery = exp.alias_(exp.Subquery(this=subquery_expr), "t")
+    outer_select = exp.select("*").from_(subquery)
+    return outer_select.sql(dialect=db_type)
+
+
+def format_like_value(value):
+    value = value.replace("_", r"\_").replace("%", r"\%").replace("[", r"\[").replace("]", r"\]")
+    return f"%{value}%"
+
+
+def format_datasource_error_message(error_msg, db_name, db_host, db_port):
+    error_msg = error_msg.replace("(pymysql.err.OperationalError) ", "")
+    if "Access denied for user" in error_msg:
+        return "用户名或密码错误"
+    elif "Unknown database" in error_msg:
+        return f"数据库{db_name}不存在"
+    elif "Can't connect to MySQL server on" in error_msg:
+        return f"无法连接到MySQL服务器: {db_host}:{db_port}"
+    return error_msg
+
+
+def calculate_row_fingerprint(values) -> str:
+    """根据一组字段值计算行指纹（md5）。"""
+    data_to_hash = []
+    for v in values:
+        if v is None or (isinstance(v, str) and v.strip() == ""):
+            v = ""
+        data_to_hash.append(str(v))
+    hash_str = "|".join(data_to_hash)
+    return hashlib.md5(hash_str.encode("utf-8")).hexdigest()
+
+
+# ============================ SQL 性能评级（移植） ============================
+
+
+grade_order = ["A1", "A2", "B1", "B2", "C1", "C2", "D1", "D2", "E1"]
+
+
+def sql_performance_grade(explain_result):
+    """根据 EXPLAIN 结果评估 SQL 性能等级。"""
+    type_ = explain_result.get("type") or ""
+    rows = explain_result.get("rows") or 0
+    extra = explain_result.get("Extra") or ""
+
+    using_temporary = "Using temporary" in extra
+    using_filesort = "Using filesort" in extra
+
+    if type_ in ["const", "eq_ref"] and not using_filesort and not using_temporary:
+        return "A1"
+    if type_ in ["ref", "range"] and not using_filesort and not using_temporary:
+        return "A2"
+    if type_ in ["ref", "range"] and using_filesort and not using_temporary:
+        return "B1"
+    if type_ in ["ref", "range"] and not using_filesort and using_temporary:
+        return "B2"
+    if type_ in ["index", "range"] and using_filesort and using_temporary:
+        return "C1"
+    if type_ == "ALL" and not using_filesort and not using_temporary:
+        return "C2"
+    if type_ == "ALL" and (using_filesort or using_temporary):
+        if rows > 100000:
+            return "E1"
+        return "D1"
+    if type_ == "ALL" and using_filesort and using_temporary:
+        return "D2"
+    return "C1"
+
+
+def get_sql_context_for_ai(sql: str, db_type, host, port, user, password, database_name, options) -> str:
+    """为 AI 收集 SQL 上下文信息（原始 SQL、EXPLAIN 结果、表/视图创建语句）。"""
+    from django.conf import settings
+    import sqlglot
+    from sqlglot import exp
+
+    context = []
+
+    context.append("--- 原始SQL ---")
+    context.append(sql)
+    context.append("\n")
+
+    db_settings = settings.DATABASES["default"]
+    db_type_local = db_settings["ENGINE"].split(".")[-1]
+    if "postgresql" in db_type_local:
+        db_type_local = "postgresql"
+    elif "mysql" in db_type_local:
+        db_type_local = "mysql"
+
+    try:
+        # SQL 现在是 jinja2 模板，先渲染成可解析的 SQL 再做 EXPLAIN
+        rendered_sql = render_sql_template(
+            sql, start_time=datetime.datetime(2020, 1, 1), end_time=datetime.datetime(2020, 1, 2)
+        )
+        parsed = parse_one(rendered_sql, dialect=db_type)
+        if not isinstance(parsed, exp.Select):
+            raise Exception("SQL不是SELECT语句")
+        if settings.EXPORT_DATA_LIMIT:
+            rendered_sql = add_limit_with_sqlglot(db_type, rendered_sql, settings.EXPORT_DATA_LIMIT)
+        try:
+            explain_sql = f"EXPLAIN FORMAT=TRADITIONAL {rendered_sql}"
+            explain_result = fetch_data_from_temp_db(
+                db_type, host, port, user, password, database_name, options, explain_sql
+            )
+        except Exception:
+            explain_sql = f"EXPLAIN {rendered_sql}"
+            explain_result = fetch_data_from_temp_db(
+                db_type, host, port, user, password, database_name, options, explain_sql
+            )
+
+        context.append(
+            "--- EXPLAIN结果(此结果是在原始SQL渲染后增加了LIMIT，优化SQL时仅针对原始SQL优化，不要增加LIMIT) ---"
+        )
+        context.append(json.dumps(explain_result, indent=2, ensure_ascii=False))
+        context.append("\n")
+    except Exception as e:
+        context.append("--- EXPLAIN结果 ---")
+        context.append(f"获取EXPLAIN失败: {e}")
+        context.append("\n")
+
+    context.append("--- 表和视图创建SQL ---")
+    try:
+        parsed = sqlglot.parse_one(sql)
+        table_names = {t.name.lower() for t in parsed.find_all(exp.Table)}
+
+        if not table_names:
+            context.append("在SQL中未发现任何表或视图。")
+        else:
+            for name in table_names:
+                try:
+                    check_type_sql = "SELECT table_type as type FROM information_schema.tables WHERE table_name = @p"
+                    result = fetch_data_from_temp_db(
+                        db_type, host, port, user, password, database_name, options, check_type_sql, params=[name]
+                    )
+                    table_type = result[0]["type"] if result else None
+
+                    if table_type == "VIEW":
+                        get_view_sql = (
+                            "SELECT view_definition as definition FROM information_schema.views WHERE table_name = @p"
+                        )
+                        view_def_result = fetch_data_from_temp_db(
+                            db_type, host, port, user, password, database_name, options, get_view_sql, params=[name]
+                        )
+                        view_def = view_def_result[0]["definition"] if view_def_result else ""
+                        context.append(f"-- 视图: {name} --")
+                        context.append(f"CREATE VIEW {name} AS {view_def};")
+                        context.append("")
+                    elif table_type == "BASE TABLE":
+                        get_create_sql = "SHOW CREATE TABLE {}".format(name)
+                        result = fetch_data_from_temp_db(
+                            db_type, host, port, user, password, database_name, options, get_create_sql
+                        )
+                        if result:
+                            context.append(f"-- 表: {name} --")
+                            context.append(result[0]["Create Table"])
+                            context.append("")
+                    else:
+                        logger.warning(f"未找到或不支持的类型: {name} ({table_type})")
+                except Exception as e:
+                    logger.exception(str(e))
+
+    except Exception as e:
+        context.append(f"解析SQL或获取表结构失败: {e}")
+
+    return "\n".join(context)

lysync/export/writer.py

@@ -0,0 +1,143 @@
+"""数据写入层（直接导入目标库，去掉了中间的 Excel 逻辑）。
+
+将处理后的行（``list[dict]``，键为 ``SyncModelField.field_name``）写入由
+:class:`~lysync.models.SyncModelType` 的 ``app_label`` / ``model`` 字符串解析出的目标 Django 模型。
+
+写入策略：
+* 以配置了 ``is_unique=True`` 的字段组合作为唯一键做 upsert（存在则更新，不存在则新建）。
+* 每条数据的“唯一标识”由其 ``is_unique`` 字段值计算 md5（``compute_unique_hash``），
+  作为落库 / 变更检测的稳定键。
+* 计算并落库数据指纹（``SyncDataFingerprint``）用于变更审计，按
+  ``(export_task, origin_recode_unique_hash)`` 唯一。
+"""
+from __future__ import annotations
+
+import hashlib
+
+from django.db.models import Q
+
+from lysync.models import SyncDataFingerprint, SyncModelField
+from lysync.export.utils import calculate_row_fingerprint
+
+
+def compute_unique_hash(row, unique_field_names):
+    """根据唯一字段值计算稳定 md5 哈希，作为数据行的唯一标识。
+
+    空值（``None`` / 空字符串）按空串处理，与 :func:`calculate_row_fingerprint` 规则一致。
+    """
+    parts = []
+    for name in unique_field_names:
+        v = row.get(name)
+        if v is None or (isinstance(v, str) and v.strip() == ""):
+            v = ""
+        parts.append(str(v))
+    return hashlib.md5("|".join(parts).encode("utf-8")).hexdigest()
+
+
+def write_rows_to_target(model_type, processed_rows, task, org_id=0):
+    """将处理后的行写入目标模型，返回 ``(created, updated, unique_pairs)``。
+
+    Args:
+        model_type: :class:`~lysync.models.SyncModelType` 实例。
+        processed_rows: 处理后的 ``list[dict]``。
+        task: :class:`~lysync.models.SyncExportTask` 实例（用于落库指纹）。
+        org_id: 组织 ID。
+
+    Returns:
+        ``unique_pairs`` 为 ``[(unique_hash, fingerprint), ...]``，顺序与
+        ``processed_rows`` 一致，便于调用方取“最后一条”唯一标识。
+    """
+    target_model = model_type.get_target_model()
+    fields = list(SyncModelField.objects.filter(model_type=model_type).order_by("id"))
+    if not fields:
+        return 0, 0, []
+
+    unique_field_names = [f.field_name for f in fields if f.is_unique]
+    if not unique_field_names:
+        # 退化为使用第一个字段作为唯一键
+        unique_field_names = [fields[0].field_name]
+
+    # 规整每行：仅保留已知字段
+    cleaned_rows = []
+    for raw in processed_rows:
+        cleaned = {f.field_name: raw.get(f.field_name) for f in fields}
+        cleaned_rows.append((raw, cleaned))
+
+    # 查询已存在记录（按唯一键组合）
+    existing_map = {}
+    if unique_field_names:
+        q_objs = []
+        for _raw, cleaned in cleaned_rows:
+            q_objs.append(Q(**{name: cleaned[name] for name in unique_field_names}))
+        if q_objs:
+            combined = q_objs[0]
+            for q in q_objs[1:]:
+                combined |= q
+            for inst in target_model.objects.filter(combined):
+                key = tuple(getattr(inst, name) for name in unique_field_names)
+                existing_map[key] = inst
+
+    to_create = []
+    to_update = []
+    unique_pairs = []
+
+    for raw, cleaned in cleaned_rows:
+        key = tuple(cleaned[name] for name in unique_field_names)
+        inst = existing_map.get(key)
+        if inst is None:
+            inst = target_model(**cleaned)
+            to_create.append(inst)
+        else:
+            changed = False
+            for name, val in cleaned.items():
+                if getattr(inst, name) != val:
+                    setattr(inst, name, val)
+                    changed = True
+            if changed:
+                to_update.append(inst)
+
+        unique_hash = compute_unique_hash(cleaned, unique_field_names)
+        fp = calculate_row_fingerprint([cleaned[name] for name in sorted(cleaned.keys())])
+        unique_pairs.append((unique_hash, fp))
+
+    created = updated = 0
+    if to_create:
+        target_model.objects.bulk_create(to_create, batch_size=1000)
+        created = len(to_create)
+    if to_update:
+        update_fields = [f.field_name for f in fields]
+        target_model.objects.bulk_update(to_update, update_fields, batch_size=1000)
+        updated = len(to_update)
+
+    # 保存指纹
+    if unique_pairs and task is not None:
+        _save_fingerprints(task, org_id, unique_pairs)
+
+    return created, updated, unique_pairs
+
+
+def _save_fingerprints(task, org_id, unique_pairs):
+    """批量 upsert 数据指纹（按 export_task + origin_recode_unique_hash 唯一）。"""
+    existing = {
+        fp.origin_recode_unique_hash: fp
+        for fp in SyncDataFingerprint.objects.filter(
+            export_task=task, origin_recode_unique_hash__in=[p[0] for p in unique_pairs]
+        )
+    }
+    to_create = []
+    to_update = []
+    for unique_hash, fp_value in unique_pairs:
+        old = existing.get(unique_hash)
+        if old is None:
+            to_create.append(
+                SyncDataFingerprint(
+                    org_id=org_id, export_task=task, origin_recode_unique_hash=unique_hash, fingerprint=fp_value
+                )
+            )
+        elif old.fingerprint != fp_value:
+            old.fingerprint = fp_value
+            to_update.append(old)
+    if to_create:
+        SyncDataFingerprint.objects.bulk_create(to_create, batch_size=1000)
+    if to_update:
+        SyncDataFingerprint.objects.bulk_update(to_update, ["fingerprint"], batch_size=1000)

lysync/funcs/__init__.py

@@ -0,0 +1,25 @@
+"""自定义函数注册包。
+
+导入本包即会触发内置 process / collect 函数的注册，确保在 Celery 任务或 Web
+请求执行 ``process_list`` / ``custom_func`` 之前，所有函数已进入注册表。
+
+自定义函数可在宿主项目中任意模块用 ``@register_func(...)`` 注册，
+只要该模块被导入过即可（推荐在 app 的 ``ready()`` 或 ``tasks.py`` 中 import）。
+"""
+
+from .registry import (  # noqa: F401
+    COLLECT_FUNCS,
+    PROCESS_FUNCS,
+    FuncType,
+    ProcessScope,
+    get_collect_func,
+    get_process_func,
+    list_registered,
+    register_collect,
+    register_func,
+    register_process,
+)
+
+# 触发内置函数注册
+from . import collect  # noqa: F401,E402
+from . import process  # noqa: F401,E402

lysync/funcs/collect.py

@@ -0,0 +1,25 @@
+"""内置采集函数（``source_type='custom'`` 时调用）。
+
+采集函数签名：``func(**kwargs) -> list[dict]``，返回的字典列表键名应与目标模型
+字段（``SyncModelField.field_name``）一致。
+
+数据源上的 ``custom_func`` 填写这里注册的配置名；``custom_func_kwargs`` 与任务上的
+``custom_func_kwargs`` 会合并后作为关键字参数传入（任务参数优先级更高）。
+"""
+
+from __future__ import annotations
+
+from .registry import FuncType, register_func
+
+
+@register_func("sample_collect", FuncType.COLLECT)
+def sample_collect(**kwargs):
+    """示例采集函数：返回两条示例数据。
+
+    仅用于演示如何编写自定义采集函数，实际项目中请替换为真实的数据拉取逻辑
+    （调用第三方 API、读取文件、内部 RPC 等）。
+    """
+    return [
+        {"id": 1, "name": "示例-A", "value": 100, "time": "2026-01-01 00:00:00"},
+        {"id": 2, "name": "示例-B", "value": 250, "time": "2026-01-02 00:00:00"},
+    ]