kibae · kibae · Apr 19, 2024
diff --git a/src/sharding-data-source.decorator.ts b/src/sharding-data-source.decorator.ts
@@ -1,12 +1,17 @@
 import { Entity, EntityOptions } from 'typeorm';
-import { ShardingType } from './types/typeorm-sharding.type';
+import { FunctionShardingRule, ShardingType } from './types/typeorm-sharding.type';
 
 export interface RangeShardingEntityOptions<ENTITY, KEY_TYPE> extends EntityOptions {
     type: ShardingType.RANGE;
     findShard: (entity: ENTITY, minKey: KEY_TYPE, maxKey: KEY_TYPE) => boolean;
     findShardById: (id: any, minKey: KEY_TYPE, maxKey: KEY_TYPE) => boolean;
 }
 
+export interface FunctionShardingEntityOptions<ENTITY> extends EntityOptions {
+    type: ShardingType.FUNCTION;
+    findShard: (entity: ENTITY, shardIndex: number, shard: FunctionShardingRule) => boolean;
+}
+
 /*
 export interface ModularShardingEntityOptions<ENTITY, KEY_TYPE> extends EntityOptions {
     type: ShardingType.MODULAR;
@@ -15,7 +20,7 @@ export interface ModularShardingEntityOptions<ENTITY, KEY_TYPE> extends EntityOp
 }
  */
 
-export type ShardingEntityOptions<ENTITY, KEY_TYPE> = RangeShardingEntityOptions<ENTITY, KEY_TYPE>;
+export type ShardingEntityOptions<ENTITY, KEY_TYPE> = RangeShardingEntityOptions<ENTITY, KEY_TYPE> | FunctionShardingEntityOptions<ENTITY>;
 // | ModularShardingEntityOptions<ENTITY, KEY_TYPE>;
 
 export function ShardingEntity<ENTITY, KEY_TYPE = number>(options: ShardingEntityOptions<ENTITY, KEY_TYPE>) {
@@ -28,14 +33,11 @@ export function ShardingEntity<ENTITY, KEY_TYPE = number>(options: ShardingEntit
             Reflect.defineMetadata('SHARDING_TYPE', type, target);
             Reflect.defineMetadata('SHARDING_FUNC', findShard, target);
             Reflect.defineMetadata('SHARDING_FUNC_BY_ID', findShardById, target);
-            /*
-        } else if (options.type === ShardingType.MODULAR) {
-            const { type, findShard, keyProvider, ...opt } = options;
+        } else if (options.type === ShardingType.FUNCTION) {
+            const { type, findShard, ...opt } = options;
             entityOptions = opt;
             Reflect.defineMetadata('SHARDING_TYPE', type, target);
             Reflect.defineMetadata('SHARDING_FUNC', findShard, target);
-            Reflect.defineMetadata('SHARDING_KEY_PROVIDER', keyProvider, target);
-             */
         } else throw new Error('ShardingEntity: Unsupported sharding type');
 
         Entity(entityOptions)(target);

diff --git a/src/sharding-manager.ts b/src/sharding-manager.ts
@@ -60,6 +60,10 @@ export class ShardingManager {
             const idx = this.options.shards.findIndex((item) => shardingFunc(entity, item.minKey, item.maxKey));
             if (idx < 0) return this.dataSources[this.dataSources.length - 1];
             return this.dataSources[idx];
+        } else if (this.options.shardingType === ShardingType.FUNCTION) {
+            const idx = this.options.shards.findIndex((item, idx) => shardingFunc(entity, idx, item));
+            if (idx < 0) return this.dataSources[this.dataSources.length - 1];
+            return this.dataSources[idx];
         } else throw new Error('ShardingManager: Unsupported sharding type');
     }
 
