pipeline/model.py

from sqlalchemy import Column, Integer, String, Float, create_engine, DateTime, PickleType
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker

Base = declarative_base()

class Stock(Base):
    __tablename__ = 'stock'

    id = Column(Integer, primary_key=True)
    ticker = Column(String(50))
    weight = Column(Float, nullable=True)
    display_name = Column(String(50), nullable=False)
    date = Column(DateTime, nullable=False)
    order = Column(Integer, autoincrement=True, nullable=True)

    def __repr__(self):
        return f"<{self.ticker}\
                    {self.date}\
                        {self.display_name}\
                            {round(self.weight * 100)}%>"


class Portfolio(Base):
    __tablename__ = 'portfolio'
    
    id = Column(Integer, primary_key=True)
    stocks = Column(PickleType, nullable=False)
    cached_result = Column(PickleType, nullable=True)
    # data = Column(PickleType, nullable=False)
    date = Column(DateTime, nullable=False)
    order = Column(Integer, autoincrement=True, nullable=True)
    

# db_url = 'sqlite:///local_db.db'  # Replace 'stocks.db' with the desired database name and location

# engine = create_engine(db_url)
# Base.metadata.create_all(engine)

# Session = sessionmaker(bind=engine)
# session = Session()