@cyber-rom/nestjs-scylladb/lib/orm/repositories/repository.ts

Based on https://www.npmjs.com/package/@ouato/nestjs-express-cassandra With support columns name mapping

import {map} from 'rxjs/operators';
import {Type} from '@nestjs/common';
import {types} from 'cassandra-driver';
import {Observable, defer, Subject} from 'rxjs';

// noinspection ES6PreferShortImport
import {
    BaseModel,
    FindQuery,
    SaveOptionsStatic,
    UpdateOptionsStatic,
    DeleteOptionsStatic,
    FindQueryOptionsStatic,
} from '../interfaces/externals/scylla.interface';
import {transformKey, transformQuery, transformQuerySelect} from '../utils/transform-query.utils';
import {transformFrom, transformTo} from '../utils/transform-entity.utils';
import {EntityNotFoundError} from '../errors';
import {ReturnQueryBuilder} from './builder/return-query.builder';

const defaultOptions = {
    findOptions: {raw: true},
    updateOptions: {if_exists: true},
    deleteOptions: {if_exists: true},
};

export class Repository<Entity = any> {
    readonly model: BaseModel<Entity>;
    readonly target: Type<Entity>;
    readonly returnQueryBuilder: ReturnQueryBuilder<Entity>;

    create(entity?: Partial<Entity>): Entity;
    create(entities: Partial<Entity>[]): Entity[];
    create(entityLike?: any): Entity | Entity[] {
        return transformFrom(this.target, entityLike);
    }

    count(query: FindQuery<Entity>, key?: String | keyof Entity): Observable<number> {
        return new Observable((subscriber) => {
            this.model.execute_query(
                `SELECT COUNT(${
                    key ? transformKey(this.target, key) : '*'
                }) FROM "${this.model.get_keyspace_name()}"."${this.model.get_table_name()}";`,
                [],
                (error, value) => {
                    if (error) {
                        return subscriber.error(error);
                    }
                    subscriber.next(parseInt(value, 10) || 0);
                    subscriber.complete();
                },
            );
        });
    }

    findOne(query: FindQuery<Entity>, options?: FindQueryOptionsStatic<Entity>): Observable<Entity>;

    findOne(query: FindQuery<Entity>, options: FindQueryOptionsStatic<Entity> = {}): Observable<Entity> {
        return defer(() =>
            this.model.findOneAsync(transformQuery(this.target, query), {
                ...options,
                ...defaultOptions.findOptions,
                select: transformQuerySelect(this.target, options?.select),
            }),
        ).pipe(map((x) => x && transformFrom(this.target, x)));
    }

    findOneOrFail(query: FindQuery<Entity>, options?: FindQueryOptionsStatic<Entity>): Observable<Entity>;

    findOneOrFail(query: FindQuery<Entity>, maybeOptions: FindQueryOptionsStatic<Entity> = {}): Observable<Entity> {
        return this.findOne(transformQuery(this.target, query), maybeOptions).pipe(
            map((entity) => {
                if (entity === undefined) {
                    throw new EntityNotFoundError(this.target, query);
                }
                return entity;
            }),
        );
    }

    find(query: FindQuery<Entity>, options?: FindQueryOptionsStatic<Entity>): Observable<Entity[]>;

    find(query: FindQuery<Entity>, options: FindQueryOptionsStatic<Entity> = {}): Observable<Entity[]> {
        return defer(() =>
            this.model.findAsync(transformQuery(this.target, query), {
                ...options,
                ...defaultOptions.findOptions,
                select: transformQuerySelect(this.target, options?.select),
            }),
        ).pipe(map((x) => transformFrom(this.target, x)));
    }

    findAndCount(
        query: FindQuery<Entity>,
        options: FindQueryOptionsStatic<Entity> = {},
    ): Observable<[Entity[], number]> {
        return defer(() =>
            this.model.findAsync(transformQuery(this.target, query), {
                ...(options as any),
                ...defaultOptions.findOptions,
                select: transformQuerySelect(this.target, options?.select),
            }),
        ).pipe(
            map((x) => transformFrom(this.target, x)),
            map((entities) => [entities, entities.length] as [Entity[], number]),
        );
    }

    save(entity: Partial<Entity>, options?: SaveOptionsStatic): Observable<Entity>;

    save(entities: Partial<Entity>[], options?: SaveOptionsStatic): Observable<Entity[]>;

