@comprehend/telemetry-node 0.1.2 → 0.1.4

package/src/ComprehendDevSpanProcessor.ts

@@ -144,6 +144,7 @@ export class ComprehendDevSpanProcessor implements SpanProcessor {
             ...(environment ? { environment } : {})
         };
         this.ingestMessage(message);
+        return newService;
     }
 
     private processHTTPRoute(service: ObservedService, route: string, method: string, span: ReadableSpan): void {

package/src/WebSocketConnection.test.ts

@@ -0,0 +1,535 @@
+import { WebSocketConnection } from './WebSocketConnection';
+import {
+    InitMessage,
+    NewObservedServiceMessage,
+    NewObservedHttpRouteMessage,
+    ObservationMessage,
+    HttpServerObservation,
+    ObservationOutputMessage
+} from './wire-protocol';
+
+// Mock the ws WebSocket library
+jest.mock('ws');
+
+import WebSocket from 'ws';
+
+const MockedWebSocket = WebSocket as jest.MockedClass<typeof WebSocket>;
+
+// Mock WebSocket instance
+interface MockWebSocketInstance {
+    on: jest.Mock;
+    send: jest.Mock;
+    close: jest.Mock;
+    readyState: number;
+    OPEN: number;
+    CLOSED: number;
+    // Event handlers that tests can trigger
+    _triggerOpen: () => void;
+    _triggerMessage: (data: string) => void;
+    _triggerClose: (code: number, reason: string) => void;
+    _triggerError: (error: Error) => void;
+}
+
+describe('WebSocketConnection', () => {
+    let mockSocket: MockWebSocketInstance;
+    let connection: WebSocketConnection;
+    let logMessages: string[];
+    let mockLogger: jest.Mock;
+
+    beforeEach(() => {
+        // Reset all mocks
+        jest.clearAllMocks();
+        jest.clearAllTimers();
+        jest.useFakeTimers();
+
+        logMessages = [];
+        mockLogger = jest.fn((message: string) => {
+            logMessages.push(message);
+        });
+
+        // Create mock WebSocket instance with event handling
+        mockSocket = {
+            on: jest.fn(),
+            send: jest.fn(),
+            close: jest.fn(),
+            readyState: WebSocket.CLOSED,
+            OPEN: WebSocket.OPEN,
+            CLOSED: WebSocket.CLOSED,
+            _triggerOpen: function() { /* will be set up below */ },
+            _triggerMessage: function(data: string) { /* will be set up below */ },
+            _triggerClose: function(code: number, reason: string) { /* will be set up below */ },
+            _triggerError: function(error: Error) { /* will be set up below */ }
+        };
+
+        // Set up event handler storage and triggering
+        const eventHandlers: Record<string, Function> = {};
+        mockSocket.on.mockImplementation((event: string, handler: Function) => {
+            eventHandlers[event] = handler;
+        });
+
+        mockSocket._triggerOpen = () => {
+            mockSocket.readyState = WebSocket.OPEN;
+            eventHandlers['open']?.();
+        };
+        mockSocket._triggerMessage = (data: string) => {
+            const buffer = Buffer.from(data);
+            eventHandlers['message']?.(buffer);
+        };
+        mockSocket._triggerClose = (code: number, reason: string) => {
+            mockSocket.readyState = WebSocket.CLOSED;
+            eventHandlers['close']?.(code, Buffer.from(reason));
+        };
+        mockSocket._triggerError = (error: Error) => {
+            eventHandlers['error']?.(error);
+        };
+
+        // Mock WebSocket constructor to return our mock instance
+        MockedWebSocket.mockImplementation(() => mockSocket as any);
+    });
+
+    afterEach(() => {
+        jest.useRealTimers();
+        if (connection) {
+            connection.close();
+        }
+    });
+
+    describe('Connection Establishment', () => {
+        it('should create WebSocket connection with correct URL', () => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+
+            expect(MockedWebSocket).toHaveBeenCalledWith('wss://ingestion.comprehend.dev/test-org/observations');
+            expect(mockSocket.on).toHaveBeenCalledWith('open', expect.any(Function));
+            expect(mockSocket.on).toHaveBeenCalledWith('message', expect.any(Function));
+            expect(mockSocket.on).toHaveBeenCalledWith('close', expect.any(Function));
+            expect(mockSocket.on).toHaveBeenCalledWith('error', expect.any(Function));
+        });
+
+        it('should send init message on connection open', () => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+
+            mockSocket._triggerOpen();
+
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify({
+                event: 'init',
+                protocolVersion: 1,
+                token: 'test-token'
+            }));
+            expect(logMessages).toContain('WebSocket connected. Sending init/auth message.');
+        });
+
+        it('should work without logger', () => {
+            connection = new WebSocketConnection('test-org', 'test-token');
+
+            mockSocket._triggerOpen();
+
+            expect(MockedWebSocket).toHaveBeenCalled();
+            expect(mockSocket.send).toHaveBeenCalled();
+        });
+    });
+
+    describe('Authorization Flow', () => {
+        beforeEach(() => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            mockSocket._triggerOpen();
+            jest.clearAllMocks(); // Clear the init message send
+        });
+
+        it('should handle authorization acknowledgment', () => {
+            const authAck: ObservationOutputMessage = {
+                type: 'ack-authorized'
+            };
+
+            mockSocket._triggerMessage(JSON.stringify(authAck));
+
+            expect(logMessages).toContain('Authorization acknowledged by server.');
+        });
+
+        it('should replay queued messages after authorization', () => {
+            // Queue some messages before authorization
+            const serviceMessage: NewObservedServiceMessage = {
+                event: 'new-entity',
+                type: 'service',
+                hash: 'test-hash-1',
+                name: 'test-service'
+            };
+
+            const observationMessage: ObservationMessage = {
+                event: 'observations',
+                seq: 1,
+                observations: [{
+                    type: 'http-server',
+                    subject: 'test-subject',
+                    timestamp: [1700000000, 0],
+                    path: '/test',
+                    status: 200,
+                    duration: [0, 100000000]
+                } as HttpServerObservation]
+            };
+
+            connection.sendMessage(serviceMessage);
+            connection.sendMessage(observationMessage);
+
+            // Should not send immediately (not authorized yet)
+            expect(mockSocket.send).not.toHaveBeenCalled();
+
+            // Now authorize
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+
+            // Should replay both messages
+            expect(mockSocket.send).toHaveBeenCalledTimes(2);
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(serviceMessage));
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(observationMessage));
+        });
+
+        it('should send messages immediately when already authorized', () => {
+            // First authorize
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            jest.clearAllMocks();
+
+            const serviceMessage: NewObservedServiceMessage = {
+                event: 'new-entity',
+                type: 'service',
+                hash: 'test-hash',
+                name: 'test-service'
+            };
+
+            connection.sendMessage(serviceMessage);
+
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(serviceMessage));
+        });
+    });
+
+    describe('Message Acknowledgment Handling', () => {
+        beforeEach(() => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            mockSocket._triggerOpen();
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            jest.clearAllMocks();
+        });
+
+        it('should handle entity/interaction acknowledgments', () => {
+            const serviceMessage: NewObservedServiceMessage = {
+                event: 'new-entity',
+                type: 'service',
+                hash: 'service-hash',
+                name: 'test-service'
+            };
+
+            const routeMessage: NewObservedHttpRouteMessage = {
+                event: 'new-entity',
+                type: 'http-route',
+                hash: 'route-hash',
+                parent: 'service-hash',
+                method: 'GET',
+                route: '/test'
+            };
+
+            // Send messages
+            connection.sendMessage(serviceMessage);
+            connection.sendMessage(routeMessage);
+
+            // Acknowledge the service message
+            mockSocket._triggerMessage(JSON.stringify({
+                type: 'ack-observed',
+                hash: 'service-hash'
+            }));
+
+            // Clear previous calls before reconnection test
+            mockSocket.send.mockClear();
+
+            // If we reconnect, only the route message should be replayed
+            mockSocket._triggerClose(1000, 'test close');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
+
+            // Clear the init message that gets sent on reconnection
+            mockSocket.send.mockClear();
+
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+
+            // Should only send the unacknowledged route message
+            expect(mockSocket.send).toHaveBeenCalledTimes(1);
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(routeMessage));
+        });
+
+        it('should handle observation acknowledgments', () => {
+            const observationMessage1: ObservationMessage = {
+                event: 'observations',
+                seq: 1,
+                observations: []
+            };
+
+            const observationMessage2: ObservationMessage = {
+                event: 'observations',
+                seq: 2,
+                observations: []
+            };
+
+            // Send messages
+            connection.sendMessage(observationMessage1);
+            connection.sendMessage(observationMessage2);
+
+            // Acknowledge first observation
+            mockSocket._triggerMessage(JSON.stringify({
+                type: 'ack-observations',
+                seq: 1
+            }));
+
+            // Clear previous calls before reconnection test
+            mockSocket.send.mockClear();
+
+            // Reconnect - only second observation should be replayed
+            mockSocket._triggerClose(1000, 'test close');
+            jest.advanceTimersByTime(1000);
+            mockSocket._triggerOpen();
+
+            // Clear the init message that gets sent on reconnection
+            mockSocket.send.mockClear();
+
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+
+            // Should only send the unacknowledged second observation message
+            expect(mockSocket.send).toHaveBeenCalledTimes(1);
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(observationMessage2));
+        });
+    });
+
+    describe('Reconnection Logic', () => {
+        beforeEach(() => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            mockSocket._triggerOpen();
+        });
+
+        it('should reconnect automatically on close', () => {
+            const setTimeoutSpy = jest.spyOn(global, 'setTimeout');
+
+            mockSocket._triggerClose(1000, 'Normal closure');
+
+            expect(logMessages).toContain('WebSocket disconnected. Code: 1000, Reason: Normal closure');
+
+            // Should schedule reconnection
+            expect(setTimeoutSpy).toHaveBeenCalledWith(expect.any(Function), 1000);
+
+            // Advance timer to trigger reconnection
+            jest.advanceTimersByTime(1000);
+
+            // Should create new WebSocket connection
+            expect(MockedWebSocket).toHaveBeenCalledTimes(2);
+
+            setTimeoutSpy.mockRestore();
+        });
+
+        it('should reset authorization state on close', () => {
+            // First authorize
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+
+            // Send a message (should send immediately)
+            const testMessage: NewObservedServiceMessage = {
+                event: 'new-entity',
+                type: 'service',
+                hash: 'test-hash',
+                name: 'test-service'
+            };
+            connection.sendMessage(testMessage);
+            expect(mockSocket.send).toHaveBeenCalled();
+
+            jest.clearAllMocks();
+
+            // Close connection
+            mockSocket._triggerClose(1000, 'test');
+
+            // Send another message - should queue (not send immediately)
+            connection.sendMessage(testMessage);
+            expect(mockSocket.send).not.toHaveBeenCalled();
+        });
+
+        it('should not reconnect when explicitly closed', () => {
+            connection.close();
+            mockSocket._triggerClose(1000, 'Explicit close');
+
+            // Should not schedule reconnection
+            jest.advanceTimersByTime(5000);
+            expect(MockedWebSocket).toHaveBeenCalledTimes(1); // Only initial connection
+        });
+    });
+
+    describe('Error Handling', () => {
+        beforeEach(() => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+        });
+
+        it('should log WebSocket errors', () => {
+            const testError = new Error('Connection failed');
+            mockSocket._triggerError(testError);
+
+            expect(logMessages).toContain('WebSocket encountered an error: Connection failed');
+        });
+
+        it('should handle malformed server messages gracefully', () => {
+            mockSocket._triggerOpen();
+            mockSocket._triggerMessage('{invalid json}');
+
+            // Check that an error message was logged (exact message varies by Node version)
+            const errorLogExists = logMessages.some(msg =>
+                msg.startsWith('Error parsing message from server:') &&
+                msg.includes('JSON')
+            );
+            expect(errorLogExists).toBe(true);
+        });
+
+        it('should handle non-string error objects in message parsing', () => {
+            mockSocket._triggerOpen();
+
+            // Mock JSON.parse to throw a non-Error object
+            const originalParse = JSON.parse;
+            JSON.parse = jest.fn().mockImplementation(() => {
+                throw 'string error';
+            });
+
+            mockSocket._triggerMessage('{}');
+
+            expect(logMessages).toContain('Error parsing message from server: string error');
+
+            // Restore original JSON.parse
+            JSON.parse = originalParse;
+        });
+
+        it('should not send messages when socket is not open', () => {
+            mockSocket.readyState = WebSocket.CLOSED;
+            mockSocket._triggerOpen();
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+
+            // Clear the init and any previous messages
+            mockSocket.send.mockClear();
+
+            // Force socket to closed state
+            mockSocket.readyState = WebSocket.CLOSED;
+
+            const testMessage: NewObservedServiceMessage = {
+                event: 'new-entity',
+                type: 'service',
+                hash: 'test-hash',
+                name: 'test-service'
+            };
+
+            connection.sendMessage(testMessage);
+
+            // Should not call send when socket is closed
+            expect(mockSocket.send).not.toHaveBeenCalled();
+        });
+    });
+
+    describe('Connection Lifecycle', () => {
+        it('should close WebSocket connection properly', () => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+
+            connection.close();
+
+            expect(mockSocket.close).toHaveBeenCalled();
+        });
+
+        it('should handle close when socket is null', () => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+
+            // Simulate socket being null
+            (connection as any).socket = null;
+
+            expect(() => connection.close()).not.toThrow();
+        });
+    });
+
+    describe('Message Types and Queuing', () => {
+        beforeEach(() => {
+            connection = new WebSocketConnection('test-org', 'test-token', mockLogger);
+            mockSocket._triggerOpen();
+        });
+
+        it('should queue new-entity messages', () => {
+            // Clear the initial init message
+            mockSocket.send.mockClear();
+
+            const serviceMessage: NewObservedServiceMessage = {
+                event: 'new-entity',
+                type: 'service',
+                hash: 'service-hash',
+                name: 'test-service'
+            };
+
+            connection.sendMessage(serviceMessage);
+
+            // Should be queued, not sent immediately (not authorized)
+            expect(mockSocket.send).not.toHaveBeenCalled();
+
+            // Authorize and check replay
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(serviceMessage));
+        });
+
+        it('should queue new-interaction messages', () => {
+            // Clear the initial init message
+            mockSocket.send.mockClear();
+
+            const httpRequestMessage = {
+                event: 'new-interaction' as const,
+                type: 'http-request' as const,
+                hash: 'request-hash',
+                from: 'service-hash',
+                to: 'http-service-hash'
+            };
+
+            connection.sendMessage(httpRequestMessage);
+
+            // Should be queued
+            expect(mockSocket.send).not.toHaveBeenCalled();
+
+            // Authorize and check replay
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(httpRequestMessage));
+        });
+
+        it('should queue observation messages', () => {
+            // Clear the initial init message
+            mockSocket.send.mockClear();
+
+            const observationMessage: ObservationMessage = {
+                event: 'observations',
+                seq: 1,
+                observations: [{
+                    type: 'http-server',
+                    subject: 'route-hash',
+                    timestamp: [1700000000, 0],
+                    path: '/test',
+                    status: 200,
+                    duration: [0, 50000000]
+                } as HttpServerObservation]
+            };
+
+            connection.sendMessage(observationMessage);
+
+            // Should be queued
+            expect(mockSocket.send).not.toHaveBeenCalled();
+
+            // Authorize and check replay
+            mockSocket._triggerMessage(JSON.stringify({ type: 'ack-authorized' }));
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(observationMessage));
+        });
+
+        it('should not queue init messages in unacknowledged maps', () => {
+            // This is implicit - init messages are sent via sendRaw, not sendMessage
+            // and don't go through the queuing logic
+
+            const initMessage: InitMessage = {
+                event: 'init',
+                protocolVersion: 1,
+                token: 'test-token'
+            };
+
+            // Manually test sendRaw behavior (private method testing)
+            mockSocket.readyState = WebSocket.OPEN;
+            (connection as any).sendRaw(initMessage);
+
+            expect(mockSocket.send).toHaveBeenCalledWith(JSON.stringify(initMessage));
+        });
+    });
+});

