feat: add index optimisations to lines

Author: harjot1singh

src/index.spec.ts

@@ -1,33 +1,72 @@
 import { describe, expect, it } from 'bun:test'
 
+import { LibSQLDatabase } from 'drizzle-orm/libsql'
+import { SQLiteRelationalQuery } from 'drizzle-orm/sqlite-core/query-builders/query'
 import { parse } from 'smol-toml'
 
 import { LineGroups } from '#collections-types/line-groups'
 
 import createDbClient from '.'
+import { getQueryPlan } from './query-utils'
+import relations from './relations'
+import * as schema from './schema'
 
 const getCollectionDoc = async <Schema>(document: string) =>
   parse(await Bun.file(`collections/${document}.toml`).text()) as Schema
 
+// Generic function to test query plan optimization
+const expectOptimizedQueryPlan = async <Q extends SQLiteRelationalQuery<'async', unknown>>(
+  db: LibSQLDatabase<typeof schema, typeof relations>,
+  query: Q,
+  {
+    maxTableScans = 0,
+    expectedIndexes = [],
+  }: { maxTableScans?: number; expectedIndexes?: string[] } = {},
+) => {
+  const queryPlan = await getQueryPlan(db, query)
+  const planText = queryPlan.join(' ')
+
+  const tableScans = queryPlan.filter(
+    (step) => typeof step === 'string' && step.includes('SCAN') && !step.includes('USING INDEX'),
+  ).length
+
+  expect(tableScans, planText).toBeLessThanOrEqual(maxTableScans)
+  expectedIndexes.forEach((indexName) => expect(planText).toContain(indexName))
+
+  return queryPlan
+}
+
 describe('Database', () => {
-  it('should return lines in order', async () => {
-    const db = createDbClient()
-    const id = 'DMP'
-
-    const lineGroup = await db.query.lineGroups.findFirst({
-      where: { id },
-      with: {
-        author: true,
-        lines: {
-          orderBy: {
-            lineGroupOrder: 'asc',
+  describe('line groups', () => {
+    it('should return, optimized, lines in order', async () => {
+      const db = createDbClient()
+      const id = 'DMP'
+
+      const lineGroupQuery = db.query.lineGroups.findFirst({
+        where: { id },
+        with: {
+          author: true,
+          lines: {
+            orderBy: {
+              lineGroupOrder: 'asc',
+            },
           },
         },
-      },
-    })
+      })
+      const lineGroup = await lineGroupQuery
 
-    const collectionDoc = await getCollectionDoc<LineGroups>(`line-groups/D/${id}`)
-    expect(lineGroup?.lines.map((line) => line.id)).toEqual(collectionDoc.lines)
-    expect(lineGroup?.author.id).toEqual(collectionDoc.author)
+      const collectionDoc = await getCollectionDoc<LineGroups>(`line-groups/D/${id}`)
+
+      expect(lineGroup?.lines.map((line) => line.id)).toEqual(collectionDoc.lines)
+      expect(lineGroup?.author.id).toEqual(collectionDoc.author)
+      await expectOptimizedQueryPlan(db, lineGroupQuery, {
+        maxTableScans: 2,
+        expectedIndexes: [
+          'sqlite_autoindex_line_groups_1',
+          'sqlite_autoindex_authors_1',
+          'line_group_id_order_index',
+        ],
+      })
+    })
   })
 })

src/schema.ts

@@ -55,7 +55,7 @@ export const lines = sqliteTable(
     lineGroupId: text(),
     lineGroupOrder: integer(),
   },
-  (t) => [index('line_group_order_index').on(t.lineGroupOrder)],
+  (t) => [index('line_group_id_order_index').on(t.lineGroupId, t.lineGroupOrder)],
 )
 
 type LinePayload = DistributedOmit<Lines['content'][number], 'asset' | 'data'>
@@ -70,7 +70,10 @@ export const assetLines = sqliteTable(
     additional: json().$type<LinePayload>(),
     priority: integer(),
   },
-  (t) => [index('priority_index').on(t.priority)],
+  (t) => [
+    index('priority_index').on(t.priority),
+    index('line_id_priority_index').on(t.lineId, t.priority),
+  ],
 )
 
 export const banis = sqliteTable('banis', {