fix: add source field to trip and change unique contraint (with rows …

…deletion) (#127)

MuoVErsi/sources/base.py

@@ -322,21 +322,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')

MuoVErsi/sources/trenitalia.py

@@ -77,14 +77,16 @@ def save_trains(self):
             total_times_count += len(stop_times)
             times_count.append(len(stop_times))
             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.origin_dep_time).first()
+                train = self.session.query(Trip).filter_by(
+                    number=stop_time.trip_id,
+                    orig_dep_date=stop_time.origin_dep_time,
+                    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.origin_dep_time,
-                                  route_name=stop_time.route_name)
+                                  route_name=stop_time.route_name, source=self.name)
                     self.session.add(train)
                     self.session.commit()
 

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'])