Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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)
)
})
}

Expand Down Expand Up @@ -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}`)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,31 @@ 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
*/
// 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
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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' }],
Expand All @@ -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 () {
Expand Down Expand Up @@ -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 () {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ import {
isRedLakeDatabase,
getTooltip,
getRedshiftTypeFromHost,
isRedLakeCatalog,
isS3TablesCatalog,
} from '../../../../sagemakerunifiedstudio/explorer/nodes/utils'
import { NodeType, ConnectionType, RedshiftType } from '../../../../sagemakerunifiedstudio/explorer/nodes/types'

Expand Down Expand Up @@ -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)
})
})
})
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For my own understanding. Is this where we should add the change log for next release? Should we starting doing this for all our PRs?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Team can changeLog if they want to inform user about the fix, feature in the new release. Just like a small announcement/information for their context. Highly recommend changeLog for user facing changes.

Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
{
"type": "Bug Fix",
"description": "SageMaker Unified Studio: Fixed s3 table catalog node showing error when it's empty"
}
Loading