diff --git a/index.ts b/index.ts index a2dc46c..ab25bf0 100644 --- a/index.ts +++ b/index.ts @@ -12,6 +12,7 @@ import { MeiliSearch } from 'meilisearch' import discordScraper from './src/scraper' import testAttachment from "./src/test.attachment" import attachmentService from "./src/attachments" +import memeService from "./src/memes" // @ts-ignore - ts don't know wtf it talking about BigInt.prototype.toJSON = function bigIntStringify() { @@ -19,7 +20,6 @@ BigInt.prototype.toJSON = function bigIntStringify() { }; const VISION_CREDENTIALS = (await import(process.env.GOOGLE_APPLICATION_CREDENTIALS, { assert: { type: "json" } })).default -const STORAGE_BUCKET = process.env.STORAGE_BUCKET const log = pino({ base: {} }); @@ -35,45 +35,20 @@ const meili = new MeiliSearch({ apiKey: process.env.MEILI_MASTER_KEY, }) -const attachments = attachmentService(log, minio, prisma) -const scraper = discordScraper(process.env.DISCORD_AUTH, vision, minio, prisma, log) - -async function ensureStorageBucketExists() { - const exists = await minio.bucketExists(STORAGE_BUCKET) - if (exists) return true - await minio.makeBucket(STORAGE_BUCKET) -} - -async function ensureMeiliIndex() { - const delTask = await meili.deleteIndex('memes') - const delResult = await meili.waitForTask(delTask.taskUid) - const task = await meili.createIndex('memes', { primaryKey: 'id' }) - const result = await meili.waitForTask(task.taskUid) - log.info(result); -} - -async function processChannels(channelIds) { - await ensureStorageBucketExists() - const msgs = await Promise.all(channelIds.map(c => scraper.getChannelAttachments(c))) - const results = await Promise.allSettled(msgs.flat().map(a => attachments.process(a))) - log.info(results.map(r => r.status)); - return results -} - -async function updateMeili() { - await ensureMeiliIndex(); - const index = meili.index('memes') - const memes = await prisma.meme.findMany(); - const task = await index.updateDocuments(memes) - const res = await meili.waitForTask(task.taskUid) - log.info(res) -} +const memes = await memeService(log, prisma, meili) +const attachments = await attachmentService(log, minio, prisma, memes) +const scraper = discordScraper(process.env.DISCORD_AUTH, vision, minio, prisma, log, attachments) const CHANNEL_IDS = [ 367936426962190336n, // TPM shitposts 1095060717117374687n, // aella memes ] -await processChannels(CHANNEL_IDS) -await updateMeili() +async function main() { + await Promise.all(CHANNEL_IDS.map(id => scraper.processChannel(id))) + await memes.updateMeili(); +} + +await main() + diff --git a/package-lock.json b/package-lock.json index 1af4978..e312c17 100644 --- a/package-lock.json +++ b/package-lock.json @@ -11,6 +11,7 @@ "dependencies": { "@google-cloud/vision": "^4.2.0", "@prisma/client": "^5.12.1", + "@theinternetfolks/snowflake": "^1.3.0", "axios": "^1.6.8", "dotenv": "^16.4.1", "lodash": "^4.17.21", @@ -1001,6 +1002,11 @@ "resolved": "https://registry.npmjs.org/@protobufjs/utf8/-/utf8-1.1.0.tgz", "integrity": "sha512-Vvn3zZrhQZkkBE8LSuW3em98c0FwgO4nxzv6OdSxPKJIEKY2bGbHn+mhGIPerzI4twdxaP8/0+06HBpwf345Lw==" }, + "node_modules/@theinternetfolks/snowflake": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/@theinternetfolks/snowflake/-/snowflake-1.3.0.tgz", + "integrity": "sha512-/5fjyE+JIa95Lm1uDJ9fx52W7/ZHBVvKo5GbpR+swpchbriLzVgFAIUOVwLsscnUTBixyQ+JsEh5pJQM/zKNzQ==" + }, "node_modules/@tootallnate/once": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/@tootallnate/once/-/once-2.0.0.tgz", @@ -3460,6 +3466,11 @@ "resolved": "https://registry.npmjs.org/@protobufjs/utf8/-/utf8-1.1.0.tgz", "integrity": "sha512-Vvn3zZrhQZkkBE8LSuW3em98c0FwgO4nxzv6OdSxPKJIEKY2bGbHn+mhGIPerzI4twdxaP8/0+06HBpwf345Lw==" }, + "@theinternetfolks/snowflake": { + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/@theinternetfolks/snowflake/-/snowflake-1.3.0.tgz", + "integrity": "sha512-/5fjyE+JIa95Lm1uDJ9fx52W7/ZHBVvKo5GbpR+swpchbriLzVgFAIUOVwLsscnUTBixyQ+JsEh5pJQM/zKNzQ==" + }, "@tootallnate/once": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/@tootallnate/once/-/once-2.0.0.tgz", diff --git a/package.json b/package.json index 8e4eb9f..45acf16 100644 --- a/package.json +++ b/package.json @@ -15,6 +15,7 @@ "dependencies": { "@google-cloud/vision": "^4.2.0", "@prisma/client": "^5.12.1", + "@theinternetfolks/snowflake": "^1.3.0", "axios": "^1.6.8", "dotenv": "^16.4.1", "lodash": "^4.17.21", diff --git a/src/attachments.ts b/src/attachments.ts index d9916ce..49a5e8e 100644 --- a/src/attachments.ts +++ b/src/attachments.ts @@ -1,13 +1,13 @@ -import sharp from 'sharp'; +import sharp from 'sharp' import axios from 'axios' import mime from 'mime-types' -import type { Logger } from 'pino'; +import type { Logger } from 'pino' import type Minio from 'minio' -import type { PrismaClient, Attachment } from '@prisma/client'; +import type { PrismaClient, Attachment, ChannelHistory } from '@prisma/client' -import * as memes from './memes' +import type { MemeService } from './memes' const STORAGE_BUCKET = process.env.STORAGE_BUCKET @@ -33,7 +33,14 @@ export function objectName(a: Attachment) { return `${a.memeId}.${mime.extension(a.contentType)}` } -export default function attachmentService(log: Logger, minio: Minio.Client, prisma: PrismaClient) { + +export default async function attachmentService(log: Logger, minio: Minio.Client, prisma: PrismaClient, memes: MemeService) { + async function ensureStorageBucketExists() { + const exists = await minio.bucketExists(STORAGE_BUCKET) + if (exists) return true + await minio.makeBucket(STORAGE_BUCKET) + } + async function storeAttachment(a: Attachment) { const download = await axios.get(a.url, { responseType: 'stream' }) @@ -53,20 +60,48 @@ export default function attachmentService(log: Logger, minio: Minio.Client, pris return minio.removeObject(STORAGE_BUCKET, objectName(a)) } - async function process(a: Attachment) { - return prisma.$transaction(async prisma => { - const meme = memes.fromAttachment(a) - const dbMeme = await prisma.meme.upsert({ create: meme, update: meme, where: { id: meme.id } }) - const dbA = await prisma.attachment.upsert({ create: a, update: a, where: { id: a.id } }) - log.info({ msg: 'pending addition to db', dbMeme }) - const upload = await storeAttachment(a) - log.info({ msg: 'successfully uploaded', upload }) - return { meme: dbMeme, upload, attachment: dbA } - }) + async function process(a: Attachment, history: Partial) { + try { + const results = await prisma.$transaction(async prisma => { + const meme = memes.fromAttachment(a) + const dbMeme = await prisma.meme.upsert({ create: meme, update: meme, where: { id: meme.id } }) + const dbA = await prisma.attachment.upsert({ create: a, update: a, where: { id: a.id } }) + log.debug({ msg: 'pending addition to db', dbMeme }) + + const upload = await storeAttachment(a) + log.debug({ msg: 'successfully uploaded', upload }) + + await memes.persist(dbMeme) + + return { meme: dbMeme, upload, attachment: dbA } + }, { timeout: 60000 }) + } catch (e) { + + // this happens pretty regularly it seems + // when ppl reply to an attachment it is included in their message + if (e.meta.modelName !== 'Meme' && e.meta.target[0] !== 'attachmentId') { + log.error(e, 'failed to process attachment') + log.error(a, 'failed to process attachment') + } + + // deletes the upload if we got that far + // (probably not) + await deleteAttachment(a) + } + + // we assume here that a failure in processing was caused by a duplicate + // so we update pg to say we've processed it + const newHistory = await prisma.channelHistory.update({ data: history, where: { id: history.id }}) + log.debug({ newHistory }, 'updated channel history') + return true } + await ensureStorageBucketExists() + return { process, } } + +export type AttachmentService = Awaited> diff --git a/src/memes.ts b/src/memes.ts index c36d98d..bc3ddd1 100644 --- a/src/memes.ts +++ b/src/memes.ts @@ -1,4 +1,6 @@ import type { PrismaClient, Attachment, Meme } from '@prisma/client'; +import type { Meilisearch } from 'meilisearch'; +import type { Logger } from 'pino'; import * as attachments from './attachments' @@ -10,23 +12,62 @@ function memeUrl(a: Attachment) { return u.href } -export function fromAttachment(a: Attachment) { - const text = a.vision.fullTextAnnotation?.text || '' - const tags = [ - a.guild.name, - a.channel.name, - a.message.author.username, - ] +export default async function memeService(baseLog: Logger, prisma: PrismaClient, meili: Meilisearch) { + const log = baseLog.child({ service: 'meme' }) - const m: Meme = { - id: a.memeId, - url: memeUrl(a), - text, - tags: tags, - attachmentId: a.id, - createdAt: undefined, - updatedAt: undefined + function fromAttachment(a: Attachment) { + const text = a.vision.fullTextAnnotation?.text || '' + const tags = [ + a.guild.name, + a.channel.name, + a.message.author.username, + ] + + const m: Meme = { + id: a.memeId, + url: memeUrl(a), + text, + tags: tags, + attachmentId: a.id, + createdAt: undefined, + updatedAt: undefined + } + + return m } - return m + async function ensureMeiliIndex() { + const delTask = await meili.deleteIndex('memes') + const delResult = await meili.waitForTask(delTask.taskUid) + const task = await meili.createIndex('memes', { primaryKey: 'id' }) + const result = await meili.waitForTask(task.taskUid) + log.debug(result); + return meili.index('memes') + } + + async function updateMeili() { + await ensureMeiliIndex(); + const index = meili.index('memes') + const memes = await prisma.meme.findMany(); + const task = await index.updateDocuments(memes) + const res = await meili.waitForTask(task.taskUid) + log.debug(res) + } + + async function persist(m: Meme) { + const index = meili.index('memes') + return index.updateDocuments([m]) + } + + await ensureMeiliIndex(); + await updateMeili(); + + return { + ensureMeiliIndex, + updateMeili, + fromAttachment, + persist, + } } + +export type MemeService = Awaited> diff --git a/src/scraper.test.ts b/src/scraper.test.ts index 1d08bb5..fa289e3 100644 --- a/src/scraper.test.ts +++ b/src/scraper.test.ts @@ -41,4 +41,4 @@ const CHANNEL_IDS = [ 1095060717117374687n, // aella memes ] -await Promise.all(CHANNEL_IDS.map(id => scraper.getChannelHistory(id))) +await Promise.all(CHANNEL_IDS.map(id => scraper.processChannel(id))) diff --git a/src/scraper.ts b/src/scraper.ts index 7d89dfa..05d972c 100644 --- a/src/scraper.ts +++ b/src/scraper.ts @@ -1,7 +1,6 @@ import axios from 'axios' -import minBy from 'lodash/minBy' -import maxBy from 'lodash/maxBy' +import _ from 'lodash'; import type { PrismaClient, Attachment } from '@prisma/client'; import type { ImageAnnotatorClient } from '@google-cloud/vision'; @@ -10,27 +9,39 @@ import type Minio from 'minio' import sharp from 'sharp'; -import * as attachments from './attachments' +import type { AttachmentService } from './attachments' +import { IMAGE_TYPES, VIDEO_TYPES } from './attachments' -export default function scraper(discordAuth, vision: ImageAnnotatorClient, minio: Minio.Client, prisma: PrismaClient, baseLog: Logger) { +// TODO +// also fetch embeds and their descriptions + +export default function scraper( + discordAuth, + vision: ImageAnnotatorClient, + minio: Minio.Client, + prisma: PrismaClient, + baseLog: Logger, + attachments: AttachmentService +) { const headers = { 'Authorization': discordAuth, } async function fetchChannelMessages(channelId) { - const url = `https://discord.com/api/v9/channels/${channelId}/messages` + const url = `https://discord.com/api/v9/channels/${channelId}/messages&limit=100` const res = await axios.get(url, { headers }) return res.data } + // returns msgs in desc id order async function fetchChannelMessagesAfter(channelId, lastMessageId) { - const url = `https://discord.com/api/v9/channels/${channelId}/messages?after=${lastMessageId}` + const url = `https://discord.com/api/v9/channels/${channelId}/messages?after=${lastMessageId}&limit=100` const res = await axios.get(url, { headers }) return res.data } async function fetchChannelMessagesBefore(channelId, earliestMessageId) { - const url = `https://discord.com/api/v9/channels/${channelId}/messages?before=${earliestMessageId}` + const url = `https://discord.com/api/v9/channels/${channelId}/messages?before=${earliestMessageId}&limit=100` const res = await axios.get(url, { headers }) return res.data } @@ -38,8 +49,8 @@ export default function scraper(discordAuth, vision: ImageAnnotatorClient, minio async function fetchGuildDetails(guildId) { const url = `https://discord.com/api/v9/guilds/${guildId}` const res = await axios.get(url, { headers }) - const keys_to_remove = ['emojis', 'stickers', 'roles'] - keys_to_remove.forEach(k => delete res[k]) + const keys_to_remove = ['emojis', 'stickers', 'roles', 'features'] + keys_to_remove.forEach(k => delete res.data[k]) return res.data } @@ -51,7 +62,7 @@ export default function scraper(discordAuth, vision: ImageAnnotatorClient, minio async function detectText(attachment: Attachment) { // bail out of anything not an image - if (!attachments.IMAGE_TYPES.includes(attachment.contentType)) return {} + if (!IMAGE_TYPES.includes(attachment.contentType)) return {} const annotateRequest = { image: { @@ -66,101 +77,79 @@ export default function scraper(discordAuth, vision: ImageAnnotatorClient, minio return res[0] || {} } - async function filterAttachments(log, messages, guild, channel) { - const acceptedContentTypes = [ - ...attachments.IMAGE_TYPES, - ...attachments.VIDEO_TYPES, - ] - - const validAttachments = messages - .filter(m => m.attachments) - // bundle message and attachment together - .flatMap(m => m.attachments.map(a => [a, m])) - .filter(([a, m]) => acceptedContentTypes.includes(a.content_type)) - - const newAttachments = validAttachments.filter(([a, m]) => ![].includes(a.id)) - log.info({ newAttachments: newAttachments.length, validAttachments: validAttachments.length }) - - return Promise.all(newAttachments - .map(async ([a, m]) => { - const attachment: Attachment = { - id: BigInt(a.id), - msgId: BigInt(m.id), - channelId: BigInt(m.channel_id), - url: a.url, - contentType: a.content_type, - message: m, - - // create this ahead of time to associate with our copy of the media - memeId: crypto.randomUUID(), - - // @ts-ignore it's fine - vision: {}, - guild, - channel, - - createdAt: undefined, - updatedAt: undefined - } - - const visionData = await detectText(attachment) - return { ...attachment, vision: visionData } - })) - } - - async function getChannelAttachments(channelId) { - try { - const channel = await fetchChannel(channelId) - const guild = await fetchGuildDetails(channel.guild_id) - const messages = await fetchChannelMessages(channelId) - - const log = baseLog.child({ guild: guild.name, channel: channel.name, channelId }) - - return filterAttachments(log, messages, guild, channel); - - // return { messages, channel, guild } - } catch (e) { - baseLog.error({ err: e.toJSON(), channelId, }) - } - } - - async function* fetchNew(channelId, log, latestMessageId) { - const messages = await fetchChannelMessagesAfter(channelId, latestMessageId) - if (messages.length === 0) { - log.info('no new messages') - return - } - - log.info({ messages: messages.length, latestMessageId }, 'fetched new messages') - yield* messages - - // not sure what order these are in - // or if guaranteed - const cursor = maxBy(messages, 'id').id - yield* fetchNew(channelId, log, cursor) - } - - async function* fetchOld(channelId, log, earliestMessageId) { - const messages = await fetchChannelMessagesBefore(channelId, earliestMessageId) - if (messages.length === 0) { - log.info('no new messages') - return - } - - log.info({ messages: messages.length, earliestMessageId }, 'fetched older messages') - yield* messages - - // not sure what order these are in - // or if guaranteed - const cursor = minBy(messages, 'id').id - yield* fetchNew(channelId, log, cursor) - } - - async function getChannelHistory(channelId) { + async function processChannel(channelId) { const channel = await fetchChannel(channelId) const guild = await fetchGuildDetails(channel.guild_id) const log = baseLog.child({ guild: guild.name, channel: channel.name, channelId }) + async function filterAttachments(messages): Promise { + const acceptedContentTypes = [ + ...IMAGE_TYPES, + ...VIDEO_TYPES, + ] + + const validAttachments = messages + .filter(m => m.attachments) + // bundle message and attachment together + .flatMap(m => m.attachments.map(a => [a, m])) + .filter(([a, m]) => acceptedContentTypes.includes(a.content_type)) + + const newAttachments = validAttachments.filter(([a, m]) => ![].includes(a.id)) + log.info({ newAttachments: newAttachments.length, validAttachments: validAttachments.length }) + + return Promise.all(newAttachments + .map(async ([a, m]) => { + const attachment: Attachment = { + id: BigInt(a.id), + msgId: BigInt(m.id), + channelId: BigInt(m.channel_id), + url: a.url, + contentType: a.content_type, + + // create this ahead of time to associate with our copy of the media + memeId: crypto.randomUUID(), + + // @ts-ignore it's fine + vision: {}, // await detectText(attachment), + guild, + channel, + message: m, + + createdAt: undefined, + updatedAt: undefined + } + return attachment + })) + } + + async function* fetchNew(channelId, log, latestMessageId): AsyncGenerator { + // reverse the messages because we want to yield them in chronological order + // the api call returns them most recent first + const messages = (await fetchChannelMessagesAfter(channelId, latestMessageId)).reverse() + if (messages.length === 0) { + log.info('no new messages') + return + } + + log.info({ messages: messages.length, latestMessageId }, 'fetched new messages') + yield* await filterAttachments(messages) + const cursor = _.last(messages).id + yield* fetchNew(channelId, log, cursor) + } + + async function* fetchOld(channelId, log, earliestMessageId): AsyncGenerator { + const messages = await fetchChannelMessagesBefore(channelId, earliestMessageId) + if (messages.length === 0) { + log.info('no new messages') + return + } + + log.info({ messages: messages.length, earliestMessageId }, 'fetched older messages') + yield* await filterAttachments(messages) + const cursor = _.last(messages).id + yield* fetchOld(channelId, log, cursor) + } + const ensureHistory = async () => { const stored = await prisma.channelHistory.findFirst({ where: { id: channelId }}) if (stored) return stored @@ -172,19 +161,24 @@ export default function scraper(discordAuth, vision: ImageAnnotatorClient, minio } const history = await ensureHistory() - history.latestMessageId = 1227182587177340938n - for await (const msg of fetchNew(channelId, log, history.latestMessageId)) { - log.info({ msg: 'new message', id: msg.id}) + // we can do this better if we yield a list of attachments instead of one at a time + // will have to only update the history if they all succeed, probably better anyway + for await (const attachment of fetchNew(channelId, log, history.latestMessageId)) { + log.debug({ msg: 'scraped new attachment', attachment: _.omit(attachment, ['guild', 'message', 'channel', 'vision']) }) + const newHistory = { id: history.id, latestMessageId: attachment.msgId } + await attachments.process(attachment, newHistory) } - // for await (const msg of fetchOld(channelId, log, history.earliestMessageId)) { - // // log.info({ msg: 'new message', id: msg.id}) - // } + // assuming history of earliestId doesn't change... + for await (const attachment of fetchOld(channelId, log, history.earliestMessageId)) { + log.debug({ msg: 'scraped old attachment', attachment: _.omit(attachment, ['guild', 'message', 'channel', 'vision']) }) + const newHistory = { id: history.id, earliestMessageId: attachment.msgId } + await attachments.process(attachment, newHistory) + } } return { - getChannelAttachments, - getChannelHistory, + processChannel, } }