Add user serial number to communal db schema

This change adds a new column to the communal database storing the
serial number of the user associated with each widget. It also
implements a migration that is run when a device is updated from the
previous version of database, populating the default value as -1.

Test: atest CommunalDatabaseMigrationsTest
Test: atest CommunalBackupHelperTest
Test: atest CommunalBackupUtilsTest
Test: atest CommunalWidgetDaoTest
Test: atest CommunalWidgetRepositoryImplTest
Bug: 330945203
Flag: com.android.systemui.communal_hub
Change-Id: I93c231ab9a25f580730628815e588e986ffd0102
diff --git a/packages/SystemUI/Android.bp b/packages/SystemUI/Android.bp
index 469b9ce..8e9f5d4 100644
--- a/packages/SystemUI/Android.bp
+++ b/packages/SystemUI/Android.bp
@@ -668,6 +668,7 @@
     ],
     asset_dirs: [
         "tests/goldens",
+        "schemas",
     ],
     static_libs: [
         "SystemUI-res",
@@ -707,6 +708,7 @@
         "androidx-constraintlayout_constraintlayout",
         "androidx.exifinterface_exifinterface",
         "androidx.room_room-runtime",
+        "androidx.room_room-testing",
         "androidx.room_room-ktx",
         "device_state_flags_lib",
         "kotlinx-coroutines-android",
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
index 6ce6cdb..a504a47 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
@@ -24,6 +24,7 @@
 import android.content.applicationContext
 import android.graphics.Bitmap
 import android.os.UserHandle
+import android.os.userManager
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
@@ -47,9 +48,6 @@
 import com.android.systemui.log.logcatLogBuffer
 import com.android.systemui.res.R
 import com.android.systemui.testKosmos
-import com.android.systemui.util.mockito.any
-import com.android.systemui.util.mockito.mock
-import com.android.systemui.util.mockito.whenever
 import com.android.systemui.util.mockito.withArgCaptor
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -59,11 +57,15 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
 import org.mockito.Mock
 import org.mockito.Mockito.eq
 import org.mockito.Mockito.never
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
+import org.mockito.kotlin.any
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -85,6 +87,9 @@
     private val kosmos = testKosmos()
     private val testScope = kosmos.testScope
     private val packageChangeRepository = kosmos.fakePackageChangeRepository
+    private val userManager = kosmos.userManager
+
+    private val mainUser = UserHandle(0)
 
     private val fakeAllowlist =
         listOf(
@@ -109,6 +114,8 @@
 
         whenever(communalWidgetDao.getWidgets()).thenReturn(fakeWidgets)
         whenever(communalWidgetHost.appWidgetProviders).thenReturn(fakeProviders)
+        whenever(userManager.getUserSerialNumber(mainUser.identifier))
+            .thenReturn(testUserSerialNumber(mainUser))
 
         underTest =
             CommunalWidgetRepositoryImpl(
@@ -121,6 +128,7 @@
                 backupManager,
                 backupUtils,
                 packageChangeRepository,
+                userManager,
             )
     }
 
@@ -128,7 +136,7 @@
     fun communalWidgets_queryWidgetsFromDb() =
         testScope.runTest {
             val communalItemRankEntry = CommunalItemRank(uid = 1L, rank = 1)
-            val communalWidgetItemEntry = CommunalWidgetItem(uid = 1L, 1, "pk_name/cls_name", 1L)
+            val communalWidgetItemEntry = CommunalWidgetItem(uid = 1L, 1, "pk_name/cls_name", 1L, 0)
             fakeWidgets.value = mapOf(communalItemRankEntry to communalWidgetItemEntry)
             fakeProviders.value = mapOf(1 to providerInfoA)
 
@@ -154,13 +162,13 @@
             fakeWidgets.value =
                 mapOf(
                     CommunalItemRank(uid = 1L, rank = 1) to
-                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L),
+                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L, 0),
                     CommunalItemRank(uid = 2L, rank = 2) to
-                        CommunalWidgetItem(uid = 2L, 2, "pk_2/cls_2", 2L),
+                        CommunalWidgetItem(uid = 2L, 2, "pk_2/cls_2", 2L, 0),
                     CommunalItemRank(uid = 3L, rank = 3) to
-                        CommunalWidgetItem(uid = 3L, 3, "pk_3/cls_3", 3L),
+                        CommunalWidgetItem(uid = 3L, 3, "pk_3/cls_3", 3L, 0),
                     CommunalItemRank(uid = 4L, rank = 4) to
-                        CommunalWidgetItem(uid = 4L, 4, "pk_4/cls_4", 4L),
+                        CommunalWidgetItem(uid = 4L, 4, "pk_4/cls_4", 4L, 0),
                 )
             fakeProviders.value =
                 mapOf(
@@ -192,9 +200,9 @@
             fakeWidgets.value =
                 mapOf(
                     CommunalItemRank(uid = 1L, rank = 1) to
-                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L),
+                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L, 0),
                     CommunalItemRank(uid = 2L, rank = 2) to