package/src/sql-analyzer.test.ts

@@ -580,4 +580,20 @@ describe('SQL Analyzer - bulk INSERT VALUES cardinality reduction', () => {
         expect(result.presentableQuery).toEqual(`INSERT INTO comments (text, author) VALUES
             (...)`);
     });
+
+    it('preserves whitespace before ON CONFLICT after VALUES clause', () => {
+        const sql = `INSERT INTO users (name, email) VALUES ('Alice', 'alice@example.com') ON CONFLICT (email) DO NOTHING`;
+        const result = analyzeSQL(sql);
+
+        expect(result.tableOperations).toEqual({ users: ['INSERT'] });
+        expect(result.presentableQuery).toEqual(`INSERT INTO users (name, email) VALUES (...) ON CONFLICT (email) DO NOTHING`);
+    });
+
+    it('preserves whitespace before ON CONFLICT with multiple VALUES tuples', () => {
+        const sql = `INSERT INTO users (name, email) VALUES ('Alice', 'alice@example.com'), ('Bob', 'bob@example.com') ON CONFLICT (email) DO UPDATE SET name = EXCLUDED.name`;
+        const result = analyzeSQL(sql);
+
+        expect(result.tableOperations).toEqual({ users: ['INSERT'] });
+        expect(result.presentableQuery).toEqual(`INSERT INTO users (name, email) VALUES (...) ON CONFLICT (email) DO UPDATE SET name = EXCLUDED.name`);
+    });
 });

