Queue routes to the TGVMax simulator to fetch updates

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!"