trackbert/classes/database.py
Kumi 1ae1a88975
feat: Add write_event method to Database class
This commit adds a new method called `write_event` to the `Database` class in `database.py`. The method takes an `Event` object as a parameter and adds it to the session before committing the changes to the database. This change simplifies the code in the `create_event` method by making use of the new `write_event` method.

The `write_event` method is used in the `Tracker` class in `tracker.py` to replace the call to `create_event`. This change improves the clarity and readability of the code.

Additionally, in the `Tracker` class, the `api` object is no longer created using `KeyDelivery.from_config`, as the `api` object is no longer referenced in the code. This change is made to remove unused code and improve code quality.

This commit also adds some new import statements to the `Tracker` class in `tracker.py` to support the new changes.

Lastly, two new tracker classes `KeyDelivery` and `PostAT` are added under the `trackers` directory. These new classes inherit from the `BaseTracker` class and implement the `get_status` method. The `KeyDelivery` class makes use of the `pykeydelivery.KeyDelivery` API to fetch tracking events, while the `PostAT` class uses the `postat.classes.api.PostAPI` API to fetch tracking events. The new tracker classes also define the supported carriers and their priorities.

Overall, these changes improve the functionality and structure of the codebase.
2023-08-24 14:46:54 +02:00

68 lines
2.3 KiB
Python

from sqlalchemy import Column, Integer, String
from sqlalchemy import create_engine, ForeignKey
from sqlalchemy.orm import sessionmaker, relationship
from sqlalchemy.ext.declarative import declarative_base
import json
Base = declarative_base()
class Shipment(Base):
__tablename__ = 'shipments'
id = Column(Integer, primary_key=True)
tracking_number = Column(String)
carrier = Column(String)
description = Column(String)
events = relationship("Event")
class Event(Base):
__tablename__ = 'events'
id = Column(Integer, primary_key=True)
shipment_id = Column(Integer, ForeignKey('shipments.id'))
event_time = Column(String)
event_description = Column(String)
raw_event = Column(String)
class Database:
def __init__(self, database_uri):
self.engine = create_engine(database_uri)
Session = sessionmaker(bind=self.engine)
self.session = Session()
def create_shipment(self, tracking_number, carrier, description=""):
new_shipment = Shipment(tracking_number=tracking_number, carrier=carrier, description=description)
self.session.add(new_shipment)
self.session.commit()
def get_shipment(self, tracking_number):
shipment = self.session.query(Shipment).filter(Shipment.tracking_number == tracking_number).first()
return shipment
def get_shipments(self):
shipments = self.session.query(Shipment).all()
return shipments
def create_event(self, shipment_id, event_time, event_description, raw_event):
if isinstance(raw_event, dict):
raw_event = json.dumps(raw_event)
new_event = Event(shipment_id=shipment_id, event_time=event_time, event_description=event_description, raw_event=raw_event)
self.write_event(new_event)
def write_event(self, event):
self.session.add(event)
self.session.commit()
def get_shipment_events(self, shipment_id):
shipment = self.session.query(Shipment).filter(Shipment.id == shipment_id).first()
return shipment.events if shipment else None
def get_latest_event(self, shipment_id):
event = self.session.query(Event).filter(Event.shipment_id == shipment_id).order_by(Event.event_time.desc()).first()
return event
def initialize_db(self):
Base.metadata.create_all(self.engine)