Setup saving drafts.

Change-Id: Id308cf902ec8ea15d2d785f21b1c821a281f8a2c
diff --git a/src/com/android/mail/compose/ComposeActivity.java b/src/com/android/mail/compose/ComposeActivity.java
index fb13ba9..290fd63 100644
--- a/src/com/android/mail/compose/ComposeActivity.java
+++ b/src/com/android/mail/compose/ComposeActivity.java
@@ -28,6 +28,7 @@
 import android.content.DialogInterface;
 import android.content.Intent;
 import android.content.pm.ActivityInfo;
+import android.database.Cursor;
 import android.net.Uri;
 import android.os.Bundle;
 import android.os.Handler;
@@ -155,11 +156,13 @@
     private boolean mReplyFromChanged;
     private MenuItem mSave;
     private MenuItem mSend;
-    private Object mDraftIdLock = new Object();
     private String mRefMessageId;
     private AlertDialog mRecipientErrorDialog;
     private AlertDialog mSendConfirmDialog;
     private Message mRefMessage;
+    private long mDraftId = UIProvider.INVALID_MESSAGE_ID;
+    private Message mDraft;
+    private Object mDraftLock = new Object();
 
     /**
      * Can be called from a non-UI thread.
@@ -670,8 +673,8 @@
 
     /*package*/ interface SendOrSaveCallback {
         public void initializeSendOrSave(SendOrSaveTask sendOrSaveTask);
-        public void notifyMessageIdAllocated(SendOrSaveMessage message, long messageId);
-        public long getMessageId();
+        public void notifyMessageIdAllocated(SendOrSaveMessage sendOrSaveMessage, Message message);
+        public Message getMessage();
         public void sendOrSaveFinished(SendOrSaveTask sendOrSaveTask, boolean success);
     }
 
@@ -690,14 +693,15 @@
 
         @Override
         public void run() {
-            final SendOrSaveMessage message = mSendOrSaveMessage;
+            final SendOrSaveMessage sendOrSaveMessage = mSendOrSaveMessage;
 
-            final Account selectedAccount = message.mSelectedAccount;
-            long messageId = mSendOrSaveCallback.getMessageId();
+            final Account selectedAccount = sendOrSaveMessage.mSelectedAccount;
+            Message message = mSendOrSaveCallback.getMessage();
+            long messageId = message != null ? message.id : UIProvider.INVALID_MESSAGE_ID;
             // If a previous draft has been saved, in an account that is different
             // than what the user wants to send from, remove the old draft, and treat this
             // as a new message
-            if (!selectedAccount.equals(message.mAccount)) {
+            if (!selectedAccount.equals(sendOrSaveMessage.mAccount)) {
                 if (messageId != UIProvider.INVALID_MESSAGE_ID) {
                     ContentResolver resolver = mContext.getContentResolver();
                     ContentValues values = new ContentValues();
@@ -712,28 +716,35 @@
             }
 
             final long messageIdToSave = messageId;
-            int newDraftId = -1;
             if (messageIdToSave != UIProvider.INVALID_MESSAGE_ID) {
+                sendOrSaveMessage.mValues.put(BaseColumns._ID, messageIdToSave);
                 mContext.getContentResolver().update(
-                        Uri.parse(message.mSave ? selectedAccount.saveDraftUri
-                                : selectedAccount.sendMessageUri), message.mValues, null, null);
+                        Uri.parse(sendOrSaveMessage.mSave ? message.saveUri : message.sendUri),
+                        sendOrSaveMessage.mValues, null, null);
             } else {
-                newDraftId = mContext.getContentResolver().update(
-                        Uri.parse(message.mSave ? selectedAccount.saveDraftUri
-                                : selectedAccount.sendMessageUri), message.mValues, null, null);
-
-                // Broadcast notification that a new message id has been
-                // allocated
-                mSendOrSaveCallback.notifyMessageIdAllocated(message, newDraftId);
+                ContentResolver resolver = mContext.getContentResolver();
+                Uri messageUri = resolver.insert(
+                        Uri.parse(sendOrSaveMessage.mSave ? selectedAccount.saveDraftUri
+                                : selectedAccount.sendMessageUri), sendOrSaveMessage.mValues);
+                if (sendOrSaveMessage.mSave && messageUri != null) {
+                    Cursor messageCursor = resolver.query(messageUri,
+                            UIProvider.MESSAGE_PROJECTION, null, null, null);
+                    if (messageCursor != null && messageCursor.moveToFirst()) {
+                        // Broadcast notification that a new message has
+                        // been allocated
+                        mSendOrSaveCallback.notifyMessageIdAllocated(sendOrSaveMessage,
+                                Message.from(messageCursor));
+                    }
+                }
             }
 
-            if (!message.mSave) {
+            if (!sendOrSaveMessage.mSave) {
                 UIProvider.incrementRecipientsTimesContacted(mContext,
-                        (String) message.mValues.get(UIProvider.MessageColumns.TO));
+                        (String) sendOrSaveMessage.mValues.get(UIProvider.MessageColumns.TO));
                 UIProvider.incrementRecipientsTimesContacted(mContext,
-                        (String) message.mValues.get(UIProvider.MessageColumns.CC));
+                        (String) sendOrSaveMessage.mValues.get(UIProvider.MessageColumns.CC));
                 UIProvider.incrementRecipientsTimesContacted(mContext,
-                        (String) message.mValues.get(UIProvider.MessageColumns.BCC));
+                        (String) sendOrSaveMessage.mValues.get(UIProvider.MessageColumns.BCC));
             }
             mSendOrSaveCallback.sendOrSaveFinished(SendOrSaveTask.this, true);
         }
