Merge "Stop throwing exceptions when a key doesn't exist when all keys is specified for SharedPreferencesMigration" into androidx-main
diff --git a/buildSrc/src/main/kotlin/androidx/build/LibraryVersions.kt b/buildSrc/src/main/kotlin/androidx/build/LibraryVersions.kt
index c25342a..3e3254a 100644
--- a/buildSrc/src/main/kotlin/androidx/build/LibraryVersions.kt
+++ b/buildSrc/src/main/kotlin/androidx/build/LibraryVersions.kt
@@ -56,7 +56,7 @@
     val CORE_ROLE = Version("1.1.0-alpha02")
     val CURSORADAPTER = Version("1.1.0-alpha01")
     val CUSTOMVIEW = Version("1.2.0-alpha01")
-    val DATASTORE = Version("1.0.0-rc01")
+    val DATASTORE = Version("1.0.0-rc02")
     val DOCUMENTFILE = Version("1.1.0-alpha01")
     val DRAWERLAYOUT = Version("1.2.0-alpha01")
     val DYNAMICANIMATION = Version("1.1.0-alpha04")
diff --git a/datastore/datastore/src/androidTest/java/androidx/datastore/migrations/SharedPreferencesMigrationTest.kt b/datastore/datastore/src/androidTest/java/androidx/datastore/migrations/SharedPreferencesMigrationTest.kt
index 52604f60..0387583 100644
--- a/datastore/datastore/src/androidTest/java/androidx/datastore/migrations/SharedPreferencesMigrationTest.kt
+++ b/datastore/datastore/src/androidTest/java/androidx/datastore/migrations/SharedPreferencesMigrationTest.kt
@@ -154,6 +154,21 @@
     }
 
     @Test
+    fun testSharedPrefsViewWithAllKeysSpecified_doesntThrowErrorWhenKeyDoesntExist() =
+        runBlockingTest {
+            assertThat(sharedPrefs.edit().putInt("unrelated_key", -123).commit()).isTrue()
+
+            val migration = SharedPreferencesMigration(
+                produceSharedPreferences = { sharedPrefs },
+            ) { prefs: SharedPreferencesView, _: Byte ->
+                prefs.getInt("this_key_doesnt_exist_yet", 123).toByte()
+            }
+
+            val dataStore = getDataStoreWithMigrations(listOf(migration))
+            assertThat(dataStore.data.first()).isEqualTo(123)
+        }
+
+    @Test
     fun producedSharedPreferencesIsUsed() = runBlockingTest {
         assertThat(sharedPrefs.edit().putInt("integer_key", 123).commit()).isTrue()
 
diff --git a/datastore/datastore/src/main/java/androidx/datastore/migrations/SharedPreferencesMigration.kt b/datastore/datastore/src/main/java/androidx/datastore/migrations/SharedPreferencesMigration.kt
index 76cfae9..309729f 100644
--- a/datastore/datastore/src/main/java/androidx/datastore/migrations/SharedPreferencesMigration.kt
+++ b/datastore/datastore/src/main/java/androidx/datastore/migrations/SharedPreferencesMigration.kt
@@ -131,20 +131,26 @@
 
     private val sharedPrefs: SharedPreferences by lazy(produceSharedPreferences)
 
-    private val keySet: MutableSet<String> by lazy {
+    /**
+     * keySet is null if the user specified [MIGRATE_ALL_KEYS].
+     */
+    private val keySet: MutableSet<String>? =
         if (keysToMigrate === MIGRATE_ALL_KEYS) {
-            sharedPrefs.all.keys
+            null
         } else {
-            keysToMigrate
-        }.toMutableSet()
-    }
+            keysToMigrate.toMutableSet()
+        }
 
     override suspend fun shouldMigrate(currentData: T): Boolean {
         if (!shouldRunMigration(currentData)) {
             return false
         }
 
-        return keySet.any(sharedPrefs::contains)
+        return if (keySet == null) {
+            sharedPrefs.all.isNotEmpty()
+        } else {
+            keySet.any(sharedPrefs::contains)
+        }
     }
 
     override suspend fun migrate(currentData: T): T =
@@ -160,8 +166,12 @@
     override suspend fun cleanUp() {
         val sharedPrefsEditor = sharedPrefs.edit()
 
-        for (key in keySet) {
-            sharedPrefsEditor.remove(key)
+        if (keySet == null) {
+            sharedPrefsEditor.clear()
+        } else {
+            keySet.forEach { key ->
+                sharedPrefsEditor.remove(key)
+            }
         }
 
         if (!sharedPrefsEditor.commit()) {
@@ -172,7 +182,7 @@
             deleteSharedPreferences(context, name)
         }
 
-        keySet.clear()
+        keySet?.clear()
     }
 
     private fun deleteSharedPreferences(context: Context, name: String) {
@@ -211,12 +221,11 @@
 }
 
 /**
- *  Read-only wrapper around SharedPreferences. This will be passed in to your migration. The
- *  constructor is public to enable easier testing of migrations.
+ *  Read-only wrapper around SharedPreferences. This will be passed in to your migration.
  */
 public class SharedPreferencesView internal constructor(
     private val prefs: SharedPreferences,
-    private val keySet: Set<String>
+    private val keySet: Set<String>?
 ) {
     /**
      * Checks whether the preferences contains a preference.
@@ -285,18 +294,22 @@
         prefs.getStringSet(checkKey(key), defValues)?.toMutableSet()
 
     /** Retrieve all values from the preferences that are in the specified keySet. */
-    public fun getAll(): Map<String, Any?> = prefs.all.filter { (key, _) ->
-        key in keySet
-    }.mapValues { (_, value) ->
-        if (value is Set<*>) {
-            value.toSet()
-        } else {
-            value
+    public fun getAll(): Map<String, Any?> =
+        prefs.all.filter { (key, _) ->
+            keySet?.contains(key) ?: true
+        }.mapValues { (_, value) ->
+            if (value is Set<*>) {
+                value.toSet()
+            } else {
+                value
+            }
         }
-    }
 
     private fun checkKey(key: String): String {
-        check(key in keySet) { "Can't access key outside migration: $key" }
+        keySet?.let {
+            check(key in it) { "Can't access key outside migration: $key" }
+        }
+
         return key
     }
 }