Merge pull request #126 from tsugumi-sys/feature/adding-polygon-sqlst…

…ores

Stores: Addigng polygon financials data of sqlstores

stocklake/polygonapi/stores.py

@@ -33,6 +33,12 @@ def save(
                 csv_file_path = os.path.join(tempdir, "financials_data.csv")
                 save_data_to_csv([d.dict() for d in data], csv_file_path)
                 repository.save_artifact(csv_file_path)
+        elif store_type == StoreType.POSTGRESQL:
+            sqlstore = PolygonFinancialsDataSQLAlchemyStore(self.sqlalchemy_session)
+            sqlstore.delete()
+            sqlstore.create(
+                [schemas.PolygonFinancialsDataCreate(**d.dict()) for d in data]
+            )
         else:
             raise NotImplementedError
 

tests/polygonapi/test_stores.py

@@ -36,6 +36,17 @@ def test_polygon_financials_store_local_artifact(
     assert os.path.exists(os.path.join(SAVE_ARTIFACTS_DIR, "financials_data.csv"))
 
 
+def test_polygon_financials_store_postgresql(
+    SessionLocal,  # noqa: F811
+    polygon_financials_data,
+):
+    store = PolygonFinancialsDataStore(SessionLocal)
+    store.save(StoreType.POSTGRESQL, polygon_financials_data)
+    with SessionLocal() as session, session.begin():
+        res = session.query(models.PolygonFinancialsData).all()
+    assert len(res) == 1
+
+
 def test_PolygonFinancialsDataSQLAlchemyStore_create(
     polygon_financials_data,
     SessionLocal,  # noqa: F811