From c46fa6ba61c85b358e39de0c2a245aa630b41c3a Mon Sep 17 00:00:00 2001
From: ymc9 <104139426+ymc9@users.noreply.github.com>
Date: Mon, 18 Nov 2024 12:50:00 -0800
Subject: [PATCH] chore: unify the logic of saving ts files

---
 packages/schema/src/plugins/enhancer/enhance/index.ts      | 7 ++++---
 .../src/plugins/enhancer/policy/policy-guard-generator.ts  | 5 +++--
 packages/sdk/src/model-meta-generator.ts                   | 5 +++--
 3 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/packages/schema/src/plugins/enhancer/enhance/index.ts b/packages/schema/src/plugins/enhancer/enhance/index.ts
index ce4e167c1..404c16cd8 100644
--- a/packages/schema/src/plugins/enhancer/enhance/index.ts
+++ b/packages/schema/src/plugins/enhancer/enhance/index.ts
@@ -11,6 +11,7 @@ import {
     isDelegateModel,
     isDiscriminatorField,
     normalizedRelative,
+    saveSourceFile,
     type PluginOptions,
 } from '@zenstackhq/sdk';
 import {
@@ -137,7 +138,7 @@ ${
                 { overwrite: true }
             );
 
-            await this.saveSourceFile(enhanceTs);
+            this.saveSourceFile(enhanceTs);
         }
 
         return {
@@ -779,9 +780,9 @@ export function enhance(prisma: any, context?: EnhancementContext<${authTypePara
         return isReferenceExpr(arg) ? (arg.target.ref as DataModelField) : undefined;
     }
 
-    private async saveSourceFile(sf: SourceFile) {
+    private saveSourceFile(sf: SourceFile) {
         if (this.options.preserveTsFiles) {
-            await sf.save();
+            saveSourceFile(sf);
         }
     }
 
diff --git a/packages/schema/src/plugins/enhancer/policy/policy-guard-generator.ts b/packages/schema/src/plugins/enhancer/policy/policy-guard-generator.ts
index 8622d13d4..8206f797b 100644
--- a/packages/schema/src/plugins/enhancer/policy/policy-guard-generator.ts
+++ b/packages/schema/src/plugins/enhancer/policy/policy-guard-generator.ts
@@ -26,6 +26,7 @@ import {
     hasValidationAttributes,
     isAuthInvocation,
     isForeignKeyField,
+    saveSourceFile,
 } from '@zenstackhq/sdk';
 import { getPrismaClientImportSpec } from '@zenstackhq/sdk/prisma';
 import { streamAst } from 'langium';
@@ -57,7 +58,7 @@ import {
 export class PolicyGenerator {
     constructor(private options: PluginOptions) {}
 
-    async generate(project: Project, model: Model, output: string) {
+    generate(project: Project, model: Model, output: string) {
         const sf = project.createSourceFile(path.join(output, 'policy.ts'), undefined, { overwrite: true });
 
         this.writeImports(model, output, sf);
@@ -86,7 +87,7 @@ export class PolicyGenerator {
         // save ts files if requested explicitly or the user provided
         const preserveTsFiles = this.options.preserveTsFiles === true || !!this.options.output;
         if (preserveTsFiles) {
-            await sf.save();
+            saveSourceFile(sf);
         }
     }
 
diff --git a/packages/sdk/src/model-meta-generator.ts b/packages/sdk/src/model-meta-generator.ts
index 4140786cb..71b4246ca 100644
--- a/packages/sdk/src/model-meta-generator.ts
+++ b/packages/sdk/src/model-meta-generator.ts
@@ -38,6 +38,7 @@ import {
     isForeignKeyField,
     isIdField,
     resolved,
+    saveSourceFile,
     TypeScriptExpressionTransformer,
 } from '.';
 
@@ -66,7 +67,7 @@ export type ModelMetaGeneratorOptions = {
     shortNameMap?: Map<string, string>;
 };
 
-export async function generate(
+export function generate(
     project: Project,
     models: DataModel[],
     typeDefs: TypeDef[],
@@ -82,7 +83,7 @@ export async function generate(
     sf.addStatements('export default metadata;');
 
     if (options.preserveTsFiles) {
-        await sf.save();
+        saveSourceFile(sf);
     }
 
     return sf;