fix(pg-meta): order primary key columns by key position and exclude INCLUDE columns

packages/mcp-server-supabase/src/pg-meta/tables.sql

@@ -26,26 +26,30 @@ FROM
   pg_namespace nc
   JOIN pg_class c ON nc.oid = c.relnamespace
   left join (
+    -- Walk indkey positionally: `= any (indkey)` discards the position of each
+    -- column within the key, so rows come back in attnum order rather than
+    -- constraint-definition order. indkey also holds the index's INCLUDE payload
+    -- columns, which are not part of the key -- indnkeyatts bounds it to the real
+    -- key columns. Same ordering guarantee the foreign key subquery below relies on.
     select
       table_id,
-      jsonb_agg(_pk.*) as primary_keys
+      jsonb_agg(to_jsonb(_pk) - 'ord' order by _pk.ord) as primary_keys
     from (
       select
         n.nspname as schema,
         c.relname as table_name,
         a.attname as name,
-        c.oid :: int8 as table_id
+        c.oid :: int8 as table_id,
+        k.ord
       from
-        pg_index i,
-        pg_class c,
-        pg_attribute a,
-        pg_namespace n
+        pg_index i
+        join pg_class c on i.indrelid = c.oid
+        join pg_namespace n on c.relnamespace = n.oid
+        cross join lateral unnest(i.indkey :: int2[]) with ordinality as k(attnum, ord)
+        join pg_attribute a on a.attrelid = c.oid and a.attnum = k.attnum
       where
-        i.indrelid = c.oid
-        and c.relnamespace = n.oid
-        and a.attrelid = c.oid
-        and a.attnum = any (i.indkey)
-        and i.indisprimary
+        i.indisprimary
+        and k.ord <= i.indnkeyatts
     ) as _pk
     group by table_id
   ) as pk

packages/mcp-server-supabase/src/server.test.ts

@@ -1559,6 +1559,91 @@ describe('tools', () => {
     ).toHaveLength(1);
   });
 
+  test('composite primary key preserves constraint-definition column order', async () => {
+    const { callTool } = await setup();
+
+    const org = await createOrganization({
+      name: 'My Org',
+      plan: 'free',
+      allowed_release_channels: ['ga'],
+    });
+
+    const project = await createProject({
+      name: 'Project 1',
+      region: 'us-east-1',
+      organization_id: org.id,
+    });
+    project.status = 'ACTIVE_HEALTHY';
+
+    // Key order is deliberately the reverse of the physical column order, so a
+    // regression to attnum ordering fails immediately. This mirrors the ordering
+    // guarantee the foreign key tests above already assert.
+    await project.db.exec(`
+      create table membership (
+        org_id int not null,
+        user_id int not null,
+        role text,
+        primary key (user_id, org_id)
+      );
+    `);
+
+    const result = await callTool({
+      name: 'list_tables',
+      arguments: {
+        project_id: project.id,
+        schemas: ['public'],
+        verbose: true,
+      },
+    });
+
+    const membership = result.tables.find(
+      (t: { name: string }) => t.name === 'public.membership'
+    );
+
+    expect(membership.primary_keys).toEqual(['user_id', 'org_id']);
+  });
+
+  test('primary key excludes non-key INCLUDE columns', async () => {
+    const { callTool } = await setup();
+
+    const org = await createOrganization({
+      name: 'My Org',
+      plan: 'free',
+      allowed_release_channels: ['ga'],
+    });
+
+    const project = await createProject({
+      name: 'Project 1',
+      region: 'us-east-1',
+      organization_id: org.id,
+    });
+    project.status = 'ACTIVE_HEALTHY';
+
+    // `b` is an INCLUDE payload column: carried in the index for index-only scans,
+    // but not part of the key and not covered by the uniqueness guarantee. Reporting
+    // it as a primary key column overstates what the database enforces.
+    await project.db.exec(`
+      create table inc (a int not null, b int not null, c int not null);
+      create unique index inc_pk_idx on inc (a) include (b);
+      alter table inc add constraint inc_pk primary key using index inc_pk_idx;
+    `);
+
+    const result = await callTool({
+      name: 'list_tables',
+      arguments: {
+        project_id: project.id,
+        schemas: ['public'],
+        verbose: true,
+      },
+    });
+
+    const inc = result.tables.find(
+      (t: { name: string }) => t.name === 'public.inc'
+    );
+
+    expect(inc.primary_keys).toEqual(['a']);
+  });
+
   test('list_tables omits advisory when all tables have RLS enabled', async () => {
     const { callTool } = await setup();