From cd529795a394963f38589cedadb281d7cef79a55 Mon Sep 17 00:00:00 2001
From: Sindre Gulseth <sgulseth@gmail.com>
Date: Tue, 2 Apr 2024 13:58:24 +0200
Subject: [PATCH] feat(schema): handle assetRequired when extracting schema
 with enforceRequiredFields (#6157)

---
 .../schema/src/sanity/extractSchema.ts        | 51 ++++++++++++++++
 .../test/extractSchema/extractSchema.test.ts  | 60 +++++++++++++++++++
 2 files changed, 111 insertions(+)

diff --git a/packages/@sanity/schema/src/sanity/extractSchema.ts b/packages/@sanity/schema/src/sanity/extractSchema.ts
index e70f909d2ae..dd101e2a5ef 100644
--- a/packages/@sanity/schema/src/sanity/extractSchema.ts
+++ b/packages/@sanity/schema/src/sanity/extractSchema.ts
@@ -208,6 +208,12 @@ export function extractSchema(
         continue
       }
 
+      // if the field sets assetRequired() we will mark the asset attribute as required
+      // also guard against the case where the field is not an object, though type validation should catch this
+      if (hasAssetRequired(field) && value.type === 'object') {
+        value.attributes.asset.optional = false
+      }
+
       // if we extract with enforceRequiredFields, we will mark the field as optional only if it is not a required field,
       // else we will always mark it as optional
       const optional = extractOptions.enforceRequiredFields ? fieldIsRequired === false : true
@@ -333,6 +339,51 @@ function isFieldRequired(field: ObjectField): boolean {
   return false
 }
 
+function hasAssetRequired(field: ObjectField): boolean {
+  const {validation} = field.type
+  if (!validation) {
+    return false
+  }
+  const rules = Array.isArray(validation) ? validation : [validation]
+  for (const rule of rules) {
+    let assetRequired = false
+
+    // hack to check if a field is required. We create a proxy that returns itself when a method is called,
+    // if the method is "required" we set a flag
+    const proxy = new Proxy(
+      {},
+      {
+        get: (target, methodName) => () => {
+          if (methodName === 'assetRequired') {
+            assetRequired = true
+          }
+          return proxy
+        },
+      },
+    ) as Rule
+
+    if (typeof rule === 'function') {
+      rule(proxy)
+      if (assetRequired) {
+        return true
+      }
+    }
+
+    if (
+      typeof rule === 'object' &&
+      rule !== null &&
+      '_rules' in rule &&
+      Array.isArray(rule._rules)
+    ) {
+      if (rule._rules.some((r) => r.flag === 'assetRequired')) {
+        return true
+      }
+    }
+  }
+
+  return false
+}
+
 function isObjectType(typeDef: SanitySchemaType): typeDef is ObjectSchemaType {
   return isType(typeDef, 'object') || typeDef.jsonType === 'object' || 'fields' in typeDef
 }
diff --git a/packages/@sanity/schema/test/extractSchema/extractSchema.test.ts b/packages/@sanity/schema/test/extractSchema/extractSchema.test.ts
index f28f5e126a2..6cf7e4c10a9 100644
--- a/packages/@sanity/schema/test/extractSchema/extractSchema.test.ts
+++ b/packages/@sanity/schema/test/extractSchema/extractSchema.test.ts
@@ -458,6 +458,66 @@ describe('Extract schema test', () => {
     expect(book.attributes.optionalTitle.optional).toBe(true)
   })
 
+  test('enforceRequiredFields handles `assetRequired`', () => {
+    const schema1 = createSchema({
+      name: 'test',
+      types: [
+        {
+          title: 'Book',
+          name: 'book',
+          type: 'document',
+          fields: [
+            {
+              title: 'Title',
+              name: 'title',
+              type: 'string',
+            },
+            defineField({
+              title: 'Required Image',
+              name: 'requiredImage',
+              type: 'image',
+              validation: (Rule) => Rule.required(),
+            }),
+            defineField({
+              title: 'Asset Required Image',
+              name: 'assetRequiredImage',
+              type: 'image',
+              validation: (Rule) => Rule.required().assetRequired(),
+            }),
+            {
+              title: 'Asset Required File Rule Spec',
+              name: 'assetRequiredFileRuleSpec',
+              type: 'file',
+              validation: {
+                _required: 'required',
+                _rules: [{flag: 'assetRequired', constraint: {assetType: 'file'}}],
+              },
+            },
+          ],
+        },
+      ],
+    })
+
+    const extracted = extractSchema(schema1, {enforceRequiredFields: true})
+    const book = extracted.find((type) => type.name === 'book')
+    expect(book).toBeDefined()
+    assert(book !== undefined) // this is a workaround for TS, but leave the expect above for clarity in case of failure
+    assert(book.type === 'document') // this is a workaround for TS, but leave the expect above for clarity in case of failure
+    expect(book.attributes.title.optional).toBe(true)
+
+    expect(book.attributes.requiredImage.optional).toBe(false)
+    assert(book.attributes.requiredImage.value.type === 'object') // this is a workaround for TS, but leave the expect above for clarity in case of failure
+    expect(book.attributes.requiredImage.value.attributes.asset.optional).toBe(true) // we dont set assetRequired(), so it should be optional
+
+    expect(book.attributes.assetRequiredImage.optional).toBe(false)
+    assert(book.attributes.assetRequiredImage.value.type === 'object') // this is a workaround for TS, but leave the expect above for clarity in case of failure
+    expect(book.attributes.assetRequiredImage.value.attributes.asset.optional).toBe(false) // with assetRequired(), it should be required
+
+    expect(book.attributes.assetRequiredFileRuleSpec.optional).toBe(false)
+    assert(book.attributes.assetRequiredFileRuleSpec.value.type === 'object') // this is a workaround for TS, but leave the expect above for clarity in case of failure
+    expect(book.attributes.assetRequiredFileRuleSpec.value.attributes.asset.optional).toBe(false) // with assetRequired defined in _rules, it should be required
+  })
+
   describe('can handle `list` option that is not an array', () => {
     const schema = createSchema(schemaFixtures.listObjectOption)
     const extracted = extractSchema(schema)