@@ -68,6 +72,8 @@ export class ShardingManager {
             const idx = this.options.shards.findIndex((item) => shardingFuncById(id, item.minKey, item.maxKey));
             if (idx < 0) return this.dataSources[this.dataSources.length - 1];
             return this.dataSources[idx];
+        } else if (this.options.shardingType === ShardingType.FUNCTION) {
+            throw new Error('ShardingManager: Function type cannot be used with getDataSourceById');
         } else throw new Error('ShardingManager: Unsupported sharding type');
     }
 

diff --git a/src/test/entity/case4func.ts b/src/test/entity/case4func.ts
@@ -0,0 +1,23 @@
+import { Column, Index, PrimaryGeneratedColumn } from 'typeorm';
+import { ShardingEntity } from '../../sharding-data-source.decorator';
+import { FunctionShardingRule, ShardingType } from '../../types/typeorm-sharding.type';
+import { ShardingBaseEntity } from '../../sharding-base-entity';
+
+@ShardingEntity<Case4, number>({
+    type: ShardingType.FUNCTION,
+    findShard: (entity: Case4, shardIndex: number, shard: FunctionShardingRule<string[]>) => shard.userData!.includes(entity.firstName),
+})
+export class Case4 extends ShardingBaseEntity {
+    @PrimaryGeneratedColumn()
+    id!: number;
+
+    @Column()
+    @Index()
+    firstName!: string;
+
+    @Column()
+    lastName!: string;
+
+    @Column()
+    age!: number;
+}
diff --git a/src/test/sharding-manager-function.spec.ts b/src/test/sharding-manager-function.spec.ts
@@ -0,0 +1,86 @@
+import { ShardingManager } from '../sharding-manager';
+import { ShardingType } from '../types/typeorm-sharding.type';
+import { DataSource } from 'typeorm';
+import { Case4 } from './entity/case4func';
+
+async function updateSeq(dataSource: DataSource, val: number): Promise<void> {
+    await dataSource.query(`INSERT INTO sqlite_sequence(name, seq) VALUES('case4', ${val})`);
+}
+
+async function getDataSource(): Promise<ShardingManager> {
+    return ShardingManager.init({
+        shardingType: ShardingType.FUNCTION,
+        type: 'sqlite',
+        synchronize: true,
+        logging: 'all',
+        entities: [Case4],
+        shards: [
+            { database: ':memory:', userData: ['Typeorm'], onInitialize: (dataSource) => updateSeq(dataSource, 0) },
+            { database: ':memory:', userData: ['John'], onInitialize: (dataSource) => updateSeq(dataSource, 1000) },
+            { database: ':memory:', userData: ['typeORM'], onInitialize: (dataSource) => updateSeq(dataSource, 2000) },
+        ],
+    });
+}
+
+let dataSource: ShardingManager;
+describe('ShardingManager-Function', () => {
+    beforeEach(async () => {
+        dataSource = await getDataSource();
+    });
+
+    afterEach(async () => {
+        if (dataSource) {
+            await dataSource.destroy();
+            dataSource.dataSources.forEach((dataSource) => expect(dataSource.isInitialized).toEqual(false));
+        }
+    });
+
+    it('DataSource', async () => {
+        expect(dataSource).toBeDefined();
+        expect(dataSource.dataSources.length).toBe(3);
+    });
+
+    it('Case1 - insert into first shard', async () => {
+        let entity;
+        entity = await Case4.save({
+            firstName: 'Typeorm',
+            lastName: 'Sharding',
+            age: 10,
+        });
+
+        const dataSources = Case4.getAllDataSource();
+        expect((await dataSources[0].manager.query(`SELECT * FROM case4`)).length).toBe(1);
+        expect((await dataSources[1].manager.query(`SELECT * FROM case4`)).length).toBe(0);
+        expect((await dataSources[2].manager.query(`SELECT * FROM case4`)).length).toBe(0);
+
+        expect((await Case4.find()).length).toBe(1);
+        expect(await Case4.findOneBy({ firstName: 'Typeorm' })).toBeDefined();
+        expect(await Case4.findOneBy({ firstName: 'John' })).toBeUndefined();
+
+        // expect((await Case4.findOneById(1))?.firstName).toBeDefined();
+        // expect((await Case4.findOneById(1001))?.firstName).toBeUndefined();
+        // expect((await Case4.findOneById(2001))?.firstName).toBeUndefined();
+    });
+
+    it('Case1 - insert into second shard', async () => {
+        let entity;
+        entity = await Case4.save({
+            firstName: 'typeORM',
+            lastName: 'Sharding',
+            age: 10,
+        });
+
+        const dataSources = Case4.getAllDataSource();
+        expect((await dataSources[0].manager.query(`SELECT * FROM case4`)).length).toBe(0);
+        expect((await dataSources[1].manager.query(`SELECT * FROM case4`)).length).toBe(0);
+        expect((await dataSources[2].manager.query(`SELECT * FROM case4`)).length).toBe(1);
+
+        expect((await Case4.find()).length).toBe(1);
+        expect(await Case4.findOneBy({ firstName: 'typeORM' })).toBeDefined();
+        expect(await Case4.findOneBy({ firstName: 'John' })).toBeUndefined();
+
+        // expect((await Case4.findOneById(1))?.firstName).toBeUndefined();
+        // expect((await Case4.findOneById(1001))?.firstName).toBeUndefined();
+        // expect((await Case4.findOneById(2001))?.firstName).toBeDefined();
+    });
+});
diff --git a/src/test/sharding-manager.spec.ts → src/test/sharding-manager-range.spec.ts b/src/test/sharding-manager.spec.ts → src/test/sharding-manager-range.spec.ts
@@ -23,7 +23,7 @@ async function getDataSource(): Promise<ShardingManager> {
 }
 
 let dataSource: ShardingManager;
