@belopash/typeorm-store 0.0.0

package/src/cacheMap.ts

@@ -0,0 +1,135 @@
+import {Entity} from '@subsquid/typeorm-store'
+import assert from 'assert'
+import {EntityManager, EntityTarget, FindOptionsRelations} from 'typeorm'
+import {copy} from './utils'
+
+export class CachedEntity<E extends Entity> {
+    value: E | null
+    relations: {[key: string]: boolean}
+
+    constructor() {
+        this.value = null
+        this.relations = {}
+    }
+}
+
+export class CacheMap {
+    private map: Map<string, Map<string, CachedEntity<any>>> = new Map()
+
+    constructor(private em: () => EntityManager) {}
+
+    exist<E extends Entity>(entityClass: EntityTarget<E>, id: string) {
+        const cacheMap = this.getEntityCache(entityClass)
+        const cachedEntity = cacheMap.get(id)
+        return cachedEntity?.value != null
+    }
+
+    get<E extends Entity>(entityClass: EntityTarget<E>, id: string) {
+        const cacheMap = this.getEntityCache(entityClass)
+        return cacheMap.get(id) as CachedEntity<E> | undefined
+    }
+
+    ensure<E extends Entity>(entityClass: EntityTarget<E>, id: string) {
+        const cacheMap = this.getEntityCache(entityClass)
+
+        let cachedEntity = cacheMap.get(id)
+        if (cachedEntity == null) {
+            cachedEntity = new CachedEntity()
+            cacheMap.set(id, cachedEntity)
+        }
+    }
+
+    delete<E extends Entity>(entityClass: EntityTarget<E>, id: string) {
+        const cacheMap = this.getEntityCache(entityClass)
+
+        const cachedEntity = new CachedEntity()
+        cacheMap.set(id, cachedEntity)
+    }
+
+    add<E extends Entity>(entity: E, mask?: FindOptionsRelations<any>): void
+    add<E extends Entity>(entities: E[], mask?: FindOptionsRelations<any>): void
+    add<E extends Entity>(e: E | E[], mask: FindOptionsRelations<any> = {}) {
+        const em = this.em()
+
+        const entities = Array.isArray(e) ? e : [e]
+        if (entities.length == 0) return
+
+        const entityClass = entities[0].constructor
+        const metadata = em.connection.getMetadata(entities[0].constructor)
+
+        const cacheMap = this.getEntityCache(metadata.target)
+
+        for (const entity of entities) {
+            let cachedEntity = cacheMap.get(entity.id)
+            if (cachedEntity == null) {
+                cachedEntity = new CachedEntity()
+                cacheMap.set(entity.id, cachedEntity)
+            }
+
+            if (cachedEntity.value == null) {
+                cachedEntity.value = em.create(entityClass)
+            }
+
+            for (const column of metadata.nonVirtualColumns) {
+                const objectColumnValue = column.getEntityValue(entity)
+                if (objectColumnValue !== undefined) {
+                    column.setEntityValue(cachedEntity.value, copy(objectColumnValue))
+                }
+            }
+
+            for (const relation of metadata.relations) {
+                const relatedMetadata = relation.inverseEntityMetadata
+                const relatedEntity = relation.getEntityValue(entity) as Entity | null | undefined
+
+                const relatedMask = mask[relation.propertyName]
+                if (relatedMask) {
+                    if (relation.isOneToMany || relation.isManyToMany) {
+                        if (Array.isArray(relatedEntity)) {
+                            for (const r of relatedEntity) {
+                                this.add(r, typeof relatedMask === 'boolean' ? {} : relatedMask)
+                            }
+                        }
+                    } else if (relatedEntity != null) {
+                        this.add(relatedEntity, typeof relatedMask === 'boolean' ? {} : relatedMask)
+                    }
+                }
+
+                if (relation.isOwning && relatedMask) {
+                    if (relatedEntity == null) {
+                        relation.setEntityValue(cachedEntity.value, null)
+                    } else {
+                        const _relationCacheMap = this.getEntityCache(relatedMetadata.target)
+                        const cachedRelation = _relationCacheMap.get(relatedEntity.id)
+                        assert(
+                            cachedRelation != null,
+                            `missing entity ${relatedMetadata.name} with id ${relatedEntity.id}`
+                        )
+
+                        const relatedEntityIdOnly = em.create(relatedMetadata.target, {id: relatedEntity.id})
+                        relation.setEntityValue(cachedEntity.value, relatedEntityIdOnly)
+                    }
+                }
+            }
+        }
+    }
+
+    clear() {
+        for (const item of this.map.values()) {
+            item.clear()
+        }
+        this.map.clear()
+    }
+
+    private getEntityCache(entityClass: EntityTarget<any>) {
+        const em = this.em()
+        const metadata = em.connection.getMetadata(entityClass)
+
+        let map = this.map.get(metadata.name)
+        if (map == null) {
+            map = new Map()
+            this.map.set(metadata.name, map)
+        }
+
+        return map
+    }
+}