-                        CommunalWidgetItem(uid = 2L, 2, "pk_2/cls_2", 2L),
+                        CommunalWidgetItem(uid = 2L, 2, "pk_2/cls_2", 2L, 0),
                 )
             fakeProviders.value =
                 mapOf(
@@ -249,7 +257,6 @@
             val provider = ComponentName("pkg_name", "cls_name")
             val id = 1
             val priority = 1
-            val user = UserHandle(0)
             whenever(communalWidgetHost.getAppWidgetInfo(id))
                 .thenReturn(PROVIDER_INFO_REQUIRES_CONFIGURATION)
             whenever(
@@ -259,11 +266,12 @@
                     )
                 )
                 .thenReturn(id)
-            underTest.addWidget(provider, user, priority, kosmos.widgetConfiguratorSuccess)
+            underTest.addWidget(provider, mainUser, priority, kosmos.widgetConfiguratorSuccess)
             runCurrent()
 
-            verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
-            verify(communalWidgetDao).addWidget(id, provider, priority)
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider, mainUser)
+            verify(communalWidgetDao)
+                .addWidget(id, provider, priority, testUserSerialNumber(mainUser))
 
             // Verify backup requested
             verify(backupManager).dataChanged()
@@ -275,7 +283,6 @@
             val provider = ComponentName("pkg_name", "cls_name")
             val id = 1
             val priority = 1
-            val user = UserHandle(0)
             whenever(communalWidgetHost.getAppWidgetInfo(id))
                 .thenReturn(PROVIDER_INFO_REQUIRES_CONFIGURATION)
             whenever(
@@ -285,11 +292,12 @@
                     )
                 )
                 .thenReturn(id)
-            underTest.addWidget(provider, user, priority, kosmos.widgetConfiguratorFail)
+            underTest.addWidget(provider, mainUser, priority, kosmos.widgetConfiguratorFail)
             runCurrent()
 
-            verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
-            verify(communalWidgetDao, never()).addWidget(id, provider, priority)
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider, mainUser)
+            verify(communalWidgetDao, never())
+                .addWidget(anyInt(), any<ComponentName>(), anyInt(), anyInt())
             verify(appWidgetHost).deleteAppWidgetId(id)
 
             // Verify backup not requested
@@ -302,7 +310,6 @@
             val provider = ComponentName("pkg_name", "cls_name")
             val id = 1
             val priority = 1
-            val user = UserHandle(0)
             whenever(communalWidgetHost.getAppWidgetInfo(id))
                 .thenReturn(PROVIDER_INFO_REQUIRES_CONFIGURATION)
             whenever(
@@ -312,13 +319,14 @@
                     )
                 )
                 .thenReturn(id)
-            underTest.addWidget(provider, user, priority) {
+            underTest.addWidget(provider, mainUser, priority) {
                 throw IllegalStateException("some error")
             }
             runCurrent()
 
-            verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
-            verify(communalWidgetDao, never()).addWidget(id, provider, priority)
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider, mainUser)
+            verify(communalWidgetDao, never())
+                .addWidget(anyInt(), any<ComponentName>(), anyInt(), anyInt())
             verify(appWidgetHost).deleteAppWidgetId(id)
 
             // Verify backup not requested
@@ -331,7 +339,6 @@
             val provider = ComponentName("pkg_name", "cls_name")
             val id = 1
             val priority = 1
-            val user = UserHandle(0)
             whenever(communalWidgetHost.getAppWidgetInfo(id))
                 .thenReturn(PROVIDER_INFO_CONFIGURATION_OPTIONAL)
             whenever(
@@ -341,11 +348,12 @@
                     )
                 )
                 .thenReturn(id)
-            underTest.addWidget(provider, user, priority, kosmos.widgetConfiguratorFail)
+            underTest.addWidget(provider, mainUser, priority, kosmos.widgetConfiguratorFail)
             runCurrent()
 
-            verify(communalWidgetHost).allocateIdAndBindWidget(provider, user)
-            verify(communalWidgetDao).addWidget(id, provider, priority)
+            verify(communalWidgetHost).allocateIdAndBindWidget(provider, mainUser)
+            verify(communalWidgetDao)
+                .addWidget(id, provider, priority, testUserSerialNumber(mainUser))
 
             // Verify backup requested
             verify(backupManager).dataChanged()
@@ -538,9 +546,9 @@
             fakeWidgets.value =
                 mapOf(
                     CommunalItemRank(uid = 1L, rank = 1) to
-                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L),
+                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L, 0),
                     CommunalItemRank(uid = 2L, rank = 2) to
-                        CommunalWidgetItem(uid = 2L, 2, "pk_2/cls_2", 2L),
+                        CommunalWidgetItem(uid = 2L, 2, "pk_2/cls_2", 2L, 0),
                 )
 
             // Widget 1 is installed
@@ -554,7 +562,7 @@
                         sessionId = 1,
                         packageName = "pk_2",
                         icon = fakeIcon,
-                        user = UserHandle.CURRENT,
+                        user = mainUser,
                     )
                 )
             )
@@ -572,7 +580,7 @@
                         priority = 2,
                         packageName = "pk_2",
                         icon = fakeIcon,
-                        user = UserHandle.CURRENT,
+                        user = mainUser,
                     ),
                 )
         }