-describe('ShardingManager', () => {
+describe('ShardingManager-Range', () => {
     beforeEach(async () => {
         dataSource = await getDataSource();
     });

diff --git a/src/types/typeorm-sharding.type.ts b/src/types/typeorm-sharding.type.ts
@@ -57,6 +57,7 @@ export type AbstractShardingDataSourceOptions<S = {}> =
 
 export enum ShardingType {
     RANGE = 'RANGE',
+    FUNCTION = 'FUNCTION',
     // MODULAR = 'MODULAR',
 }
 
@@ -84,11 +85,18 @@ export type RangeShardingDataSourceOptions<T = DefaultShardingKeyType> = Abstrac
     shardingType: ShardingType.RANGE;
 };
 
-/*
-export type ModularShardingDataSourceOptions<T = DefaultShardingKeyType> = AbstractShardingDataSourceOptions & {
-    shardingType: ShardingType.MODULAR;
+//  _______  __    __  .__   __.   ______ .___________. __    ______   .__   __.
+// |   ____||  |  |  | |  \ |  |  /      ||           ||  |  /  __  \  |  \ |  |
+// |  |__   |  |  |  | |   \|  | |  ,----'`---|  |----`|  | |  |  |  | |   \|  |
+// |   __|  |  |  |  | |  . `  | |  |         |  |     |  | |  |  |  | |  . `  |
+// |  |     |  `--'  | |  |\   | |  `----.    |  |     |  | |  `--'  | |  |\   |
+// |__|      \______/  |__| \__|  \______|    |__|     |__|  \______/  |__| \__|
+export interface FunctionShardingRule<U = any> {
+    userData?: U;
+}
+
+export type FunctionShardingDataSourceOptions<U = any> = AbstractShardingDataSourceOptions<FunctionShardingRule<U>> & {
+    shardingType: ShardingType.FUNCTION;
 };
 
-export type ShardingDataSourceOptions = RangeShardingDataSourceOptions | ModularShardingDataSourceOptions;
- */
-export type ShardingDataSourceOptions = RangeShardingDataSourceOptions;
+export type ShardingDataSourceOptions<T = DefaultShardingKeyType> = RangeShardingDataSourceOptions<T> | FunctionShardingDataSourceOptions;