fix: continue work on db pull

svetch · svetch · commit 8e40d9a998b6 · 2025-09-25T00:53:50.000+02:00
diff --git a/packages/cli/package.json b/packages/cli/package.json
@@ -29,6 +29,7 @@
         "pack": "pnpm pack"
     },
     "dependencies": {
+        "@dotenvx/dotenvx": "^1.51.0",
         "@zenstackhq/common-helpers": "workspace:*",
         "@zenstackhq/language": "workspace:*",
         "@zenstackhq/sdk": "workspace:*",
diff --git a/packages/cli/src/actions/action-utils.ts b/packages/cli/src/actions/action-utils.ts
@@ -48,7 +48,7 @@ export async function loadSchemaDocument(schemaFile: string) {
 }
 
 export async function loadSchemaDocumentWithServices(schemaFile: string) {
-    const loadResult = await loadDocument(schemaFile);
+    const loadResult = await loadDocument(schemaFile, [], true);
     if (!loadResult.success) {
         loadResult.errors.forEach((err) => {
             console.error(colors.red(err));
diff --git a/packages/cli/src/actions/db.ts b/packages/cli/src/actions/db.ts
@@ -1,5 +1,6 @@
 import { ZModelCodeGenerator } from '@zenstackhq/sdk';
 import fs from 'node:fs';
+import path from 'node:path';
 import { execPackage } from '../utils/exec-utils';
 import { generateTempPrismaSchema, getSchemaFile, handleSubProcessError, loadSchemaDocumentWithServices } from './action-utils';
 import { syncEnums, syncRelation, syncTable, type Relation } from './pull';
@@ -14,6 +15,7 @@ type PushOptions = {
 
 type PullOptions = {
     schema?: string;
+    out?: string;
 };
 
 /**
@@ -60,7 +62,7 @@ async function runPush(options: PushOptions) {
 async function runPull(options: PullOptions) {
     const schemaFile = getSchemaFile(options.schema);
     const { model, services } = await loadSchemaDocumentWithServices(schemaFile);
-
+    await import("@dotenvx/dotenvx/config")
     const SUPPORTED_PROVIDERS = ['sqlite', 'postgresql']
     const datasource = getDatasource(model)
 
@@ -82,16 +84,16 @@ async function runPull(options: PullOptions) {
 
     const { enums, tables } = await provider.introspect(datasource.url)
 
-    syncEnums(enums, model)
+    syncEnums({ dbEnums: enums, model, services })
 
     const resolveRelations: Relation[] = []
     for (const table of tables) {
-        const relations = syncTable({ table, model, provider })
+        const relations = syncTable({ table, model, provider, services })
         resolveRelations.push(...relations)
     }
 
-    for (const rel of resolveRelations) {
-        syncRelation(model, rel, services);
+    for (const relation of resolveRelations) {
+        syncRelation({ model, relation, services });
     }
 
     for (const d of model.declarations) {
@@ -104,6 +106,14 @@ async function runPull(options: PullOptions) {
 
     model.declarations = model.declarations.filter((d) => d !== undefined)
 
-    const zmpdelSchema = await new ZModelCodeGenerator().generate(model)
-    fs.writeFileSync(schemaFile, zmpdelSchema)
+    const generator = await new ZModelCodeGenerator();
+
+    const zmodelSchema = await generator.generate(model)
+
+    console.log(options.out ? `Writing to ${options.out}` : schemaFile);
+
+    const outPath = options.out ? path.resolve(options.out) : schemaFile;
+    console.log(outPath);
+
+    fs.writeFileSync(outPath, zmodelSchema)
 }
diff --git a/packages/cli/src/actions/pull/index.ts b/packages/cli/src/actions/pull/index.ts
@@ -1,5 +1,6 @@
 import type { ZModelServices } from '@zenstackhq/language'
 import type {
+    ArrayExpr,
     Attribute,
     AttributeArg,
     DataField,
@@ -9,18 +10,20 @@ import type {
     Enum,
     EnumField,
     Model,
+    ReferenceExpr,
+    StringLiteral,
     UnsupportedFieldType
 } from '@zenstackhq/language/ast'
+import { getStringLiteral } from '@zenstackhq/language/utils'
 import type { IntrospectedEnum, IntrospectedTable, IntrospectionProvider } from './provider'
-import { getAttributeRef, getDbName } from './utils'
+import { getAttributeRef, getDbName, getEnumRef, getModelRef } from './utils'
 
-export function syncEnums(dbEnums: IntrospectedEnum[], model: Model) {
+export function syncEnums({ dbEnums, model, services }: { dbEnums: IntrospectedEnum[], model: Model, services: ZModelServices }) {
     for (const dbEnum of dbEnums) {
-        let schemaEnum = model.declarations.find(
-            (d) => d.$type === 'Enum' && getDbName(d) === dbEnum.enum_type
-        ) as Enum | undefined
+        let schemaEnum = getEnumRef(dbEnum.enum_type, services);
 
         if (!schemaEnum) {
+            console.log(`Adding enum for type ${dbEnum.enum_type}`);
             schemaEnum = {
                 $type: 'Enum' as const,
                 $container: model,
@@ -66,17 +69,29 @@ export function syncTable({
     model,
     provider,
     table,
+    services
 }: {
     table: IntrospectedTable
     model: Model
     provider: IntrospectionProvider
+    services: ZModelServices
 }) {
+    const idAttribute = getAttributeRef('@id', services)
+    const uniqueAttribute = getAttributeRef('@unique', services)
+    const relationAttribute = getAttributeRef('@relation', services)
+    const fieldMapAttribute = getAttributeRef('@map', services)
+    const tableMapAttribute = getAttributeRef('@@map', services)
+
+    if (!idAttribute || !uniqueAttribute || !relationAttribute || !fieldMapAttribute || !tableMapAttribute) {
+        throw new Error('Cannot find required attributes in the model.')
+    }
+
     const relations: Relation[] = []
-    let modelTable = model.declarations.find(
-        (d) => d.$type === 'DataModel' && getDbName(d) === table.name
-    ) as DataModel | undefined
+    let modelTable = getModelRef(table.name, services)
 
     if (!modelTable) {
+        console.log(`Adding model for table ${table.name}`);
+
         modelTable = {
             $type: 'DataModel' as const,
             $container: model,
@@ -96,7 +111,7 @@ export function syncTable({
                 schema: table.schema,
                 table: table.name,
                 column: col.name,
-                type: col.unique ? 'one' : 'many',
+                type: 'one',
                 fk_name: col.foreign_key_name!,
                 nullable: col.nullable,
                 references: {
@@ -115,85 +130,68 @@ export function syncTable({
         )
         if (!existingField) {
             const builtinType = provider.getBuiltinType(col.datatype)
-            const unsupported: UnsupportedFieldType = {
-                get $container() {
-                    return type
-                },
-                $type: 'UnsupportedFieldType' as const,
-                value: {
-                    get $container() {
-                        return unsupported
-                    },
-                    $type: 'StringLiteral',
-                    value: col.datatype,
-                },
-            }
-
-            const type: DataFieldType = {
-                get $container() {
-                    return field
-                },
-                $type: 'DataFieldType' as const,
-                type: builtinType.type === 'Unsupported' ? undefined : builtinType.type,
-                array: builtinType.isArray,
-                unsupported:
-                    builtinType.type === 'Unsupported' ? unsupported : undefined,
-                optional: col.nullable,
-                reference: col.options.length
-                    ? {
+            const field: DataField = {
+                $type: 'DataField' as const,
+                get type() {
+                    return {
+                        $container: this,
+                        $type: 'DataFieldType' as const,
+                        type: builtinType.type === 'Unsupported' ? undefined : builtinType.type,
+                        array: builtinType.isArray,
+                        get unsupported() {
+                            return builtinType.type === 'Unsupported' ? {
+                                $container: this,
+                                $type: 'UnsupportedFieldType' as const,
+                                get value() {
+                                    return {
+                                        $container: this,
+                                        $type: 'StringLiteral',
+                                        value: col.datatype,
+                                    } satisfies StringLiteral
+                                },
+                            } satisfies UnsupportedFieldType : undefined
+                        },
+                        optional: col.nullable,
+                        reference: col.options.length
+                            ? {
                         $refText: col.datatype,
                         ref: model.declarations.find(
                             (d) => d.$type === 'Enum' && getDbName(d) === col.datatype
-                        ) as Enum | undefined,
-                    }
-                    : undefined,
-            }
-
-            const field: DataField = {
-                $type: 'DataField' as const,
-                type,
+                                ) as Enum | undefined,
+                            }
+                            : undefined,
+                    } satisfies DataFieldType
+                },
                 $container: modelTable!,
                 name: fieldName,
                 get attributes() {
                     if (fieldPrefix !== '') return []
 
-                    const attr: DataFieldAttribute = {
+                    return [{
                         $type: 'DataFieldAttribute' as const,
-                        get $container() {
-                            return field
-                        },
+                        $container: this,
                         decl: {
                             $refText: '@map',
-                            ref: model.$document?.references.find(
-                                (r) =>
-                                    //@ts-ignore
-                                    r.ref.$type === 'Attribute' && r.ref.name === '@map'
-                            )?.ref as Attribute,
+                            ref: fieldMapAttribute,
                         },
                         get args() {
-                            const arg: AttributeArg = {
+                            return [{
                                 $type: 'AttributeArg' as const,
-                                get $container() {
-                                    return attr
-                                },
+                                $container: this,
                                 name: 'name',
                                 $resolvedParam: {
                                     name: 'name',
                                 },
                                 get value() {
                                     return {
                                         $type: 'StringLiteral' as const,
-                                        $container: arg,
+                                        $container: this,
                                         value: col.name,
                                     }
                                 },
-                            }
-
-                            return [arg]
+                            }] satisfies AttributeArg[]
                         },
-                    }
-
-                    return [attr]
+                    }] satisfies DataFieldAttribute[]
                 },
                 comments: [],
             }
@@ -205,10 +203,16 @@ export function syncTable({
     return relations
 }
 
-export function syncRelation(model: Model, relation: Relation, services: ZModelServices) {
+export function syncRelation({ model, relation, services }: { model: Model, relation: Relation, services: ZModelServices }) {
     const idAttribute = getAttributeRef('@id', services)
     const uniqueAttribute = getAttributeRef('@unique', services)
     const relationAttribute = getAttributeRef('@relation', services)
+    const fieldMapAttribute = getAttributeRef('@map', services)
+    const tableMapAttribute = getAttributeRef('@@map', services)
+
+    if (!idAttribute || !uniqueAttribute || !relationAttribute || !fieldMapAttribute || !tableMapAttribute) {
+        throw new Error('Cannot find required attributes in the model.')
+    }
 
     if (!idAttribute || !uniqueAttribute || !relationAttribute) {
         throw new Error('Cannot find required attributes in the model.')
diff --git a/packages/cli/src/actions/pull/utils.ts b/packages/cli/src/actions/pull/utils.ts
@@ -1,11 +1,13 @@
 import type { ZModelServices } from '@zenstackhq/language'
 import {
+  AbstractDeclaration,
   DataField,
+  DataModel,
+  Enum,
   EnumField,
   isInvocationExpr,
-  type AbstractDeclaration,
   type Attribute,
-  type Model,
+  type Model
 } from '@zenstackhq/language/ast'
 import { getStringLiteral } from '@zenstackhq/language/utils'
 import type {
@@ -28,10 +30,20 @@ export function getDatasource(model: Model) {
   }
 
   const urlField = datasource.fields.find((f) => f.name === 'url')!
+
   let url = getStringLiteral(urlField.value)
 
   if (!url && isInvocationExpr(urlField.value)) {
-    url = process.env[getStringLiteral(urlField.value.args[0]) as string]!
+    const envName = getStringLiteral(urlField.value.args[0]?.value)
+    if (!envName) {
+      throw new Error('The url field must be a string literal or an env().')
+    }
+    if (!process.env[envName]) {
+      throw new Error(
+        `Environment variable ${envName} is not set, please set it to the database connection string.`
+      )
+    }
+    url = process.env[envName]
   }
 
   if (!url) {
@@ -62,6 +74,19 @@ export function getDbName(
   return attrValue.value
 }
 
+
+export function getDeclarationRef<T extends AbstractDeclaration>(type: T["$type"], name: string, services: ZModelServices) {
+  return services.shared.workspace.IndexManager.allElements(type).find((m) => m.node && getDbName(m.node as T) === name)?.node as T | undefined
+}
+
+export function getEnumRef(name: string, services: ZModelServices) {
+  return getDeclarationRef<Enum>('Enum', name, services);
+}
+
+export function getModelRef(name: string, services: ZModelServices) {
+  return getDeclarationRef<DataModel>('DataModel', name, services);
+}
+
 export function getAttributeRef(name: string, services: ZModelServices) {
-  return services.shared.workspace.IndexManager.allElements("Attribute").find(a => a.name === name) as Attribute | undefined
+  return getDeclarationRef<Attribute>('Attribute', name, services);
 }
diff --git a/packages/cli/src/index.ts b/packages/cli/src/index.ts
@@ -121,6 +121,14 @@ function createProgram() {
         .addOption(new Option('--force-reset', 'force a reset of the database before push'))
         .action((options) => dbAction('push', options));
 
+    dbCommand
+        .command('pull')
+        .description('Introspect your database.')
+        .addOption(schemaOption)
+        .addOption(noVersionCheckOption)
+        .addOption(new Option('--out <path>', 'add custom output path for the introspected schema'))
+        .action((options) => dbAction('pull', options));
+
     program
         .command('info')
         .description('Get information of installed ZenStack packages.')
diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml

Original file line number	Diff line number	Diff line change
`@@ -48,7 +48,7 @@ export async function loadSchemaDocument(schemaFile: string) {`
`48`	`48`	`}`
`49`	`49`
`50`	`50`	`export async function loadSchemaDocumentWithServices(schemaFile: string) {`
`51`		`- const loadResult = await loadDocument(schemaFile);`
	`51`	`+ const loadResult = await loadDocument(schemaFile, [], true);`
`52`	`52`	`if (!loadResult.success) {`
`53`	`53`	`loadResult.errors.forEach((err) => {`
`54`	`54`	`console.error(colors.red(err));`