From d685c902349ef111e017996f9dab719dac59241c Mon Sep 17 00:00:00 2001
From: Tony Aiuto <aiuto@google.com>
Date: Thu, 5 Jan 2012 16:58:38 -0500
Subject: [PATCH] Add a 'clear' command to the Google taskqueue sample. Fix the
 handling of per-command flags so that the help displays correctly.

---
 .../gtaskqueue_sample/gtaskqueue/gtaskqueue   |  0
 .../gtaskqueue_sample/gtaskqueue/task_cmds.py | 80 ++++++++++++++++---
 .../gtaskqueue/taskqueue_cmds.py              |  5 +-
 3 files changed, 74 insertions(+), 11 deletions(-)
 mode change 100644 => 100755 samples/gtaskqueue_sample/gtaskqueue/gtaskqueue

diff --git a/samples/gtaskqueue_sample/gtaskqueue/gtaskqueue b/samples/gtaskqueue_sample/gtaskqueue/gtaskqueue
old mode 100644
new mode 100755
diff --git a/samples/gtaskqueue_sample/gtaskqueue/task_cmds.py b/samples/gtaskqueue_sample/gtaskqueue/task_cmds.py
index 4adea2b..5bbc2ad 100644
--- a/samples/gtaskqueue_sample/gtaskqueue/task_cmds.py
+++ b/samples/gtaskqueue_sample/gtaskqueue/task_cmds.py
@@ -55,18 +55,21 @@ class LeaseTaskCommand(GoogleTaskCommand):
     """Lease a new task from the queue."""
 
     def __init__(self, name, flag_values):
+        flags.DEFINE_integer('lease_secs',
+                             None,
+                             'The lease for the task in seconds',
+                             flag_values=flag_values)
+        flags.DEFINE_integer('num_tasks',
+                             1,
+                             'The number of tasks to lease',
+                             flag_values=flag_values)
+        flags.DEFINE_integer('payload_size_to_display',
+                             2 * 1024 * 1024,
+                             'Size of the payload for leased tasks to show',
+                             flag_values=flag_values)
         super(LeaseTaskCommand, self).__init__(name,
                                                flag_values,
                                                need_task_flag=False)
-        flags.DEFINE_integer('lease_secs',
-                             None,
-                             'The lease for the task in seconds')
-        flags.DEFINE_integer('num_tasks',
-                             1,
-                             'The number of tasks to lease')
-        flags.DEFINE_integer('payload_size_to_display',
-                             2 * 1024 * 1024,
-                             'Size of the payload for leased tasks to show')
 
     def build_request(self, task_api, flag_values):
         """Build a request to lease a pending task from the TaskQueue.
@@ -143,8 +146,67 @@ class ListTasksCommand(GoogleTaskCommand):
                              taskqueue=flag_values.taskqueue_name)
 
 
+class ClearTaskQueueCommand(GoogleTaskCommand):
+    """Deletes all tasks in a queue (default to a max of 100)."""
+
+
+    def __init__(self, name, flag_values):
+        flags.DEFINE_integer('max_delete', 100, 'How many to clear at most',
+                             flag_values=flag_values)
+        super(ClearTaskQueueCommand, self).__init__(name,
+                                                    flag_values,
+                                                    need_task_flag=False)
+
+    def run_with_api_and_flags(self, api, flag_values):
+        """Run the command, returning the result.
+
+        Args:
+            api: The handle to the Google TaskQueue API.
+            flag_values: The parsed command flags.
+        Returns:
+            The result of running the command.
+        """
+        tasks_api = api.tasks()
+        self._flag_values = flag_values
+        self._to_delete = flag_values.max_delete
+        total_deleted = 0
+        while self._to_delete > 0:
+          n_deleted = self._delete_a_batch(tasks_api)
+          if n_deleted <= 0:
+            break
+          total_deleted += n_deleted
+        return {'deleted': total_deleted}
+
+    def _delete_a_batch(self, tasks):
+        """Delete a batch of tasks.
+
+        Since the list method only gives us back 100 at a time, we may have
+        to call it several times to clear the entire queue.
+
+        Args:
+            tasks: The handle to the Google TaskQueue API Tasks resource.
+        Returns:
+            The number of tasks deleted.
+        """
+        list_request = tasks.list(project=self._flag_values.project_name,
+                                  taskqueue=self._flag_values.taskqueue_name)
+        result = list_request.execute()
+        n_deleted = 0
+        if result:
+          for task in result.get('items', []):
+            if self._to_delete > 0:
+              self._to_delete -= 1
+              n_deleted += 1
+              print 'Deleting: %s' % task['id']
+              tasks.delete(project=self._flag_values.project_name,
+                           taskqueue=self._flag_values.taskqueue_name,
+                           task=task['id']).execute()
+        return n_deleted
+
+
 def add_commands():
     appcommands.AddCmd('listtasks', ListTasksCommand)
     appcommands.AddCmd('gettask', GetTaskCommand)
     appcommands.AddCmd('deletetask', DeleteTaskCommand)
     appcommands.AddCmd('leasetask', LeaseTaskCommand)
+    appcommands.AddCmd('clear', ClearTaskQueueCommand)
diff --git a/samples/gtaskqueue_sample/gtaskqueue/taskqueue_cmds.py b/samples/gtaskqueue_sample/gtaskqueue/taskqueue_cmds.py
index 99c3f2a..f21840a 100644
--- a/samples/gtaskqueue_sample/gtaskqueue/taskqueue_cmds.py
+++ b/samples/gtaskqueue_sample/gtaskqueue/taskqueue_cmds.py
@@ -33,10 +33,11 @@ class GetTaskQueueCommand(GoogleTaskQueueCommand):
     """Get properties of an existing task queue."""
 
     def __init__(self, name, flag_values):
-        super(GetTaskQueueCommand, self).__init__(name, flag_values)
         flags.DEFINE_boolean('get_stats',
                              False,
-                             'Whether to get Stats')
+                             'Whether to get Stats',
+                             flag_values=flag_values)
+        super(GetTaskQueueCommand, self).__init__(name, flag_values)
 
     def build_request(self, taskqueue_api, flag_values):
         """Build a request to get properties of a TaskQueue.