Severale adjustments based on reviewers' suggestions

dev-requirements.txt

@@ -1,5 +1,5 @@
 black
 graphviz >= 0.8
 jinja2 >= 2.0, < 4.0
-mypy >= 0.600, < 0.931
+mypy >= 0.600, < 0.813
 Pygments >= 2.2, < 3.0

src/prefect/tasks/cubejs/cubejs_tasks.py

@@ -40,6 +40,9 @@ class CubeJSQueryTask(Task):
             If the security context does not contain an expiration period,
             then a 7-day expiration period is added automatically.
             More info at: https://cube.dev/docs/security/context.
+        - wait_time_between_api_calls (int, optional): The number of seconds to
+            wait between API calls.
+            Default to 10.
         - max_wait_time (int, optional): The number of seconds to wait for the
             Cube.js load API to return a response.
         - **kwargs (optional): Additional keyword arguments to pass to the
@@ -57,6 +60,7 @@ def __init__(
         api_secret_env_var: str = "CUBEJS_API_SECRET",
         query: Union[Dict, List[Dict]] = None,
         security_context: Union[str, Dict] = None,
+        wait_time_between_api_calls: int = 10,
         max_wait_time: int = None,
         **kwargs,
     ):
@@ -66,6 +70,7 @@ def __init__(
         self.api_secret_env_var = api_secret_env_var
         self.query = query
         self.security_context = security_context
+        self.wait_time_between_api_calls = wait_time_between_api_calls
         self.max_wait_time = max_wait_time
         super().__init__(**kwargs)
 
@@ -76,6 +81,7 @@ def __init__(
         "api_secret_env_var",
         "query",
         "security_context",
+        "wait_time_between_api_calls",
         "max_wait_time",
     )
     def run(
@@ -86,6 +92,7 @@ def run(
         api_secret_env_var: str = "CUBEJS_API_SECRET",
         query: Union[Dict, List[Dict]] = None,
         security_context: Union[str, Dict] = None,
+        wait_time_between_api_calls: int = 10,
         max_wait_time: int = None,
     ):
         """
@@ -114,6 +121,9 @@ def run(
                 If the security context does not contain an expiration period,
                 then a 7-day expiration period is added automatically.
                 More info at https://cube.dev/docs/security/context.
+            - wait_time_between_api_calls (int, optional): The number of seconds to
+                wait between API calls.
+                Default to 10.
             - max_wait_time (int, optional): The number of seconds to wait for the
                 Cube.js load API to return a response.
 
@@ -154,7 +164,7 @@ def run(
 
             extended_context = security_context
             if "exp" not in security_context and "expiresIn" not in security_context:
-                extended_context["epxiresIn"] = "7d"
+                extended_context["expiresIn"] = "7d"
             api_token = jwt.encode(
                 payload=extended_context, key=secret, algorithm="HS256"
             )
@@ -172,7 +182,9 @@ def run(
 
         params = {"query": json.dumps(query)}
 
-        wait_time_between_api_calls = 10
+        wait_api_call_secs = (
+            wait_time_between_api_calls if wait_time_between_api_calls > 0 else 10
+        )
         elapsed_wait_time = 0
         while not max_wait_time or elapsed_wait_time <= max_wait_time:
 
@@ -185,19 +197,21 @@ def run(
                     if "error" in data.keys() and "Continue wait" in data["error"]:
                         msg = (
                             "Cube.js load API still running."
-                            "Waiting {wait_time_between_api_calls} seconds before retrying"
+                            "Waiting {wait_api_call_secs} seconds before retrying"
                         )
-                        self.logger.warning(msg)
-                        time.sleep(wait_time_between_api_calls)
-                        elapsed_wait_time += wait_time_between_api_calls
+                        self.logger.info(msg)
+                        time.sleep(wait_api_call_secs)
+                        elapsed_wait_time += wait_api_call_secs
                         continue
 
                     else:
                         return data
 
                 else:
                     raise FAIL(
-                        message=f"Cube.js load API failed!. Error is: {response.reason}"
+                        message=f"Cube.js load API failed! Error is: {response.reason}"
                     )
 
-        raise FAIL(message="Cube.js load API took too long to provide a response.")
+        raise FAIL(
+            message=f"Cube.js load API took longer than {max_wait_time} seconds to provide a response."
+        )

tests/tasks/cubejs/test_cubejs_tasks.py

@@ -1,4 +1,5 @@
 import logging
+import jwt
 import pytest
 from prefect.engine.signals import FAIL
 from prefect.tasks.cubejs import CubeJSQueryTask
@@ -65,7 +66,7 @@ def test_run_with_failing_api_raises(self):
         with pytest.raises(FAIL) as exc:
             cubejs_task.run(subdomain="test", api_secret="foo", query="query")
 
-        assert "Cube.js load API failed!." in str(exc)
+        assert "Cube.js load API failed!" in str(exc)
 
     @responses.activate
     def test_run_with_continue_waiting(self, caplog):
@@ -110,7 +111,12 @@ def test_run_with_security_context(self, caplog):
             security_context={"foo": "bar"},
         )
 
-        assert "JWT token generated with security context." in caplog.text
+        expected_jwt = jwt.encode(
+            payload={"foo": "bar", "expiresIn": "7d"}, key="foo", algorithm="HS256"
+        )
+
+        # assert "JWT token generated with security context." in caplog.text
+        assert responses.calls[0].request.headers["Authorization"] == expected_jwt
 
     @responses.activate
     def test_run_with_max_wait_time_raises(self):
@@ -129,7 +135,11 @@ def test_run_with_max_wait_time_raises(self):
                 api_secret="foo",
                 query="query",
                 security_context={"foo": "bar"},
-                max_wait_time=15,
+                wait_time_between_api_calls=1,
+                max_wait_time=3,
             )
 
-        assert "Cube.js load API took too long to provide a response." in str(exc)
+        assert (
+            "Cube.js load API took longer than 3 seconds to provide a response."
+            in str(exc)
+        )