added the primary key field to source streams

src/components/Connections/CreateConnectionForm.tsx

@@ -1,7 +1,16 @@
 import React, { useContext, useEffect, useMemo, useRef, useState } from 'react';
 import useSWR from 'swr';
 import CustomDialog from '../Dialog/CustomDialog';
-import { Autocomplete, Box, Button, Switch, Select, MenuItem, TextField } from '@mui/material';
+import {
+  Autocomplete,
+  Box,
+  Button,
+  Switch,
+  Select,
+  MenuItem,
+  TextField,
+  Checkbox,
+} from '@mui/material';
 import { Table, TableBody, TableCell, TableHead, TableRow } from '@mui/material';
 import { Controller, useForm } from 'react-hook-form';
 import { httpGet, httpPost, httpPut } from '@/helpers/http';
@@ -25,6 +34,10 @@
   sourceDefinedCursor: boolean;
   cursorFieldOptions: string[];
 };
+type PrimayKeyConfig = {
+  sourceDefinedPrimaryKey: boolean;
+  primaryKeyOptions: string[];
+};
 
 interface SourceStream {
   name: string;
@@ -34,6 +47,8 @@
   destinationSyncMode: string; // append | overwrite | append_dedup
   cursorFieldConfig: CursorFieldConfig; // this will not be posted to backend
   cursorField: string;
+  primaryKeyConfig: PrimayKeyConfig;
+  primaryKey: string;
 }
 
 const CreateConnectionForm = ({
@@ -61,7 +76,6 @@
   const isAnyCursorAbsent = useMemo(() => {
     return filteredSourceStreams.some((stream) => !stream.cursorField);
   }, [filteredSourceStreams]);
-  console.log(filteredSourceStreams, 'filtered source stream');
   const [loading, setLoading] = useState<boolean>(false);
   const [someStreamSelected, setSomeStreamSelected] = useState<boolean>(false);
   const [normalize, setNormalize] = useState<boolean>(false);
@@ -91,6 +105,11 @@
           cursorFieldOptions: [],
         },
         cursorField: '',
+        primaryKeyConfig: {
+          sourceDefinedPrimaryKey: false,
+          primaryKeyOptions: [],
+        },
+        primaryKey: [], // eg.[[id]], [[id], [airbyte_raw]]etc. this can be multiple hence we have to make it an array. This can be a composite primary key.
       };
 
       const cursorFieldObj = stream.cursorFieldConfig;
@@ -119,6 +138,26 @@
           stream.cursorField = el.config.cursorField.length > 0 ? el.config.cursorField[0] : '';
       }
 
+      const primaryKeyObj = stream.primaryKeyConfig;
+
+      if ('sourceDefinedPrimaryKey' in el.stream && el.stream.sourceDefinedPrimaryKey.length > 0)
+        primaryKeyObj.sourceDefinedPrimaryKey = true;
+
+      if (primaryKeyObj.sourceDefinedPrimaryKey) {
+        stream.primaryKey = el.config.primaryKey.flat();
+        primaryKeyObj.primaryKeyOptions = el.config.primaryKey.flat();
+      } else {
+        // user needs to define the primary key
+        // available options are picked from the stream's jsonSchema (cols)
+        if ('jsonSchema' in el.stream)
+          primaryKeyObj.primaryKeyOptions = Object.keys(el.stream.jsonSchema.properties) as any;
+
+        // overwrite default if the primary key field is set
+        if ('primaryKey' in el.config) {
+          stream.primaryKey = el.config.primaryKey.length > 0 ? el.config.primaryKey.flat() : [];
+        }
+      }
+
       return stream;
     });
 
@@ -240,6 +279,7 @@
           syncMode: stream.syncMode, // incremental | full_refresh
           destinationSyncMode: stream.destinationSyncMode, // append | overwrite | append_dedup
           cursorField: stream.cursorField,
+          primaryKey: stream.primaryKey,
         };
       }),
       normalize,
@@ -326,6 +366,10 @@
     updateThisStreamTo_(stream, { ...stream, cursorField: value });
   };
 
+  const updatePrimaryKey = (value: string, stream: SourceStream) => {
+    updateThisStreamTo_(stream, { ...stream, primaryKey: value });
+  };
+
   const handleSyncAllStreams = (checked: boolean) => {
     setSelectAllStreams(checked);
     if (!checked && incrementalAllStreams) {
@@ -477,6 +521,9 @@
                     <TableCell key="cursorfield" align="center">
                       Cursor Field
                     </TableCell>
+                    <TableCell key="primarykey" align="center">
+                      Primary Key
+                    </TableCell>
                   </TableRow>
                   <TableRow>
                     <TableCell key="searchstream" align="center">
@@ -593,6 +640,44 @@
                               )}
                             </Select>
                           </TableCell>
