109 lines
3.8 KiB
Python
109 lines
3.8 KiB
Python
from datetime import date, timedelta
|
|
|
|
from sqlalchemy import FLOAT, TEXT, BigInteger, Column, Date, Integer, String, cast, create_engine
|
|
from sqlalchemy.dialects.postgresql import TIMESTAMP
|
|
from sqlalchemy.ext.declarative import declarative_base
|
|
from sqlalchemy.orm import Session, aliased
|
|
from sqlalchemy.pool import NullPool
|
|
|
|
from regiojet_search.models import Fare, Result
|
|
from regiojet_search.settings import CACHE_SURNAME, DATABASE_URL
|
|
|
|
Base = declarative_base()
|
|
|
|
|
|
class Journey(Base):
|
|
__tablename__ = f"journeys_{CACHE_SURNAME}"
|
|
|
|
id = Column(Integer, primary_key=True)
|
|
slug = Column(TEXT)
|
|
source = Column(TEXT)
|
|
destination = Column(TEXT)
|
|
departure_datetime = Column(TIMESTAMP)
|
|
arrival_datetime = Column(TIMESTAMP)
|
|
carrier = Column(TEXT)
|
|
vehicle_type = Column(TEXT)
|
|
price = Column(FLOAT)
|
|
currency = Column(String(3))
|
|
source_id = Column(BigInteger)
|
|
destination_id = Column(BigInteger)
|
|
free_seats = Column(Integer)
|
|
|
|
|
|
engine = create_engine(DATABASE_URL, poolclass=NullPool)
|
|
|
|
|
|
def save_results(results: list[Result]):
|
|
with Session(engine) as session:
|
|
existing_slugs: list[str] = [row.slug for row in session.query(Journey).distinct().all()]
|
|
|
|
journeys = [
|
|
Journey(
|
|
slug=result.slug,
|
|
source=result.origin,
|
|
destination=result.destination,
|
|
departure_datetime=result.departure,
|
|
arrival_datetime=result.arrival,
|
|
carrier=result.carrier,
|
|
vehicle_type=result.type,
|
|
price=result.fare.amount,
|
|
currency=result.fare.currency,
|
|
source_id=result.source_id,
|
|
destination_id=result.destination_id,
|
|
free_seats=result.free_seats,
|
|
)
|
|
for result in results
|
|
if result.slug not in existing_slugs
|
|
]
|
|
|
|
session.add_all(journeys)
|
|
session.commit()
|
|
|
|
|
|
def fetch_results(slug: str) -> list[Result]:
|
|
with Session(engine) as session:
|
|
journeys: list[Journey] = session.query(Journey).filter(Journey.slug == slug).all()
|
|
return [
|
|
Result(
|
|
departure=journey.departure_datetime,
|
|
arrival=journey.arrival_datetime,
|
|
origin=journey.source,
|
|
destination=journey.destination,
|
|
fare=Fare(
|
|
amount=journey.price,
|
|
currency=journey.currency,
|
|
),
|
|
type=journey.vehicle_type,
|
|
source_id=journey.source_id,
|
|
destination_id=journey.destination_id,
|
|
free_seats=journey.free_seats,
|
|
carrier=journey.carrier,
|
|
)
|
|
for journey in journeys
|
|
]
|
|
|
|
|
|
def fetch_combinations(from_city_name: str, to_city_name: str, departure_date: date) -> None:
|
|
leg1: Journey = aliased(Journey, name="leg1")
|
|
leg2: Journey = aliased(Journey, name="leg2")
|
|
|
|
with Session(engine) as session:
|
|
raw_combinations = (
|
|
session.query(leg1, leg2)
|
|
.join(leg2, leg1.destination == leg2.source)
|
|
.filter(
|
|
cast(leg1.departure_datetime, Date) == departure_date,
|
|
leg1.source == from_city_name,
|
|
leg2.destination == to_city_name,
|
|
leg1.arrival_datetime < leg2.departure_datetime,
|
|
(leg2.departure_datetime - leg1.arrival_datetime) >= timedelta(hours=1),
|
|
(leg2.departure_datetime - leg1.arrival_datetime) <= timedelta(hours=8),
|
|
)
|
|
.all()
|
|
)
|
|
for row in raw_combinations:
|
|
print(
|
|
f"Found: {row.leg1.source} {row.leg1.destination} {row.leg1.departure_datetime} {row.leg1.arrival_datetime} | "
|
|
f"{row.leg2.source} {row.leg2.destination} {row.leg2.departure_datetime} {row.leg2.arrival_datetime}"
|
|
)
|