@@ -744,7 +755,6 @@
     /* package for testing */
     ArrayList<SendOrSaveTask> mActiveTasks = Lists.newArrayList();
     private int mRequestId;
-    private long mDraftId;
 
     /*package*/ static class SendOrSaveMessage {
         final Account mAccount;
@@ -859,7 +869,7 @@
      * Returns true if we need to save the current draft.
      */
     private boolean shouldSave() {
-        synchronized (mDraftIdLock) {
+        synchronized (mDraftLock) {
             // The message should only be saved if:
             // It hasn't been sent AND
             // Some text has been added to the message OR
@@ -1102,7 +1112,6 @@
 
 
        SendOrSaveCallback callback = new SendOrSaveCallback() {
-               private long mDraftId;
             private int mRestoredRequestId;
 
             public void initializeSendOrSave(SendOrSaveTask sendOrSaveTask) {
@@ -1121,63 +1130,55 @@
                    }
                }
 
-               public void notifyMessageIdAllocated(SendOrSaveMessage message, long messageId) {
-                   synchronized(mDraftIdLock) {
-                       mDraftId = messageId;
-                       sRequestMessageIdMap.put(message.requestId(), messageId);
-
+            public void notifyMessageIdAllocated(SendOrSaveMessage sendOrSaveMessage,
+                    Message message) {
+                   synchronized(mDraftLock) {
+                       mDraftId = message.id;
+                       mDraft = message;
+                    if (sRequestMessageIdMap != null) {
+                        sRequestMessageIdMap.put(sendOrSaveMessage.requestId(), mDraftId);
+                    }
                        // Cache request message map, in case the process is killed
                        saveRequestMap();
                    }
                    if (sTestSendOrSaveCallback != null) {
-                       sTestSendOrSaveCallback.notifyMessageIdAllocated(message, messageId);
+                    sTestSendOrSaveCallback.notifyMessageIdAllocated(sendOrSaveMessage, message);
                    }
                }
 
-               public long getMessageId() {
-                   synchronized(mDraftIdLock) {
-                       if (mDraftId == UIProvider.INVALID_MESSAGE_ID) {
-                           // We don't have the message Id, check to see if we have a restored
-                           // request id, and see if we have a message for that request.
-                           if (mRestoredRequestId != 0) {
-                               Long retrievedMessageId =
-                                       sRequestMessageIdMap.get(mRestoredRequestId);
-                               if (retrievedMessageId != null) {
-                                   mDraftId = retrievedMessageId.longValue();
-                               }
-                           }
-                       }
-                       return mDraftId;
-                   }
-               }
+            public Message getMessage() {
+                synchronized (mDraftLock) {
+                    return mDraft;
+                }
+            }
 
-               public void sendOrSaveFinished(SendOrSaveTask task, boolean success) {
-                   if (success) {
-                       // Successfully sent or saved so reset change markers
-                       discardChanges();
-                   } else {
-                       // A failure happened with saving/sending the draft
-                       // TODO(pwestbro): add a better string that should be used when failing to
-                       // send or save
-                       Toast.makeText(ComposeActivity.this, R.string.send_failed,
-                               Toast.LENGTH_SHORT).show();
-                   }
+            public void sendOrSaveFinished(SendOrSaveTask task, boolean success) {
+                if (success) {
+                    // Successfully sent or saved so reset change markers
+                    discardChanges();
+                } else {
+                    // A failure happened with saving/sending the draft
+                    // TODO(pwestbro): add a better string that should be used
+                    // when failing to send or save
+                    Toast.makeText(ComposeActivity.this, R.string.send_failed, Toast.LENGTH_SHORT)
+                            .show();
+                }
 
-                   int numTasks;
-                   synchronized(mActiveTasks) {
-                       // Remove the task from the list of active tasks
-                       mActiveTasks.remove(task);
-                       numTasks = mActiveTasks.size();
-                   }
+                int numTasks;
+                synchronized (mActiveTasks) {
+                    // Remove the task from the list of active tasks
+                    mActiveTasks.remove(task);
+                    numTasks = mActiveTasks.size();
+                }
 
-                   if (numTasks == 0) {
-                       // Stop service so we can be killed.
-                       stopService(new Intent(ComposeActivity.this, EmptyService.class));
-                   }
-                   if (sTestSendOrSaveCallback != null) {
-                       sTestSendOrSaveCallback.sendOrSaveFinished(task, success);
-                   }
-               }
+                if (numTasks == 0) {
+                    // Stop service so we can be killed.
+                    stopService(new Intent(ComposeActivity.this, EmptyService.class));
+                }
+                if (sTestSendOrSaveCallback != null) {
+                    sTestSendOrSaveCallback.sendOrSaveFinished(task, success);
+                }
+            }
          };
 
        // Get the selected account if the from spinner has been setup.
@@ -1408,7 +1409,7 @@
      * @param showToast show "Message discarded" toast if true
      */
     private void doDiscardWithoutConfirmation(boolean showToast) {
-        synchronized (mDraftIdLock) {
+        synchronized (mDraftLock) {
             if (mDraftId != UIProvider.INVALID_MESSAGE_ID) {
                 ContentValues values = new ContentValues();
                 values.put(MessageColumns.SERVER_ID, mDraftId);
diff --git a/src/com/android/mail/providers/AccountCacheProvider.java b/src/com/android/mail/providers/AccountCacheProvider.java
index 82ae24a..40f1f99 100644
--- a/src/com/android/mail/providers/AccountCacheProvider.java
+++ b/src/com/android/mail/providers/AccountCacheProvider.java
@@ -91,8 +91,8 @@
                     builder.add(account.mAccountFromAddressesUri);
                 } else if (TextUtils.equals(column, UIProvider.AccountColumns.SAVE_DRAFT_URI)) {
                     builder.add(account.mSaveDraftUri);
-                } else if (TextUtils.equals(column, UIProvider.AccountColumns.SEND_MESSAGE_URI)) {
-                    builder.add(account.mSendMessageUri);
+                } else if (TextUtils.equals(column, UIProvider.AccountColumns.SEND_MAIL_URI)) {
+                    builder.add(account.mSendMailUri);
                 }
             }
 
@@ -149,11 +149,12 @@
         private final String mSearchUri;
         private final String mAccountFromAddressesUri;
         private final String mSaveDraftUri;
-        private final String mSendMessageUri;
+        private final String mSendMailUri;
+        private final String mExpungeMessageUri;
 
         public CachedAccount(long id, String name, String uri, long capabilities,
                 String folderListUri, String searchUri, String fromAddressesUri,
-                String saveDraftUri, String sendMessageUri) {
+                String saveDraftUri, String sendMailUri, String expungeMessageUri) {
             mId = id;
             mName = name;
             mUri = uri;
@@ -162,7 +163,8 @@
             mSearchUri = searchUri;
             mAccountFromAddressesUri = fromAddressesUri;
             mSaveDraftUri = saveDraftUri;
-            mSendMessageUri = sendMessageUri;
+            mSendMailUri = sendMailUri;
+            mExpungeMessageUri = expungeMessageUri;
         }
 
         @Override
@@ -182,13 +184,14 @@
                     TextUtils.equals(mSearchUri, other.mSearchUri) &&
                     TextUtils.equals(mAccountFromAddressesUri, other.mAccountFromAddressesUri) &&
                     TextUtils.equals(mSaveDraftUri, other.mSaveDraftUri) &&
-                    TextUtils.equals(mSendMessageUri, other.mSendMessageUri);
+                    TextUtils.equals(mSendMailUri, other.mSendMailUri) &&
+                    TextUtils.equals(mExpungeMessageUri, other.mExpungeMessageUri);
         }
 
         @Override
         public int hashCode() {
             return Objects.hashCode(mId, mName, mUri, mCapabilities, mFolderListUri, mSearchUri,
-                    mAccountFromAddressesUri, mSaveDraftUri, mSendMessageUri);
+                    mAccountFromAddressesUri, mSaveDraftUri, mSendMailUri, mExpungeMessageUri);
         }
     }
 }
\ No newline at end of file
diff --git a/src/com/android/mail/providers/Message.java b/src/com/android/mail/providers/Message.java
index 9a428c6..0153cda 100644
--- a/src/com/android/mail/providers/Message.java
+++ b/src/com/android/mail/providers/Message.java
@@ -44,6 +44,8 @@
     public String attachmentListUri;
     public long messageFlags;
     public String joinedAttachmentInfos;
+    public String saveUri;
+    public String sendUri;
 
     @Override
     public int describeContents() {
@@ -74,6 +76,8 @@
         dest.writeString(attachmentListUri);
         dest.writeLong(messageFlags);
         dest.writeString(joinedAttachmentInfos);
+        dest.writeString(saveUri);
+        dest.writeString(sendUri);
     }
 
     public Message() {
@@ -102,6 +106,8 @@
         attachmentListUri = in.readString();
         messageFlags = in.readLong();
         joinedAttachmentInfos = in.readString();
+        saveUri = in.readString();
+        sendUri = in.readString();
     }
 
     @Override
@@ -154,6 +160,10 @@
             messageFlags = cursor.getLong(UIProvider.MESSAGE_FLAGS_COLUMN);
             joinedAttachmentInfos = cursor
                     .getString(UIProvider.MESSAGE_JOINED_ATTACHMENT_INFOS_COLUMN);
+            saveUri = cursor
+                    .getString(UIProvider.MESSAGE_SAVE_URI_COLUMN);
+            sendUri = cursor
+                    .getString(UIProvider.MESSAGE_SEND_URI_COLUMN);
         }
     }
 }
