laas-fog/source/database.py

   1 """
   2 #############################################################################
   3 #Copyright 2017 Parker Berberian and others                                 #
   4 #                                                                           #
   5 #Licensed under the Apache License, Version 2.0 (the "License");            #
   6 #you may not use this file except in compliance with the License.           #
   7 #You may obtain a copy of the License at                                    #
   8 #                                                                           #
   9 #    http://www.apache.org/licenses/LICENSE-2.0                             #
  10 #                                                                           #
  11 #Unless required by applicable law or agreed to in writing, software        #
  12 #distributed under the License is distributed on an "AS IS" BASIS,          #
  13 #WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.   #
  14 #See the License for the specific language governing permissions and        #
  15 #limitations under the License.                                             #
  16 #############################################################################
  17 """
  18 import sqlite3
  19 import sys
  20 import time
  21
  22
  23 class HostDataBase:
  24     """
  25     This class talks with a simple sqlite database and can select a free host
  26     when one is needed.
  27     The layout of the database is:
  28         TABLE host:
  29             name <hostname> status <status_code> book_start
  30                 <Unix timestamp> book_end <Unix timestamp>
  31         status_codes:
  32         0 - idle
  33         1 - deploying
  34         2 - deployed, in use
  35         3 - expired, ready to be reset
  36     """
  37
  38     def __init__(self, path):
  39         """
  40         init function. Will create the file at the end of path
  41         if it doesnt already exist
  42         """
  43         self.database = sqlite3.connect(path)
  44         self.cursor = self.database.cursor()
  45
  46     def resetHosts(self, hosts):
  47         """
  48         Recreates the host table in the database.
  49         WILL ERASE ALL DATA. USE WITH CAUTION.
  50         """
  51         try:
  52             self.cursor.execute("DROP TABLE hosts")
  53             self.createTable()
  54         except:
  55             pass
  56
  57         for host in hosts:
  58             self.addHost(host)
  59
  60     def createTable(self):
  61         """
  62         This method creates the table hosts with
  63         a name and status field
  64         """
  65         self.cursor.execute("CREATE TABLE hosts (name text, status integer)")
  66         self.database.commit()
  67
  68     def addHost(self, name):
  69         """
  70         Adds a host with name to the available hosts.
  71         When first added, the host is assumed to be idle.
  72         """
  73         host = (name, )
  74         self.cursor.execute("INSERT INTO hosts VALUES (?, 0) ", host)
  75         self.database.commit()
  76
  77     def getHost(self, requested=None):
  78         """
  79         Returns the name of an available host.
  80         If a host is specifically requested,
  81         that host is returned.
  82         If the requested host is not available,
  83         this method will throw an error.
  84         If no host is specificaly requested,
  85         the next available host is returned.
  86         """
  87         self.cursor.execute("SELECT name FROM hosts WHERE status = 0")
  88         hostList = self.cursor.fetchall()
  89         if len(hostList) < 1:
  90             # throw and exception
  91             sys.exit(1)
  92         host = None
  93         if requested is not None:
  94             if (requested, ) in hostList and self.hostIsIdle(requested):
  95                 host = requested  # If requested, exists, and idle, return it
  96             else:
  97                 sys.exit(1)
  98         else:
  99             host = hostList[0][0]
 100         self.makeHostBusy(host)
 101         return host
 102
 103     def makeHostBusy(self, name):
 104         """
 105         makes the status of host 'name' equal 1,
 106         making it 'busy'
 107         """
 108         host = (name, )
 109         self.cursor.execute("UPDATE hosts SET status = 1 WHERE name=?", host)
 110         self.database.commit()
 111
 112     def makeHostDeployed(self, name):
 113         """
 114         makes the status of host 'name' equal 2,
 115         making it 'deployed' and/or in use
 116         """
 117         host = (name, )
 118         self.cursor.execute("UPDATE hosts SET status = 2 WHERE name=?", host)
 119         self.database.commit()
 120
 121     def makeHostExpired(self, name):
 122         """
 123         makes the status of host 'name' equal 3,
 124         meaning its booking has ended and needs to be cleaned.
 125         """
 126         host = (name, )
 127         self.cursor.execute("UPDATE hosts SET status = 3 WHERE name=?", host)
 128         self.database.commit()
 129
 130     def getExpiredHosts(self):
 131         """
 132         returns a list of all hosts with an expired booking that
 133         need to be cleaned.
 134         """
 135         self.cursor.execute("SELECT name FROM hosts where status = 3")
 136         host_tuples = self.cursor.fetchall()
 137         hosts = []
 138         for host in host_tuples:
 139             hosts.append(host[0])
 140         return hosts  # returns list of strings, not tuples
 141
 142     def hostIsBusy(self, name):
 143         """
 144         returns True if the host is not idle
 145         """
 146         host = (name, )
 147         self.cursor.execute("SELECT status FROM hosts WHERE name=?", host)
 148         stat = self.cursor.fetchone()[0]
 149         if stat < 1:
 150             return False
 151         return True
 152
 153     def hostIsIdle(self, name):
 154         """
 155         returns True if the host is idle.
 156         """
 157         return not self.hostIsBusy(name)
 158
 159     def getAllHosts(self):
 160         """
 161         returns the whole host database.
 162         """
 163         self.cursor.execute("SELECT * FROM hosts")
 164         return self.cursor.fetchall()
 165
 166     def close(self):
 167         """
 168         commits and closes connection to the database file.
 169         """
 170         self.database.commit()
 171         self.database.close()
 172
 173
 174 class BookingDataBase:
 175     """
 176     Database to hold all active bookings for our servers.
 177     Database contains table bookings - can be same or different
 178     db file as the host database
 179     bookings contains a field for every json key from the pharos dashboard,
 180     plus a "status" integer which is either
 181     0   -   waiting to start
 182     1   -   started
 183     2   -   booking over
 184
 185     As written, the pharos listener will immediately store all bookings that
 186     are both for your dev pods and not
 187     yet over, regardless of when the booking starts. Once the booking ends
 188     and the dev pod is cleaned, the booking is deleted to save space and cpu.
 189     """
 190
 191     def __init__(self, path):
 192         """
 193         creates a BookingDataBase object with the database located
 194         at path. if path does not yet exist, it will be created.
 195         """
 196         self.database = sqlite3.connect(path)
 197         self.cursor = self.database.cursor()
 198
 199     def createTable(self):
 200         """
 201         Creates table in the database to store booking information
 202         """
 203         try:
 204             self.cursor.execute("DROP TABLE bookings")
 205         except:
 206             pass
 207         self.cursor.execute("""CREATE TABLE bookings
 208         (id integer, resource_id integer, start double, end double,
 209             installer_name text, scenario_name text,
 210             purpose text, status integer, vpn text)""")
 211         self.database.commit()
 212
 213     def checkAddBooking(self, booking):
 214         """
 215         This method accepts a JSON booking definition from the dashboard
 216         api and adds it to the database if it does not already exist.
 217         """
 218         # first, check if booking is already expired
 219         if time.time() > booking['end']:
 220             return
 221         # check if booking is in database already
 222         b_id = (booking['id'], )
 223         self.cursor.execute("SELECT * FROM bookings WHERE id=?", b_id)
 224         if len(self.cursor.fetchall()) > 0:  # booking already in the db
 225             return
 226         tup = (
 227                 booking['id'],
 228                 booking['resource_id'],
 229                 booking['start'],
 230                 booking['end'],
 231                 booking['installer_name'],
 232                 booking['scenario_name'],
 233                 booking['purpose'],
 234                 0,
 235                 ''
 236                 )
 237         self.cursor.execute(
 238                 "INSERT INTO bookings VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)", tup)
 239         self.database.commit()
 240
 241     def removeBooking(self, idNum):
 242         """
 243         deletes booking with given id from the database.
 244         """
 245         booking_id = (idNum, )
 246         self.cursor.execute("DELETE FROM bookings WHERE id=?", booking_id)
 247
 248     def getBookings(self):
 249         """
 250         returns a list of all bookings.
 251         """
 252         self.cursor.execute("SELECT * FROM bookings")
 253         return self.cursor.fetchall()
 254
 255     def setStatus(self, booking_id, status):
 256         """
 257         sets the status of the booking with booking id booking_id.
 258         as noted above, the status codes are:
 259         0 - not yet started
 260         1 - started, but not yet over
 261         2 - over, expired
 262         """
 263         data = (status, booking_id)
 264         self.cursor.execute("UPDATE bookings SET status=? WHERE id=?", data)
 265         self.database.commit()
 266
 267     def setVPN(self, resource, uid):
 268         data = (uid, resource, 1)
 269         self.cursor.execute(
 270                 "UPDATE bookings SET vpn=? WHERE resource_id=? AND status=?",
 271                 data
 272             )
 273         self.database.commit()
 274
 275     def getVPN(self):
 276         """
 277         returns a list of all vpn users associated with current
 278         bookings.
 279         """
 280         self.cursor.execute("SELECT vpn FROM bookings WHERE status=1")
 281         users_messy = self.cursor.fetchall()
 282         users = []
 283         for user in users_messy:
 284             user = user[0]  # get string rather than tuple
 285             user = user.strip()
 286             if len(user) < 1:
 287                 continue
 288             users.append(user)  # a list of non-empty strings
 289         return users
 290
 291     def close(self):
 292         """
 293         commits changes and closes connection to db file.
 294         """
 295         self.database.commit()
 296         self.database.close()