diff --git a/stocklake/environment_variables.py b/stocklake/environment_variables.py index 3488720..f8037da 100644 --- a/stocklake/environment_variables.py +++ b/stocklake/environment_variables.py @@ -13,6 +13,10 @@ def __init__(self, name: str, default: Any): def defined(self): return self.name in os.environ + @property + def env_name(self) -> str: + return self.name + def get(self): if (val := os.getenv(self.name)) is not None: return val diff --git a/stocklake/polygonapi/stores.py b/stocklake/polygonapi/stores.py index fb1044d..45116d0 100644 --- a/stocklake/polygonapi/stores.py +++ b/stocklake/polygonapi/stores.py @@ -4,12 +4,13 @@ from sqlalchemy import orm +from stocklake.core.base_sqlalchemy_store import SQLAlchemyStore from stocklake.core.base_store import BaseStore from stocklake.core.constants import DATA_DIR from stocklake.stores.artifact.local_artifact_repo import LocalArtifactRepository from stocklake.stores.constants import StoreType +from stocklake.stores.db import models, schemas from stocklake.stores.db.database import LocalSession # noqa: E402 -from stocklake.stores.db.schemas import PreprocessedPolygonFinancialsData from stocklake.utils.file_utils import save_data_to_csv SAVE_ARTIFACTS_DIR = os.path.join(DATA_DIR, "nasdaqapi") @@ -22,7 +23,9 @@ def __init__( self.sqlalchemy_session = sqlalchemy_session def save( - self, store_type: StoreType, data: List[PreprocessedPolygonFinancialsData] + self, + store_type: StoreType, + data: List[schemas.PreprocessedPolygonFinancialsData], ): if store_type == StoreType.LOCAL_ARTIFACT: repository = LocalArtifactRepository(SAVE_ARTIFACTS_DIR) @@ -32,3 +35,24 @@ def save( repository.save_artifact(csv_file_path) else: raise NotImplementedError + + +class PolygonFinancialsDataSQLAlchemyStore(SQLAlchemyStore): + def __init__(self, session: orm.sessionmaker[orm.session.Session]): + self.session = session + + def create(self, data: List[schemas.PolygonFinancialsDataCreate]): + with self.session() as session, session.begin(): + session.add_all( + [models.PolygonFinancialsData(**d.model_dump()) for d in data] + ) + + def read(self): + raise NotImplementedError() + + def update(self): + raise NotImplementedError() + + def delete(self): + with self.session() as session, session.begin(): + session.query(models.PolygonFinancialsData).delete() diff --git a/stocklake/stores/db/models.py b/stocklake/stores/db/models.py index 8051d79..79cb94e 100644 --- a/stocklake/stores/db/models.py +++ b/stocklake/stores/db/models.py @@ -55,13 +55,14 @@ class PolygonFinancialsData(Base): # net_cash_flow_from_operating_activities = Column(Float) # net_cash_flow_from_financing_activities_continuing = Column(Float) # - comprehensive income - loss_attributable_to_noncontrolling_interest = Column(Float) - loss_attributable_to_parent = Column(Float) + # loss_attributable_to_noncontrolling_interest = Column(Float) + # loss_attributable_to_parent = Column(Float) + comprehensive_income_loss_attributable_to_parent = Column(Float) other_comprehensive_income_loss = Column(Float) # other_comprehensive_income_loss_attributable_to_parent = Column(Float) comprehensive_income_loss = Column(Float) # - income statement - income_loss_before_equity_method_investments = Column(Float) + # income_loss_before_equity_method_investments = Column(Float) # diluted_earnings_per_share = Column(Float) # income_loss_from_equity_method_investments = Column(Float) operating_expenses = Column(Float) diff --git a/stocklake/stores/db_migrations/versions/3ae9ebc67811_remove_unnecessary_columns_from_polygon_.py b/stocklake/stores/db_migrations/versions/3ae9ebc67811_remove_unnecessary_columns_from_polygon_.py new file mode 100644 index 0000000..83bc744 --- /dev/null +++ b/stocklake/stores/db_migrations/versions/3ae9ebc67811_remove_unnecessary_columns_from_polygon_.py @@ -0,0 +1,72 @@ +"""Remove unnecessary columns from polygon financials table + +Revision ID: 3ae9ebc67811 +Revises: 29fe33d1bd8a +Create Date: 2024-06-06 14:06:43.827791 + +""" +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "3ae9ebc67811" +down_revision: Union[str, None] = "29fe33d1bd8a" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "polygonapi_financials_data", + sa.Column( + "comprehensive_income_loss_attributable_to_parent", + sa.Float(), + nullable=True, + ), + ) + op.drop_column( + "polygonapi_financials_data", "income_loss_before_equity_method_investments" + ) + op.drop_column( + "polygonapi_financials_data", "loss_attributable_to_noncontrolling_interest" + ) + op.drop_column("polygonapi_financials_data", "loss_attributable_to_parent") + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column( + "polygonapi_financials_data", + sa.Column( + "loss_attributable_to_parent", + sa.DOUBLE_PRECISION(precision=53), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "polygonapi_financials_data", + sa.Column( + "loss_attributable_to_noncontrolling_interest", + sa.DOUBLE_PRECISION(precision=53), + autoincrement=False, + nullable=True, + ), + ) + op.add_column( + "polygonapi_financials_data", + sa.Column( + "income_loss_before_equity_method_investments", + sa.DOUBLE_PRECISION(precision=53), + autoincrement=False, + nullable=True, + ), + ) + op.drop_column( + "polygonapi_financials_data", "comprehensive_income_loss_attributable_to_parent" + ) + # ### end Alembic commands ### diff --git a/tests/polygonapi/test_stores.py b/tests/polygonapi/test_stores.py index e4f485c..7873788 100644 --- a/tests/polygonapi/test_stores.py +++ b/tests/polygonapi/test_stores.py @@ -1,20 +1,89 @@ import os +import pytest + +from stocklake.environment_variables import STOCKLAKE_POLYGON_API_KEY from stocklake.polygonapi.data_loader import PolygonFinancialsDataLoader from stocklake.polygonapi.preprocessor import PolygonFinancialsDataPreprocessor -from stocklake.polygonapi.stores import SAVE_ARTIFACTS_DIR, PolygonFinancialsDataStore +from stocklake.polygonapi.stores import ( + SAVE_ARTIFACTS_DIR, + PolygonFinancialsDataSQLAlchemyStore, + PolygonFinancialsDataStore, +) from stocklake.stores.constants import StoreType +from stocklake.stores.db import models, schemas from tests.polygonapi.test_data_loader import MockPolygonAPIServer # noqa: F401 +from tests.stores.db.utils import SessionLocal # noqa: F401 -def test_polygon_financials_store_local_artifact( +@pytest.fixture +def polygon_financials_data( MockPolygonAPIServer, # noqa: F811 monkeypatch, ): - monkeypatch.setenv("STOCKLAKE_POLYGON_API_KEY", "dummy_key") + monkeypatch.setenv(STOCKLAKE_POLYGON_API_KEY.env_name, "dummy_key") dataloader = PolygonFinancialsDataLoader() preprocessor = PolygonFinancialsDataPreprocessor() data = preprocessor.process(dataloader.download(["MSFT"])) + yield data + + +def test_polygon_financials_store_local_artifact( + polygon_financials_data, +): store = PolygonFinancialsDataStore() - store.save(StoreType.LOCAL_ARTIFACT, data) + store.save(StoreType.LOCAL_ARTIFACT, polygon_financials_data) assert os.path.exists(os.path.join(SAVE_ARTIFACTS_DIR, "financials_data.csv")) + + +def test_PolygonFinancialsDataSQLAlchemyStore_create( + polygon_financials_data, + SessionLocal, # noqa: F811 +): + data_length = len(polygon_financials_data) + with SessionLocal() as session, session.begin(): + res = session.query(models.PolygonFinancialsData).all() + assert len(res) == 0 + + store = PolygonFinancialsDataSQLAlchemyStore(SessionLocal) + store.create( + [ + schemas.PolygonFinancialsDataCreate(**d.dict()) + for d in polygon_financials_data + ] + ) + + with SessionLocal() as session, session.begin(): + res = session.query(models.PolygonFinancialsData).all() + assert len(res) == data_length + + +def test_PolygonFinancialsDataSQLAlchemyStore_delete( + polygon_financials_data, + SessionLocal, # noqa: F811 +): + data_length = len(polygon_financials_data) + with SessionLocal() as session, session.begin(): + res = session.query(models.PolygonFinancialsData).all() + assert len(res) == 0 + + store = PolygonFinancialsDataSQLAlchemyStore(SessionLocal) + store.create( + [ + schemas.PolygonFinancialsDataCreate(**d.dict()) + for d in polygon_financials_data + ] + ) + + # check data is created + with SessionLocal() as session, session.begin(): + res = session.query(models.PolygonFinancialsData).all() + assert len(res) == data_length + + # delete all rows + store.delete() + + # check successfully deleted. + with SessionLocal() as session, session.begin(): + res = session.query(models.PolygonFinancialsData).all() + assert len(res) == 0