@@ -583,7 +591,7 @@
             fakeWidgets.value =
                 mapOf(
                     CommunalItemRank(uid = 1L, rank = 1) to
-                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L),
+                        CommunalWidgetItem(uid = 1L, 1, "pk_1/cls_1", 1L, 0),
                 )
 
             // Widget 1 is pending install
@@ -594,7 +602,7 @@
                         sessionId = 1,
                         packageName = "pk_1",
                         icon = fakeIcon,
-                        user = UserHandle.CURRENT,
+                        user = mainUser,
                     )
                 )
             )
@@ -607,7 +615,7 @@
                         priority = 1,
                         packageName = "pk_1",
                         icon = fakeIcon,
-                        user = UserHandle.CURRENT,
+                        user = mainUser,
                     ),
                 )
 
@@ -633,6 +641,13 @@
         whenever(appWidgetHost.appWidgetIds).thenReturn(ids.toIntArray())
     }
 
+    // Commonly the user id and user serial number are the same, but for testing purposes use a
+    // simple algorithm to map a user id to a different user serial number to make sure the correct
+    // value is used.
+    private fun testUserSerialNumber(user: UserHandle): Int {
+        return user.identifier + 100
+    }
+
     private companion object {
         val PROVIDER_INFO_REQUIRES_CONFIGURATION =
             AppWidgetProviderInfo().apply { configure = ComponentName("test.pkg", "test.cmp") }
@@ -650,11 +665,13 @@
                                 widgetId = 1
                                 componentName = "pk_name/fake_widget_1"
                                 rank = 1
+                                userSerialNumber = 0
                             },
                             CommunalHubState.CommunalWidgetItem().apply {
                                 widgetId = 2
                                 componentName = "pk_name/fake_widget_2"
                                 rank = 2
+                                userSerialNumber = 0
                             },
                         )
                         .toTypedArray()
diff --git a/packages/SystemUI/schemas/com.android.systemui.communal.data.db.CommunalDatabase/2.json b/packages/SystemUI/schemas/com.android.systemui.communal.data.db.CommunalDatabase/2.json
new file mode 100644
index 0000000..f10d92a
--- /dev/null
+++ b/packages/SystemUI/schemas/com.android.systemui.communal.data.db.CommunalDatabase/2.json
@@ -0,0 +1,81 @@
+{
+  "formatVersion": 1,
+  "database": {
+    "version": 2,
+    "identityHash": "02e2da2d36e6955200edd5fb49e63c72",
+    "entities": [
+      {
+        "tableName": "communal_widget_table",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`uid` INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, `widget_id` INTEGER NOT NULL, `component_name` TEXT NOT NULL, `item_id` INTEGER NOT NULL, `user_serial_number` INTEGER NOT NULL DEFAULT -1)",
+        "fields": [
+          {
+            "fieldPath": "uid",
+            "columnName": "uid",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "widgetId",
+            "columnName": "widget_id",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "componentName",
+            "columnName": "component_name",
+            "affinity": "TEXT",
+            "notNull": true
+          },
+          {
+            "fieldPath": "itemId",
+            "columnName": "item_id",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "userSerialNumber",
+            "columnName": "user_serial_number",
+            "affinity": "INTEGER",
+            "notNull": true,
+            "defaultValue": "-1"
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": true,
+          "columnNames": [
+            "uid"
+          ]
+        }
+      },
+      {
+        "tableName": "communal_item_rank_table",
+        "createSql": "CREATE TABLE IF NOT EXISTS `${TABLE_NAME}` (`uid` INTEGER PRIMARY KEY AUTOINCREMENT NOT NULL, `rank` INTEGER NOT NULL DEFAULT 0)",
+        "fields": [
+          {
+            "fieldPath": "uid",
+            "columnName": "uid",
+            "affinity": "INTEGER",
+            "notNull": true
+          },
+          {
+            "fieldPath": "rank",
+            "columnName": "rank",
+            "affinity": "INTEGER",
+            "notNull": true,
+            "defaultValue": "0"
+          }
+        ],
+        "primaryKey": {
+          "autoGenerate": true,
+          "columnNames": [
+            "uid"
+          ]
+        }
+      }
+    ],
+    "setupQueries": [
+      "CREATE TABLE IF NOT EXISTS room_master_table (id INTEGER PRIMARY KEY,identity_hash TEXT)",
+      "INSERT OR REPLACE INTO room_master_table (id,identity_hash) VALUES(42, '02e2da2d36e6955200edd5fb49e63c72')"
+    ]
+  }
+}
\ No newline at end of file
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt b/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt
index a8e5174..c3d2683 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/backup/CommunalBackupUtils.kt
@@ -43,11 +43,13 @@
         val widgetsFromDb = runBlocking { database.communalWidgetDao().getWidgets().first() }
         val widgetsState = mutableListOf<CommunalHubState.CommunalWidgetItem>()
         widgetsFromDb.keys.forEach { rankItem ->
+            val widget = widgetsFromDb[rankItem]!!
             widgetsState.add(
                 CommunalHubState.CommunalWidgetItem().apply {
                     rank = rankItem.rank
-                    widgetId = widgetsFromDb[rankItem]!!.widgetId
-                    componentName = widgetsFromDb[rankItem]?.componentName
+                    widgetId = widget.widgetId
+                    componentName = widget.componentName
+                    userSerialNumber = widget.userSerialNumber
                 }
             )
         }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt
