41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
# helper.py
|
|
from sqlmodel import SQLModel, Session, create_engine, select
|
|
from proxy import Hotel, DATABASE_URL
|
|
|
|
engine = create_engine(DATABASE_URL, echo=False)
|
|
|
|
def add_hotel(hotel_id, email):
|
|
with Session(engine) as session:
|
|
existing = session.exec(select(Hotel).where(Hotel.hotel_id == hotel_id)).first()
|
|
if existing:
|
|
print(f"Hotel '{hotel_id}' already exists with email: {existing.email}")
|
|
return
|
|
hotel = Hotel(hotel_id=hotel_id, email=email)
|
|
session.add(hotel)
|
|
session.commit()
|
|
print(f"Added: {hotel_id} -> {email}")
|
|
|
|
def remove_hotel(hotel_id):
|
|
with Session(engine) as session:
|
|
hotel = session.exec(select(Hotel).where(Hotel.hotel_id == hotel_id)).first()
|
|
if not hotel:
|
|
print(f"Hotel '{hotel_id}' not found.")
|
|
return
|
|
session.delete(hotel)
|
|
session.commit()
|
|
print(f"Removed: {hotel_id}")
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
if len(sys.argv) < 3:
|
|
print("Usage: python helper.py add <hotel_id> <email>")
|
|
print(" or: python helper.py remove <hotel_id>")
|
|
exit(1)
|
|
command = sys.argv[1].lower()
|
|
if command == "add" and len(sys.argv) == 4:
|
|
add_hotel(sys.argv[2], sys.argv[3])
|
|
elif command == "remove" and len(sys.argv) == 3:
|
|
remove_hotel(sys.argv[2])
|
|
else:
|
|
print("Invalid command or arguments.")
|