More calendar optimizations

This commit is contained in:
Emmy D'Anello 2024-05-12 14:03:48 +02:00
parent 2b6523c728
commit b65dc10bc6
Signed by: ynerant
GPG Key ID: 3A75C55819C8CF85
1 changed files with 15 additions and 12 deletions

View File

@ -225,21 +225,23 @@ class Command(BaseCommand):
calendars.clear()
calendar_dates = []
all_calendars = {calendar.id: calendar for calendar in Calendar.objects.filter(gtfs_feed_id=gtfs_code)}
new_calendars = {}
with transaction.atomic():
for calendar_date_dict in read_csv("calendar_dates.txt"):
calendar_date_dict: dict
service_id = f"{gtfs_code}-{calendar_date_dict['service_id']}"
date = calendar_date_dict['date']
date = datetime.fromisoformat(calendar_date_dict['date']).date()
calendar_date = CalendarDate(
id=f"{gtfs_code}-{calendar_date_dict['service_id']}-{calendar_date_dict['date']}",
service_id=service_id,
date=date,
date=calendar_date_dict['date'],
exception_type=calendar_date_dict['exception_type'],
)
calendar_dates.append(calendar_date)
if not Calendar.objects.filter(id=calendar_date.service_id).exists():
if service_id not in all_calendars:
calendar = Calendar(
id=service_id,
monday=False,
@ -249,26 +251,27 @@ class Command(BaseCommand):
friday=False,
saturday=False,
sunday=False,
start_date=calendar_date_dict['date'],
end_date=calendar_date_dict['date'],
start_date=date,
end_date=date,
gtfs_feed_id=gtfs_code,
)
calendar.save()
all_calendars[service_id] = calendar
new_calendars[service_id] = calendar
else:
calendar = Calendar.objects.get(id=service_id)
if calendar.start_date.isoformat() > date:
calendar = all_calendars[service_id]
if calendar.start_date > date:
calendar.start_date = date
calendar.save()
if calendar.end_date.isoformat() < date:
if calendar.end_date < date:
calendar.end_date = date
calendar.save()
if len(calendar_dates) >= bulk_size and not dry_run:
CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size)
calendar_dates.clear()
if calendar_dates and not dry_run:
if (calendar_dates or new_calendars) and not dry_run:
Calendar.objects.bulk_create(new_calendars.values(), batch_size=bulk_size)
CalendarDate.objects.bulk_create(calendar_dates, batch_size=bulk_size)
new_calendars.clear()
calendar_dates.clear()
trips = []