diff --git a/src/com/android/mail/providers/UIProvider.java b/src/com/android/mail/providers/UIProvider.java
index 7f872de..8fcfe2b 100644
--- a/src/com/android/mail/providers/UIProvider.java
+++ b/src/com/android/mail/providers/UIProvider.java
@@ -50,7 +50,7 @@
             AccountColumns.SEARCH_URI,
             AccountColumns.ACCOUNT_FROM_ADDRESSES_URI,
             AccountColumns.SAVE_DRAFT_URI,
-            AccountColumns.SEND_MESSAGE_URI,
+            AccountColumns.SEND_MAIL_URI,
             AccountColumns.EXPUNGE_MESSAGE_URI
     };
 
@@ -134,7 +134,7 @@
          * a message for this account.
          * NOTE: This might be better to be an update operation on the messageUri.
          */
-        public static final String SEND_MESSAGE_URI = "sendMessageUri";
+        public static final String SEND_MAIL_URI = "sendMailUri";
 
         /**
          * This string column contains the content provider uri that can be used
@@ -376,7 +376,9 @@
         MessageColumns.HAS_ATTACHMENTS,
         MessageColumns.ATTACHMENT_LIST_URI,
         MessageColumns.MESSAGE_FLAGS,
-        MessageColumns.JOINED_ATTACHMENT_INFOS
+        MessageColumns.JOINED_ATTACHMENT_INFOS,
+        MessageColumns.SAVE_MESSAGE_URI,
+        MessageColumns.SEND_MESSAGE_URI
     };
 
     /** Separates attachment info parts in strings in a message. */
