pokemans/models.py
2024-03-03 14:38:44 -05:00

79 lines
2.2 KiB
Python

from sqlalchemy import create_engine, Column, Integer, String, Float
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import os
from contextlib import contextmanager
Base = declarative_base()
Session = sessionmaker()
@contextmanager
def session_scope():
session = get_session()
try:
yield session
session.commit()
except:
session.rollback()
raise
finally:
session.close()
class Submission(Base):
__tablename__ = 'submissions'
id = Column(String, primary_key=True)
title = Column(String)
name = Column(String)
url = Column(String)
score = Column(Integer)
num_comments = Column(Integer)
created_utc = Column(Float)
selftext = Column(String)
permalink = Column(String)
upvote_ratio = Column(Float)
def get_engine(database_url=os.getenv("POKEMANS_DB_URL", "sqlite:///pokemans.db")):
engine = create_engine(database_url)
Session.configure(bind=engine)
return engine
def create_db():
engine = get_engine()
Base.metadata.create_all(engine)
def reset_db():
engine = get_engine()
Base.metadata.drop_all(engine)
Base.metadata.create_all(engine)
def get_session():
return Session()
def insert_submission(session, submission):
session.add(submission)
session.commit()
def submission_exists(session, submission_id):
return session.query(Submission).filter(Submission.id == submission_id).first() is not None
def get_all_submissions(session):
return session.query(Submission).all()
def delete_submission(session, submission_id):
session.query(Submission).filter(Submission.id == submission_id).delete()
session.commit()
def update_submission(session, submission):
session.query(Submission).filter(Submission.id == submission.id).update({
'title': submission.title,
'name': submission.name,
'url': submission.url,
'score': submission.score,
'num_comments': submission.num_comments,
'created_utc': submission.created_utc,
'selftext': submission.selftext,
'permalink': submission.permalink,
'upvote_ratio': submission.upvote_ratio
})
session.commit()