Merge "Add call to reset quota usage"

4 files changed, 97 insertions, 9 deletions

diff --git a/nova/db/api.py b/nova/db/api.py
index 82b7e895b..757f101b3 100644
--- a/nova/db/api.py
+++ b/nova/db/api.py
@@ -994,11 +994,9 @@ def quota_usage_get_all_by_project(context, project_id):
     return IMPL.quota_usage_get_all_by_project(context, project_id)
 
 
-def quota_usage_update(context, project_id, resource, in_use, reserved,
-                       until_refresh):
+def quota_usage_update(context, project_id, resource, **kwargs):
     """Update a quota usage or raise if it does not exist."""
-    return IMPL.quota_usage_update(context, project_id, resource,
-                                   in_use, reserved, until_refresh)
+    return IMPL.quota_usage_update(context, project_id, resource, **kwargs)
 
 
 def quota_usage_destroy(context, project_id, resource):
diff --git a/nova/db/sqlalchemy/api.py b/nova/db/sqlalchemy/api.py
index d68c39f6e..f5f7f3b15 100644
--- a/nova/db/sqlalchemy/api.py
+++ b/nova/db/sqlalchemy/api.py
@@ -2569,14 +2569,16 @@ def quota_usage_create(context, project_id, resource, in_use, reserved,
 
 
 @require_admin_context
-def quota_usage_update(context, project_id, resource, in_use, reserved,
-                       until_refresh, session=None):
+def quota_usage_update(context, project_id, resource, session=None, **kwargs):
     def do_update(session):
         quota_usage_ref = quota_usage_get(context, project_id, resource,
                                           session=session)
-        quota_usage_ref.in_use = in_use
-        quota_usage_ref.reserved = reserved
-        quota_usage_ref.until_refresh = until_refresh
+        if 'in_use' in kwargs:
+            quota_usage_ref.in_use = kwargs['in_use']
+        if 'reserved' in kwargs:
+            quota_usage_ref.reserved = kwargs['reserved']
+        if 'until_refresh' in kwargs:
+            quota_usage_ref.until_refresh = kwargs['until_refresh']
         quota_usage_ref.save(session=session)
 
     if session:
diff --git a/nova/quota.py b/nova/quota.py
index d3ba0aa02..fa0b813de 100644
--- a/nova/quota.py
+++ b/nova/quota.py
@@ -357,6 +357,35 @@ class DbQuotaDriver(object):
 
         db.reservation_rollback(context, reservations)
 
+    def usage_reset(self, context, resources):
+        """
+        Reset the usage records for a particular user on a list of
+        resources.  This will force that user's usage records to be
+        refreshed the next time a reservation is made.
+
+        Note: this does not affect the currently outstanding
+        reservations the user has; those reservations must be
+        committed or rolled back (or expired).
+
+        :param context: The request context, for access checks.
+        :param resources: A list of the resource names for which the
+                          usage must be reset.
+        """
+
+        # We need an elevated context for the calls to
+        # quota_usage_update()
+        elevated = context.elevated()
+
+        for resource in resources:
+            try:
+                # Reset the usage to -1, which will force it to be
+                # refreshed
+                db.quota_usage_update(elevated, context.project_id,
+                                      resource, in_use=-1)
+            except exception.QuotaUsageNotFound:
+                # That means it'll be refreshed anyway
+                pass
+
     def destroy_all_by_project(self, context, project_id):
         """
         Destroy all quotas, usages, and reservations associated with a
@@ -734,6 +763,23 @@ class QuotaEngine(object):
             LOG.exception(_("Failed to roll back reservations "
                             "%(reservations)s") % locals())
 
+    def usage_reset(self, context, resources):
+        """
+        Reset the usage records for a particular user on a list of
+        resources.  This will force that user's usage records to be
+        refreshed the next time a reservation is made.
+
+        Note: this does not affect the currently outstanding
+        reservations the user has; those reservations must be
+        committed or rolled back (or expired).
+
+        :param context: The request context, for access checks.
+        :param resources: A list of the resource names for which the
+                          usage must be reset.
+        """
+
+        self._driver.usage_reset(context, resources)
+
     def destroy_all_by_project(self, context, project_id):
         """
         Destroy all quotas, usages, and reservations associated with a
diff --git a/nova/tests/test_quota.py b/nova/tests/test_quota.py
index 5ec753efe..5cc5dedde 100644
--- a/nova/tests/test_quota.py
+++ b/nova/tests/test_quota.py
@@ -304,6 +304,9 @@ class FakeDriver(object):
     def rollback(self, context, reservations):
         self.called.append(('rollback', context, reservations))
 
+    def usage_reset(self, context, resources):
+        self.called.append(('usage_reset', context, resources))
+
     def destroy_all_by_project(self, context, project_id):
         self.called.append(('destroy_all_by_project', context, project_id))
 
@@ -664,6 +667,16 @@ class QuotaEngineTestCase(test.TestCase):
                 ('rollback', context, ['resv-01', 'resv-02', 'resv-03']),
                 ])
 
+    def test_usage_reset(self):
+        context = FakeContext(None, None)
+        driver = FakeDriver()
+        quota_obj = self._make_quota_obj(driver)
+        quota_obj.usage_reset(context, ['res1', 'res2', 'res3'])
+
+        self.assertEqual(driver.called, [
+                ('usage_reset', context, ['res1', 'res2', 'res3']),
+                ])
+
     def test_destroy_all_by_project(self):
         context = FakeContext(None, None)
         driver = FakeDriver()
@@ -1362,6 +1375,35 @@ class DbQuotaDriverTestCase(test.TestCase):
                 ])
         self.assertEqual(result, ['resv-1', 'resv-2', 'resv-3'])
 
+    def test_usage_reset(self):
+        calls = []
+
+        def fake_quota_usage_update(context, project_id, resource, **kwargs):
+            calls.append(('quota_usage_update', context, project_id,
+                          resource, kwargs))
+            if resource == 'nonexist':
+                raise exception.QuotaUsageNotFound()
+        self.stubs.Set(db, 'quota_usage_update', fake_quota_usage_update)
+
+        ctx = FakeContext('test_project', 'test_class')
+        resources = ['res1', 'res2', 'nonexist', 'res4']
+        self.driver.usage_reset(ctx, resources)
+
+        # Make sure we had some calls
+        self.assertEqual(len(calls), len(resources))
+
+        # Extract the elevated context that was used and do some
+        # sanity checks
+        elevated = calls[0][1]
+        self.assertEqual(elevated.project_id, ctx.project_id)
+        self.assertEqual(elevated.quota_class, ctx.quota_class)
+        self.assertEqual(elevated.is_admin, True)
+
+        # Now check that all the expected calls were made
+        exemplar = [('quota_usage_update', elevated, 'test_project',
+                     res, dict(in_use=-1)) for res in resources]
+        self.assertEqual(calls, exemplar)
+
 
 class FakeSession(object):
     def begin(self):