From f2a21599d5da7a91f4bf2636440087e896ff03e2 Mon Sep 17 00:00:00 2001
From: Anton Beloglazov <anton.beloglazov@gmail.com>
Date: Sat, 22 Sep 2012 17:47:30 +1000
Subject: [PATCH] Added a hosts db table, added update_host

---
 neat/db.py             | 38 ++++++++++++++++++++++++++++++++++++--
 neat/db_utils.py       |  8 +++++++-
 tests/test_db.py       | 19 +++++++++++++++++++
 tests/test_db_utils.py |  3 +++
 4 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/neat/db.py b/neat/db.py
index dc5f2b7..c96d4c0 100644
--- a/neat/db.py
+++ b/neat/db.py
@@ -21,15 +21,20 @@ class Database(object):
     """ A class representing the database, where fields are tables.
     """
 
-    @contract(connection=Connection, vms=Table, vm_resource_usage=Table)
-    def __init__(self, connection, vms, vm_resource_usage):
+    @contract(connection=Connection,
+              hosts=Table,
+              vms=Table,
+              vm_resource_usage=Table)
+    def __init__(self, connection, hosts, vms, vm_resource_usage):
         """ Initialize the database.
 
         :param connection: A database connection table.
+        :param hosts: The hosts table.
         :param vms: The vms table.
         :param vm_resource_usage: The vm_resource_usage table.
         """
         self.connection = connection
+        self.hosts = hosts
         self.vms = vms
         self.vm_resource_usage = vm_resource_usage
 
@@ -86,3 +91,32 @@ class Database(object):
                 query.append({'vm_id': vm_id,
                               'cpu_mhz': cpu_mhz})
             self.vm_resource_usage.insert().execute(query)
+
+    @contract
+    def update_host(self, hostname, cpu_mhz, ram):
+        """ Insert new or update the corresponding host record.
+
+        :param hostname: A host name.
+         :type hostname: str
+
+        :param cpu_mhz: The total CPU frequency of the host in MHz.
+         :type cpu_mhz: int,>0
+
+        :param ram: The total amount of RAM of the host in MB.
+         :type ram: int,>0
+
+        :return: The ID of the host.
+         :rtype: int
+        """
+        sel = select([self.hosts.c.id]).where(self.hosts.c.hostname == hostname)
+        row = self.connection.execute(sel).fetchone()
+        if row == None:
+            return self.hosts.insert().execute(hostname=hostname,
+                                               cpu_mhz=cpu_mhz,
+                                               ram=ram).inserted_primary_key[0]
+        else:
+            self.connection.execute(self.hosts.update().
+                                    where(self.hosts.c.id == row['id']).
+                                    values(cpu_mhz=cpu_mhz,
+                                           ram=ram))
+            return row['id']
diff --git a/neat/db_utils.py b/neat/db_utils.py
index 86c2d82..5da24b9 100644
--- a/neat/db_utils.py
+++ b/neat/db_utils.py
@@ -34,6 +34,12 @@ def init_db(sql_connection):
     metadata = MetaData()
     metadata.bind = engine
 
+    hosts = Table('hosts', metadata,
+                Column('id', Integer, primary_key=True),
+                Column('hostname', String(255), nullable=False),
+                Column('cpu_mhz', Integer, nullable=False),
+                Column('ram', Integer, nullable=False))
+
     vms = Table('vms', metadata,
                 Column('id', Integer, primary_key=True),
                 Column('uuid', String(36), nullable=False))
@@ -48,4 +54,4 @@ def init_db(sql_connection):
     metadata.create_all()
     connection = engine.connect()
 
-    return Database(connection, vms, vm_resource_usage)
+    return Database(connection, hosts, vms, vm_resource_usage)
diff --git a/tests/test_db.py b/tests/test_db.py
index 55e01c8..ea3c4b7 100644
--- a/tests/test_db.py
+++ b/tests/test_db.py
@@ -84,3 +84,22 @@ class Db(TestCase):
 
         for uuid, data in final_data.items():
             assert db.select_cpu_mhz_for_vm(uuid, 11) == data
+
+    @qc(1)
+    def update_host():
+        db = db_utils.init_db('sqlite:///:memory:')
+        db.update_host('host1', 3000, 4000)
+        hosts = db.hosts.select().execute().fetchall()
+        assert len(hosts) == 1
+        host = hosts[0]
+        assert host['hostname'] == 'host1'
+        assert host['cpu_mhz'] == 3000
+        assert host['ram'] == 4000
+
+        db.update_host('host1', 3500, 8000)
+        hosts = db.hosts.select().execute().fetchall()
+        assert len(hosts) == 1
+        host = hosts[0]
+        assert host['hostname'] == 'host1'
+        assert host['cpu_mhz'] == 3500
+        assert host['ram'] == 8000
diff --git a/tests/test_db_utils.py b/tests/test_db_utils.py
index a0b1e0d..1a05427 100644
--- a/tests/test_db_utils.py
+++ b/tests/test_db_utils.py
@@ -27,8 +27,11 @@ class DbUtils(TestCase):
     def init_db():
         db = db_utils.init_db('sqlite:///:memory:')
         assert type(db) is neat.db.Database
+        assert isinstance(db.hosts, Table)
         assert isinstance(db.vms, Table)
         assert isinstance(db.vm_resource_usage, Table)
+        assert db.hosts.c.keys() == \
+            ['id', 'hostname', 'cpu_mhz', 'ram']
         assert db.vms.c.keys() == \
             ['id', 'uuid']
         assert db.vm_resource_usage.c.keys() == \