feat: add X-Cloudscheduler header support to TaskMetadata

django_cloud_tasks/tasks/task.py

@@ -36,6 +36,9 @@ class TaskMetadata:
     previous_failure: str | None = None
     project_id: str | None = None
     custom_headers: dict | None = None
+    is_cloud_scheduler: bool | None = None
+    cloud_scheduler_schedule_time: datetime | None = None
+    cloud_scheduler_job_name: str | None = None
 
     def __post_init__(self):
         self.custom_headers = get_current_headers()
@@ -44,6 +47,7 @@ def __post_init__(self):
     def from_headers(cls, headers: dict) -> Self:
         # Available data: https://cloud.google.com/tasks/docs/creating-http-target-tasks#handler
         cloud_tasks_prefix = "X-Cloudtasks-"
+        cloud_scheduler_prefix = "X-Cloudscheduler"
 
         if (attempt_str := headers.get(f"{cloud_tasks_prefix}Taskexecutioncount")) is not None:
             execution_number = int(attempt_str)
@@ -60,6 +64,18 @@ def from_headers(cls, headers: dict) -> Self:
         else:
             eta = None
 
+        cloud_scheduler_job_name = headers.get(f"{cloud_scheduler_prefix}-Jobname")
+
+        if schedule_time_str := headers.get(f"{cloud_scheduler_prefix}-Scheduletime"):
+            try:
+                schedule_time = datetime.fromisoformat(schedule_time_str)
+            except ValueError:
+                schedule_time = None
+        else:
+            schedule_time = None
+
+        is_cloud_scheduler = headers.get(cloud_scheduler_prefix) == "true"
+
         return cls(
             project_id=headers.get(f"{cloud_tasks_prefix}Projectname"),
             queue_name=headers.get(f"{cloud_tasks_prefix}Queuename"),
@@ -69,11 +85,14 @@ def from_headers(cls, headers: dict) -> Self:
             eta=eta,
             previous_response=headers.get(f"{cloud_tasks_prefix}TaskPreviousResponse"),
             previous_failure=headers.get(f"{cloud_tasks_prefix}TaskRetryReason"),
+            is_cloud_scheduler=is_cloud_scheduler,
+            cloud_scheduler_schedule_time=schedule_time,
+            cloud_scheduler_job_name=cloud_scheduler_job_name,
         )
 
     def to_headers(self) -> dict:
         cloud_tasks_prefix = "X-Cloudtasks-"
-        return {
+        cloud_tasks_headers = {
             f"{cloud_tasks_prefix}Taskname": self.task_id,
             f"{cloud_tasks_prefix}Queuename": self.queue_name,
             f"{cloud_tasks_prefix}Projectname": self.project_id,
@@ -84,6 +103,17 @@ def to_headers(self) -> dict:
             f"{cloud_tasks_prefix}TaskRetryReason": self.previous_failure,
         }
 
+        if self.is_cloud_scheduler:
+            cloud_scheduler_prefix = "X-Cloudscheduler"
+            cloud_scheduler_headers = {
+                f"{cloud_scheduler_prefix}-Jobname": self.cloud_scheduler_job_name,
+                f"{cloud_scheduler_prefix}-Scheduletime": self.cloud_scheduler_schedule_time.isoformat(),
+                f"{cloud_scheduler_prefix}": "true",
+            }
+            return cloud_tasks_headers | cloud_scheduler_headers
+
+        return cloud_tasks_headers
+
     @classmethod
     def from_task_obj(cls, task_obj: GoogleCloudTask) -> Self:
         _, project_id, _, _, _, queue_name, _, task_id = task_obj.name.split("/")  # TODO: use regex

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "django-google-cloud-tasks"
-version = "2.5.0"
+version = "2.6.0"
 description = "Async Tasks with HTTP endpoints"
 authors = ["Joao Daher <[email protected]>"]
 packages = [

sample_project/sample_app/tests/tests_tasks/tests_routine_executor_tasks.py

@@ -152,6 +152,21 @@ def sample_metadata(self) -> TaskMetadata:
             execution_number=7,
             dispatch_number=1,
             eta=self.some_date,
+            is_cloud_scheduler=False,
+        )
+
+    @property
+    def sample_cloud_scheduler_metadata(self) -> TaskMetadata:
+        return TaskMetadata(
+            project_id="wizard-project",
+            queue_name="wizard-queue",
+            task_id="hp-1234567",
+            execution_number=7,
+            dispatch_number=1,
+            eta=self.some_date,
+            is_cloud_scheduler=True,
+            cloud_scheduler_job_name="wizard-api--LevitationTask",
+            cloud_scheduler_schedule_time=datetime(2023, 11, 3, 15, 27, 0, tzinfo=UTC),
         )
 
     def test_create_from_headers(self):
@@ -164,6 +179,23 @@ def test_create_from_headers(self):
         self.assertEqual("wizard-project", metadata.project_id)
         self.assertEqual("wizard-queue", metadata.queue_name)
         self.assertEqual("hp-1234567", metadata.task_id)
+        self.assertFalse(metadata.is_cloud_scheduler)
+        self.assertIsNone(metadata.cloud_scheduler_job_name)
+        self.assertIsNone(metadata.cloud_scheduler_schedule_time)
+
+    def test_create_from_cloud_schedule_headers(self):
+        metadata = TaskMetadata.from_headers(
+            headers=self.sample_headers
+            | {
+                "X-Cloudscheduler": "true",
+                "X-Cloudscheduler-Scheduletime": "2023-11-03T08:27:00-07:00",
+                "X-Cloudscheduler-Jobname": "wizard-api--LevitationTask",
+            }
+        )
+
+        self.assertTrue(metadata.is_cloud_scheduler)
+        self.assertEqual("wizard-api--LevitationTask", metadata.cloud_scheduler_job_name)
+        self.assertEqual(datetime(2023, 11, 3, 15, 27, 0, tzinfo=UTC), metadata.cloud_scheduler_schedule_time)
 
     def test_build_headers(self):
         headers = self.sample_metadata.to_headers()
@@ -174,6 +206,22 @@ def test_build_headers(self):
         self.assertEqual("wizard-project", headers["X-Cloudtasks-Projectname"])
         self.assertEqual("wizard-queue", headers["X-Cloudtasks-Queuename"])
         self.assertEqual("hp-1234567", headers["X-Cloudtasks-Taskname"])
+        self.assertNotIn("X-Cloudscheduler", headers)
+        self.assertNotIn("X-Cloudscheduler-Scheduletime", headers)
+        self.assertNotIn("X-Cloudscheduler-Jobname", headers)
+
+    def test_build_cloud_scheduler_headers(self):
+        headers = self.sample_cloud_scheduler_metadata.to_headers()
+
+        self.assertEqual("7", headers["X-Cloudtasks-Taskexecutioncount"])
+        self.assertEqual("1", headers["X-Cloudtasks-Taskretrycount"])
+        self.assertEqual(str(int(self.some_date.timestamp())), headers["X-Cloudtasks-Tasketa"])
+        self.assertEqual("wizard-project", headers["X-Cloudtasks-Projectname"])
+        self.assertEqual("wizard-queue", headers["X-Cloudtasks-Queuename"])
+        self.assertEqual("hp-1234567", headers["X-Cloudtasks-Taskname"])
+        self.assertEqual("true", headers["X-Cloudscheduler"])
+        self.assertEqual("2023-11-03T15:27:00+00:00", headers["X-Cloudscheduler-Scheduletime"])
+        self.assertEqual("wizard-api--LevitationTask", headers["X-Cloudscheduler-Jobname"])
 
     def test_comparable(self):
         reference = self.sample_metadata

sample_project/sample_app/tests/tests_tasks/tests_tasks.py

@@ -279,6 +279,21 @@ def sample_metadata(self) -> TaskMetadata:
             execution_number=7,
             dispatch_number=1,
             eta=self.some_date,
+            is_cloud_scheduler=False,
+        )
+
+    @property
+    def sample_cloud_scheduler_metadata(self) -> TaskMetadata:
+        return TaskMetadata(
+            project_id="wizard-project",
+            queue_name="wizard-queue",
+            task_id="hp-1234567",
+            execution_number=7,
+            dispatch_number=1,
+            eta=self.some_date,
+            is_cloud_scheduler=True,
+            cloud_scheduler_job_name="wizard-api--LevitationTask",
+            cloud_scheduler_schedule_time=datetime(2023, 11, 3, 15, 27, 0, tzinfo=UTC),
         )
 
     def test_create_from_headers(self):
