Support env variables

README.md

@@ -10,7 +10,7 @@ this project when I encounter issues in my own work.
 
 1. Install with `yarn add zero-codegen` or `npm install zero-codegen`
 2. Find your Postgres database connection string.
-3. Run `zero-codegen --database-url <connection-string>`
+3. Run `zero-codegen --database-url <connection-string>`. You can also use the `DATABASE_URL` environment variable.
 4. Optionally, you can specify the file path to the tables and relationships files with `--tables-path` and
    `--relationships-path` parameters.
 

package.json

@@ -14,6 +14,7 @@
   },
   "dependencies": {
     "commander": "^13.1.0",
+    "dotenv": "^16.4.7",
     "pg": "^8.13.3",
     "ts-morph": "^25.0.1"
   },

src/cli.ts

@@ -2,6 +2,10 @@
 
 import { Command } from "commander";
 import { generate } from "./generate";
+import { config } from "dotenv";
+
+const envPaths = [".env", ".env.local"];
+envPaths.forEach((path) => config({ path: path, override: true }));
 
 const program = new Command();
 
@@ -10,14 +14,19 @@ program
   .option("--database-url <url>", "Connection string to your Postgres database")
   .option("--tables-path <filepath>", "Path to the tables file (default ./tables.ts)")
   .option("--relationships-path <filepath>", "Path to the relationships file (default ./relationships.ts)")
+  .option(
+    "--force",
+    "Force option overwrites the existing files. By using this you will lose all changes made to the files."
+  )
   .parse(process.argv);
 
 const options = program.opts();
 
 generate({
-  databaseUrl: options.databaseUrl,
+  databaseUrl: options.databaseUrl ?? process.env["DATABASE_URL"],
   tablesPath: options.tablesPath,
   relationshipsPath: options.relationshipsPath,
+  forceGenerate: !!options.force,
 }).then(() => {
   console.log("Zero tables and relationships generated");
 });

src/generate.ts

