Add helpers for get/set of base/current vlan on ports

Change-Id: I528fcb83e5b7e8e252b0970dfad2c902fab92df9
diff --git a/db/db.py b/db/db.py
index d71f46d..3e33133 100644
--- a/db/db.py
+++ b/db/db.py
@@ -288,6 +288,12 @@
     def get_port_by_switch_and_name(self, switch_id, name):
         return self._get_element2("port_id", "port", "switch_id", switch_id, "name", name)
 
+    def get_current_vlan_id_by_port(self, port_id):
+        return self._get_element("current_vlan_id", "port", "port_id", port_id)
+
+    def get_base_vlan_id_by_port(self, port_id):
+        return self._get_element("base_vlan_id", "port", "port_id", port_id)
+
     def get_ports_by_current_vlan(self, vlan_id):
         return self._get_multi_elements("port_id", "port", "current_vlan_id", vlan_id)
 
@@ -390,7 +396,32 @@
             raise
         return port_id
 
-    def restore_default_vlan(self, port_id):
+    def set_base_vlan(self, port_id, vlan_id):
+        port = self.get_port_by_id(port_id)
+        if port is None:
+            raise InputError("Port ID %d does not exist" % int(port_id))
+
+        if port.is_trunk or port.is_locked:
+            raise CriticalError("The port is locked")
+
+        vlan = self.get_vlan_by_id(vlan_id)
+        if vlan is None:
+            raise InputError("VLAN ID %d does not exist" % int(vlan_id))
+        if not vlan.is_base_vlan:
+            raise InputError("VLAN ID %d is not a base VLAN" % int(vlan_id))
+
+        try:
+            sql = "UPDATE port SET base_vlan_id=%s WHERE port_id=%s RETURNING port_id"
+            data = (vlan_id, 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 restore_base_vlan(self, port_id):
         port = self.get_port_by_id(port_id)
         if port is None:
             raise InputError("Port ID %d does not exist" % int(port_id))