From 6f5e0a8aa584536b81c0abb67cb93b2802dbc57f Mon Sep 17 00:00:00 2001
From: Anton Beloglazov <anton.beloglazov@gmail.com>
Date: Thu, 18 Oct 2012 10:55:03 +1100
Subject: [PATCH] Added insert_vm_migration

---
 neat/db.py             | 15 +++++++++++++++
 neat/db_utils.py       |  4 ++--
 tests/test_db.py       | 11 +++++++++++
 tests/test_db_utils.py |  4 +++-
 4 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/neat/db.py b/neat/db.py
index 7c71434..58eadd0 100644
--- a/neat/db.py
+++ b/neat/db.py
@@ -300,3 +300,18 @@ class Database(object):
         self.host_overload.insert().execute(
             host_id=self.select_host_id(hostname),
             overload=int(overload))
+
+    @contract
+    def insert_vm_migration(self, vm, hostname):
+        """ Insert a VM migration.
+
+        :param hostname: A VM UUID.
+         :type hostname: str[36]
+
+        :param hostname: A host name.
+         :type hostname: str
+        """
+        self.vm_migrations.insert().execute(
+            vm_id=self.select_vm_id(vm),
+            host_id=self.select_host_id(hostname))
+        
diff --git a/neat/db_utils.py b/neat/db_utils.py
index a9d0ce2..ad8efd9 100644
--- a/neat/db_utils.py
+++ b/neat/db_utils.py
@@ -60,8 +60,8 @@ def init_db(sql_connection):
         Table('vm_migrations', metadata,
               Column('id', Integer, primary_key=True),
               Column('vm_id', Integer, ForeignKey('vms.id'), nullable=False),
-              Column('timestamp', DateTime, default=func.now()),
-              Column('hostname', String(255), nullable=False))
+              Column('host_id', Integer, ForeignKey('hosts.id'), nullable=False),
+              Column('timestamp', DateTime, default=func.now()))
 
     host_states = \
         Table('host_states', metadata,
diff --git a/tests/test_db.py b/tests/test_db.py
index fe84d30..2141b7f 100644
--- a/tests/test_db.py
+++ b/tests/test_db.py
@@ -279,3 +279,14 @@ class Db(TestCase):
                     lambda x: x[1] == hosts['host2'], 
                     result), key=lambda x: x[0])]
         self.assertEqual(host2, [0, 1])        
+
+    @qc(1)
+    def insert_select():
+        db = db_utils.init_db('sqlite:///:memory:')
+        db.vms.insert().execute(uuid='x' * 36).inserted_primary_key[0]
+        vm_id = db.vms.insert().execute(uuid='vm' * 18).inserted_primary_key[0]
+        host_id = db.update_host('host', 1, 1, 1)
+        db.insert_vm_migration('vm' * 18, 'host')
+        result = db.vm_migrations.select().execute().first()
+        assert result[1] == vm_id
+        assert result[2] == host_id
diff --git a/tests/test_db_utils.py b/tests/test_db_utils.py
index 08e6a67..0b6c3d8 100644
--- a/tests/test_db_utils.py
+++ b/tests/test_db_utils.py
@@ -43,9 +43,11 @@ class DbUtils(TestCase):
         assert list(db.vm_resource_usage.foreign_keys)[0].target_fullname \
             == 'vms.id'
         assert db.vm_migrations.c.keys() == \
-            ['id', 'vm_id', 'timestamp', 'hostname']
+            ['id', 'vm_id', 'host_id', 'timestamp']
         assert list(db.vm_migrations.foreign_keys)[0].target_fullname \
             == 'vms.id'
+        assert list(db.vm_migrations.foreign_keys)[1].target_fullname \
+            == 'hosts.id'
         assert db.host_states.c.keys() == \
             ['id', 'host_id', 'timestamp', 'state']
         assert list(db.host_states.foreign_keys)[0].target_fullname \