tracking: Allow identical tracking codes across types.

1 files changed, 10 insertions, 8 deletions

diff --git a/modules/tracking.py b/modules/tracking.py
index 40e7e44..90b3555 100644
--- a/modules/tracking.py
+++ b/modules/tracking.py
@@ -10,7 +10,7 @@ import urllib2, datetime, re, pytz, json, urllib
 from xml.etree import ElementTree as ET
 from twisted.internet import reactor
 from twisted.internet.task import LoopingCall
-from sqlalchemy import create_engine, Column, Integer, String, DateTime, ForeignKey
+from sqlalchemy import create_engine, Column, Integer, String, DateTime, ForeignKey, UniqueConstraint
 from sqlalchemy.ext.declarative import declarative_base
 from sqlalchemy.orm import sessionmaker, relation, backref
 from sqlalchemy.orm.exc import NoResultFound
@@ -27,14 +27,16 @@ class Consignment(Base):
 	server = Column(String, nullable = False)
 	nick = Column(String, nullable = False)
 	channel = Column(String) # NULL for private
-	code = Column(String, nullable = False, unique = True)
-	label = Column(String)
+	code = Column(String, nullable = False, index = True)
+	label = Column(String, index = True)
 	added = Column(DateTime, nullable = False)
 	last = Column(DateTime)
 	type = Column(String, nullable = False) # Tracking module
 
 	packages = relation('Package', backref = 'consignment', primaryjoin = 'Package.consignment_id == Consignment.id')
 
+	__table_args__ = (UniqueConstraint('code', 'type', name = 'consignment_code_type_uc'),)
+
 	def __init__(self, server, nick, channel, type, code, label, added):
 		self.server = server
 		self.nick = nick
@@ -58,7 +60,7 @@ class Package(Base):
 
 	id = Column(Integer, primary_key = True)
 	consignment_id = Column(Integer, ForeignKey('consignment.id'))
-	code = Column(String, nullable = False, unique = True)
+	code = Column(String, nullable = False, index = True)
 	last = Column(DateTime) # Last update
 	status = Column(String) # Last status
 
@@ -290,7 +292,7 @@ class Module:
 			session.add(consignment)
 			session.commit()
 			msg = 'Now tracking \002%s\002.' % code
-			reactor.callLater(1, self.track_update, code, True)
+			reactor.callLater(1, self.track_update, type, code, True)
 		except IntegrityError as e:
 			msg = 'Already tracking \002%s\002.' % code
 		finally:
@@ -340,7 +342,7 @@ class Module:
 			results = []
 			for row in consignments:
 				try:
-					self.track_update(row.code, propagate_error = True)
+					self.track_update(row.type, row.code, propagate_error = True)
 				except NoPackageFound:
 					results.append('No packages found for %s' % unicode(row))
 					continue
@@ -507,10 +509,10 @@ url    [TRACKINGNO] - Lists URLs for matches to the corresponding website''')
 			self.irc.msg(channel if not channel == self.irc.nickname else nick.split('!')[0], 'No data returned (this is a bug).')
 
 	# called by start and status
-	def track_update(self, code, announce = False, propagate_error = False):
+	def track_update(self, type, code, announce = False, propagate_error = False):
 		try:
 			session = Session()
-			consignment = session.query(Consignment).filter(and_(Consignment.server == self.irc.factory.server, Consignment.code == code)).one()
+			consignment = session.query(Consignment).filter(and_(Consignment.server == self.irc.factory.server, Consignment.type == type, Consignment.code == code)).one()
 			self.update_consignment(session, consignment, announce, propagate_error)
 			session.commit()
 		finally: