db/db.py

#! /usr/bin/python

#  Copyright 2014 Linaro Limited
#  Author: Dave Pigott <dave.pigott@linaro.org>
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.

import psycopg2
import psycopg2.extras
import datetime
from errors import CriticalError, InputError

class VlanDB:
    def __init__(self, db_name="vland", username="vland"):
        try:
            self.connection = psycopg2.connect(database=db_name, user=username)
            self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.DictCursor)
        except Exception as e:
            print "Failed to access database: %s" % e

    def __del__(self):
        self.cursor.close()
        self.connection.close()

    # Create a new switch in the database. Switches are really simple
    # devices - they're just containers for ports.
    #
    # Constraints:
    # Switches must be uniquely named
    def create_switch(self, name):

        switch_id = self.get_switch_id(name)
        if switch_id is not None:
            raise InputError("Switch name %s already exists" % name)

        try:
            sql = "INSERT INTO switch (name) VALUES (%s) RETURNING switch_id"
            data = (name, )
            self.cursor.execute(sql, data)
            switch_id = self.cursor.fetchone()[0]
            self.connection.commit()
        except:
            self.connection.rollback()
            raise

        return switch_id

    # Create a new port in the database. Two of the fields are created
    # with default values (is_locked, is_trunk) here, and should be
    # updated separately if desired. For the current_vlan_id and
    # base_vlan_id fields, *BE CAREFUL* that you have already looked
    # up the correct VLAN_ID for each. This is *NOT* the same as the
    # VLAN tag (likely to be 1).
    # You Have Been Warned!
    #
    # Constraints:
    # 1. The switch referred to must already exist
    # 2. The VLANs mentioned here must already exist
    def create_port(self, name, switch_id, current_vlan_id, base_vlan_id):

        switch = self.get_switch(switch_id)
        if switch is None:
            raise InputError("Switch id %s does not exist" % switch_id)

        for vlan_id in (current_vlan_id, base_vlan_id):
            vlan = self.get_vlan(vlan_id)
            if vlan is None:
                raise InputError("VLAN id %s does not exist" % vlan_id)
                
        try:
            sql = "INSERT INTO port (name, switch_id, is_locked, is_trunk, current_vlan_id, base_vlan_id) VALUES (%s, %s, %s, %s, %s, %s) RETURNING port_id"
            data = (name, switch_id,
                    False, False,
                    current_vlan_id, base_vlan_id)
            self.cursor.execute(sql, data)
            port_id = self.cursor.fetchone()[0]
            self.connection.commit()
        except:
            self.connection.rollback()
            raise

        return port_id

    # Create a new vlan in the database. We locally add a creation
    # timestamp, for debug purposes. If vlans seems to be sticking
    # around, we'll be able to see when they were created.
    #
    # Constraints:
    # Names and tags must be unique
    def create_vlan(self, name, tag, is_base_vlan):

        vlan_id = self.get_vlan_id_from_name(name)
        if vlan_id is not None:
            raise InputError("VLAN name %s is already in use" % name)

        vlan_id = self.get_vlan_id_from_tag(tag)
        if vlan_id is not None:
            raise InputError("VLAN tag %d is already in use" % int(tag))

        for vlan_id in (current_vlan_id, base_vlan_id):
            vlan = self.get_vlan(vlan_id)
            if vlan is None:
                raise InputError("VLAN id %s does not exist" % vlan_id)        

        try:
            dt = datetime.datetime.now()
            sql = "INSERT INTO vlan (name, tag, is_base_vlan, creation_time) VALUES (%s, %s, %s, %s) RETURNING vlan_id"
            data = (name, tag, is_base_vlan, dt)
            self.cursor.execute(sql, data)
            vlan_id = self.cursor.fetchone()[0]
            self.connection.commit()
        except:
            self.connection.rollback()
            raise

        return vlan_id

    def _delete_row(self, table, field, value):
        try:
            sql = "DELETE FROM %s WHERE %s = %s"
            data = (table, field, value)
            self.cursor.execute(sql, data)
            self.connection.commit()
        except:
            self.connection.rollback()
            raise

    def delete_switch(self, switch_id):
        self._delete_row("switch", "switch_id", switch_id)

    def delete_port(self, port_id):
        self._delete_row("port", "port_id", port_id)

    def delete_vlan(self, vlan_id):
        self._delete_row("vlan", "vlan_id", vlan_id)

    def _get_element(self, select_field, table, compare_field, value):

        # We really want to use psycopg's type handling deal with the
        # (potentially) user-supplied data in the value field, so we
        # have to pass (sql,data) through to cursor.execute. However,
        # we can't have psycopg do all the argument substitution here
        # as it will quote all the params like the table name. That
        # doesn't work. So, we substitute a "%s" for "%s" here so we
        # keep it after python's own string substitution.
        sql = "SELECT %s FROM %s WHERE %s = %s" % (select_field, table, compare_field, "%s")

        # Now, the next icky thing: we need to make sure that we're
        # passing a dict so that psycopg2 can pick it apart properly
        # for its own substitution code. We force this with the
        # trailing comma here
        data = (value, )
        self.cursor.execute(sql, data)

        if self.cursor.rowcount > 0:
            return self.cursor.fetchone()[0]
        else:
            return None

    # Grab one column from multiple rows of a query; useful as a quick wrapper
    def _get_multi_elements(self, select_field, table, compare_field, value):

        # We really want to use psycopg's type handling deal with the
        # (potentially) user-supplied data in the value field, so we
        # have to pass (sql,data) through to cursor.execute. However,
        # we can't have psycopg do all the argument substitution here
        # as it will quote all the params like the table name. That
        # doesn't work. So, we substitute a "%s" for "%s" here so we
        # keep it after python's own string substitution.
        sql = "SELECT %s FROM %s WHERE %s = %s" % (select_field, table, compare_field, "%s")

        # Now, the next icky thing: we need to make sure that we're
        # passing a dict so that psycopg2 can pick it apart properly
        # for its own substitution code. We force this with the
        # trailing comma here
        data = (value, )
        self.cursor.execute(sql, data)

        if self.cursor.rowcount > 0:
            results = []
            for record in self.cursor:
                results.append(record[0])
        else:
            return None

    def get_switch_id(self, name):
        return self._get_element("switch_id", "switch", "name", name)

    def get_port_id(self, name):
        return self._get_element("port_id", "port", "name", name)

    def get_vlan_id_from_name(self, name):
        return self._get_element("vlan_id", "vlan", "name", name)

    def get_vlan_id_from_tag(self, tag):
        return self._get_element("vlan_id", "vlan", "tag", tag)

    def get_switch_name(self, switch_id):
        return self._get_element("name", "switch", "switch_id", switch_id)

    def get_port_name(self, port_id):
        return self._get_element("port_name", "port", "port_id", port_id)

    def get_vlan_name(self, vlan_id):
        return self._get_element("vlan_name", "vlan", "vlan_id", vlan_id)

    def _get_row(self, table, field, value):

        # We really want to use psycopg's type handling deal with the
        # (potentially) user-supplied data in the value field, so we
        # have to pass (sql,data) through to cursor.execute. However,
        # we can't have psycopg do all the argument substitution here
        # as it will quote all the params like the table name. That
        # doesn't work. So, we substitute a "%s" for "%s" here so we
        # keep it after python's own string substitution.
        sql = "SELECT * FROM %s WHERE %s = %s" % (table, field, "%s")

        # Now, the next icky thing: we need to make sure that we're
        # passing a dict so that psycopg2 can pick it apart properly
        # for its own substitution code. We force this with the
        # trailing comma here
        data = (value, )
        self.cursor.execute(sql, data)
        return self.cursor.fetchone()

    def get_switch(self, switch_id):
        return self._get_row("switch", "switch_id", switch_id)

    def get_port(self, port_id):
        return self._get_row("port", "port_id", port_id)

    def get_vlan(self, vlan_id):
        return self._get_row("vlan", "vlan_id", vlan_id)

    # (Un)Lock a port in the database. This can only be done through
    # the admin interface, and will stop API users from modifying
    # settings on the port. Use this to lock down ports that are used
    # for PDUs and other core infrastructure
    def set_port_is_locked(self, port_id, is_locked):
        try:
            sql = "UPDATE port SET is_locked=%s WHERE port_id=%s"
            data = (is_locked, port_id)
            self.cursor.execute(sql, data)
            port_id = self.cursor.fetchone()[0]
            self.connection.commit()
        except:
            self.connection.rollback()
            raise
        return port_id        

    def set_vlan(self, port_id, vlan_id):
        port = self.get_port(port_id)
        if port == None:
            raise("Port %s does not exist" % port_id)

        if port["is_trunk"] or port["is_locked"]:
            raise CriticalError("The port is locked")

        vlan = self.get_vlan(vlan_id)
        if vlan == None:
            raise CriticalError("VLAN %s does not exist" % vlan_id)

        try:
            sql = "UPDATE port SET current_vlan_id=%s WHERE port_id=%s"
            data = (vlan_id, port_id)
            self.cursor.execute(sql, data)
        except:
            self.connection.rollback()
            raise

    def restore_default_vlan(self, port_id):
        port = self.get_port(port_id)
        if port == None:
            raise CriticalError("Port %s does not exist")

        if port["is_trunk"] or port["is_locked"]:
            raise CriticalError("The port is locked")

        try:
            sql = "UPDATE port SET current_vlan_id=base_vlan_id WHERE port_id=%d"
            data = port_id
            self.cursor.execute(sql, data)
        except:
            self.connection.rollback()
            raise

    def _dump_table(self, table):
        result = []
        self.cursor.execute("SELECT * FROM %s" % table)
        record = self.cursor.fetchone()
        while record != None:
            result.append(record)
            record = self.cursor.fetchone()
        return result

    def all_switches(self):
        return self._dump_table("switch")

    def all_ports(self):
        return self._dump_table("port")

    def all_vlans(self):
        return self._dump_table("vlan")