Support quota for keypairs in Mogan(part-three)

Introduce the quota management in keypair creation and deletion.

Change-Id: I8241b0bfb38a9c41b564df07a7799e4bc5438c0d
Implements: blueprint quota-support

mogan/conf/quota.py

@@ -32,6 +32,12 @@ quota_opts = [
     cfg.IntOpt('max_age',
                default=0,
                help=_('Number of seconds between subsequent usage refreshes')),
+    cfg.IntOpt('servers_hard_limit',
+               default=10,
+               help=_('Number of servers quota hard limit.')),
+    cfg.IntOpt('keypairs_hard_limit',
+               default=100,
+               help=_('Number of keypairs quota hard limit.')),
 ]
 
 opt_quota_group = cfg.OptGroup(name='quota',

mogan/db/sqlalchemy/alembic/versions/91941bf1ebc9_initial_migration.py

@@ -180,6 +180,7 @@ def upgrade():
         sa.Column('id', sa.Integer(), primary_key=True, nullable=False),
         sa.Column('name', sa.String(length=255), nullable=False),
         sa.Column('user_id', sa.String(length=255), nullable=True),
+        sa.Column('project_id', sa.String(length=36), nullable=True),
         sa.Column('fingerprint', sa.String(255)),
         sa.Column('public_key', sa.Text()),
         sa.Column('type', sa.Enum('ssh', 'x509'), nullable=False,

mogan/db/sqlalchemy/api.py

@@ -104,7 +104,8 @@ class Connection(api.Connection):
     """SqlAlchemy connection."""
 
     def __init__(self):
-        self.QUOTA_SYNC_FUNCTIONS = {'_sync_servers': self._sync_servers}
+        self.QUOTA_SYNC_FUNCTIONS = {'_sync_servers': self._sync_servers,
+                                     '_sync_keypairs': self._sync_keypairs}
         pass
 
     def _add_servers_filters(self, context, query, filters):
@@ -515,6 +516,11 @@ class Connection(api.Connection):
             filter_by(project_id=project_id).all()
         return {'servers': len(query) or 0}
 
+    def _sync_keypairs(self, context, project_id):
+        query = model_query(context, models.KeyPair).\
+            filter_by(project_id=project_id).all()
+        return {'keypairs': len(query) or 0}
+
     @oslo_db_api.retry_on_deadlock
     def quota_reserve(self, context, resources, quotas, deltas, expire,
                       until_refresh, max_age, project_id,

mogan/db/sqlalchemy/models.py

@@ -260,6 +260,8 @@ class KeyPair(Base):
 
     user_id = Column(String(255), nullable=False)
 
+    project_id = Column(String(255), nullable=False)
+
     fingerprint = Column(String(255))
     public_key = Column(Text())
     type = Column(Enum('ssh', 'x509', name='keypair_types'),

mogan/engine/api.py

@@ -70,6 +70,7 @@ class API(object):
         self.network_api = network.API()
         self.quota = quota.Quota()
         self.quota.register_resource(objects.quota.ServerResource())
+        self.quota.register_resource(objects.quota.KeyPairResource())
         self.consoleauth_rpcapi = consoleauth_rpcapi.ConsoleAuthAPI()
 
     def _get_image(self, context, image_uuid):
@@ -182,6 +183,21 @@ class API(object):
         else:
             raise exception.OverQuota(overs='servers')
 
+    def _check_num_keypairs_quota(self, context, count):
+        keypair_resource = self.quota.resources['keypairs']
+        quotas = self.quota.get_quota_limit_and_usage(context,
+                                                      {'keyparis':
+                                                       keypair_resource},
+                                                      context.tenant)
+        limit = quotas['keypairs']['limit']
+        in_use = quotas['keypairs']['in_use']
+        reserved = quotas['keypairs']['reserved']
+        available_quota = limit - in_use - reserved
+        if count <= available_quota:
+            return count
+        else:
+            raise exception.OverQuota(overs='keypairs')
+
     def _decode_files(self, injected_files):
         """Base64 decode the list of files to inject."""
         if not injected_files:
@@ -470,13 +486,21 @@ class API(object):
         self._validate_new_key_pair(context, user_id, key_name, key_type)
         private_key, public_key, fingerprint = self._generate_key_pair(
             user_id, key_type)
+        # Create the keypair reservations
+        num_keypairs = self._check_num_keypairs_quota(context, 1)
+        reserve_opts = {'keypairs': num_keypairs}
+        reservations = self.quota.reserve(context, **reserve_opts)
         keypair = objects.KeyPair(context)
         keypair.user_id = user_id
         keypair.name = key_name
         keypair.type = key_type
         keypair.fingerprint = fingerprint
         keypair.public_key = public_key
+        keypair.project_id = context.tenant
         keypair.create()
+        # Commit keypairs reservations
+        if reservations:
+            self.quota.commit(context, reservations)
         return keypair, private_key
 
     def _generate_fingerprint(self, public_key, key_type):
@@ -494,6 +518,10 @@ class API(object):
     def delete_key_pair(self, context, user_id, key_name):
         """Delete a keypair by name."""
         objects.KeyPair.destroy_by_name(context, user_id, key_name)
+        reserve_opts = {'keypairs': -1}
+        reservations = self.quota.reserve(context, **reserve_opts)
+        if reservations:
+            self.quota.commit(context, reservations)
 
     def get_key_pairs(self, context, user_id):
         """List key pairs."""

mogan/objects/keypair.py

@@ -36,6 +36,7 @@ class KeyPair(base.MoganObject):
         'id': fields.IntegerField(),
         'name': fields.StringField(nullable=False),
         'user_id': fields.StringField(nullable=True),
+        'project_id': fields.StringField(nullable=True),
         'fingerprint': fields.StringField(nullable=True),
         'public_key': fields.StringField(nullable=True),
         'type': fields.StringField(nullable=False),

mogan/objects/quota.py

@@ -185,11 +185,18 @@ class DbQuotaDriver(object):
         for p_quota in res:
             project_quotas[p_quota.resource_name] = p_quota.hard_limit
         if project_quotas == {}:
+            servers_hard_limit = CONF.quota.servers_hard_limit
             self.dbapi.quota_create(context, {'resource_name': 'servers',
                                               'project_id': project_id,
-                                              'hard_limit': 10,
+                                              'hard_limit': servers_hard_limit,
                                               'allocated': 0})
-            project_quotas['servers'] = 10
+            project_quotas['servers'] = servers_hard_limit
+            kpairs_hard_limit = CONF.quota.keypairs_hard_limit
+            self.dbapi.quota_create(context, {'resource_name': 'keypairs',
+                                              'project_id': project_id,
+                                              'hard_limit': kpairs_hard_limit,
+                                              'allocated': 0})
+            project_quotas['keypairs'] = kpairs_hard_limit
         allocated_quotas = None
         if usages:
             project_usages = self.dbapi.quota_usage_get_all_by_project(
@@ -410,9 +417,22 @@ class BaseResource(object):
 class ServerResource(BaseResource):
     """ReservableResource for a specific server."""
 
-    def __init__(self, name='servers'):
+    def __init__(self):
         """Initializes a ServerResource.
 
-        :param name: The kind of resource, i.e., "servers".
         """
-        super(ServerResource, self).__init__(name, "_sync_%s" % name)
+        self.name = 'servers'
+        super(ServerResource, self).__init__(self.name,
+                                             "_sync_%s" % self.name)
+
+
+class KeyPairResource(BaseResource):
+    """ReservableResource for a specific keypair."""
+
+    def __init__(self):
+        """Initializes a KeyPairResource.
+
+        """
+        self.name = 'keypairs'
+        super(KeyPairResource, self).__init__(self.name,
+                                              "_sync_%s" % self.name)

mogan/tests/unit/engine/test_engine_api.py

@@ -365,3 +365,33 @@ class ComputeAPIUnitTest(base.DbTestCase):
         self.engine_api.detach_interface(self.context, fake_server_obj,
                                          fake_server_obj['nics'][0]['port_id'])
         self.assertTrue(mock_detach_interface.called)
+
+    def test_create_key_pairs_with_quota(self):
+        res = self.dbapi._get_quota_usages(self.context, self.project_id)
+        before_in_use = 0
+        if res.get('keypairs') is not None:
+            before_in_use = res.get('keypairs').in_use
+        self.engine_api.create_key_pair(self.context, self.user_id,
+                                        'test_keypair')
+        res = self.dbapi._get_quota_usages(self.context, self.project_id)
+        after_in_use = res.get('keypairs').in_use
+        self.assertEqual(before_in_use + 1, after_in_use)
+
+    def test_create_key_pairs_with_over_quota_limit(self):
+        res = self.dbapi._get_quota_usages(self.context, self.project_id)
+        before_in_use = 0
+        if res.get('keypairs') is not None:
+            before_in_use = res.get('keypairs').in_use
+        for i in range(1, 101):
+            key_name = 'test_keypair_%s' % str(i)
+            self.engine_api.create_key_pair(self.context, self.user_id,
+                                            key_name)
+        res = self.dbapi._get_quota_usages(self.context, self.project_id)
+        after_in_use = res.get('keypairs').in_use
+        self.assertEqual(before_in_use + 100, after_in_use)
+        self.assertRaises(
+            exception.OverQuota,
+            self.engine_api.create_key_pair,
+            self.context,
+            self.user_id,
+            'test_keypair')

mogan/tests/unit/objects/test_keypairs.py

@@ -28,6 +28,7 @@ class TestKeyPairObject(base.DbTestCase):
             "id": 1,
             "public_key": "fake-publick-key",
             "user_id": "e78b60069fc9467e97fb4b74de9cadc1",
+            "project_id": "c18e8a1a870d4c08a0b51ced6e0b6459",
             "name": "test_key",
             "fingerprint": "f1:83:34:02:f9:63:79:d4:bd:2a:1d:50:16:61:1b:cc",
             "type": "ssh",

mogan/tests/unit/objects/test_objects.py

@@ -390,7 +390,7 @@ expected_object_fingerprints = {
     'ServerNic': '1.0-0494306157ef437802260ff8b51cf5cf',
     'ServerNics': '1.0-33a2e1bb91ad4082f9f63429b77c1244',
     'Quota': '1.0-c8caa082f4d726cb63fdc5943f7cd186',
-    'KeyPair': '1.0-c6820166e307676c5900f7801831b84c',
+    'KeyPair': '1.0-1a1ea1f9b4d03503f5c13b52d1432fa9',
     'KeyPairList': '1.0-33a2e1bb91ad4082f9f63429b77c1244',
     'Aggregate': '1.0-e2f3060705a18e1de8daae3f1af6b104',
     'AggregateList': '1.0-33a2e1bb91ad4082f9f63429b77c1244'