diff --git a/app.py b/app.py index 24f0193..a771fdc 100644 --- a/app.py +++ b/app.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 import csv -from datetime import date, datetime, time +from datetime import date, datetime, time, timedelta import os import json from pytz import timezone @@ -13,6 +13,7 @@ from flask.cli import AppGroup from flask_migrate import Migrate from flask_sqlalchemy import SQLAlchemy from sqlalchemy import Boolean, Column, Date, DateTime, Integer, String, Time +from sqlalchemy.sql import func from tqdm import tqdm import config @@ -49,6 +50,16 @@ class Train(db.Model): expiration_time = Column(DateTime) +class RouteQueue(db.Model): + id = Column(Integer, autoincrement=True, primary_key=True) + queue_time = Column(DateTime(timezone=True), server_default=func.now()) + day = Column(Date) + origin = Column(String(5)) + destination = Column(String(5)) + response_time = Column(DateTime(timezone=True), nullable=True, default=None) + expiration_time = Column(DateTime(timezone=True), nullable=True, default=None) + + @cli.command("update-dataset") def update_dataset(): """ @@ -197,6 +208,78 @@ def print_route(route: list[Train]): print(s[:-2]) +@cli.command('queue-route') +@click.argument('day', type=click.DateTime(formats=['%Y-%m-%d'])) +@click.argument('origin', type=str) +@click.argument('destination', type=str) +def queue_route(day: date | datetime, origin: str, destination: str): + """ + Fetch the TGVMax simulator to refresh data. + + DAY: The day to query, in format YYYY-MM-DD. + + ORIGIN: The origin of the route. + + DESTINATION: The destination of the route. + """ + if isinstance(day, datetime): + day = day.date() + + query = db.session.query(RouteQueue).filter_by(day=day, origin=origin, destination=destination, response_time=None) + if query.count(): + print("Already queued") + return + + db.session.add(RouteQueue(day=day, origin=origin, destination=destination)) + db.session.commit() + + +@cli.command('process-queue', help="Process the waiting list to refresh from the simulator.") +@click.argument('number', default=5, type=int) +def process_queue(number: int): + queue = db.session.query(RouteQueue).filter_by(response_time=None).order_by(RouteQueue.queue_time) + if number > 0: + queue = queue[:number] + + URL = "https://www.maxjeune-tgvinoui.sncf/api/public/refdata/search-freeplaces-proposals" + + for req in queue: + req: RouteQueue + resp = requests.post(URL, json={ + 'departureDateTime': req.day.isoformat(), + 'origin': req.origin, + 'destination': req.destination, + }) + + if resp.status_code == 404: + # No travel found + req.response_time = datetime.now() + req.expiration_time = datetime.now() + timedelta(hours=1) + db.session.add(req) + continue + + resp.raise_for_status() + + data = resp.json() + req.response_time = datetime.utcfromtimestamp(data['updatedAt'] // 1000).replace(tzinfo=timezone('UTC')) + req.expiration_time = datetime.utcfromtimestamp(data['expiresAt'] // 1000).replace(tzinfo=timezone('UTC')) + db.session.add(req) + + db.session.query(Train).filter_by(day=req.day, orig_iata=req.origin, dest_iata=req.destination)\ + .update(dict(tgvmax=False, remaining_seats=-1)) + + for proposal in data['proposals']: + train = db.session.query(Train).filter_by(day=req.day, number=int(proposal['trainNumber']), + orig_iata=req.origin, dest_iata=req.destination).first() + train.tgvmax = True + train.remaining_seats = proposal['freePlaces'] + train.last_modification = req.response_time + train.expiration_time = req.expiration_time + db.session.add(train) + + db.session.commit() + + @app.get('/') def index(): return "Hello world!"