sopel/db.py

import json
import os.path
import sys
import sqlite3

from tools import Identifier


def _deserialize(value):
	if value is None:
		return None
	# sqlite likes to return ints for strings that look like ints, even though
	# the column type is string. That's how you do dynamic typing wrong.
	value = str(value)
	# Just in case someone's mucking with the DB in a way we can't account for,
	# ignore json parsing errors
	try:
		value = json.loads(value)
	except:
		pass
	return value


class SopelDB(object):
	"""*Availability: 5.0+*

	This defines an interface for basic, common operations on a sqlite
	database. It simplifies those common operations, and allows direct access
	to the database, wherever the user has configured it to be.

	When configured with a relative filename, it is assumed to be in the same
	directory as the config."""

	def __init__(self, config):
		path = config.core.db_filename
		config_dir, config_file = os.path.split(config.filename)
		config_name, _ = os.path.splitext(config_file)
		if path is None:
			path = os.path.join(config_dir, config_name + '.db')
		path = os.path.expanduser(path)
		if not os.path.isabs(path):
			path = os.path.normpath(os.path.join(config_dir, path))
		self.filename = path
		self._create()

	def connect(self):
		"""Return a raw database connection object."""
		return sqlite3.connect(self.filename, timeout=10)

	def execute(self, *args, **kwargs):
		"""Execute an arbitrary SQL query against the database.

		Returns a cursor object, on which things like `.fetchall()` can be
		called per PEP 249."""
		with self.connect() as conn:
			cur = conn.cursor()
			return cur.execute(*args, **kwargs)

	def _create(self):
		"""Create the basic database structure."""
		# Do nothing if the db already exists.
		try:
			self.execute('SELECT * FROM nick_ids;')
			self.execute('SELECT * FROM nicknames;')
			self.execute('SELECT * FROM nick_values;')
			self.execute('SELECT * FROM channel_values;')
		except:
			pass
		else:
			return

		self.execute(
			'CREATE TABLE nick_ids (nick_id INTEGER PRIMARY KEY AUTOINCREMENT)'
		)
		self.execute(
			'CREATE TABLE nicknames '
			'(nick_id INTEGER REFERENCES nick_ids, '
			'slug STRING PRIMARY KEY, canonical string)'
		)
		self.execute(
			'CREATE TABLE nick_values '
			'(nick_id INTEGER REFERENCES nick_ids(nick_id), '
			'key STRING, value STRING, '
			'PRIMARY KEY (nick_id, key))'
		)
		self.execute(
			'CREATE TABLE channel_values '
			'(channel STRING, key STRING, value STRING, '
			'PRIMARY KEY (channel, key))'
		)

	def get_uri(self):
		"""Returns a URL for the database, usable to connect with SQLAlchemy.
		"""
		return 'sqlite://{}'.format(self.filename)

	# NICK FUNCTIONS

	def get_nick_id(self, nick, create=True):
		"""Return the internal identifier for a given nick.

		This identifier is unique to a user, and shared across all of that
		user's aliases. If create is True, a new ID will be created if one does
		not already exist"""
		slug = nick.lower()
		nick_id = self.execute('SELECT nick_id from nicknames where slug = ?',
							   [slug]).fetchone()
		if nick_id is None:
			if not create:
				raise ValueError('No ID exists for the given nick')
			with self.connect() as conn:
				cur = conn.cursor()
				cur.execute('INSERT INTO nick_ids VALUES (NULL)')
				nick_id = cur.execute('SELECT last_insert_rowid()').fetchone()[0]
				cur.execute(
					'INSERT INTO nicknames (nick_id, slug, canonical) VALUES '
					'(?, ?, ?)',
					[nick_id, slug, nick]
				)
			nick_id = self.execute('SELECT nick_id from nicknames where slug = ?',
								   [slug]).fetchone()
		return nick_id[0]

	def alias_nick(self, nick, alias):
		"""Create an alias for a nick.

		Raises ValueError if the alias already exists. If nick does not already
		exist, it will be added along with the alias."""
		nick = Identifier(nick)
		alias = Identifier(alias)
		nick_id = self.get_nick_id(nick)
		sql = 'INSERT INTO nicknames (nick_id, slug, canonical) VALUES (?, ?, ?)'
		values = [nick_id, alias.lower(), alias]
		try:
			self.execute(sql, values)
		except sqlite3.IntegrityError:
			raise ValueError('Alias already exists.')

	def set_nick_value(self, nick, key, value):
		"""Sets the value for a given key to be associated with the nick."""
		nick = Identifier(nick)
		value = json.dumps(value, ensure_ascii=False)
		nick_id = self.get_nick_id(nick)
		self.execute('INSERT OR REPLACE INTO nick_values VALUES (?, ?, ?)',
					 [nick_id, key, value])

	def get_nick_value(self, nick, key):
		"""Retrieves the value for a given key associated with a nick."""
		nick = Identifier(nick)
		result = self.execute(
			'SELECT value FROM nicknames JOIN nick_values '
			'ON nicknames.nick_id = nick_values.nick_id '
			'WHERE slug = ? AND key = ?',
			[nick.lower(), key]
		).fetchone()
		if result is not None:
			result = result[0]
		return _deserialize(result)

	def unalias_nick(self, alias):
		"""Removes an alias.

		Raises ValueError if there is not at least one other nick in the group.
		To delete an entire group, use `delete_group`.
		"""
		alias = Identifier(alias)
		nick_id = self.get_nick_id(alias, False)
		count = self.execute('SELECT COUNT(*) FROM nicknames WHERE nick_id = ?',
							 [nick_id]).fetchone()[0]
		if count <= 1:
			raise ValueError('Given alias is the only entry in its group.')
		self.execute('DELETE FROM nicknames WHERE slug = ?', [alias.lower()])

	def delete_nick_group(self, nick):
		"""Removes a nickname, and all associated aliases and settings.
		"""
		nick = Identifier(nick)
		nick_id = self.get_nick_id(nick, False)
		self.execute('DELETE FROM nicknames WHERE nick_id = ?', [nick_id])
		self.execute('DELETE FROM nick_values WHERE nick_id = ?', [nick_id])

	def merge_nick_groups(self, first_nick, second_nick):
		"""Merges the nick groups for the specified nicks.

		Takes two nicks, which may or may not be registered.  Unregistered
		nicks will be registered. Keys which are set for only one of the given
		nicks will be preserved. Where multiple nicks have values for a given
		key, the value set for the first nick will be used.

		Note that merging of data only applies to the native key-value store.
		If modules define their own tables which rely on the nick table, they
		will need to have their merging done separately."""
		first_id = self.get_nick_id(Identifier(first_nick))
		second_id = self.get_nick_id(Identifier(second_nick))
		self.execute(
			'UPDATE OR IGNORE nick_values SET nick_id = ? WHERE nick_id = ?',
			[first_id, second_id])
		self.execute('DELETE FROM nick_values WHERE nick_id = ?', [second_id])
		self.execute('UPDATE nicknames SET nick_id = ? WHERE nick_id = ?',
					 [first_id, second_id])

	# CHANNEL FUNCTIONS

	def set_channel_value(self, channel, key, value):
		channel = Identifier(channel).lower()
		value = json.dumps(value, ensure_ascii=False)
		self.execute('INSERT OR REPLACE INTO channel_values VALUES (?, ?, ?)',
					 [channel, key, value])

	def get_channel_value(self, channel, key):
		"""Retrieves the value for a given key associated with a channel."""
		channel = Identifier(channel).lower()
		result = self.execute(
			'SELECT value FROM channel_values WHERE channel = ? AND key = ?',
			[channel, key]
		).fetchone()
		if result is not None:
			result = result[0]
		return _deserialize(result)

	# NICK AND CHANNEL FUNCTIONS

	def get_nick_or_channel_value(self, name, key):
		"""Gets the value `key` associated to the nick or channel  `name`.
		"""
		name = Identifier(name)
		if name.is_nick():
			return self.get_nick_value(name, key)
		else:
			return self.get_channel_value(name, key)

	def get_preferred_value(self, names, key):
		"""Gets the value for the first name which has it set.

		`names` is a list of channel and/or user names. Returns None if none of
		the names have the key set."""
		for name in names:
			value = self.get_nick_or_channel_value(name, key)
			if value is not None:
				return value