SQLiteKeyValueStore.cs

using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Data;
using System.Data.SQLite;
using System.Diagnostics;
using System.Text;

namespace Ramp.Data
{
    /// <summary>
    ///     An SQLite-backed key value store.
    /// </summary>
    public sealed class SQLiteKeyValueStore : IKeyValueStore, IDisposable
    {
        /// <summary>
        ///     The maximum length of a bucket string.
        /// </summary>
        public const uint MaxBucketLength = 64;

        /// <summary>
        ///     The maximum length of a key string.
        /// </summary>
        public const uint MaxKeyLength = 192;

        public const string DefaultBucketName = "";

        public const string MetadataBucketName = "_M";

        private const long CurrentVersion = 1;

        private readonly SQLiteCommand _get;
        private readonly SQLiteCommand _set;
        private readonly SQLiteCommand _delete;
        private readonly SQLiteCommand _exists;
        private readonly SQLiteCommand _getItems;
        private readonly SQLiteCommand _getKeys;

        private readonly Func<Type, object, string> _convertToStringFunc;
        private readonly Func<Type, string, object> _convertFromStringFunc;

        private readonly Dictionary<Type, TypeConverter> _converters = new Dictionary<Type, TypeConverter>();

        public SQLiteKeyValueStore(string connectionString,
                                   Func<Type, object, string> convertToStringFunc = null,
                                   Func<Type, string, object> convertFromStringFunc = null)
        {
            if (connectionString == null)
                throw new ArgumentNullException("connectionString");

            _convertToStringFunc = convertToStringFunc ?? DefaultConvertToString;
            _convertFromStringFunc = convertFromStringFunc ?? DefaultConvertFromString;

            // Create/open the database
            Connection = new SQLiteConnection(connectionString);
            Connection.Open();

            // Update the database schema to latest version
            Update();

            // Prepare all needed commands
            _get = new SQLiteCommand("SELECT Value FROM Store WHERE Bucket = ? AND Key = ?", Connection);
            _get.Parameters.Add(null, DbType.String);
            _get.Parameters.Add(null, DbType.String);
            _get.Prepare();

            _set = new SQLiteCommand("INSERT OR REPLACE INTO Store (Bucket, Key, Value) VALUES (?, ?, ?)", Connection);
            _set.Parameters.Add(null, DbType.String);
            _set.Parameters.Add(null, DbType.String);
            _set.Parameters.Add(null, DbType.Binary);
            _set.Prepare();

            _delete = new SQLiteCommand("DELETE FROM Store WHERE Bucket = ? AND Key = ?", Connection);
            _delete.Parameters.Add(null, DbType.String);
            _delete.Parameters.Add(null, DbType.String);
            _delete.Prepare();

            _exists = new SQLiteCommand("SELECT COUNT(Key) FROM Store WHERE Bucket = ? AND Key = ?", Connection);
            _exists.Parameters.Add(null, DbType.String);
            _exists.Parameters.Add(null, DbType.String);
            _exists.Prepare();

            _getItems = new SQLiteCommand("SELECT Key, Value FROM Store WHERE Bucket = ?", Connection);
            _getItems.Parameters.Add(null, DbType.String);
            _getItems.Prepare();

            _getKeys = new SQLiteCommand("SELECT Key FROM Store WHERE Bucket = ?", Connection);
            _getKeys.Parameters.Add(null, DbType.String);
            _getKeys.Prepare();
        }

        /// <summary>
        ///     The SQLite connection.
        /// </summary>
        public SQLiteConnection Connection { get; private set; }

        public Encoding Encoding { get { return Encoding.UTF8; } }

        string IKeyValueStore.DefaultBucketName
        {
            get { return DefaultBucketName; }
        }

        string IKeyValueStore.MetadataBucketName
        {
            get { return MetadataBucketName; }
        }

        public void Dispose()
        {
            _getKeys.Dispose();
            _getItems.Dispose();
            _delete.Dispose();
            _set.Dispose();
            _get.Dispose();
            Connection.Dispose();
        }

        public bool TryGetBytes(string bucket, string key, out byte[] value)
        {
            ValidateArgs(bucket, key);

            _get.Parameters[0].Value = bucket ?? DefaultBucketName;
            _get.Parameters[1].Value = key;
            object resultBase = _get.ExecuteScalar();

            // Should only be DBNull if the entry doesn't exist in the database.
            if (Convert.IsDBNull(resultBase))
            {
                value = null;
                return false;
            }

            value = (byte[]) resultBase;
            return true;
        }

        public bool TryGetString(string bucket, string key, out string value)
        {
            ValidateArgs(bucket, key);

            byte[] resultBytes;
            if (!TryGetBytes(bucket, key, out resultBytes))
            {
                value = null;
                return false;
            }

            value = Encoding.GetString(resultBytes);
            return true;
        }

        public bool TryGetValue<T>(string bucket, string key, out T value)
        {
            ValidateArgs(bucket, key);

            string resultString;
            if (!TryGetString(bucket, key, out resultString))
            {
                value = default(T);
                return false;
            }

            try
            {
                value = (T) _convertFromStringFunc(typeof(T), resultString);
                return true;
            }
            catch (Exception)
            {
                value = default(T);
                return false;
            }
        }

        public byte[] GetBytesOrDefault(string bucket, string key, byte[] defaultValue = null)
        {
            byte[] result;
            return TryGetBytes(bucket, key, out result) ? result : defaultValue;
        }