@@ -13,12 +13,14 @@ import { ZeroColumn, ZeroRelationship } from "./types";
 
 export async function generate({
   databaseUrl,
-  tablesPath,
-  relationshipsPath,
+  tablesPath = "./tables.ts",
+  relationshipsPath = "./relationships.ts",
+  forceGenerate,
 }: {
   databaseUrl: string;
   tablesPath?: string;
   relationshipsPath?: string;
+  forceGenerate?: boolean;
 }) {
   // Get columns from database
   const pool = new Pool({ connectionString: databaseUrl });
@@ -27,16 +29,25 @@ export async function generate({
 
   // Initialize ts-morph project
   const project = new Project();
-  const tablesFile = project.createSourceFile(tablesPath ?? "./tables.ts", {}, { overwrite: true });
-  const relationshipsFile = project.createSourceFile(
-    relationshipsPath ?? "./relationships.ts",
-    {},
-    { overwrite: true }
-  );
+  let tablesFile = project.addSourceFileAtPathIfExists(tablesPath);
+  let relationshipsFile = project.addSourceFileAtPathIfExists(relationshipsPath);
+
+  // Create tables file if it doesn't exist or forceGenerate is true
+  if (!tablesFile || forceGenerate) {
+    tablesFile = project.createSourceFile(tablesPath, {}, { overwrite: true });
+  }
+
+  // Create relationships file if it doesn't exist or forceGenerate is true
+  if (!relationshipsFile || forceGenerate) {
+    relationshipsFile = project.createSourceFile(relationshipsPath, {}, { overwrite: true });
+  }
 
   const tables = columns.reduce<
     Record<string, { primaryKey: string; columns: ZeroColumn[]; relationships: ZeroRelationship[] }>
   >((acc, column) => {
+    const isPrimaryKey = isPrimaryKeyColumn(column);
+    const isForeignKey = isForeignKeyColumn(column);
+
     // Add table to accumulator if it doesn't exist
     if (!acc[column.table_name]) {
       acc[column.table_name] = { primaryKey: "", columns: [], relationships: [] };
@@ -47,19 +58,17 @@ export async function generate({
       acc[column.foreign_table_name] = { primaryKey: "", columns: [], relationships: [] };
     }
 
+    if (isPrimaryKey) {
+      acc[column.table_name].primaryKey = column.column_name;
+    }
+
+    // TODO: CHECK IF COLUMN EXISTS IN TABLE
     acc[column.table_name].columns.push({
       name: column.column_name,
       isOptional: column.is_nullable === "YES",
       type: columnDataTypeToZeroType(column.data_type),
     });
 
-    const isPrimaryKey = isPrimaryKeyColumn(column);
-    const isForeignKey = isForeignKeyColumn(column);
-
-    if (isPrimaryKey) {
-      acc[column.table_name].primaryKey = column.column_name;
-    }
-
     if (isForeignKey) {
       const foreignColumn = columns.find(
         ({ table_name, column_name }) =>
@@ -73,6 +82,7 @@ export async function generate({
         ({ destSchema }) => destSchema === column.foreign_table_name
       );
 
+      // TODO: CHECK IF RELATIONSHIP EXISTS IN TABLE, FIND A MATCH WITH SAME SOURCE FIELD, DEST SCHEMA AND DEST FIELD
       acc[column.table_name].relationships.push({
         name:
           relationshipsWithSameForeignKeyTable.length >= 1
@@ -101,7 +111,7 @@ export async function generate({
   // Add import statements to tables file
   tablesFile.addImportDeclaration({
     moduleSpecifier: "@rocicorp/zero",
-    namedImports: ["boolean", "number", "string", "table"], // TODO: Add only required imports
+    namedImports: ["boolean", "number", "string", "json", "table"], // TODO: Add only required imports
   });
 
   // Add table definitions to tables file

src/types.ts

@@ -12,9 +12,11 @@ export interface Column {
   foreign_column_name: string;
 }
 
+export type ZeroColumnType = "string" | "number" | "boolean" | "json";
+
 export interface ZeroColumn {
   name: string;
-  type: "string" | "number" | "boolean";
+  type: ZeroColumnType;
   isOptional: boolean;
 }
 

src/utils.ts

@@ -1,4 +1,4 @@
-import { Column, ZeroColumn, ZeroRelationship } from "./types";
+import { Column, ZeroColumn, ZeroColumnType, ZeroRelationship } from "./types";
 import { Pool } from "pg";
 
 export async function getColumns(pool: Pool): Promise<Column[]> {
@@ -51,7 +51,7 @@ export async function getColumns(pool: Pool): Promise<Column[]> {
   return result.rows;
 }
 
-export function columnDataTypeToZeroType(dataType: string) {
+export function columnDataTypeToZeroType(dataType: string): ZeroColumnType {
   switch (dataType) {
     case "text":
     case "character varying":
@@ -80,6 +80,9 @@ export function columnDataTypeToZeroType(dataType: string) {
       return "number";
     case "boolean":
       return "boolean";
+    case "json":
+    case "jsonb":
+      return "json";
     default:
       return "string";
   }

yarn.lock

@@ -958,7 +958,7 @@ detect-libc@^2.0.0:
   resolved "https://registry.yarnpkg.com/detect-libc/-/detect-libc-2.0.3.tgz#f0cd503b40f9939b894697d19ad50895e30cf700"
   integrity sha512-bwy0MGW55bG41VqxxypOsdSdGqLwXPI/focwgTYCFMbdUiBAxLg9CFzG08sz2aqzknwiX7Hkl0bQENjg8iLByw==
 
-dotenv@^16.4.5:
+dotenv@^16.4.5, dotenv@^16.4.7:
   version "16.4.7"
   resolved "https://registry.yarnpkg.com/dotenv/-/dotenv-16.4.7.tgz#0e20c5b82950140aa99be360a8a5f52335f53c26"
   integrity sha512-47qPchRCykZC03FhkYAhrvwU4xDBFIj1QPqaarj6mdM/hgUzfPHcpkHJOn3mJAufFeeAxAzeGsr5X0M4k6fLZQ==