@@ -408,6 +410,8 @@
     public static final int MESSAGE_ATTACHMENT_LIST_URI_COLUMN = 19;
     public static final int MESSAGE_FLAGS_COLUMN = 20;
     public static final int MESSAGE_JOINED_ATTACHMENT_INFOS_COLUMN = 21;
+    public static final int MESSAGE_SAVE_URI_COLUMN = 22;
+    public static final int MESSAGE_SEND_URI_COLUMN = 23;
 
     public static final class MessageFlags {
         public static final int STARRED =       1 << 0;
@@ -492,8 +496,8 @@
          */
         public static final String HAS_ATTACHMENTS = "hasAttachments";
         /**
-         * This string column contains content provider URI for the list of attachments associated
-         * with this message.
+         * This string column contains the content provider URI for the list of
+         * attachments associated with this message.
          */
         public static final String ATTACHMENT_LIST_URI = "attachmentListUri";
         /**
@@ -505,6 +509,16 @@
          * attachments that we added to a message that is being sent or saved.
          */
         public static final String JOINED_ATTACHMENT_INFOS = "joinedAttachmentInfos";
+        /**
+         * This string column contains the content provider URI for saving this
+         * message.
+         */
+        public static final String SAVE_MESSAGE_URI = "saveMessageUri";
+        /**
+         * This string column contains content provider URI for sending this
+         * message.
+         */
+        public static final String SEND_MESSAGE_URI = "sendMessageUri";
 
         private MessageColumns() {}
     }
