implements extends to ObjectSchema (#68067) (#68301)

* implements `extends` to `ObjectSchema`

* add unit tests

* use expectType for types assertions

* allow to extends options

* add comment about deep extend

packages/kbn-config-schema/package.json

@@ -10,7 +10,8 @@
     "kbn:bootstrap": "yarn build"
   },
   "devDependencies": {
-    "typescript": "3.7.2"
+    "typescript": "3.7.2",
+    "tsd": "^0.7.4"
   },
   "peerDependencies": {
     "joi": "^13.5.2",

packages/kbn-config-schema/src/index.ts

@@ -44,6 +44,7 @@ import {
   ObjectType,
   ObjectTypeOptions,
   Props,
+  NullableProps,
   RecordOfOptions,
   RecordOfType,
   StringOptions,
@@ -57,7 +58,7 @@ import {
   StreamType,
 } from './types';
 
-export { ObjectType, TypeOf, Type };
+export { ObjectType, TypeOf, Type, Props, NullableProps };
 export { ByteSizeValue } from './byte_size_value';
 export { SchemaTypeError, ValidationError } from './errors';
 export { isConfigSchema } from './typeguards';

packages/kbn-config-schema/src/types/index.ts

@@ -29,7 +29,7 @@ export { LiteralType } from './literal_type';
 export { MaybeType } from './maybe_type';
 export { MapOfOptions, MapOfType } from './map_type';
 export { NumberOptions, NumberType } from './number_type';
-export { ObjectType, ObjectTypeOptions, Props, TypeOf } from './object_type';
+export { ObjectType, ObjectTypeOptions, Props, NullableProps, TypeOf } from './object_type';
 export { RecordOfOptions, RecordOfType } from './record_type';
 export { StreamType } from './stream_type';
 export { StringOptions, StringType } from './string_type';

packages/kbn-config-schema/src/types/object_type.test.ts

@@ -17,6 +17,7 @@
  * under the License.
  */
 
+import { expectType } from 'tsd';
 import { schema } from '..';
 import { TypeOf } from './object_type';
 
@@ -360,17 +361,142 @@ test('handles optional properties', () => {
 
   type SchemaType = TypeOf<typeof type>;
 
-  let foo: SchemaType = {
+  expectType<SchemaType>({
     required: 'foo',
-  };
-  foo = {
+  });
+  expectType<SchemaType>({
     required: 'hello',
     optional: undefined,
-  };
-  foo = {
+  });
+  expectType<SchemaType>({
     required: 'hello',
     optional: 'bar',
-  };
-
-  expect(foo).toBeDefined();
+  });
+});
+
+describe('#extends', () => {
+  it('allows to extend an existing schema by adding new properties', () => {
+    const origin = schema.object({
+      initial: schema.string(),
+    });
+
+    const extended = origin.extends({
+      added: schema.number(),
+    });
+
+    expect(() => {
+      extended.validate({ initial: 'foo' });
+    }).toThrowErrorMatchingInlineSnapshot(
+      `"[added]: expected value of type [number] but got [undefined]"`
+    );
+
+    expect(() => {
+      extended.validate({ initial: 'foo', added: 42 });
+    }).not.toThrowError();
+
+    expectType<TypeOf<typeof extended>>({
+      added: 12,
+      initial: 'foo',
+    });
+  });
+
+  it('allows to extend an existing schema by removing properties', () => {
+    const origin = schema.object({
+      string: schema.string(),
+      number: schema.number(),
+    });
+
+    const extended = origin.extends({ number: undefined });
+
+    expect(() => {
+      extended.validate({ string: 'foo', number: 12 });
+    }).toThrowErrorMatchingInlineSnapshot(`"[number]: definition for this key is missing"`);
+
+    expect(() => {
+      extended.validate({ string: 'foo' });
+    }).not.toThrowError();
+
+    expectType<TypeOf<typeof extended>>({
+      string: 'foo',
+    });
+  });
+
+  it('allows to extend an existing schema by overriding an existing properties', () => {
+    const origin = schema.object({
+      string: schema.string(),
+      mutated: schema.number(),
+    });
+
+    const extended = origin.extends({
+      mutated: schema.string(),
+    });
+
+    expect(() => {
+      extended.validate({ string: 'foo', mutated: 12 });
+    }).toThrowErrorMatchingInlineSnapshot(
+      `"[mutated]: expected value of type [string] but got [number]"`
+    );
+
+    expect(() => {
+      extended.validate({ string: 'foo', mutated: 'bar' });
+    }).not.toThrowError();
+
+    expectType<TypeOf<typeof extended>>({
+      string: 'foo',
+      mutated: 'bar',
+    });
+  });
+
+  it('properly infer the type from optional properties', () => {
+    const origin = schema.object({
+      original: schema.maybe(schema.string()),
+      mutated: schema.maybe(schema.number()),
+      removed: schema.maybe(schema.string()),
+    });
+
+    const extended = origin.extends({
+      removed: undefined,
+      mutated: schema.string(),
+    });
+
+    expect(() => {
+      extended.validate({ original: 'foo' });
+    }).toThrowErrorMatchingInlineSnapshot(
+      `"[mutated]: expected value of type [string] but got [undefined]"`
+    );
+    expect(() => {
+      extended.validate({ original: 'foo' });
+    }).toThrowErrorMatchingInlineSnapshot(
+      `"[mutated]: expected value of type [string] but got [undefined]"`
+    );
+    expect(() => {
+      extended.validate({ original: 'foo', mutated: 'bar' });
+    }).not.toThrowError();
+
+    expectType<TypeOf<typeof extended>>({
+      original: 'foo',
+      mutated: 'bar',
+    });
+    expectType<TypeOf<typeof extended>>({
+      mutated: 'bar',
+    });
+  });
+
+  it(`allows to override the original schema's options`, () => {
+    const origin = schema.object(
+      {
+        initial: schema.string(),
+      },
+      { defaultValue: { initial: 'foo' } }
+    );
+
+    const extended = origin.extends(
+      {
+        added: schema.number(),
+      },
+      { defaultValue: { initial: 'bar', added: 42 } }
+    );
+
+    expect(extended.validate(undefined)).toEqual({ initial: 'bar', added: 42 });
+  });
 });

packages/kbn-config-schema/src/types/object_type.ts

@@ -24,6 +24,8 @@ import { ValidationError } from '../errors';
 
 export type Props = Record<string, Type<any>>;
 
+export type NullableProps = Record<string, Type<any> | undefined | null>;
+
 export type TypeOf<RT extends Type<any>> = RT['type'];
 
 type OptionalProperties<Base extends Props> = Pick<
@@ -47,6 +49,24 @@ export type ObjectResultType<P extends Props> = Readonly<
     { [K in keyof RequiredProperties<P>]: TypeOf<P[K]> }
 >;
 
+type DefinedProperties<Base extends NullableProps> = Pick<
+  Base,
+  {
+    [Key in keyof Base]: undefined extends Base[Key] ? never : null extends Base[Key] ? never : Key;
+  }[keyof Base]
+>;
+
+type ExtendedProps<P extends Props, NP extends NullableProps> = Omit<P, keyof NP> &
+  { [K in keyof DefinedProperties<NP>]: NP[K] };
+
+type ExtendedObjectType<P extends Props, NP extends NullableProps> = ObjectType<
+  ExtendedProps<P, NP>
+>;
+
+type ExtendedObjectTypeOptions<P extends Props, NP extends NullableProps> = ObjectTypeOptions<
+  ExtendedProps<P, NP>
+>;
+
 interface UnknownOptions {
   /**
    * Options for dealing with unknown keys:
@@ -61,10 +81,13 @@ export type ObjectTypeOptions<P extends Props = any> = TypeOptions<ObjectResultT
   UnknownOptions;
 
 export class ObjectType<P extends Props = any> extends Type<ObjectResultType<P>> {
-  private props: Record<string, AnySchema>;
+  private props: P;
+  private options: ObjectTypeOptions<P>;
+  private propSchemas: Record<string, AnySchema>;
 
-  constructor(props: P, { unknowns = 'forbid', ...typeOptions }: ObjectTypeOptions<P> = {}) {
+  constructor(props: P, options: ObjectTypeOptions<P> = {}) {
     const schemaKeys = {} as Record<string, AnySchema>;
+    const { unknowns = 'forbid', ...typeOptions } = options;
     for (const [key, value] of Object.entries(props)) {
       schemaKeys[key] = value.getSchema();
     }
@@ -77,7 +100,93 @@ export class ObjectType<P extends Props = any> extends Type<ObjectResultType<P>>
       .options({ stripUnknown: { objects: unknowns === 'ignore' } });
 
     super(schema, typeOptions);
-    this.props = schemaKeys;
+    this.props = props;
+    this.propSchemas = schemaKeys;
+    this.options = options;
+  }
+
+  /**
+   * Return a new `ObjectType` instance extended with given `newProps` properties.
+   * Original properties can be deleted from the copy by passing a `null` or `undefined` value for the key.
+   *
+   * @example
+   * How to add a new key to an object schema
+   * ```ts
+   * const origin = schema.object({
+   *   initial: schema.string(),
+   * });
+   *
+   * const extended = origin.extends({
+   *   added: schema.number(),
+   * });
+   * ```
+   *
+   * How to remove an existing key from an object schema
+   * ```ts
+   * const origin = schema.object({
+   *   initial: schema.string(),
+   *   toRemove: schema.number(),
+   * });
+   *
+   * const extended = origin.extends({
+   *   toRemove: undefined,
+   * });
+   * ```
+   *
+   * How to override the schema's options
+   * ```ts
+   * const origin = schema.object({
+   *   initial: schema.string(),
+   * }, { defaultValue: { initial: 'foo' }});
+   *
+   * const extended = origin.extends({
+   *   added: schema.number(),
+   * }, { defaultValue: { initial: 'foo', added: 'bar' }});
+   *
+   * @remarks
+   * `extends` only support extending first-level properties. It's currently not possible to perform deep/nested extensions.
+   *
+   * ```ts
+   * const origin = schema.object({
+   *   foo: schema.string(),
+   *   nested: schema.object({
+   *     a: schema.string(),
+   *     b: schema.string(),
+   *   }),
+   * });
+   *
+   * const extended = origin.extends({
+   *   nested: schema.object({
+   *     c: schema.string(),
+   *   }),
+   * });
+   *
+   * // TypeOf<typeof extended> is `{ foo: string; nested: { c: string } }`
+   * ```
+   */
+  public extends<NP extends NullableProps>(
+    newProps: NP,
+    newOptions?: ExtendedObjectTypeOptions<P, NP>
+  ): ExtendedObjectType<P, NP> {
+    const extendedProps = Object.entries({
+      ...this.props,
+      ...newProps,
+    }).reduce((memo, [key, value]) => {
+      if (value !== null && value !== undefined) {
+        return {
+          ...memo,
+          [key]: value,
+        };
+      }
+      return memo;
+    }, {} as ExtendedProps<P, NP>);
+
+    const extendedOptions = {
+      ...this.options,
+      ...newOptions,
+    } as ExtendedObjectTypeOptions<P, NP>;
+
+    return new ObjectType(extendedProps, extendedOptions);
   }
 
   protected handleError(type: string, { reason, value }: Record<string, any>) {
@@ -95,10 +204,10 @@ export class ObjectType<P extends Props = any> extends Type<ObjectResultType<P>>
   }
 
   validateKey(key: string, value: any) {
-    if (!this.props[key]) {
+    if (!this.propSchemas[key]) {
       throw new Error(`${key} is not a valid part of this schema`);
     }
-    const { value: validatedValue, error } = this.props[key].validate(value);
+    const { value: validatedValue, error } = this.propSchemas[key].validate(value);
     if (error) {
       throw new ValidationError(error as any, key);
     }