ai_member_xiaoxi/scripts/process_xiaolong_dij.py
2026-06-05 08:00:01 +08:00

236 lines
8.2 KiB
Python

#!/usr/bin/env python3
"""
Process 小龙 sheet: fill D (trial counts), I (reg dates), J (channels).
H column already written. Now fill D/I/J for matched rows.
"""
import sys
sys.path.insert(0, '/root/.openclaw/workspace/scripts')
from phone_encrypt import encrypt_phone
import psycopg2
import requests
import time
PG_CONFIG = {
'host': 'bj-postgres-16pob4sg.sql.tencentcdb.com',
'port': 28591,
'user': 'ai_member',
'password': 'LdfjdjL83h3h3^$&**YGG*',
'database': 'vala_bi',
}
SPREADSHEET_TOKEN = 'DU4dsUOJThfbPStMcgBcsMH5nyb'
SHEET_ID = 'dff8c7'
def get_token():
resp = requests.post('https://open.feishu.cn/open-apis/auth/v3/tenant_access_token/internal',
json={'app_id': 'cli_a929ae22e0b8dcc8', 'app_secret': 'OtFjMy7p3qE3VvLbMdcWidwgHOnGD4FJ'})
return resp.json()['tenant_access_token']
def read_sheet(token):
url = f'https://open.feishu.cn/open-apis/sheets/v2/spreadsheets/{SPREADSHEET_TOKEN}/values/{SHEET_ID}!A3:J1142?valueRenderOption=ToString'
resp = requests.get(url, headers={'Authorization': f'Bearer {token}'})
return resp.json().get('data', {}).get('valueRange', {}).get('values', [])
def write_batch(token, range_str, values, max_retries=3):
url = f'https://open.feishu.cn/open-apis/sheets/v2/spreadsheets/{SPREADSHEET_TOKEN}/values'
body = {"valueRange": {"range": f"{SHEET_ID}!{range_str}", "values": values}}
for attempt in range(max_retries):
resp = requests.put(url, headers={
'Authorization': f'Bearer {token}',
'Content-Type': 'application/json'
}, json=body)
result = resp.json()
code = result.get('code', -1)
if code == 0:
return True
elif code == 90217: # too many request
wait = 1.0 * (attempt + 1)
print(f" Rate limited on {range_str}, waiting {wait}s...")
time.sleep(wait)
else:
print(f" ERROR {range_str}: {result}")
return False
print(f" FAILED after {max_retries} retries: {range_str}")
return False
def batch_writes(token, col, writes):
"""Write sorted writes in consecutive batches."""
if not writes:
return 0
writes.sort(key=lambda x: x[0])
batches = []
batch_start = writes[0][0]
batch_vals = [[writes[0][1]]]
for j in range(1, len(writes)):
if writes[j][0] == writes[j-1][0] + 1:
batch_vals.append([writes[j][1]])
else:
batches.append((batch_start, batch_vals))
batch_start = writes[j][0]
batch_vals = [[writes[j][1]]]
batches.append((batch_start, batch_vals))
written = 0
for start_idx, vals in batches:
start_row = start_idx + 3
end_row = start_row + len(vals) - 1
range_str = f'{col}{start_row}:{col}{end_row}'
if write_batch(token, range_str, vals):
written += len(vals)
print(f" {range_str} ({len(vals)} cells)")
time.sleep(0.08)
return written
def main():
token = get_token()
print("Reading sheet...")
rows = read_sheet(token)
print(f" {len(rows)} rows")
# Extract phones
print("\nEncrypting phones...")
phone_map = {} # enc -> (row_idx, phone)
for i, row in enumerate(rows):
if len(row) > 4 and row[4]:
phone = row[4].strip()
if len(phone) == 11 and phone.isdigit():
enc = encrypt_phone(phone)
phone_map[enc] = (i, phone)
print(f" {len(phone_map)} phones")
# Query DB
print("\nQuerying PostgreSQL...")
conn = psycopg2.connect(**PG_CONFIG)
cur = conn.cursor()
enc_list = list(phone_map.keys())
phone_to_uid = {}
phone_to_created = {}
phone_to_channel = {}
for chunk_start in range(0, len(enc_list), 500):
chunk = enc_list[chunk_start:chunk_start+500]
placeholders = ','.join(['%s'] * len(chunk))
cur.execute(f"""
SELECT id, tel_encrypt, created_at, download_channel
FROM bi_vala_app_account
WHERE tel_encrypt IN ({placeholders})
AND status = 1 AND deleted_at IS NULL
""", chunk)
for row in cur.fetchall():
uid, enc, created, channel = row
phone_to_uid[enc] = str(uid)
phone_to_created[enc] = created
phone_to_channel[enc] = channel or ''
print(f" Matched {len(phone_to_uid)} UIDs")
# Trial counts
all_uids = list(set(phone_to_uid.values()))
uid_to_trial = {}
for chunk_start in range(0, len(all_uids), 500):
chunk = all_uids[chunk_start:chunk_start+500]
placeholders = ','.join(['%s'] * len(chunk))
cur.execute(f"""
SELECT account_id, COUNT(*)
FROM bi_user_course_detail
WHERE account_id IN ({placeholders})
AND expire_time IS NULL AND deleted_at IS NULL
GROUP BY account_id
""", [int(x) for x in chunk])
for row in cur.fetchall():
uid_to_trial[str(row[0])] = row[1]
print(f" Trial counts for {len(uid_to_trial)} users")
# Also get trial counts for rows that already have UIDs in H but empty D
existing_uids = set()
for i, row in enumerate(rows):
if len(row) > 7 and row[7]:
h_val = row[7].strip()
if h_val.isdigit():
existing_uids.add(h_val)
# Query trial counts for existing UIDs not already covered
missing_uids = existing_uids - set(uid_to_trial.keys())
if missing_uids:
print(f" Querying trial counts for {len(missing_uids)} existing UIDs...")
uid_list = list(missing_uids)
for chunk_start in range(0, len(uid_list), 500):
chunk = uid_list[chunk_start:chunk_start+500]
placeholders = ','.join(['%s'] * len(chunk))
cur.execute(f"""
SELECT account_id, COUNT(*)
FROM bi_user_course_detail
WHERE account_id IN ({placeholders})
AND expire_time IS NULL AND deleted_at IS NULL
GROUP BY account_id
""", [int(x) for x in chunk])
for row in cur.fetchall():
uid_to_trial[str(row[0])] = row[1]
print(f" Added {len(missing_uids & set(uid_to_trial.keys()))} more")
cur.close()
conn.close()
# Build writes
row_data = {}
for enc, (row_idx, phone) in phone_map.items():
uid = phone_to_uid.get(enc)
if uid:
trial = uid_to_trial.get(uid, 0)
created = phone_to_created.get(enc)
channel = phone_to_channel.get(enc, '')
reg_date = created.strftime('%Y-%m-%d') if created else ''
row_data[row_idx] = {
'D': str(trial) if trial > 0 else '',
'I': reg_date,
'J': channel,
}
# Also fill D for rows with existing H UID but empty D
for i, row in enumerate(rows):
if i in row_data:
continue
if len(row) > 7 and row[7]:
h_val = row[7].strip()
if h_val.isdigit():
d_empty = len(row) <= 3 or not row[3]
if d_empty:
trial = uid_to_trial.get(h_val, 0)
if trial > 0:
if i not in row_data:
row_data[i] = {}
row_data[i]['D'] = str(trial)
# Prepare column writes
d_writes = [(idx, data['D']) for idx, data in row_data.items() if data.get('D')]
i_writes = [(idx, data['I']) for idx, data in row_data.items() if data.get('I')]
j_writes = [(idx, data['J']) for idx, data in row_data.items() if data.get('J')]
print(f"\nPrepared: D={len(d_writes)}, I={len(i_writes)}, J={len(j_writes)}")
# Write D
print("\n=== Writing D column ===")
d_written = batch_writes(token, 'D', d_writes)
# Write I
print("\n=== Writing I column ===")
i_written = batch_writes(token, 'I', i_writes)
# Write J
print("\n=== Writing J column ===")
j_written = batch_writes(token, 'J', j_writes)
# Report
print("\n" + "="*60)
print("FINAL REPORT")
print("="*60)
print(f"Phones matched to UIDs: {len(phone_to_uid)}")
print(f"D column written: {d_written}")
print(f"I column written: {i_written}")
print(f"J column written: {j_written}")
if __name__ == '__main__':
main()