482 lines
23 KiB
Python
482 lines
23 KiB
Python
#!/usr/bin/env python3
|
||
"""Transactional DingDrive -> cloud -> LAN replication. No app/network monkeypatch imports."""
|
||
import argparse
|
||
import csv
|
||
import fcntl
|
||
import gzip
|
||
import hashlib
|
||
import importlib.util
|
||
import io
|
||
import json
|
||
import os
|
||
from pathlib import Path
|
||
import re
|
||
import subprocess
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
from datetime import datetime, date, timedelta
|
||
from decimal import Decimal
|
||
from zoneinfo import ZoneInfo
|
||
|
||
import pymysql
|
||
|
||
ROOT = Path(__file__).resolve().parents[1]
|
||
STATE = ROOT / 'var/ding-sync'
|
||
SSH = ['ssh', '-i', '/home/dieya/.ssh/id_ed25519_dycg_restore', '-o', 'BatchMode=yes',
|
||
'-o', 'ConnectTimeout=15', 'ubuntu@82.156.135.91']
|
||
REMOTE = '/opt/dy-dashboard/var/ding-sync'
|
||
DWS = '/home/ubuntu/.npm-global/lib/node_modules/dingtalk-workspace-cli/vendor/dws'
|
||
SPACE = '28814565790'
|
||
FOLDERS = {'inventory': 'NkDwLng8ZLGrXmrnFxYq3vRqVKMEvZBY',
|
||
'outbound': '3NwLYZXWynyBjRB4tZjwNxDvVkyEqBQm'}
|
||
SOURCES = {'wdt', 'jdbz', 'jdgj', 'jdjwt', 'tmgj', 'dymuto', 'dyqsimelr', 'jdsmj', 'jdmrt'}
|
||
OUTBOUND_WAREHOUSES = {
|
||
'tmgj': '9（外仓-海外仓）天猫菜鸟仓',
|
||
'jdgj': '11（外仓-海外仓）京东自营保税仓',
|
||
'jdjwt': '14（外仓-国内仓）京东POP仓',
|
||
'jdbz': '13（外仓-国内仓）京东自营仓',
|
||
'jdmrt': '14（外仓-国内仓）京东POP仓',
|
||
'jdsmj': '13（外仓-国内仓）京东自营仓',
|
||
'dymuto': '026 (外仓-海外仓)抖音自营仓',
|
||
'dyqsimelr': '026 (外仓-海外仓)抖音自营仓',
|
||
}
|
||
PATTERNS = {'inventory': re.compile(r'^(20\d{6})-wdt-kcmx-([\w]+)\.csv$'),
|
||
'outbound': re.compile(r'^(20\d{6})-([a-z0-9]+)-cksj\.csv$')}
|
||
COLS = {'inventory': ['data_date', 'source', 'warehouse_code', 'goods_code', 'product_code',
|
||
'product_name', 'brand', 'cost_price', 'available_qty', 'production_in_transit',
|
||
'purchase_in_transit', 'category1', 'category2', 'category3', 'created_time'],
|
||
'outbound': ['货品编号', '货品名称', '品牌', '仓库', '店铺', '一级类目', '二级类目',
|
||
'三级类目', '货品数量', '货品成本', '日期', '数据来源', '数据类型']}
|
||
TABLES = {'inventory': 'inventory_detail', 'outbound': 'daily_outbound'}
|
||
SPECIALS = {'POZ-FM-0184-TM-SM': '.055', 'POZ-XFS-9014-TM-MD': '.062',
|
||
'POZ-FM-9742-TM-MD': '.181', 'POZ-MM-0580-TM-MD': '.22',
|
||
**{f'POZ-SKS-{n}-LD': '.1' for n in ('3374', '3367', '3350', '3343')}}
|
||
DWS_RETRY_DELAYS = (2, 5, 10)
|
||
|
||
|
||
def now():
|
||
return datetime.now(ZoneInfo('Asia/Shanghai'))
|
||
|
||
|
||
def command(args, timeout=180, retry_delays=()):
|
||
"""Run a command, optionally retrying without exposing credential-bearing output."""
|
||
attempts = len(retry_delays) + 1
|
||
failure = 'failed'
|
||
for attempt in range(attempts):
|
||
try:
|
||
result = subprocess.run(args, capture_output=True, timeout=timeout)
|
||
if not result.returncode:
|
||
return result.stdout
|
||
failure = f'failed (exit {result.returncode})'
|
||
except subprocess.TimeoutExpired:
|
||
failure = f'timed out after {timeout}s'
|
||
if attempt < len(retry_delays):
|
||
time.sleep(retry_delays[attempt])
|
||
# Never echo command output: authentication errors may contain tokens.
|
||
raise RuntimeError(f'{Path(args[0]).name} {failure} after {attempts} attempts')
|
||
|
||
|
||
def dws_command(args, timeout=180):
|
||
# A single daily job may recover from brief DingDrive/IPaaS failures. This
|
||
# is bounded in-process retrying, not polling or an additional schedule.
|
||
return command(args, timeout=timeout, retry_delays=DWS_RETRY_DELAYS)
|
||
|
||
|
||
def encode(value):
|
||
return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(',', ':'), default=str).encode()
|
||
|
||
|
||
def atomic(path, value):
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
fd, tmp = tempfile.mkstemp(dir=path.parent)
|
||
try:
|
||
with os.fdopen(fd, 'wb') as out:
|
||
out.write(value if isinstance(value, bytes) else encode(value))
|
||
out.flush()
|
||
os.fsync(out.fileno())
|
||
os.replace(tmp, path)
|
||
finally:
|
||
if os.path.exists(tmp):
|
||
os.unlink(tmp)
|
||
|
||
|
||
def read_json(path, default=None):
|
||
return json.loads(path.read_bytes()) if path.exists() else default
|
||
|
||
|
||
def clean(v):
|
||
return str(v or '').replace('"', '').replace('=', '').strip()
|
||
|
||
|
||
def number(v, integer=False):
|
||
value = Decimal(clean(v).replace(',', '') or '0')
|
||
if not value.is_finite() or (integer and value != value.to_integral_value()):
|
||
raise ValueError('Invalid numeric CSV field')
|
||
return int(value) if integer else format(value.quantize(Decimal('.0001')), 'f')
|
||
|
||
|
||
def bad(code):
|
||
return clean(code).upper().rstrip(':：') in {'', 'NA', 'N/A', '合计'}
|
||
|
||
|
||
def csv_reader(path):
|
||
raw = path.read_bytes()
|
||
for encoding in ('utf-8-sig', 'gb18030'):
|
||
try:
|
||
content = raw.decode(encoding)
|
||
break
|
||
except UnicodeDecodeError:
|
||
continue
|
||
else:
|
||
raise ValueError('Unrecognized CSV encoding')
|
||
reader = csv.DictReader(io.StringIO(content))
|
||
if not reader.fieldnames:
|
||
raise ValueError('Missing CSV header')
|
||
reader.fieldnames = [clean(x).lstrip('\ufeff') for x in reader.fieldnames]
|
||
return reader
|
||
|
||
|
||
def pick(row, keys):
|
||
return next((clean(row[k]) for k in keys if clean(row.get(k))), '')
|
||
|
||
|
||
def parse(path, kind, day, key, products):
|
||
reader = csv_reader(path)
|
||
headers = set(reader.fieldnames)
|
||
if kind == 'inventory':
|
||
required = {'货品编号', '货品名称', '成本价', '采购在途', '生产在途量'}
|
||
qty_col = '库存' if '库存' in headers else '可用库存'
|
||
required.add(qty_col)
|
||
elif key == 'wdt':
|
||
required = {'货品编号', '货品名称', '货品数量', '货品成本'}
|
||
else:
|
||
required = set()
|
||
if not headers.intersection({'货品编号', '产品编码', '商品编码', '编码', 'product_code', 'goods_code'}):
|
||
raise ValueError(f'{path.name}: Missing product code column')
|
||
if not headers.intersection({'货品数量', '数量', '出库量', 'qty', 'quantity'}):
|
||
raise ValueError(f'{path.name}: Missing quantity column')
|
||
if required - headers:
|
||
raise ValueError(f'{path.name}: missing columns {sorted(required - headers)}; retained existing data')
|
||
rows = []
|
||
valid_codes = 0
|
||
for row in reader:
|
||
code = pick(row, ('货品编号', '产品编码', '商品编码', '编码', 'product_code', 'goods_code'))
|
||
if bad(code):
|
||
continue
|
||
valid_codes += 1
|
||
if kind == 'inventory':
|
||
name = clean(row.get('货品名称'))
|
||
if name in {'-', '--'}:
|
||
continue
|
||
ct = clean(row.get('单品创建时间'))
|
||
created = None
|
||
if ct:
|
||
created = datetime.fromisoformat(ct).strftime('%Y-%m-%d %H:%M:%S')
|
||
rows.append([day, 'wdt', key, clean(row.get('商家编码')), code, name,
|
||
clean(row.get('品牌')), number(SPECIALS.get(code, row.get('成本价'))),
|
||
number(row.get(qty_col), True), number(row.get('生产在途量'), True),
|
||
number(row.get('采购在途'), True), clean(row.get('一级类目')),
|
||
clean(row.get('二级类目')), clean(row.get('三级类目')), created])
|
||
else:
|
||
if key == 'wdt' and clean(row.get('日期')):
|
||
row_day = clean(row['日期'])[:10].replace('/', '-')
|
||
if date.fromisoformat(row_day).isoformat() != day:
|
||
raise ValueError(f'{path.name}: row date does not match filename')
|
||
qty = number(pick(row, ('货品数量', '数量', '出库量', 'qty', 'quantity')), True)
|
||
p = row if key == 'wdt' else products.get(code, {})
|
||
name = clean(p.get('货品名称') if key == 'wdt' else p.get('product_name'))
|
||
brand = clean(p.get('品牌') if key == 'wdt' else p.get('brand'))
|
||
if qty <= 0 or name in {'-', '--'} or '潘思宇代发' in brand:
|
||
continue
|
||
cats = ['一级类目', '二级类目', '三级类目'] if key == 'wdt' else ['category1', 'category2', 'category3']
|
||
warehouse = clean(row.get('仓库')) if key == 'wdt' else OUTBOUND_WAREHOUSES.get(key, '')
|
||
rows.append([code, name or code, brand, warehouse,
|
||
clean(row.get('店铺')) if key == 'wdt' else '', *[clean(p.get(x)) for x in cats],
|
||
qty, number(p.get('货品成本') if key == 'wdt' else p.get('cost_price')),
|
||
day, key, '出库数据'])
|
||
if not rows and (kind == 'inventory' or not valid_codes):
|
||
raise ValueError(f'{path.name}: no valid records; refusing replacement')
|
||
return rows
|
||
|
||
|
||
def connect(local):
|
||
spec = importlib.util.spec_from_file_location('sync_config', ROOT / 'config.py')
|
||
mod = importlib.util.module_from_spec(spec)
|
||
spec.loader.exec_module(mod)
|
||
config = dict(mod.DB_CONFIG)
|
||
if local:
|
||
config.update(user='dy_test', password=(ROOT / '.local-test-db-password').read_text().strip(),
|
||
unix_socket='/run/mysqld/mysqld.sock')
|
||
config.update(cursorclass=pymysql.cursors.DictCursor, autocommit=False)
|
||
return pymysql.connect(**config)
|
||
|
||
|
||
def where(packet):
|
||
if packet['kind'] == 'inventory':
|
||
return 'data_date=%s AND source=%s AND warehouse_code=%s', (packet['date'], 'wdt', packet['key'])
|
||
return '`日期`=%s AND `数据来源`=%s', (packet['date'], packet['key'])
|
||
|
||
|
||
def normalize(rows, kind):
|
||
money_index = 7 if kind == 'inventory' else 9
|
||
result = []
|
||
for row in rows:
|
||
r = [str(v) if isinstance(v, (date, datetime)) else v for v in row]
|
||
r[money_index] = number(r[money_index])
|
||
result.append(r)
|
||
return result
|
||
|
||
|
||
def digest(rows, kind):
|
||
return hashlib.sha256(b'\n'.join(sorted(encode(r) for r in normalize(rows, kind)))).hexdigest()
|
||
|
||
|
||
def db_rows(conn, p):
|
||
condition, params = where(p)
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT `' + '`,`'.join(COLS[p['kind']]) + '` FROM ' + TABLES[p['kind']] + ' WHERE ' + condition, params)
|
||
return [[r[c] for c in COLS[p['kind']]] for r in cur.fetchall()]
|
||
|
||
|
||
def refresh():
|
||
# Load cache modules directly, avoiding services/__init__ and any DB/AI side effects.
|
||
for name, function in (('cache', 'refresh_all_caches'), ('simple_cache', 'clear')):
|
||
spec = importlib.util.spec_from_file_location('sync_' + name, ROOT / 'services' / (name + '.py'))
|
||
module = importlib.util.module_from_spec(spec)
|
||
spec.loader.exec_module(module)
|
||
getattr(module, function)()
|
||
|
||
|
||
def apply(conn, packets, dry_run=False):
|
||
changed = []
|
||
for p in packets:
|
||
if p['kind'] not in COLS or p['date'] != date.fromisoformat(p['date']).isoformat():
|
||
raise ValueError('Invalid packet target')
|
||
if any(len(r) != len(COLS[p['kind']]) for r in p['rows']):
|
||
raise ValueError('Invalid packet shape')
|
||
old = db_rows(conn, p)
|
||
if digest(old, p['kind']) != digest(p['rows'], p['kind']):
|
||
changed.append(p)
|
||
if not dry_run:
|
||
condition, params = where(p)
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT * FROM ' + TABLES[p['kind']] + ' WHERE ' + condition, params)
|
||
atomic(STATE / 'backups' / (now().strftime('%Y%m%dT%H%M%S%f') + '.json.gz'),
|
||
gzip.compress(encode({'target': {k: p[k] for k in ('kind', 'date', 'key')}, 'rows': cur.fetchall()})))
|
||
cur.execute('DELETE FROM ' + TABLES[p['kind']] + ' WHERE ' + condition, params)
|
||
sql = 'INSERT INTO ' + TABLES[p['kind']] + ' (`' + '`,`'.join(COLS[p['kind']]) + '`) VALUES (' + ','.join(['%s'] * len(COLS[p['kind']])) + ')'
|
||
if p['rows']:
|
||
cur.executemany(sql, p['rows'])
|
||
if digest(db_rows(conn, p), p['kind']) != digest(p['rows'], p['kind']):
|
||
raise RuntimeError('Database readback mismatch; transaction rolled back')
|
||
if dry_run:
|
||
conn.rollback()
|
||
else:
|
||
conn.commit()
|
||
# Refresh even after interrupted commit; publication/ack happens only after refresh succeeds.
|
||
if packets:
|
||
refresh()
|
||
return [{'kind': p['kind'], 'date': p['date'], 'key': p['key'], 'rows': len(p['rows'])} for p in changed]
|
||
|
||
|
||
def list_files(kind):
|
||
items, cursor, seen = [], None, set()
|
||
while True:
|
||
args = [DWS, 'drive', 'list', '--space-id', SPACE, '--folder', FOLDERS[kind], '--limit', '50', '--format', 'json']
|
||
if cursor:
|
||
args.extend(['--cursor', cursor])
|
||
response = json.loads(dws_command(args))
|
||
if response.get('success') is False or 'result' not in response:
|
||
raise RuntimeError('DingDrive list failed')
|
||
items.extend(response['result'].get('items', []))
|
||
cursor = response['result'].get('nextToken')
|
||
if not cursor:
|
||
return items
|
||
if cursor in seen:
|
||
raise RuntimeError('Repeated DingDrive pagination cursor')
|
||
seen.add(cursor)
|
||
|
||
|
||
def cloud(conn, args):
|
||
config = read_json(STATE / 'settings.json')
|
||
target = args.target_date or now().date()
|
||
old = read_json(STATE / 'manifest.json', {'entries': {}})
|
||
# Publish only this run's dates. Missing files never become a catch-up queue.
|
||
entries = {}
|
||
packets, missing = [], []
|
||
with conn.cursor() as cur:
|
||
cur.execute('SELECT product_code,product_name,brand,category1,category2,category3,cost_price FROM products')
|
||
products = {r['product_code']: r for r in cur.fetchall()}
|
||
for kind in FOLDERS:
|
||
files = {}
|
||
for item in list_files(kind):
|
||
m = PATTERNS[kind].fullmatch(item.get('name', ''))
|
||
if not m or item.get('type') != 'FILE':
|
||
continue
|
||
day, key = m.groups()
|
||
day = datetime.strptime(day, '%Y%m%d').date()
|
||
key = 'dymuto' if key == 'dymtuo' else key
|
||
if kind == 'outbound' and key not in SOURCES:
|
||
continue
|
||
identity = (day, key)
|
||
if identity not in files or item['modifyTime'] > files[identity]['modifyTime']:
|
||
files[identity] = item
|
||
begin = target - timedelta(days=kind == 'outbound')
|
||
end = target - timedelta(days=kind == 'outbound')
|
||
day = begin
|
||
while day <= end:
|
||
expected = set(config['warehouses']) if kind == 'inventory' else SOURCES
|
||
uploaded = {k for d, k in files if d == day}
|
||
for key in sorted(expected - uploaded):
|
||
missing.append(f'{kind}:{day}:{key}')
|
||
for key in sorted(uploaded):
|
||
item = files[(day, key)]
|
||
identity = f'{kind}:{day}:{key}'
|
||
fingerprint = hashlib.sha256(encode([item['fileId'], item.get('md5'), item['modifyTime'], item['fileSize']])).hexdigest()
|
||
existing = old['entries'].get(identity)
|
||
if existing and existing['source_fingerprint'] == fingerprint:
|
||
entries[identity] = existing
|
||
continue
|
||
dest = STATE / 'downloads' / item['name']
|
||
dest.parent.mkdir(parents=True, exist_ok=True)
|
||
dws_command([DWS, 'drive', 'download', '--node', item['fileId'], '--space-id', SPACE, '--output', str(dest), '--format', 'json'])
|
||
if dest.stat().st_size != item['fileSize']:
|
||
raise ValueError(f'{item["name"]}: incomplete download')
|
||
p = {'kind': kind, 'date': str(day), 'key': key,
|
||
'rows': parse(dest, kind, str(day), key, products)}
|
||
blob = gzip.compress(encode(p), mtime=0)
|
||
sha = hashlib.sha256(blob).hexdigest()
|
||
atomic(STATE / 'packets' / (sha + '.json.gz'), blob)
|
||
entries[identity] = {'packet': sha, 'source_fingerprint': fingerprint,
|
||
'rows': len(p['rows']), 'source_name': item['name']}
|
||
packets.append(p)
|
||
day += timedelta(days=1)
|
||
changed = apply(conn, packets, args.dry_run)
|
||
if not args.dry_run and not packets and not (STATE / 'manifest.json').exists():
|
||
refresh()
|
||
coverage = {kind: sum(identity.startswith(f'{kind}:{target - timedelta(days=kind == "outbound")}:')
|
||
for identity in entries) for kind in FOLDERS}
|
||
status = {'checked_at': now().isoformat(), 'target_inventory': str(target),
|
||
'target_outbound': str(target - timedelta(days=1)), 'missing': missing,
|
||
'state': 'success', 'missing_policy': 'skip', 'changed': changed,
|
||
'cache_refreshed': not args.dry_run, 'files': len(entries), 'coverage': coverage}
|
||
if not args.dry_run:
|
||
manifest = {'entries': entries, 'missing': missing, 'missing_policy': 'skip', 'target_inventory': str(target),
|
||
'target_outbound': status['target_outbound'], 'coverage': coverage}
|
||
manifest['version'] = hashlib.sha256(encode(manifest)).hexdigest()
|
||
status['version'] = manifest['version']
|
||
atomic(STATE / 'status.json', status)
|
||
atomic(STATE / 'manifest.json', manifest)
|
||
print(json.dumps(status, ensure_ascii=False))
|
||
|
||
|
||
def local(conn, args):
|
||
manifest = json.loads(command(SSH + ['cat', REMOTE + '/manifest.json']))
|
||
cloud_status = json.loads(command(SSH + ['cat', REMOTE + '/status.json']))
|
||
target = (args.target_date or now().date()).isoformat()
|
||
if manifest.get('target_inventory') != target or manifest.get('missing_policy') != 'skip':
|
||
raise RuntimeError('No requested-date single-run publication; other dates will not be imported')
|
||
if not cloud_status.get('cache_refreshed') or cloud_status.get('version') != manifest['version']:
|
||
raise RuntimeError('Cloud update incomplete or failed; waiting for a validated publication')
|
||
old = read_json(STATE / 'local-status.json', {})
|
||
packets = []
|
||
ledger = read_json(STATE / 'local-ledger.json', {})
|
||
for identity, entry in sorted(manifest['entries'].items()):
|
||
sha = entry['packet']
|
||
if not re.fullmatch('[a-f0-9]{64}', sha):
|
||
raise ValueError('Invalid packet hash')
|
||
if ledger.get(identity) == sha:
|
||
continue
|
||
path = STATE / 'packets' / (sha + '.json.gz')
|
||
if not path.exists():
|
||
atomic(path, command(SSH + ['cat', REMOTE + '/packets/' + path.name]))
|
||
blob = path.read_bytes()
|
||
if hashlib.sha256(blob).hexdigest() != sha:
|
||
# Keep the corrupt artifact for diagnosis; let the next attempt fetch a clean copy.
|
||
path.rename(path.with_suffix('.corrupt-' + now().strftime('%Y%m%dT%H%M%S%f')))
|
||
raise ValueError('Packet checksum mismatch')
|
||
p = json.loads(gzip.decompress(blob))
|
||
if identity != f'{p["kind"]}:{p["date"]}:{p["key"]}':
|
||
raise ValueError('Packet identity mismatch')
|
||
packets.append(p)
|
||
if not packets and old.get('version') == manifest['version'] and old.get('cache_refreshed'):
|
||
# Keep online evidence current, without re-invalidating valid caches.
|
||
status = {**old, 'checked_at': now().isoformat()}
|
||
else:
|
||
changed = apply(conn, packets, args.dry_run)
|
||
status = {'checked_at': now().isoformat(), 'version': manifest['version'],
|
||
'target_inventory': manifest['target_inventory'], 'target_outbound': manifest['target_outbound'],
|
||
'missing': manifest['missing'], 'state': 'success', 'missing_policy': 'skip',
|
||
'changed': changed, 'cache_refreshed': not args.dry_run, 'files': len(manifest['entries']),
|
||
'coverage': manifest['coverage']}
|
||
if not args.dry_run:
|
||
atomic(STATE / 'local-ledger.json', {k: v['packet'] for k, v in manifest['entries'].items()})
|
||
atomic(STATE / 'local-status.json', status)
|
||
# Exact local readback + cache completion are acknowledged to cloud.
|
||
result = subprocess.run(SSH + ['umask 077; tee ' + REMOTE + '/local-ack.json.tmp >/dev/null && mv ' + REMOTE + '/local-ack.json.tmp ' + REMOTE + '/local-ack.json'], input=encode(status), capture_output=True, timeout=45)
|
||
if result.returncode:
|
||
raise RuntimeError('Cloud acknowledgement failed; no automatic retry')
|
||
if old.get('version') != status['version']:
|
||
env = dict(os.environ, DBUS_SESSION_BUS_ADDRESS=f'unix:path=/run/user/{os.getuid()}/bus', XDG_RUNTIME_DIR=f'/run/user/{os.getuid()}')
|
||
msg = (f"库存 {status['target_inventory']}：{status['coverage']['inventory']}仓库；"
|
||
f"出库 {status['target_outbound']}：{status['coverage']['outbound']}来源。两端已同步这些文件并刷新缓存。")
|
||
if status['missing']:
|
||
msg += f" 本次未上传的{len(status['missing'])}份文件已跳过，不等待补录。"
|
||
try:
|
||
notification = subprocess.run(['notify-send', '采购数据同步结果', msg], env=env, capture_output=True, timeout=10)
|
||
status['desktop_notification'] = 'delivered' if notification.returncode == 0 else 'desktop_unavailable'
|
||
except (OSError, subprocess.TimeoutExpired):
|
||
status['desktop_notification'] = 'desktop_unavailable'
|
||
atomic(STATE / 'local-status.json', status)
|
||
print(json.dumps(status, ensure_ascii=False))
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('mode', choices=['cloud', 'local'])
|
||
parser.add_argument('--dry-run', action='store_true')
|
||
parser.add_argument('--target-date', type=date.fromisoformat,
|
||
help='Manual recovery date (YYYY-MM-DD); scheduled runs omit this option')
|
||
args = parser.parse_args()
|
||
if args.target_date and args.target_date > now().date():
|
||
parser.error('--target-date cannot be in the future')
|
||
os.umask(0o077)
|
||
STATE.mkdir(parents=True, exist_ok=True)
|
||
with open(STATE / (args.mode + '.lock'), 'w') as lock:
|
||
try:
|
||
fcntl.flock(lock, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
except BlockingIOError:
|
||
print('Already running; skipped')
|
||
return
|
||
conn = None
|
||
try:
|
||
conn = connect(args.mode == 'local')
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT GET_LOCK('dy_ding_daily_sync', 0) AS acquired")
|
||
if not cur.fetchone()['acquired']:
|
||
raise RuntimeError('Another database import is running')
|
||
(cloud if args.mode == 'cloud' else local)(conn, args)
|
||
except Exception as error:
|
||
if conn:
|
||
conn.rollback()
|
||
failure = {'checked_at': now().isoformat(), 'state': 'error', 'error': str(error), 'cache_refreshed': False}
|
||
if not args.dry_run:
|
||
atomic(STATE / ('status.json' if args.mode == 'cloud' else 'local-status.json'), failure)
|
||
if args.mode == 'local':
|
||
# One best-effort failure receipt allows cloud to distinguish failure from offline.
|
||
try:
|
||
subprocess.run(SSH + ['umask 077; tee ' + REMOTE + '/local-ack.json.tmp >/dev/null && mv ' + REMOTE + '/local-ack.json.tmp ' + REMOTE + '/local-ack.json'],
|
||
input=encode(failure), capture_output=True, timeout=30)
|
||
except (OSError, subprocess.TimeoutExpired):
|
||
pass
|
||
print(json.dumps(failure, ensure_ascii=False), file=sys.stderr)
|
||
sys.exit(1)
|
||
finally:
|
||
if conn:
|
||
conn.close()
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main()
|