Merge branch 'master'

3 files changed, 29 insertions, 26 deletions

diff --git a/f1elo/db.py b/f1elo/db.py
index f4ed307..554a061 100644
--- a/f1elo/db.py
+++ b/f1elo/db.py
@@ -2,13 +2,16 @@ import json
 from os import path
 
 import __main__
-from sqlalchemy import create_engine
+from sqlalchemy import create_engine, event
 from sqlalchemy.orm import sessionmaker
 
 config = json.load(open(path.join(path.dirname(__main__.__file__), 'config', 'db.json')))
 if config['engine'] == 'mysql':
     engine = create_engine("mysql://{0[user]}:{0[pass]}@{0[host]}/{0[db]}?charset=utf8".format(config))
-else:
+elif config['engine'] == 'sqlite':
     engine = create_engine("sqlite:///{0[file]}".format(config))
+    def fk_pragma(conn, record):
+        conn.execute('PRAGMA FOREIGN_KEYS=ON');
+    event.listen(engine, 'connect', fk_pragma)
 
 Session = sessionmaker(bind=engine)
diff --git a/f1elo/model.py b/f1elo/model.py
index 7865198..3c27ac9 100644
--- a/f1elo/model.py
+++ b/f1elo/model.py
@@ -12,7 +12,7 @@ class Driver(Base):
     driver = Column(String(1024))
     country = Column(String(255))
 
-    rankings = relationship('Ranking', order_by='Ranking.rank_date', back_populates='driver')
+    rankings = relationship('Ranking', order_by='Ranking.rank_date', back_populates='driver', cascade="all", passive_deletes=True)
 
     def __repr__(self):
         return (u"<%s (#%d)>" % (self.driver, self.id)).encode('utf8')
@@ -25,10 +25,20 @@ class Driver(Base):
             return ranks[-1]
         return None
 
+    @staticmethod
+    def fetch(name, country, session):
+        driver = session.query(Driver).filter(Driver.driver==name).first()
+        if not driver:
+            driver = Driver()
+            driver.driver = name
+            driver.country = country
+            session.add(driver)
+        return driver
+
 driver_entry = Table('driver_entries', Base.metadata,
-                     Column('_driver', Integer, ForeignKey('drivers.id')),
-                     Column('_entry', Integer, ForeignKey('entries.id')),
-                     Column('id', Integer))
+                     Column('_driver', Integer, ForeignKey('drivers.id', onupdate="CASCADE", ondelete="CASCADE")),
+                     Column('_entry', Integer, ForeignKey('entries.id', onupdate="CASCADE", ondelete="CASCADE")),
+                     Column('id', Integer, primary_key=True))
 
 class Entry(Base):
     __tablename__ = 'entries'
@@ -37,10 +47,11 @@ class Entry(Base):
     result = Column(String(255))
     car_no = Column(String(255))
     result_group = Column(Integer)
-    _race = Column(Integer, ForeignKey('races.id'))
 
+    _race = Column(Integer, ForeignKey('races.id', onupdate="CASCADE", ondelete="CASCADE"))
     race = relationship('Race', back_populates='entries', order_by=result_group)
-    drivers = relationship('Driver', secondary=driver_entry)
+
+    drivers = relationship('Driver', secondary=driver_entry, cascade="all", passive_deletes=True)
 
     def __repr__(self):
         return ('#%s (%s) %s[%d]' % (self.car_no, ', '.join([driver.__repr__().decode('utf8') for driver in self.drivers]), self.result, self.result_group)).encode('utf8')
@@ -53,9 +64,10 @@ class Race(Base):
     date = Column(Date)
     ranked = Column(Boolean, default=False)
 
-    _type = Column(Integer, ForeignKey('race_types.id'))
+    _type = Column(Integer, ForeignKey('race_types.id', onupdate="CASCADE", ondelete="CASCADE"))
     type = relationship('RaceType', back_populates='races', order_by='Race.date')
-    entries = relationship('Entry', back_populates='race', order_by='Entry.result_group')
+
+    entries = relationship('Entry', back_populates='race', order_by='Entry.result_group', cascade="all", passive_deletes=True)
 
     def __repr__(self):
         return ('%s (%s)' % (self.race, self.date)).encode('utf8')
@@ -67,7 +79,7 @@ class RaceType(Base):
     code = Column(String(255))
     description = Column(String(1024))
 
-    races = relationship('Race', back_populates='type')
+    races = relationship('Race', back_populates='type', cascade="all", passive_deletes=True)
 
     def __repr__(self):
         return ('%s (%s)' % (self.description, self.code)).encode('utf8')
@@ -79,22 +91,11 @@ class Ranking(Base):
     rank_date = Column(Date)
     ranking = Column(Float)
 
-    _driver = Column(Integer, ForeignKey('drivers.id'))
+    _driver = Column(Integer, ForeignKey('drivers.id', onupdate="CASCADE", ondelete="CASCADE"))
     driver = relationship('Driver', back_populates='rankings', order_by=rank_date)
 
     def __repr__(self):
         return ("%s: %0.2f (%s)" % (self.driver.__repr__().decode('utf8'), self.ranking, self. rank_date)).encode('utf8')
 
-def find_driver(name, country, session):
-    driver = session.query(Driver).filter(Driver.driver==name).first()
-    if driver:
-        return driver
-    else:
-        driver = Driver()
-        driver.driver = name
-        driver.country = country
-        session.add(driver)
-        return driver
-
 
 __all__ = ['Driver', 'Entry', 'Ranking', 'Race', 'RaceType']
diff --git a/import-csv.py b/import-csv.py
index 49b699e..199fa5e 100644
--- a/import-csv.py
+++ b/import-csv.py
@@ -2,7 +2,6 @@ import csv
 import sys
 
 from f1elo.db import Session
-from f1elo.model import find_driver
 from f1elo.model import *
 
 session = Session()
@@ -18,10 +17,10 @@ with open(sys.argv[1]) as f:
             entry.car_no = row[2]
             entry.result_group = row[5]
             session.add(entry)
-            driver = find_driver(row[4].strip(), row[3].strip(), session)
+            driver = Driver.fetch(row[4].strip(), row[3].strip(), session)
             entry.drivers.append(driver)
         elif len(row) == 2:
-            driver = find_driver(row[1].strip(), row[0].strip(), session)
+            driver = Driver.fetch(row[1].strip(), row[0].strip(), session)
             entry.drivers.append(driver)
         else:
             print row