feat: Add custom client class support

packages/consul/src/manager.ts

@@ -45,6 +45,13 @@ export class ConsulServiceFactory extends ServiceFactory<ConsulClient> {
     config: ConsulOptions,
     clientName: string
   ): Promise<InstanceType<typeof Consul>> {
+    const { customClientClass, ...otherConfig } = config as any;
+    if (customClientClass) {
+      const client = new customClientClass(otherConfig);
+      this.bindTraceContext(client as any, clientName);
+      return client;
+    }
+
     this.logger.info(
       '[midway:consul] init %s at %s:%s',
       clientName,

packages/consul/test/customClient.test.ts

@@ -0,0 +1,103 @@
+import { close, createLightApp } from '@midwayjs/mock';
+import * as consul from '../src';
+
+describe('/test/customClient.test.ts', () => {
+  class CustomConsul {
+    constructor(public config: any) {}
+
+    request(options: any) {
+      return Promise.resolve(options.method);
+    }
+
+    destroy() {}
+  }
+
+  it('should create a custom client from merged default config', async () => {
+    const app = await createLightApp({
+      imports: [consul],
+      globalConfig: {
+        consul: {
+          default: {
+            customClientClass: CustomConsul,
+          },
+          client: {
+            host: 'default.local',
+            port: 8500,
+          },
+        },
+      },
+    });
+
+    const factory = await app
+      .getApplicationContext()
+      .getAsync(consul.ConsulServiceFactory);
+
+    const client = factory.get<CustomConsul>();
+    expect(client).toBeInstanceOf(CustomConsul);
+    expect(client.config).toEqual({
+      host: 'default.local',
+      port: 8500,
+    });
+
+    await close(app);
+  });
+
+  it('should create custom clients and keep factory/service injection', async () => {
+    const app = await createLightApp({
+      imports: [consul],
+      globalConfig: {
+        consul: {
+          clients: {
+            default: {
+              customClientClass: CustomConsul,
+              host: 'default.local',
+              port: 8500,
+            },
+            backup: {
+              customClientClass: CustomConsul,
+              host: 'backup.local',
+              port: 8500,
+            },
+          },
+        },
+      },
+    });
+
+    const factory = await app
+      .getApplicationContext()
+      .getAsync(consul.ConsulServiceFactory);
+    const service = await app
+      .getApplicationContext()
+      .getAsync(consul.ConsulService);
+    const defaultClient = factory.get<CustomConsul>();
+    const backupClient = factory.get<CustomConsul>('backup');
+
+    expect(defaultClient).toBeInstanceOf(CustomConsul);
+    expect(backupClient).toBeInstanceOf(CustomConsul);
+    expect(defaultClient.config.customClientClass).toBeUndefined();
+    expect(backupClient.config.host).toBe('backup.local');
+    expect((service as any).instance).toBe(defaultClient);
+
+    await close(app);
+  });
+
+  it('should bind trace context to custom client', async () => {
+    const factory = new consul.ConsulServiceFactory();
+    const runWithExitSpan = jest.fn(async (_name, _options, callback) => {
+      return callback();
+    });
+    (factory as any).traceService = { runWithExitSpan };
+
+    const client: any = await factory.createClient(
+      {
+        customClientClass: CustomConsul,
+        host: 'localhost',
+        port: 8500,
+      } as any,
+      'default'
+    );
+
+    expect(await client.request({ method: 'GET' })).toBe('GET');
+    expect(runWithExitSpan).toHaveBeenCalledTimes(1);
+  });
+});

packages/core/src/interface.ts

@@ -489,11 +489,15 @@ export type TraceMetaResolver =
         | ((args: TraceMetaResolverArgs) => TraceMetaRecord);
     };
 
+export type BaseServiceFactoryConfigOption<OPTIONS> = PowerPartial<OPTIONS> & {
+  customClientClass?: any;
+};
+
 export type ServiceFactoryConfigOption<OPTIONS> = {
-  default?: PowerPartial<OPTIONS>;
-  client?: PowerPartial<OPTIONS>;
+  default?: BaseServiceFactoryConfigOption<OPTIONS>;
+  client?: BaseServiceFactoryConfigOption<OPTIONS>;
   clients?: {
-    [key: string]: PowerPartial<OPTIONS>;
+    [key: string]: BaseServiceFactoryConfigOption<OPTIONS>;
   };
   defaultClientName?: string;
   clientPriority?: {

packages/core/test/common/serviceFactory.test.ts

@@ -1,4 +1,10 @@
-import { ServiceFactory, DEFAULT_PRIORITY, MidwayPriorityManager, sleep } from '../../src';
+import {
+  ServiceFactory,
+  ServiceFactoryConfigOption,
+  DEFAULT_PRIORITY,
+  MidwayPriorityManager,
+  sleep,
+} from '../../src';
 
 describe('test/common/serviceFactory.test.ts', () => {
 
@@ -74,6 +80,30 @@ describe('test/common/serviceFactory.test.ts', () => {
     expect(instance.get('default')).toBeDefined();
   });
 
+  it('should support custom client class in service factory config types', () => {
+    class CustomClient {}
+
+    const options: ServiceFactoryConfigOption<{ host: string }> = {
+      default: {
+        customClientClass: CustomClient,
+      },
+      client: {
+        customClientClass: CustomClient,
+        host: '127.0.0.1',
+      },
+      clients: {
+        backup: {
+          customClientClass: CustomClient,
+          host: '127.0.0.2',
+        },
+      },
+    };
+
+    expect(options.default.customClientClass).toBe(CustomClient);
+    expect(options.client.customClientClass).toBe(CustomClient);
+    expect(options.clients.backup.customClientClass).toBe(CustomClient);
+  });
+
   it('should test default name', async () => {
     const instance = new TestServiceFactory();
     await instance.initClients({

packages/cos/src/manager.ts

@@ -43,6 +43,13 @@ export class COSServiceFactory extends ServiceFactory<COS> {
   protected traceInjector;
 
   async createClient(config: COS.COSOptions): Promise<COS> {
+    const { customClientClass, ...otherConfig } = config as any;
+    if (customClientClass) {
+      const client = new customClientClass(otherConfig);
+      this.bindTraceContext(client as any);
+      return client;
+    }
+
     assert.ok(
       config.SecretKey && config.SecretId,
       '[@midwayjs/cos] secretId secretKey is required on config'

packages/cos/test/customClient.test.ts

@@ -0,0 +1,89 @@
+import { close, createLightApp } from '@midwayjs/mock';
+import * as cos from '../src';
+
+describe('/test/customClient.test.ts', () => {
+  class CustomCOS {
+    constructor(public config: any) {}
+
+    request(options: any) {
+      return Promise.resolve(options.Action);
+    }
+  }
+
+  it('should create a custom client from client config', async () => {
+    const app = await createLightApp('', {
+      imports: [cos],
+      globalConfig: {
+        cos: {
+          client: {
+            customClientClass: CustomCOS,
+            SecretId: 'default-id',
+            SecretKey: 'secret',
+          },
+        },
+      },
+    });
+
+    const factory = await app
+      .getApplicationContext()
+      .getAsync(cos.COSServiceFactory);
+
+    expect(factory.get()).toBeInstanceOf(CustomCOS);
+
+    await close(app);
+  });
+
+  it('should create custom clients and keep factory/service injection', async () => {
+    const app = await createLightApp('', {
+      imports: [cos],
+      globalConfig: {
+        cos: {
+          clients: {
+            default: {
+              customClientClass: CustomCOS,
+              SecretId: 'default-id',
+              SecretKey: 'secret',
+            },
+            backup: {
+              customClientClass: CustomCOS,
+              SecretId: 'backup-id',
+              SecretKey: 'secret',
+            },
+          },
+        },
+      },
+    });
+
+    const factory = await app
+      .getApplicationContext()
+      .getAsync(cos.COSServiceFactory);
+    const service = await app.getApplicationContext().getAsync(cos.COSService);
+    const defaultClient = factory.get<CustomCOS>();
+    const backupClient = factory.get<CustomCOS>('backup');
+
+    expect(defaultClient).toBeInstanceOf(CustomCOS);
+    expect(backupClient).toBeInstanceOf(CustomCOS);
+    expect(defaultClient.config.customClientClass).toBeUndefined();
+    expect(backupClient.config.SecretId).toBe('backup-id');
+    expect((service as any).instance).toBe(defaultClient);
+
+    await close(app);
+  });
+
+  it('should bind trace context to custom client', async () => {
+    const factory = new cos.COSServiceFactory();
+    const runWithExitSpan = jest.fn(async (_name, _options, callback) => {
+      return callback();
+    });
+    (factory as any).traceService = { runWithExitSpan };
+
+    const client: any = await factory.createClient({
+      customClientClass: CustomCOS,
+      SecretId: 'test',
+      SecretKey: 'secret',
+    } as any);
+
+    expect(await client.request({ Action: 'GetObject' })).toBe('GetObject');
+    expect(runWithExitSpan).toHaveBeenCalledTimes(1);
+  });
+});

packages/etcd/src/manager.ts

@@ -42,6 +42,13 @@ export class ETCDServiceFactory extends ServiceFactory<Etcd3> {
   protected traceInjector;
 
   async createClient(config: IOptions): Promise<Etcd3> {
+    const { customClientClass, ...otherConfig } = config as any;
+    if (customClientClass) {
+      const client = new customClientClass(otherConfig);
+      this.bindTraceContext(client as any);
+      return client;
+    }
+
     this.logger.info('[midway:etcd] init %s', config.hosts);
     const client = new Etcd3(config);
     this.bindTraceContext(client);

packages/etcd/test/customClient.test.ts

@@ -0,0 +1,87 @@
+import { close, createLightApp } from '@midwayjs/mock';
+import * as etcd from '../src';
+
+describe('/test/customClient.test.ts', () => {
+  class CustomEtcd3 {
+    public pool = {
+      exec: jest.fn(async (_service, method) => method),
+    };
+
+    constructor(public config: any) {}
+
+    async close() {}
+  }
+
+  it('should create a custom client from client config', async () => {
+    const app = await createLightApp('', {
+      imports: [etcd],
+      globalConfig: {
+        etcd: {
+          client: {
+            customClientClass: CustomEtcd3,
+            hosts: ['default.local:2379'],
+          },
+        },
+      },
+    });
+
+    const factory = await app
+      .getApplicationContext()
+      .getAsync(etcd.ETCDServiceFactory);
+
+    expect(factory.get()).toBeInstanceOf(CustomEtcd3);
+
+    await close(app);
+  });
+
+  it('should create custom clients and keep factory/service injection', async () => {
+    const app = await createLightApp('', {
+      imports: [etcd],
+      globalConfig: {
+        etcd: {
+          clients: {
+            default: {
+              customClientClass: CustomEtcd3,
+              hosts: ['default.local:2379'],
+            },
+            backup: {
+              customClientClass: CustomEtcd3,
+              hosts: ['backup.local:2379'],
+            },
+          },
+        },
+      },
+    });
+
+    const factory = await app
+      .getApplicationContext()
+      .getAsync(etcd.ETCDServiceFactory);
+    const service = await app.getApplicationContext().getAsync(etcd.ETCDService);
+    const defaultClient = factory.get<CustomEtcd3>();
+    const backupClient = factory.get<CustomEtcd3>('backup');
+
+    expect(defaultClient).toBeInstanceOf(CustomEtcd3);
+    expect(backupClient).toBeInstanceOf(CustomEtcd3);
+    expect(defaultClient.config.customClientClass).toBeUndefined();
+    expect(backupClient.config.hosts).toEqual(['backup.local:2379']);
+    expect((service as any).instance).toBe(defaultClient);
+
+    await close(app);
+  });
+
+  it('should bind trace context to custom client', async () => {
+    const factory = new etcd.ETCDServiceFactory();
+    const runWithExitSpan = jest.fn(async (_name, _options, callback) => {
+      return callback();
+    });
+    (factory as any).traceService = { runWithExitSpan };
+
+    const client: any = await factory.createClient({
+      customClientClass: CustomEtcd3,
+      hosts: ['localhost:2379'],
+    } as any);
+
+    expect(await client.pool.exec('KV', 'put', {})).toBe('put');
+    expect(runWithExitSpan).toHaveBeenCalledTimes(1);
+  });
+});

packages/oss/src/manager.ts

@@ -59,6 +59,13 @@ export class OSSServiceFactory<
   async createClient(
     config: OSSServiceFactoryCreateClientConfigType
   ): Promise<T> {
+    const { customClientClass, ...otherConfig } = config as any;
+    if (customClientClass) {
+      const client = new customClientClass(otherConfig);
+      this.bindTraceContext(client as any);
+      return client;
+    }
+
     if (config['cluster'] && !config.clusters) {
       config.clusters = config['cluster'];
     }