Path: blob/main/SignalServiceKit/Storage/MediaGallery/MediaGalleryAttachmentFinderTest.swift
1 views
//
// Copyright 2024 Signal Messenger, LLC
// SPDX-License-Identifier: AGPL-3.0-only
//
import Foundation
import GRDB
import XCTest
@testable import SignalServiceKit
class MediaGalleryAttachmentFinderTest: XCTestCase {
private let attachmentStore = AttachmentStore()
private var db: InMemoryDB!
override func setUp() async throws {
db = InMemoryDB()
}
// MARK: - Queries
func testQueryDateRange() throws {
let (thread, messageRowId) = insertThreadAndInteraction()
let threadRowId = thread.sqliteRowId!
// Insert one matching content type before the date range
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 100,
contentType: .image,
orderInMessage: 0,
)
// ...and one matching content type after the date range
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 300,
contentType: .image,
orderInMessage: 1,
)
// ...and one non-matching content type within the date range
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 200,
contentType: .audio,
orderInMessage: 2,
)
// ...and two matching content type within the date range
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 200,
contentType: .image,
orderInMessage: 3,
)
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 200,
contentType: .image,
orderInMessage: 4,
)
// ...and one within the date range that we will exclude.
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 200,
contentType: .image,
orderInMessage: 5,
)
// ...and a view once attachment that will be excluded.
try insertAttachment(
messageRowId: messageRowId,
threadRowId: threadRowId,
receivedAtTimestamp: 200,
contentType: .image,
isViewOnce: true,
orderInMessage: 6,
)
let exclusionSet = Set<AttachmentReferenceId>([
.init(ownerId: .messageBodyAttachment(messageRowId: messageRowId), orderInMessage: 5),
])
let finder = MediaGalleryAttachmentFinder(threadId: thread.grdbId!.int64Value, filter: .allPhotoVideoCategory)
// Should get two results with offset 0
var query = finder.galleryItemQuery(
in: .init(
start: .init(millisecondsSince1970: 150),
end: .init(millisecondsSince1970: 250),
),
excluding: exclusionSet,
offset: 0,
ascending: true,
)
var results = try db.read { tx in
return try query.fetchAll(tx.database)
}
XCTAssertEqual(results.count, 2)
XCTAssertEqual(results[0].receivedAtTimestamp, 200)
XCTAssertEqual(results[0].orderInMessage, 3)
XCTAssertEqual(results[1].receivedAtTimestamp, 200)
XCTAssertEqual(results[1].orderInMessage, 4)
// Should get just the second result with offset 1
query = finder.galleryItemQuery(
in: .init(
start: .init(millisecondsSince1970: 150),
end: .init(millisecondsSince1970: 250),
),
excluding: exclusionSet,
offset: 1,
ascending: true,
)
results = try db.read { tx in
return try query.fetchAll(tx.database)
}
XCTAssertEqual(results.count, 1)
XCTAssertEqual(results[0].receivedAtTimestamp, 200)
XCTAssertEqual(results[0].orderInMessage, 4)
}
// MARK: - Index Usage
func testAllQueriesUseIndex() throws {
let (thread, _) = insertThreadAndInteraction()
// Set up some parametrized values for tests.
// Specific values for many things don't matter, just presence
// and combinations thereof.
let dateIntervals: [DateInterval?] = [
nil,
.init(
start: .init(millisecondsSince1970: 100),
end: .init(millisecondsSince1970: 200),
),
]
let exclusionSets: [Set<AttachmentReferenceId>] = [
Set(),
Set([.init(ownerId: .messageBodyAttachment(messageRowId: 100), orderInMessage: nil)]),
Set([.init(ownerId: .messageBodyAttachment(messageRowId: 200), orderInMessage: 5)]),
Set([
.init(ownerId: .messageBodyAttachment(messageRowId: 100), orderInMessage: nil),
.init(ownerId: .messageBodyAttachment(messageRowId: 200), orderInMessage: nil),
.init(ownerId: .messageBodyAttachment(messageRowId: 300), orderInMessage: 5),
]),
]
let offsets: [Int] = [0, 5]
let limits: [Int] = [5, 100]
let ascendings: [Bool] = [true, false]
for filter in AllMediaFilter.allCases {
if filter == .gifs {
// Skip gif filter; it uses a b-tree for sorting and doesn't
// use a simple index.
continue
}
let finder = MediaGalleryAttachmentFinder(threadId: thread.grdbId!.int64Value, filter: filter)
var queries = [QueryInterfaceRequest<RecordType>]()
for dateInterval in dateIntervals {
for exclusionSet in exclusionSets {
for offset in offsets {
for ascending in ascendings {
queries.append(finder.galleryItemQuery(
in: dateInterval,
excluding: exclusionSet,
offset: offset,
ascending: ascending,
))
}
}
}
}
for dateInterval in dateIntervals.compacted() {
for exclusionSet in exclusionSets {
for offset in offsets {
for limit in limits {
queries.append(finder.enumerateMediaAttachmentsQuery(
in: dateInterval,
excluding: exclusionSet,
range: .init(location: offset, length: limit),
))
}
}
}
}
for dateInterval in dateIntervals {
for exclusionSet in exclusionSets {
for limit in limits {
queries.append(finder.enumerateTimestampsQuery(
beforeDate: dateInterval?.end,
afterDate: nil,
excluding: exclusionSet,
count: limit,
ascending: false,
))
queries.append(finder.enumerateTimestampsQuery(
beforeDate: nil,
afterDate: dateInterval?.start,
excluding: exclusionSet,
count: limit,
ascending: true,
))
}
}
}
for limit in limits {
queries.append(finder.recentMediaAttachmentsQuery(limit: limit))
}
try db.read { tx in
for query in queries {
let preparedStatement = try query.makePreparedRequest(tx.database).statement
let queryPlan: [String] = try Row.fetchAll(
tx.database,
sql: "EXPLAIN QUERY PLAN \(preparedStatement.sql);",
arguments: preparedStatement.arguments,
).map { $0["detail"] }
// Ensure we use the relevant indexes and...
// * we use all the columns up to the ordering columns
// * we DONT use expensive B trees for ordering
let allowedQueryPlans: [String] = [
"SEARCH MessageAttachmentReference USING INDEX message_attachment_reference_media_gallery_single_content_type_index (threadRowId=? AND ownerType=? AND contentType=?",
"SEARCH MessageAttachmentReference USING INDEX message_attachment_reference_media_gallery_visualMedia_content_type_index (threadRowId=? AND ownerType=? AND isVisualMediaContentType=?",
"SEARCH MessageAttachmentReference USING INDEX message_attachment_reference_media_gallery_fileOrInvalid_content_type_index (threadRowId=? AND ownerType=? AND isInvalidOrFileContentType=?",
]
XCTAssert(queryPlan.allSatisfy { queryPlan in
for allowedQueryPlan in allowedQueryPlans {
if queryPlan.hasPrefix(allowedQueryPlan) {
return true
}
}
return false
})
// There should NOT be expensive B-TREE usage.
XCTAssert(queryPlan.allSatisfy { $0.contains("USE TEMP B-TREE").negated })
}
}
}
}
// MARK: - Helpers
typealias RecordType = MediaGalleryAttachmentFinder.RecordType
private func insertThreadAndInteraction() -> (thread: TSThread, interactionRowId: Int64) {
let thread = TSThread(uniqueId: UUID().uuidString)
let interaction = TSInteraction(timestamp: 0, receivedAtTimestamp: 0, thread: thread)
db.write { tx in
try! thread.insert(tx.database)
try! interaction.asRecord().insert(tx.database)
}
return (thread, interaction.sqliteRowId!)
}
@discardableResult
private func insertAttachment(
messageRowId: Int64,
threadRowId: Int64,
receivedAtTimestamp: UInt64,
contentType: AttachmentReference.ContentType,
caption: String? = nil,
renderingFlag: AttachmentReference.RenderingFlag = .default,
isViewOnce: Bool = false,
isPastEditRevision: Bool = false,
orderInMessage: UInt32,
idInOwner: UUID? = nil,
) throws -> Attachment.IDType {
let attachmentParams = Attachment.Record.mockStream(
streamInfo: .mock(
contentType: {
switch contentType {
case .invalid:
return .invalid
case .animatedImage:
return .animatedImage(pixelSize: .square(100))
case .audio:
return .audio(duration: 1, waveformRelativeFilePath: nil)
case .file:
return .file
case .image:
return .image(pixelSize: .square(100))
case .video:
return .video(duration: 1, pixelSize: .square(100), stillFrameRelativeFilePath: nil)
}
}(),
),
)
let referenceParams = AttachmentReference.ConstructionParams.mock(
owner: .message(.bodyAttachment(.init(
messageRowId: messageRowId,
receivedAtTimestamp: receivedAtTimestamp,
threadRowId: threadRowId,
contentType: contentType,
isPastEditRevision: isPastEditRevision,
caption: caption,
renderingFlag: renderingFlag,
orderInMessage: orderInMessage,
idInOwner: idInOwner,
isViewOnce: isViewOnce,
))),
)
var attachmentRecord = attachmentParams
try db.write { tx in
try attachmentRecord.insert(tx.database)
attachmentStore.addReference(
referenceParams,
attachmentRowId: attachmentRecord.sqliteId!,
tx: tx,
)
}
return attachmentRecord.sqliteId!
}
}