index 3ce8109..dff6352 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalDatabase.kt
@@ -17,17 +17,21 @@
 package com.android.systemui.communal.data.db
 
 import android.content.Context
+import android.util.Log
 import androidx.annotation.VisibleForTesting
 import androidx.room.Database
 import androidx.room.Room
 import androidx.room.RoomDatabase
+import androidx.room.migration.Migration
+import androidx.sqlite.db.SupportSQLiteDatabase
 import com.android.systemui.res.R
 
-@Database(entities = [CommunalWidgetItem::class, CommunalItemRank::class], version = 1)
+@Database(entities = [CommunalWidgetItem::class, CommunalItemRank::class], version = 2)
 abstract class CommunalDatabase : RoomDatabase() {
     abstract fun communalWidgetDao(): CommunalWidgetDao
 
     companion object {
+        private const val TAG = "CommunalDatabase"
         private var instance: CommunalDatabase? = null
 
         /**
@@ -51,7 +55,8 @@
                             context.resources.getString(R.string.config_communalDatabase)
                         )
                         .also { builder ->
-                            builder.fallbackToDestructiveMigration(dropAllTables = false)
+                            builder.addMigrations(MIGRATION_1_2)
+                            builder.fallbackToDestructiveMigration(dropAllTables = true)
                             callback?.let { callback -> builder.addCallback(callback) }
                         }
                         .build()
@@ -64,5 +69,23 @@
         fun setInstance(database: CommunalDatabase) {
             instance = database
         }
+
+        /**
+         * This migration adds a user_serial_number column and sets its default value as
+         * [CommunalWidgetItem.USER_SERIAL_NUMBER_UNDEFINED]. Work profile widgets added before the
+         * migration still work as expected, but they would be backed up as personal.
+         */
+        @VisibleForTesting
+        val MIGRATION_1_2 =
+            object : Migration(1, 2) {
+                override fun migrate(db: SupportSQLiteDatabase) {
+                    Log.i(TAG, "Migrating from version 1 to 2")
+                    db.execSQL(
+                        "ALTER TABLE communal_widget_table " +
+                            "ADD COLUMN user_serial_number INTEGER NOT NULL DEFAULT " +
+                            "${CommunalWidgetItem.USER_SERIAL_NUMBER_UNDEFINED}"
+                    )
+                }
+            }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalEntities.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalEntities.kt
index 0d5336a..e33aead 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalEntities.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalEntities.kt
@@ -29,7 +29,28 @@
     @ColumnInfo(name = "component_name") val componentName: String,
     /** Reference the id of an item persisted in the glanceable hub */
     @ColumnInfo(name = "item_id") val itemId: Long,
-)
+    /**
+     * A serial number of the user that the widget provider is associated with. For example, a work
+     * profile widget.
+     *
+     * A serial number may be different from its user id in that user ids may be recycled but serial
+     * numbers are unique until the device is wiped.
+     *
+     * Most commonly, this value would be 0 for the primary user, and 10 for the work profile.
+     */
+    @ColumnInfo(name = "user_serial_number", defaultValue = "$USER_SERIAL_NUMBER_UNDEFINED")
+    val userSerialNumber: Int,
+) {
+    companion object {
+        /**
+         * The user serial number associated with the widget is undefined.
+         *
+         * This should only happen for widgets migrated from V1 before user serial number was
+         * included in the schema.
+         */
+        const val USER_SERIAL_NUMBER_UNDEFINED = -1
+    }
+}
 
 @Entity(tableName = "communal_item_rank_table")
 data class CommunalItemRank(
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
index d174fd1..4dcd9bf 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
@@ -17,6 +17,7 @@
 package com.android.systemui.communal.data.db
 
 import android.content.ComponentName
+import android.os.UserManager
 import androidx.room.Dao
 import androidx.room.Delete
 import androidx.room.Query
@@ -53,23 +54,34 @@
     private val communalWidgetDaoProvider: Provider<CommunalWidgetDao>,
     @Named(DEFAULT_WIDGETS) private val defaultWidgets: Array<String>,
     @CommunalLog logBuffer: LogBuffer,
+    private val userManager: UserManager,
 ) : RoomDatabase.Callback() {
     companion object {
         private const val TAG = "DefaultWidgetPopulation"
     }
+
     private val logger = Logger(logBuffer, TAG)
 
     override fun onCreate(db: SupportSQLiteDatabase) {
         super.onCreate(db)
-        applicationScope.launch {
-            addDefaultWidgets()
-            logger.i("Default widgets were populated in the database.")
-        }
+        applicationScope.launch { addDefaultWidgets() }
     }
 
     // Read default widgets from config.xml and populate the database.
     private suspend fun addDefaultWidgets() =
         withContext(bgDispatcher) {
+            // Default widgets should be associated with the main user.
+            val userSerialNumber =
+                userManager.mainUser?.let { mainUser ->
+                    userManager.getUserSerialNumber(mainUser.identifier)
+                }
+            if (userSerialNumber == null) {
+                logger.w(
+                    "Skipped populating default widgets because device does not have a main user"
+                )
+                return@withContext
+            }
+
             defaultWidgets.forEachIndexed { index, name ->
                 val provider = ComponentName.unflattenFromString(name)
                 provider?.let {
@@ -80,11 +92,14 @@
                             .addWidget(
                                 widgetId = id,
                                 provider = provider,
-                                priority = defaultWidgets.size - index
+                                priority = defaultWidgets.size - index,
+                                userSerialNumber = userSerialNumber,
                             )
                     }
                 }
             }
+
+            logger.i("Populated default widgets in the database.")
         }
 }
 
