79 lines
2.2 KiB
Python
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() |