Merge pull request #122 from tsugumi-sys/feature/imple-polygon-financ…

…ials-store

Store: Implement polygon financials sqlalchemy sotre

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

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()

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)

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 ###

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