Use Tableau PAT for automated refresh

2026-09-10 14:42:38 +08:00
parent 1d3c6a919c
commit 96bddd1096
4 changed files with 125 additions and 11 deletions
@@ -1,2 +1,4 @@
 .DS_Store
 *.tmp
+.runtime/
+hyperd.log
@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+import csv
+import json
+import os
+import re
+import sys
+import tempfile
+import urllib.request
+import zipfile
+from pathlib import Path
+
+from tableauhyperapi import Connection, HyperProcess, Telemetry
+
+
+SERVER = os.environ.get("TABLEAU_SERVER", "https://mars.ushareit.org").rstrip("/")
+API_VERSION = os.environ.get("TABLEAU_API_VERSION", "3.15")
+SITE_CONTENT_URL = os.environ.get("TABLEAU_SITE_CONTENT_URL", "")
+WORKBOOK_ID = os.environ.get("TABLEAU_WORKBOOK_ID", "07c673fb-bac8-4c78-83e6-c60f28e04e13")
+
+
+def request_json(url, *, body=None, headers=None):
+    request = urllib.request.Request(url, data=body, headers=headers or {})
+    with urllib.request.urlopen(request, timeout=120) as response:
+        return json.load(response)
+
+
+def main():
+    if len(sys.argv) != 3:
+        raise SystemExit("usage: refresh_tableau.py YYYY-MM-DD OUTPUT.csv")
+    target_date, output_path = sys.argv[1:]
+    if not re.fullmatch(r"\d{4}-\d{2}-\d{2}", target_date):
+        raise SystemExit("invalid target date")
+
+    pat_name = os.environ.get("TABLEAU_PAT_NAME", "")
+    pat_secret = os.environ.get("TABLEAU_PAT_SECRET", "")
+    if not pat_name or not pat_secret:
+        raise SystemExit("TABLEAU_PAT_NAME and TABLEAU_PAT_SECRET are required")
+
+    signin_body = json.dumps({"credentials": {
+        "personalAccessTokenName": pat_name,
+        "personalAccessTokenSecret": pat_secret,
+        "site": {"contentUrl": SITE_CONTENT_URL},
+    }}).encode()
+    signin = request_json(
+        f"{SERVER}/api/{API_VERSION}/auth/signin",
+        body=signin_body,
+        headers={"Content-Type": "application/json", "Accept": "application/json"},
+    )["credentials"]
+    auth_token = signin["token"]
+    site_id = signin["site"]["id"]
+
+    try:
+        with tempfile.TemporaryDirectory(prefix="manta-tableau-") as temp_dir:
+            archive_path = Path(temp_dir) / "workbook.twbx"
+            download = urllib.request.Request(
+                f"{SERVER}/api/{API_VERSION}/sites/{site_id}/workbooks/{WORKBOOK_ID}/content",
+                headers={"X-Tableau-Auth": auth_token},
+            )
+            with urllib.request.urlopen(download, timeout=180) as response:
+                archive_path.write_bytes(response.read())
+            with zipfile.ZipFile(archive_path) as archive:
+                hyper_members = [name for name in archive.namelist() if name.endswith(".hyper")]
+                if len(hyper_members) != 1:
+                    raise SystemExit(f"expected one Hyper extract, found {len(hyper_members)}")
+                archive.extract(hyper_members[0], temp_dir)
+            hyper_path = Path(temp_dir) / hyper_members[0]
+
+            query = f'''
+                SELECT "pkg", "source", "dt", "campaign_name",
+                       SUM("impression"), SUM("click"), SUM("conversions"), SUM("spend")
+                FROM "Extract"."Extract"
+                WHERE "dt" = '{target_date}'
+                  AND LOWER("source") IN ('bigo', 'kwai', 'mintegral')
+                GROUP BY "pkg", "source", "dt", "campaign_name"
+                ORDER BY "pkg", "source", "campaign_name"
+            '''
+            with HyperProcess(Telemetry.DO_NOT_SEND_USAGE_DATA_TO_TABLEAU) as process:
+                with Connection(process.endpoint, str(hyper_path)) as connection:
+                    rows = connection.execute_list_query(query)
+            if not rows:
+                raise SystemExit(f"no Tableau rows for {target_date}")
+            with open(output_path, "w", encoding="utf-8", newline="") as handle:
+                writer = csv.writer(handle)
+                writer.writerow(("pkg", "source", "dt", "campaign_name",
+                                 "impression", "click", "conversions", "spend"))
+                writer.writerows(rows)
+            print(json.dumps({"date": target_date, "rawRows": len(rows)}))
+    finally:
+        try:
+            signout = urllib.request.Request(
+                f"{SERVER}/api/{API_VERSION}/auth/signout",
+                data=b"",
+                headers={"X-Tableau-Auth": auth_token},
+            )
+            urllib.request.urlopen(signout, timeout=30).close()
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    main()
@@ -1,19 +1,30 @@
 #!/usr/bin/env bash
 set -euo pipefail

