stock / data /database.py
feliponi's picture
Release 0.002
2293f58
# data/database.py
from sqlalchemy import create_engine, Column, Integer, String, JSON
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from datetime import datetime
Base = declarative_base()
class CompanyData(Base):
__tablename__ = 'company_data'
id = Column(Integer, primary_key=True)
ticker = Column(String)
data_type = Column(String)
data = Column(JSON)
date = Column(String)
class DatabaseManager:
def __init__(self, db_url='sqlite:///financial_data.db'):
self.engine = create_engine(db_url)
self.Session = sessionmaker(bind=self.engine)
Base.metadata.create_all(self.engine)
def save_data(self, ticker, data_type, data):
session = self.Session()
try:
new_entry = CompanyData(
ticker=ticker,
data_type=data_type,
data=data,
date=datetime.now().strftime('%Y-%m-%d %H:%M:%S')
)
session.add(new_entry)
session.commit()
except Exception as e:
print(f"Error saving data: {e}")
finally:
session.close()
def get_historical_data(self, ticker):
session = self.Session()
try:
financials = session.query(CompanyData).filter(
CompanyData.ticker == ticker,
CompanyData.data_type == 'financials'
).order_by(CompanyData.date.desc()).first()
news = session.query(CompanyData).filter(
CompanyData.ticker == ticker,
CompanyData.data_type == 'news'
).order_by(CompanyData.date.desc()).all()
return {
'financials': financials.data if financials else None,
'news': [n.data for n in news]
}
finally:
session.close()