Files
ai-security-xdr/haobang-security-dm/scripts/dm_seed_syslog_normal_data.py
王绍宏 1c07acc01d feat(syslog-consumer-rule): 增加多节点原始日志采集并修复天眼/微步JSON规则匹配
一、增加现场原始日志采集功能

1. 新增按设备采集原始日志的API:
   - 指定设备ID;
   - 采集时长最长10分钟;
   - 单次最多返回50条样本。

2. 样本选择策略:
   - 优先保留时间窗口内最早到达的N条日志;
   - 后续出现JSON解析异常、非标、规则执行异常、InfluxDB/达梦写入异常或Retry/DLT时,
     使用异常日志逐条替换最早的正常样本;
   - 异常样本同样遵循先到先得。

3. 支持三个syslog-consumer-rule实例联合采集:
   - 采集请求只需发送到任意一个实例;
   - 三个实例通过公共Redis同步采集任务并汇总候选样本;
   - 不需要配置其他实例IP;
   - Redis只保存有界候选集,不保存时间窗口内的全部日志。

4. 生命周期和资源控制:
   - 采集到期后自动停止接收新日志;
   - 默认进入30秒DRAINING阶段,等待异步处理异常完成样本替换;
   - 最终结果默认保留30分钟,之后自动删除;
   - 单实例、单会话最多保留“前N条+前N条异常”;
   - Redis故障不阻断Kafka主消费链路。

5. 返回规则解析和非标表截断之前的完整原始报文,不受
   HB_XDR_TRUNCATED字段截断影响。

6. 原始日志采集和多实例汇聚默认启用:
   SYSLOG_RAW_CAPTURE_ENABLED=true
   SYSLOG_RAW_CAPTURE_CLUSTER_ENABLED=true

二、修复天眼/微步JSON规则匹配问题

1. JSON规则配置了cropper时,先判断原始报文是否包含对应head_key;
   不匹配的规则不再将当前报文误判为非法JSON。

2. 同一设备绑定多条规则时,单条不适用或解析失败的JSON规则不再立即终止整个规则链,
   程序会继续匹配后续适用规则。

3. 只有全部适用JSON规则均解析失败、且没有其他规则成功处理时,才写入
   “JSON解析异常”非标记录。

4. 对cropper截取后仍包含Syslog或设备前缀的内容,在JSON专用路径中重新收敛到
   第一个“{”至最后一个“}”,确保交给Jackson的是完整外层JSON。

5. 修复以下现场场景:
   - 天眼规则129截取结果仍带有alarm前缀;
   - 微步设备同时绑定固定偏移规则和Threatbook/TDP JSON规则;
   - 不适用的固定偏移规则抢先产生JSON解析异常;
   - Syslog头加JSON正文不能继续匹配后续适用规则。

三、增加现场采集工具及运维说明

新增upgrade-doc/v5:
   - README.md:升级、配置、生命周期和故障排查说明;
   - capture-device-raw-logs.sh:一键创建采集会话、等待完成并自动落盘。

