Add convenience category for YapDatabaseConnection.

// FREEBIE

SignalServiceKit/src/Profiles/OWSProfilesManager.m

@@ -87,8 +87,8 @@ static const NSInteger kProfileKeyLength = 16;
 
 @interface OWSProfilesManager ()
 
-@property (nonatomic, readonly) TSStorageManager *storageManager;
 @property (nonatomic, readonly) OWSMessageSender *messageSender;
+@property (nonatomic, readonly) YapDatabaseConnection *dbConnection;
 
 @property (atomic, readonly, nullable) NSData *localProfileKey;
 
@@ -136,8 +136,8 @@ static const NSInteger kProfileKeyLength = 16;
     OWSAssert(storageManager);
     OWSAssert(messageSender);
 
-    _storageManager = storageManager;
     _messageSender = messageSender;
+    _dbConnection = storageManager.newDatabaseConnection;
 
     OWSSingletonAssert();
 
@@ -146,13 +146,13 @@ static const NSInteger kProfileKeyLength = 16;
     [messageSender setProfilesManager:self];
 
     // Try to load.
-    _localProfileKey = [self.storageManager objectForKey:kOWSProfilesManager_LocalProfileSecretKey
+    _localProfileKey = [self.dbConnection objectForKey:kOWSProfilesManager_LocalProfileSecretKey
                                             inCollection:kOWSProfilesManager_Collection];
     if (!_localProfileKey) {
         // Generate
         _localProfileKey = [OWSProfilesManager generateLocalProfileKey];
         // Persist
-        [self.storageManager setObject:_localProfileKey
+        [self.dbConnection setObject:_localProfileKey
                                 forKey:kOWSProfilesManager_LocalProfileSecretKey
                           inCollection:kOWSProfilesManager_Collection];
     }
@@ -214,20 +214,20 @@ static const NSInteger kProfileKeyLength = 16;
     self.localProfileAvatarMetadata = localProfileAvatarMetadata;
 
     if (localProfileName) {
-        [self.storageManager setObject:localProfileName
-                                forKey:kOWSProfilesManager_LocalProfileNameKey
-                          inCollection:kOWSProfilesManager_Collection];
+        [self.dbConnection setObject:localProfileName
+                              forKey:kOWSProfilesManager_LocalProfileNameKey
+                        inCollection:kOWSProfilesManager_Collection];
     } else {
-        [self.storageManager removeObjectForKey:kOWSProfilesManager_LocalProfileNameKey
-                                   inCollection:kOWSProfilesManager_Collection];
+        [self.dbConnection removeObjectForKey:kOWSProfilesManager_LocalProfileNameKey
+                                 inCollection:kOWSProfilesManager_Collection];
     }
     if (localProfileAvatarMetadata) {
-        [self.storageManager setObject:localProfileAvatarMetadata
-                                forKey:kOWSProfilesManager_LocalProfileAvatarMetadataKey
-                          inCollection:kOWSProfilesManager_Collection];
+        [self.dbConnection setObject:localProfileAvatarMetadata
+                              forKey:kOWSProfilesManager_LocalProfileAvatarMetadataKey
+                        inCollection:kOWSProfilesManager_Collection];
     } else {
-        [self.storageManager removeObjectForKey:kOWSProfilesManager_LocalProfileAvatarMetadataKey
-                                   inCollection:kOWSProfilesManager_Collection];
+        [self.dbConnection removeObjectForKey:kOWSProfilesManager_LocalProfileAvatarMetadataKey
+                                 inCollection:kOWSProfilesManager_Collection];
     }
 
     [[NSNotificationCenter defaultCenter] postNotificationName:kNSNotificationName_LocalProfileDidChange
@@ -378,11 +378,11 @@ static const NSInteger kProfileKeyLength = 16;
 - (void)loadLocalProfileAsync
 {
     dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-        NSString *_Nullable localProfileName = [self.storageManager objectForKey:kOWSProfilesManager_LocalProfileNameKey
-                                                                    inCollection:kOWSProfilesManager_Collection];
+        NSString *_Nullable localProfileName = [self.dbConnection objectForKey:kOWSProfilesManager_LocalProfileNameKey
+                                                                  inCollection:kOWSProfilesManager_Collection];
         AvatarMetadata *_Nullable localProfileAvatarMetadata =
-            [self.storageManager objectForKey:kOWSProfilesManager_LocalProfileAvatarMetadataKey
-                                 inCollection:kOWSProfilesManager_Collection];
+            [self.dbConnection objectForKey:kOWSProfilesManager_LocalProfileAvatarMetadataKey
+                               inCollection:kOWSProfilesManager_Collection];
         UIImage *_Nullable localProfileAvatarImage = nil;
         if (localProfileAvatarMetadata) {
             localProfileAvatarImage = [self loadProfileAvatarWithFilename:localProfileAvatarMetadata.fileName];

SignalServiceKit/src/Storage/TSStorageManager.h

@@ -3,7 +3,7 @@
 //
 
 #import "TSStorageKeys.h"
-
+#import "YapDatabaseConnection+OWS.h"
 #import <Foundation/Foundation.h>
 #import <YapDatabase/YapDatabase.h>
 

SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.h

@@ -0,0 +1,37 @@
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
+
+#import <YapDatabase/YapDatabase.h>
+
+@class ECKeyPair;
+@class PreKeyRecord;
+@class SignedPreKeyRecord;
+
+NS_ASSUME_NONNULL_BEGIN
+
+@interface YapDatabaseConnection (OWS)
+
+- (BOOL)boolForKey:(NSString *)key inCollection:(NSString *)collection;
+- (int)intForKey:(NSString *)key inCollection:(NSString *)collection;
+- (id)objectForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSDate *)dateForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSDictionary *)dictionaryForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSString *)stringForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable NSData *)dataForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable ECKeyPair *)keyPairForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable PreKeyRecord *)preKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection;
+- (nullable SignedPreKeyRecord *)signedPreKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection;
+
+#pragma mark -
+
+- (void)setObject:(id)object forKey:(NSString *)key inCollection:(NSString *)collection;
+- (void)removeObjectForKey:(NSString *)string inCollection:(NSString *)collection;
+- (void)setInt:(int)integer forKey:(NSString *)key inCollection:(NSString *)collection;
+- (void)setDate:(NSDate *)value forKey:(NSString *)key inCollection:(NSString *)collection;
+- (void)purgeCollection:(NSString *)collection;
+- (int)incrementIntForKey:(NSString *)key inCollection:(NSString *)collection;
+
+@end
+
+NS_ASSUME_NONNULL_END

