feat(store/labels): add ability to store and query labels

src/Store/make-in-memory-store.ts

@@ -1,13 +1,16 @@
-import type KeyedDB from '@adiwajshing/keyed-db'
+import KeyedDB from '@adiwajshing/keyed-db'
 import type { Comparable } from '@adiwajshing/keyed-db/lib/Types'
 import type { Logger } from 'pino'
 import { proto } from '../../WAProto'
 import { DEFAULT_CONNECTION_CONFIG } from '../Defaults'
 import type makeMDSocket from '../Socket'
 import type { BaileysEventEmitter, Chat, ConnectionState, Contact, GroupMetadata, PresenceData, WAMessage, WAMessageCursor, WAMessageKey } from '../Types'
+import { Label } from '../Types/Label'
+import { LabelAssociation, LabelAssociationType, MessageLabelAssociation } from '../Types/LabelAssociation'
 import { toNumber, updateMessageWithReaction, updateMessageWithReceipt } from '../Utils'
 import { jidNormalizedUser } from '../WABinary'
 import makeOrderedDictionary from './make-ordered-dictionary'
+import { ObjectRepository } from './object-repository'
 
 type WASocket = ReturnType<typeof makeMDSocket>
 
@@ -18,26 +21,74 @@ export const waChatKey = (pin: boolean) => ({
 
 export const waMessageID = (m: WAMessage) => m.key.id || ''
 
+export const waLabelAssociationKey: Comparable<LabelAssociation, string> = {
+	key: (la: LabelAssociation) => (la.type === LabelAssociationType.Chat ? la.chatId + la.labelId : la.chatId + la.messageId + la.labelId),
+	compare: (k1: string, k2: string) => k2.localeCompare(k1)
+}
+
 export type BaileysInMemoryStoreConfig = {
 	chatKey?: Comparable<Chat, string>
+	labelAssociationKey?: Comparable<LabelAssociation, string>
 	logger?: Logger
 }
 
 const makeMessagesDictionary = () => makeOrderedDictionary(waMessageID)
 
-export default (
+const predefinedLabels = Object.freeze<Record<string, Label>>({
-	{ logger: _logger, chatKey }: BaileysInMemoryStoreConfig
+	'0': {
-) => {
+		id: '0',
-	const logger = _logger || DEFAULT_CONNECTION_CONFIG.logger.child({ stream: 'in-mem-store' })
+		name: 'New customer',
-	chatKey = chatKey || waChatKey(true)
+		predefinedId: '0',
-	const KeyedDB = require('@adiwajshing/keyed-db').default as new (...args: any[]) => KeyedDB<Chat, string>
+		color: 0,
+		deleted: false
+	},
+	'1': {
+		id: '1',
+		name: 'New order',
+		predefinedId: '1',
+		color: 1,
+		deleted: false
+	},
+	'2': {
+		id: '2',
+		name: 'Pending payment',
+		predefinedId: '2',
+		color: 2,
+		deleted: false
+	},
+	'3': {
+		id: '3',
+		name: 'Paid',
+		predefinedId: '3',
+		color: 3,
+		deleted: false
+	},
+	'4': {
+		id: '4',
+		name: 'Order completed',
+		predefinedId: '4',
+		color: 4,
+		deleted: false
+	}
+})
 
-	const chats = new KeyedDB(chatKey, c => c.id)
+export default (
+	{ logger: _logger, chatKey, labelAssociationKey }: BaileysInMemoryStoreConfig
+) => {
+	// const logger = _logger || DEFAULT_CONNECTION_CONFIG.logger.child({ stream: 'in-mem-store' })
+	chatKey = chatKey || waChatKey(true)
+	labelAssociationKey = labelAssociationKey || waLabelAssociationKey
+	const logger = _logger || DEFAULT_CONNECTION_CONFIG.logger.child({ stream: 'in-mem-store' })
+	// const KeyedDB = require('@adiwajshing/keyed-db').default as new (...args: any[]) => KeyedDB<Chat, string>
+
+	const chats = new KeyedDB<Chat, string>(chatKey, c => c.id)
 	const messages: { [_: string]: ReturnType<typeof makeMessagesDictionary> } = {}
 	const contacts: { [_: string]: Contact } = {}
 	const groupMetadata: { [_: string]: GroupMetadata } = {}
 	const presences: { [id: string]: { [participant: string]: PresenceData } } = {}
 	const state: ConnectionState = { connection: 'close' }
+	const labels = new ObjectRepository<Label>(predefinedLabels)
+	const labelAssociations = new KeyedDB<LabelAssociation, string>(labelAssociationKey, labelAssociationKey.key)
 
 	const assertMessageList = (jid: string) => {
 		if (!messages[jid]) {
@@ -60,6 +111,12 @@ export default (
 		return oldContacts
 	}
 
+	const labelsUpsert = (newLabels: Label[]) => {
+		for (const label of newLabels) {
+			labels.upsertById(label.id, label)
+		}
+	}
+
 	/**
 	 * binds to a BaileysEventEmitter.
 	 * It listens to all events and constructs a state that you can query accurate data from.
@@ -131,6 +188,33 @@ export default (
 				}
 			}
 		})
+
+		ev.on('labels.edit', (label: Label) => {
+			if (label.deleted) {
+				return labels.deleteById(label.id)
+			}
+
+			// WhatsApp can store only up to 20 labels
+			if (labels.count() < 20) {
+				return labels.upsertById(label.id, label)
+			}
+
+			logger.error('Labels count exceed')
+		})
+
+		ev.on('labels.association', ({ type, association }) => {
+			switch (type) {
+				case 'add':
+					labelAssociations.upsert(association)
+					break
+				case 'remove':
+					labelAssociations.delete(association)
+					break
+				default:
+					console.error(`unknown operation type [${type}]`)
+			}
+		})
+
 		ev.on('presence.update', ({ id, presences: update }) => {
 			presences[id] = presences[id] || {}
 			Object.assign(presences[id], update)
@@ -246,12 +330,22 @@ export default (
 	const toJSON = () => ({
 		chats,
 		contacts,
-		messages
+		messages,
+		labels,
+		labelAssociations
 	})
 
-	const fromJSON = (json: { chats: Chat[], contacts: { [id: string]: Contact }, messages: { [id: string]: WAMessage[] } }) => {
+	const fromJSON = (json: {
+		chats: Chat[],
+		contacts: { [id: string]: Contact },
+		messages: { [id: string]: WAMessage[] },
+		labels: { [labelId: string]: Label },
+		labelAssociations: LabelAssociation[]
+	}) => {
 		chats.upsert(...json.chats)
+		labelAssociations.upsert(...json.labelAssociations || [])
 		contactsUpsert(Object.values(json.contacts))
+		labelsUpsert(Object.values(json.labels || {}))
 		for (const jid in json.messages) {
 			const list = assertMessageList(jid)
 			for (const msg of json.messages[jid]) {
@@ -268,6 +362,8 @@ export default (
 		groupMetadata,
 		state,
 		presences,
+		labels,
+		labelAssociations,
 		bind,
 		/** loads messages from the store, if not found -- uses the legacy connection */
 		loadMessages: async (jid: string, count: number, cursor: WAMessageCursor) => {
@@ -295,6 +391,38 @@ export default (
 
 			return messages
 		},
+		/**
+		 * Get all available labels for profile
+		 *
+		 * Keep in mind that the list is formed from predefined tags and tags
+		 * that were "caught" during their editing.
+		 */
+		getLabels: () => {
+			return labels
+		},
+
+		/**
+		 * Get labels for chat
+		 *
+		 * @returns Label IDs
+		 **/
+		getChatLabels: (chatId: string) => {
+			return labelAssociations.filter((la) => la.chatId === chatId).all()
+		},
+
+		/**
+		 * Get labels for message
+		 *
+		 * @returns Label IDs
+		 **/
+		getMessageLabels: (messageId: string) => {
+			const associations = labelAssociations
+				.filter((la: MessageLabelAssociation) => la.messageId === messageId)
+				.all()
+
+			return associations.map(({ labelId }) => labelId)
+
+		},
 		loadMessage: async (jid: string, id: string) => messages[jid]?.get(id),
 		mostRecentMessage: async (jid: string) => {
 			const message: WAMessage | undefined = messages[jid]?.array.slice(-1)[0]

src/Store/object-repository.ts

@@ -0,0 +1,31 @@
+export class ObjectRepository<T extends object> {
+    readonly entityMap: Map<string, T>
+
+    constructor(entities: Record<string, T> = {}) {
+        this.entityMap = new Map(Object.entries(entities))
+    }
+
+    findById(id: string) {
+        return this.entityMap.get(id)
+    }
+
+    findAll() {
+        return Array.from(this.entityMap.values())
+    }
+
+    upsertById(id: string, entity: T) {
+        return this.entityMap.set(id, { ...entity })
+    }
+
+    deleteById(id: string) {
+        return this.entityMap.delete(id)
+    }
+
+    count() {
+        return this.entityMap.size
+    }
+
+    toJSON() {
+        return this.findAll()
+    }
+}