diff --git a/src/com/android/mail/providers/protos/boot/EmailAccountService.java b/src/com/android/mail/providers/protos/boot/EmailAccountService.java
index f34ee03..292b722 100644
--- a/src/com/android/mail/providers/protos/boot/EmailAccountService.java
+++ b/src/com/android/mail/providers/protos/boot/EmailAccountService.java
@@ -79,7 +79,8 @@
                             (String)mockAccountMap.get(AccountColumns.SEARCH_URI),
                             (String)mockAccountMap.get(AccountColumns.ACCOUNT_FROM_ADDRESSES_URI),
                             (String)mockAccountMap.get(AccountColumns.SAVE_DRAFT_URI),
-                            getUriString("uisendmail", accountName));
+                            getUriString("uisendmail", accountName),
+                            (String)mockAccountMap.get(AccountColumns.EXPUNGE_MESSAGE_URI));
 
                 AccountCacheProvider.addAccount(cachedAccount);
                 i++;
diff --git a/src/com/android/mail/providers/protos/boot/GmailAccountService.java b/src/com/android/mail/providers/protos/boot/GmailAccountService.java
index 3be59b2..2723bbe 100644
--- a/src/com/android/mail/providers/protos/boot/GmailAccountService.java
+++ b/src/com/android/mail/providers/protos/boot/GmailAccountService.java
@@ -50,8 +50,12 @@
         return Uri.parse(GMAIL_UI_PROVIDER_BASE_URI_STRING + "/" + account + "/labels");
     }
 
-    private static Uri getAccountSendMessageUri(String account) {
-        return Uri.parse(GMAIL_UI_PROVIDER_BASE_URI_STRING + "/" + account + "/messages");
+    private static Uri getAccountSendMailUri(String account) {
+        return Uri.parse(GMAIL_UI_PROVIDER_BASE_URI_STRING + "/" + account + "/sendNewMessage");
+    }
+
+    private static Uri getAccountSaveDraftUri(String account) {
+        return Uri.parse(GMAIL_UI_PROVIDER_BASE_URI_STRING + "/" + account + "/saveNewMessage");
     }
 
     public GmailAccountService() {
@@ -111,8 +115,9 @@
                             getAccountFoldersUri(account.name).toString(),
                             (String)mockAccountMap.get(AccountColumns.SEARCH_URI),
                             (String)mockAccountMap.get(AccountColumns.ACCOUNT_FROM_ADDRESSES_URI),
-                            (String)mockAccountMap.get(AccountColumns.SAVE_DRAFT_URI),
-                            getAccountSendMessageUri(account.name).toString());
+                            getAccountSaveDraftUri(account.name).toString(),
+                            getAccountSendMailUri(account.name).toString(),
+                            (String)mockAccountMap.get(AccountColumns.EXPUNGE_MESSAGE_URI));
 
             AccountCacheProvider.addAccount(cachedAccount);
         }
diff --git a/src/com/android/mail/providers/protos/mock/MockUiProvider.java b/src/com/android/mail/providers/protos/mock/MockUiProvider.java
index 1fa9c1f..8ef5eec 100644
--- a/src/com/android/mail/providers/protos/mock/MockUiProvider.java
+++ b/src/com/android/mail/providers/protos/mock/MockUiProvider.java
@@ -264,7 +264,7 @@
         accountMap.put(AccountColumns.SEARCH_URI, accountUri + "/search");
         accountMap.put(AccountColumns.ACCOUNT_FROM_ADDRESSES_URI, accountUri + "/fromAddresses");
         accountMap.put(AccountColumns.SAVE_DRAFT_URI, accountUri + "/saveDraft");
-        accountMap.put(AccountColumns.SEND_MESSAGE_URI, accountUri + "/sendMessage");
+        accountMap.put(AccountColumns.SEND_MAIL_URI, accountUri + "/sendMail");
         accountMap.put(AccountColumns.EXPUNGE_MESSAGE_URI, accountUri + "/expungeMessage");
 
         addAccountInfoToAccountCache(accountMap);
@@ -349,7 +349,8 @@
                         (String)accountInfo.get(AccountColumns.SEARCH_URI),
                         (String)accountInfo.get(AccountColumns.ACCOUNT_FROM_ADDRESSES_URI),
                         (String)accountInfo.get(AccountColumns.SAVE_DRAFT_URI),
-                        (String)accountInfo.get(AccountColumns.SEND_MESSAGE_URI));
+                        (String)accountInfo.get(AccountColumns.SEND_MAIL_URI),
+                        (String)accountInfo.get(AccountColumns.EXPUNGE_MESSAGE_URI));
 
 
         AccountCacheProvider.addAccount(account);