Merge branch 'staging' into feat/upload_gtfs_stop_times_to_postgres

MuoVErsi/sources/base.py

@@ -340,14 +340,16 @@ def get_source_stations(self) -> list[Station]:
 
     def upload_trip_stop_times_to_postgres(self, stop_times: list[TripStopTime]):
         for stop_time in stop_times:
-            train = self.session.query(Trip).filter_by(orig_id=stop_time.origin_id,
-                                                        number=stop_time.trip_id,
-                                                        orig_dep_date=stop_time.orig_dep_date).first()
+            train = self.session.query(Trip).filter_by(
+                number=stop_time.trip_id,
+                orig_dep_date=stop_time.orig_dep_date,
+                source=self.name
+            ).first()
 
             if not train:
                 train = Trip(orig_id=stop_time.origin_id, dest_text=stop_time.destination,
                                 number=stop_time.trip_id, orig_dep_date=stop_time.orig_dep_date,
-                                route_name=stop_time.route_name)
+                                route_name=stop_time.route_name, source=self.name)
                 self.session.add(train)
                 self.session.commit()
 
@@ -368,21 +370,22 @@ class Trip(Base):
     __tablename__ = 'trips'
 
     id: Mapped[int] = mapped_column(primary_key=True)
-    orig_id: Mapped[str]
+    orig_id: Mapped[Optional[str]]
     dest_text: Mapped[str]
     number: Mapped[int]
     orig_dep_date: Mapped[date]
     route_name: Mapped[str]
-    stop_times = relationship('StopTime', back_populates='trip')
+    source: Mapped[str] = mapped_column(server_default='treni')
+    stop_times = relationship('StopTime', back_populates='trip', cascade='all, delete-orphan', passive_deletes=True)
 
-    __table_args__ = (UniqueConstraint('orig_id', 'number', 'orig_dep_date'),)
+    __table_args__ = (UniqueConstraint('source', 'number', 'orig_dep_date'),)
 
 
 class StopTime(Base):
     __tablename__ = 'stop_times'
 
     id: Mapped[int] = mapped_column(primary_key=True)
-    trip_id: Mapped[int] = mapped_column(ForeignKey('trips.id'))
+    trip_id: Mapped[int] = mapped_column(ForeignKey('trips.id', ondelete='CASCADE'))
     trip: Mapped[Trip] = relationship('Trip', back_populates='stop_times')
     stop_id: Mapped[str] = mapped_column(ForeignKey('stops.id'))
     stop: Mapped[Stop] = relationship('Stop', back_populates='stop_times')

alembic/versions/1f2c7b1eec8b_add_source_field_to_trips.py

@@ -0,0 +1,45 @@
+"""add source field to trips
+
+Revision ID: 1f2c7b1eec8b
+Revises: e07191853dcb
+Create Date: 2023-10-01 16:46:36.372782
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '1f2c7b1eec8b'
+down_revision = 'e07191853dcb'
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.add_column('trips', sa.Column('source', sa.String(), server_default='treni', nullable=False))
+    op.alter_column('trips', 'orig_id',
+               existing_type=sa.VARCHAR(),
+               nullable=True)
+    op.drop_constraint('trains_codOrigine_numeroTreno_dataPartenzaTreno_key', 'trips', type_='unique')
+
+    # update foreign key of stop_times.trip_id to cascade on delete
+    op.drop_constraint('stop_times_train_id_fkey', 'stop_times', type_='foreignkey')
+    op.create_foreign_key('stop_times_trip_id_fkey', 'stop_times', 'trips', ['trip_id'], ['id'], ondelete='CASCADE')
+
+    # remove duplicates of source, number, orig_dep_date from trips table
+    op.execute('DELETE FROM trips WHERE id IN (SELECT id FROM (SELECT id, ROW_NUMBER() OVER (partition BY source, number, orig_dep_date ORDER BY id) AS rnum FROM trips) t WHERE t.rnum > 1);')
+
+    op.create_unique_constraint(None, 'trips', ['source', 'number', 'orig_dep_date'])
+
+
+
+def downgrade() -> None:
+    op.drop_constraint(None, 'trips', type_='unique')
+    op.create_unique_constraint('trains_codOrigine_numeroTreno_dataPartenzaTreno_key', 'trips', ['orig_id', 'number', 'orig_dep_date'])
+    op.alter_column('trips', 'orig_id',
+               existing_type=sa.VARCHAR(),
+               nullable=False)
+    op.drop_column('trips', 'source')
+    op.drop_constraint('stop_times_trip_id_fkey', 'stop_times', type_='foreignkey')
+    op.create_foreign_key('stop_times_train_id_fkey', 'stop_times', 'trips', ['trip_id'], ['id'])

alembic/versions/6c9ef3a680e3_create_stops_table.py

@@ -7,6 +7,9 @@
 """
 import sqlalchemy as sa
 from alembic import op
+from sqlalchemy.orm import sessionmaker
+
+from MuoVErsi.sources.base import Station
 
 # revision identifiers, used by Alembic.
 revision = '6c9ef3a680e3'
@@ -16,8 +19,7 @@
 
 
 def upgrade() -> None:
-    # ### commands auto generated by Alembic - please adjust! ###
-    op.create_table('stops',
+    stops_table = op.create_table('stops',
                     sa.Column('id', sa.String(), nullable=False),
                     sa.Column('platform', sa.String(), nullable=True),
                     sa.Column('lat', sa.Float(), nullable=False),
@@ -26,7 +28,18 @@ def upgrade() -> None:
                     sa.ForeignKeyConstraint(['station_id'], ['stations.id'], ),
                     sa.PrimaryKeyConstraint('id')
                     )
-    # ### end Alembic commands ###
+
+    # populate stops table from stations table
+    session = sessionmaker(bind=op.get_bind())()
+    bulk_inserts = []
+    for station in session.scalars(sa.select(Station)).all():
+        bulk_inserts.append({
+            'id': station.id,
+            'lat': station.lat,
+            'lon': station.lon,
+            'station_id': station.id,
+        })
+    op.bulk_insert(stops_table, bulk_inserts)
 
 
 def downgrade() -> None: