mongodbexporter.py

from datetime import date
from hashlib import md5
from jsonexporter import jsonizer as _jsonizer
import json
import os
import pymongo
import urlparse


def jsonizer(obj):
	return _jsonizer(obj, specify_object_type=False)


class _MongoImportFile(object):
	def __init__(self, name, path):
		self._path = os.path.abspath(os.path.expanduser(path))
		self._fname = os.path.join(self._path, name + '.json')
		self._previous_md5_fname = os.path.join(self._path, name + '_previous.md5')
		self._f = open(self._fname, 'a')

	def update(self, id_dict, content, **kwargs):
		line = json.dumps(content)
		lines = line.splitlines()
		if (len(lines) > 1):
			line = ' '.join(lines)
		self._f.writelines((line, '\n'))
		#print '>%s: %s' % (self._fname, line)

	def ensure_index(self, name, *args, **kwargs):
		pass

	def close(self):
		pass
		if self._f is not None and not self._f.closed:
			self._f.close()
		if self._md5 is not None and not self._md5.closed:
			self._md5.close()


class _MongoImportFileSetConnectionMock(object):
	def __init__(self, *args):
		self._files = args

	def disconnect(self, *args):
		for open_file in (f for f in self._files if f is not None):
			# close the files
			try:
				open_file.close()
			except ValueError:
				pass


class _MongoImportFileSet(object):
	'''Interface to a set of mongo import files providing almost the same interface as mongoConnection[db_name] does.'''
	def __init__(self, base_path):
		self._path = base_path
		self._files = {}
		self._artists = None
		self._labels = None
		self._releases = None
		self._masters = None
		self.collection_names = lambda: ('artists', 'labels', 'releases', 'masters')
		self.connection = _MongoImportFileSetConnectionMock(*(self._artists, self._labels, self._releases, self._masters))

	def __getitem__(self, key):
		if key in self.collection_names():
			attr = self.__dict__['_' + key]
			if attr is None:
				attr = _MongoImportFile(key, self._path)
			return attr
		raise IndexError('Invalid key: %s' % key)

	def __getattr__(self, name):
		if name in self.collection_names():
			return self[name]
		else:
			raise AttributeError("Unknown attribute '%s'." % name)


class _IdHashPairs(object):
	'''Cache of the current records in the database indexed by id'''
	def __init__(self, path):
		self._path = os.path.abspath(os.path.expanduser(path))
		self._hashes = {'artists': None, 'labels': None, 'releases': None, 'masters': None}
		self._name = lambda x: os.path.join(self._path, '%s.md5' % x)

	def _load(self, name):
		if self._hashes[name] is None:
			try:
				fname = self._name(name)
				with open(fname, 'rb') as f:
					name_hash = {}
					while 1:
						lines = f.readlines(10000)
						if not lines:
							break
						for line in lines:
							line = line.rstrip('\r\n')
							old_id, md5 = line.split(':')
							name_hash[int(old_id)] = md5
					self._hashes[name] = name_hash
			except IOError as e:
				self._hashes[name] = {}

	def _one_line(self, line):
		lines = line.splitlines()
		return ' '.join(lines) if (len(lines) > 1) else line

	def is_uniq(self, collection, id, json_string):
		self._load(collection)
		new_md5 = md5(self._one_line(json_string)).hexdigest()
		i_id = int(id)
		if i_id in self._hashes[collection]:
			old_md5 = self._hashes[collection][i_id]
			if new_md5 == old_md5:
				return (False, new_md5)
			else:
				return (True, new_md5)
		else:
			return (True, new_md5)

	def process(self, collection, id, md5_digest):
		self._load(collection)
		i_id = int(id)
		self._hashes[collection][i_id] = md5_digest

	def finish(self):
		# write hashes back to disk
		for name in self._hashes:
			try:
				fname = self._name(name)
				if self._hashes[name]:
					with open(fname, 'w') as f:
						for id, hash in self._hashes[name].iteritems():
							f.writelines(('%s:%s' % (id, hash), '\n'))
			except IOError as e:
				print 'IOError writing out %s: %s' % (name, e)


class MongoDbExporter(object):
	def __init__(self, mongo_uri, data_quality=[]):
		'''mongo_uri: mongodb://[username:password@]host1[:port1],...[,hostN[:portN]][/[database][?options]]'''
		# TODO: if uri is file://path/ - create a json dump for using with mongo import
		self.min_data_quality = data_quality
		self._options = {}
		self._quick_uniq = None
		self.connect(mongo_uri)
		# don't submit to mongo values that already exist - faster to compute it here then in mongo

	def connect(self, mongo_uri):
		db_name, options = None, {}
		u = urlparse.urlparse(mongo_uri)
		if u.scheme == 'file':
			path = u.path
			self._options = urlparse.parse_qs(u.query) if u.query else {}
			path = u.netloc + path
			self.db = _MongoImportFileSet(path)
			if 'uniq' in self._options and 'md5' in self._options['uniq']:
				self._quick_uniq = _IdHashPairs(path)
		elif u.scheme == 'mongodb':
			if '?' in u.path and u.query == '':
				#url didn't parse it properly u.path is '/dbname?options
				db_name, self._options = u.path.split('?', 1)
				self._options = urlparse.parse_qs(self._options) if self._options else {}
			else:
				db_name = u.path
			if db_name.startswith('/'):
				db_name = db_name[1:]
			#print 'Connecting to db %s on %s with options.' % (db_name, mongo_uri, options)
			mongo = pymongo.Connection(mongo_uri)
			self.db = mongo[db_name]
			if 'uniq' in self._options and 'md5' in self._options['uniq']:
				self._quick_uniq = False
		else:
			raise ValueError("Invalid URI scheme: '%s'. Can only accept 'file' or 'mongodb'" % u.scheme)

	def _is_uniq(self, collection, id, json_string):
		if self._quick_uniq is not None:
			return self._quick_uniq.is_uniq(collection, id, json_string)
		return (True, None)

	def _store_processed(self, collection, id, md5_digest):
		if self._quick_uniq is not None:
			self._quick_uniq.process(collection, id, md5_digest)

	def good_quality(self, what):
		if len(self.min_data_quality):
			return what.data_quality.lower() in self.min_data_quality
		return True

	def execute(self, collection, what):
		if not self.good_quality(what):
			# print "Bad quality: %s for %s" % (what.data_quality, what.id)
			return
		# have to convert it to json and back because
		# on simple objects couchdb-python throws:
		# TypeError: argument of type 'instance' is not iterable
		# and on dicts:
		# AttributeError: 'dict' object has no attribute 'read'
		json_string = json.dumps(what, default=jsonizer)
		uniq, md5 = self._is_uniq(collection, what.id, json_string)
		if uniq:
			doc = json.loads(json_string)
			doc['updated_on'] = "%s" % date.today()
			self.db[collection].update({'id': what.id}, doc, upsert=True)
			self._store_processed(collection, what.id, md5)

	def finish(self, completely_done=False):
		collections = self.db.collection_names()
		if 'artists' in collections:
			#self.db.artists.('id', background=True)
			self.db.artists.ensure_index('id', background=True, unique=True)
			# should be unique=True, but can't seem to offer a true guarantee
			self.db.artists.ensure_index('l_name', background=True)
		if 'labels' in collections:
			self.db.labels.ensure_index('id', background=True, unique=True)
			self.db.labels.ensure_index('l_name', background=True)
		if 'releases' in collections:
			self.db.releases.ensure_index('id', background=True, unique=True)
			self.db.releases.ensure_index([('l_artist', pymongo.ASCENDING),
				('l_title', pymongo.ASCENDING)],
				background=True)
			self.db.releases.ensure_index('format.name', background=True)
		if 'masters' in collections:
			self.db.masters.ensure_index('id', background=True, unique=True)
			self.db.masters.ensure_index('l_title', background=True)
			self.db.masters.ensure_index('main_release', background=True, unique=True)
		self.db.connection.disconnect()
		if self._quick_uniq is not None:
			self._quick_uniq.finish()

	def storeLabel(self, label):
		label.l_name = label.name.lower()
		self.execute('labels', label)

	def storeArtist(self, artist):
		artist.l_name = artist.name.lower()
		self.execute('artists', artist)

	def storeRelease(self, release):
		if hasattr(release, 'artist'):
			release.l_artist = release.artist.lower()
		else:
			release.l_artist = None
		release.l_title = release.title.lower()
		self.execute('releases', release)

	def storeMaster(self, master):
		master.l_artist = master.artist.lower()
		master.l_title = master.title.lower()
		self.execute('masters', master)