sql.py 8.69 KB
Newer Older
1
import mysql.connector
Daniele Venzano's avatar
Daniele Venzano committed
2
import mysql.connector.errors
3

Daniele Venzano's avatar
Daniele Venzano committed
4
from utils.config import get_database_config
5

Daniele Venzano's avatar
Daniele Venzano committed
6 7

class CAaaState:
8 9 10
    def __init__(self):
        self.cnx = None

Daniele Venzano's avatar
Daniele Venzano committed
11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27
    def _reconnect(self):
        if self.cnx is not None:
            self.cnx.disconnect()
        db_config = get_database_config()
        self.cnx = mysql.connector.connect(**db_config)

    def _get_cursor(self, dictionary=False):
        try:
            cursor = self.cnx.cursor(dictionary=dictionary)
        except (mysql.connector.errors.OperationalError, AttributeError):
            self._reconnect()
            cursor = self.cnx.cursor(dictionary=dictionary)
        return cursor

    def _close_cursor(self, cursor):
        self.cnx.commit()
        cursor.close()
28 29

    def _check_user(self, username):
Daniele Venzano's avatar
Daniele Venzano committed
30
        cursor = self._get_cursor(dictionary=True)
31 32 33 34
        q = "SELECT id FROM users WHERE username=%s"

        cursor.execute(q, (username,))
        if cursor.rowcount == 0:
Daniele Venzano's avatar
Daniele Venzano committed
35
            self._close_cursor(cursor)
36 37 38
            return self._create_user(username)
        else:
            row = cursor.fetchone()
Daniele Venzano's avatar
Daniele Venzano committed
39
            self._close_cursor(cursor)
40 41 42
            return row["id"]

    def _create_user(self, username):
Daniele Venzano's avatar
Daniele Venzano committed
43
        cursor = self._get_cursor()
44
        q = "INSERT INTO users (username) VALUES (%s)"
45
        cursor.execute(q, (username,))
46
        user_id = cursor.lastrowid
Daniele Venzano's avatar
Daniele Venzano committed
47
        self._close_cursor(cursor)
48 49
        return user_id

50 51 52 53
    def get_user_id(self, username):
        return self._check_user(username)

    def get_all_users(self):
Daniele Venzano's avatar
Daniele Venzano committed
54
        cursor = self._get_cursor()
55 56 57 58 59 60
        q = "SELECT id, username FROM users"

        user_list = []
        cursor.execute(q)
        for row in cursor:
            user_list.append(row)
Daniele Venzano's avatar
Daniele Venzano committed
61
        self._close_cursor(cursor)
62 63 64
        return user_list

    def count_clusters(self, user_id=None):
Daniele Venzano's avatar
Daniele Venzano committed
65
        cursor = self._get_cursor()
66 67 68 69 70 71 72
        if user_id is None:
            q = "SELECT COUNT(*) FROM clusters"
            cursor.execute(q)
        else:
            q = "SELECT COUNT(*) FROM clusters WHERE user_id=%s"
            cursor.execute(q, (user_id,))
        row = cursor.fetchone()
Daniele Venzano's avatar
Daniele Venzano committed
73
        self._close_cursor(cursor)
74 75
        return row[0]

76
    def count_containers(self, user_id=None, cluster_id=None):
Daniele Venzano's avatar
Daniele Venzano committed
77
        cursor = self._get_cursor()
78
        if user_id is None and cluster_id is None:
79 80
            q = "SELECT COUNT(*) FROM containers"
            cursor.execute(q)
81
        elif user_id is not None and cluster_id is None:
82 83
            q = "SELECT COUNT(*) FROM containers WHERE user_id=%s"
            cursor.execute(q, (user_id,))
84 85 86 87 88 89 90
        elif user_id is None and cluster_id is not None:
            q = "SELECT COUNT(*) FROM containers WHERE cluster_id=%s"
            cursor.execute(q, (cluster_id,))
        elif user_id is not None and cluster_id is not None:
            q = "SELECT COUNT(*) FROM containers WHERE user_id=%s AND cluster_id=%s"
            cursor.execute(q, (user_id, cluster_id))

91
        row = cursor.fetchone()
Daniele Venzano's avatar
Daniele Venzano committed
92
        self._close_cursor(cursor)
