Adds new method to update existing documents in MtpDatabase.

BUG=25162822

Change-Id: I7aa63fc272aa7b57d6a9672565f842774e898a00

2 files changed, 367 insertions, 83 deletions

diff --git a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
index 0f31e2c39a1a..0ee6460d8402 100644
--- a/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
+++ b/packages/MtpDocumentsProvider/src/com/android/mtp/MtpDatabase.java
@@ -27,10 +27,11 @@ import android.database.sqlite.SQLiteQueryBuilder;
 import android.mtp.MtpObjectInfo;
 import android.provider.DocumentsContract.Document;
 import android.provider.DocumentsContract.Root;
-import android.util.Log;
 
 import com.android.internal.annotations.VisibleForTesting;
 
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 
 /**
@@ -46,6 +47,27 @@ import java.util.Objects;
  * remembers the map of document ID and object handle, and remaps new object handle with document ID
  * by comparing the directory structure and object name.
  *
+ * To start putting documents into the database, the client needs to call
+ * {@link #startAddingChildDocuments(String)} with the parent document ID. Also it needs to call
+ * {@link #stopAddingChildDocuments(String)} after putting all child documents to the database.
+ * (All explanations are same for root documents)
+ *
+ * database.startAddingChildDocuments();
+ * database.putChildDocuments();
+ * database.stopAddingChildDocuments();
+ *
+ * To update the existing documents, the client code can repeat to call the three methods again.
+ * The newly added rows update corresponding existing rows that have same MTP identifier like
+ * objectHandle.
+ *
+ * The client can call putChildDocuments multiple times to add documents by chunk, but it needs to
+ * put all documents under the parent before calling stopAddingChildDocuments. Otherwise missing
+ * documents are regarded as deleted, and will be removed from the database.
+ *
+ * If the client calls clearMtpIdentifier(), it clears MTP identifier in the database. In this case,
+ * the database tries to find corresponding rows by using document's name instead of MTP identifier
+ * at the next update cycle.
+ *
  * TODO: Remove @VisibleForTesting annotation when we start to use this class.
  * TODO: Improve performance by SQL optimization.
  */
@@ -79,22 +101,34 @@ class MtpDatabase {
     /**
      * The state represents that the row has a valid object handle.
      */
-    static final int ROW_STATE_MAPPED = 0;
+    static final int ROW_STATE_VALID = 0;
+
+    /**
+     * The state represents that the rows added at the previous cycle and need to be updated with
+     * fresh values.
+     * The row may not have valid object handle. External application can still fetch the documents.
+     * If the external application tries to fetch object handle, the provider resolves pending
+     * documents with invalidated documents ahead.
+     */
+    static final int ROW_STATE_INVALIDATED = 1;
+
+    /**
+     * The state represents the raw has a valid object handle but it may be going to be mapped with
+     * another rows invalidated. After fetching all documents under the parent, the database tries
+     * to map the pending documents and the invalidated documents in order to keep old document ID
+     * alive.
+     */
+    static final int ROW_STATE_PENDING = 2;
 
     /**
-     * The state represents that the object handle was cleared because the MTP session closed.
-     * External application can still fetch the unmapped documents. If the external application
-     * tries to open an unmapped document, the provider resolves the document with new object handle
-     * ahead.
+     * Mapping mode that uses MTP identifier to find corresponding rows.
      */
-    static final int ROW_STATE_UNMAPPED = 1;
+    static final int MAP_BY_MTP_IDENTIFIER = 0;
 
     /**
-     * The state represents the raw has a valid object handle but it may be going to be merged into
-     * another unmapped row. After fetching all documents under the parent, the database tries to
-     * map the mapping document and the unmapped document in order to keep old document ID alive.
+     * Mapping mode that uses name to find corresponding rows.
      */
-    static final int ROW_STATE_MAPPING = 2;
+    static final int MAP_BY_NAME = 1;
 
     private static final String SELECTION_DOCUMENT_ID = Document.COLUMN_DOCUMENT_ID + " = ?";
     private static final String SELECTION_ROOT_ID = Root.COLUMN_ROOT_ID + " = ?";
@@ -175,6 +209,7 @@ class MtpDatabase {
         }
     }
 
+    private final Map<String, Integer> mMappingMode = new HashMap<>();
     private final SQLiteDatabase mDatabase;
 
     @VisibleForTesting