package/src/sql-analyzer.ts

@@ -28,6 +28,7 @@ export function analyzeSQL(sql: string): SQLAnalysisResult {
     let skippingValues = false;
     let lookingForCommaOrEnd = false;
     let valuesDepth = 0;
+    let skippedWhitespace: Token[] = [];
     for (let token of tokenizeSQL(sql)) {
         switch (token.type) {
             case "whitespace":
@@ -140,22 +141,30 @@ export function analyzeSQL(sql: string): SQLAnalysisResult {
             switch (token.type) {
                 case "comment":
                 case "whitespace":
-                    // Skip whitespace/comments while looking for comma or end
+                    // Collect whitespace/comments while looking for comma or end
+                    skippedWhitespace.push(token);
                     break;
                 case "punct":
                     if (token.value === ",") {
-                        // More tuples coming, continue skipping
+                        // More tuples coming, clear skipped whitespace and continue skipping
+                        skippedWhitespace = [];
                         lookingForCommaOrEnd = false;
                         skippingValues = true;
                     } else {
                         // Not a comma, so VALUES clause is done
+                        // Add back the skipped whitespace, then the current token
+                        presentableTokens.push(...skippedWhitespace);
                         presentableTokens.push(token);
+                        skippedWhitespace = [];
                         lookingForCommaOrEnd = false;
                     }
                     break;
                 default:
                     // VALUES clause is done, resume normal processing
+                    // Add back the skipped whitespace, then the current token
+                    presentableTokens.push(...skippedWhitespace);
                     presentableTokens.push(token);
+                    skippedWhitespace = [];
                     lookingForCommaOrEnd = false;
                     break;
             }