From 510677f9c1a264f0a055a1c08a10afdfb7acd4fa Mon Sep 17 00:00:00 2001 From: Jim Schneidereit Date: Thu, 2 Jan 2025 10:32:56 -0800 Subject: [PATCH] Metadata fixes for system schema and lakehouse identifier (#6) These changes exclude the dataspace from the lakehouse identifier when it is null as well as excluding system schemas when making the get_schemas_query. --- .../jdbc/core/DataCloudDatabaseMetadata.java | 2 +- .../jdbc/core/QueryMetadataUtil.java | 21 +++-- src/main/resources/sql/get_schemas_query.sql | 4 +- .../jdbc/core/QueryMetadataUtilTest.java | 76 +++++++++++++++++++ 4 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 src/test/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtilTest.java diff --git a/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudDatabaseMetadata.java b/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudDatabaseMetadata.java index 4d7cedc..a1153fd 100644 --- a/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudDatabaseMetadata.java +++ b/src/main/java/com/salesforce/datacloud/jdbc/core/DataCloudDatabaseMetadata.java @@ -674,7 +674,7 @@ public ResultSet getSchemas() throws SQLException { @Override public ResultSet getCatalogs() throws SQLException { - return QueryMetadataUtil.createCatalogsResultSet(tokenProcessor, client); + return QueryMetadataUtil.createCatalogsResultSet(tokenProcessor); } @Override diff --git a/src/main/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtil.java b/src/main/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtil.java index 6557d0f..83870fa 100644 --- a/src/main/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtil.java +++ b/src/main/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtil.java @@ -40,6 +40,7 @@ import java.util.Optional; import java.util.TimeZone; import java.util.stream.Collectors; +import lombok.SneakyThrows; import lombok.experimental.UtilityClass; import lombok.extern.slf4j.Slf4j; import lombok.val; @@ -355,13 +356,23 @@ private static List constructTableTypesData() { return data; } - public static ResultSet createCatalogsResultSet(Optional tokenProcessor, OkHttpClient client) - throws SQLException { + @SneakyThrows + static List getLakehouse(Optional tokenProcessor) { + if (!tokenProcessor.isPresent()) { + return ImmutableList.of(); + } + val tenantId = tokenProcessor.get().getDataCloudToken().getTenantId(); - val dataspaceName = tokenProcessor.get().getSettings().getDataspace(); - List data = ImmutableList.of(ImmutableList.of("lakehouse:" + tenantId + ";" + dataspaceName)); + val dataspace = tokenProcessor.get().getSettings().getDataspace(); + val lakehouse = + "lakehouse:" + tenantId + ";" + Optional.ofNullable(dataspace).orElse(""); + + return ImmutableList.of(ImmutableList.of(lakehouse)); + } - QueryDBMetadata queryDbMetadata = QueryDBMetadata.GET_CATALOGS; + public static ResultSet createCatalogsResultSet(Optional tokenProcessor) throws SQLException { + val data = getLakehouse(tokenProcessor); + val queryDbMetadata = QueryDBMetadata.GET_CATALOGS; return getMetadataResultSet(queryDbMetadata, NUM_CATALOG_METADATA_COLUMNS, data); } diff --git a/src/main/resources/sql/get_schemas_query.sql b/src/main/resources/sql/get_schemas_query.sql index 21effa1..84f668d 100644 --- a/src/main/resources/sql/get_schemas_query.sql +++ b/src/main/resources/sql/get_schemas_query.sql @@ -4,4 +4,6 @@ WHERE nspname <> 'pg_toast' AND (nspname !~ '^pg_temp_' OR nspname = (pg_catalog.current_schemas(true))[1]) AND (nspname !~ '^pg_toast_temp_' - OR nspname = replace((pg_catalog.current_schemas(true))[1], 'pg_temp_', 'pg_toast_temp_')) \ No newline at end of file + OR nspname = replace((pg_catalog.current_schemas(true))[1], 'pg_temp_', 'pg_toast_temp_')) + AND (nspname !~ 'tableau*') + AND (nspname !~ 'pg_*') \ No newline at end of file diff --git a/src/test/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtilTest.java b/src/test/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtilTest.java new file mode 100644 index 0000000..8ad2169 --- /dev/null +++ b/src/test/java/com/salesforce/datacloud/jdbc/core/QueryMetadataUtilTest.java @@ -0,0 +1,76 @@ +/* + * Copyright (c) 2024, Salesforce, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.salesforce.datacloud.jdbc.core; + +import static org.assertj.core.api.AssertionsForClassTypes.assertThat; +import static org.mockito.Mockito.when; + +import com.google.common.collect.ImmutableList; +import com.salesforce.datacloud.jdbc.auth.AuthenticationSettings; +import com.salesforce.datacloud.jdbc.auth.DataCloudToken; +import com.salesforce.datacloud.jdbc.auth.TokenProcessor; +import java.util.Optional; +import java.util.UUID; +import lombok.SneakyThrows; +import lombok.val; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +@ExtendWith(MockitoExtension.class) +class QueryMetadataUtilTest { + @Mock + private TokenProcessor tokenProcessor; + + @Mock + private DataCloudToken dataCloudToken; + + @Mock + private AuthenticationSettings authenticationSettings; + + @Test + @SneakyThrows + void excludesDataspaceWhenNull() { + val tenantId = UUID.randomUUID().toString(); + + when(dataCloudToken.getTenantId()).thenReturn(tenantId); + when(authenticationSettings.getDataspace()).thenReturn(null); + when(tokenProcessor.getSettings()).thenReturn(authenticationSettings); + when(tokenProcessor.getDataCloudToken()).thenReturn(dataCloudToken); + + val actual = QueryMetadataUtil.getLakehouse(Optional.of(tokenProcessor)); + val expected = ImmutableList.of(ImmutableList.of("lakehouse:" + tenantId + ";")); + assertThat(actual).isEqualTo(expected); + } + + @Test + @SneakyThrows + void includesDataspaceWhenNotNull() { + val tenantId = UUID.randomUUID().toString(); + val dataspace = UUID.randomUUID().toString(); + + when(dataCloudToken.getTenantId()).thenReturn(tenantId); + when(authenticationSettings.getDataspace()).thenReturn(dataspace); + when(tokenProcessor.getSettings()).thenReturn(authenticationSettings); + when(tokenProcessor.getDataCloudToken()).thenReturn(dataCloudToken); + + val actual = QueryMetadataUtil.getLakehouse(Optional.of(tokenProcessor)); + val expected = ImmutableList.of(ImmutableList.of("lakehouse:" + tenantId + ";" + dataspace)); + + assertThat(actual).isEqualTo(expected); + } +}