code refactoring

Author: qcdyx

api/src/shared/database/database.py

@@ -20,6 +20,7 @@
 )
 from sqlalchemy.orm import sessionmaker
 import logging
+from google.cloud import tasks_v2
 
 from shared.common.logging_utils import get_env_logging_level
 
@@ -124,9 +125,6 @@ def create_refresh_materialized_view_task():
     Returns:
         dict: Response message and status code.
     """
-    from google.auth.transport.requests import Request
-    from google.oauth2 import id_token
-    from google.cloud import tasks_v2
     from google.protobuf import timestamp_pb2
     from datetime import datetime, timedelta
 
@@ -159,33 +157,25 @@ def create_refresh_materialized_view_task():
             f"tasks-executor-{os.getenv('ENVIRONMENT_NAME')}"
         )
 
-        parent = client.queue_path(project, location, queue)
         task_name = client.task_path(project, location, queue, task_name)
 
         # Convert to protobuf timestamp
         proto_time = timestamp_pb2.Timestamp()
         proto_time.FromDatetime(bucket_time)
 
-        # Fetch an identity token for the target URL
-        auth_req = Request()
-        token = id_token.fetch_id_token(auth_req, url)
-
-        task = {
-            "name": task_name,
-            "http_request": {
-                "http_method": tasks_v2.HttpMethod.GET,
-                "url": url,
-                "headers": {
-                    "Content-Type": "application/json",
-                    "Authorization": f"Bearer {token}",
-                },
-            },
-            "schedule_time": proto_time,
-        }
-
         # Enqueue the task
         try:
-            client.create_task(request={"parent": parent, "task": task})
+            create_http_task_with_name(
+                client=client,
+                body=b"",
+                url=url,
+                project_id=project,
+                gcp_region=location,
+                queue_name=queue,
+                task_name=task_name,
+                task_time=proto_time,
+                http_method=tasks_v2.HttpMethod.GET,
+            )
             logging.info(f"Scheduled refresh materialized view task for {timestamp_str}")
             return {"message": f"Refresh task for {timestamp_str} scheduled."}, 200
         except Exception as e:
@@ -211,6 +201,45 @@ def refresh_materialized_view(session: "Session", view_name: str) -> bool:
         return False
 
 
+def create_http_task_with_name(
+    client,  # type: tasks_v2.CloudTasksClient
+    body: bytes,
+    url: str,
+    project_id: str,
+    gcp_region: str,
+    queue_name: str,
+    task_name: str,
+    task_time: str,
+    http_method: tasks_v2.HttpMethod,
+):
+    """Creates a GCP Cloud Task."""
+
+    token = (tasks_v2.OidcToken(service_account_email=os.getenv("SERVICE_ACCOUNT_EMAIL")),)
+
+    task = tasks_v2.Task(
+        http_request=tasks_v2.HttpRequest(
+            url=url,
+            http_method=http_method,
+            oidc_token=token,
+            body=body,
+            headers={"Content-Type": "application/json"},
+        )
+    )
+    task = {
+        "name": task_name,
+        "http_request": {
+            "http_method": http_method,
+            "url": url,
+            "headers": {
+                "Content-Type": "application/json",
+                "Authorization": f"Bearer {token}",
+            },
+        },
+        "schedule_time": task_time,
+    }
+    client.create_task(parent=client.queue_path(project_id, gcp_region, queue_name), task=task)
+
+
 def with_db_session(func=None, db_url: str | None = None):
     """
     Decorator to handle the session management for the decorated function.

functions-python/helpers/utils.py

@@ -151,20 +151,13 @@ def create_http_task(
     gcp_region: str,
     queue_name: str,
 ) -> None:
-    """Creates a GCP Cloud Task."""
-    from google.cloud import tasks_v2
+    from shared.database.database import create_http_task_with_name
 
-    task = tasks_v2.Task(
-        http_request=tasks_v2.HttpRequest(
-            url=url,
-            http_method=tasks_v2.HttpMethod.POST,
-            oidc_token=tasks_v2.OidcToken(
-                service_account_email=os.getenv("SERVICE_ACCOUNT_EMAIL")
-            ),
-            body=body,
-            headers={"Content-Type": "application/json"},
-        )
-    )
-    client.create_task(
-        parent=client.queue_path(project_id, gcp_region, queue_name), task=task
+    create_http_task_with_name(
+        client=client,
+        body=body,
+        url=url,
+        project_id=project_id,
+        gcp_region=gcp_region,
+        queue_name=queue_name,
     )