    save(
        entityLike: Partial<Entity> | Partial<Entity>[],
        options: SaveOptionsStatic = {},
    ): Observable<Entity> | Observable<Entity[]> {
        const saveFunc = async (entity) => {
            const model = new this.model(transformTo(this.target, entity));
            await model.saveAsync(options);
            return transformFrom(this.target, model.toJSON());
        };
        const saveMultipleFunc = (arrayLike: Entity[]) => Promise.all(arrayLike.map((x) => saveFunc(x)));
        return Array.isArray(entityLike)
            ? defer(() => saveMultipleFunc(entityLike as any))
            : defer(() => saveFunc(entityLike as any));
    }

    update(
        query: FindQuery<Entity>,
        updateValue: Partial<Entity>,
        options?: UpdateOptionsStatic<Entity>,
    ): Observable<any>;

    update(
        query: FindQuery<Entity>,
        updateValue: Partial<Entity>,
        options: UpdateOptionsStatic<Entity> = {},
    ): Observable<any> {
        return defer(() =>
            this.model.updateAsync(transformQuery(this.target, query), transformQuery(this.target, updateValue), {
                ...defaultOptions.updateOptions,
                ...options,
            }),
        );
    }

    remove(entity: Entity, options?: DeleteOptionsStatic): Observable<Entity>;

    remove(entity: Entity[], options?: DeleteOptionsStatic): Observable<Entity[]>;

    remove(entityOrEntities: Entity | Entity[], options: DeleteOptionsStatic = {}): Observable<Entity | Entity[]> {
        const removeFunc = (entity) =>
            new this.model(transformTo(this.target, entity)).deleteAsync({
                ...defaultOptions.deleteOptions,
                ...options,
            });
        const promiseArray = Array.isArray(entityOrEntities)
            ? entityOrEntities.map((x) => removeFunc(x))
            : [removeFunc(entityOrEntities)];

        return defer(() => Promise.all(promiseArray)).pipe(map(() => entityOrEntities));
    }

    delete(query: FindQuery<Entity>, options?: DeleteOptionsStatic): Observable<any>;

    delete(query = {}, options = {}) {
        return defer(() =>
            this.model.deleteAsync(transformQuery(this.target, query), {
                ...defaultOptions.deleteOptions,
                ...options,
            }),
        );
    }

    truncate(): Observable<any> {
        return defer(() => this.model.truncateAsync());
    }

    stream(query: FindQuery<Entity>, options: FindQueryOptionsStatic<Entity> = {}): Observable<Entity> {
        const reader$ = new Subject<any>();

        const onRead = (reader): void => {
            while (true) {
                const row = reader.readRow();
                if (row === null) {
                    break;
                }
                reader$.next(transformFrom(this.target, row));
            }
        };

        const onDone = (error): void => {
            if (error) {
                reader$.error(error);
            }
            reader$.complete();
            return;
        };

        this.model.stream(
            transformQuery(this.target, query),
            {
                ...options,
                ...defaultOptions.findOptions,
                select: transformQuerySelect(this.target, options?.select),
            },
            onRead,
            onDone,
        );

        return reader$.asObservable();
    }

    eachRow(query: FindQuery<Entity>, options: FindQueryOptionsStatic<Entity> = {}): EachRowArgument {
        const reader$ = new Subject<any>();
        const done$ = new Subject<any>();
        const getReader = () => reader$.asObservable();
        const getDone = () => done$.asObservable();

        const onRow = (n, row): void => reader$.next(transformFrom(this.target, row));
        const onDone = (err: Error, result: any): void => {
            if (err) {
                reader$.error(err);
                done$.error(err);
            } else {
                done$.next(result);
            }
            reader$.complete();
            done$.complete();
        };

        this.model.eachRow(
            transformQuery(this.target, query),
            {
                ...options,
                ...defaultOptions.findOptions,
                select: transformQuerySelect(this.target, options?.select),
            },
            onRow,
            onDone,
        );

        return {getReader, getDone};
    }

    get getModelRef(): BaseModel<Entity> {
        return this.model;
    }

    getReturnQueryBuilder(): ReturnQueryBuilder<Entity> {
        return this.returnQueryBuilder;
    }

    doBatch(queries): Promise<any> {
        return this.model.execute_batchAsync(queries);
    }
}

export interface EachRowArgument {
    getReader<T = any>(): Observable<T>;

    getDone(): Observable<types.ResultSet>;
}