From 8b0e69eaa2ccef85f35e7458182bf41c5e2e3938 Mon Sep 17 00:00:00 2001
From: Marcio Mazza <mazza@interlegis.leg.br>
Date: Fri, 5 Jun 2015 21:07:58 -0300
Subject: [PATCH] Add ForeignKey legacy migration and some models

---
 legacy/legacy_migration.py | 37 +++++++++++++++++++++++++++++++------
 parlamentares/models.py    | 16 ++++++++++++++++
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/legacy/legacy_migration.py b/legacy/legacy_migration.py
index e38c04e5c..dabdb00c7 100644
--- a/legacy/legacy_migration.py
+++ b/legacy/legacy_migration.py
@@ -1,20 +1,39 @@
 from django.db import connection
+from django.db import models
+
 
 import legacy
-from parlamentares.models import NivelInstrucao
+from parlamentares.models import (
+    Legislatura, SessaoLegislativa, NivelInstrucao)
 
 
 mappings = (
-    (NivelInstrucao,
-        'cod_nivel_instrucao',
-     [('des_nivel_instrucao', 'nivel_instrucao')]),
+    (Legislatura, 'num_legislatura', [
+        ('dat_inicio', 'data_inicio'),
+        ('dat_fim', 'data_fim'),
+        ('dat_eleicao', 'data_eleicao'), ]
+     ),
+    (SessaoLegislativa, 'cod_sessao_leg', [
+        ('num_legislatura', 'legislatura'),
+        ('num_sessao_leg', 'numero'),
+        ('tip_sessao_leg', 'tipo'),
+        ('dat_inicio', 'data_inicio'),
+        ('dat_fim', 'data_fim'),
+        ('dat_inicio_intervalo', 'data_inicio_intervalo'),
+        ('dat_fim_intervalo', 'data_fim_intervalo'), ]
+     ),
+    (NivelInstrucao, 'cod_nivel_instrucao', [
+        ('des_nivel_instrucao', 'nivel_instrucao')]
+     ),
 )
 
 
-def run_legacy_migration():
+def migrate():
 
     for model, pk, field_pairs in mappings:
 
+        print 'Migrating %s...' % model.__name__
+
         # clear all model entries
         model.objects.all().delete()
         # resets id sequence
@@ -28,9 +47,15 @@ def run_legacy_migration():
             new = model()
             while not new.id:
                 for old_field, new_field in field_pairs:
-                    setattr(new, new_field, getattr(old, old_field))
+                    value = getattr(old, old_field)
+                    # check for a relation
+                    model_field = model._meta.get_field(new_field)
+                    if isinstance(model_field, models.ForeignKey):
+                        value = model_field.related_model.objects.get(id=value)
+                    setattr(new, new_field, value)
                 new.save()
                 assert new.id <= old_id, 'New id exceeds old one. Be sure your new table was just created!'
                 if new.id < old_id:
                     new.delete()
                     new = model()
+
diff --git a/parlamentares/models.py b/parlamentares/models.py
index 79e9f4aff..c3258c4a6 100644
--- a/parlamentares/models.py
+++ b/parlamentares/models.py
@@ -1,5 +1,21 @@
 from django.db import models
 
 
+class Legislatura(models.Model):
+    data_inicio = models.DateField()
+    data_fim = models.DateField()
+    data_eleicao = models.DateField()
+
+
+class SessaoLegislativa(models.Model):
+    legislatura = models.ForeignKey(Legislatura)
+    numero = models.IntegerField()
+    tipo = models.CharField(max_length=1)
+    data_inicio = models.DateField()
+    data_fim = models.DateField()
+    data_inicio_intervalo = models.DateField(blank=True, null=True)
+    data_fim_intervalo = models.DateField(blank=True, null=True)
+
+
 class NivelInstrucao(models.Model):
     nivel_instrucao = models.CharField(max_length=50)