Skip to content

Commit 18a7cac

Browse files
committed
feat: add ball raw data pulls
1 parent 157e37f commit 18a7cac

3 files changed

Lines changed: 417 additions & 0 deletions

File tree

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,212 @@
1+
"""
2+
Fetch ball data recordings for a session and download raw JSON for each recording
3+
that has a URL. Uses GraphQL for session/recording metadata and HTTP GET for the
4+
raw data files.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
import asyncio
10+
import json
11+
import os
12+
import sys
13+
from datetime import datetime, timedelta, timezone
14+
15+
import httpx
16+
17+
from playerdatapy.constants import API_BASE_URL
18+
from playerdatapy.gqlauth import AuthenticationType, GraphqlAuth
19+
from playerdatapy.gqlclient import Client
20+
21+
# -----------------------------------------------------------------------------
22+
# Config (env or override below)
23+
# -----------------------------------------------------------------------------
24+
CLIENT_ID = os.environ.get("CLIENT_ID")
25+
CLUB_ID = os.environ.get("CLUB_ID")
26+
SESSION_DAYS = 30 # sessions from last N days
27+
28+
# -----------------------------------------------------------------------------
29+
# GraphQL
30+
# -----------------------------------------------------------------------------
31+
SESSIONS_QUERY = """
32+
query($clubIdEq: ID!, $startTimeGteq: ISO8601DateTime!, $endTimeLteq: ISO8601DateTime!) {
33+
sessions(filter: { clubIdEq: $clubIdEq, startTimeGteq: $startTimeGteq, endTimeLteq: $endTimeLteq }) {
34+
id
35+
startTime
36+
endTime
37+
}
38+
}
39+
"""
40+
41+
SESSION_BALL_DATA_QUERY = """
42+
query($sessionId: ID!) {
43+
session(id: $sessionId) {
44+
id
45+
startTime
46+
endTime
47+
ballDataRecordings {
48+
id
49+
url(format: json)
50+
ball { id serialNumber }
51+
}
52+
}
53+
}
54+
"""
55+
56+
57+
def _record_count(data: list | dict) -> int:
58+
if isinstance(data, list):
59+
return len(data)
60+
return len(data.get("records", []))
61+
62+
63+
def _format_session_line(i: int, s: dict) -> str:
64+
"""One line for a session: number, start–end, id."""
65+
start = s.get("startTime", "")[:19].replace("T", " ")
66+
end = s.get("endTime", "")[:19].replace("T", " ")
67+
sid = s.get("id", "")
68+
return f" {i}. {start}{end} {sid}"
69+
70+
71+
def _choose_session(sessions: list[dict]) -> dict | None:
72+
"""
73+
Let the user choose a session when running interactively; otherwise use latest.
74+
Returns the chosen session dict or None if invalid/abort.
75+
"""
76+
if not sessions:
77+
return None
78+
79+
print("Sessions (most recent first):")
80+
for i, s in enumerate(sessions, start=1):
81+
print(_format_session_line(i, s))
82+
83+
if not sys.stdin.isatty():
84+
chosen = sessions[0]
85+
print(f"Using latest session: {chosen['id']}")
86+
return chosen
87+
88+
n = len(sessions)
89+
try:
90+
raw = input(f"Select session (1–{n}, or Enter for latest): ").strip()
91+
if not raw:
92+
return sessions[0]
93+
idx = int(raw)
94+
if 1 <= idx <= n:
95+
return sessions[idx - 1]
96+
except (ValueError, EOFError):
97+
pass
98+
print("Invalid choice; using latest session.")
99+
return sessions[0]
100+
101+
102+
async def fetch_recordings_for_session(
103+
client: Client,
104+
session_id: str,
105+
) -> list[dict] | None:
106+
"""Return session dict with ballDataRecordings, or None if not found."""
107+
resp = await client.execute(
108+
query=SESSION_BALL_DATA_QUERY,
109+
variables={"sessionId": session_id},
110+
)
111+
data = client.get_data(resp)
112+
return data.get("session")
113+
114+
115+
async def download_recording(
116+
http_client: httpx.AsyncClient,
117+
recording: dict,
118+
out_dir: str,
119+
) -> bool:
120+
"""Download one recording's raw JSON to out_dir. Returns True if saved, False if skipped."""
121+
url = recording.get("url")
122+
if not url:
123+
return False
124+
if url.startswith("/"):
125+
url = f"{API_BASE_URL.rstrip('/')}{url}"
126+
127+
ball = recording.get("ball") or {}
128+
serial = ball.get("serialNumber", "?")
129+
130+
try:
131+
r = await http_client.get(url)
132+
r.raise_for_status()
133+
raw = r.json()
134+
except httpx.HTTPStatusError as e:
135+
print(f" Skip {recording['id']} (Ball {serial}): {e.response.status_code}")
136+
return False
137+
except httpx.RequestError as e:
138+
print(f" Skip {recording['id']} (Ball {serial}): {e}")
139+
return False
140+
141+
if _record_count(raw) == 0:
142+
print(f" Skip {recording['id']} (Ball {serial}): empty data")
143+
return False
144+
145+
path = os.path.join(out_dir, f"{recording['id']}.json")
146+
with open(path, "w") as f:
147+
json.dump(raw, f, indent=2)
148+
print(f" Ball {serial}: {_record_count(raw)} records -> {path}")
149+
return True
150+
151+
152+
async def main() -> None:
153+
auth = GraphqlAuth(
154+
client_id=CLIENT_ID,
155+
type=AuthenticationType.AUTHORISATION_CODE_FLOW_PCKE,
156+
)
157+
client = Client(
158+
url=f"{API_BASE_URL}/api/graphql",
159+
headers={"Authorization": f"Bearer {auth._get_authentication_token()}"},
160+
)
161+
162+
now = datetime.now(timezone.utc)
163+
list_vars = {
164+
"clubIdEq": CLUB_ID,
165+
"startTimeGteq": (now - timedelta(days=SESSION_DAYS)).isoformat(),
166+
"endTimeLteq": now.isoformat(),
167+
}
168+
169+
resp = await client.execute(query=SESSIONS_QUERY, variables=list_vars)
170+
sessions = client.get_data(resp).get("sessions") or []
171+
172+
if not sessions:
173+
print("No sessions found.")
174+
return
175+
print(f"Found {len(sessions)} session(s) in last {SESSION_DAYS} days.")
176+
177+
chosen = _choose_session(sessions)
178+
if not chosen:
179+
return
180+
181+
session = await fetch_recordings_for_session(client, chosen["id"])
182+
if not session:
183+
print(f"Session {chosen['id']} not found.")
184+
return
185+
186+
recordings_with_url = [
187+
r for r in (session.get("ballDataRecordings") or []) if r.get("url")
188+
]
189+
if not recordings_with_url:
190+
print(f"No ball data recordings with URLs for session {session['id']}.")
191+
return
192+
193+
out_dir = session["id"]
194+
os.makedirs(out_dir, exist_ok=True)
195+
print(f"Session {session['id']} ({session['startTime']}{session['endTime']})")
196+
print(f"Downloading {len(recordings_with_url)} recording(s) to {out_dir}/")
197+
198+
headers = {"Authorization": f"Bearer {auth._get_authentication_token()}"}
199+
async with httpx.AsyncClient(headers=headers) as http_client:
200+
ok = sum(
201+
await asyncio.gather(
202+
*[
203+
download_recording(http_client, r, out_dir)
204+
for r in recordings_with_url
205+
]
206+
)
207+
)
208+
print(f"Done: {ok}/{len(recordings_with_url)} saved.")
209+
210+
211+
if __name__ == "__main__":
212+
asyncio.run(main())
Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
"""
2+
Fetch ball data recordings for a session and download raw JSON for each recording
3+
that has a URL. Uses the Pydantic API (PlayerDataAPI + query builders) for
4+
session/recording metadata and HTTP GET for the raw data files.
5+
"""
6+
7+
from __future__ import annotations
8+
9+
import asyncio
10+
import json
11+
import os
12+
import sys
13+
from datetime import datetime, timedelta, timezone
14+
15+
import httpx
16+
17+
from playerdatapy.constants import API_BASE_URL
18+
from playerdatapy.gqlauth import AuthenticationType
19+
from playerdatapy.playerdata_api import PlayerDataAPI
20+
21+
from queries.club_sessions_filtered_by_time_range import (
22+
club_sessions_filtered_by_time_range,
23+
)
24+
from queries.session_ball_data import session_ball_data
25+
26+
# -----------------------------------------------------------------------------
27+
# Config (env or override below)
28+
# -----------------------------------------------------------------------------
29+
CLIENT_ID = os.environ.get("CLIENT_ID")
30+
CLUB_ID = os.environ.get("CLUB_ID")
31+
SESSION_DAYS = 30 # sessions from last N days
32+
33+
34+
def _record_count(data: list | dict) -> int:
35+
if isinstance(data, list):
36+
return len(data)
37+
return len(data.get("records", []))
38+
39+
40+
def _format_session_line(i: int, s: dict) -> str:
41+
"""One line for a session: number, start–end, id."""
42+
start = (s.get("startTime") or "")[:19].replace("T", " ")
43+
end = (s.get("endTime") or "")[:19].replace("T", " ")
44+
sid = s.get("id", "")
45+
return f" {i}. {start}{end} {sid}"
46+
47+
48+
def _choose_session(sessions: list[dict]) -> dict | None:
49+
"""
50+
Let the user choose a session when running interactively; otherwise use latest.
51+
Returns the chosen session dict or None if invalid/abort.
52+
"""
53+
if not sessions:
54+
return None
55+
56+
print("Sessions (most recent first):")
57+
for i, s in enumerate(sessions, start=1):
58+
print(_format_session_line(i, s))
59+
60+
if not sys.stdin.isatty():
61+
chosen = sessions[0]
62+
print(f"Using latest session: {chosen['id']}")
63+
return chosen
64+
65+
n = len(sessions)
66+
try:
67+
raw = input(f"Select session (1–{n}, or Enter for latest): ").strip()
68+
if not raw:
69+
return sessions[0]
70+
idx = int(raw)
71+
if 1 <= idx <= n:
72+
return sessions[idx - 1]
73+
except (ValueError, EOFError):
74+
pass
75+
print("Invalid choice; using latest session.")
76+
return sessions[0]
77+
78+
79+
async def download_recording(
80+
http_client: httpx.AsyncClient,
81+
recording: dict,
82+
out_dir: str,
83+
) -> bool:
84+
"""Download one recording's raw JSON to out_dir. Returns True if saved, False if skipped."""
85+
url = recording.get("url")
86+
if not url:
87+
return False
88+
if url.startswith("/"):
89+
url = f"{API_BASE_URL.rstrip('/')}{url}"
90+
91+
ball = recording.get("ball") or {}
92+
serial = ball.get("serialNumber", "?")
93+
94+
try:
95+
r = await http_client.get(url)
96+
r.raise_for_status()
97+
raw = r.json()
98+
except httpx.HTTPStatusError as e:
99+
print(f" Skip {recording['id']} (Ball {serial}): {e.response.status_code}")
100+
return False
101+
except httpx.RequestError as e:
102+
print(f" Skip {recording['id']} (Ball {serial}): {e}")
103+
return False
104+
105+
if _record_count(raw) == 0:
106+
print(f" Skip {recording['id']} (Ball {serial}): empty data")
107+
return False
108+
109+
path = os.path.join(out_dir, f"{recording['id']}.json")
110+
with open(path, "w") as f:
111+
json.dump(raw, f, indent=2)
112+
print(f" Ball {serial}: {_record_count(raw)} records -> {path}")
113+
return True
114+
115+
116+
async def main() -> None:
117+
api = PlayerDataAPI(
118+
client_id=CLIENT_ID,
119+
client_secret="",
120+
authentication_type=AuthenticationType.AUTHORISATION_CODE_FLOW_PCKE,
121+
)
122+
123+
now = datetime.now(timezone.utc)
124+
start = now - timedelta(days=SESSION_DAYS)
125+
126+
sessions_response = await api.run_queries(
127+
"ClubSessionsFilteredByTimeRangeQuery",
128+
club_sessions_filtered_by_time_range(
129+
club_id=CLUB_ID,
130+
start_time_gteq=start,
131+
end_time_lteq=now,
132+
),
133+
)
134+
sessions = sessions_response.get("sessions") or []
135+
136+
if not sessions:
137+
print("No sessions found.")
138+
return
139+
print(f"Found {len(sessions)} session(s) in last {SESSION_DAYS} days.")
140+
141+
chosen = _choose_session(sessions)
142+
if not chosen:
143+
return
144+
145+
session_response = await api.run_queries(
146+
"SessionBallDataQuery",
147+
session_ball_data(chosen["id"]),
148+
)
149+
session = session_response.get("session")
150+
151+
if not session:
152+
print(f"Session {chosen['id']} not found.")
153+
return
154+
155+
recordings_with_url = [
156+
r for r in (session.get("ballDataRecordings") or []) if r.get("url")
157+
]
158+
if not recordings_with_url:
159+
print(f"No ball data recordings with URLs for session {session['id']}.")
160+
return
161+
162+
out_dir = session["id"]
163+
os.makedirs(out_dir, exist_ok=True)
164+
print(f"Session {session['id']} ({session['startTime']}{session['endTime']})")
165+
print(f"Downloading {len(recordings_with_url)} recording(s) to {out_dir}/")
166+
167+
headers = {"Authorization": f"Bearer {api._get_authentication_token()}"}
168+
async with httpx.AsyncClient(headers=headers) as http_client:
169+
ok = sum(
170+
await asyncio.gather(
171+
*[
172+
download_recording(http_client, r, out_dir)
173+
for r in recordings_with_url
174+
]
175+
)
176+
)
177+
print(f"Done: {ok}/{len(recordings_with_url)} saved.")
178+
179+
180+
if __name__ == "__main__":
181+
asyncio.run(main())

0 commit comments

Comments
 (0)