pipeline/db_operation.py

from model import Stock
from sqlalchemy import create_engine
from sqlalchemy.orm import Session


class db_operator():
    def __init__(self, db_url):
        self.engine = create_engine(db_url)

    def get_stocks_between(self, d1, d2):
        with Session(self.engine) as session:
            return session.query(Stock).filter(
                Stock.date.between(d1, d2)
                ).all()
    
    def add_stock(self, stock_data: dict):
        with Session(self.engine) as session:
            new_stock = Stock(**stock_data)
            session.add(new_stock)
            session.commit()
    
    def delete_stocks_between(self, d1, d2):
        with Session(self.engine) as session:
            session.query(Stock).filter(
                Stock.date.between(d1, d2)
                ).delete()
            session.commit()
    
    def delete_all_stocks(self):
        with Session(self.engine) as session:
            session.query(Stock).delete()
            session.commit()