improve session handling around quotas

blueprint db-session-cleanup

Change-Id: I7606f69c0945b810deae098000717c7bcf96e77d

2 files changed, 47 insertions, 42 deletions

diff --git a/nova/db/sqlalchemy/api.py b/nova/db/sqlalchemy/api.py
index ba5a7150e..3b7c17a14 100644
--- a/nova/db/sqlalchemy/api.py
+++ b/nova/db/sqlalchemy/api.py
@@ -2323,9 +2323,8 @@ def iscsi_target_create_safe(context, values):
 
 
 @require_context
-def quota_get(context, project_id, resource, session=None):
-    result = model_query(context, models.Quota, session=session,
-                         read_deleted="no").\
+def quota_get(context, project_id, resource):
+    result = model_query(context, models.Quota, read_deleted="no").\
                      filter_by(project_id=project_id).\
                      filter_by(resource=resource).\
                      first()
@@ -2363,20 +2362,21 @@ def quota_create(context, project_id, resource, limit):
 
 @require_admin_context
 def quota_update(context, project_id, resource, limit):
-    session = get_session()
-    with session.begin():
-        quota_ref = quota_get(context, project_id, resource, session=session)
-        quota_ref.hard_limit = limit
-        quota_ref.save(session=session)
+    result = model_query(context, models.Quota, read_deleted="no").\
+                     filter_by(project_id=project_id).\
+                     filter_by(resource=resource).\
+                     update({'hard_limit': limit})
+
+    if not result:
+        raise exception.ProjectQuotaNotFound(project_id=project_id)
 
 
 ###################
 
 
 @require_context
-def quota_class_get(context, class_name, resource, session=None):
-    result = model_query(context, models.QuotaClass, session=session,
-                         read_deleted="no").\
+def quota_class_get(context, class_name, resource):
+    result = model_query(context, models.QuotaClass, read_deleted="no").\
                      filter_by(class_name=class_name).\
                      filter_by(resource=resource).\
                      first()
@@ -2414,21 +2414,21 @@ def quota_class_create(context, class_name, resource, limit):
 
 @require_admin_context
 def quota_class_update(context, class_name, resource, limit):
-    session = get_session()
-    with session.begin():
-        quota_class_ref = quota_class_get(context, class_name, resource,
-                                          session=session)
-        quota_class_ref.hard_limit = limit
-        quota_class_ref.save(session=session)
+    result = model_query(context, models.QuotaClass, read_deleted="no").\
+                     filter_by(class_name=class_name).\
+                     filter_by(resource=resource).\
+                     update({'hard_limit': limit})
+
+    if not result:
+        raise exception.QuotaClassNotFound(class_name=class_name)
 
 
 ###################
 
 
 @require_context
-def quota_usage_get(context, project_id, resource, session=None):
-    result = model_query(context, models.QuotaUsage, session=session,
-                         read_deleted="no").\
+def quota_usage_get(context, project_id, resource):
+    result = model_query(context, models.QuotaUsage, read_deleted="no").\
                      filter_by(project_id=project_id).\
                      filter_by(resource=resource).\
                      first()
@@ -2456,6 +2456,13 @@ def quota_usage_get_all_by_project(context, project_id):
 
 @require_admin_context
 def quota_usage_create(context, project_id, resource, in_use, reserved,
+                       until_refresh):
+    return _quota_usage_create(context, project_id, resource, in_use,
+                                reserved, until_refresh)
+
+
+@require_admin_context
+def _quota_usage_create(context, project_id, resource, in_use, reserved,
                        until_refresh, session=None):
     quota_usage_ref = models.QuotaUsage()
     quota_usage_ref.project_id = project_id
@@ -2463,31 +2470,29 @@ def quota_usage_create(context, project_id, resource, in_use, reserved,
     quota_usage_ref.in_use = in_use
     quota_usage_ref.reserved = reserved
     quota_usage_ref.until_refresh = until_refresh
+
     quota_usage_ref.save(session=session)
 
     return quota_usage_ref
 
 
 @require_admin_context
-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)
-        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:
-        # Assume caller started a transaction
-        do_update(session)
-    else:
-        session = get_session()
-        with session.begin():
-            do_update(session)
+def quota_usage_update(context, project_id, resource, **kwargs):
+    updates = {}
+    if 'in_use' in kwargs:
+        updates['in_use'] = kwargs['in_use']
+    if 'reserved' in kwargs:
+        updates['reserved'] = kwargs['reserved']
+    if 'until_refresh' in kwargs:
+        updates['until_refresh'] = kwargs['until_refresh']
+
+    result = model_query(context, models.QuotaUsage, read_deleted="no").\
+                     filter_by(project_id=project_id).\
+                     filter_by(resource=resource).\
+                     update(updates)
+
+    if not result:
+        raise exception.QuotaUsageNotFound(project_id=project_id)
 
 
 ###################
@@ -2564,7 +2569,7 @@ def quota_reserve(context, resources, quotas, deltas, expire,
             # Do we need to refresh the usage?
             refresh = False
             if resource not in usages:
-                usages[resource] = quota_usage_create(elevated,
+                usages[resource] = _quota_usage_create(elevated,
                                                       context.project_id,
                                                       resource,
                                                       0, 0,
@@ -2592,7 +2597,7 @@ def quota_reserve(context, resources, quotas, deltas, expire,
                 for res, in_use in updates.items():
                     # Make sure we have a destination for the usage!
                     if res not in usages:
-                        usages[res] = quota_usage_create(elevated,
+                        usages[res] = _quota_usage_create(elevated,
                                                          context.project_id,
                                                          res,
                                                          0, 0,
diff --git a/nova/tests/test_quota.py b/nova/tests/test_quota.py
index 26346276f..da064df1e 100644
--- a/nova/tests/test_quota.py
+++ b/nova/tests/test_quota.py
@@ -1418,7 +1418,7 @@ class QuotaReserveSqlAlchemyTestCase(test.TestCase):
 
         self.stubs.Set(sqa_api, 'get_session', fake_get_session)
         self.stubs.Set(sqa_api, '_get_quota_usages', fake_get_quota_usages)
-        self.stubs.Set(sqa_api, 'quota_usage_create', fake_quota_usage_create)
+        self.stubs.Set(sqa_api, '_quota_usage_create', fake_quota_usage_create)
         self.stubs.Set(sqa_api, 'reservation_create', fake_reservation_create)
 
         timeutils.set_time_override()