@@ -106,10 +121,16 @@
     fun deleteItemRankById(itemId: Long)
 
     @Query(
-        "INSERT INTO communal_widget_table(widget_id, component_name, item_id) " +
-            "VALUES(:widgetId, :componentName, :itemId)"
+        "INSERT INTO communal_widget_table" +
+            "(widget_id, component_name, item_id, user_serial_number) " +
+            "VALUES(:widgetId, :componentName, :itemId, :userSerialNumber)"
     )
-    fun insertWidget(widgetId: Int, componentName: String, itemId: Long): Long
+    fun insertWidget(
+        widgetId: Int,
+        componentName: String,
+        itemId: Long,
+        userSerialNumber: Int,
+    ): Long
 
     @Query("INSERT INTO communal_item_rank_table(rank) VALUES(:rank)")
     fun insertItemRank(rank: Int): Long
@@ -132,28 +153,41 @@
     }
 
     @Transaction
-    fun addWidget(widgetId: Int, provider: ComponentName, priority: Int): Long {
+    fun addWidget(
+        widgetId: Int,
+        provider: ComponentName,
+        priority: Int,
+        userSerialNumber: Int,
+    ): Long {
         return addWidget(
             widgetId = widgetId,
             componentName = provider.flattenToString(),
             priority = priority,
+            userSerialNumber = userSerialNumber,
         )
     }
 
     @Transaction
-    fun addWidget(widgetId: Int, componentName: String, priority: Int): Long {
+    fun addWidget(
+        widgetId: Int,
+        componentName: String,
+        priority: Int,
+        userSerialNumber: Int,
+    ): Long {
         return insertWidget(
             widgetId = widgetId,
             componentName = componentName,
             itemId = insertItemRank(priority),
+            userSerialNumber = userSerialNumber,
         )
     }
 
     @Transaction
     fun deleteWidgetById(widgetId: Int): Boolean {
         val widget =
-            getWidgetByIdNow(widgetId) ?: // no entry to delete from db
-            return false
+            getWidgetByIdNow(widgetId)
+                ?: // no entry to delete from db
+                return false
 
         deleteItemRankById(widget.itemId)
         deleteWidgets(widget)
@@ -166,6 +200,8 @@
         clearCommunalWidgetsTable()
         clearCommunalItemRankTable()
 
-        state.widgets.forEach { addWidget(it.widgetId, it.componentName, it.rank) }
+        state.widgets.forEach {
+            addWidget(it.widgetId, it.componentName, it.rank, it.userSerialNumber)
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
index fdb797d..1af4dc2 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
@@ -20,6 +20,7 @@
 import android.appwidget.AppWidgetProviderInfo
 import android.content.ComponentName
 import android.os.UserHandle
+import android.os.UserManager
 import com.android.systemui.common.data.repository.PackageChangeRepository
 import com.android.systemui.common.shared.model.PackageInstallSession
 import com.android.systemui.communal.data.backup.CommunalBackupUtils
@@ -98,6 +99,7 @@
     private val backupManager: BackupManager,
     private val backupUtils: CommunalBackupUtils,
     packageChangeRepository: PackageChangeRepository,
+    private val userManager: UserManager,
 ) : CommunalWidgetRepository {
     companion object {
         const val TAG = "CommunalWidgetRepository"
@@ -185,6 +187,7 @@
                     widgetId = id,
                     provider = provider,
                     priority = priority,
+                    userSerialNumber = userManager.getUserSerialNumber(user.identifier),
                 )
                 backupManager.dataChanged()
             } else {
@@ -255,6 +258,7 @@
                         widgetId = newWidgetId
                         componentName = restoredWidget.componentName
                         rank = restoredWidget.rank
+                        userSerialNumber = restoredWidget.userSerialNumber
                     }
                 }
             val newState = CommunalHubState().apply { widgets = newWidgets.toTypedArray() }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto b/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto
index 0816259..bc14ae1 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto
+++ b/packages/SystemUI/src/com/android/systemui/communal/proto/communal_hub_state.proto
@@ -35,5 +35,8 @@
 
         // Rank or order of the widget in the communal hub.
         int32 rank = 3;
+
+        // Serial number of the user associated with the widget.
+        int32 user_serial_number = 4;
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
index 7094848..e60848b 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupHelperTest.kt
@@ -120,11 +120,26 @@
 
     private fun setUpDatabase(): List<FakeWidgetMetadata> {
         return listOf(
-                FakeWidgetMetadata(11, "com.android.fakePackage1/fakeWidget1", 3),
-                FakeWidgetMetadata(12, "com.android.fakePackage2/fakeWidget2", 2),
-                FakeWidgetMetadata(13, "com.android.fakePackage3/fakeWidget3", 1),
+                FakeWidgetMetadata(
+                    widgetId = 11,
+                    componentName = "com.android.fakePackage1/fakeWidget1",
+                    rank = 3,
+                    userSerialNumber = 0,
+                ),
+                FakeWidgetMetadata(
+                    widgetId = 12,
+                    componentName = "com.android.fakePackage2/fakeWidget2",
+                    rank = 2,
+                    userSerialNumber = 0,
+                ),
+                FakeWidgetMetadata(
+                    widgetId = 13,
+                    componentName = "com.android.fakePackage3/fakeWidget3",
+                    rank = 1,
+                    userSerialNumber = 10,
+                ),
             )
-            .onEach { dao.addWidget(it.widgetId, it.componentName, it.rank) }
+            .onEach { dao.addWidget(it.widgetId, it.componentName, it.rank, it.userSerialNumber) }
     }
 
     private fun getBackupDataInputStream(): BackupDataInputStream {
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt
index cde7a0e..983a435 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/backup/CommunalBackupUtilsTest.kt
@@ -66,11 +66,28 @@
         // Set up database
         val expectedWidgets =
             listOf(
-                FakeWidgetMetadata(11, "com.android.fakePackage1/fakeWidget1", 3),
-                FakeWidgetMetadata(12, "com.android.fakePackage2/fakeWidget2", 2),
-                FakeWidgetMetadata(13, "com.android.fakePackage3/fakeWidget3", 1),
+                FakeWidgetMetadata(
+                    widgetId = 11,
+                    componentName = "com.android.fakePackage1/fakeWidget1",
+                    rank = 3,
+                    userSerialNumber = 0,
+                ),
+                FakeWidgetMetadata(
+                    widgetId = 12,
+                    componentName = "com.android.fakePackage2/fakeWidget2",
+                    rank = 2,
+                    userSerialNumber = 0,
+                ),
+                FakeWidgetMetadata(
+                    widgetId = 13,
+                    componentName = "com.android.fakePackage3/fakeWidget3",
+                    rank = 1,
+                    userSerialNumber = 10,
+                ),
             )
-        expectedWidgets.forEach { dao.addWidget(it.widgetId, it.componentName, it.rank) }
+        expectedWidgets.forEach {
+            dao.addWidget(it.widgetId, it.componentName, it.rank, it.userSerialNumber)
+        }
 
         // Get communal hub state
         val state = underTest.getCommunalHubState()
@@ -128,7 +145,12 @@
         assertThat(underTest.fileExists()).isFalse()
     }
 
-    data class FakeWidgetMetadata(val widgetId: Int, val componentName: String, val rank: Int)
+    data class FakeWidgetMetadata(
+        val widgetId: Int,
+        val componentName: String,
+        val rank: Int,
+        val userSerialNumber: Int,
+    )
 
     companion object {
         /**
@@ -140,7 +162,8 @@
                 { actual, expected ->
                     actual?.widgetId == expected?.widgetId &&
                         actual?.componentName == expected?.componentName &&
-                        actual?.rank == expected?.rank
+                        actual?.rank == expected?.rank &&
+                        actual?.userSerialNumber == expected?.userSerialNumber
                 },
                 "represents",
             )
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalDatabaseMigrationsTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalDatabaseMigrationsTest.kt
new file mode 100644
index 0000000..eb0ab78
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalDatabaseMigrationsTest.kt
@@ -0,0 +1,149 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.systemui.communal.data.db
+
+import androidx.room.testing.MigrationTestHelper
+import androidx.sqlite.db.SupportSQLiteDatabase
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.lifecycle.InstantTaskExecutorRule
+import com.google.common.truth.Truth.assertThat
+import org.junit.Rule
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class CommunalDatabaseMigrationsTest : SysuiTestCase() {
+
+    @JvmField @Rule val instantTaskExecutor = InstantTaskExecutorRule()
+
+    @get:Rule
+    val migrationTestHelper =
+        MigrationTestHelper(
+            InstrumentationRegistry.getInstrumentation(),
+            CommunalDatabase::class.java.canonicalName,
+        )
+
+    @Test
+    fun migrate1To2() {
+        // Create a communal database in version 1
+        val databaseV1 = migrationTestHelper.createDatabase(DATABASE_NAME, version = 1)
+
+        // Populate some fake data
+        val fakeWidgetsV1 =
+            listOf(
+                FakeCommunalWidgetItemV1(1, "test_widget_1", 11),
+                FakeCommunalWidgetItemV1(2, "test_widget_2", 12),
+                FakeCommunalWidgetItemV1(3, "test_widget_3", 13),
+            )
+        databaseV1.insertWidgetsV1(fakeWidgetsV1)
+
+        // Verify fake widgets populated
+        databaseV1.verifyWidgetsV1(fakeWidgetsV1)
+
+        // Run migration and get database V2, the migration test helper verifies that the schema is
+        // updated correctly
+        val databaseV2 =
+            migrationTestHelper.runMigrationsAndValidate(
+                name = DATABASE_NAME,
+                version = 2,
+                validateDroppedTables = false,
+                CommunalDatabase.MIGRATION_1_2,
+            )
+
+        // Verify data is migrated correctly
+        databaseV2.verifyWidgetsV2(fakeWidgetsV1.map { it.getV2() })
+    }
+
+    private fun SupportSQLiteDatabase.insertWidgetsV1(widgets: List<FakeCommunalWidgetItemV1>) {
+        widgets.forEach { widget ->
+            execSQL(
+                "INSERT INTO communal_widget_table(widget_id, component_name, item_id) " +
+                    "VALUES(${widget.widgetId}, '${widget.componentName}', ${widget.itemId})"
+            )
+        }
+    }
+
+    private fun SupportSQLiteDatabase.verifyWidgetsV1(widgets: List<FakeCommunalWidgetItemV1>) {
+        val cursor = query("SELECT * FROM communal_widget_table")
+        assertThat(cursor.moveToFirst()).isTrue()
+
+        widgets.forEach { widget ->
+            assertThat(cursor.getInt(cursor.getColumnIndex("widget_id"))).isEqualTo(widget.widgetId)
+            assertThat(cursor.getString(cursor.getColumnIndex("component_name")))
+                .isEqualTo(widget.componentName)
+            assertThat(cursor.getInt(cursor.getColumnIndex("item_id"))).isEqualTo(widget.itemId)
+
+            cursor.moveToNext()
+        }
+
+        // Verify there is no more columns
+        assertThat(cursor.isAfterLast).isTrue()
+    }
+
+    private fun SupportSQLiteDatabase.verifyWidgetsV2(widgets: List<FakeCommunalWidgetItemV2>) {
+        val cursor = query("SELECT * FROM communal_widget_table")
+        assertThat(cursor.moveToFirst()).isTrue()
+
+        widgets.forEach { widget ->
+            assertThat(cursor.getInt(cursor.getColumnIndex("widget_id"))).isEqualTo(widget.widgetId)
+            assertThat(cursor.getString(cursor.getColumnIndex("component_name")))
+                .isEqualTo(widget.componentName)
+            assertThat(cursor.getInt(cursor.getColumnIndex("item_id"))).isEqualTo(widget.itemId)
+            assertThat(cursor.getInt(cursor.getColumnIndex("user_serial_number")))
+                .isEqualTo(widget.userSerialNumber)
+
+            cursor.moveToNext()
+        }
+
+        // Verify there is no more columns
+        assertThat(cursor.isAfterLast).isTrue()
+    }
+
+    /**
+     * Returns the expected data after migration from V1 to V2, which is simply that the new user
+     * serial number field is now set to [CommunalWidgetItem.USER_SERIAL_NUMBER_UNDEFINED].
+     */
+    private fun FakeCommunalWidgetItemV1.getV2(): FakeCommunalWidgetItemV2 {
+        return FakeCommunalWidgetItemV2(
+            widgetId,
+            componentName,
+            itemId,
+            CommunalWidgetItem.USER_SERIAL_NUMBER_UNDEFINED,
+        )
+    }
+
+    private data class FakeCommunalWidgetItemV1(
+        val widgetId: Int,
+        val componentName: String,
+        val itemId: Int,
+    )
+
+    private data class FakeCommunalWidgetItemV2(
+        val widgetId: Int,
+        val componentName: String,
+        val itemId: Int,
+        val userSerialNumber: Int,
+    )
+
+    companion object {
+        private const val DATABASE_NAME = "communal_db"
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
index f77c7a6..d670508 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
@@ -67,11 +67,12 @@
     @Test
     fun addWidget_readValueInDb() =
         testScope.runTest {
-            val (widgetId, provider, priority) = widgetInfo1
+            val (widgetId, provider, priority, userSerialNumber) = widgetInfo1
             communalWidgetDao.addWidget(
                 widgetId = widgetId,
                 provider = provider,
                 priority = priority,
+                userSerialNumber = userSerialNumber,
             )
             val entry = communalWidgetDao.getWidgetByIdNow(id = 1)
             assertThat(entry).isEqualTo(communalWidgetItemEntry1)
@@ -80,11 +81,12 @@
     @Test
     fun deleteWidget_notInDb_returnsFalse() =
         testScope.runTest {
-            val (widgetId, provider, priority) = widgetInfo1
+            val (widgetId, provider, priority, userSerialNumber) = widgetInfo1
             communalWidgetDao.addWidget(
                 widgetId = widgetId,
                 provider = provider,
                 priority = priority,
+                userSerialNumber = userSerialNumber,
             )
             assertThat(communalWidgetDao.deleteWidgetById(widgetId = 123)).isFalse()
         }
@@ -95,11 +97,12 @@
             val widgetsToAdd = listOf(widgetInfo1, widgetInfo2)
             val widgets = collectLastValue(communalWidgetDao.getWidgets())
             widgetsToAdd.forEach {
-                val (widgetId, provider, priority) = it
+                val (widgetId, provider, priority, userSerialNumber) = it
                 communalWidgetDao.addWidget(
                     widgetId = widgetId,
                     provider = provider,
                     priority = priority,
+                    userSerialNumber = userSerialNumber
                 )
             }
             assertThat(widgets())
@@ -118,11 +121,12 @@
             val widgets = collectLastValue(communalWidgetDao.getWidgets())
 
             widgetsToAdd.forEach {
-                val (widgetId, provider, priority) = it
+                val (widgetId, provider, priority, userSerialNumber) = it
                 communalWidgetDao.addWidget(
                     widgetId = widgetId,
                     provider = provider,
                     priority = priority,
+                    userSerialNumber = userSerialNumber,
                 )
             }
             assertThat(widgets())
@@ -144,11 +148,12 @@
             val widgets = collectLastValue(communalWidgetDao.getWidgets())
 
             widgetsToAdd.forEach {
-                val (widgetId, provider, priority) = it
+                val (widgetId, provider, priority, userSerialNumber) = it
                 communalWidgetDao.addWidget(
                     widgetId = widgetId,
                     provider = provider,
                     priority = priority,
+                    userSerialNumber = userSerialNumber,
                 )
             }
             assertThat(widgets())
@@ -180,11 +185,12 @@
             val widgets = collectLastValue(communalWidgetDao.getWidgets())
 
             existingWidgets.forEach {
-                val (widgetId, provider, priority) = it
+                val (widgetId, provider, priority, userSerialNumber) = it
                 communalWidgetDao.addWidget(
                     widgetId = widgetId,
                     provider = provider,
                     priority = priority,
+                    userSerialNumber = userSerialNumber,
                 )
             }
             assertThat(widgets())
@@ -212,6 +218,7 @@
                 widgetId = widgetInfo3.widgetId,
                 provider = widgetInfo3.provider,
                 priority = 2,
+                userSerialNumber = widgetInfo3.userSerialNumber,
             )
             assertThat(widgets())
                 .containsExactly(
@@ -246,6 +253,7 @@
                         widgetId = fakeWidget.widgetId,
                         componentName = fakeWidget.componentName,
                         itemId = rank.uid,
+                        userSerialNumber = fakeWidget.userSerialNumber,
                     )
                 expected[rank] = widget
             }
@@ -258,13 +266,15 @@
             widgetId = metadata.widgetId,
             provider = metadata.provider,
             priority = priority ?: metadata.priority,
+            userSerialNumber = metadata.userSerialNumber,
         )
     }
 
     data class FakeWidgetMetadata(
         val widgetId: Int,
         val provider: ComponentName,
-        val priority: Int
+        val priority: Int,
+        val userSerialNumber: Int,
     )
 
     companion object {
@@ -272,19 +282,22 @@
             FakeWidgetMetadata(
                 widgetId = 1,
                 provider = ComponentName("pk_name", "cls_name_1"),
-                priority = 1
+                priority = 1,
+                userSerialNumber = 0,
             )
         val widgetInfo2 =
             FakeWidgetMetadata(
                 widgetId = 2,
                 provider = ComponentName("pk_name", "cls_name_2"),
-                priority = 2
+                priority = 2,
+                userSerialNumber = 0,
             )
         val widgetInfo3 =
             FakeWidgetMetadata(
                 widgetId = 3,
                 provider = ComponentName("pk_name", "cls_name_3"),
-                priority = 3
+                priority = 3,
+                userSerialNumber = 10,
             )
         val communalItemRankEntry1 = CommunalItemRank(uid = 1L, rank = widgetInfo1.priority)
         val communalItemRankEntry2 = CommunalItemRank(uid = 2L, rank = widgetInfo2.priority)
@@ -295,6 +308,7 @@
                 widgetId = widgetInfo1.widgetId,
                 componentName = widgetInfo1.provider.flattenToString(),
                 itemId = communalItemRankEntry1.uid,
+                userSerialNumber = widgetInfo1.userSerialNumber,
             )
         val communalWidgetItemEntry2 =
             CommunalWidgetItem(
@@ -302,6 +316,7 @@
                 widgetId = widgetInfo2.widgetId,
                 componentName = widgetInfo2.provider.flattenToString(),
                 itemId = communalItemRankEntry2.uid,
+                userSerialNumber = widgetInfo2.userSerialNumber,
             )
         val communalWidgetItemEntry3 =
             CommunalWidgetItem(
@@ -309,6 +324,7 @@
                 widgetId = widgetInfo3.widgetId,
                 componentName = widgetInfo3.provider.flattenToString(),
                 itemId = communalItemRankEntry3.uid,
+                userSerialNumber = widgetInfo3.userSerialNumber,
             )
         val fakeState =
             CommunalHubState().apply {
@@ -318,11 +334,13 @@
                                 widgetId = 1
                                 componentName = "pk_name/fake_widget_1"
                                 rank = 1
+                                userSerialNumber = 0
                             },
                             CommunalHubState.CommunalWidgetItem().apply {
                                 widgetId = 2
                                 componentName = "pk_name/fake_widget_2"
                                 rank = 2
+                                userSerialNumber = 10
                             },
                         )
                         .toTypedArray()