93 94 95
        return row[0]

    def get_notebook(self, user_id):
Daniele Venzano's avatar
Daniele Venzano committed
96
        cursor = self._get_cursor(dictionary=True)
97 98 99
        q = "SELECT * FROM notebooks WHERE user_id=%s"
        cursor.execute(q, (user_id,))
        if cursor.rowcount == 0:
Daniele Venzano's avatar
Daniele Venzano committed
100
            self.cnx.commit()
101 102 103 104
            cursor.close()
            return None
        else:
            row = cursor.fetchone()
Daniele Venzano's avatar
Daniele Venzano committed
105
            self._close_cursor(cursor)
106 107 108 109 110
            return row

    def has_notebook(self, user_id):
        ret = self.get_notebook(user_id)
        return ret is not None
111

112
    def get_url_proxy(self, proxy_id):
Daniele Venzano's avatar
Daniele Venzano committed
113
        cursor = self._get_cursor()
114 115 116
        q = "SELECT url FROM proxy WHERE proxy_id=%s"
        cursor.execute(q, (proxy_id,))
        if cursor.rowcount == 0:
Daniele Venzano's avatar
Daniele Venzano committed
117
            self._close_cursor(cursor)
118 119 120
            return None
        else:
            row = cursor.fetchone()
Daniele Venzano's avatar
Daniele Venzano committed
121
            self._close_cursor(cursor)
122 123 124
            return row[0]

    def get_all_proxy(self):
Daniele Venzano's avatar
Daniele Venzano committed
125
        cursor = self._get_cursor()
126 127 128 129 130
        q = "SELECT proxy_id, url, proxy_type, container_id FROM proxy"
        cursor.execute(q)
        proxy_list = []
        for proxy_id, url, proxy_type, container_id in cursor:
            proxy_list.append((proxy_id, url, proxy_type, container_id))
Daniele Venzano's avatar
Daniele Venzano committed
131
        self._close_cursor(cursor)
132 133 134
        return proxy_list

    def new_cluster(self, user_id, name):
Daniele Venzano's avatar
Daniele Venzano committed
135
        cursor = self._get_cursor()
136 137 138
        q = "INSERT INTO clusters (user_id, name) VALUES (%s, %s)"
        cursor.execute(q, (user_id, name))
        cluster_id = cursor.lastrowid
Daniele Venzano's avatar
Daniele Venzano committed
139
        self._close_cursor(cursor)
140 141 142
        return cluster_id

    def set_master_address(self, cluster_id, address):
Daniele Venzano's avatar
Daniele Venzano committed
143
        cursor = self._get_cursor()
144 145
        q = "UPDATE clusters SET master_address=%s WHERE clusters.id=%s"
        print(address, cluster_id)
Daniele Venzano's avatar
Daniele Venzano committed
146
        self._close_cursor(cursor)
147 148 149
        cursor.close()

    def new_container(self, cluster_id, user_id, docker_id, ip_address, contents):
Daniele Venzano's avatar
Daniele Venzano committed
150
        cursor = self._get_cursor()
Daniele Venzano's avatar
Daniele Venzano committed
151
        q = "INSERT INTO containers (user_id, cluster_id, docker_id, ip_address, contents) VALUES (%s, %s, %s, %s, %s)"
152 153
        cursor.execute(q, (user_id, cluster_id, docker_id, ip_address, contents))
        cont_id = cursor.lastrowid
Daniele Venzano's avatar
Daniele Venzano committed
154
        self._close_cursor(cursor)
155 156 157
        return cont_id

    def new_proxy_entry(self, proxy_id, cluster_id, address, proxy_type, container_id):
Daniele Venzano's avatar
Daniele Venzano committed
158
        cursor = self._get_cursor()
Daniele Venzano's avatar
Daniele Venzano committed
159
        q = "INSERT INTO proxy (proxy_id, url, cluster_id, proxy_type, container_id)  VALUES (%s, %s, %s, %s, %s)"
160
        cursor.execute(q, (proxy_id, address, cluster_id, proxy_type, container_id))
Daniele Venzano's avatar
Daniele Venzano committed
161
        self._close_cursor(cursor)
162 163 164
        return proxy_id

    def new_notebook(self, cluster_id, address, user_id, container_id):
