Complete rewrite with Polars. No more single loop to fetch all data; posts, queued posts, drafts, and activity are all fetched separately. Extra caution when querying API to prevent HTTP 429
This commit is contained in:
357
fetcher.py
Normal file
357
fetcher.py
Normal file
@@ -0,0 +1,357 @@
|
||||
from datetime import datetime as dt
|
||||
import json
|
||||
import sys
|
||||
from time import sleep
|
||||
import traceback
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import pytumblr
|
||||
import polars as pl
|
||||
|
||||
|
||||
def filter_for_after(df: pl.DataFrame,
|
||||
after: dt) -> pl.DataFrame:
|
||||
if df.is_empty() or not after: return pl.DataFrame()
|
||||
base_col = pl.col('date')
|
||||
if df.get_column('date').dtype == pl.String:
|
||||
base_col = pl.when(base_col.str.contains('GMT')) \
|
||||
.then(base_col.str.to_date('%Y-%m-%d %H:%M:%S GMT')) \
|
||||
.otherwise(base_col.str.to_date('%Y-%m-%d %H:%M:%S'))
|
||||
return df.filter(base_col >= after.date())
|
||||
|
||||
|
||||
def filter_for_before(df: pl.DataFrame,
|
||||
before: dt) -> pl.DataFrame:
|
||||
if df.is_empty() or not before: return pl.DataFrame()
|
||||
base_col = pl.col('date')
|
||||
if df.get_column('date').dtype == pl.String:
|
||||
base_col = pl.when(base_col.str.contains('GMT')) \
|
||||
.then(base_col.str.to_date('%Y-%m-%d %H:%M:%S GMT')) \
|
||||
.otherwise(base_col.str.to_date('%Y-%m-%d %H:%M:%S'))
|
||||
return df.filter(base_col <= before.date())
|
||||
|
||||
|
||||
def handle_error(e: Exception, data: Any, df: pl.DataFrame) -> None:
|
||||
print(e)
|
||||
traceback.print_exc()
|
||||
schema_dict = {col: str(dtype) for col, dtype in df.schema.items()}
|
||||
with open("schema.json", "w") as f:
|
||||
json.dump(schema_dict, f, indent=4)
|
||||
with open("bad_posts.json", "w") as f:
|
||||
json.dump(data, f, indent=4, default=str)
|
||||
|
||||
|
||||
def handle_limit_exceeded(data) -> bool:
|
||||
if 'meta' in data and data['meta']['status'] == 429:
|
||||
print("WARN: API limit exceeded. Sleeping for 10 minutes...")
|
||||
sleep(600)
|
||||
return True
|
||||
sleep(3)
|
||||
return False
|
||||
|
||||
|
||||
def fetch_minimized_post(blog_name: str,
|
||||
client: pytumblr.TumblrRestClient,
|
||||
id_string: str) -> Dict[str, Any]:
|
||||
api_url = f"/v2/blog/{blog_name}/posts"
|
||||
data: Dict[str, Any] = client.send_api_request(method='get',
|
||||
url=api_url,
|
||||
params={'id': id_string},
|
||||
valid_parameters=['id'])
|
||||
|
||||
if 'posts' in data:
|
||||
posts = data['posts']
|
||||
if handle_limit_exceeded(data):
|
||||
data = client.send_api_request(method='get', url=api_url, params={'id': id_string}, valid_parameters=['id'])
|
||||
posts = data['posts']
|
||||
if 'posts' not in data or not data['posts']:
|
||||
print(f"WARN: Post with ID {id_string} could not be retrieved!")
|
||||
print(data)
|
||||
return {}
|
||||
|
||||
minimized_post: Dict[str, Any] = {}
|
||||
if posts:
|
||||
minimized_post = posts[0]
|
||||
minimized_post = {
|
||||
'id_string': minimized_post['id_string'],
|
||||
'url': minimized_post['post_url'],
|
||||
'date_str': minimized_post['date'],
|
||||
'tags': minimized_post['tags'],
|
||||
'is_og': 'parent_post_url' not in minimized_post
|
||||
}
|
||||
return minimized_post
|
||||
|
||||
|
||||
def fetch_queued_posts(client: pytumblr.TumblrRestClient,
|
||||
args: Dict[str, Any]) -> pl.DataFrame:
|
||||
queued_posts: pl.DataFrame = pl.DataFrame()
|
||||
blog_name: str = args['blog']
|
||||
|
||||
offset: int = 0
|
||||
limit: int = 20
|
||||
|
||||
while True:
|
||||
# Begin LOOP
|
||||
# Get me some queued posts via REST! 😈🍪🍪🍪
|
||||
data: Dict[str, Any]
|
||||
data = client.queue(f"{blog_name}.tumblr.com",
|
||||
offset=offset,
|
||||
limit=limit)
|
||||
|
||||
if handle_limit_exceeded(data):
|
||||
data = client.queue(f"{blog_name}.tumblr.com",
|
||||
offset=offset,
|
||||
limit=limit)
|
||||
|
||||
# This struct breaks Polars for some reason 🤷
|
||||
posts_json = data['posts']
|
||||
for post in posts_json:
|
||||
if 'trail' in post:
|
||||
del post['trail']
|
||||
|
||||
try:
|
||||
curr_posts: pl.DataFrame = pl.DataFrame(posts_json)
|
||||
except Exception as e:
|
||||
handle_error(e, data, queued_posts)
|
||||
sys.exit(0)
|
||||
|
||||
# Stop the presses if we found no posts.
|
||||
if curr_posts.is_empty() or len(curr_posts) < 1:
|
||||
print('Stopping queued post retrieval, as no more queued posts were found.')
|
||||
break
|
||||
|
||||
queued_posts = pl.concat([queued_posts, curr_posts], how="diagonal_relaxed")
|
||||
|
||||
offset += limit
|
||||
# End LOOP
|
||||
|
||||
# Empty check.
|
||||
if queued_posts.is_empty():
|
||||
print("No queued posts were found.")
|
||||
|
||||
return queued_posts
|
||||
|
||||
|
||||
def fetch_drafts(client: pytumblr.TumblrRestClient,
|
||||
args: Dict[str, Any]) -> pl.DataFrame:
|
||||
drafts: pl.DataFrame = pl.DataFrame()
|
||||
blog_name: str = args['blog']
|
||||
draft_url = f"/v2/blog/{blog_name}/posts/draft"
|
||||
|
||||
while True:
|
||||
# Begin LOOP
|
||||
# Get me some drafts via REST! 😈🍪🍪🍪
|
||||
data: Dict[str, Any]
|
||||
data = client.send_api_request("get", draft_url)
|
||||
|
||||
if handle_limit_exceeded(data):
|
||||
data = client.send_api_request("get", draft_url)
|
||||
|
||||
# This struct breaks Polars for some reason 🤷
|
||||
drafts_json = data['posts']
|
||||
for draft in drafts_json:
|
||||
if 'trail' in draft:
|
||||
del draft['trail']
|
||||
|
||||
curr_drafts: pl.DataFrame = pl.DataFrame(drafts_json)
|
||||
drafts = pl.concat([drafts, curr_drafts], how="diagonal_relaxed")
|
||||
|
||||
if '_links' in data:
|
||||
draft_url = data['_links']['next']['href']
|
||||
continue
|
||||
else:
|
||||
print('Stopping draft retrieval, as no more drafts were found.')
|
||||
break
|
||||
# End LOOP
|
||||
|
||||
# Empty check.
|
||||
if drafts.is_empty():
|
||||
print("No drafts were found.")
|
||||
|
||||
return drafts
|
||||
|
||||
|
||||
def fetch_posts(client: pytumblr.TumblrRestClient,
|
||||
args: Dict[str, Any]) -> pl.DataFrame:
|
||||
posts: pl.DataFrame = pl.DataFrame()
|
||||
blog_name: str = args['blog']
|
||||
|
||||
# We populate params, starting with any tags for filtering.
|
||||
params = {}
|
||||
if args['tags']:
|
||||
params.update({'tag': ','.join(args['tags'])})
|
||||
|
||||
total: int = 0
|
||||
offset: int = 0
|
||||
limit: int = 20
|
||||
|
||||
# The request loop that pulls all data from the APIs.
|
||||
while True:
|
||||
# Begin LOOP
|
||||
# Get me some posts via REST! 😈🍪🍪🍪
|
||||
data: Dict[str, Any]
|
||||
data = client.posts(f"{blog_name}.tumblr.com",
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
**params)
|
||||
|
||||
if handle_limit_exceeded(data):
|
||||
data = client.posts(f"{blog_name}.tumblr.com",
|
||||
offset=offset,
|
||||
limit=limit,
|
||||
**params)
|
||||
|
||||
# These structs break Polars for some reason 🤷
|
||||
posts_json = data['posts']
|
||||
for post in posts_json:
|
||||
if 'trail' in post:
|
||||
del post['trail']
|
||||
if 'player' in post:
|
||||
del post['player']
|
||||
|
||||
try:
|
||||
curr_posts: pl.DataFrame = pl.DataFrame(posts_json, strict=False, infer_schema_length=None)
|
||||
except Exception as e:
|
||||
handle_error(e, data, curr_posts)
|
||||
sys.exit(1)
|
||||
|
||||
# Stop the presses if we found no posts.
|
||||
if curr_posts.is_empty() or len(curr_posts) < 1:
|
||||
print('Stopping, as no more posts were found.')
|
||||
break
|
||||
|
||||
next_off: int = 0
|
||||
|
||||
# Total init check for the first iteration, but always checked for sanity.
|
||||
if not total and 'total_posts' in data:
|
||||
total_posts = data['total_posts']
|
||||
print(f"I'm working with {total_posts} total posts...")
|
||||
total = total_posts
|
||||
|
||||
if args['since'] and not args['until']:
|
||||
after: dt = args['since']
|
||||
curr_posts = filter_for_after(curr_posts, after)
|
||||
if curr_posts.is_empty():
|
||||
print(
|
||||
f"All posts since {after.year}-{after.month}-{after.day} retrieved.")
|
||||
return posts
|
||||
if args['until'] and not args['since']:
|
||||
before: dt = args['until']
|
||||
curr_posts = filter_for_before(curr_posts, before)
|
||||
if curr_posts.is_empty():
|
||||
print(
|
||||
f"All posts until {before.year}-{before.month}-{before.day} retrieved.")
|
||||
return posts
|
||||
if args['since'] and args['until']:
|
||||
if not posts.is_empty() and curr_posts.is_empty():
|
||||
print(f"All posts since {after.year}-{after.month}-{after.day}, "
|
||||
+ f"until {before.year}-{before.month}-{before.day} retrieved.")
|
||||
return posts
|
||||
|
||||
# Update the frame with what we found.
|
||||
try:
|
||||
posts = pl.concat([posts, curr_posts], how="diagonal_relaxed")
|
||||
except Exception as e:
|
||||
handle_error(e, data, posts)
|
||||
sys.exit(1)
|
||||
|
||||
# The increment and status printing.
|
||||
if next_off != 0 and next_off != offset:
|
||||
offset = next_off
|
||||
else:
|
||||
offset += limit
|
||||
if not (args['since'] or args['until']) and total:
|
||||
print(
|
||||
f"Processed batch {offset // limit} of {(total // 20) + 1}...")
|
||||
# End LOOP
|
||||
|
||||
if posts.is_empty():
|
||||
print("No posts were found.")
|
||||
|
||||
# Return posts (both original and not original).
|
||||
return posts.unique(subset=["id_string"])
|
||||
|
||||
|
||||
def fetch_activity(client: pytumblr.TumblrRestClient,
|
||||
args: Dict[str, Any],
|
||||
posts: pl.DataFrame = pl.DataFrame()) -> pl.DataFrame:
|
||||
activity: pl.DataFrame = pl.DataFrame(infer_schema_length=None)
|
||||
blog_name: str = args['blog']
|
||||
|
||||
before: dt = args['until']
|
||||
after: dt = args['since']
|
||||
|
||||
activity_url = f"/v2/blog/{blog_name}/notifications"
|
||||
if before:
|
||||
activity_url += f"?before={str(int(before.timestamp()))}"
|
||||
|
||||
params = {}
|
||||
params.update({'rollups': 'false'})
|
||||
|
||||
post_map: Dict[str, Any] = {}
|
||||
if not posts.is_empty():
|
||||
post_map = { row["id"]: row for row in posts.to_dicts() }
|
||||
|
||||
data: Dict[str, Any]
|
||||
while True:
|
||||
# Begin LOOP
|
||||
# Get me some activity via REST! 😈🍪🍪🍪
|
||||
data = client.send_api_request("get", activity_url)
|
||||
if handle_limit_exceeded(data):
|
||||
data = client.send_api_request("get", activity_url)
|
||||
|
||||
try:
|
||||
curr_activity: pl.DataFrame = pl.DataFrame(data['notifications'], infer_schema_length=None)
|
||||
except Exception as e:
|
||||
handle_error(e, data, activity)
|
||||
sys.exit(0)
|
||||
curr_activity = curr_activity.with_columns(pl.from_epoch('timestamp').alias('date'))
|
||||
if before:
|
||||
curr_activity = filter_for_before(curr_activity, before)
|
||||
if after:
|
||||
curr_activity = filter_for_after(curr_activity, after)
|
||||
|
||||
# Add additional post info to notifications.
|
||||
# Yeah, df to dictionary list. Sorry. :c
|
||||
curr_activity_dicts: List[Dict[str, Any]] = curr_activity.to_dicts()
|
||||
for notification in curr_activity_dicts:
|
||||
if 'target_post_id' in notification and notification['target_post_id']:
|
||||
target_post_id = notification['target_post_id']
|
||||
post: Any
|
||||
if target_post_id in post_map:
|
||||
post = post_map.get(target_post_id)
|
||||
notification['post'] = post
|
||||
continue
|
||||
elif 'type' in notification and notification['type'] == 'post_attribution':
|
||||
notification['post'] = {}
|
||||
continue
|
||||
else:
|
||||
post = fetch_minimized_post(blog_name, client, notification['target_post_id'])
|
||||
notification['post'] = post
|
||||
post_map.update({target_post_id: post})
|
||||
continue
|
||||
|
||||
curr_activity = pl.DataFrame(curr_activity_dicts, strict=False, infer_schema_length=None)
|
||||
activity = pl.concat([activity, curr_activity], how="diagonal_relaxed")
|
||||
|
||||
if '_links' in data:
|
||||
next_time = dt.fromtimestamp(int(data['_links']['next']['query_params']['before']))
|
||||
if before or after:
|
||||
print(f"Reached activity time: {next_time}...")
|
||||
if (before and next_time >= before) or (after and next_time <= after):
|
||||
print('Stopping activity retrieval, as no more activity is left in date range.')
|
||||
break
|
||||
activity_url = data['_links']['next']['href']
|
||||
continue
|
||||
else:
|
||||
print('Stopping activity retrieval, as no more activity was found.')
|
||||
break
|
||||
# End LOOP
|
||||
|
||||
# Empty check.
|
||||
if activity.is_empty():
|
||||
print("No activity was found.")
|
||||
else:
|
||||
activity = activity.unique(subset=["id"])
|
||||
|
||||
return activity
|
||||
Reference in New Issue
Block a user