Merge pull request #3280 from omnivore-app/fix/slow-db-query

execute update query in batch to avoid slow bulk action query
This commit is contained in:
Hongbo Wu 2023-12-29 16:48:18 +08:00 committed by GitHub
commit 42502f9ff1
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 158 additions and 39 deletions

View file

@ -1,5 +1,5 @@
import * as httpContext from 'express-http-context2'
import { EntityManager, EntityTarget, Repository } from 'typeorm'
import { EntityManager, EntityTarget, QueryBuilder, Repository } from 'typeorm'
import { appDataSource } from '../data_source'
import { Claims } from '../resolvers/types'
import { SetClaimsRole } from '../utils/dictionary'
@ -45,3 +45,59 @@ export const authTrx = async <T>(
export const getRepository = <T>(entity: EntityTarget<T>) => {
return appDataSource.getRepository(entity)
}
export const queryBuilderToRawSql = <T>(q: QueryBuilder<T>): string => {
const queryAndParams = q.getQueryAndParameters()
let sql = queryAndParams[0]
const params = queryAndParams[1]
params.forEach((value, index) => {
if (typeof value === 'string') {
sql = sql.replace(`$${index + 1}`, `'${value}'`)
} else if (typeof value === 'object') {
if (Array.isArray(value)) {
sql = sql.replace(
`$${index + 1}`,
value
.map((element) => {
if (typeof element === 'string') {
return `'${element}'`
}
if (typeof element === 'number' || typeof element === 'boolean') {
return element.toString()
}
})
.join(',')
)
} else if (value instanceof Date) {
sql = sql.replace(`$${index + 1}`, `'${value.toISOString()}'`)
}
} else if (typeof value === 'number' || typeof value === 'boolean') {
sql = sql.replace(`$${index + 1}`, value.toString())
}
})
return sql
}
export const valuesToRawSql = (
values: Record<string, string | number | boolean>
): string => {
let sql = ''
Object.keys(values).forEach((key, index) => {
const value = values[key]
if (typeof value === 'string') {
sql += `${key} = '${value}'`
} else {
sql += `${key} = ${value.toString()}`
}
if (index < Object.keys(values).length - 1) {
sql += ', '
}
})
return sql
}

View file

@ -46,6 +46,7 @@ import {
SearchErrorCode,
SearchSuccess,
SetBookmarkArticleError,
SetBookmarkArticleErrorCode,
SetBookmarkArticleSuccess,
SetFavoriteArticleError,
SetFavoriteArticleErrorCode,
@ -70,6 +71,7 @@ import {
findOrCreateLabels,
} from '../../services/labels'
import {
batchUpdateLibraryItems,
createLibraryItem,
findLibraryItemById,
findLibraryItemByUrl,
@ -78,7 +80,6 @@ import {
sortParamsToSort,
updateLibraryItem,
updateLibraryItemReadingProgress,
updateLibraryItems,
} from '../../services/library_item'
import { parsedContentToLibraryItem } from '../../services/save_page'
import {
@ -551,6 +552,10 @@ export const setBookmarkArticleResolver = authorized<
SetBookmarkArticleError,
MutationSetBookmarkArticleArgs
>(async (_, { input: { articleID } }, { uid, log, pubsub }) => {
if (!articleID) {
return { errorCodes: [SetBookmarkArticleErrorCode.NotFound] }
}
// delete the item and its metadata
const deletedLibraryItem = await updateLibraryItem(
articleID,
@ -860,7 +865,7 @@ export const bulkActionResolver = authorized<
labels = await findLabelsByIds(labelIds, uid)
}
await updateLibraryItems(
await batchUpdateLibraryItems(
action,
{
query,

View file

@ -8,7 +8,12 @@ import { Label } from '../entity/label'
import { LibraryItem, LibraryItemState } from '../entity/library_item'
import { BulkActionType, InputMaybe, SortParams } from '../generated/graphql'
import { createPubSubClient, EntityType } from '../pubsub'
import { authTrx, getColumns } from '../repository'
import {
authTrx,
getColumns,
queryBuilderToRawSql,
valuesToRawSql,
} from '../repository'
import { libraryItemRepository } from '../repository/library_item'
import { wordsCount } from '../utils/helpers'
import { parseSearchQuery } from '../utils/search'
@ -886,7 +891,7 @@ export const countByCreatedAt = async (
)
}
export const updateLibraryItems = async (
export const batchUpdateLibraryItems = async (
action: BulkActionType,
searchArgs: SearchArgs,
userId: string,
@ -901,20 +906,21 @@ export const updateLibraryItems = async (
return 'folder' in args
}
const now = new Date().toISOString()
// build the script
let values: QueryDeepPartialEntity<LibraryItem> = {}
let values: Record<string, string | number> = {}
let addLabels = false
switch (action) {
case BulkActionType.Archive:
values = {
archivedAt: new Date(),
archived_at: now,
state: LibraryItemState.Archived,
}
break
case BulkActionType.Delete:
values = {
state: LibraryItemState.Deleted,
deletedAt: new Date(),
deleted_at: now,
}
break
case BulkActionType.AddLabels:
@ -922,9 +928,9 @@ export const updateLibraryItems = async (
break
case BulkActionType.MarkAsRead:
values = {
readAt: new Date(),
readingProgressTopPercent: 100,
readingProgressBottomPercent: 100,
read_at: now,
reading_progress_top_percent: 100,
reading_progress_bottom_percent: 100,
}
break
case BulkActionType.MoveToFolder:
@ -934,7 +940,7 @@ export const updateLibraryItems = async (
values = {
folder: args.folder,
savedAt: new Date(),
saved_at: now,
}
break
@ -973,10 +979,11 @@ export const updateLibraryItems = async (
.map((label) => ({
labelId: label.id,
libraryItemId: libraryItem.id,
name: label.name,
}))
.filter((entityLabel) => {
const existingLabel = libraryItem.labels?.find(
(l) => l.id === entityLabel.labelId
const existingLabel = libraryItem.labelNames?.find(
(l) => l.toLowerCase() === entityLabel.name.toLowerCase()
)
return !existingLabel
})
@ -984,7 +991,34 @@ export const updateLibraryItems = async (
return tx.getRepository(EntityLabel).save(labelsToAdd)
}
return queryBuilder.update(LibraryItem).set(values).execute()
// generate raw sql because postgres doesn't support prepared statements in DO blocks
const countSql = queryBuilderToRawSql(queryBuilder.select('COUNT(1)'))
const subQuery = queryBuilderToRawSql(queryBuilder.select('id'))
const valuesSql = valuesToRawSql(values)
const start = new Date().toISOString()
const batchSize = 1000
const sql = `
-- Set batch size
DO $$
DECLARE
batch_size INT := ${batchSize};
BEGIN
-- Loop through batches
FOR i IN 0..CEIL((${countSql}) * 1.0 / batch_size) - 1 LOOP
-- Update the batch
UPDATE omnivore.library_item
SET ${valuesSql}
WHERE id = ANY(
${subQuery}
AND updated_at < '${start}'
LIMIT batch_size
);
END LOOP;
END $$
`
return tx.query(sql)
})
}

View file

@ -2107,31 +2107,31 @@ describe('Article API', () => {
}
`
before(async () => {
// Create some test items
for (let i = 0; i < 5; i++) {
await createLibraryItem(
{
user,
itemType: i == 0 ? PageType.Article : PageType.File,
title: 'test item',
readableContent: '<p>test</p>',
slug: '',
state:
i == 0 ? LibraryItemState.Failed : LibraryItemState.Succeeded,
originalUrl: `https://blog.omnivore.app/p/bulk-action-${i}`,
},
user.id
)
}
})
after(async () => {
// Delete all items
await deleteLibraryItemsByUserId(user.id)
})
context('when action is MarkAsRead and query is in:unread', () => {
before(async () => {
// Create some test items
for (let i = 0; i < 5; i++) {
await createLibraryItem(
{
user,
itemType: i == 0 ? PageType.Article : PageType.File,
title: 'test item',
readableContent: '<p>test</p>',
slug: '',
state:
i == 0 ? LibraryItemState.Failed : LibraryItemState.Succeeded,
originalUrl: `https://blog.omnivore.app/p/bulk-action-${i}`,
},
user.id
)
}
})
after(async () => {
// Delete all items
await deleteLibraryItemsByUserId(user.id)
})
it('marks unread items as read', async () => {
const res = await graphqlRequest(
bulkActionQuery(BulkActionType.MarkAsRead, 'is:unread'),
@ -2199,6 +2199,30 @@ describe('Article API', () => {
)
context('when action is Delete', () => {
before(async () => {
// Create some test items
for (let i = 0; i < 5; i++) {
await createLibraryItem(
{
user,
itemType: i == 0 ? PageType.Article : PageType.File,
title: 'test item',
readableContent: '<p>test</p>',
slug: '',
state:
i == 0 ? LibraryItemState.Failed : LibraryItemState.Succeeded,
originalUrl: `https://blog.omnivore.app/p/bulk-action-${i}`,
},
user.id
)
}
})
after(async () => {
// Delete all items
await deleteLibraryItemsByUserId(user.id)
})
it('deletes all items', async () => {
const res = await graphqlRequest(
bulkActionQuery(BulkActionType.Delete),