Daniele Venzano's avatar
Daniele Venzano committed
165
        cursor = self._get_cursor()
Daniele Venzano's avatar
Daniele Venzano committed
166
        q = "INSERT INTO notebooks (cluster_id, address, user_id, container_id)  VALUES (%s, %s, %s, %s)"
167 168
        cursor.execute(q, (cluster_id, address, user_id, container_id))
        nb_id = cursor.lastrowid
Daniele Venzano's avatar
Daniele Venzano committed
169
        self._close_cursor(cursor)
170
        return nb_id
171 172

    def get_clusters(self, user_id=None):
Daniele Venzano's avatar
Daniele Venzano committed
173
        cursor = self._get_cursor(dictionary=True)
174 175 176 177 178 179 180 181 182 183 184 185 186
        res = {}
        if user_id is None:
            q = "SELECT id, user_id, master_address, name FROM clusters"
            cursor.execute(q)
        else:
            q = "SELECT id, user_id, master_address, name FROM clusters WHERE user_id=%s"
            cursor.execute(q, (user_id,))
        for row in cursor:
            res[str(row["id"])] = {
                "user_id": row["user_id"],
                "master_address": row["master_address"],
                "name": row["name"]
            }
Daniele Venzano's avatar
Daniele Venzano committed
187
        self._close_cursor(cursor)
188 189 190
        return res

    def get_containers(self, user_id=None, cluster_id=None):
Daniele Venzano's avatar
Daniele Venzano committed
191
        cursor = self._get_cursor(dictionary=True)
192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213
        res = {}
        if user_id is None and cluster_id is None:
            q = "SELECT id, docker_id, cluster_id, user_id, ip_address, contents FROM containers"
            cursor.execute(q)
        elif user_id is not None and cluster_id is None:
            q = "SELECT id, docker_id, cluster_id, user_id, ip_address, contents FROM containers WHERE user_id=%s"
            cursor.execute(q, (user_id,))
        elif user_id is None and cluster_id is not None:
            q = "SELECT id, docker_id, cluster_id, user_id, ip_address, contents FROM containers WHERE cluster_id=%s"
            cursor.execute(q, (cluster_id,))
        elif user_id is not None and cluster_id is not None:
            q = "SELECT id, docker_id, cluster_id, user_id, ip_address, contents FROM containers WHERE user_id=%s AND cluster_id=%s"
            cursor.execute(q, (user_id, cluster_id))

        for row in cursor:
            res[str(row["id"])] = {
                "docker_id": row["docker_id"],
                "cluster_id": row["cluster_id"],
                "user_id": row["user_id"],
                "ip_address": row["ip_address"],
                "contents": row["contents"],
            }
Daniele Venzano's avatar
Daniele Venzano committed
214
        self._close_cursor(cursor)
215 216 217
        return res

    def remove_proxy(self, container_id):
Daniele Venzano's avatar
Daniele Venzano committed
218
        cursor = self._get_cursor()
219 220
        q = "DELETE FROM proxy WHERE container_id=%s"
        cursor.execute(q, (container_id,))
Daniele Venzano's avatar
Daniele Venzano committed
221
        self._close_cursor(cursor)
222 223

    def remove_notebook(self, container_id):
Daniele Venzano's avatar
Daniele Venzano committed
224
        cursor = self._get_cursor()
225 226
        q = "DELETE FROM notebooks WHERE container_id=%s"
        cursor.execute(q, (container_id,))
Daniele Venzano's avatar
Daniele Venzano committed
227
        self._close_cursor(cursor)
228 229

    def remove_container(self, container_id):
Daniele Venzano's avatar
Daniele Venzano committed
230
        cursor = self._get_cursor()
231 232
        q = "DELETE FROM containers WHERE id=%s"
        cursor.execute(q, (container_id,))
Daniele Venzano's avatar
Daniele Venzano committed
233
        self._close_cursor(cursor)
234 235

    def remove_cluster(self, cluster_id):
Daniele Venzano's avatar
Daniele Venzano committed
236
        cursor = self._get_cursor()
237 238
        q = "DELETE FROM clusters WHERE id=%s"
        cursor.execute(q, (cluster_id,))
Daniele Venzano's avatar
Daniele Venzano committed
239
        self._close_cursor(cursor)