simple-login/app/oneshot/schedule_sync_user_job.py
MrMeeb d09b3b992c
Some checks failed
Build-Release-Image / Build-Image (linux/arm64) (push) Failing after 7m18s
Build-Release-Image / Build-Image (linux/amd64) (push) Has been cancelled
Build-Release-Image / Merge-Images (push) Has been cancelled
Build-Release-Image / Create-Release (push) Has been cancelled
Build-Release-Image / Notify (push) Has been cancelled
4.66.0
2025-03-03 12:00:09 +00:00

124 lines
4.0 KiB
Python

#!/usr/bin/env python3
import argparse
import sys
import time
from sqlalchemy import func
from typing import Optional
from app.jobs.send_event_job import SendEventToWebhookJob
from app.db import Session
from app.events.generated.event_pb2 import UserPlanChanged, EventContent
from app.models import PartnerUser, User
def process(start_pu_id: int, end_pu_id: int, step: int, only_lifetime: bool):
print(
f"Checking partner user {start_pu_id} to {end_pu_id} (step={step}) (only_lifetime={only_lifetime})"
)
start_time = time.time()
with_lifetime = 0
with_plan = 0
with_free = 0
for batch_start in range(start_pu_id, end_pu_id, step):
query = (
Session.query(User)
.join(PartnerUser, PartnerUser.user_id == User.id)
.filter(PartnerUser.id >= batch_start, PartnerUser.id < batch_start + step)
)
if only_lifetime:
query = query.filter(
User.lifetime == True, # noqa :E712
)
users = query.all()
for user in users:
# Just in case the == True cond is wonky
if user.lifetime:
event = UserPlanChanged(lifetime=True)
with_lifetime += 1
else:
plan_end = user.get_active_subscription_end(
include_partner_subscription=False
)
if plan_end:
event = UserPlanChanged(plan_end_time=plan_end.timestamp)
with_plan += 1
else:
event = UserPlanChanged()
with_free += 1
job = SendEventToWebhookJob(
user=user, event=EventContent(user_plan_change=event)
)
job.store_job_in_db(run_at=None, commit=False)
Session.flush()
Session.commit()
elapsed = time.time() - start_time
last_batch_id = batch_start + step
time_per_user = elapsed / last_batch_id
remaining = end_pu_id - last_batch_id
time_remaining = remaining / time_per_user
hours_remaining = time_remaining / 60.0
print(
f"PartnerUser {batch_start}/{end_pu_id} lifetime {with_lifetime} paid {with_plan} free {with_free} {hours_remaining:.2f} mins remaining"
)
print(f"Sent lifetime {with_lifetime} paid {with_plan} free {with_free}")
def main():
parser = argparse.ArgumentParser(
prog="Schedule Sync User Jobs", description="Create jobs to sync users"
)
parser.add_argument(
"-s", "--start_pu_id", default=0, type=int, help="Initial partner_user_id"
)
parser.add_argument(
"-e", "--end_pu_id", default=0, type=int, help="Last partner_user_id"
)
parser.add_argument("-t", "--step", default=10000, type=int, help="Step to use")
parser.add_argument("-u", "--user", default="", type=str, help="User to sync")
parser.add_argument(
"-l", "--lifetime", action="store_true", help="Only sync lifetime users"
)
args = parser.parse_args()
start_pu_id = args.start_pu_id
end_pu_id = args.end_pu_id
user_id = args.user
only_lifetime = args.lifetime
step = args.step
if not end_pu_id:
end_pu_id = Session.query(func.max(PartnerUser.id)).scalar()
if user_id:
try:
user_id = int(user_id)
except ValueError:
user = User.get_by(email=user_id)
if not user:
print(f"User {user_id} not found")
sys.exit(1)
user_id = user.id
print(f"Limiting to user {user_id}")
partner_user: Optional[PartnerUser] = PartnerUser.get_by(user_id=user_id)
if not partner_user:
print(f"Could not find PartnerUser for user_id={user_id}")
sys.exit(1)
# So we only have one loop
step = 1
start_pu_id = partner_user.id
end_pu_id = partner_user.id
process(
start_pu_id=start_pu_id,
end_pu_id=end_pu_id,
step=step,
only_lifetime=only_lifetime,
)
if __name__ == "__main__":
main()