From 995e8895bab6d436ab84743b9404229866f533a9 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Mon, 2 Dec 2019 17:56:03 +0300
Subject: [PATCH] Transparent return of a correct scheduled server message.

GitOrigin-RevId: c260cefdcef0a919f6aeb4a025500cfc3cd5af00
---
 td/telegram/MessagesDb.cpp      | 15 +++++++++++++--
 td/telegram/MessagesManager.cpp | 13 +++++++++----
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/td/telegram/MessagesDb.cpp b/td/telegram/MessagesDb.cpp
index 54f0569eb..8dac32da9 100644
--- a/td/telegram/MessagesDb.cpp
+++ b/td/telegram/MessagesDb.cpp
@@ -248,6 +248,9 @@ class MessagesDbImpl : public MessagesDbSyncInterface {
     TRY_RESULT_ASSIGN(
         get_scheduled_message_stmt_,
         db_.get_statement("SELECT data FROM scheduled_messages WHERE dialog_id = ?1 AND message_id = ?2"));
+    TRY_RESULT_ASSIGN(
+        get_scheduled_server_message_stmt_,
+        db_.get_statement("SELECT data FROM scheduled_messages WHERE dialog_id = ?1 AND server_message_id = ?2"));
     TRY_RESULT_ASSIGN(delete_scheduled_message_stmt_,
                       db_.get_statement("DELETE FROM scheduled_messages WHERE dialog_id = ?1 AND message_id = ?2"));
 
@@ -412,13 +415,20 @@ class MessagesDbImpl : public MessagesDbSyncInterface {
     auto message_id = full_message_id.get_message_id();
     CHECK(dialog_id.is_valid());
     CHECK(message_id.is_valid() || message_id.is_valid_scheduled());
-    auto &stmt = message_id.is_scheduled() ? get_scheduled_message_stmt_ : get_message_stmt_;
+    bool is_scheduled = message_id.is_scheduled();
+    bool is_scheduled_server = is_scheduled && message_id.is_scheduled_server();
+    auto &stmt = is_scheduled ? (is_scheduled_server ? get_scheduled_server_message_stmt_ : get_scheduled_message_stmt_)
+                              : get_message_stmt_;
     SCOPE_EXIT {
       stmt.reset();
     };
 
     stmt.bind_int64(1, dialog_id.get()).ensure();
-    stmt.bind_int64(2, message_id.get()).ensure();
+    if (is_scheduled_server) {
+      stmt.bind_int32(2, message_id.get_scheduled_server_message_id().get()).ensure();
+    } else {
+      stmt.bind_int64(2, message_id.get()).ensure();
+    }
     stmt.step().ensure();
     if (!stmt.has_row()) {
       return Status::Error("Not found");
@@ -783,6 +793,7 @@ class MessagesDbImpl : public MessagesDbSyncInterface {
 
   SqliteStatement add_scheduled_message_stmt_;
   SqliteStatement get_scheduled_message_stmt_;
+  SqliteStatement get_scheduled_server_message_stmt_;
   SqliteStatement delete_scheduled_message_stmt_;
 
   Result<std::vector<BufferSlice>> get_messages_impl(GetMessagesStmt &stmt, DialogId dialog_id,
diff --git a/td/telegram/MessagesManager.cpp b/td/telegram/MessagesManager.cpp
index 091395e6c..0a6814ceb 100644
--- a/td/telegram/MessagesManager.cpp
+++ b/td/telegram/MessagesManager.cpp
@@ -24542,6 +24542,15 @@ const MessagesManager::Message *MessagesManager::get_message(const Dialog *d, Me
   CHECK(d != nullptr);
   LOG(DEBUG) << "Search for " << message_id << " in " << d->dialog_id;
   bool is_scheduled = message_id.is_scheduled();
+  if (is_scheduled && message_id.is_scheduled_server()) {
+    auto server_message_id = message_id.get_scheduled_server_message_id();
+    auto it = d->scheduled_message_date.find(server_message_id);
+    if (it != d->scheduled_message_date.end()) {
+      int32 date = it->second;
+      message_id = MessageId(server_message_id, date);
+      CHECK(message_id.is_scheduled_server());
+    }
+  }
   auto result = treap_find_message(is_scheduled ? &d->scheduled_messages : &d->messages, message_id)->get();
   if (result != nullptr && !is_scheduled) {
     result->last_access_date = G()->unix_time_cached();
@@ -24569,10 +24578,6 @@ MessagesManager::Message *MessagesManager::get_message_force(Dialog *d, MessageI
 
   LOG(INFO) << "Trying to load " << FullMessageId{d->dialog_id, message_id} << " from database from " << source;
 
-  if (message_id.is_scheduled_server()) {
-    // TODO load scheduled message by server message_id
-  }
-
   auto r_value = G()->td_db()->get_messages_db_sync()->get_message({d->dialog_id, message_id});
   if (r_value.is_error()) {
     return nullptr;