sqlalchemy-utils

import os from sqlalchemy import create_engine, Column, Integer, String from sqlalchemy.orm import declarative_base, sessionmaker from sqlalchemy_utils import database_exists, create_database, drop_database, ChoiceType # --- Database Setup --- DB_USER = os.environ.get('DB_USER', 'testuser') DB_PASS = os.environ.get('DB_PASS', 'testpass') DB_HOST = os.environ.get('DB_HOST', 'localhost') DB_NAME = os.environ.get('DB_NAME', 'test_db') DATABASE_URL = f"postgresql://{DB_USER}:{DB_PASS}@{DB_HOST}/{DB_NAME}" engine = create_engine(DATABASE_URL) # Create database if it doesn't exist if not database_exists(engine.url): create_database(engine.url) print(f"Database '{DB_NAME}' created.") else: print(f"Database '{DB_NAME}' already exists.") Session = sessionmaker(bind=engine) Base = declarative_base() # --- Define a Model with ChoiceType --- class TaskPriority: LOW = 'low' MEDIUM = 'medium' HIGH = 'high' class Task(Base): __tablename__ = 'tasks' id = Column(Integer, primary_key=True) name = Column(String(50), nullable=False) priority = Column(ChoiceType(TaskPriority, impl=String(10))) def __repr__(self): return f"<Task(id={self.id}, name='{self.name}', priority='{self.priority}')>" # --- ORM Operations --- Base.metadata.create_all(engine) # Create tables session = Session() # Add tasks task1 = Task(name="Write documentation", priority=TaskPriority.HIGH) task2 = Task(name="Review code", priority=TaskPriority.MEDIUM) task3 = Task(name="Deploy update", priority=TaskPriority.LOW) session.add_all([task1, task2, task3]) session.commit() # Query tasks print("\nAll tasks:") for task in session.query(Task).all(): print(task) # Clean up (optional: uncomment to drop the database) # session.close() # drop_database(engine.url) # print(f"Database '{DB_NAME}' dropped.")