Short-circuit storyId predicate to optimize query

ts/sql/Server.ts

@@ -2056,6 +2056,16 @@ async function getMessageBySender({
   return jsonToObject(rows[0].json);
 }
 
+export function _storyIdPredicate(storyId: string | undefined): string {
+  if (storyId === undefined) {
+    // We could use 'TRUE' here, but it is better to require `$storyId`
+    // parameter
+    return '$storyId IS NULL';
+  }
+
+  return 'storyId IS $storyId';
+}
+
 async function getUnreadByConversationAndMarkRead({
   conversationId,
   newestUnreadAt,
@@ -2086,7 +2096,7 @@ async function getUnreadByConversationAndMarkRead({
         ) AND
         expireTimer > 0 AND
         conversationId = $conversationId AND
-        ($storyId IS NULL OR storyId IS $storyId) AND
+        (${_storyIdPredicate(storyId)}) AND
         received_at <= $newestUnreadAt;
       `
     ).run({
@@ -2105,7 +2115,7 @@ async function getUnreadByConversationAndMarkRead({
         WHERE
           readStatus = ${ReadStatus.Unread} AND
           conversationId = $conversationId AND
-          ($storyId IS NULL OR storyId IS $storyId) AND
+          (${_storyIdPredicate(storyId)}) AND
           received_at <= $newestUnreadAt
         ORDER BY received_at DESC, sent_at DESC;
         `
@@ -2125,7 +2135,7 @@ async function getUnreadByConversationAndMarkRead({
         WHERE
           readStatus = ${ReadStatus.Unread} AND
           conversationId = $conversationId AND
-          ($storyId IS NULL OR storyId IS $storyId) AND
+          (${_storyIdPredicate(storyId)}) AND
           received_at <= $newestUnreadAt;
         `
     ).run({
@@ -2360,7 +2370,7 @@ function getOlderMessagesByConversationSync(
         conversationId = $conversationId AND
         ($messageId IS NULL OR id IS NOT $messageId) AND
         isStory IS 0 AND
-        ($storyId IS NULL OR storyId IS $storyId) AND
+        (${_storyIdPredicate(storyId)}) AND
         (
           (received_at = $received_at AND sent_at < $sent_at) OR
           received_at < $received_at
@@ -2453,7 +2463,7 @@ function getNewerMessagesByConversationSync(
       SELECT json FROM messages WHERE
         conversationId = $conversationId AND
         isStory IS 0 AND
-        ($storyId IS NULL OR storyId IS $storyId) AND
+        (${_storyIdPredicate(storyId)}) AND
         (
           (received_at = $received_at AND sent_at > $sent_at) OR
           received_at > $received_at
@@ -2483,7 +2493,7 @@ function getOldestMessageForConversation(
       SELECT * FROM messages WHERE
         conversationId = $conversationId AND
         isStory IS 0 AND
-        ($storyId IS NULL OR storyId IS $storyId)
+        (${_storyIdPredicate(storyId)})
       ORDER BY received_at ASC, sent_at ASC
       LIMIT 1;
       `
@@ -2510,7 +2520,7 @@ function getNewestMessageForConversation(
       SELECT * FROM messages WHERE
         conversationId = $conversationId AND
         isStory IS 0 AND
-        ($storyId IS NULL OR storyId IS $storyId)
+        (${_storyIdPredicate(storyId)})
       ORDER BY received_at DESC, sent_at DESC
       LIMIT 1;
       `
@@ -2651,7 +2661,7 @@ function getOldestUnreadMessageForConversation(
         conversationId = $conversationId AND
         readStatus = ${ReadStatus.Unread} AND
         isStory IS 0 AND
-        ($storyId IS NULL OR storyId IS $storyId)
+        (${_storyIdPredicate(storyId)})
       ORDER BY received_at ASC, sent_at ASC
       LIMIT 1;
       `
@@ -2688,7 +2698,7 @@ function getTotalUnreadForConversationSync(
         conversationId = $conversationId AND
         readStatus = ${ReadStatus.Unread} AND
         isStory IS 0 AND
-        ($storyId IS NULL OR storyId IS $storyId);
+        (${_storyIdPredicate(storyId)})
       `
     )
     .get({

ts/sql/migrations/52-optimize-stories.ts

@@ -0,0 +1,38 @@
+// Copyright 2022 Signal Messenger, LLC
+// SPDX-License-Identifier: AGPL-3.0-only
+
+import type { Database } from 'better-sqlite3';
+
+import type { LoggerType } from '../../types/Logging';
+
+export default function updateToSchemaVersion52(
+  currentVersion: number,
+  db: Database,
+  logger: LoggerType
+): void {
+  if (currentVersion >= 52) {
+    return;
+  }
+
+  db.transaction(() => {
+    db.exec(
+      `
+      -- Create indices that don't have storyId in them so that
+      -- '_storyIdPredicate' could be optimized.
+
+      -- See migration 47
+      CREATE INDEX messages_conversation_no_story_id ON messages
+        (conversationId, isStory, received_at, sent_at);
+
+      -- See migration 50
+      CREATE INDEX messages_unread_no_story_id ON messages
+        (conversationId, readStatus, isStory, received_at, sent_at)
+        WHERE readStatus IS NOT NULL;
+      `
+    );
+
+    db.pragma('user_version = 52');
+  })();
+
+  logger.info('updateToSchemaVersion52: success!');
+}

ts/sql/migrations/index.ts

@@ -27,6 +27,7 @@ import updateToSchemaVersion48 from './48-fix-user-initiated-index';
 import updateToSchemaVersion49 from './49-fix-preview-index';
 import updateToSchemaVersion50 from './50-fix-messages-unread-index';
 import updateToSchemaVersion51 from './51-centralize-conversation-jobs';
+import updateToSchemaVersion52 from './52-optimize-stories';
 
 function updateToSchemaVersion1(
   currentVersion: number,
@@ -1917,6 +1918,7 @@ export const SCHEMA_VERSIONS = [
   updateToSchemaVersion49,
   updateToSchemaVersion50,
   updateToSchemaVersion51,
+  updateToSchemaVersion52,
 ];
 
 export function updateSchema(db: Database, logger: LoggerType): void {

ts/test-node/sql_migrations_test.ts

@@ -8,7 +8,11 @@ import { v4 as generateGuid } from 'uuid';
 
 import { SCHEMA_VERSIONS } from '../sql/migrations';
 import { consoleLogger } from '../util/consoleLogger';
-import { getJobsInQueueSync, insertJobSync } from '../sql/Server';
+import {
+  getJobsInQueueSync,
+  insertJobSync,
+  _storyIdPredicate,
+} from '../sql/Server';
 
 const OUR_UUID = generateGuid();
 
@@ -1600,4 +1604,70 @@ describe('SQL migrations test', () => {
       ]);
     });
   });
+
+  describe('updateToSchemaVersion52', () => {
+    const queries = [
+      {
+        query: `
+          EXPLAIN QUERY PLAN
+          SELECT * FROM messages WHERE
+            conversationId = 'conversation' AND
+            readStatus = 'something' AND
+            isStory IS 0 AND
+            :story_id_predicate:
+          ORDER BY received_at ASC, sent_at ASC
+          LIMIT 1;
+        `,
+        index: 'messages_unread',
+      },
+      {
+        query: `
+          EXPLAIN QUERY PLAN
+          SELECT json FROM messages WHERE
+            conversationId = 'd8b05bb1-36b3-4478-841b-600af62321eb' AND
+            (NULL IS NULL OR id IS NOT NULL) AND
+            isStory IS 0 AND
+            :story_id_predicate: AND
+            (
+              (received_at = 17976931348623157 AND sent_at < NULL) OR
+              received_at < 17976931348623157
+            )
+          ORDER BY received_at DESC, sent_at DESC
+          LIMIT 10;
+        `,
+        index: 'messages_conversation',
+      },
+    ];
+
+    function insertPredicate(
+      query: string,
+      storyId: string | undefined
+    ): string {
+      return query.replaceAll(
+        ':story_id_predicate:',
+        _storyIdPredicate(storyId)
+      );
+    }
+
+    it('produces optimizable queries for present and absent storyId', () => {
+      updateToVersion(52);
+
+      for (const storyId of ['123', undefined]) {
+        for (const { query, index } of queries) {
+          const details = db
+            .prepare(insertPredicate(query, storyId))
+            .all({ storyId })
+            .map(({ detail }) => detail)
+            .join('\n');
+
+          const postfixedIndex = index + (storyId ? '' : '_no_story_id');
+
+          // Intentional trailing whitespace
+          assert.include(details, `USING INDEX ${postfixedIndex} `);
+          assert.notInclude(details, 'TEMP B-TREE');
+          assert.notInclude(details, 'SCAN');
+        }
+      }
+    });
+  });
 });