Skip to content

Commit 8f80bc7

Browse files
vkarpov15vkarpov15
authored
Merge pull request #15574 from Automattic/vkarpov15/gh-10894
feat(schema): support for union types
2 parents a59c33b + 56ac298 commit 8f80bc7

File tree

10 files changed

+345
-21
lines changed

10 files changed

+345
-21
lines changed

‎lib/options/schemaUnionOptions.js

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
'use strict';
2+
3+
const SchemaTypeOptions = require('./schemaTypeOptions');
4+
5+
/**
6+
* The options defined on a Union schematype.
7+
*
8+
* @api public
9+
* @inherits SchemaTypeOptions
10+
* @constructor SchemaUnionOptions
11+
*/
12+
13+
class SchemaUnionOptions extends SchemaTypeOptions {}
14+
15+
const opts = require('./propertyOptions');
16+
17+
/**
18+
* If set, specifies the types that this union can take. Mongoose will cast
19+
* the value to one of the given types.
20+
*
21+
* If not set, Mongoose will not cast the value to any specific type.
22+
*
23+
* @api public
24+
* @property of
25+
* @memberOf SchemaUnionOptions
26+
* @type {Function|Function[]|string|string[]}
27+
* @instance
28+
*/
29+
30+
Object.defineProperty(SchemaUnionOptions.prototype, 'of', opts);
31+
32+
module.exports = SchemaUnionOptions;

