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..c76f2f73543 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,26 @@ 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) => { + if (!catalog) { + return false + } + 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/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) + }) + }) })