update_companies: switch to /company/all endpoint

py/update_companies.py

@@ -1,114 +1,43 @@
 from __future__ import annotations
-import time
 
-import itertools
 import json
-import sys
 import tomllib
 import typing
 
 import httpx
 
 def main() -> None:
-	(month,) = sys.argv[1:]
-
 	with open('config.toml', 'rb') as f:
 		config = tomllib.load(f)
 	fio_api_key = config['fio_api_key']
 
-	with open(f'www/data/company-data-{month}.json', 'r') as f:
-		company_ids = frozenset(json.load(f)['totals'])
-	print(len(company_ids), 'companies in', month, 'data')
-
 	with open('www/data/knownCompanies.json', 'r') as f:
 		companies = json.load(f)
 
-	now = int(time.time())
-	cached: set[str] = set()
-	for co in companies.values():
-		co_last_update = co.get('LastUpdated')
-		if co_last_update is not None and now - co_last_update < 24 * 60 * 60:
-			cached.add(co['Username'].casefold())
-	print(len(cached), 'cached companies')
-
-	all_users: typing.Collection[str] = httpx.get('https://rest.fnar.net/user/allusers',
+	all_companies: typing.Collection[FIOCompany] = httpx.get('https://rest.fnar.net/company/all',
 			headers={'Authorization': fio_api_key}).raise_for_status().json()
 	try:
-		for username in all_users:
-			update_company_from_fio(companies, cached, now, username)
-	finally:
-		with open('www/data/knownCompanies.json', 'w') as f:
-			json.dump(companies, f)
-
-	unknown_companies = company_ids - companies.keys()
-	print(len(unknown_companies), 'unknown companies')
-	try:
-		for batch in itertools.batched(unknown_companies, 10):
-			company_data: typing.Collection[FIOCompany | None] = httpx.get('https://api.fnar.net/company/lookup',
-					params={'company': batch}).raise_for_status().json()
-			for fio_co in company_data:
-				if fio_co is None:
-					continue
-				known_user = {'Username': fio_co['UserName'], 'LastUpdated': now}
-				if (corp := fio_co.get('CorporationCode')) and corp != 'UNKN':
-					known_user['Corporation'] = corp
-				print(known_user)
-				assert fio_co['CompanyId'] not in companies
-				companies[fio_co['CompanyId']] = known_user
+		for company in all_companies:
+			username = company['UserName']
+			if username is None:
+				print(company['CompanyId'], 'has no username:', company)
+				continue
+			known_company = {'Username': username}
+			if corp := company.get('CorporationCode'):
+				known_company['Corporation'] = corp
+
+			if known_company == companies.get(company['CompanyId']):
+				print(username, 'had no change')
+			else:
+				print(known_company)
+				companies[company['CompanyId']] = known_company
 	finally:
 		with open('www/data/knownCompanies.json', 'w') as f:
 			json.dump(companies, f)
 
-	fio_rest_users = frozenset(user.casefold() for user in all_users)
-	known_usernames = frozenset(co['Username'].casefold() for co in companies.values() if co['Username'] is not None)
-	users_without_fio_accounts = known_usernames - fio_rest_users
-	print(len(users_without_fio_accounts), 'users without FIO REST accounts')
-	try:
-		for username in users_without_fio_accounts:
-			update_company_from_fio(companies, cached, now, username)
-	finally:
-		with open('www/data/knownCompanies.json', 'w') as f:
-			json.dump(companies, f)
-
-def update_company_from_fio(companies: dict[str, dict[str, typing.Any]], cached: typing.AbstractSet[str], now: int,
-		username: str) -> None:
-	if username.casefold() in cached:
-		return
-
-	response = get_with_retry('https://rest.fnar.net/user/' + username)
-	if response.status_code in (204, 404):
-		print(username, 'not found')
-		return
-	fio_user = response.json()
-	assert username.casefold() == fio_user['UserName'].casefold(), f'{username} != {fio_user["UserName"]}'
-	known_user = {'Username': username, 'LastUpdated': now}
-	if corp := fio_user.get('CorporationCode'):
-		known_user['Corporation'] = corp
-	if known_user == companies.get(fio_user['CompanyId']):
-		print(username, 'had no change')
-	else:
-		print(known_user)
-		companies[fio_user['CompanyId']] = known_user
-
-def get_with_retry(url) -> httpx.Response:
-	for attempt in range(10):
-		try:
-			response = httpx.get(url)
-		except httpx.TransportError as e:
-			print(e, 'retrying in', attempt)
-			time.sleep(attempt)
-			continue
-		if response.status_code in (204, 404) or response.is_success:
-			return response
-		elif attempt == 9:
-			response.raise_for_status()
-		print(response.status_code, 'retrying in', attempt)
-		time.sleep(attempt)
-	raise AssertionError
-
 class FIOCompany(typing.TypedDict):
 	CompanyId: str
-	UserName: str
+	UserName: str | None
 	CorporationCode: str | None
 
 if __name__ == '__main__':