first
This commit is contained in:
Your Name
@@ -9,9 +9,11 @@
|
||||
'''
|
||||
|
||||
import psycopg2
|
||||
from psycopg2.extras import Json
|
||||
import json
|
||||
import redis
|
||||
import threading
|
||||
import socket
|
||||
from loguru import logger
|
||||
from app.config import config
|
||||
from app.utils.status_codes import (
|
||||
@@ -19,6 +21,7 @@ from app.utils.status_codes import (
|
||||
DETECT_STATUS_COMPLETED,
|
||||
DETECT_STATUS_FAILED,
|
||||
DETECT_STATUS_PENDING,
|
||||
DETECT_STATUS_RUNNING,
|
||||
REGISTER_STATUS_AVAILABLE,
|
||||
REGISTER_STATUS_REGISTERED,
|
||||
REVIEW_STATUS_PENDING,
|
||||
@@ -114,7 +117,7 @@ class Database:
|
||||
self.use_bloom_filter = True
|
||||
else:
|
||||
# 如果命令不存在，说明 Redis 没有加载布隆过滤器模块
|
||||
logger.warning(f"Redis 布隆过滤器不可用: {e}，将使用普通缓存")
|
||||
logger.info(f"Redis 布隆过滤器不可用: {e}，将使用普通缓存")
|
||||
self.use_bloom_filter = False
|
||||
except Exception as e:
|
||||
logger.warning(f"初始化布隆过滤器失败: {e}")
|
||||
@@ -367,8 +370,464 @@ class Database:
|
||||
pass
|
||||
return False
|
||||
finally:
|
||||
# 将连接放回连接池
|
||||
self.close(conn, cur)
|
||||
|
||||
def ensure_cluster_runtime_tables(self):
|
||||
sql = """
|
||||
CREATE TABLE IF NOT EXISTS detect_worker_nodes (
|
||||
node_code VARCHAR(64) PRIMARY KEY,
|
||||
region VARCHAR(32) NOT NULL DEFAULT 'unknown',
|
||||
role VARCHAR(32) NOT NULL DEFAULT 'worker',
|
||||
hostname VARCHAR(255) NOT NULL DEFAULT '',
|
||||
ip VARCHAR(64) NOT NULL DEFAULT '',
|
||||
status VARCHAR(32) NOT NULL DEFAULT 'unknown',
|
||||
worker_version VARCHAR(32) NOT NULL DEFAULT '',
|
||||
current_load INTEGER NOT NULL DEFAULT 0,
|
||||
metadata_json JSONB,
|
||||
last_heartbeat_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS detect_jobs (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
job_code VARCHAR(64) NOT NULL UNIQUE,
|
||||
source VARCHAR(64) NOT NULL DEFAULT 'manual',
|
||||
plan_hash VARCHAR(128) NOT NULL DEFAULT '',
|
||||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||||
remark TEXT NOT NULL DEFAULT '',
|
||||
created_by VARCHAR(64) NOT NULL DEFAULT '',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
started_at TIMESTAMP,
|
||||
finished_at TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS detect_job_items (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
job_id BIGINT NOT NULL REFERENCES detect_jobs(id) ON DELETE CASCADE,
|
||||
domain_id BIGINT NOT NULL,
|
||||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||||
claimed_by VARCHAR(64) NOT NULL DEFAULT '',
|
||||
claim_token VARCHAR(64) NOT NULL DEFAULT '',
|
||||
lease_expires_at TIMESTAMP,
|
||||
attempt_count INTEGER NOT NULL DEFAULT 0,
|
||||
last_error TEXT NOT NULL DEFAULT '',
|
||||
result_version VARCHAR(64) NOT NULL DEFAULT '',
|
||||
started_at TIMESTAMP,
|
||||
finished_at TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
CONSTRAINT uq_detect_job_items_job_domain UNIQUE (job_id, domain_id)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_detect_job_items_status_lease
|
||||
ON detect_job_items(status, lease_expires_at);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS detect_run_events (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
job_id BIGINT REFERENCES detect_jobs(id) ON DELETE SET NULL,
|
||||
job_item_id BIGINT REFERENCES detect_job_items(id) ON DELETE SET NULL,
|
||||
node_code VARCHAR(64) NOT NULL DEFAULT '',
|
||||
event_type VARCHAR(64) NOT NULL DEFAULT '',
|
||||
level VARCHAR(16) NOT NULL DEFAULT 'info',
|
||||
message TEXT NOT NULL DEFAULT '',
|
||||
payload_json JSONB,
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS detect_sync_records (
|
||||
id BIGSERIAL PRIMARY KEY,
|
||||
sync_type VARCHAR(32) NOT NULL DEFAULT '',
|
||||
source_region VARCHAR(32) NOT NULL DEFAULT '',
|
||||
target_region VARCHAR(32) NOT NULL DEFAULT '',
|
||||
status VARCHAR(32) NOT NULL DEFAULT 'pending',
|
||||
payload_json JSONB,
|
||||
error_message TEXT NOT NULL DEFAULT '',
|
||||
created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
|
||||
);
|
||||
"""
|
||||
return self.execute(sql)
|
||||
|
||||
def register_cluster_node(self, node_code, region, role, status='online', current_load=0, metadata=None):
|
||||
try:
|
||||
ip_addr = ''
|
||||
try:
|
||||
ip_addr = socket.gethostbyname(socket.gethostname())
|
||||
except Exception:
|
||||
pass
|
||||
sql = """
|
||||
INSERT INTO detect_worker_nodes (
|
||||
node_code, region, role, hostname, ip, status, worker_version, current_load, metadata_json, last_heartbeat_at, update_time
|
||||
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s::jsonb, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)
|
||||
ON CONFLICT (node_code) DO UPDATE SET
|
||||
region = EXCLUDED.region,
|
||||
role = EXCLUDED.role,
|
||||
hostname = EXCLUDED.hostname,
|
||||
ip = EXCLUDED.ip,
|
||||
status = EXCLUDED.status,
|
||||
worker_version = EXCLUDED.worker_version,
|
||||
current_load = EXCLUDED.current_load,
|
||||
metadata_json = EXCLUDED.metadata_json,
|
||||
last_heartbeat_at = CURRENT_TIMESTAMP,
|
||||
update_time = CURRENT_TIMESTAMP
|
||||
"""
|
||||
return self.execute(
|
||||
sql,
|
||||
(
|
||||
node_code,
|
||||
region,
|
||||
role,
|
||||
socket.gethostname(),
|
||||
ip_addr,
|
||||
status,
|
||||
'0.1.0',
|
||||
int(current_load or 0),
|
||||
json.dumps(metadata or {}, ensure_ascii=False),
|
||||
),
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(f"注册检测节点失败: {e}")
|
||||
return False
|
||||
|
||||
def claim_detect_job_items(self, node_code, limit=1000, lease_seconds=3600):
|
||||
"""
|
||||
领取一批待执行的任务项。
|
||||
"""
|
||||
conn = None
|
||||
cur = None
|
||||
claim_token = f"{node_code}-{int(threading.current_thread().ident or 0)}-{int(__import__('time').time())}"
|
||||
try:
|
||||
conn, cur = self.connect()
|
||||
if not conn or not cur:
|
||||
logger.error("领取检测任务失败: 无法获取数据库连接")
|
||||
return []
|
||||
cur.execute(
|
||||
"""
|
||||
WITH picked AS (
|
||||
SELECT id
|
||||
FROM detect_job_items
|
||||
WHERE status IN ('pending', 'failed')
|
||||
AND (lease_expires_at IS NULL OR lease_expires_at < CURRENT_TIMESTAMP)
|
||||
ORDER BY create_time ASC, id ASC
|
||||
FOR UPDATE SKIP LOCKED
|
||||
LIMIT %s
|
||||
),
|
||||
updated AS (
|
||||
UPDATE detect_job_items AS item
|
||||
SET status = 'claimed',
|
||||
claimed_by = %s,
|
||||
claim_token = %s,
|
||||
lease_expires_at = CURRENT_TIMESTAMP + (%s || ' seconds')::interval,
|
||||
attempt_count = item.attempt_count + 1,
|
||||
started_at = COALESCE(item.started_at, CURRENT_TIMESTAMP),
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
FROM picked
|
||||
WHERE item.id = picked.id
|
||||
RETURNING item.id, item.job_id, item.domain_id, item.claim_token
|
||||
)
|
||||
SELECT
|
||||
updated.id,
|
||||
updated.job_id,
|
||||
updated.domain_id,
|
||||
updated.claim_token,
|
||||
domains.domain,
|
||||
domains.source_type,
|
||||
domains.register_status,
|
||||
domains.detect_status,
|
||||
domains.use_status,
|
||||
domains.expire_date,
|
||||
domains.jucha_status,
|
||||
domains.juziseo_status
|
||||
FROM updated
|
||||
JOIN domains ON domains.id = updated.domain_id
|
||||
ORDER BY updated.id ASC
|
||||
""",
|
||||
(limit, node_code, claim_token, max(60, int(lease_seconds or 3600))),
|
||||
)
|
||||
rows = cur.fetchall() or []
|
||||
touched_job_ids = sorted({row[1] for row in rows})
|
||||
for job_id in touched_job_ids:
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE detect_jobs
|
||||
SET status = CASE WHEN status = 'pending' THEN 'running' ELSE status END,
|
||||
started_at = COALESCE(started_at, CURRENT_TIMESTAMP)
|
||||
WHERE id = %s
|
||||
""",
|
||||
(job_id,),
|
||||
)
|
||||
conn.commit()
|
||||
return [
|
||||
{
|
||||
"job_item_id": row[0],
|
||||
"job_id": row[1],
|
||||
"id": row[2],
|
||||
"claim_token": row[3],
|
||||
"domain": row[4],
|
||||
"source_type": row[5],
|
||||
"register_status": row[6],
|
||||
"detect_status": row[7],
|
||||
"use_status": row[8],
|
||||
"expire_date": row[9],
|
||||
"jucha_status": row[10],
|
||||
"juziseo_status": row[11],
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
except Exception as e:
|
||||
try:
|
||||
if conn:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"领取检测任务失败: {e}")
|
||||
return []
|
||||
finally:
|
||||
self.close(conn, cur)
|
||||
|
||||
def recycle_expired_detect_job_items(self):
|
||||
"""
|
||||
回收租约过期但未完成的任务项，重新放回 pending。
|
||||
"""
|
||||
conn = None
|
||||
cur = None
|
||||
try:
|
||||
conn, cur = self.connect()
|
||||
if not conn or not cur:
|
||||
return 0
|
||||
cur.execute(
|
||||
"""
|
||||
WITH recycled AS (
|
||||
UPDATE detect_job_items
|
||||
SET status = 'pending',
|
||||
claimed_by = '',
|
||||
claim_token = '',
|
||||
lease_expires_at = NULL,
|
||||
updated_at = CURRENT_TIMESTAMP,
|
||||
last_error = CASE
|
||||
WHEN status = 'running' THEN 'lease expired while running'
|
||||
WHEN status = 'claimed' THEN 'lease expired before running'
|
||||
ELSE last_error
|
||||
END
|
||||
WHERE status IN ('claimed', 'running')
|
||||
AND lease_expires_at IS NOT NULL
|
||||
AND lease_expires_at < CURRENT_TIMESTAMP
|
||||
RETURNING job_id
|
||||
)
|
||||
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
|
||||
FROM recycled
|
||||
"""
|
||||
)
|
||||
row = cur.fetchone()
|
||||
recycled_count = int((row or [0])[0] or 0)
|
||||
touched_job_ids = (row or [0, []])[1] or []
|
||||
for job_id in touched_job_ids:
|
||||
self._refresh_detect_job_status_with_cursor(cur, job_id)
|
||||
conn.commit()
|
||||
return recycled_count
|
||||
except Exception as e:
|
||||
try:
|
||||
if conn:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"回收过期检测任务失败: {e}")
|
||||
return 0
|
||||
finally:
|
||||
self.close(conn, cur)
|
||||
|
||||
def release_detect_job_items_for_node(self, node_code):
|
||||
"""
|
||||
释放指定节点遗留的 claimed/running 任务项，供节点重启后快速恢复。
|
||||
"""
|
||||
conn = None
|
||||
cur = None
|
||||
try:
|
||||
conn, cur = self.connect()
|
||||
if not conn or not cur:
|
||||
return 0
|
||||
cur.execute(
|
||||
"""
|
||||
WITH released AS (
|
||||
UPDATE detect_job_items
|
||||
SET status = 'pending',
|
||||
claimed_by = '',
|
||||
claim_token = '',
|
||||
lease_expires_at = NULL,
|
||||
updated_at = CURRENT_TIMESTAMP,
|
||||
last_error = CASE
|
||||
WHEN status = 'running' THEN 'released after worker restart'
|
||||
WHEN status = 'claimed' THEN 'released before execution after worker restart'
|
||||
ELSE last_error
|
||||
END
|
||||
WHERE claimed_by = %s
|
||||
AND status IN ('claimed', 'running')
|
||||
RETURNING job_id
|
||||
)
|
||||
SELECT count(*), array_remove(array_agg(DISTINCT job_id), NULL)
|
||||
FROM released
|
||||
""",
|
||||
(node_code,),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
released_count = int((row or [0])[0] or 0)
|
||||
touched_job_ids = (row or [0, []])[1] or []
|
||||
for job_id in touched_job_ids:
|
||||
self._refresh_detect_job_status_with_cursor(cur, job_id)
|
||||
conn.commit()
|
||||
return released_count
|
||||
except Exception as e:
|
||||
try:
|
||||
if conn:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"释放节点遗留任务失败: {e}")
|
||||
return 0
|
||||
finally:
|
||||
self.close(conn, cur)
|
||||
|
||||
def mark_detect_job_item_running(self, job_item_id, claim_token):
|
||||
return self.execute(
|
||||
"""
|
||||
UPDATE detect_job_items
|
||||
SET status = 'running',
|
||||
updated_at = CURRENT_TIMESTAMP,
|
||||
lease_expires_at = CURRENT_TIMESTAMP + interval '1 hour'
|
||||
WHERE id = %s AND claim_token = %s
|
||||
""",
|
||||
(job_item_id, claim_token),
|
||||
)
|
||||
|
||||
def renew_detect_job_item_lease(self, job_item_id, claim_token, lease_seconds=3600):
|
||||
return self.execute(
|
||||
"""
|
||||
UPDATE detect_job_items
|
||||
SET lease_expires_at = CURRENT_TIMESTAMP + (%s || ' seconds')::interval,
|
||||
updated_at = CURRENT_TIMESTAMP
|
||||
WHERE id = %s AND claim_token = %s AND status IN ('claimed', 'running')
|
||||
""",
|
||||
(max(60, int(lease_seconds or 3600)), job_item_id, claim_token),
|
||||
)
|
||||
|
||||
def append_detect_run_event(self, job_id, job_item_id, node_code, event_type, message, level='info', payload=None):
|
||||
return self.execute(
|
||||
"""
|
||||
INSERT INTO detect_run_events (job_id, job_item_id, node_code, event_type, level, message, payload_json)
|
||||
VALUES (%s, %s, %s, %s, %s, %s, %s)
|
||||
""",
|
||||
(job_id, job_item_id, node_code, event_type, level, message, Json(payload or {})),
|
||||
)
|
||||
|
||||
def complete_detect_job_item(self, job_item_id, claim_token, final_status='completed'):
|
||||
conn = None
|
||||
cur = None
|
||||
try:
|
||||
conn, cur = self.connect()
|
||||
if not conn or not cur:
|
||||
return False
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE detect_job_items
|
||||
SET status = %s,
|
||||
finished_at = CURRENT_TIMESTAMP,
|
||||
updated_at = CURRENT_TIMESTAMP,
|
||||
lease_expires_at = NULL
|
||||
WHERE id = %s AND claim_token = %s
|
||||
RETURNING job_id
|
||||
""",
|
||||
(final_status, job_item_id, claim_token),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if row:
|
||||
self._refresh_detect_job_status_with_cursor(cur, row[0])
|
||||
conn.commit()
|
||||
return bool(row)
|
||||
except Exception as e:
|
||||
try:
|
||||
if conn:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"完成检测任务项失败: {e}")
|
||||
return False
|
||||
finally:
|
||||
self.close(conn, cur)
|
||||
|
||||
def fail_detect_job_item(self, job_item_id, claim_token, reason=''):
|
||||
conn = None
|
||||
cur = None
|
||||
try:
|
||||
conn, cur = self.connect()
|
||||
if not conn or not cur:
|
||||
return False
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE detect_job_items
|
||||
SET status = 'failed',
|
||||
last_error = %s,
|
||||
finished_at = CURRENT_TIMESTAMP,
|
||||
updated_at = CURRENT_TIMESTAMP,
|
||||
lease_expires_at = NULL
|
||||
WHERE id = %s AND claim_token = %s
|
||||
RETURNING job_id
|
||||
""",
|
||||
(str(reason or '')[:1000], job_item_id, claim_token),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
if row:
|
||||
self._refresh_detect_job_status_with_cursor(cur, row[0])
|
||||
conn.commit()
|
||||
return bool(row)
|
||||
except Exception as e:
|
||||
try:
|
||||
if conn:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"标记检测任务项失败: {e}")
|
||||
return False
|
||||
finally:
|
||||
self.close(conn, cur)
|
||||
|
||||
def _refresh_detect_job_status_with_cursor(self, cur, job_id):
|
||||
cur.execute(
|
||||
"""
|
||||
SELECT
|
||||
count(*) FILTER (WHERE status IN ('pending', 'claimed', 'running')) AS active_count,
|
||||
count(*) FILTER (WHERE status = 'failed') AS failed_count,
|
||||
count(*) FILTER (WHERE status IN ('completed', 'blacklisted')) AS done_count
|
||||
FROM detect_job_items
|
||||
WHERE job_id = %s
|
||||
""",
|
||||
(job_id,),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
active_count = int(row[0] or 0)
|
||||
failed_count = int(row[1] or 0)
|
||||
done_count = int(row[2] or 0)
|
||||
if active_count > 0:
|
||||
cur.execute(
|
||||
"UPDATE detect_jobs SET status = 'running', started_at = COALESCE(started_at, CURRENT_TIMESTAMP) WHERE id = %s",
|
||||
(job_id,),
|
||||
)
|
||||
return
|
||||
final_status = 'completed'
|
||||
if failed_count > 0 and done_count > 0:
|
||||
final_status = 'partial_failed'
|
||||
elif failed_count > 0:
|
||||
final_status = 'failed'
|
||||
cur.execute(
|
||||
"""
|
||||
UPDATE detect_jobs
|
||||
SET status = %s,
|
||||
finished_at = CURRENT_TIMESTAMP,
|
||||
started_at = COALESCE(started_at, CURRENT_TIMESTAMP)
|
||||
WHERE id = %s
|
||||
""",
|
||||
(final_status, job_id),
|
||||
)
|
||||
|
||||
def fetch_one(self, sql, params=None):
|
||||
"""
|
||||
@@ -812,6 +1271,35 @@ class Database:
|
||||
sql = "UPDATE domains SET detect_status = %s WHERE id = %s"
|
||||
return self.execute(sql, (status, domain_id))
|
||||
|
||||
def recycle_running_domains(self, target_status):
|
||||
"""
|
||||
回收异常中断后遗留的“检测中”状态。
|
||||
|
||||
:param target_status: 目标状态，通常使用 DETECT_STATUS_FAILED
|
||||
:return: int - 影响行数
|
||||
"""
|
||||
conn, cur = self.connect()
|
||||
if not conn or not cur:
|
||||
logger.error("回收检测中状态失败: 无法获取数据库连接")
|
||||
return 0
|
||||
try:
|
||||
cur.execute(
|
||||
"UPDATE domains SET detect_status = %s WHERE detect_status = %s",
|
||||
(target_status, DETECT_STATUS_RUNNING)
|
||||
)
|
||||
affected = cur.rowcount or 0
|
||||
conn.commit()
|
||||
return affected
|
||||
except Exception as e:
|
||||
try:
|
||||
conn.rollback()
|
||||
except Exception:
|
||||
pass
|
||||
logger.error(f"回收检测中状态失败: {e}")
|
||||
return 0
|
||||
finally:
|
||||
self.close(conn, cur)
|
||||
|
||||
def update_domain_third_party_status(self, domain_id, field_name, status):
|
||||
"""
|
||||
更新第三方平台检测状态。
|
||||
|
||||
Reference in New Issue
Block a user