        public string GetStringOrDefault(string bucket, string key, string defaultValue = null)
        {
            string result;
            return TryGetString(bucket, key, out result) ? result : defaultValue;
        }

        public T GetValueOrDefault<T>(string bucket, string key, T defaultValue = default(T))
        {
            T result;
            return TryGetValue(bucket, key, out result) ? result : defaultValue;
        }

        public void Set(string bucket, string key, string value)
        {
            Set(bucket, key, value == null ? null : Encoding.GetBytes(value));
        }

        public void Set(string bucket, string key, byte[] value)
        {
            ValidateArgs(bucket, key);

            _set.Parameters[0].Value = bucket ?? DefaultBucketName;
            _set.Parameters[1].Value = key;
            _set.Parameters[2].Value = value;
            _set.ExecuteNonQuery();
            _set.Parameters[2].Value = null; // value could be large, don't prevent a GC
        }

        public void Set<T>(string bucket, string key, T value)
        {
            if (value is DBNull)
                throw new ArgumentOutOfRangeException("value", "value cannot be DBNull");

            string stringValue = _convertToStringFunc(typeof(T), value);

            Set(bucket, key, stringValue);
        }

        public bool Delete(string bucket, string key)
        {
            ValidateArgs(bucket, key);

            _delete.Parameters[0].Value = bucket ?? DefaultBucketName;
            _delete.Parameters[1].Value = key;
            return _delete.ExecuteNonQuery() != 0;
        }

        public bool Exists(string bucket, string key)
        {
            ValidateArgs(bucket, key);

            _exists.Parameters[0].Value = bucket ?? DefaultBucketName;
            _exists.Parameters[1].Value = key;
            return (long) _exists.ExecuteScalar() != 0;
        }

        public IEnumerable<KeyValuePair<string, byte[]>> GetItemBytes(string bucket)
        {
            ValidateBucket(bucket);

            foreach (var item in GetItems(bucket))
            {
                yield return new KeyValuePair<string, byte[]>(item.Key, (byte[]) item.Value);
            }
        }

        public IEnumerable<KeyValuePair<string, string>> GetItemStrings(string bucket)
        {
            ValidateBucket(bucket);

            foreach (KeyValuePair<string, object> item in GetItems(bucket))
            {
                string stringValue = item.Value != null ? Encoding.GetString((byte[]) item.Value) : null;
                yield return new KeyValuePair<string, string>(item.Key, stringValue);
            }
        }

        public IEnumerable<string> GetKeys(string bucket)
        {
            ValidateBucket(bucket);

            _getItems.Parameters[0].Value = bucket ?? DefaultBucketName;
            using (SQLiteDataReader reader = _getItems.ExecuteReader())
                while (reader.Read())
                    yield return reader.GetString(0);
        }

        public IDbTransaction BeginTransaction()
        {
            return Connection.BeginTransaction();
        }

        public IDbTransaction BeginTransaction(IsolationLevel isolationLevel)
        {
            return Connection.BeginTransaction(isolationLevel);
        }

        private string DefaultConvertToString(Type type, object value)
        {
            return GetConverter(type).ConvertToString(value);
        }

        private object DefaultConvertFromString(Type type, string value)
        {
            return GetConverter(type).ConvertFromString(value);
        }

        // Caches type converters to avoid entering the type descriptor machinery each time.
        private TypeConverter GetConverter(Type type)
        {
            TypeConverter result;
            if (!_converters.TryGetValue(type, out result))
            {
                result = TypeDescriptor.GetConverter(type);
                _converters.Add(type, result);
            }
            return result;
        }

        private IEnumerable<KeyValuePair<string, object>> GetItems(string bucket)
        {
            _getItems.Parameters[0].Value = bucket ?? DefaultBucketName;
            using (SQLiteDataReader reader = _getItems.ExecuteReader())
            {
                while (reader.Read())
                {
                    string key = reader.GetString(0);
                    object value = reader.GetValue(1);
                    Debug.Assert(!Convert.IsDBNull(value));
                    yield return new KeyValuePair<string, object>(key, value);
                }
            }
        }

        // Updates the store to the latest schema
        private void Update()
        {
            using (SQLiteTransaction transaction = Connection.BeginTransaction())
            using (SQLiteCommand cmd = Connection.CreateCommand())
            {
                cmd.CommandText = "PRAGMA user_version";
                var version = (long) cmd.ExecuteScalar();

                if (version < 1)
                {
                    cmd.CommandText =
                        "CREATE TABLE Store (Bucket TEXT NOT NULL, Key TEXT NOT NULL, Value BLOB, PRIMARY KEY (Bucket, Key))";
                    cmd.ExecuteNonQuery();
                }

                cmd.CommandText = "PRAGMA user_version = " + CurrentVersion;
                cmd.ExecuteNonQuery();

                transaction.Commit();
            }
        }

        private void ValidateBucket(string bucket)
        {
            if (bucket != null && bucket.Length > MaxBucketLength)
                throw new ArgumentOutOfRangeException("bucket");
        }

        private void ValidateArgs(string bucket, string key)
        {
            if (bucket != null && bucket.Length > MaxBucketLength)
                throw new ArgumentOutOfRangeException("bucket");
            if (key == null)
                throw new ArgumentNullException("key");
            if (key.Length > MaxKeyLength)
                throw new ArgumentOutOfRangeException("key");
        }
    }
}