Skip to content

Commit f1e7c3c

Browse files
madhav-dbmadhav-db
authored
Merge pull request #191 from databricks/madhav-db/more-metadata-fixes
[PECO-1685] More fixes for metadata calls
2 parents 5945677 + 140518b commit f1e7c3c

File tree

4 files changed

+58
-17
lines changed

4 files changed

+58
-17
lines changed

‎src/main/java/com/databricks/jdbc/client/impl/helper/MetadataResultConstants.java‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,9 @@ public class MetadataResultConstants {
88
public static final String[] DEFAULT_TABLE_TYPES = {"TABLE", "VIEW", "SYSTEM TABLE"};
99
private static final ResultColumn CATALOG_COLUMN =
1010
new ResultColumn("TABLE_CAT", "catalogName", Types.VARCHAR);
11+
12+
private static final ResultColumn CATALOG_COLUMN_FOR_GET_CATALOGS =
13+
new ResultColumn("TABLE_CAT", "catalog", Types.VARCHAR);
1114
private static final ResultColumn TYPE_CATALOG_COLUMN =
1215
new ResultColumn("TYPE_CAT", "TYPE_CATALOG_COLUMN", Types.VARCHAR);
1316
private static final ResultColumn TYPE_SCHEMA_COLUMN =
@@ -26,6 +29,10 @@ public class MetadataResultConstants {
2629
new ResultColumn("TYPE_NAME", "TYPE_NAME", Types.VARCHAR);
2730
private static final ResultColumn SCHEMA_COLUMN =
2831
new ResultColumn("TABLE_SCHEM", "namespace", Types.VARCHAR);
32+
33+
private static final ResultColumn SCHEMA_COLUMN_FOR_GET_SCHEMA =
34+
new ResultColumn("TABLE_SCHEM", "databaseName", Types.VARCHAR);
35+
2936
private static final ResultColumn TABLE_NAME_COLUMN =
3037
new ResultColumn("TABLE_NAME", "tableName", Types.VARCHAR);
3138
private static final ResultColumn TABLE_TYPE_COLUMN =
@@ -133,8 +140,9 @@ public class MetadataResultConstants {
133140
ORDINAL_POSITION_COLUMN,
134141
IS_AUTO_INCREMENT_COLUMN,
135142
IS_GENERATED_COLUMN);
136-
public static List<ResultColumn> CATALOG_COLUMNS = List.of(CATALOG_COLUMN);
137-
public static List<ResultColumn> SCHEMA_COLUMNS = List.of(SCHEMA_COLUMN, CATALOG_COLUMN);
143+
public static List<ResultColumn> CATALOG_COLUMNS = List.of(CATALOG_COLUMN_FOR_GET_CATALOGS);
144+
public static List<ResultColumn> SCHEMA_COLUMNS =
145+
List.of(SCHEMA_COLUMN_FOR_GET_SCHEMA, CATALOG_COLUMN);
138146
public static List<ResultColumn> TABLE_COLUMNS =
139147
List.of(
140148
CATALOG_COLUMN,

‎src/main/java/com/databricks/jdbc/client/impl/helper/MetadataResultSetBuilder.java‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,9 @@ public static DatabricksResultSet getCatalogsResult(ResultSet resultSet) throws
3030
return buildResultSet(CATALOG_COLUMNS, rows, GET_CATALOGS_STATEMENT_ID);
3131
}
3232

33-
public static DatabricksResultSet getSchemasResult(ResultSet resultSet) throws SQLException {
34-
List<List<Object>> rows = getRows(resultSet, SCHEMA_COLUMNS);
33+
public static DatabricksResultSet getSchemasResult(ResultSet resultSet, String catalog)
34+
throws SQLException {
35+
List<List<Object>> rows = getRowsForSchemas(resultSet, SCHEMA_COLUMNS, catalog);
3536
return buildResultSet(SCHEMA_COLUMNS, rows, METADATA_STATEMENT_ID);
3637
}
3738

@@ -85,6 +86,34 @@ private static List<List<Object>> getRows(ResultSet resultSet, List<ResultColumn
8586
return rows;
8687
}
8788

89+
private static List<List<Object>> getRowsForSchemas(
90+
ResultSet resultSet, List<ResultColumn> columns, String catalog) throws SQLException {
91+
// TODO(PECO-1677): Remove this method once the server side ResultSet metadata contains catalogs
92+
List<List<Object>> rows = new ArrayList<>();
93+
while (resultSet.next()) {
94+
List<Object> row = new ArrayList<>();
95+
for (ResultColumn column : columns) {
96+
if (column.getColumnName().equals("TABLE_CAT")) {
97+
row.add(catalog);
98+
continue;
99+
}
100+
Object object;
101+
try {
102+
object = resultSet.getObject(column.getResultSetColumnName());
103+
if (object == null) {
104+
object = NULL_STRING;
105+
}
106+
} catch (DatabricksSQLException e) {
107+
// Remove non-relevant columns from the obtained result set
108+
object = NULL_STRING;
109+
}
110+
row.add(object);
111+
}
112+
rows.add(row);
113+
}
114+
return rows;
115+
}
116+
88117
private static DatabricksResultSet buildResultSet(
89118
List<ResultColumn> columns, List<List<Object>> rows, String statementId) {
90119
return new DatabricksResultSet(

‎src/main/java/com/databricks/jdbc/client/impl/sdk/DatabricksNewMetadataSdkClient.java‎

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,8 @@ public DatabricksResultSet listCatalogs(IDatabricksSession session) throws SQLEx
4242
CommandBuilder commandBuilder = new CommandBuilder(session);
4343
String SQL = commandBuilder.getSQLString(CommandName.LIST_CATALOGS);
4444
LOGGER.debug("SQL command to fetch catalogs: {}", SQL);
45-
return MetadataResultSetBuilder.getCatalogsResult(getResultSet(SQL, session));
45+
return MetadataResultSetBuilder.getCatalogsResult(
46+
getResultSet(SQL, session, StatementType.METADATA));
4647
}
4748

4849
@Override
@@ -52,7 +53,8 @@ public DatabricksResultSet listSchemas(
5253
new CommandBuilder(catalog, session).setSchemaPattern(schemaNamePattern);
5354
String SQL = commandBuilder.getSQLString(CommandName.LIST_SCHEMAS);
5455
LOGGER.debug("SQL command to fetch schemas: {}", SQL);
55-
return MetadataResultSetBuilder.getSchemasResult(getResultSet(SQL, session));
56+
return MetadataResultSetBuilder.getSchemasResult(
57+
getResultSet(SQL, session, StatementType.METADATA), catalog);
5658
}
5759

5860
@Override
@@ -72,7 +74,8 @@ public DatabricksResultSet listTables(
7274
.setSchemaPattern(schemaNamePattern)
7375
.setTablePattern(tableNamePattern);
7476
String SQL = commandBuilder.getSQLString(CommandName.LIST_TABLES);
75-
return MetadataResultSetBuilder.getTablesResult(getResultSet(SQL, session), tableTypes);
77+
return MetadataResultSetBuilder.getTablesResult(
78+
getResultSet(SQL, session, StatementType.METADATA), tableTypes);
7679
}
7780

7881
@Override
@@ -95,7 +98,8 @@ public DatabricksResultSet listColumns(
9598
.setTablePattern(tableNamePattern)
9699
.setColumnPattern(columnNamePattern);
97100
String SQL = commandBuilder.getSQLString(CommandName.LIST_COLUMNS);
98-
return MetadataResultSetBuilder.getColumnsResult(getResultSet(SQL, session));
101+
return MetadataResultSetBuilder.getColumnsResult(
102+
getResultSet(SQL, session, StatementType.QUERY));
99103
}
100104

101105
@Override
@@ -111,7 +115,8 @@ public DatabricksResultSet listFunctions(
111115
.setFunctionPattern(functionNamePattern);
112116
String SQL = commandBuilder.getSQLString(CommandName.LIST_FUNCTIONS);
113117
LOGGER.debug("SQL command to fetch functions: {}", SQL);
114-
return MetadataResultSetBuilder.getFunctionsResult(getResultSet(SQL, session));
118+
return MetadataResultSetBuilder.getFunctionsResult(
119+
getResultSet(SQL, session, StatementType.METADATA));
115120
}
116121

117122
@Override
@@ -121,15 +126,17 @@ public DatabricksResultSet listPrimaryKeys(
121126
new CommandBuilder(catalog, session).setSchema(schema).setTable(table);
122127
String SQL = commandBuilder.getSQLString(CommandName.LIST_PRIMARY_KEYS);
123128
LOGGER.debug("SQL command to fetch primary keys: {}", SQL);
124-
return MetadataResultSetBuilder.getPrimaryKeysResult(getResultSet(SQL, session));
129+
return MetadataResultSetBuilder.getPrimaryKeysResult(
130+
getResultSet(SQL, session, StatementType.METADATA));
125131
}
126132

127-
private ResultSet getResultSet(String SQL, IDatabricksSession session) throws SQLException {
133+
private ResultSet getResultSet(
134+
String SQL, IDatabricksSession session, StatementType statementType) throws SQLException {
128135
return sdkClient.executeStatement(
129136
SQL,
130137
session.getComputeResource(),
131138
new HashMap<Integer, ImmutableSqlParameter>(),
132-
StatementType.METADATA,
139+
statementType,
133140
session,
134141
null /* parentStatement */);
135142
}

‎src/test/java/com/databricks/jdbc/client/impl/DatabricksNewMetadataSdkClientTest.java‎

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,7 @@ void testListColumns(
234234
sqlStatement,
235235
mockedComputeResource,
236236
new HashMap<Integer, ImmutableSqlParameter>(),
237-
StatementType.METADATA,
237+
StatementType.QUERY,
238238
session,
239239
null))
240240
.thenReturn(mockedResultSet);
@@ -267,10 +267,7 @@ void testListSchemas(String sqlStatement, String schema, String description) thr
267267
null))
268268
.thenReturn(mockedResultSet);
269269
when(mockedResultSet.next()).thenReturn(true, false);
270-
for (ResultColumn resultColumn : SCHEMA_COLUMNS) {
271-
when(mockedResultSet.getObject(resultColumn.getResultSetColumnName()))
272-
.thenReturn(TEST_COLUMN);
273-
}
270+
when(mockedResultSet.getObject("databaseName")).thenReturn(TEST_COLUMN);
274271
DatabricksResultSet actualResult = metadataClient.listSchemas(session, TEST_CATALOG, schema);
275272
assertEquals(
276273
actualResult.getStatementStatus().getState(), StatementState.SUCCEEDED, description);

0 commit comments

Comments
 (0)