-# AgentStage runs this hourly. The guard below refreshes once after 12:00 Asia/Shanghai.
+# AgentStage runs this hourly. Refresh exactly once during the 12:00 hour.
 export TZ=Asia/Shanghai
 hour="$(date +%H)"
 target_date="$(date -v-1d +%F 2>/dev/null || date -d yesterday +%F)"
-if [ "$hour" -lt 12 ]; then exit 0; fi
+if [ "$hour" -ne 12 ]; then exit 0; fi

-: "${TABLEAU_COOKIE:?Set TABLEAU_COOKIE in the AgentStage runner secret environment}"
-base_url="${TABLEAU_CSV_URL:-https://mars.ushareit.org/views/DSPSpend/DSPDailyData.csv}"
-tmp_file="$(mktemp)"
-trap 'rm -f "$tmp_file"' EXIT
+: "${TABLEAU_PAT_NAME:?Set TABLEAU_PAT_NAME in the AgentStage runner secret environment}"
+: "${TABLEAU_PAT_SECRET:?Set TABLEAU_PAT_SECRET in the AgentStage runner secret environment}"

-pkg_csv="$(python3 -c 'import json; d=json.load(open("campaign_config.json")); print(",".join(sorted({r["pkg"] for r in d if r["status"].lower()=="active"})))')"
-query="?:showVizHome=no&dt=${target_date}&pkg=${pkg_csv}&维度1=pkg&维度2=campaign_name&Source=展开&source=bigo,kwai,mintegral"
-curl -fsSL --cookie "$TABLEAU_COOKIE" "${base_url}${query}" -o "$tmp_file"
+repo_dir="$(CDPATH= cd -- "$(dirname -- "$0")" && pwd)"
+cd "$repo_dir"
+runtime_dir="${TABLEAU_HYPER_RUNTIME:-$repo_dir/.runtime/tableau-hyperapi}"
+success_marker="$repo_dir/.runtime/last-successful-date"
+if [ -f "$success_marker" ] && [ "$(cat "$success_marker")" = "$target_date" ]; then
+  exit 0
+fi
+if ! PYTHONPATH="$runtime_dir" python3 -c 'import tableauhyperapi' 2>/dev/null; then
+  mkdir -p "$runtime_dir"
+  python3 -m pip install --disable-pip-version-check --target "$runtime_dir" tableauhyperapi
+fi

-python3 transform_tableau.py "$tmp_file" "$target_date"
+tmp_dir="$(mktemp -d)"
+trap 'rm -rf "$tmp_dir"' EXIT
+raw_file="$tmp_dir/tableau.csv"
+PYTHONPATH="$runtime_dir" python3 refresh_tableau.py "$target_date" "$raw_file"
+python3 transform_tableau.py "$raw_file" "$target_date"
+printf '%s\n' "$target_date" > "$success_marker"