fix(zmodel): require implicitly used "public" schema to be declared in "schemas" config (#492)

* fix(zmodel): require implicitly used "public" schema to be declared in "schemas" config

* Update packages/language/src/validators/datasource-validator.ts

Co-authored-by: Copilot <[email protected]>

* test: verify "public" schema not required when all models have explicit @@Schema (#496)

* Initial plan

* Add test for explicit schema usage without public

Co-authored-by: ymc9 <[email protected]>

---------

Co-authored-by: copilot-swe-agent[bot] <[email protected]>
Co-authored-by: ymc9 <[email protected]>

---------

Co-authored-by: Copilot <[email protected]>
Co-authored-by: Copilot <[email protected]>

Author: 3 people

packages/language/src/validators/datasource-validator.ts

@@ -1,6 +1,6 @@
 import type { ValidationAcceptor } from 'langium';
 import { SUPPORTED_PROVIDERS } from '../constants';
-import { DataSource, isConfigArrayExpr, isInvocationExpr, isLiteralExpr } from '../generated/ast';
+import { DataSource, isConfigArrayExpr, isDataModel, isEnum, isInvocationExpr, isLiteralExpr } from '../generated/ast';
 import { getStringLiteral } from '../utils';
 import { validateDuplicatedDeclarations, type AstValidator } from './common';
 
@@ -70,14 +70,28 @@ export default class DataSourceValidator implements AstValidator<DataSource> {
                 accept('error', '"schemas" must be an array of string literals', {
                     node: schemasField,
                 });
-            } else if (
-                // validate `defaultSchema` is included in `schemas`
-                defaultSchemaValue &&
-                !schemasValue.items.some((e) => getStringLiteral(e) === defaultSchemaValue)
-            ) {
-                accept('error', `"${defaultSchemaValue}" must be included in the "schemas" array`, {
-                    node: schemasField,
-                });
+            } else {
+                const schemasArray = schemasValue.items.map((e) => getStringLiteral(e)!);
+
+                if (defaultSchemaValue) {
+                    // validate `defaultSchema` is included in `schemas`
+                    if (!schemasArray.includes(defaultSchemaValue)) {
+                        accept('error', `"${defaultSchemaValue}" must be included in the "schemas" array`, {
+                            node: schemasField,
+                        });
+                    }
+                } else {
+                    // if no explicit default schema is specified, and there are models or enums without '@@schema',
+                    // "public" is implicitly used, so it must be included in the "schemas" array
+                    const hasImplicitPublicSchema = ds.$container.declarations.some(
+                        (d) => (isDataModel(d) || isEnum(d)) && !d.attributes.some((a) => a.decl.$refText === '@@schema'),
+                    );
+                    if (hasImplicitPublicSchema && !schemasArray.includes('public')) {
+                        accept('error', `"public" must be included in the "schemas" array`, {
+                            node: schemasField,
+                        });
+                    }
+                }
             }
         }
     }

tests/e2e/orm/client-api/pg-custom-schema.test.ts

@@ -247,6 +247,75 @@ model Foo {
         ).rejects.toThrow('"mySchema" must be included in the "schemas" array');
     });
 
+    it('requires implicit public schema to be included in schemas', async () => {
+        await expect(
+            createTestClient(
+                `
+datasource db {
+    provider = 'postgresql'
+    schemas = ['mySchema']
+    url = '$DB_URL'
+}
+
+enum Role {
+    ADMIN
+    USER
+}
+
+model Foo {
+    id Int @id
+    name String
+    role Role
+    @@schema('mySchema')
+}
+
+model Bar {
+    id Int @id
+    name String
+}
+`,
+            ),
+        ).rejects.toThrow('"public" must be included in the "schemas" array');
+    });
+
+    it('does not require public schema when all models and enums have explicit schema', async () => {
+        const db = await createTestClient(
+            `
+datasource db {
+    provider = 'postgresql'
+    schemas = ['mySchema']
+    url = '$DB_URL'
+}
+
+enum Role {
+    ADMIN
+    USER
+    @@schema('mySchema')
+}
+
+model Foo {
+    id Int @id
+    name String
+    role Role
+    @@schema('mySchema')
+}
+
+model Bar {
+    id Int @id
+    name String
+    @@schema('mySchema')
+}
+`,
+            {
+                provider: 'postgresql',
+                usePrismaPush: true,
+            },
+        );
+
+        await expect(db.foo.create({ data: { id: 1, name: 'test', role: 'ADMIN' } })).toResolveTruthy();
+        await expect(db.bar.create({ data: { id: 1, name: 'test' } })).toResolveTruthy();
+    });
+
     it('allows specifying schema only on a few models', async () => {
         let fooQueriesVerified = false;
         let barQueriesVerified = false;