feat: Support fetching distributed table metadata with cluster()

Author: pulpdrew

.changeset/tidy-phones-brake.md

@@ -0,0 +1,5 @@
+---
+"@hyperdx/common-utils": patch
+---
+
+feat: Support fetching distribbuted table metadata with cluster()

packages/common-utils/src/__tests__/metadata.test.ts

@@ -192,6 +192,62 @@ describe('Metadata', () => {
       expect(result!.partition_key).toEqual('column1');
     });
 
+    it('should query via cluster() for Distributed table underlying metadata', async () => {
+      const distributedMetadata = {
+        database: 'test_db',
+        name: 'dist_table',
+        engine: 'Distributed',
+        engine_full:
+          "Distributed('my_cluster', 'test_db', 'local_table', rand())",
+        partition_key: '',
+        sorting_key: '',
+        primary_key: '',
+        sampling_key: '',
+        create_table_query: 'CREATE TABLE test_db.dist_table ...',
+      };
+
+      const localMetadata = {
+        database: 'test_db',
+        name: 'local_table',
+        engine: 'MergeTree',
+        engine_full: 'MergeTree() ORDER BY id',
+        partition_key: 'toYYYYMM(timestamp)',
+        sorting_key: 'id, timestamp',
+        primary_key: 'id',
+        sampling_key: '',
+        create_table_query: 'CREATE TABLE test_db.local_table ...',
+      };
+
+      let callCount = 0;
+      (mockClickhouseClient.query as jest.Mock).mockImplementation(() => {
+        callCount++;
+        return Promise.resolve({
+          json: jest.fn().mockResolvedValue({
+            data: [callCount === 1 ? distributedMetadata : localMetadata],
+          }),
+        });
+      });
+
+      const result = await metadata.getTableMetadata({
+        databaseName: 'test_db',
+        tableName: 'dist_table',
+        connectionId: 'test_connection',
+      });
+
+      // Two queries: one for the distributed table, one via cluster() for the local table
+      expect(callCount).toBe(2);
+      expect(result!.engine).toBe('MergeTree');
+      expect(result!.sorting_key).toBe('id, timestamp');
+      expect(result!.create_local_table_query).toBe(
+        'CREATE TABLE test_db.local_table ...',
+      );
+      // The second query should use cluster() - verify it references system.tables via cluster
+      const secondQuery = (mockClickhouseClient.query as jest.Mock).mock
+        .calls[1][0].query;
+      expect(secondQuery).toContain('cluster(');
+      expect(secondQuery).toContain('system.tables');
+    });
+
     it('should use the cache when retrieving table metadata', async () => {
       // Setup the mock implementation
       mockCache.getOrFetch.mockReset();

packages/common-utils/src/__tests__/utils.test.ts

@@ -17,8 +17,8 @@ import {
   findJsonExpressions,
   formatDate,
   getAlignedDateRange,
+  getDistributedTableArgs,
   getFirstOrderingItem,
-  getLocalTableFromDistributedTable,
   isFirstOrderByAscending,
   isJsonExpression,
   isTimestampExpressionInFirstOrderBy,
@@ -30,7 +30,6 @@ import {
   replaceJsonExpressions,
   splitAndTrimCSV,
   splitAndTrimWithBracket,
-  TextIndexTokenizer,
 } from '../core/utils';
 
 describe('utils', () => {
@@ -1996,61 +1995,85 @@ describe('utils', () => {
       ({ engine_full: engineFull }) as any;
 
     it('parses a simple Distributed engine_full', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata("Distributed('default', 'mydb', 'local_table', rand())"),
       );
-      expect(result).toEqual({ database: 'mydb', table: 'local_table' });
+      expect(result).toEqual({
+        cluster: 'default',
+        database: 'mydb',
+        table: 'local_table',
+      });
     });
 
     it('parses without a sharding key', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata("Distributed('cluster', 'db', 'tbl')"),
       );
-      expect(result).toEqual({ database: 'db', table: 'tbl' });
+      expect(result).toEqual({
+        cluster: 'cluster',
+        database: 'db',
+        table: 'tbl',
+      });
     });
 
     it('handles double-quoted identifiers', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata('Distributed("cluster", "my_database", "my_table")'),
       );
-      expect(result).toEqual({ database: 'my_database', table: 'my_table' });
+      expect(result).toEqual({
+        cluster: 'cluster',
+        database: 'my_database',
+        table: 'my_table',
+      });
     });
 
     it('handles backtick-quoted identifiers', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata("Distributed('cluster', `mydb`, `local_tbl`, rand())"),
       );
-      expect(result).toEqual({ database: 'mydb', table: 'local_tbl' });
+      expect(result).toEqual({
+        cluster: 'cluster',
+        database: 'mydb',
+        table: 'local_tbl',
+      });
     });
 
     it('handles unquoted identifiers', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata('Distributed(cluster, mydb, local_tbl, rand())'),
       );
-      expect(result).toEqual({ database: 'mydb', table: 'local_tbl' });
+      expect(result).toEqual({
+        cluster: 'cluster',
+        database: 'mydb',
+        table: 'local_tbl',
+      });
     });
 
     it('returns undefined when engine_full has fewer than 3 args', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata("Distributed('cluster', 'db')"),
       );
       expect(result).toBeUndefined();
     });
 
     it('returns undefined when engine_full does not match Distributed pattern', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata('MergeTree() ORDER BY id'),
       );
       expect(result).toBeUndefined();
     });
 
     it('handles a complex sharding expression with nested parentheses', () => {
-      const result = getLocalTableFromDistributedTable(
+      const result = getDistributedTableArgs(
         makeMetadata(
           "Distributed('cluster', 'db', 'tbl', sipHash64(UserID, EventDate))",
         ),
       );
-      expect(result).toEqual({ database: 'db', table: 'tbl' });
+      expect(result).toEqual({
+        cluster: 'cluster',
+        database: 'db',
+        table: 'tbl',
+      });
     });
   });
 });