fix: retain composite PK columns

Author: IderAghbal

src/tables.ts

@@ -264,7 +264,28 @@ const createZeroTableBuilder = <
   const tableColumns = getTableColumns(table);
   const tableConfig = getTableConfigForDatabase(table);
 
-  const primaryKeysFromColumns: string[] = [];
+  const columnNameToStableKey = new Map<string, string>(
+    typedEntries(tableColumns).map(([key, column]) => [
+      column.name,
+      String(key),
+    ]),
+  );
+
+  const primaryKeys = new Set<string>();
+  for (const [key, column] of typedEntries(tableColumns)) {
+    if (column.primary) {
+      primaryKeys.add(String(key));
+    }
+  }
+
+  for (const pk of tableConfig.primaryKeys) {
+    for (const pkColumn of pk.columns) {
+      const key = columnNameToStableKey.get(pkColumn.name);
+      if (key) {
+        primaryKeys.add(String(key));
+      }
+    }
+  }
 
   const isColumnBuilder = (value: unknown): value is ColumnBuilder<any> =>
     typeof value === "object" && value !== null && "schema" in value;
@@ -299,7 +320,7 @@ const createZeroTableBuilder = <
         if (
           columnConfig !== true &&
           !isColumnConfigOverride &&
-          !column.primary
+          !primaryKeys.has(String(key))
         ) {
           debugLog(
             debug,
@@ -335,10 +356,6 @@ const createZeroTableBuilder = <
             ? columnConfig.schema.optional
             : false;
 
-      if (column.primary) {
-        primaryKeysFromColumns.push(String(key));
-      }
-
       if (columnConfig && typeof columnConfig !== "boolean") {
         return {
           ...acc,
@@ -371,19 +388,7 @@ const createZeroTableBuilder = <
     {} as Record<string, any>,
   );
 
-  const primaryKeys = [
-    ...primaryKeysFromColumns,
-    ...tableConfig.primaryKeys.flatMap((k) =>
-      k.columns.map((c) =>
-        getDrizzleColumnKeyFromColumnName({
-          columnName: c.name,
-          table: c.table,
-        }),
-      ),
-    ),
-  ];
-
-  if (!primaryKeys.length) {
+  if (primaryKeys.size === 0) {
     throw new Error(
       `drizzle-zero: No primary keys found in table - ${actualTableName}. Did you forget to define a primary key?`,
     );

tests/config.test.ts

@@ -64,12 +64,11 @@ describe("drizzleZeroConfig with explicit table and column configuration", () =>
   test("should handle explicit table and column configurations", () => {
     const schema = drizzleZeroConfig(drizzleSchema, {
       tables: {
-        users: {
-          name: true,
-          email: false,
-          // phone is not mentioned, should be excluded
+        users: true, // include all columns
+        usersToPosts: {
+          userId: false, // will be included anyway because it is part of the primary key
+          // role is not mentioned, should be excluded
         },
-        usersToPosts: true, // include all columns
         posts: false,
         // comments table is not mentioned, should be excluded
       },
@@ -88,11 +87,11 @@ describe("drizzleZeroConfig with explicit table and column configuration", () =>
     // `users` table should have `id` (pk) and `name`
     expect(
       new Set(Object.keys((schema.tables as any).users.columns)),
-    ).toStrictEqual(new Set(["id", "name"]));
+    ).toStrictEqual(new Set(["id", "name", "email", "phone"]));
 
     // `usersToPosts` table should have all its columns
     expect(
       new Set(Object.keys((schema.tables as any).usersToPosts.columns)),
-    ).toStrictEqual(new Set(["userId", "postId", "role"]));
+    ).toStrictEqual(new Set(["userId", "postId"]));
   });
 });