From 95214e2a3f07856dee7483bbf92a7ec0266b630d Mon Sep 17 00:00:00 2001
From: Tomas Hubelbauer <tomas@hubelbauer.net>
Date: Wed, 26 Jun 2024 17:32:58 +0200
Subject: [PATCH] Fail on error, support array operations, CompositeTypes,
 quoted identifiers and single-member enums

This makes the library capable of handling the TypeScript types as generated by recent versions of the Supabase CLI as well as makes it able to cover some more exotic scenarios like the aforementioned single-member enums and quoted identifiers.
---
 src/lib/.gitignore         |  2 ++
 src/lib/get-node-name.ts   |  5 +++
 src/lib/transform-types.ts | 74 ++++++++++++++++++++++++++++++++++++--
 src/supabase-to-zod.ts     |  6 +++-
 4 files changed, 83 insertions(+), 4 deletions(-)
 create mode 100644 src/lib/.gitignore

diff --git a/src/lib/.gitignore b/src/lib/.gitignore
new file mode 100644
index 0000000..dc0d34c
--- /dev/null
+++ b/src/lib/.gitignore
@@ -0,0 +1,2 @@
+temp.ts
+temp
diff --git a/src/lib/get-node-name.ts b/src/lib/get-node-name.ts
index d8ad573..3d17d83 100644
--- a/src/lib/get-node-name.ts
+++ b/src/lib/get-node-name.ts
@@ -6,6 +6,11 @@ export const getNodeName = (n: ts.Node) => {
     if (ts.isIdentifier(n)) {
       name = n.text;
     }
+
+    // Handle quoted identifiers in case they contain special characters
+    if (ts.isStringLiteral(n)) {
+      name = n.text;
+    }
   });
   if (!name) throw new Error('Cannot get name of node');
   return name;
diff --git a/src/lib/transform-types.ts b/src/lib/transform-types.ts
index d02fa5f..e83f992 100644
--- a/src/lib/transform-types.ts
+++ b/src/lib/transform-types.ts
@@ -3,6 +3,10 @@ import { z } from 'zod';
 import { getNodeName } from './get-node-name';
 
 const enumFormatterSchema = z.function().args(z.string()).returns(z.string());
+const compositeTypeFormatterSchema = z
+  .function()
+  .args(z.string())
+  .returns(z.string());
 
 const functionFormatterSchema = z
   .function()
@@ -18,6 +22,9 @@ export const transformTypesOptionsSchema = z.object({
   sourceText: z.string(),
   schema: z.string().default('public'),
   enumFormatter: enumFormatterSchema.default(() => (name: string) => name),
+  compositeTypeFormatter: compositeTypeFormatterSchema.default(
+    () => (name: string) => name
+  ),
   functionFormatter: functionFormatterSchema.default(
     () => (name: string, type: string) => `${name}${type}`
   ),
@@ -33,8 +40,13 @@ export const transformTypes = z
   .args(transformTypesOptionsSchema)
   .returns(z.string())
   .implement((opts) => {
-    const { schema, tableOrViewFormatter, enumFormatter, functionFormatter } =
-      opts;
+    const {
+      schema,
+      tableOrViewFormatter,
+      enumFormatter,
+      compositeTypeFormatter,
+      functionFormatter,
+    } = opts;
     const sourceFile = ts.createSourceFile(
       'index.ts',
       opts.sourceText,
@@ -43,6 +55,7 @@ export const transformTypes = z
 
     const typeStrings: string[] = [];
     const enumNames: { name: string; formattedName: string }[] = [];
+    const compositeTypeNames: { name: string; formattedName: string }[] = [];
 
     sourceFile.forEachChild((n) => {
       const processDatabase = (n: ts.Node | ts.TypeNode) => {
@@ -68,7 +81,11 @@ export const transformTypes = z
                                       const operation = getNodeName(n);
                                       if (operation) {
                                         n.forEachChild((n) => {
-                                          if (ts.isTypeLiteralNode(n)) {
+                                          if (
+                                            ts.isTypeLiteralNode(n) ||
+                                            // Handle `Relationships` operation which is an array
+                                            ts.isTupleTypeNode(n)
+                                          ) {
                                             typeStrings.push(
                                               `export type ${tableOrViewFormatter(
                                                 tableOrViewName,
@@ -106,6 +123,46 @@ export const transformTypes = z
                                     name: enumName,
                                   });
                                 }
+
+                                // Handle single-member enums
+                                if (ts.isIdentifier(n)) {
+                                  const formattedName = enumFormatter(enumName);
+                                  typeStrings.push(
+                                    `export type ${formattedName} = '${n.getText(
+                                      sourceFile,
+                                    )}'`,
+                                  );
+                                  enumNames.push({
+                                    formattedName,
+                                    name: enumName,
+                                  });
+                                }
+                              });
+                            }
+                          });
+                        }
+                      });
+                    }
+                    if ('CompositeTypes' === n.name.text) {
+                      n.forEachChild((n) => {
+                        if (ts.isTypeLiteralNode(n)) {
+                          n.forEachChild((n) => {
+                            const enumName = getNodeName(n);
+                            if (ts.isPropertySignature(n)) {
+                              n.forEachChild((n) => {
+                                if (ts.isTypeLiteralNode(n)) {
+                                  const formattedName =
+                                    compositeTypeFormatter(enumName);
+                                  typeStrings.push(
+                                    `export type ${formattedName} = ${n.getText(
+                                      sourceFile,
+                                    )}`,
+                                  );
+                                  compositeTypeNames.push({
+                                    formattedName,
+                                    name: enumName,
+                                  });
+                                }
                               });
                             }
                           });
@@ -181,5 +238,16 @@ export const transformTypes = z
       );
     }
 
+    for (const { name, formattedName } of compositeTypeNames) {
+      parsedTypes = parsedTypes.replaceAll(
+        `Database["${schema}"]["CompositeTypes"]["${name}"]`,
+        formattedName,
+      );
+      parsedTypes = parsedTypes.replaceAll(
+        `Database['${schema}']['CompositeTypes']['${name}']`,
+        formattedName,
+      );
+    }
+
     return parsedTypes;
   });
diff --git a/src/supabase-to-zod.ts b/src/supabase-to-zod.ts
index 9683e92..e7263cf 100644
--- a/src/supabase-to-zod.ts
+++ b/src/supabase-to-zod.ts
@@ -49,11 +49,15 @@ export default async function supabaseToZod(opts: SupabaseToZodOptions) {
 
   const parsedTypes = transformTypes({ sourceText, ...opts });
 
-  const { getZodSchemasFile } = generate({
+  const { getZodSchemasFile, errors } = generate({
     sourceText: parsedTypes,
     ...opts,
   });
 
+  if (errors.length > 0) {
+    throw new Error(errors.join('\n'));
+  }
+
   const zodSchemasFile = getZodSchemasFile(
     getImportPath(outputPath, inputPath)
   );