@@ -291,6 +306,23 @@ def test_create_from_headers(self):
         self.assertEqual("wizard-project", metadata.project_id)
         self.assertEqual("wizard-queue", metadata.queue_name)
         self.assertEqual("hp-1234567", metadata.task_id)
+        self.assertFalse(metadata.is_cloud_scheduler)
+        self.assertIsNone(metadata.cloud_scheduler_job_name)
+        self.assertIsNone(metadata.cloud_scheduler_schedule_time)
+
+    def test_create_from_cloud_schedule_headers(self):
+        metadata = TaskMetadata.from_headers(
+            headers=self.sample_headers
+            | {
+                "X-Cloudscheduler": "true",
+                "X-Cloudscheduler-Scheduletime": "2023-11-03T08:27:00-07:00",
+                "X-Cloudscheduler-Jobname": "wizard-api--LevitationTask",
+            }
+        )
+
+        self.assertTrue(metadata.is_cloud_scheduler)
+        self.assertEqual("wizard-api--LevitationTask", metadata.cloud_scheduler_job_name)
+        self.assertEqual(datetime(2023, 11, 3, 15, 27, 0, tzinfo=UTC), metadata.cloud_scheduler_schedule_time)
 
     def test_build_headers(self):
         headers = self.sample_metadata.to_headers()
@@ -301,6 +333,22 @@ def test_build_headers(self):
         self.assertEqual("wizard-project", headers["X-Cloudtasks-Projectname"])
         self.assertEqual("wizard-queue", headers["X-Cloudtasks-Queuename"])
         self.assertEqual("hp-1234567", headers["X-Cloudtasks-Taskname"])
+        self.assertNotIn("X-Cloudscheduler", headers)
+        self.assertNotIn("X-Cloudscheduler-Scheduletime", headers)
+        self.assertNotIn("X-Cloudscheduler-Jobname", headers)
+
+    def test_build_cloud_scheduler_headers(self):
+        headers = self.sample_cloud_scheduler_metadata.to_headers()
+
+        self.assertEqual("7", headers["X-Cloudtasks-Taskexecutioncount"])
+        self.assertEqual("1", headers["X-Cloudtasks-Taskretrycount"])
+        self.assertEqual(str(int(self.some_date.timestamp())), headers["X-Cloudtasks-Tasketa"])
+        self.assertEqual("wizard-project", headers["X-Cloudtasks-Projectname"])
+        self.assertEqual("wizard-queue", headers["X-Cloudtasks-Queuename"])
+        self.assertEqual("hp-1234567", headers["X-Cloudtasks-Taskname"])
+        self.assertEqual("true", headers["X-Cloudscheduler"])
+        self.assertEqual("2023-11-03T15:27:00+00:00", headers["X-Cloudscheduler-Scheduletime"])
+        self.assertEqual("wizard-api--LevitationTask", headers["X-Cloudscheduler-Jobname"])
 
     def test_comparable(self):
         reference = self.sample_metadata