@@ -194,7 +229,7 @@ class MtpDatabase {
                 VIEW_ROOTS,
                 columnNames,
                 COLUMN_ROW_STATE + " IN (?, ?)",
-                strings(ROW_STATE_MAPPED, ROW_STATE_UNMAPPED),
+                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED),
                 null,
                 null,
                 null);
@@ -206,7 +241,7 @@ class MtpDatabase {
                 TABLE_DOCUMENTS,
                 columnNames,
                 COLUMN_ROW_STATE + " IN (?, ?)",
-                strings(ROW_STATE_MAPPED, ROW_STATE_UNMAPPED),
+                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED),
                 null,
                 null,
                 null);
@@ -218,13 +253,71 @@ class MtpDatabase {
                 TABLE_DOCUMENTS,
                 columnNames,
                 COLUMN_ROW_STATE + " IN (?, ?) AND " + COLUMN_PARENT_DOCUMENT_ID + " = ?",
-                strings(ROW_STATE_MAPPED, ROW_STATE_UNMAPPED, parentDocumentId),
+                strings(ROW_STATE_VALID, ROW_STATE_INVALIDATED, parentDocumentId),
                 null,
                 null,
                 null);
     }
 
     @VisibleForTesting
+    void startAddingRootDocuments(int deviceId) {
+        final String mappingStateKey = getRootDocumentsMappingStateKey(deviceId);
+        if (mMappingMode.containsKey(mappingStateKey)) {
+            throw new Error("Mapping for the root has already started.");
+        }
+        mMappingMode.put(
+                mappingStateKey,
+                startAddingDocuments(SELECTION_ROOT_DOCUMENTS, Integer.toString(deviceId)));
+    }
+
+    @VisibleForTesting
+    void startAddingChildDocuments(String parentDocumentId) {
+        final String mappingStateKey = getChildDocumentsMappingStateKey(parentDocumentId);
+        if (mMappingMode.containsKey(mappingStateKey)) {
+            throw new Error("Mapping for the root has already started.");
+        }
+        mMappingMode.put(
+                mappingStateKey,
+                startAddingDocuments(SELECTION_CHILD_DOCUMENTS, parentDocumentId));
+    }
+
+    /**
+     * Starts adding new documents.
+     * The methods decides mapping mode depends on if all documents under the given parent have MTP
+     * identifier or not. If all the documents have MTP identifier, it uses the identifier to find
+     * a corresponding existing row. Otherwise it does heuristic.
+     *
+     * @param selection Query matches valid documents.
+     * @param arg Argument for selection.
+     * @return Mapping mode.
+     */
+    @VisibleForTesting
+    private int startAddingDocuments(String selection, String arg) {
+        mDatabase.beginTransaction();
+        try {
+            // Delete all pending rows.
+            deleteDocumentsAndRoots(
+                    selection + " AND " + COLUMN_ROW_STATE + "=?", strings(arg, ROW_STATE_PENDING));
+
+            // Set all documents as invalidated.
+            final ContentValues values = new ContentValues();
+            values.put(COLUMN_ROW_STATE, ROW_STATE_INVALIDATED);
+            mDatabase.update(TABLE_DOCUMENTS, values, selection, new String[] { arg });
+
+            // If we have rows that does not have MTP identifier, do heuristic mapping by name.
+            final boolean useNameForResolving = DatabaseUtils.queryNumEntries(
+                    mDatabase,
+                    TABLE_DOCUMENTS,
+                    selection + " AND " + COLUMN_STORAGE_ID + " IS NULL",
+                    new String[] { arg }) > 0;
+            mDatabase.setTransactionSuccessful();
+            return useNameForResolving ? MAP_BY_NAME : MAP_BY_MTP_IDENTIFIER;
+        } finally {
+            mDatabase.endTransaction();
+        }
+    }
+
+    @VisibleForTesting
     void putRootDocuments(int deviceId, Resources resources, MtpRoot[] roots) {
         mDatabase.beginTransaction();
         try {
@@ -236,20 +329,37 @@ class MtpDatabase {
                 valuesList[i] = new ContentValues();
                 getRootDocumentValues(valuesList[i], resources, roots[i]);
             }
-            final long[] documentIds =
-                    putDocuments(valuesList, SELECTION_ROOT_DOCUMENTS, Integer.toString(deviceId));
+            boolean heuristic;
+            String mapColumn;
+            switch (mMappingMode.get(getRootDocumentsMappingStateKey(deviceId))) {
+                case MAP_BY_MTP_IDENTIFIER:
+                    heuristic = false;
+                    mapColumn = COLUMN_STORAGE_ID;
+                    break;
+                case MAP_BY_NAME:
+                    heuristic = true;
+                    mapColumn = Document.COLUMN_DISPLAY_NAME;
+                    break;
+                default:
+                    throw new Error("Unexpected map mode.");
+            }
+            final long[] documentIds = putDocuments(
+                    valuesList,
+                    SELECTION_ROOT_DOCUMENTS,
+                    Integer.toString(deviceId),
+                    heuristic,
+                    mapColumn);
             final ContentValues values = new ContentValues();
             int i = 0;
             for (final MtpRoot root : roots) {
                 // Use the same value for the root ID and the corresponding document ID.
                 values.put(Root.COLUMN_ROOT_ID, documentIds[i++]);
-                values.put(Root.COLUMN_FLAGS,
-                        Root.FLAG_SUPPORTS_IS_CHILD |
-                        Root.FLAG_SUPPORTS_CREATE);
+                values.put(
+                        Root.COLUMN_FLAGS, Root.FLAG_SUPPORTS_IS_CHILD | Root.FLAG_SUPPORTS_CREATE);
                 values.put(Root.COLUMN_AVAILABLE_BYTES, root.mFreeSpace);
                 values.put(Root.COLUMN_CAPACITY_BYTES, root.mMaxCapacity);
                 values.put(Root.COLUMN_MIME_TYPES, "");
-                mDatabase.insert(TABLE_ROOT_EXTRA, null, values);
+                mDatabase.replace(TABLE_ROOT_EXTRA, null, values);
             }
             mDatabase.setTransactionSuccessful();
         } finally {
@@ -264,71 +374,142 @@ class MtpDatabase {
             valuesList[i] = new ContentValues();
             getChildDocumentValues(valuesList[i], deviceId, parentId, documents[i]);
         }
-        putDocuments(valuesList, SELECTION_CHILD_DOCUMENTS, parentId);
+        boolean heuristic;
+        String mapColumn;
+        switch (mMappingMode.get(getChildDocumentsMappingStateKey(parentId))) {
+            case MAP_BY_MTP_IDENTIFIER:
+                heuristic = false;
+                mapColumn = COLUMN_STORAGE_ID;
+                break;
+            case MAP_BY_NAME:
+                heuristic = true;
+                mapColumn = Document.COLUMN_DISPLAY_NAME;
+                break;
+            default:
+                throw new Error("Unexpected map mode.");
+        }
+        putDocuments(valuesList, SELECTION_CHILD_DOCUMENTS, parentId, heuristic, mapColumn);
     }
 
     /**
      * Clears MTP related identifier.
      * It clears MTP's object handle and storage ID that are not stable over MTP sessions and mark
-     * the all documents as 'unmapped'. It also remove 'mapping' rows as mapping is cancelled now.
+     * the all documents as 'invalidated'. It also remove 'pending' rows as adding is cancelled
+     * now.
      */
     @VisibleForTesting
     void clearMapping() {
         mDatabase.beginTransaction();
         try {
-            deleteDocumentsAndRoots(COLUMN_ROW_STATE + " = ?", strings(ROW_STATE_MAPPING));
+            deleteDocumentsAndRoots(COLUMN_ROW_STATE + " = ?", strings(ROW_STATE_PENDING));
             final ContentValues values = new ContentValues();
             values.putNull(COLUMN_OBJECT_HANDLE);
             values.putNull(COLUMN_STORAGE_ID);
-            values.put(COLUMN_ROW_STATE, ROW_STATE_UNMAPPED);
+            values.put(COLUMN_ROW_STATE, ROW_STATE_INVALIDATED);
             mDatabase.update(TABLE_DOCUMENTS, values, null, null);
             mDatabase.setTransactionSuccessful();
+            mMappingMode.clear();
         } finally {
             mDatabase.endTransaction();
         }
     }
 
     @VisibleForTesting
-    void resolveRootDocuments(int deviceId) {
-        resolveDocuments(SELECTION_ROOT_DOCUMENTS, Integer.toString(deviceId));
+    void stopAddingRootDocuments(int deviceId) {
+        final String mappingModeKey = getRootDocumentsMappingStateKey(deviceId);
+        switch (mMappingMode.get(mappingModeKey)) {
+            case MAP_BY_MTP_IDENTIFIER:
+                stopAddingDocuments(
+                        SELECTION_ROOT_DOCUMENTS,
+                        Integer.toString(deviceId),
+                        COLUMN_STORAGE_ID);
+                break;
+            case MAP_BY_NAME:
+                stopAddingDocuments(
+                        SELECTION_ROOT_DOCUMENTS,
+                        Integer.toString(deviceId),
+                        Document.COLUMN_DISPLAY_NAME);
+                break;
+            default:
+                throw new Error("Unexpected mapping state.");
+        }
+        mMappingMode.remove(mappingModeKey);
     }
 
     @VisibleForTesting
-    void resolveChildDocuments(String parentId) {
-        resolveDocuments(SELECTION_CHILD_DOCUMENTS, parentId);
+    void stopAddingChildDocuments(String parentId) {
+        final String mappingModeKey = getChildDocumentsMappingStateKey(parentId);
+        switch (mMappingMode.get(mappingModeKey)) {
+            case MAP_BY_MTP_IDENTIFIER:
+                stopAddingDocuments(
+                        SELECTION_CHILD_DOCUMENTS,
+                        parentId,
+                        COLUMN_OBJECT_HANDLE);
+                break;
+            case MAP_BY_NAME:
+                stopAddingDocuments(
+                        SELECTION_CHILD_DOCUMENTS,
+                        parentId,
+                        Document.COLUMN_DISPLAY_NAME);
+                break;
+            default:
+                throw new Error("Unexpected mapping state.");
+        }
+        mMappingMode.remove(mappingModeKey);
     }
 
     /**
      * Puts the documents into the database.
-     * If the database found another unmapped document that shares the same name and parent,
-     * the document may be merged into the unmapped document. In that case, the database marks the
-     * root as 'mapping' and wait for {@link #resolveRootDocuments(int)} is invoked.
+     * If the mapping mode is not heuristic, it just adds the rows to the database or updates the
+     * existing rows with the new values. If the mapping mode is heuristic, it adds some new rows as
+     * 'pending' state when that rows may be corresponding to existing 'invalidated' rows. Then
+     * {@link #stopAddingDocuments(String, String, String)} turns the pending rows into 'valid'
+     * rows.
+     *
      * @param valuesList Values that are stored in the database.
      * @param selection SQL where closure to select rows that shares the same parent.
      * @param arg Argument for selection SQL.
+     * @param heuristic Whether the mapping mode is heuristic.
      * @return List of Document ID inserted to the table.
      */
-    private long[] putDocuments(ContentValues[] valuesList, String selection, String arg) {
+    private long[] putDocuments(
+            ContentValues[] valuesList,
+            String selection,
+            String arg,
+            boolean heuristic,
+            String mappingKey) {
         mDatabase.beginTransaction();
         try {
             final long[] documentIds = new long[valuesList.length];
             int i = 0;
             for (final ContentValues values : valuesList) {
-                final String displayName =
-                        values.getAsString(Document.COLUMN_DISPLAY_NAME);
-                final long numUnmapped = DatabaseUtils.queryNumEntries(
-                        mDatabase,
+                final Cursor candidateCursor = mDatabase.query(
                         TABLE_DOCUMENTS,
+                        strings(Document.COLUMN_DOCUMENT_ID),
                         selection + " AND " +
-                        COLUMN_ROW_STATE + " = ? AND " +
-                        Document.COLUMN_DISPLAY_NAME + " = ?",
-                        strings(arg, ROW_STATE_UNMAPPED, displayName));
-                if (numUnmapped != 0) {
-                    values.put(COLUMN_ROW_STATE, ROW_STATE_MAPPING);
+                        COLUMN_ROW_STATE + "=? AND " +
+                        mappingKey + "=?",
+                        strings(arg, ROW_STATE_INVALIDATED, values.getAsString(mappingKey)),
+                        null,
+                        null,
+                        null,
+                        "1");
+                final long rowId;
+                if (candidateCursor.getCount() == 0) {
+                    rowId = mDatabase.insert(TABLE_DOCUMENTS, null, values);
+                } else if (!heuristic) {
+                    candidateCursor.moveToNext();
+                    final String documentId = candidateCursor.getString(0);
+                    rowId = mDatabase.update(
+                            TABLE_DOCUMENTS, values, SELECTION_DOCUMENT_ID, strings(documentId));
+                } else {
+                    values.put(COLUMN_ROW_STATE, ROW_STATE_PENDING);
+                    rowId = mDatabase.insert(TABLE_DOCUMENTS, null, values);
                 }
-                // Document ID is a primary integer key of the table. So the returned row IDs should
-                // be same with the document ID.
-                documentIds[i++] = mDatabase.insert(TABLE_DOCUMENTS, null, values);
+                // Document ID is a primary integer key of the table. So the returned row
+                // IDs should be same with the document ID.
+                documentIds[i++] = rowId;
+                candidateCursor.close();
             }
 
             mDatabase.setTransactionSuccessful();
@@ -339,21 +520,21 @@ class MtpDatabase {
     }
 
     /**
-     * Maps 'unmapped' document and 'mapping' document that don't have document but shares the same
-     * name.
-     * If the database does not find corresponding 'mapping' document, it just removes 'unmapped'
-     * document from the database.
+     * Maps 'pending' document and 'invalidated' document that shares the same column of groupKey.
+     * If the database does not find corresponding 'invalidated' document, it just removes
+     * 'invalidated' document from the database.
      * @param selection Query to select rows for resolving.
      * @param arg Argument for selection SQL.
+     * @param groupKey Column name used to find corresponding rows.
      */
-    private void resolveDocuments(String selection, String arg) {
+    private void stopAddingDocuments(String selection, String arg, String groupKey) {
         mDatabase.beginTransaction();
         try {
-            // Get 1-to-1 mapping of unmapped document and mapping document.
-            final String unmappedIdQuery = createStateFilter(
-                    ROW_STATE_UNMAPPED, Document.COLUMN_DOCUMENT_ID);
-            final String mappingIdQuery = createStateFilter(
-                    ROW_STATE_MAPPING, Document.COLUMN_DOCUMENT_ID);
+            // Get 1-to-1 mapping of invalidated document and pending document.
+            final String invalidatedIdQuery = createStateFilter(
+                    ROW_STATE_INVALIDATED, Document.COLUMN_DOCUMENT_ID);
+            final String pendingIdQuery = createStateFilter(
+                    ROW_STATE_PENDING, Document.COLUMN_DOCUMENT_ID);
             // SQL should be like:
             // SELECT group_concat(CASE WHEN raw_state = 1 THEN document_id ELSE NULL END),
             //        group_concat(CASE WHEN raw_state = 2 THEN document_id ELSE NULL END)
@@ -364,38 +545,38 @@ class MtpDatabase {
             final Cursor mergingCursor = mDatabase.query(
                     TABLE_DOCUMENTS,
                     new String[] {
-                            "group_concat(" + unmappedIdQuery + ")",
-                            "group_concat(" + mappingIdQuery + ")"
+                            "group_concat(" + invalidatedIdQuery + ")",
+                            "group_concat(" + pendingIdQuery + ")"
                     },
                     selection,
                     strings(arg),
-                    Document.COLUMN_DISPLAY_NAME,
-                    "count(" + unmappedIdQuery + ") = 1 AND count(" + mappingIdQuery + ") = 1",
+                    groupKey,
+                    "count(" + invalidatedIdQuery + ") = 1 AND count(" + pendingIdQuery + ") = 1",
                     null);
 
             final ContentValues values = new ContentValues();
             while (mergingCursor.moveToNext()) {
-                final String unmappedId = mergingCursor.getString(0);
-                final String mappingId = mergingCursor.getString(1);
+                final String invalidatedId = mergingCursor.getString(0);
+                final String pendingId = mergingCursor.getString(1);
 
                 // Obtain the new values including the latest object handle from mapping row.
                 getFirstRow(
                         TABLE_DOCUMENTS,
                         SELECTION_DOCUMENT_ID,
-                        new String[] { mappingId },
+                        new String[] { pendingId },
                         values);
                 values.remove(Document.COLUMN_DOCUMENT_ID);
-                values.put(COLUMN_ROW_STATE, ROW_STATE_MAPPED);
+                values.put(COLUMN_ROW_STATE, ROW_STATE_VALID);
                 mDatabase.update(
                         TABLE_DOCUMENTS,
                         values,
                         SELECTION_DOCUMENT_ID,
-                        new String[] { unmappedId });
+                        new String[] { invalidatedId });
 
                 getFirstRow(
                         TABLE_ROOT_EXTRA,
                         SELECTION_ROOT_ID,
-                        new String[] { mappingId },
+                        new String[] { pendingId },
                         values);
                 if (values.size() > 0) {
                     values.remove(Root.COLUMN_ROOT_ID);
@@ -403,29 +584,29 @@ class MtpDatabase {
                             TABLE_ROOT_EXTRA,
                             values,
                             SELECTION_ROOT_ID,
-                            new String[] { unmappedId });
+                            new String[] { invalidatedId });
                 }
 
-                // Delete 'mapping' row.
-                deleteDocumentsAndRoots(SELECTION_DOCUMENT_ID, new String[] { mappingId });
+                // Delete 'pending' row.
+                deleteDocumentsAndRoots(SELECTION_DOCUMENT_ID, new String[] { pendingId });
             }
             mergingCursor.close();
 
-            // Delete all unmapped rows that cannot be mapped.
+            // Delete all invalidated rows that cannot be mapped.
             deleteDocumentsAndRoots(
                     COLUMN_ROW_STATE + " = ? AND " + selection,
-                    strings(ROW_STATE_UNMAPPED, arg));
+                    strings(ROW_STATE_INVALIDATED, arg));
 
-            // The database cannot find old document ID for the mapping rows.
-            // Turn the all mapping rows into mapped state, which means the rows become to be
+            // The database cannot find old document ID for the pending rows.
+            // Turn the all pending rows into valid state, which means the rows become to be
             // valid with new document ID.
             values.clear();
-            values.put(COLUMN_ROW_STATE, ROW_STATE_MAPPED);
+            values.put(COLUMN_ROW_STATE, ROW_STATE_VALID);
             mDatabase.update(
                     TABLE_DOCUMENTS,
                     values,
                     COLUMN_ROW_STATE + " = ? AND " + selection,
-                    strings(ROW_STATE_MAPPING, arg));
+                    strings(ROW_STATE_PENDING, arg));
             mDatabase.setTransactionSuccessful();
         } finally {
             mDatabase.endTransaction();
@@ -445,7 +626,7 @@ class MtpDatabase {
         values.put(COLUMN_STORAGE_ID, root.mStorageId);
         values.putNull(COLUMN_OBJECT_HANDLE);
         values.putNull(COLUMN_PARENT_DOCUMENT_ID);
-        values.put(COLUMN_ROW_STATE, ROW_STATE_MAPPED);
+        values.put(COLUMN_ROW_STATE, ROW_STATE_VALID);
         values.put(Document.COLUMN_MIME_TYPE, Document.MIME_TYPE_DIR);
         values.put(Document.COLUMN_DISPLAY_NAME, root.getRootName(resources));
         values.putNull(Document.COLUMN_SUMMARY);
@@ -482,7 +663,7 @@ class MtpDatabase {
         values.put(COLUMN_STORAGE_ID, info.getStorageId());
         values.put(COLUMN_OBJECT_HANDLE, info.getObjectHandle());
         values.put(COLUMN_PARENT_DOCUMENT_ID, parentId);
-        values.put(COLUMN_ROW_STATE, ROW_STATE_MAPPED);
+        values.put(COLUMN_ROW_STATE, ROW_STATE_VALID);
         values.put(Document.COLUMN_MIME_TYPE, mimeType);
         values.put(Document.COLUMN_DISPLAY_NAME, info.getName());
         values.putNull(Document.COLUMN_SUMMARY);
@@ -539,6 +720,14 @@ class MtpDatabase {
         }
     }
 
+    private String getRootDocumentsMappingStateKey(int deviceId) {
+        return "RootDocuments/" + deviceId;
+    }
+
+    private String getChildDocumentsMappingStateKey(String parentDocumentId) {
+        return "ChildDocuments/" + parentDocumentId;
+    }
+
     /**
      * Converts values into string array.
      * @param args Values converted into string array.
diff --git a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
index 3878ba637d7c..90999b1c11d4 100644
--- a/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
+++ b/packages/MtpDocumentsProvider/tests/src/com/android/mtp/MtpDatabaseTest.java
@@ -23,6 +23,7 @@ import android.provider.DocumentsContract;
 import android.provider.DocumentsContract.Root;
 import android.test.AndroidTestCase;
 import android.test.suitebuilder.annotation.SmallTest;
+import android.util.Log;
 
 @SmallTest
 public class MtpDatabaseTest extends AndroidTestCase {
@@ -49,6 +50,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
 
     public void testPutRootDocuments() throws Exception {
         final MtpDatabase database = new MtpDatabase(getContext());
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 1, "Device", "Storage", 1000, 2000, ""),
                 new MtpRoot(0, 2, "Device", "Storage", 2000, 4000, ""),
@@ -141,7 +143,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
 
     public void testPutChildDocuments() throws Exception {
         final MtpDatabase database = new MtpDatabase(getContext());
-
+        database.startAddingChildDocuments("parentId");
         database.putChildDocuments(0, "parentId", new MtpObjectInfo[] {
                 createDocument(100, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
                 createDocument(101, "image.jpg", MtpConstants.FORMAT_EXIF_JPEG, 2 * 1024 * 1024),
@@ -216,6 +218,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_ROOT_ID,
                 Root.COLUMN_AVAILABLE_BYTES
         };
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage A", 1000, 0, ""),
                 new MtpRoot(0, 101, "Device", "Storage B", 1001, 0, "")
@@ -275,6 +278,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage A", 2000, 0, ""),
                 new MtpRoot(0, 202, "Device", "Storage C", 2002, 0, "")
@@ -313,7 +317,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        database.resolveRootDocuments(0);
+        database.stopAddingRootDocuments(0);
 
         {
             final Cursor cursor = database.queryRootDocuments(columns);
@@ -349,6 +353,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 MtpDatabase.COLUMN_OBJECT_HANDLE,
                 DocumentsContract.Document.COLUMN_DISPLAY_NAME
         };
+        database.startAddingChildDocuments("parentId");
         database.putChildDocuments(0, "parentId", new MtpObjectInfo[] {
                 createDocument(100, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
                 createDocument(101, "image.jpg", MtpConstants.FORMAT_EXIF_JPEG, 2 * 1024 * 1024),
@@ -378,6 +383,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
+        database.startAddingChildDocuments("parentId");
         database.putChildDocuments(0, "parentId", new MtpObjectInfo[] {
                 createDocument(200, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
                 createDocument(203, "video.mp4", MtpConstants.FORMAT_MP4_CONTAINER, 1024),
@@ -395,7 +401,7 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
 
-        database.resolveChildDocuments("parentId");
+        database.stopAddingChildDocuments("parentId");
 
         {
             final Cursor cursor = database.queryChildDocuments(columns, "parentId");
@@ -425,6 +431,8 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_ROOT_ID,
                 Root.COLUMN_AVAILABLE_BYTES
         };
+        database.startAddingRootDocuments(0);
+        database.startAddingRootDocuments(1);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, "")
         });
@@ -460,14 +468,16 @@ public class MtpDatabaseTest extends AndroidTestCase {
 
         database.clearMapping();
 
+        database.startAddingRootDocuments(0);
+        database.startAddingRootDocuments(1);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, "")
         });
         database.putRootDocuments(1, resources, new MtpRoot[] {
                 new MtpRoot(1, 300, "Device", "Storage", 3000, 0, "")
         });
-        database.resolveRootDocuments(0);
-        database.resolveRootDocuments(1);
+        database.stopAddingRootDocuments(0);
+        database.stopAddingRootDocuments(1);
 
         {
             final Cursor cursor = database.queryRootDocuments(columns);
@@ -502,6 +512,9 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 DocumentsContract.Document.COLUMN_DOCUMENT_ID,
                 MtpDatabase.COLUMN_OBJECT_HANDLE
         };
+
+        database.startAddingChildDocuments("parentId1");
+        database.startAddingChildDocuments("parentId2");
         database.putChildDocuments(0, "parentId1", new MtpObjectInfo[] {
                 createDocument(100, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
         });
@@ -509,13 +522,16 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 createDocument(101, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
         });
         database.clearMapping();
+
+        database.startAddingChildDocuments("parentId1");
+        database.startAddingChildDocuments("parentId2");
         database.putChildDocuments(0, "parentId1", new MtpObjectInfo[] {
                 createDocument(200, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
         });
         database.putChildDocuments(0, "parentId2", new MtpObjectInfo[] {
                 createDocument(201, "note.txt", MtpConstants.FORMAT_TEXT, 1024),
         });
-        database.resolveChildDocuments("parentId1");
+        database.stopAddingChildDocuments("parentId1");
 
         {
             final Cursor cursor = database.queryChildDocuments(columns, "parentId1");
@@ -546,18 +562,25 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_ROOT_ID,
                 Root.COLUMN_AVAILABLE_BYTES
         };
+
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, ""),
         });
         database.clearMapping();
+
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, ""),
         });
         database.clearMapping();
+
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 300, "Device", "Storage", 3000, 0, ""),
         });
-        database.resolveRootDocuments(0);
+        database.stopAddingRootDocuments(0);
+
         {
             final Cursor cursor = database.queryRootDocuments(columns);
             assertEquals(1, cursor.getCount());
@@ -588,15 +611,20 @@ public class MtpDatabaseTest extends AndroidTestCase {
                 Root.COLUMN_ROOT_ID,
                 Root.COLUMN_AVAILABLE_BYTES
         };
+
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 100, "Device", "Storage", 0, 0, ""),
         });
         database.clearMapping();
+
+        database.startAddingRootDocuments(0);
         database.putRootDocuments(0, resources, new MtpRoot[] {
                 new MtpRoot(0, 200, "Device", "Storage", 2000, 0, ""),
                 new MtpRoot(0, 201, "Device", "Storage", 2001, 0, ""),
         });
-        database.resolveRootDocuments(0);
+        database.stopAddingRootDocuments(0);
+
         {
             final Cursor cursor = database.queryRootDocuments(columns);
             assertEquals(2, cursor.getCount());
@@ -622,4 +650,71 @@ public class MtpDatabaseTest extends AndroidTestCase {
             cursor.close();
         }
     }
+
+    public void testReplaceExistingRoots() {
+        // The client code should be able to replace exisitng rows with new information.
+        final MtpDatabase database = new MtpDatabase(getContext());
+        // Add one.
+        database.startAddingRootDocuments(0);
+        database.putRootDocuments(0, resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
+        });
+        database.stopAddingRootDocuments(0);
+        // Replace it.
+        database.startAddingRootDocuments(0);
+        database.putRootDocuments(0, resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device", "Storage B", 1000, 1000, ""),
+        });
+        database.stopAddingRootDocuments(0);
+        {
+            final String[] columns = new String[] {
+                    DocumentsContract.Document.COLUMN_DOCUMENT_ID,
+                    MtpDatabase.COLUMN_STORAGE_ID,
+                    DocumentsContract.Document.COLUMN_DISPLAY_NAME
+            };
+            final Cursor cursor = database.queryRootDocuments(columns);
+            assertEquals(1, cursor.getCount());
+            cursor.moveToNext();
+            assertEquals("documentId", 1, cursor.getInt(0));
+            assertEquals("storageId", 100, cursor.getInt(1));
+            assertEquals("name", "Device Storage B", cursor.getString(2));
+            cursor.close();
+        }
+        {
+            final String[] columns = new String[] {
+                    Root.COLUMN_ROOT_ID,
+                    Root.COLUMN_AVAILABLE_BYTES
+            };
+            final Cursor cursor = database.queryRoots(columns);
+            assertEquals(1, cursor.getCount());
+            cursor.moveToNext();
+            assertEquals("rootId", 1, cursor.getInt(0));
+            assertEquals("availableBytes", 1000, cursor.getInt(1));
+            cursor.close();
+        }
+    }
+
+    public void _testFailToReplaceExisitingUnmappedRoots() {
+        // The client code should not be able to replace rows before resolving 'unmapped' rows.
+        final MtpDatabase database = new MtpDatabase(getContext());
+        // Add one.
+        database.startAddingRootDocuments(0);
+        database.putRootDocuments(0, resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device", "Storage A", 0, 0, ""),
+        });
+        database.clearMapping();
+        // Add one.
+        database.putRootDocuments(0, resources, new MtpRoot[] {
+                new MtpRoot(0, 100, "Device", "Storage B", 1000, 1000, ""),
+        });
+        // Add one more before resolving unmapped documents.
+        try {
+            database.putRootDocuments(0, resources, new MtpRoot[] {
+                    new MtpRoot(0, 100, "Device", "Storage B", 1000, 1000, ""),
+            });
+            fail();
+        } catch (Throwable e) {
+            assertTrue(e instanceof Error);
+        }
+    }
 }