Issue/sfs3 (#95)

* add s3fs to requirements * tidy up logs * add to tests * lint * only check if forecasts have been deleted * lint
openclimatefix · Nov 16, 2024 · d402492 · d402492
1 parent bf92627
commit d402492
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 15 deletions.
diff --git a/database-cleanup/database_cleanup/app.py b/database-cleanup/database_cleanup/app.py
@@ -21,7 +21,16 @@
 from sqlalchemy.orm import Session, sessionmaker
 import pandas as pd
 
+
+logging.basicConfig(
+    level=getattr(logging, os.getenv("LOGLEVEL", "INFO")),
+    format="[%(asctime)s] {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s",
+)
 _log = logging.getLogger(__name__)
+# Get rid of the verbose logs
+logging.getLogger("sqlalchemy").setLevel(logging.ERROR)
+logging.getLogger("aiobotocore").setLevel(logging.ERROR)
+
 
 version = importlib.metadata.version("database-cleanup")
 
@@ -167,20 +176,20 @@ def main(
                 limit=batch_size,
             )
 
+            if len(forecast_uuids) == 0:
+                _log.info(f"Done deleting forecasts made before {date}")
+                _log.info(
+                    f"A total of {num_forecast_deleted} (and corresponding values) "
+                    f"were deleted from the database."
+                )
+                _log.info("Exiting.")
+                return
+
             if save_dir is not None:
                 save_forecast_and_values(
                     session=session, forecast_uuids=forecast_uuids, directory=save_dir
                 )
 
-        if len(forecast_uuids) == 0:
-            _log.info(f"Done deleting forecasts made before {date}")
-            _log.info(
-                f"A total of {num_forecast_deleted} (and corresponding values) were deleted from"
-                " the database."
-            )
-            _log.info("Exiting.")
-            return
-
         if do_delete:
             # Not that it is important to run this in a transaction for atomicity.
             with Session.begin() as session:

diff --git a/database-cleanup/tests/test_app.py b/database-cleanup/tests/test_app.py
@@ -2,6 +2,7 @@
 import uuid
 import os
 import tempfile
+import pandas as pd
 
 import pytest
 import sqlalchemy as sa
@@ -130,12 +131,18 @@ def test_app(session: Session, site, batch_size: int, date_str: str | None, expe
         ).one()
         assert num_values_left == expected * num_values
 
-        # check that forecast.csv and forecast_values.csv are saved
-        date = format_date(date_str).isoformat()
-        assert os.path.exists(f"{tmpdirname}")
-        assert os.path.exists(f"{tmpdirname}/{date}")
-        assert os.path.exists(f"{tmpdirname}/{date}/forecast.csv")
-        assert os.path.exists(f"{tmpdirname}/{date}/forecast_value.csv")
+        if num_forecasts_left < num_forecasts:
+            # check that forecast.csv and forecast_values.csv are saved
+            date = format_date(date_str).isoformat()
+            assert os.path.exists(f"{tmpdirname}")
+            assert os.path.exists(f"{tmpdirname}/{date}")
+            assert os.path.exists(f"{tmpdirname}/{date}/forecast.csv")
+            assert os.path.exists(f"{tmpdirname}/{date}/forecast_value.csv")
+
+            forecast_df = pd.read_csv(f"{tmpdirname}/{date}/forecast.csv")
+            forecast_value_df = pd.read_csv(f"{tmpdirname}/{date}/forecast_value.csv")
+            for data in [forecast_df, forecast_value_df]:
+                assert len(data) > 0
 
 
 @freeze_time("2020-01-11 00:01")