Let plugins register websocket endpoints

.vscode/settings.json

@@ -1,7 +1,5 @@
 {
-    "spellright.language": [
-        "en_US"
-    ],
+    "spellright.language": ["en_US"],
     "spellright.documentTypes": [
         "markdown",
         "latex",
@@ -14,10 +12,10 @@
     "python.linting.flake8Enabled": true,
     "workbench.editor.wrapTabs": true,
     "python.formatting.provider": "black",
-    "python.testing.pytestArgs": [
-        "flexmeasures"
-    ],
+    "python.testing.pytestArgs": ["flexmeasures"],
     "python.testing.unittestEnabled": false,
     "python.testing.pytestEnabled": true,
-    "python.analysis.autoImportCompletions": true
+    "python.analysis.autoImportCompletions": true,
+    "stm32-for-vscode.openOCDPath": false,
+    "stm32-for-vscode.armToolchainPath": false
 }

flexmeasures/api/common/routes.py

@@ -1,4 +1,9 @@
-from flask_security import auth_token_required
+from collections import deque
+import os
+
+from flask import current_app, request, Response
+from flask_security import auth_token_required, login_required
+from werkzeug.exceptions import NotFound, abort
 
 from flexmeasures.auth.decorators import roles_required
 from flexmeasures.api.common import flexmeasures_api as flexmeasures_api_ops
@@ -20,3 +25,20 @@ def get_task_run():
 @roles_required("task-runner")
 def post_task_run():
     return ops_impl.post_task_run()
+
+
+@flexmeasures_api_ops.route("/logs")
+@login_required
+@roles_required("debugger")
+def show_logs():
+    """Show server logs for debugging."""
+    if current_app.config.get("LOGGING_LEVEL") != "DEBUG":
+        raise NotFound
+
+    log_file = "flexmeasures.log"
+    n = int(request.args.get("tail", 200))
+    if not os.path.exists(log_file):
+        abort(404, "Log file not found")
+    with open(log_file, "r") as f:
+        last_n_lines = deque(f, maxlen=n)
+    return Response("".join(last_n_lines), mimetype="text/plain")

flexmeasures/app.py

@@ -20,6 +20,7 @@
 from rq import Queue
 
 from flexmeasures.data.services.job_cache import JobCache
+from flexmeasures.ws import sock
 
 
 def create(  # noqa C901
@@ -56,6 +57,7 @@ def create(  # noqa C901
     cfg_location = find_flexmeasures_cfg()  # Find flexmeasures.cfg location
     # Create app
     app = Flask("flexmeasures")
+    sock.init_app(app)
 
     if env is not None:  # overwrite
         app.config["FLEXMEASURES_ENV"] = env

flexmeasures/conftest.py

@@ -309,6 +309,7 @@ def create_roles_users(db, test_accounts) -> dict[str, User]:
             password="testtest",
         )
     )
+    db.session.commit()
     return {user.username: user.id for user in new_users}
 
 

flexmeasures/data/models/data_sources.py

@@ -178,8 +178,8 @@ def get_data_source_info(cls: type) -> dict:
     @property
     def data_source(self) -> "DataSource":
         """DataSource property derived from the `source_info`: `source_type` (scheduler, forecaster or reporter), `model` (e.g AggregatorReporter)
-        and `attributes`. It looks for a data source in the database the marges the `source_info` and, in case of not finding any, it creates a new one.
-        This property gets created once and it's cached for the rest of the lifetime of the DataGenerator object.
+        and `attributes`. It looks for a data source in the database that matches the `source_info` and, in case of not finding any, it creates a new one.
+        This property gets created once, and it's cached for the rest of the lifetime of the DataGenerator object.
         """
 
         from flexmeasures.data.services.data_sources import get_or_create_source

flexmeasures/data/services/utils.py

@@ -7,7 +7,6 @@
 from copy import deepcopy
 import inspect
 
-import click
 from sqlalchemy import JSON, String, cast, literal
 from flask import current_app
 from rq import Queue
@@ -130,7 +129,6 @@ def get_or_create_model(
         model = model_class(**init_kwargs)
         db.session.add(model)
         db.session.flush()  # assign ID
-        click.echo(f"Created {repr(model)}")
     return model
 
 

flexmeasures/ui/static/openapi-specs.json

@@ -100,6 +100,7 @@
     "/api/ops/ping": {},
     "/api/ops/getLatestTaskRun": {},
     "/api/ops/postLatestTaskRun": {},
+    "/api/ops/logs": {},
     "/api/v3_0/sensors/{id}": {
       "delete": {
         "summary": "Delete a sensor",

flexmeasures/utils/coding_utils.py

@@ -231,3 +231,42 @@ def __hash__(self):
                 f"Unhashable object: {self} has no ID. Consider calling `db.session.flush()` before using {type(self).__name__} objects in sets or as dictionary keys."
             )
         return hash(self.id)
+
+
+def only_if_timer_due(*kwarg_names):
+    """
+    Decorator that runs a method only if its timer is due.
+
+    Timer name is derived from:
+        "<function_name> for <kwarg1_value>, <kwarg2_value>, ..."
+
+    :param kwarg_names: names of kwargs to include in the timer name
+    """
+
+    def decorator(func):
+        sig = inspect.signature(func)  # get function signature
+
+        @functools.wraps(func)
+        def wrapper(self, *args, **kwargs):
+            timer_name = func.__name__  # start with the function name
+
+            if kwarg_names:
+                bound_args = sig.bind(self, *args, **kwargs)
+                bound_args.apply_defaults()
+
+                values = []
+                for name in kwarg_names:
+                    if name in bound_args.arguments:
+                        values.append(str(bound_args.arguments[name]))
+                    else:
+                        raise ValueError(f"Keyword '{name}' not found in function call")
+
+                timer_name += " for " + ", ".join(values)
+
+            if not self._is_timer_due(timer_name):
+                return  # skip if timer not due
+            return func(self, *args, **kwargs)
+
+        return wrapper
+
+    return decorator

flexmeasures/utils/config_utils.py

@@ -27,6 +27,12 @@
     "version": 1,
     "formatters": {
         "default": {"format": "[FLEXMEASURES][%(asctime)s] %(levelname)s: %(message)s"},
+        "s2python": {
+            "format": "\033[94m[FLEXMEASURES][%(asctime)s][%(name)s] %(levelname)s: %(message)s\033[0m"
+        },  # blue
+        "s2rm": {
+            "format": "\033[91m[FLEXMEASURES][%(asctime)s][%(name)s] %(levelname)s: %(message)s\033[0m"
+        },
         "detail": {
             "format": "[FLEXMEASURES][%(asctime)s] %(levelname)s: %(message)s [logged in %(pathname)s:%(lineno)d]"
         },
@@ -37,6 +43,16 @@
             "stream": sys.stdout,
             "formatter": "default",
         },
+        "s2console": {  # handler specific to s2python
+            "class": "logging.StreamHandler",
+            "stream": sys.stdout,
+            "formatter": "s2python",
+        },
+        "s2rmconsole": {  # handler specific to s2python-rm
+            "class": "logging.StreamHandler",
+            "stream": sys.stdout,
+            "formatter": "s2rm",
+        },
         "file": {
             "class": "logging.handlers.RotatingFileHandler",
             "level": "INFO",
@@ -46,7 +62,15 @@
             "backupCount": 6,
         },
     },
-    "root": {"level": "INFO", "handlers": ["console", "file"], "propagate": True},
+    "loggers": {
+        "root": {"level": "INFO", "handlers": ["console", "file"], "propagate": True},
+        "s2python": {"level": "DEBUG", "handlers": ["s2console"], "propagate": False},
+        "flexmeasures_s2.rm": {
+            "level": "DEBUG",
+            "handlers": ["s2rmconsole"],
+            "propagate": False,
+        },
+    },
 }
 
 

flexmeasures/utils/time_utils.py

@@ -21,6 +21,13 @@ def server_now() -> datetime:
     return datetime.now(get_timezone())
 
 
+def floored_server_now(resolution: timedelta) -> datetime:
+    """Return the current server time floored to the nearest multiple of `resolution`."""
+    ref = pytz.utc.localize(datetime.min)
+    _tz = get_timezone()
+    return (ref + ((server_now() - ref) // resolution) * resolution).astimezone(_tz)
+
+
 def ensure_local_timezone(
     dt: pd.Timestamp | datetime, tz_name: str = "Europe/Amsterdam"
 ) -> pd.Timestamp | datetime:

flexmeasures/ws/__init__.py

@@ -0,0 +1,18 @@
+import importlib
+import pkgutil
+from flask import Blueprint, current_app
+from flask_security import auth_token_required
+
+from flask_sock import Sock
+
+sock = Sock()
+
+
+def import_all_modules(package_name):
+    package = importlib.import_module(package_name)
+    for _, name, _ in pkgutil.iter_modules(package.__path__):
+        importlib.import_module(f"{package_name}.{name}")
+
+
+# we need to import all the modules to run the route decorators
+import_all_modules("flexmeasures.ws")

flexmeasures/ws/ping1.py

@@ -0,0 +1,31 @@
+import logging
+from flexmeasures.ws import sock
+from flask import current_app
+from flexmeasures import Sensor
+from sqlalchemy import select, func
+import uuid
+import json
+
+logger = logging.getLogger(__name__)
+
+
+@sock.route("/ping1")
+async def echo1(ws):
+    headers = ws.environ  # Access all headers from the connection
+    client_id = str(uuid.uuid4())
+
+    logger.info("-----------------------------------------")
+    logger.info(f"Received headers: {headers}")
+    logger.info("-----------------------------------------")
+    logger.info(f"Type of ws: {type(ws)}")
+    logger.info(f"Client ID: {client_id}")
+    await ws.send(
+        json.dumps({"type": "metadata", "headers": {"X-Server-Header": "ServerValue"}})
+    )
+    while True:
+        data = await ws.receive()
+        logger.error("ping1>" + data)
+        if data == "close":
+            break
+        # sensors = current_app.db.session.execute(select(func.count(Sensor.id))).scalar()
+        await ws.send(data)

flexmeasures/ws/ping2.py

@@ -0,0 +1,14 @@
+import logging
+from flexmeasures.ws import sock
+
+logger = logging.Logger(__name__)
+
+
+@sock.route("/ping2")
+def echo2(ws):
+    while True:
+        data = ws.receive()
+        logger.error("ping2>" + data)
+        if data == "close":
+            break
+        ws.send(data)

flexmeasures/ws/tests/conftest.py

@@ -0,0 +1,31 @@
+import threading
+import time
+
+import pytest
+import pytest_asyncio
+import websockets
+from werkzeug.serving import make_server
+
+
+@pytest.fixture(scope="module")
+def server(app):
+    """Run Flask app with Sock in a thread for testing WebSocket"""
+    srv = make_server("127.0.0.1", 5005, app)
+    thread = threading.Thread(target=srv.serve_forever)
+    thread.start()
+    time.sleep(0.1)  # wait for server to start
+    yield "ws://127.0.0.1:5005"
+    srv.shutdown()
+    thread.join()
+
+
+@pytest_asyncio.fixture
+async def connect_to_ws(server):
+    """Yield a callable to connect to a given WS endpoint by name."""
+
+    async def connect(endpoint_name):
+        url = f"{server}/{endpoint_name}"
+        conn = await websockets.connect(url)
+        return conn
+
+    yield connect

flexmeasures/ws/tests/test_s2_client_rm.py

@@ -0,0 +1,24 @@
+import pytest
+
+import logging
+import websockets
+
+logger = logging.getLogger("s2python")
+SERVER_URL = "ws://127.0.0.1:5000"
+
+
+@pytest.mark.asyncio
+async def test_ping2_echo(connect_to_ws):
+
+    # Connect to WS endpoint
+    ws = await connect_to_ws("ping2")
+
+    # Send a message
+    await ws.send("hello")
+    resp = await ws.recv()
+    assert resp == "hello", "echo should return the same message"
+
+    # Trigger server-side close
+    await ws.send("close")
+    with pytest.raises(websockets.exceptions.ConnectionClosedOK):
+        await ws.recv(), "expected that, after sending 'close', server breaks loop; connection closes"

flexmeasures/ws/v1.py

@@ -0,0 +1,38 @@
+import logging
+from flexmeasures.ws import sock
+from flask import current_app
+from flexmeasures import Sensor
+from sqlalchemy import select, func
+import json
+
+logger = logging.Logger(__name__)
+
+
+@sock.route("/v1")
+def header_test(ws):
+    # Get all headers
+    all_headers = {
+        k[5:].lower().replace("_", "-"): v
+        for k, v in ws.environ.items()
+        if k.startswith("HTTP_")
+    }
+
+    # Get specific header if needed
+    custom_header = ws.environ.get("HTTP_X_CUSTOM_HEADER")
+    # show the type of ws
+    logger.info(f"Type of ws: {type(ws)}")
+    logger.info(f"All headers: {all_headers}")
+    logger.info(f"Custom header: {custom_header}")
+
+    # Send initial message with metadata
+    ws.send(
+        json.dumps({"type": "metadata", "headers": {"X-Server-Header": "ServerValue"}})
+    )
+
+    while True:
+        data = ws.receive()
+        logger.error("v1>" + data)
+        if data == "close":
+            break
+        sensors = current_app.db.session.execute(select(func.count(Sensor.id))).scalar()
+        ws.send(str(sensors))