package/src/database.ts

@@ -0,0 +1,40 @@
+import {IsolationLevel, TypeormDatabase, TypeormDatabaseOptions} from '@subsquid/typeorm-store'
+import {ChangeTracker} from '@subsquid/typeorm-store/lib/hot'
+import {FinalTxInfo, HotTxInfo, HashAndHeight} from '@subsquid/typeorm-store/lib/interfaces'
+import assert from 'assert'
+import {EntityManager} from 'typeorm'
+import {StoreWithCache} from './store'
+
+export {IsolationLevel, TypeormDatabaseOptions}
+
+// @ts-ignore
+export class TypeormDatabaseWithCache extends TypeormDatabase {
+    // @ts-ignore
+    transact(info: FinalTxInfo, cb: (store: StoreWithCache) => Promise<void>): Promise<void> {
+        return super.transact(info, cb as any)
+    }
+
+    // @ts-ignore
+    transactHot(info: HotTxInfo, cb: (store: StoreWithCache, block: HashAndHeight) => Promise<void>): Promise<void> {
+        return super.transactHot(info, cb as any)
+    }
+
+    private async performUpdates(
+        cb: (store: StoreWithCache) => Promise<void>,
+        em: EntityManager,
+        changeTracker?: ChangeTracker
+    ): Promise<void> {
+        let running = true
+
+        let store = new StoreWithCache(() => {
+            assert(running, `too late to perform db updates, make sure you haven't forgot to await on db query`)
+            return em
+        }, changeTracker)
+
+        try {
+            await cb(store)
+        } finally {
+            running = false
+        }
+    }
+}

package/src/index.ts

@@ -0,0 +1,2 @@
+export * from './database'
+export {EntityClass, Entity, FindManyOptions, FindOneOptions, StoreWithCache} from './store'

package/src/store.ts

