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
@@ -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()