+                          <TableCell key="primarykey" align="center">
+                            <Select
+                              data-testid={`stream-primarykey-${idx}`}
+                              disabled={
+                                !stream.selected ||
+                                !stream.supportsIncremental ||
+                                stream.syncMode !== 'incremental' ||
+                                stream.destinationSyncMode !== 'append_dedup'
+                              }
+                              required={ifIncremental}
+                              onInvalid={(e: any) =>
+                                e.target.setCustomValidity(
+                                  'Primary Key is required for incremental streams'
+                                )
+                              }
+                              multiple
+                              value={stream.primaryKey}
+                              onChange={(event) => {
+                                if (!stream.primaryKeyConfig.sourceDefinedPrimaryKey) {
+                                  updatePrimaryKey(event.target.value, stream);
+                                }
+                              }}
+                              renderValue={(selected: any) => selected.join(', ')}
+                            >
+                              {stream.primaryKeyConfig?.primaryKeyOptions?.length > 0 &&
+                                stream.primaryKeyConfig.primaryKeyOptions.map(
+                                  (option: string, index: number) => (
+                                    <MenuItem key={option} value={option}>
+                                      <Checkbox
+                                        checked={stream.primaryKey.indexOf(option) > -1}
+                                        disabled={stream.primaryKeyConfig.sourceDefinedPrimaryKey}
+                                      />
+                                      {option}
+                                    </MenuItem>
+                                  )
+                                )}
+                            </Select>
+                          </TableCell>
                         </TableRow>
                       );
                     })}

src/components/Connections/__tests__/CreateConnection.test.tsx

@@ -19,6 +19,16 @@ jest.mock('@/helpers/websocket', () => ({
   generateWebsocketUrl: jest.fn(),
 }));
 describe('Create connection', () => {
+  let mockServer: Server;
+
+  afterEach(async () => {
+    if (mockServer) {
+      await new Promise((resolve) => {
+        mockServer.stop(resolve); // Fully stop the WebSocket server
+      });
+    }
+  });
+
   const mockSession: Session = {
     expires: 'false',
     user: { email: 'a' },
@@ -47,7 +57,7 @@ describe('Create connection', () => {
   ];
 
   it('mocks generateWebsocketUrl and tests WebSocket connection', async () => {
-    const mockServer = new Server('wss://mock-websocket-url');
+    mockServer = new Server('wss://mock-websocket-url');
     const mockGenerateWebsocketUrl = generateWebsocketUrl;
 
     // Mock the generateWebsocketUrl function
@@ -98,7 +108,7 @@ describe('Create connection', () => {
     );
 
     // Clean up the mock WebSocket server
-    mockServer.close();
+    //mockServer.close();
   });
 
   it('renders the form', () => {
@@ -132,7 +142,7 @@ describe('Create connection', () => {
   });
 
   it('checks source stream selection and WebSocket interactions', async () => {
-    const mockServer = new Server('wss://mock-websocket-url');
+    mockServer = new Server('wss://mock-websocket-url');
     const mockGenerateWebsocketUrl = generateWebsocketUrl;
 
     // Mock the generateWebsocketUrl function
@@ -224,20 +234,21 @@ describe('Create connection', () => {
       expect(sourceStreamTableRows.length).toBe(STREAMS.length + 2); // Header + Streams
       // Step 4: Validate table headers
       const headerCells = within(sourceStreamTableRows[0]).getAllByRole('columnheader');
-      expect(headerCells.length).toBe(5);
+      expect(headerCells.length).toBe(6);
       expect(headerCells[0].textContent).toBe('Stream');
       expect(headerCells[1].textContent).toBe('Sync?');
       expect(headerCells[2].textContent).toBe('Incremental?');
       expect(headerCells[3].textContent).toBe('Destination');
       expect(headerCells[4].textContent).toBe('Cursor Field');
+      expect(headerCells[5].textContent).toBe('Primary Key');
 
       // Clean up the mock WebSocket server
-      mockServer.close();
+      //mockServer.close();
     });
   });
 
   it('create connection success with WebSocket and fetch interactions', async () => {
-    const mockServer = new Server('wss://mock-websocket-url');
+    mockServer = new Server('wss://mock-websocket-url');
     const mockGenerateWebsocketUrl = generateWebsocketUrl;
 
     // Mock the generateWebsocketUrl function
@@ -345,11 +356,11 @@ describe('Create connection', () => {
     expect(mockCreateConnectionFetch).toHaveBeenCalledTimes(1);
 
     // Clean up mock server
-    mockServer.close();
+    //mockServer.close();
   });
 
   it('create connection failed with WebSocket and fetch interactions', async () => {
-    const mockServer = new Server('wss://mock-websocket-url');
+    mockServer = new Server('wss://mock-websocket-url');
     const mockGenerateWebsocketUrl = generateWebsocketUrl;
 
     // Mock the generateWebsocketUrl function
@@ -457,6 +468,6 @@ describe('Create connection', () => {
     expect(mockCreateConnectionFetch).toHaveBeenCalledTimes(1);
 
     //check  how to show this test failed.
-    mockServer.close();
+    //mockServer.close();
   });
 });