@@ -0,0 +1,474 @@
+import {Entity as _Entity, Entity, EntityClass, FindManyOptions, FindOneOptions, Store} from '@subsquid/typeorm-store'
+import {ChangeTracker} from '@subsquid/typeorm-store/lib/hot'
+import {def} from '@subsquid/util-internal'
+import assert from 'assert'
+import {Graph} from 'graph-data-structure'
+import {EntityManager, EntityTarget, FindOptionsRelations, FindOptionsWhere, In} from 'typeorm'
+import {copy, splitIntoBatches} from './utils'
+import {CacheMap} from './cacheMap'
+import {UpdateMap, UpdateType} from './updateMap'
+import {RelationMetadata} from 'typeorm/metadata/RelationMetadata'
+
+export {EntityClass, FindManyOptions, FindOneOptions, Entity}
+
+export type DeferMap = Map<string, {ids: Set<string>; relations: FindOptionsRelations<any>}>
+export interface ChangeSet {
+    inserts: Entity[]
+    upserts: Entity[]
+    delayedUpserts: Entity[]
+    removes: Entity[]
+}
+
+// @ts-ignore
+export class StoreWithCache extends Store {
+    private deferMap: DeferMap = new Map()
+    private updates: Map<string, UpdateMap> = new Map()
+    private cache: CacheMap
+
+    constructor(private em: () => EntityManager, changes?: ChangeTracker) {
+        super(em, changes)
+        this.cache = new CacheMap(em)
+    }
+
+    async insert<E extends _Entity>(entity: E): Promise<void>
+    async insert<E extends _Entity>(entities: E[]): Promise<void>
+    async insert<E extends _Entity>(e: E | E[]): Promise<void> {
+        const em = this.em()
+
+        const entities = Array.isArray(e) ? e : [e]
+        if (entities.length == 0) return
+
+        const entityClass = entities[0].constructor
+        const metadata = em.connection.getMetadata(entityClass)
+
+        const relationMask: FindOptionsRelations<any> = {}
+        for (const relation of metadata.relations) {
+            if (relation.isOwning) {
+                relationMask[relation.propertyName] = true
+            }
+        }
+
+        const updateMap = this.getUpdateMap(entityClass)
+        for (const entity of entities) {
+            updateMap.insert(entity.id)
+            this.cache.add(entity, relationMask)
+        }
+    }
+
+    async upsert<E extends _Entity>(entity: E): Promise<void>
+    async upsert<E extends _Entity>(entities: E[]): Promise<void>
+    async upsert<E extends _Entity>(e: E | E[]): Promise<void> {
+        const em = this.em()
+
+        let entities = Array.isArray(e) ? e : [e]
+        if (entities.length == 0) return
+
+        const entityClass = entities[0].constructor
+        const metadata = em.connection.getMetadata(entityClass)
+
+        const updateMap = this.getUpdateMap(entityClass)
+        for (const entity of entities) {
+            const relationMask: FindOptionsRelations<any> = {}
+            for (const relation of metadata.relations) {
+                const relatedEntity = relation.getEntityValue(entity) as Entity | null | undefined
+
+                if (relation.isOwning && relatedEntity !== undefined) {
+                    relationMask[relation.propertyName] = true
+                }
+            }
+
+            updateMap.upsert(entity.id)
+            this.cache.add(entity, relationMask)
+        }
+    }
+
+    async save<E extends _Entity>(entity: E): Promise<void>
+    async save<E extends _Entity>(entities: E[]): Promise<void>
+    async save<E extends _Entity>(e: E | E[]): Promise<void> {
+        return await this.upsert(e as any)
+    }
+
+    async remove<E extends Entity>(entity: E): Promise<void>
+    async remove<E extends Entity>(entities: E[]): Promise<void>
+    async remove<E extends Entity>(entityClass: EntityTarget<E>, id: string | string[]): Promise<void>
+    async remove<E extends Entity>(e: E | E[] | EntityTarget<E>, id?: string | string[]): Promise<void> {
+        const em = this.em()
+
+        if (id == null) {
+            const entities = Array.isArray(e) ? e : [e as E]
+            if (entities.length == 0) return
+
+            const entityClass = entities[0].constructor
+            const updateMap = this.getUpdateMap(entityClass)
+
+            for (const entity of entities) {
+                updateMap.remove(entity.id)
+                this.cache.delete(entityClass, entity.id)
+            }
+        } else {
+            const ids = Array.isArray(id) ? id : [id]
+            if (ids.length == 0) return
+
+            const entityClass = e as EntityTarget<E>
+            const updateMap = this.getUpdateMap(entityClass)
+
+            for (const i of ids) {
+                updateMap.remove(i)
+                this.cache.delete(entityClass, i)
+            }
+        }
+    }
+
+    async count<E extends Entity>(entityClass: EntityTarget<E>, options?: FindManyOptions<E>): Promise<number> {
+        await this.persist()
+        return await super.count(entityClass as EntityClass<E>, options)
+    }
+
+    async countBy<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        where: FindOptionsWhere<E> | FindOptionsWhere<E>[]
+    ): Promise<number> {
+        await this.persist()
+        return await super.countBy(entityClass as EntityClass<E>, where)
+    }
+
+    async find<E extends Entity>(entityClass: EntityTarget<E>, options: FindManyOptions<E>): Promise<E[]> {
+        await this.persist()
+        const res = await super.find(entityClass as EntityClass<E>, options)
+        if (res != null) this.cache.add(res, options.relations)
+        return res
+    }
+
+    async findBy<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        where: FindOptionsWhere<E> | FindOptionsWhere<E>[]
+    ): Promise<E[]> {
+        await this.persist()
+        const res = await super.findBy(entityClass as EntityClass<E>, where)
+        if (res != null) this.cache.add(res)
+        return res
+    }
+
+    async findOne<E extends Entity>(entityClass: EntityTarget<E>, options: FindOneOptions<E>): Promise<E | undefined> {
+        await this.persist()
+        const res = await super.findOne(entityClass as EntityClass<E>, options)
+        if (res != null) this.cache.add(res, options.relations)
+        return res
+    }
+
+    async findOneOrFail<E extends Entity>(entityClass: EntityTarget<E>, options: FindOneOptions<E>): Promise<E> {
+        await this.persist()
+        const res = await super.findOneOrFail(entityClass as EntityClass<E>, options)
+        if (res != null) this.cache.add(res, options.relations)
+        return res
+    }
+
+    async findOneBy<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        where: FindOptionsWhere<E> | FindOptionsWhere<E>[]
+    ): Promise<E | undefined> {
+        await this.persist()
+        const res = await super.findOneBy(entityClass as EntityClass<E>, where)
+        if (res != null) this.cache.add(res)
+        return res
+    }
+
+    async findOneByOrFail<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        where: FindOptionsWhere<E> | FindOptionsWhere<E>[]
+    ): Promise<E> {
+        await this.persist()
+        const res = await super.findOneByOrFail(entityClass as EntityClass<E>, where)
+        if (res != null) this.cache.add(res)
+        return res
+    }
+
+    async get<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        id: string,
+        relations?: FindOptionsRelations<E>
+    ): Promise<E | undefined> {
+        await this.load()
+
+        const entity = this.getCached(entityClass, id, relations)
+
+        if (entity !== undefined) {
+            return entity == null ? undefined : entity
+        } else {
+            return await this.findOne(entityClass, {where: {id} as any, relations})
+        }
+    }
+
+    async getOrFail<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        id: string,
+        relations?: FindOptionsRelations<E>
+    ): Promise<E> {
+        let e = await this.get(entityClass, id, relations)
+
+        if (e == null) {
+            const metadata = this.em().connection.getMetadata(entityClass)
+            throw new Error(`Missing entity ${metadata.name} with id "${id}"`)
+        }
+
+        return e
+    }
+
+    private getCached<E extends Entity>(entityClass: EntityTarget<E>, id: string, mask: FindOptionsRelations<E> = {}) {
+        const em = this.em()
+        const metadata = em.connection.getMetadata(entityClass)
+
+        const cachedEntity = this.cache.get(entityClass, id)
+
+        if (cachedEntity == null) {
+            return undefined
+        } else if (cachedEntity.value == null) {
+            return null
+        } else {
+            const clonedEntity = em.create(entityClass)
+
+            for (const column of metadata.nonVirtualColumns) {
+                const objectColumnValue = column.getEntityValue(cachedEntity.value)
+                if (objectColumnValue !== undefined) {
+                    column.setEntityValue(clonedEntity, copy(objectColumnValue))
+                }
+            }
+
+            for (const relation of metadata.relations) {
+                let relatedMask = mask[relation.propertyName as keyof E]
+                if (!relatedMask) continue
+
+                const relatedEntity = relation.getEntityValue(cachedEntity.value)
+
+                if (relatedEntity === undefined) {
+                    return undefined // relation is missing, but required
+                } else if (relatedEntity == null) {
+                    relation.setEntityValue(clonedEntity, null)
+                } else {
+                    const cachedRelatedEntity = this.getCached(
+                        relation.inverseEntityMetadata.target,
+                        relatedEntity.id,
+                        typeof relatedMask === 'boolean' ? {} : relatedMask
+                    )
+                    assert(cachedRelatedEntity != null)
+
+                    relation.setEntityValue(clonedEntity, cachedRelatedEntity)
+                }
+            }
+
+            return clonedEntity
+        }
+    }
+
+    defer<E extends Entity>(
+        entityClass: EntityTarget<E>,
+        id: string,
+        relations?: FindOptionsRelations<E>
+    ): DeferredEntity<E> {
+        const _deferredList = this.getDeferData(entityClass)
+
+        _deferredList.ids.add(id)
+
+        if (relations != null) {
+            _deferredList.relations = mergeRelataions(_deferredList.relations, relations)
+        }
+
+        return new DeferredEntity({
+            get: async () => this.get(entityClass, id, relations),
+            getOrFail: async () => this.getOrFail(entityClass, id, relations),
+        })
+    }
+
+    private async persist(): Promise<void> {
+        const em = this.em()
+
+        const entityOrder = this.getTopologicalOrder()
+        const entityOrderReversed = [...entityOrder].reverse()
+
+        const changeSets: Map<string, ChangeSet> = new Map()
+        for (const name of entityOrder) {
+            const updateMap = this.getUpdateMap(name)
+
+            const inserts: Entity[] = []
+            const upserts: Entity[] = []
+            const delayedUpserts: Entity[] = []
+            const removes: Entity[] = []
+            for (const {id, type} of updateMap) {
+                const cached = this.cache.get(name, id)
+
+                switch (type) {
+                    case UpdateType.Insert: {
+                        assert(cached != null && cached.value != null)
+                        inserts.push(cached.value)
+                        break
+                    }
+                    case UpdateType.Upsert: {
+                        assert(cached != null && cached.value != null)
+
+                        let isDelayed = false
+                        for (const relation of this.getSelfRelations(name)) {
+                            const relatedEntity = relation.getEntityValue(cached.value)
+                            const relatedUpdateType = updateMap.get(relatedEntity.id)
+
+                            if (relatedUpdateType === UpdateType.Insert) {
+                                isDelayed = true
+                                break
+                            }
+                        }
+
+                        if (isDelayed) {
+                            delayedUpserts.push(cached.value)
+                        } else {
+                            upserts.push(cached.value)
+                        }
+                        break
+                    }
+                    case UpdateType.Remove: {
+                        const e = em.create(name, {id})
+                        removes.push(e)
+                        break
+                    }
+                }
+            }
+
+            changeSets.set(name, {
+                inserts,
+                upserts,
+                delayedUpserts,
+                removes,
+            })
+        }
+
+        for (const name of entityOrder) {
+            const changeSet = changeSets.get(name)
+            if (changeSet == null) continue
+
+            await super.upsert(changeSet.upserts)
+            await super.insert(changeSet.inserts)
+            await super.upsert(changeSet.delayedUpserts)
+        }
+
+        for (const name of entityOrderReversed) {
+            const changeSet = changeSets.get(name)
+            if (changeSet == null) continue
+
+            await super.remove(changeSet.removes)
+        }
+
+        this.updates.clear()
+    }
+
+    async flush(): Promise<void> {
+        await this.persist()
+        this.cache.clear()
+    }
+
+    private async load(): Promise<void> {
+        const em = this.em()
+
+        for (const [name, deferData] of this.deferMap) {
+            const metadata = em.connection.getMetadata(name)
+
+            for (const id of deferData.ids) {
+                this.cache.ensure(metadata.target, id)
+            }
+
+            for (let batch of splitIntoBatches([...deferData.ids], 30000)) {
+                if (batch.length == 0) continue
+                await this.find<any>(metadata.target, {where: {id: In(batch)}, relations: deferData.relations})
+            }
+        }
+
+        this.deferMap.clear()
+    }
+
+    private knownSelfRelations: Record<string, RelationMetadata[]> = {}
+    private getSelfRelations<E extends Entity>(entityClass: EntityTarget<E>) {
+        const em = this.em()
+        const metadata = em.connection.getMetadata(entityClass)
+
+        if (this.knownSelfRelations[metadata.name] == null) {
+            this.knownSelfRelations[metadata.name] = metadata.relations.filter(
+                (r) => r.inverseEntityMetadata.name === metadata.name
+            )
+        }
+        return this.knownSelfRelations[metadata.name]
+    }
+
+    @def
+    private getTopologicalOrder() {
+        const em = this.em()
+        const graph = Graph()
+        for (const metadata of em.connection.entityMetadatas) {
+            graph.addNode(metadata.name)
+            for (const foreignKey of metadata.foreignKeys) {
+                if (foreignKey.referencedEntityMetadata === metadata) continue // don't add self-relations
+
+                graph.addEdge(metadata.name, foreignKey.referencedEntityMetadata.name)
+            }
+        }
+
+        return graph.topologicalSort().reverse()
+    }
+
+    private getDeferData(entityClass: EntityTarget<any>) {
+        const em = this.em()
+        const metadata = em.connection.getMetadata(entityClass)
+
+        let list = this.deferMap.get(metadata.name)
+        if (list == null) {
+            list = {ids: new Set(), relations: {}}
+            this.deferMap.set(metadata.name, list)
+        }
+
+        return list
+    }
+
+    private getUpdateMap(entityClass: EntityTarget<any>) {
+        const em = this.em()
+        const metadata = em.connection.getMetadata(entityClass)
+
+        let list = this.updates.get(metadata.name)
+        if (list == null) {
+            list = new UpdateMap()
+            this.updates.set(metadata.name, list)
+        }
+
+        return list
+    }
+}
+
+function mergeRelataions<E extends Entity>(
+    a: FindOptionsRelations<E>,
+    b: FindOptionsRelations<E>
+): FindOptionsRelations<E> {
+    const mergedObject: FindOptionsRelations<E> = {}
+
+    for (const key in a) {
+        mergedObject[key] = a[key]
+    }
+
+    for (const key in b) {
+        const bValue = b[key]
+        const value = mergedObject[key]
+        if (typeof bValue === 'object') {
+            mergedObject[key] = (typeof value === 'object' ? mergeRelataions(value, bValue) : bValue) as any
+        } else {
+            mergedObject[key] = value || bValue
+        }
+    }
+
+    return mergedObject
+}
+
+export class DeferredEntity<E extends Entity> {
+    constructor(private opts: {get: () => Promise<E | undefined>; getOrFail: () => Promise<E>}) {}
+
+    async get(): Promise<E | undefined> {
+        return await this.opts.get()
+    }
+
+    async getOrFail(): Promise<E> {
+        return await this.opts.getOrFail()
+    }
+}