diff --git a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.ts b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.ts index 8e399e1bea6..4627eea586b 100644 --- a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.ts +++ b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.ts @@ -30,6 +30,8 @@ import { createColumnTreeItem, getColumnType, createErrorItem, + isRedLakeCatalog, + isS3TablesCatalog, } from './utils' import { createPlaceholderItem } from '../../../shared/treeview/utils' import { Column, Database, Table } from '@aws-sdk/client-glue' @@ -353,7 +355,15 @@ async function getCatalogs( } // For catalogs without children, create regular catalog node - return createCatalogNode(parentCatalog.CatalogId || '', parentCatalog, glueClient, parent, false) + // For RedLake and S3TableCatalog, they are supposed to have children + // Pass isParent as true + return createCatalogNode( + parentCatalog.CatalogId || '', + parentCatalog, + glueClient, + parent, + isRedLakeCatalog(parentCatalog) || isS3TablesCatalog(parentCatalog) + ) }) } @@ -385,7 +395,7 @@ function createCatalogNode( }, // Child catalogs load databases, parent catalogs will have their children provider overridden isParent - ? async () => [] // Placeholder, will be overridden for parent catalogs with children + ? async () => [createPlaceholderItem(NO_DATA_FOUND_MESSAGE) as LakehouseNode] : async (node) => { try { logger.info(`Loading databases for catalog ${catalogId}`) diff --git a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/utils.ts b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/utils.ts index ddf0ed5f3ea..43df51ee911 100644 --- a/packages/core/src/sagemakerunifiedstudio/explorer/nodes/utils.ts +++ b/packages/core/src/sagemakerunifiedstudio/explorer/nodes/utils.ts @@ -24,6 +24,7 @@ import { getContext } from '../../../shared/vscode/setContext' import { SmusAuthenticationProvider } from '../../auth/providers/smusAuthenticationProvider' import { SmusIamConnection } from '../../auth/model' import { ConnectionStatus } from '@aws-sdk/client-datazone' +import { GlueCatalog } from '../../shared/client/glueCatalogClient' /** * Polling interval in milliseconds for checking space status updates @@ -31,6 +32,23 @@ import { ConnectionStatus } from '@aws-sdk/client-datazone' // eslint-disable-next-line @typescript-eslint/naming-convention export const PENDING_NODE_POLLING_INTERVAL_MS = 5000 +/** + * Check if a catalog is a RedLake catalog + */ +export const isRedLakeCatalog = (catalog?: GlueCatalog) => { + return ( + catalog?.FederatedCatalog?.ConnectionName === 'aws:redshift' || + catalog?.CatalogProperties?.DataLakeAccessProperties?.CatalogType === 'aws:redshift' + ) +} + +/** + * Check if a catalog is a S3 table catalog + */ +export const isS3TablesCatalog = (catalog?: GlueCatalog) => { + return catalog?.FederatedCatalog?.ConnectionName === 'aws:s3tables' +} + /** * Gets the label for a node based on its data */ diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts index 63e87c25f23..b12fad6296b 100644 --- a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts +++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/lakehouseStrategy.test.ts @@ -220,9 +220,16 @@ describe('LakehouseStrategy', function () { }) describe('Catalog nodes', function () { - it('should create catalog nodes from API', async function () { + it('should create regular catalog nodes and load databases', async function () { + const regularCatalog = { + CatalogId: 'regular-catalog', + Name: 'regular-catalog', + CatalogType: 'HIVE', + } + mockGlueCatalogClient.getCatalogs.resolves({ - catalogs: [{ CatalogId: 'test-catalog', CatalogType: 'HIVE' }], + catalogs: [regularCatalog], + nextToken: undefined, }) mockGlueClient.getDatabases.resolves({ databases: [{ Name: 'test-db' }], @@ -236,8 +243,13 @@ describe('LakehouseStrategy', function () { ) const children = await node.getChildren() - assert.ok(children.length > 0) + const catalogNode = children.find((child) => child.id === 'regular-catalog') as LakehouseNode + assert.ok(catalogNode) assert.ok(mockGlueCatalogClient.getCatalogs.called) + + const catalogChildren = await catalogNode.getChildren() + assert.strictEqual(catalogChildren.length, 1) + assert.strictEqual((catalogChildren[0] as LakehouseNode).data.nodeType, NodeType.GLUE_DATABASE) }) it('should handle catalog database pagination', async function () { @@ -277,6 +289,64 @@ describe('LakehouseStrategy', function () { assert.strictEqual(children.length, 2) assert.ok(mockGlueClient.getDatabases.calledTwice) }) + + it('should treat RedLake catalogs as parent catalogs', async function () { + const redLakeCatalog = { + CatalogId: 'redlake-catalog', + Name: 'redlake-catalog', + FederatedCatalog: { + ConnectionName: 'aws:redshift', + }, + } + + mockGlueCatalogClient.getCatalogs.resolves({ + catalogs: [redLakeCatalog], + nextToken: undefined, + }) + + const node = createLakehouseConnectionNode( + mockConnection as any, + mockCredentialsProvider as ConnectionCredentialsProvider, + 'us-east-1' + ) + const children = await node.getChildren() + + const catalogNode = children.find((child) => child.id === 'redlake-catalog') as LakehouseNode + assert.ok(catalogNode) + + const catalogChildren = await catalogNode.getChildren() + assert.strictEqual(catalogChildren.length, 1) + assert.ok(catalogChildren[0].resource === '[No data found]') + }) + + it('should treat S3Tables catalogs as parent catalogs', async function () { + const s3TablesCatalog = { + CatalogId: 's3tables-catalog', + Name: 's3tables-catalog', + FederatedCatalog: { + ConnectionName: 'aws:s3tables', + }, + } + + mockGlueCatalogClient.getCatalogs.resolves({ + catalogs: [s3TablesCatalog], + nextToken: undefined, + }) + + const node = createLakehouseConnectionNode( + mockConnection as any, + mockCredentialsProvider as ConnectionCredentialsProvider, + 'us-east-1' + ) + const children = await node.getChildren() + + const catalogNode = children.find((child) => child.id === 's3tables-catalog') as LakehouseNode + assert.ok(catalogNode) + + const catalogChildren = await catalogNode.getChildren() + assert.strictEqual(catalogChildren.length, 1) + assert.ok(catalogChildren[0].resource === '[No data found]') + }) }) describe('Database nodes', function () { diff --git a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts index cd92aa42981..bc2785315f2 100644 --- a/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts +++ b/packages/core/src/test/sagemakerunifiedstudio/explorer/nodes/utils.test.ts @@ -15,6 +15,8 @@ import { isRedLakeDatabase, getTooltip, getRedshiftTypeFromHost, + isRedLakeCatalog, + isS3TablesCatalog, } from '../../../../sagemakerunifiedstudio/explorer/nodes/utils' import { NodeType, ConnectionType, RedshiftType } from '../../../../sagemakerunifiedstudio/explorer/nodes/types' @@ -249,4 +251,63 @@ describe('utils', function () { assert.strictEqual(getRedshiftTypeFromHost(unknownHost), undefined) }) }) + + describe('isRedLakeCatalog', function () { + it('should return true for RedLake catalogs with FederatedCatalog connection', function () { + const catalog = { + FederatedCatalog: { + ConnectionName: 'aws:redshift', + }, + } + assert.strictEqual(isRedLakeCatalog(catalog), true) + }) + + it('should return true for RedLake catalogs with CatalogProperties', function () { + const catalog = { + CatalogProperties: { + DataLakeAccessProperties: { + CatalogType: 'aws:redshift', + }, + }, + } + assert.strictEqual(isRedLakeCatalog(catalog), true) + }) + + it('should return false for non-RedLake catalogs', function () { + const catalog = { + FederatedCatalog: { + ConnectionName: 'aws:s3tables', + }, + } + assert.strictEqual(isRedLakeCatalog(catalog), false) + }) + + it('should return false for undefined catalog', function () { + assert.strictEqual(isRedLakeCatalog(undefined), false) + }) + }) + + describe('isS3TablesCatalog', function () { + it('should return true for S3 Tables catalogs', function () { + const catalog = { + FederatedCatalog: { + ConnectionName: 'aws:s3tables', + }, + } + assert.strictEqual(isS3TablesCatalog(catalog), true) + }) + + it('should return false for non-S3 Tables catalogs', function () { + const catalog = { + FederatedCatalog: { + ConnectionName: 'aws:redshift', + }, + } + assert.strictEqual(isS3TablesCatalog(catalog), false) + }) + + it('should return false for undefined catalog', function () { + assert.strictEqual(isS3TablesCatalog(undefined), false) + }) + }) }) diff --git a/packages/toolkit/.changes/next-release/Bug Fix-490ee89b-01ae-4e60-a2fe-ec26e7dae6fe.json b/packages/toolkit/.changes/next-release/Bug Fix-490ee89b-01ae-4e60-a2fe-ec26e7dae6fe.json new file mode 100644 index 00000000000..774702af724 --- /dev/null +++ b/packages/toolkit/.changes/next-release/Bug Fix-490ee89b-01ae-4e60-a2fe-ec26e7dae6fe.json @@ -0,0 +1,4 @@ +{ + "type": "Bug Fix", + "description": "SageMaker Unified Studio: Fixed s3 table catalog node showing error when it's empty" +}