‎lib/schema.js

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1525,7 +1525,7 @@ Object.defineProperty(Schema.prototype, 'base', {
15251525
*
15261526
* @param {String} path
15271527
* @param {Object} obj constructor
1528-
* @param {Object} options
1528+
* @param {Object} options schema options
15291529
* @api private
15301530
*/
15311531

@@ -1539,7 +1539,6 @@ Schema.prototype.interpretAsType = function(path, obj, options) {
15391539
return clone;
15401540
}
15411541

1542-
15431542
// If this schema has an associated Mongoose object, use the Mongoose object's
15441543
// copy of SchemaTypes re: gh-7158 gh-6933
15451544
const MongooseTypes = this.base != null ? this.base.Schema.Types : Schema.Types;
@@ -1740,7 +1739,10 @@ Schema.prototype.interpretAsType = function(path, obj, options) {
17401739
'https://bit.ly/mongoose-schematypes for a list of valid schema types.');
17411740
}
17421741

1743-
const schemaType = new MongooseTypes[name](path, obj);
1742+
if (name === 'Union') {
1743+
obj.parentSchema = this;
1744+
}
1745+
const schemaType = new MongooseTypes[name](path, obj, options);
17441746

17451747
if (schemaType.$isSchemaMap) {
17461748
createMapNestedSchemaType(this, schemaType, path, obj, options);

‎lib/schema/index.js

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,15 +12,16 @@ exports.Buffer = require('./buffer');
1212
exports.Date = require('./date');
1313
exports.Decimal128 = exports.Decimal = require('./decimal128');
1414
exports.DocumentArray = require('./documentArray');
15+
exports.Double = require('./double');
16+
exports.Int32 = require('./int32');
1517
exports.Map = require('./map');
1618
exports.Mixed = require('./mixed');
1719
exports.Number = require('./number');
1820
exports.ObjectId = require('./objectId');
1921
exports.String = require('./string');
2022
exports.Subdocument = require('./subdocument');
2123
exports.UUID = require('./uuid');
22-
exports.Double = require('./double');
23-
exports.Int32 = require('./int32');
24+
exports.Union = require('./union');
2425

2526
// alias
2627

‎lib/schema/union.js

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
'use strict';
2+
3+
/*!
4+
* ignore
5+
*/
6+
7+
const SchemaUnionOptions = require('../options/schemaUnionOptions');
8+
const SchemaType = require('../schemaType');
9+
10+
const firstValueSymbol = Symbol('firstValue');
11+
12+
/*!
13+
* ignore
14+
*/
15+
16+
class Union extends SchemaType {
17+
constructor(key, options, schemaOptions = {}) {
18+
super(key, options, 'Union');
19+
if (!options || !Array.isArray(options.of) || options.of.length === 0) {
20+
throw new Error('Union schema type requires an array of types');
21+
}
22+
this.schemaTypes = options.of.map(obj => options.parentSchema.interpretAsType(key, obj, schemaOptions));
23+
}
24+
25+
cast(val, doc, init, prev, options) {
26+
let firstValue = firstValueSymbol;
27+
let lastError;
28+
// Loop through each schema type in the union. If one of the schematypes returns a value that is `=== val`, then
29+
// use `val`. Otherwise, if one of the schematypes casted successfully, use the first successfully casted value.
30+
// Finally, if none of the schematypes casted successfully, throw the error from the last schema type in the union.
31+
// The `=== val` check is a workaround to ensure that the original value is returned if it matches one of the schema types,
32+
// avoiding cases like where numbers are casted to strings or dates even if the schema type is a number.
33+
for (let i = 0; i < this.schemaTypes.length; ++i) {
34+
try {
35+
const casted = this.schemaTypes[i].cast(val, doc, init, prev, options);
36+
if (casted === val) {
37+
return casted;
38+
}
39+
if (firstValue === firstValueSymbol) {
40+
firstValue = casted;
41+
}
42+
} catch (error) {
43+
lastError = error;
44+
}
45+
}
46+
if (firstValue !== firstValueSymbol) {
47+
return firstValue;
48+
}
49+
throw lastError;
50+
}
51+
52+
// Setters also need to be aware of casting - we need to apply the setters of the entry in the union we choose.
53+
applySetters(val, doc, init, prev, options) {
54+
let firstValue = firstValueSymbol;
55+
let lastError;
56+
// Loop through each schema type in the union. If one of the schematypes returns a value that is `=== val`, then
57+
// use `val`. Otherwise, if one of the schematypes casted successfully, use the first successfully casted value.
58+
// Finally, if none of the schematypes casted successfully, throw the error from the last schema type in the union.
59+
// The `=== val` check is a workaround to ensure that the original value is returned if it matches one of the schema types,
60+
// avoiding cases like where numbers are casted to strings or dates even if the schema type is a number.
61+
for (let i = 0; i < this.schemaTypes.length; ++i) {
62+
try {
63+
let castedVal = this.schemaTypes[i]._applySetters(val, doc, init, prev, options);
64+
if (castedVal == null) {
65+
castedVal = this.schemaTypes[i]._castNullish(castedVal);
66+
} else {
67+
castedVal = this.schemaTypes[i].cast(castedVal, doc, init, prev, options);
68+
}
69+
if (castedVal === val) {
70+
return castedVal;
71+
}
72+
if (firstValue === firstValueSymbol) {
73+
firstValue = castedVal;
74+
}
75+
} catch (error) {
76+
lastError = error;
77+
}
78+
}
79+
if (firstValue !== firstValueSymbol) {
80+
return firstValue;
81+
}
82+
throw lastError;
83+
}
84+
85+
clone() {
86+
const schematype = super.clone();
87+
88+
schematype.schemaTypes = this.schemaTypes.map(schemaType => schemaType.clone());
89+
return schematype;
90+
}
91+
}
92+
93+
/**
94+
* This schema type's name, to defend against minifiers that mangle
95+
* function names.
96+
*
97+
* @api public
98+
*/
99+
Union.schemaName = 'Union';
100+
101+
Union.defaultOptions = {};
102+
103+
Union.prototype.OptionsConstructor = SchemaUnionOptions;
104+
105+
module.exports = Union;

‎test/schema.union.test.js

Lines changed: 137 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,137 @@
1+
'use strict';
2+
3+
const start = require('./common');
4+
const util = require('./util');
5+
6+
const assert = require('assert');
7+
8+
const mongoose = start.mongoose;
9+
const Schema = mongoose.Schema;
10+
11+
describe('Union', function() {
12+
let db;
13+
14+
before(async function() {
15+
db = await start().asPromise();
16+
});
17+
18+
after(async function() {
19+
await db.close();
20+
});
21+
22+
afterEach(() => db.deleteModel(/Test/));
23+
afterEach(() => util.clearTestData(db));
24+
afterEach(() => util.stopRemainingOps(db));
25+
26+
it('basic functionality should work', async function() {
27+
const schema = new Schema({
28+
test: {
29+
type: 'Union',
30+
of: [Number, String]
31+
}
32+
});
33+
const TestModel = db.model('Test', schema);
34+
35+
const doc1 = new TestModel({ test: 1 });
36+
assert.strictEqual(doc1.test, 1);
37+
await doc1.save();
38+
39+
const doc1FromDb = await TestModel.collection.findOne({ _id: doc1._id });
40+
assert.strictEqual(doc1FromDb.test, 1);
41+
42+
const doc2 = new TestModel({ test: 'abc' });
43+
assert.strictEqual(doc2.test, 'abc');
44+
await doc2.save();
45+
46+
const doc2FromDb = await TestModel.collection.findOne({ _id: doc2._id });
47+
assert.strictEqual(doc2FromDb.test, 'abc');
48+
});
49+
50+
it('should report last cast error', async function() {
51+
const schema = new Schema({
52+
test: {
53+
type: 'Union',
54+
of: [Number, Boolean]
55+
}
56+
});
57+
const TestModel = db.model('Test', schema);
58+
59+
const doc1 = new TestModel({ test: 'taco tuesday' });
60+
assert.strictEqual(doc1.test, undefined);
61+
await assert.rejects(
62+
doc1.save(),
63+
'ValidationError: test: Cast to Boolean failed for value "taco tuesday" (type string) at path "test" because of "CastError"'
64+
);
65+
});
66+
67+
it('should cast for query', async function() {
68+
const schema = new Schema({
69+
test: {
70+
type: 'Union',
71+
of: [Number, Date]
72+
}
73+
});
74+
const TestModel = db.model('Test', schema);
75+
76+
const doc1 = new TestModel({ test: 1 });
77+
assert.strictEqual(doc1.test, 1);
78+
await doc1.save();
79+
80+
let res = await TestModel.findOne({ test: 1 });
81+
assert.strictEqual(res.test, 1);
82+
83+
res = await TestModel.findOne({ test: '1' });
84+
assert.strictEqual(res.test, 1);
85+
86+
await TestModel.create({ test: new Date('2025-06-01') });
87+
res = await TestModel.findOne({ test: '2025-06-01' });
88+
assert.strictEqual(res.test.valueOf(), new Date('2025-06-01').valueOf());
89+
});
90+
91+
it('should cast updates', async function() {
92+
const schema = new Schema({
93+
test: {
94+
type: 'Union',
95+
of: [Number, Date]
96+
}
97+
});
98+
const TestModel = db.model('Test', schema);
99+
100+
const doc1 = new TestModel({ test: 1 });
101+
assert.strictEqual(doc1.test, 1);
102+
await doc1.save();
103+
104+
let res = await TestModel.findOneAndUpdate({ _id: doc1._id }, { test: '1' }, { returnDocument: 'after' });
105+
assert.strictEqual(res.test, 1);
106+
107+
res = await TestModel.findOneAndUpdate({ _id: doc1._id }, { test: new Date('2025-06-01') }, { returnDocument: 'after' });
108+
assert.strictEqual(res.test.valueOf(), new Date('2025-06-01').valueOf());
109+
});
110+
111+
it('should handle setters', async function() {
112+
const schema = new Schema({
113+
test: {
114+
type: 'Union',
115+
of: [
116+
Number,
117+
{
118+
type: String,
119+
trim: true
120+
}
121+
]
122+
}
123+
});
124+
const TestModel = db.model('Test', schema);
125+
126+
const doc1 = new TestModel({ test: 1 });
127+
assert.strictEqual(doc1.test, 1);
128+
await doc1.save();
129+
130+
const doc2 = new TestModel({ test: ' bbb ' });
131+
assert.strictEqual(doc2.test, 'bbb');
132+
await doc2.save();
133+
134+
const doc2FromDb = await TestModel.collection.findOne({ _id: doc2._id });
135+
assert.strictEqual(doc2FromDb.test, 'bbb');
136+
});
137+
});

‎test/schematype.test.js

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -266,7 +266,7 @@ describe('schematype', function() {
266266
});
267267

268268
const typesToTest = Object.values(mongoose.SchemaTypes).
269-
filter(t => t.name !== 'SchemaSubdocument' && t.name !== 'SchemaDocumentArray');
269+
filter(t => t.name !== 'SchemaSubdocument' && t.name !== 'SchemaDocumentArray' && t.name !== 'Union');
270270

271271
typesToTest.forEach((type) => {
272272
it(type.name + ', when given a default option, set its', () => {

‎test/types/schema.test.ts

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1920,3 +1920,33 @@ function gh15536() {
19201920
const user3 = new UserModelNameRequiredCustom({ name: null });
19211921
expectType<string>(user3.name);
19221922
}
1923+
1924+
function gh10894() {
1925+
function autoInferred() {
1926+
const schema = new Schema({
1927+
testProp: {
1928+
type: 'Union',
1929+
of: [String, Number]
1930+
}
1931+
});
1932+
const TestModel = model('Test', schema);
1933+
1934+
type InferredDocType = InferSchemaType<typeof schema>;
1935+
expectType<string | number | null | undefined>({} as InferredDocType['testProp']);
1936+
1937+
const doc = new TestModel({ testProp: 42 });
1938+
expectType<string | number | null | undefined>(doc.testProp);
1939+
1940+
const toObject = doc.toObject();
1941+
expectType<string | number | null | undefined>(toObject.testProp);
1942+
1943+
const schemaDefinition = {
1944+
testProp: {
1945+
type: 'Union',
1946+
of: ['String', 'Number']
1947+
}
1948+
} as const;
1949+
type RawDocType = InferRawDocType<typeof schemaDefinition>;
1950+
expectType<string | number | null | undefined>({} as RawDocType['testProp']);
1951+
}
1952+
}

‎types/inferrawdoctype.d.ts

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,10 @@ declare module 'mongoose' {
3535
TypeKey
3636
>;
3737

38+
type UnionToRawPathType<T extends readonly any[]> = T[number] extends infer U
39+
? ResolveRawPathType<U>
40+
: never;
41+
3842
/**
3943
* Same as inferSchemaType, except:
4044
*
@@ -109,11 +113,12 @@ declare module 'mongoose' {
109113
IfEquals<PathValueType, Schema.Types.UUID> extends true ? Buffer :
110114
PathValueType extends MapConstructor | 'Map' ? Map<string, ResolveRawPathType<Options['of']>> :
111115
IfEquals<PathValueType, typeof Schema.Types.Map> extends true ? Map<string, ResolveRawPathType<Options['of']>> :
112-
PathValueType extends ArrayConstructor ? any[] :
113-
PathValueType extends typeof Schema.Types.Mixed ? any:
114-
IfEquals<PathValueType, ObjectConstructor> extends true ? any:
115-
IfEquals<PathValueType, {}> extends true ? any:
116-
PathValueType extends typeof SchemaType ? PathValueType['prototype'] :
117-
PathValueType extends Record<string, any> ? InferRawDocType<PathValueType> :
118-
unknown;
116+
PathValueType extends 'Union' | 'union' | typeof Schema.Types.Union ? Options['of'] extends readonly any[] ? UnionToRawPathType<Options['of']> : never :
117+
PathValueType extends ArrayConstructor ? any[] :
118+
PathValueType extends typeof Schema.Types.Mixed ? any:
119+
IfEquals<PathValueType, ObjectConstructor> extends true ? any:
120+
IfEquals<PathValueType, {}> extends true ? any:
121+
PathValueType extends typeof SchemaType ? PathValueType['prototype'] :
122+
PathValueType extends Record<string, any> ? InferRawDocType<PathValueType> :
123+
unknown;
119124
}

0 commit comments

Comments
 (0)