SignalServiceKit/src/Storage/YapDatabaseConnection+OWS.m

@@ -0,0 +1,138 @@
+//
+//  Copyright (c) 2017 Open Whisper Systems. All rights reserved.
+//
+
+#import "YapDatabaseConnection+OWS.h"
+#import <YapDatabase/YapDatabaseTransaction.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+@implementation YapDatabaseConnection (OWS)
+
+- (id)objectForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    __block NSString *object;
+
+    [self readWithBlock:^(YapDatabaseReadTransaction *transaction) {
+        object = [transaction objectForKey:key inCollection:collection];
+    }];
+
+    return object;
+}
+
+- (nullable NSDictionary *)dictionaryForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    __block NSDictionary *object;
+
+    [self readWithBlock:^(YapDatabaseReadTransaction *transaction) {
+        object = [transaction objectForKey:key inCollection:collection];
+    }];
+
+    return object;
+}
+
+- (nullable NSString *)stringForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    NSString *string = [self objectForKey:key inCollection:collection];
+
+    return string;
+}
+
+- (BOOL)boolForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    NSNumber *boolNum = [self objectForKey:key inCollection:collection];
+
+    return [boolNum boolValue];
+}
+
+- (nullable NSData *)dataForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    NSData *data = [self objectForKey:key inCollection:collection];
+    return data;
+}
+
+- (nullable ECKeyPair *)keyPairForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    ECKeyPair *keyPair = [self objectForKey:key inCollection:collection];
+
+    return keyPair;
+}
+
+- (nullable PreKeyRecord *)preKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    PreKeyRecord *preKeyRecord = [self objectForKey:key inCollection:collection];
+
+    return preKeyRecord;
+}
+
+- (nullable SignedPreKeyRecord *)signedPreKeyRecordForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    SignedPreKeyRecord *preKeyRecord = [self objectForKey:key inCollection:collection];
+
+    return preKeyRecord;
+}
+
+- (int)intForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    int integer = [[self objectForKey:key inCollection:collection] intValue];
+
+    return integer;
+}
+
+- (nullable NSDate *)dateForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    NSNumber *value = [self objectForKey:key inCollection:collection];
+    if (value) {
+        return [NSDate dateWithTimeIntervalSince1970:value.doubleValue];
+    } else {
+        return nil;
+    }
+}
+
+#pragma mark -
+
+- (void)purgeCollection:(NSString *)collection
+{
+    [self readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        [transaction removeAllObjectsInCollection:collection];
+    }];
+}
+
+- (void)setObject:(id)object forKey:(NSString *)key inCollection:(NSString *)collection
+{
+    [self readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        [transaction setObject:object forKey:key inCollection:collection];
+    }];
+}
+
+- (void)removeObjectForKey:(NSString *)string inCollection:(NSString *)collection
+{
+    [self readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        [transaction removeObjectForKey:string inCollection:collection];
+    }];
+}
+
+- (void)setInt:(int)integer forKey:(NSString *)key inCollection:(NSString *)collection
+{
+    [self setObject:[NSNumber numberWithInt:integer] forKey:key inCollection:collection];
+}
+
+- (int)incrementIntForKey:(NSString *)key inCollection:(NSString *)collection
+{
+    __block int value = 0;
+    [self readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
+        value = [[transaction objectForKey:key inCollection:collection] intValue];
+        value++;
+        [transaction setObject:@(value) forKey:key inCollection:collection];
+    }];
+    return value;
+}
+
+- (void)setDate:(NSDate *)value forKey:(NSString *)key inCollection:(NSString *)collection
+{
+    [self setObject:@(value.timeIntervalSince1970) forKey:key inCollection:collection];
+}
+
+@end
+
+NS_ASSUME_NONNULL_END