quota.py 4.87 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
# Copyright (c) 2017, Daniele Venzano
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
# implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Interface to PostgresQL for Zoe state."""

import logging

from zoe_lib.state.base import BaseTable, BaseRecord

log = logging.getLogger(__name__)


class Quota(BaseRecord):
    """A quota object describes limits imposed to users on resource usage."""

    def __init__(self, d, sql_manager):
        super().__init__(d, sql_manager)

        self.name = d['name']
        self.concurrent_executions = d['concurrent_executions']
        self.memory = d['memory']
        self.cores = d['cores']
35
        self.runtime_limit = d['runtime_limit']
36
37
38
39
40
41
42
43

    def serialize(self):
        """Generates a dictionary that can be serialized in JSON."""
        return {
            'id': self.id,
            'name': self.name,
            'concurrent_executions': self.concurrent_executions,
            'cores': self.cores,
44
45
            'memory': self.memory,
            'runtime_limit': self.runtime_limit
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
        }

    def set_concurrent_executions(self, value):
        """Setter for concurrent execution limit."""
        self.concurrent_executions = value
        self.sql_manager.quota_update(self.id, concurrent_executions=value)

    def set_memory(self, value):
        """Setter for memory limit."""
        self.memory = value
        self.sql_manager.quota_update(self.id, memory=value)

    def set_cores(self, value):
        """Setter for cores limit."""
        self.cores = value
        self.sql_manager.quota_update(self.id, cores=value)

63
64
65
66
67
    def set_runtime_limit(self, value):
        """Setter for the runtime limit."""
        self.runtime_limit = value
        self.sql_manager.quota_update(self.id, runtime_limit=value)

68
69
70
71
72
73
74
75
76
    def __repr__(self):
        return self.name

    def __eq__(self, other):
        if isinstance(other, Quota):
            return self.id == other.id
        else:
            return False

77
78
79
80
81
82
83
84
85
86
87
88
89

class QuotaTable(BaseTable):
    """Abstraction for the quota table in the database."""
    def __init__(self, sql_manager):
        super().__init__(sql_manager, "quota")

    def create(self):
        """Create the quota table."""
        self.cursor.execute('''CREATE TABLE quota (
            id SERIAL PRIMARY KEY,
            name TEXT NOT NULL,
            concurrent_executions INT NOT NULL,
            memory BIGINT NOT NULL,
90
91
            cores INT NOT NULL,
            runtime_limit INT NOT NULL
92
        )''')
93
        self.cursor.execute('''INSERT INTO quota (id, name, concurrent_executions, memory, cores, runtime_limit) VALUES (DEFAULT, 'default', 5, 34359738368, 20, 24)''')
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

    def select(self, only_one=False, **kwargs):
        """
        Return a list of quotas.

        :param only_one: only one result is expected
        :type only_one: bool
        :param kwargs: filter services based on their fields/columns
        :return: one or more ports
        """
        q_base = 'SELECT * FROM quota'
        if len(kwargs) > 0:
            q = q_base + " WHERE "
            filter_list = []
            args_list = []
            for key, value in kwargs.items():
                filter_list.append('{} = %s'.format(key))
                args_list.append(value)
            q += ' AND '.join(filter_list)
            query = self.cursor.mogrify(q, args_list)
        else:
            query = self.cursor.mogrify(q_base)

        self.cursor.execute(query)
        if only_one:
            row = self.cursor.fetchone()
            if row is None:
                return None
122
            return Quota(row, self.sql_manager)
123
        else:
124
            return [Quota(x, self.sql_manager) for x in self.cursor]
125

126
    def insert(self, name, concurrent_executions, memory, cores, runtime_limit):
127
        """Adds a new quota to the state."""
128
        query = self.cursor.mogrify('INSERT INTO quota (id, name, concurrent_executions, memory, cores, runtime_limit) VALUES (DEFAULT, %s, %s, %s, %s, %s) RETURNING id', (name, concurrent_executions, memory, cores, runtime_limit))
129
130
131
132
133
134
135
136
137
138
139
        self.cursor.execute(query)
        self.sql_manager.commit()
        return self.cursor.fetchone()[0]

    def delete(self, record_id):
        """Delete a quota from the state."""
        query = 'UPDATE "user" SET quota_id = (SELECT id from quota WHERE name=\'default\') WHERE quota_id=%s'
        self.cursor.execute(query, (record_id,))
        query = "DELETE FROM quota WHERE id = %s"
        self.cursor.execute(query, (record_id,))
        self.sql_manager.commit()