脚本输出:
   - result.json:完整采集结果;
   - samples.jsonl:每行一条样本及处理状态;
   - raw/*.log:每条未截断原始报文;
   - summary.txt:采集数量、异常数量和替换数量;
   - latest.json:中断或超时时的最后有效快照。

四、兼容性与升级范围

- 兼容JDK 8;
- 无需新增或修改达梦数据库表;
- 无需创建新的Kafka Topic;
- 需要滚动升级全部syslog-consumer-rule实例;
- 三个实例必须连接相同的Redis;
- 本次提交不包含syslog-serve的高内存、WAL或JVM参数调整。
2026-09-14 23:11:45 +08:00

1017 lines
34 KiB
Python
Executable File
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
#!/usr/bin/env python3
"""向达梦 syslog_normal_data 写入 2026-09-20 的大记录测试数据。
单条记录的目标大小可由用户设置,默认每一整行的所有字符字段内容合计
512 KiB,而不是每个字段 512 KiB。脚本从达梦系统视图读取实际列结构,
为除 created_at(使用数据库默认值)外的字段生成非空值。为避免一次占用约 5 GiB
内存及产生超大事务,10,000 条记录按小批次生成、写入和提交。
脚本会自动识别本机达梦安装目录、设置运行库路径并加载 dmPython;如果只有
达梦自带的驱动源码,则在 /tmp 中离线编译后直接加载,不执行联网安装。
"""
from __future__ import annotations
import concurrent.futures
import getpass
import importlib
import os
import queue
import random
import re
import subprocess
import sys
import tempfile
import threading
import time
import uuid
from dataclasses import dataclass
from datetime import datetime, timedelta
from pathlib import Path
from typing import Dict, List, Optional, Sequence, Tuple
from urllib.parse import urlparse
TARGET_OWNER = "ECOSYS"
TARGET_TABLE_NAME = "syslog_normal_data"
TARGET_TABLE = '"ECOSYS"."syslog_normal_data"'
TARGET_PARTITION = '"ECOSYS"."syslog_normal_data_p_20260920"'
TARGET_DAY_START = datetime(2026, 9, 20, 0, 0, 0)
TARGET_DAY_END = datetime(2026, 9, 21, 0, 0, 0)
DEFAULT_RECORD_SIZE_KIB = 512
MIN_RECORD_SIZE_KIB = 1
MAX_RECORD_SIZE_KIB = 10 * 1024
BATCH_RECORD_COUNT = 10_000
DEFAULT_BATCH_SIZE = 200
DEFAULT_CONCURRENCY = 1
MAX_CONCURRENCY = 32
PROGRESS_EVERY = 100
DM_RUNTIME_MARKER = "DM_SEED_RUNTIME_READY"
DM_HOME_CANDIDATES = (
Path("/dmdba/dmdbms"),
Path("/home/dmdba/dmdbms"),
Path("/opt/dmdbms"),
)
@dataclass(frozen=True)
class ConnectionTarget:
host: str
port: int
schema: Optional[str] = None
@dataclass(frozen=True)
class ColumnInfo:
name: str
data_type: str
data_length: Optional[int]
@dataclass(frozen=True)
class WorkerStats:
generation_elapsed: float
insert_elapsed: float
commit_elapsed: float
committed_count: int
def find_dm_home() -> Optional[Path]:
"""查找本机达梦安装目录,优先采用已有 DM_HOME。"""
candidates = []
configured_home = os.environ.get("DM_HOME")
if configured_home:
candidates.append(Path(configured_home))
candidates.extend(DM_HOME_CANDIDATES)
for candidate in candidates:
resolved = candidate.expanduser().resolve()
if (resolved / "bin" / "disql").exists() or (
resolved / "drivers" / "python" / "dmPython"
).exists():
return resolved
return None
def prepare_dm_runtime_environment() -> Optional[Path]:
"""自动设置达梦运行库路径,并通过重启 Python 让动态加载器生效。"""
dm_home = find_dm_home()
if dm_home is None or os.environ.get(DM_RUNTIME_MARKER) == "1":
return dm_home
library_directories = [
path
for path in (dm_home / "bin", dm_home / "drivers" / "dpi")
if path.is_dir()
]
child_environment = os.environ.copy()
child_environment["DM_HOME"] = str(dm_home)
child_environment[DM_RUNTIME_MARKER] = "1"
current_library_path = child_environment.get("LD_LIBRARY_PATH", "")
library_path_parts = [str(path) for path in library_directories]
if current_library_path:
library_path_parts.append(current_library_path)
child_environment["LD_LIBRARY_PATH"] = os.pathsep.join(library_path_parts)
script_path = str(Path(__file__).resolve())
print(
f"[准备] 已识别达梦目录 {dm_home},正在自动加载达梦运行库环境……",
flush=True,
)
os.execve(
sys.executable,
[sys.executable, script_path, *sys.argv[1:]],
child_environment,
)
return dm_home # pragma: no cover - os.execve 成功后不会返回
def try_import_dm_python(search_paths: Sequence[Path] = ()):
for path in search_paths:
path_text = str(path)
if path.is_dir() and path_text not in sys.path:
sys.path.insert(0, path_text)
importlib.invalidate_caches()
sys.modules.pop("dmPython", None)
sys.modules.pop("dpi", None)
return importlib.import_module("dmPython")
def existing_dm_python_paths(dm_home: Path) -> List[Path]:
source_directory = dm_home / "drivers" / "python" / "dmPython"
paths = [source_directory]
if source_directory.is_dir():
paths.extend(
path
for path in source_directory.glob("build/lib*")
if path.is_dir()
)
paths.extend(
module.parent
for module in source_directory.rglob("dmPython*.so")
)
# 保持顺序并去重。
return list(dict.fromkeys(paths))
def build_dm_python_offline(dm_home: Path) -> Path:
"""使用达梦安装介质中的源码在 /tmp 离线编译扩展,不执行 pip。"""
source_directory = dm_home / "drivers" / "python" / "dmPython"
setup_file = source_directory / "setup.py"
if not setup_file.is_file():
raise RuntimeError(f"未找到达梦 Python 驱动源码: {setup_file}")
user_id = os.getuid() if hasattr(os, "getuid") else 0
python_tag = f"py{sys.version_info.major}{sys.version_info.minor}"
build_root = (
Path(tempfile.gettempdir())
/ f"dm_seed_dmpython_{user_id}_{python_tag}"
)
build_library = build_root / "lib"
build_temporary = build_root / "temp"
build_library.mkdir(parents=True, exist_ok=True)
build_temporary.mkdir(parents=True, exist_ok=True)
command = [
sys.executable,
str(setup_file),
"build_ext",
"--build-lib",
str(build_library),
"--build-temp",
str(build_temporary),
]
print(
f"[准备] 未检测到可直接加载的 dmPython,正在使用 {setup_file} "
"离线编译(以下为实时编译输出)……",
flush=True,
)
try:
result = subprocess.run(
command,
cwd=str(source_directory),
env=os.environ.copy(),
check=False,
timeout=300,
)
except subprocess.TimeoutExpired as exc:
raise RuntimeError(
"达梦自带 dmPython 离线编译超过 300 秒,已停止等待。"
"请检查编译器或 Python 开发环境是否卡住"
) from exc
if result.returncode != 0:
raise RuntimeError(
"达梦自带 dmPython 离线编译失败。服务器可能缺少 gcc、"
f"Python 开发头文件或 setuptools,退出码={result.returncode}。"
"具体原因见上方实时编译输出"
)
print(f"[成功] dmPython 离线编译完成,输出目录: {build_library}", flush=True)
return build_library
def load_dm_python(dm_home: Optional[Path]):
"""自动加载已安装、达梦自带或现场离线编译的 dmPython。"""
initial_error: Optional[Exception] = None
try:
module = try_import_dm_python()
print(f"[成功] 已加载现有 dmPython: {module.__file__}", flush=True)
return module
except Exception as exc:
initial_error = exc
if dm_home is None:
raise RuntimeError(
"未找到 dmPython,也未找到达梦安装目录。已自动检查 "
"/dmdba/dmdbms、/home/dmdba/dmdbms 和 /opt/dmdbms"
) from initial_error
try:
module = try_import_dm_python(existing_dm_python_paths(dm_home))
print(
f"[成功] 已从达梦安装目录加载 dmPython: {module.__file__}",
flush=True,
)
return module
except Exception as exc:
print(
f"[准备] 达梦目录中的驱动尚不能直接加载: "
f"{type(exc).__name__}: {exc}",
flush=True,
)
build_library = build_dm_python_offline(dm_home)
try:
module = try_import_dm_python([build_library])
print(
f"[成功] 已从本机达梦安装介质离线加载 dmPython: {module.__file__}",
flush=True,
)
return module
except Exception as exc:
raise RuntimeError(
f"dmPython 已离线编译,但加载失败: {exc}"
) from exc
def parse_connection_address(raw_address: str) -> ConnectionTarget:
"""解析 host、host:port 或 jdbc:dm://host:port[/schema]。"""
address = raw_address.strip()
if not address:
raise ValueError("连接地址不能为空")
if address.lower().startswith("jdbc:dm://"):
address = address[len("jdbc:") :]
elif "://" not in address:
address = "dm://" + address
parsed = urlparse(address)
if parsed.scheme.lower() != "dm" or not parsed.hostname:
raise ValueError(
"连接地址格式不正确,请使用 IP:端口,"
"例如 192.168.4.117:5237"
)
try:
port = parsed.port or 5236
except ValueError as exc:
raise ValueError("端口必须是 1 到 65535 的整数") from exc
if not 1 <= port <= 65535:
raise ValueError("端口必须在 1 到 65535 之间")
schema = parsed.path.strip("/") or None
if schema and "/" in schema:
raise ValueError("连接地址中只能包含一个 schema 名称")
return ConnectionTarget(host=parsed.hostname, port=port, schema=schema)
def parse_record_size(raw_value: str) -> int:
"""解析 512、512KB、512KiB、1MB 或 1MiB,返回字节数。"""
value = raw_value.strip()
if not value:
return DEFAULT_RECORD_SIZE_KIB * 1024
match = re.fullmatch(r"(\d+)\s*(KB|KIB|MB|MIB)?", value, re.IGNORECASE)
if match is None:
raise ValueError("整行大小格式错误,请输入例如 256、512KB 或 1MB")
number = int(match.group(1))
unit = (match.group(2) or "KIB").upper()
multiplier = 1024**2 if unit in {"MB", "MIB"} else 1024
size_bytes = number * multiplier
minimum = MIN_RECORD_SIZE_KIB * 1024
maximum = MAX_RECORD_SIZE_KIB * 1024
if not minimum <= size_bytes <= maximum:
raise ValueError(
f"单条整行目标大小必须在 {MIN_RECORD_SIZE_KIB}KB 到 "
f"{MAX_RECORD_SIZE_KIB // 1024}MB 之间"
)
return size_bytes
def format_size(size_bytes: int) -> str:
if size_bytes % (1024**2) == 0:
return f"{size_bytes // (1024**2):,} MiB"
return f"{size_bytes / 1024:,.2f} KiB"
def prompt_connection() -> Tuple[ConnectionTarget, str, str, int, int, int]:
print("达梦 syslog_normal_data 大记录测试数据写入工具")
print("连接地址支持:IP:端口 或 jdbc:dm://IP:端口[/schema]")
address = input("达梦连接地址: ").strip()
username = input("用户名: ").strip()
password = getpass.getpass("密码(输入内容不会显示): ")
raw_record_size = input(
f"单条整行目标大小(所有字符字段合计;直接回车默认 "
f"{DEFAULT_RECORD_SIZE_KIB}KB,"
f"范围 {MIN_RECORD_SIZE_KIB}KB~{MAX_RECORD_SIZE_KIB // 1024}MB,"
"支持如 256KB、1MB): "
).strip()
raw_batch_size = input(
f"每批写入条数(直接回车默认 {DEFAULT_BATCH_SIZE}): "
).strip()
raw_concurrency = input(
f"并发数(直接回车默认 {DEFAULT_CONCURRENCY},1 表示不并发,"
f"范围 1~{MAX_CONCURRENCY}): "
).strip()
if not username:
raise ValueError("用户名不能为空")
if not password:
raise ValueError("密码不能为空")
record_size_bytes = parse_record_size(raw_record_size)
try:
batch_size = (
DEFAULT_BATCH_SIZE if not raw_batch_size else int(raw_batch_size)
)
except ValueError as exc:
raise ValueError("每批写入条数必须是整数") from exc
if not 1 <= batch_size <= BATCH_RECORD_COUNT:
raise ValueError(
f"每批写入条数必须在 1 到 {BATCH_RECORD_COUNT:,} 之间"
)
try:
concurrency = (
DEFAULT_CONCURRENCY if not raw_concurrency else int(raw_concurrency)
)
except ValueError as exc:
raise ValueError("并发数必须是整数") from exc
if not 1 <= concurrency <= MAX_CONCURRENCY:
raise ValueError(f"并发数必须在 1 到 {MAX_CONCURRENCY} 之间")
return (
parse_connection_address(address),
username,
password,
batch_size,
record_size_bytes,
concurrency,
)
def confirm_execution(
target: ConnectionTarget,
username: str,
batch_size: int,
record_size_bytes: int,
concurrency: int,
) -> bool:
schema_info = f",连接 schema={target.schema}" if target.schema else ""
payload_gib = (
(BATCH_RECORD_COUNT + 1) * record_size_bytes / (1024**3)
)
batch_payload_mib = batch_size * record_size_bytes / (1024**2)
batch_count = (BATCH_RECORD_COUNT + batch_size - 1) // batch_size
effective_concurrency = min(concurrency, batch_count)
inflight_payload_mib = batch_payload_mib * effective_concurrency
print("\n即将执行以下操作:")
print(f" 数据库: {target.host}:{target.port}{schema_info}")
print(f" 用户名: {username}")
print(f" 目标表: {TARGET_TABLE}")
print(f" 目标日分区: {TARGET_PARTITION}")
print(" log_time: 全部位于 2026-09-20")
print(" 写入数量: 1 条验证记录 + 10,000 条批量记录")
print(
f" 单条整行目标内容: {format_size(record_size_bytes)} "
"(所有字符字段合计,不是每个字段)"
)
print(
f" 允许的整行大小范围: {MIN_RECORD_SIZE_KIB} KiB~"
f"{MAX_RECORD_SIZE_KIB // 1024} MiB"
)
print(" 其他字段: 尽量全部非空")
print(
f" 10,001 行字符内容总量: 约 {payload_gib:.2f} GiB;"
"数据库实际占用还包含数值字段、行和 LOB 开销"
)
print(
f" 提交策略: 每 {batch_size:,} 条提交一次,"
f"约 {batch_payload_mib:,.2f} MiB 字符内容/批"
)
print(
f" 写入并发: {effective_concurrency} "
f"(用户设置 {concurrency},每个并发使用独立连接;"
"另保留 1 个验证连接)"
)
print(
f" 并发批次字符内容上限: 约 {inflight_payload_mib:,.2f} MiB,"
"不含 Python 和驱动额外内存"
)
print(" 中断处理: 当前批次回滚,此前已提交批次会保留")
answer = input("确认表空间充足并继续?请输入 yes: ").strip().lower()
return answer in {"yes", "y"}
def connect_dm(dm_python, target: ConnectionTarget, username: str, password: str):
properties = {
"user": username,
"password": password,
"server": target.host,
"port": target.port,
"autoCommit": False,
}
if target.schema:
properties["schema"] = target.schema
return dm_python.connect(**properties)
def quote_identifier(identifier: str) -> str:
return '"' + identifier.replace('"', '""') + '"'
def load_table_columns(cursor) -> List[ColumnInfo]:
sql = (
"SELECT COLUMN_NAME, DATA_TYPE, DATA_LENGTH "
"FROM ALL_TAB_COLUMNS "
"WHERE OWNER = ? AND TABLE_NAME = ? ORDER BY COLUMN_ID"
)
candidates = (
(TARGET_OWNER, TARGET_TABLE_NAME),
(TARGET_OWNER, TARGET_TABLE_NAME.upper()),
)
for owner, table_name in candidates:
cursor.execute(sql, (owner, table_name))
rows = cursor.fetchall()
if rows:
return [
ColumnInfo(
name=str(row[0]),
data_type=str(row[1]).upper(),
data_length=None if row[2] is None else int(row[2]),
)
for row in rows
]
raise RuntimeError(
f"无法从 ALL_TAB_COLUMNS 读取 {TARGET_TABLE} 的列结构;"
"请确认表名、权限和当前连接实例"
)
def test_connection_and_schema(cursor) -> List[ColumnInfo]:
cursor.execute("SELECT 1 FROM DUAL")
row = cursor.fetchone()
if row is None or int(row[0]) != 1:
raise RuntimeError(f"连接测试查询返回异常结果: {row!r}")
try:
cursor.execute(
f'SELECT "id", "log_time" FROM {TARGET_PARTITION} WHERE 1 = 0'
)
except Exception as exc:
raise RuntimeError(
f"2026-09-20 日分区 {TARGET_PARTITION} 不存在或当前用户无权访问;"
"脚本不会自动执行 DDL,请先确认分区已创建"
) from exc
columns = load_table_columns(cursor)
names = {column.name.lower() for column in columns}
missing = {"id", "created_at", "log_time"} - names
if missing:
raise RuntimeError(f"目标表缺少必要字段: {', '.join(sorted(missing))}")
return columns
def is_character(column: ColumnInfo) -> bool:
data_type = column.data_type
return "CHAR" in data_type or "TEXT" in data_type or "CLOB" in data_type
def is_large_text(column: ColumnInfo) -> bool:
return "TEXT" in column.data_type or "CLOB" in column.data_type
def make_ascii_payload(length: int, label: str) -> str:
if length <= 0:
return ""
seed = f"dm-load-test|{label}|"
repeats = length // len(seed) + 1
return (seed * repeats)[:length]
def make_short_character_value(column: ColumnInfo, run_token: str) -> str:
name = column.name.lower()
max_length = max(1, column.data_length or 64)
if "ip" in name:
value = "192.0.2.10"
elif name == "syslog_uuid":
value = f"{run_token}-0000-4000-8000-000000000000"
elif "port" in name:
value = "5237"
elif "md5" in name:
value = "0123456789abcdef0123456789abcdef"
else:
value = f"load-test-{run_token}-{column.name}"
return value.encode("ascii", errors="ignore").decode("ascii")[:max_length]
def build_character_values(
columns: Sequence[ColumnInfo],
run_token: str,
example_id: str,
record_size_bytes: int,
) -> Dict[str, str]:
"""生成字符字段模板,使字符内容总量等于用户指定大小。"""
values: Dict[str, str] = {}
large_text_columns: List[ColumnInfo] = []
used_bytes = len(example_id.encode("utf-8"))
for column in columns:
name = column.name.lower()
if name in {"id", "created_at"} or not is_character(column):
continue
if is_large_text(column):
large_text_columns.append(column)
continue
value = make_short_character_value(column, run_token)
values[column.name] = value
used_bytes += len(value.encode("utf-8"))
if not large_text_columns:
raise RuntimeError("目标表没有 TEXT/CLOB 字段,无法生成指定大小的记录")
remaining = record_size_bytes - used_bytes
if remaining < len(large_text_columns):
raise RuntimeError("字符字段固定内容已经超过目标记录大小")
base_size, extra_count = divmod(remaining, len(large_text_columns))
for index, column in enumerate(large_text_columns):
length = base_size + (1 if index < extra_count else 0)
values[column.name] = make_ascii_payload(
length, f"{run_token}|{column.name}"
)
return values
def build_insert_sql(columns: Sequence[ColumnInfo]) -> str:
names = ", ".join(quote_identifier(column.name) for column in columns)
placeholders = ", ".join("?" for _ in columns)
return f"INSERT INTO {TARGET_TABLE} ({names}) VALUES ({placeholders})"
def value_for_column(
column: ColumnInfo,
record_id: str,
log_time: datetime,
index: int,
character_values: Dict[str, str],
):
name = column.name.lower()
data_type = column.data_type
if name == "id":
return record_id
if is_character(column):
return character_values[column.name]
if "TIMESTAMP" in data_type or "DATETIME" in data_type or data_type == "DATE":
return log_time
if "BIT" in data_type or "BOOL" in data_type:
return index % 2
if any(kind in data_type for kind in ("INT", "NUMBER", "DECIMAL", "NUMERIC")):
if "port" in name:
return 10_000 + index % 50_000
return 1_000_000 + index
if any(kind in data_type for kind in ("DOUBLE", "FLOAT", "REAL")):
return float(index) + 0.125
if "BLOB" in data_type or "BINARY" in data_type:
return b"dm-load-test"
raise RuntimeError(
f"暂不支持字段 {column.name} 的数据类型 {column.data_type},"
"为避免静默写入 NULL,脚本已停止"
)
def build_row(
columns: Sequence[ColumnInfo],
character_values: Dict[str, str],
record_id: str,
log_time: datetime,
index: int,
) -> Tuple[object, ...]:
return tuple(
value_for_column(column, record_id, log_time, index, character_values)
for column in columns
)
def character_bytes_in_row(
columns: Sequence[ColumnInfo], row: Sequence[object]
) -> int:
return sum(
len(value.encode("utf-8"))
for column, value in zip(columns, row)
if is_character(column) and isinstance(value, str)
)
def insert_and_verify_one(
connection,
cursor,
columns: Sequence[ColumnInfo],
insert_sql: str,
run_token: str,
record_size_bytes: int,
) -> None:
test_id = f"dm_test_{run_token}"
test_log_time = datetime(2026, 9, 20, 12, 34, 56, 123456)
character_values = build_character_values(
columns, run_token, test_id, record_size_bytes
)
row_values = build_row(
columns, character_values, test_id, test_log_time, -1
)
try:
cursor.execute(insert_sql, row_values)
connection.commit()
except Exception:
connection.rollback()
raise
cursor.execute(
f'SELECT "id", "log_time", "created_at" '
f'FROM {TARGET_TABLE} '
'WHERE "id" = ? AND "log_time" >= ? AND "log_time" < ?',
(test_id, TARGET_DAY_START, TARGET_DAY_END),
)
result = cursor.fetchone()
if result is None:
raise RuntimeError("单条记录已提交,但提交后未查询到该记录")
payload_bytes = character_bytes_in_row(columns, row_values)
print(f"[成功] 单条记录写入并查询成功")
print(f" id={result[0]}, log_time={result[1]}, created_at={result[2]}")
print(f" 该整行所有字符字段合计={payload_bytes:,} 字节")
def query_batch_summary(cursor, run_token: str):
id_prefix = f"dm_seed_{run_token}_"
cursor.execute(
f'SELECT COUNT(*), MIN("log_time"), MAX("log_time") '
f'FROM {TARGET_TABLE} WHERE "id" BETWEEN ? AND ?',
(f"{id_prefix}00000", f"{id_prefix}09999"),
)
return cursor.fetchone()
def insert_and_verify_batch(
verification_connection,
verification_cursor,
dm_python,
target: ConnectionTarget,
username: str,
password: str,
columns: Sequence[ColumnInfo],
insert_sql: str,
run_token: str,
batch_size: int,
record_size_bytes: int,
concurrency: int,
) -> None:
example_id = f"dm_seed_{run_token}_00000"
character_values = build_character_values(
columns, run_token, example_id, record_size_bytes
)
random_source = random.Random(run_token)
day_microseconds = 24 * 60 * 60 * 1_000_000
time_offsets = [
random_source.randrange(day_microseconds)
for _ in range(BATCH_RECORD_COUNT)
]
batch_starts = list(range(0, BATCH_RECORD_COUNT, batch_size))
effective_concurrency = min(concurrency, len(batch_starts))
work_queue = queue.Queue()
for batch_start in batch_starts:
work_queue.put(batch_start)
stop_event = threading.Event()
progress_lock = threading.Lock()
progress = {"committed": 0, "next": PROGRESS_EVERY}
def report_progress(batch_count: int) -> None:
with progress_lock:
progress["committed"] += batch_count
committed = progress["committed"]
if committed >= progress["next"] or committed == BATCH_RECORD_COUNT:
print(
f" 已提交 {committed:,}/{BATCH_RECORD_COUNT:,} 条",
flush=True,
)
while progress["next"] <= committed:
progress["next"] += PROGRESS_EVERY
def run_worker(worker_number: int) -> WorkerStats:
worker_connection = None
worker_cursor = None
generation_elapsed = 0.0
insert_elapsed = 0.0
commit_elapsed = 0.0
committed_count = 0
try:
worker_connection = connect_dm(
dm_python, target, username, password
)
worker_cursor = worker_connection.cursor()
while not stop_event.is_set():
try:
batch_start = work_queue.get_nowait()
except queue.Empty:
break
batch_count = min(
batch_size, BATCH_RECORD_COUNT - batch_start
)
generation_started_at = time.perf_counter()
rows = []
for index in range(batch_start, batch_start + batch_count):
log_time = TARGET_DAY_START + timedelta(
microseconds=time_offsets[index]
)
record_id = f"dm_seed_{run_token}_{index:05d}"
rows.append(
build_row(
columns,
character_values,
record_id,
log_time,
index,
)
)
generation_elapsed += (
time.perf_counter() - generation_started_at
)
insert_started_at = time.perf_counter()
worker_cursor.executemany(insert_sql, rows)
insert_elapsed += time.perf_counter() - insert_started_at
commit_started_at = time.perf_counter()
worker_connection.commit()
commit_elapsed += time.perf_counter() - commit_started_at
committed_count += batch_count
report_progress(batch_count)
return WorkerStats(
generation_elapsed=generation_elapsed,
insert_elapsed=insert_elapsed,
commit_elapsed=commit_elapsed,
committed_count=committed_count,
)
except Exception as exc:
stop_event.set()
if worker_connection is not None:
try:
worker_connection.rollback()
except Exception:
pass
raise RuntimeError(
f"并发工作线程 {worker_number} 写入失败: {exc}"
) from exc
finally:
if worker_cursor is not None:
try:
worker_cursor.close()
except Exception:
pass
if worker_connection is not None:
try:
worker_connection.close()
except Exception:
pass
total_started_at = time.perf_counter()
print(
f"\n开始流式生成并写入 {BATCH_RECORD_COUNT:,} 条数据,"
f"每批 {batch_size:,} 条,并发数 {effective_concurrency}……"
)
worker_stats: List[WorkerStats] = []
first_error: Optional[Exception] = None
write_phase_started_at = time.perf_counter()
with concurrent.futures.ThreadPoolExecutor(
max_workers=effective_concurrency,
thread_name_prefix="dm-seed",
) as executor:
futures = [
executor.submit(run_worker, worker_number)
for worker_number in range(1, effective_concurrency + 1)
]
for future in concurrent.futures.as_completed(futures):
try:
worker_stats.append(future.result())
except Exception as exc:
stop_event.set()
if first_error is None:
first_error = exc
write_phase_elapsed = time.perf_counter() - write_phase_started_at
if first_error is not None:
print(
f"[提示] 失败线程的当前批次已回滚,共已提交 "
f"{progress['committed']:,} 条;"
f"ID 前缀为 dm_seed_{run_token}_",
file=sys.stderr,
)
raise first_error
# 结束验证连接上此前的只读事务,确保看到并发连接刚提交的数据。
verification_connection.commit()
summary = query_batch_summary(verification_cursor, run_token)
if summary is None or int(summary[0]) != BATCH_RECORD_COUNT:
actual = 0 if summary is None else int(summary[0])
raise RuntimeError(
f"提交后校验失败:期望 {BATCH_RECORD_COUNT:,} 条,实际 {actual:,} 条"
)
total_elapsed = time.perf_counter() - total_started_at
generation_elapsed = sum(item.generation_elapsed for item in worker_stats)
insert_elapsed = sum(item.insert_elapsed for item in worker_stats)
commit_elapsed = sum(item.commit_elapsed for item in worker_stats)
rows_per_second = BATCH_RECORD_COUNT / write_phase_elapsed
milliseconds_per_row = write_phase_elapsed * 1_000 / BATCH_RECORD_COUNT
content_mib_per_second = (
BATCH_RECORD_COUNT * record_size_bytes
/ (1024**2)
/ write_phase_elapsed
)
sample_row = build_row(
columns,
character_values,
example_id,
TARGET_DAY_START,
0,
)
sample_payload_bytes = character_bytes_in_row(columns, sample_row)
print(f"[成功] {BATCH_RECORD_COUNT:,} 条大记录写入并查询成功")
print(f" 数量={int(summary[0]):,}")
print(f" 每一整行所有字符字段合计={sample_payload_bytes:,} 字节")
print(f" 最早 log_time={summary[1]}")
print(f" 最晚 log_time={summary[2]}")
print(f" 实际写入并发数={effective_concurrency}")
print(f" 各线程累计数据生成耗时={generation_elapsed:.3f} 秒")
print(f" 各线程累计批量插入耗时={insert_elapsed:.3f} 秒")
print(f" 各线程累计事务提交耗时={commit_elapsed:.3f} 秒")
print(
f" 并发写入墙钟耗时(连接、生成、插入、提交)="
f"{write_phase_elapsed:.3f} 秒"
)
print(f" 平均每条墙钟耗时={milliseconds_per_row:.3f} 毫秒")
print(f" 平均写入速度={rows_per_second:,.2f} 条/秒")
print(f" 字段内容吞吐={content_mib_per_second:,.2f} MiB/秒")
print(f" 生成、写入、提交及校验总耗时={total_elapsed:.3f} 秒")
print(f" 本次批量 ID 前缀=dm_seed_{run_token}_")
def main() -> int:
dm_home = prepare_dm_runtime_environment()
if dm_home is not None and os.environ.get(DM_RUNTIME_MARKER) == "1":
print(
f"[成功] 达梦运行库环境已自动就绪: {dm_home}",
flush=True,
)
try:
(
target,
username,
password,
batch_size,
record_size_bytes,
concurrency,
) = prompt_connection()
except (EOFError, KeyboardInterrupt):
print("\n操作已取消")
return 130
except ValueError as exc:
print(f"[错误] {exc}", file=sys.stderr)
return 2
try:
if not confirm_execution(
target,
username,
batch_size,
record_size_bytes,
concurrency,
):
print("操作已取消,未连接数据库,也未写入任何数据")
return 0
except (EOFError, KeyboardInterrupt):
print("\n操作已取消,未连接数据库,也未写入任何数据")
return 130
try:
dm_python = load_dm_python(dm_home)
except Exception as exc:
print(f"[错误] 无法自动加载达梦 Python 驱动: {exc}", file=sys.stderr)
return 3
connection = None
cursor = None
run_token = uuid.uuid4().hex[:12]
try:
print("\n[下一步] 正在连接达梦数据库并读取实际表结构……", flush=True)
connection = connect_dm(dm_python, target, username, password)
cursor = connection.cursor()
all_columns = test_connection_and_schema(cursor)
insert_columns = [
column
for column in all_columns
if column.name.lower() != "created_at"
]
insert_sql = build_insert_sql(insert_columns)
print(
f"[成功] 连接、日分区及表结构检查通过,共读取 "
f"{len(all_columns)} 个字段,本次填写 {len(insert_columns)} 个字段;"
"created_at 使用数据库默认值"
)
insert_and_verify_one(
connection,
cursor,
insert_columns,
insert_sql,
run_token,
record_size_bytes,
)
insert_and_verify_batch(
connection,
cursor,
dm_python,
target,
username,
password,
insert_columns,
insert_sql,
run_token,
batch_size,
record_size_bytes,
concurrency,
)
print(f"\n全部完成:本次共写入 {BATCH_RECORD_COUNT + 1:,} 条大记录。")
return 0
except KeyboardInterrupt:
if connection is not None:
connection.rollback()
print(
"\n操作被中断,仅当前未提交批次已回滚;此前已提交记录会保留。",
file=sys.stderr,
)
return 130
except Exception as exc:
if connection is not None:
try:
connection.rollback()
except Exception:
pass
print(f"\n[失败] {type(exc).__name__}: {exc}", file=sys.stderr)
print(
"仅当前未提交事务已回滚;单条验证记录和此前成功提交的批次会保留。",
file=sys.stderr,
)
return 1
finally:
if cursor is not None:
try:
cursor.close()
except Exception:
pass
if connection is not None:
try:
connection.close()
except Exception:
pass
if __name__ == "__main__":
raise SystemExit(main())