gguptp commented on code in PR #206: URL: https://github.com/apache/flink-connector-aws/pull/206#discussion_r4097254361
########## flink-catalog-aws/flink-catalog-aws-glue/src/main/java/org/apache/flink/table/catalog/glue/GlueCatalog.java: ########## @@ -0,0 +1,1818 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.table.catalog.glue; + +import org.apache.flink.annotation.VisibleForTesting; +import org.apache.flink.table.api.Schema; +import org.apache.flink.table.catalog.AbstractCatalog; +import org.apache.flink.table.catalog.CatalogBaseTable; +import org.apache.flink.table.catalog.CatalogDatabase; +import org.apache.flink.table.catalog.CatalogFunction; +import org.apache.flink.table.catalog.CatalogPartition; +import org.apache.flink.table.catalog.CatalogPartitionImpl; +import org.apache.flink.table.catalog.CatalogPartitionSpec; +import org.apache.flink.table.catalog.CatalogTable; +import org.apache.flink.table.catalog.CatalogView; +import org.apache.flink.table.catalog.ObjectPath; +import org.apache.flink.table.catalog.ResolvedCatalogBaseTable; +import org.apache.flink.table.catalog.exceptions.CatalogException; +import org.apache.flink.table.catalog.exceptions.DatabaseAlreadyExistException; +import org.apache.flink.table.catalog.exceptions.DatabaseNotEmptyException; +import org.apache.flink.table.catalog.exceptions.DatabaseNotExistException; +import org.apache.flink.table.catalog.exceptions.FunctionAlreadyExistException; +import org.apache.flink.table.catalog.exceptions.FunctionNotExistException; +import org.apache.flink.table.catalog.exceptions.PartitionAlreadyExistsException; +import org.apache.flink.table.catalog.exceptions.PartitionNotExistException; +import org.apache.flink.table.catalog.exceptions.PartitionSpecInvalidException; +import org.apache.flink.table.catalog.exceptions.TableAlreadyExistException; +import org.apache.flink.table.catalog.exceptions.TableNotExistException; +import org.apache.flink.table.catalog.exceptions.TableNotPartitionedException; +import org.apache.flink.table.catalog.exceptions.TablePartitionedException; +import org.apache.flink.table.catalog.glue.operator.GlueDatabaseOperator; +import org.apache.flink.table.catalog.glue.operator.GlueFunctionOperator; +import org.apache.flink.table.catalog.glue.operator.GluePartitionOperator; +import org.apache.flink.table.catalog.glue.operator.GlueTableOperator; +import org.apache.flink.table.catalog.glue.util.GlueCatalogConstants; +import org.apache.flink.table.catalog.glue.util.GlueTableUtils; +import org.apache.flink.table.catalog.glue.util.GlueTypeConverter; +import org.apache.flink.table.catalog.stats.CatalogColumnStatistics; +import org.apache.flink.table.catalog.stats.CatalogTableStatistics; +import org.apache.flink.table.expressions.Expression; +import org.apache.flink.table.functions.FunctionIdentifier; +import org.apache.flink.util.Preconditions; +import org.apache.flink.util.StringUtils; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.Partition; +import software.amazon.awssdk.services.glue.model.PartitionInput; +import software.amazon.awssdk.services.glue.model.StorageDescriptor; +import software.amazon.awssdk.services.glue.model.Table; +import software.amazon.awssdk.services.glue.model.TableInput; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * GlueCatalog is an implementation of the Flink AbstractCatalog that interacts with AWS Glue. This + * class allows Flink to perform various catalog operations such as creating, deleting, and + * retrieving databases and tables from Glue. It encapsulates AWS Glue's API and provides a + * Flink-compatible interface. + * + * <p>This catalog uses GlueClient to interact with AWS Glue services, and operations related to + * databases and tables are delegated to respective helper classes like GlueDatabaseOperations and + * GlueTableOperations. + */ +public class GlueCatalog extends AbstractCatalog { + + private static final Logger LOG = LoggerFactory.getLogger(GlueCatalog.class); + + private GlueClient glueClient; + private GlueTypeConverter glueTypeConverter; + private GlueDatabaseOperator glueDatabaseOperations; + private GlueTableOperator glueTableOperations; + private GlueFunctionOperator glueFunctionsOperations; + private GluePartitionOperator gluePartitionOperations; + private GlueTableUtils glueTableUtils; + + /** + * Constructs a GlueCatalog with a provided Glue client. + * + * @param name the name of the catalog + * @param defaultDatabase the default database for the catalog + * @param region the AWS region to be used for Glue operations + * @param glueClient Glue Client so we can decide which one to use for testing + */ + @VisibleForTesting + GlueCatalog(String name, String defaultDatabase, String region, GlueClient glueClient) { + super(name, defaultDatabase); + + // Validate region parameter + Preconditions.checkNotNull(region, "region cannot be null"); + Preconditions.checkArgument(!region.trim().isEmpty(), "region cannot be empty"); + + // Initialize GlueClient in the constructor + if (glueClient != null) { + setup(glueClient); + } else { + // If no GlueClient is provided, initialize it using the default region + GlueClient client = GlueClient.builder().region(Region.of(region)).build(); + setup(client); + } + } + + /** + * Constructs a GlueCatalog with default client. + * + * @param name the name of the catalog + * @param defaultDatabase the default database for the catalog + * @param region the AWS region to be used for Glue operations + */ + public GlueCatalog(String name, String defaultDatabase, String region) { + super(name, defaultDatabase); + + // Validate region parameter + Preconditions.checkNotNull(region, "region cannot be null"); + Preconditions.checkArgument(!region.trim().isEmpty(), "region cannot be empty"); + + // Create a synchronized client builder to avoid concurrent modification exceptions + GlueClient client = + GlueClient.builder() + .region(Region.of(region)) + .credentialsProvider( + software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider + .create()) + .build(); + setup(client); + } + + /** + * Private helper method to set up the GlueCatalog with a GlueClient instance. This method + * initializes all the necessary components and operators. + * + * @param glueClient the GlueClient to use for AWS Glue operations + */ + private void setup(GlueClient glueClient) { + this.glueClient = glueClient; + this.glueTypeConverter = new GlueTypeConverter(); + this.glueTableUtils = new GlueTableUtils(glueTypeConverter); + this.glueDatabaseOperations = new GlueDatabaseOperator(glueClient, getName()); + this.glueTableOperations = new GlueTableOperator(glueClient, getName()); + this.glueFunctionsOperations = new GlueFunctionOperator(glueClient, getName()); + this.gluePartitionOperations = new GluePartitionOperator(glueClient, getName()); + } + + /** + * Validates that a database exists, throwing DatabaseNotExistException if it doesn't. + * + * @param databaseName the name of the database to validate + * @throws DatabaseNotExistException if the database does not exist + * @throws CatalogException if an error occurs while checking database existence + */ + private void validateDatabaseExists(String databaseName) + throws DatabaseNotExistException, CatalogException { + Preconditions.checkArgument( + !StringUtils.isNullOrWhitespaceOnly(databaseName), + "databaseName cannot be null or empty"); + + if (!databaseExists(databaseName)) { + throw new DatabaseNotExistException(getName(), databaseName); + } + } + + /** + * Opens the GlueCatalog and initializes necessary resources. + * + * @throws CatalogException if an error occurs during the opening process + */ + @Override + public void open() throws CatalogException { + LOG.info("Opening GlueCatalog with client: {}", glueClient); + } + + /** + * Closes the GlueCatalog and releases resources. + * + * @throws CatalogException if an error occurs during the closing process + */ + @Override + public void close() throws CatalogException { + if (glueClient != null) { + LOG.info("Closing GlueCatalog client"); + // The AWS SDK close() is best-effort and does not surface exceptions, + // so no retry logic is required here. + glueClient.close(); + } + } + + /** + * Lists all the databases available in the Glue catalog. + * + * @return a list of database names + * @throws CatalogException if an error occurs while listing the databases + */ + @Override + public List<String> listDatabases() throws CatalogException { + return glueDatabaseOperations.listDatabases(); + } + + /** + * Retrieves a specific database by its name. + * + * @param databaseName the name of the database to retrieve + * @return the database if found + * @throws DatabaseNotExistException if the database does not exist + * @throws CatalogException if an error occurs while retrieving the database + */ + @Override + public CatalogDatabase getDatabase(String databaseName) + throws DatabaseNotExistException, CatalogException { + Preconditions.checkArgument( + !StringUtils.isNullOrWhitespaceOnly(databaseName), + "databaseName cannot be null or empty"); + + // Use case-insensitive database name resolution + String glueDatabaseName = findGlueDatabaseName(databaseName); + if (glueDatabaseName == null) { + throw new DatabaseNotExistException(getName(), databaseName); + } + + return glueDatabaseOperations.getDatabase(glueDatabaseName); + } + + /** + * Checks if a database exists in Glue. + * + * @param databaseName the name of the database + * @return true if the database exists, false otherwise + * @throws CatalogException if an error occurs while checking the database + */ + @Override + public boolean databaseExists(String databaseName) throws CatalogException { + Preconditions.checkArgument( + !StringUtils.isNullOrWhitespaceOnly(databaseName), + "databaseName cannot be null or empty"); + + // Use case-insensitive database name resolution + return findGlueDatabaseName(databaseName) != null; + } + + /** + * Creates a new database in Glue. + * + * @param databaseName the name of the database to create + * @param catalogDatabase the catalog database object containing database metadata + * @param ifNotExists flag indicating whether to ignore the error if the database already exists + * @throws DatabaseAlreadyExistException if the database already exists and ifNotExists is false + * @throws CatalogException if an error occurs while creating the database + */ + @Override + public void createDatabase( + String databaseName, CatalogDatabase catalogDatabase, boolean ifNotExists) + throws DatabaseAlreadyExistException, CatalogException { + Preconditions.checkArgument( + !StringUtils.isNullOrWhitespaceOnly(databaseName), + "databaseName cannot be null or empty"); + Preconditions.checkNotNull(catalogDatabase, "CatalogDatabase cannot be null"); + + // Check for exact case match first + boolean exactExists = databaseExists(databaseName); + if (exactExists && !ifNotExists) { + throw new DatabaseAlreadyExistException(getName(), databaseName); + } + if (exactExists) { + return; // Database exists with exact case, and IF NOT EXISTS is true + } + + // Check for case-insensitive collision (Glue limitation) + String conflictingDatabase = findCaseInsensitiveConflict(databaseName); + if (conflictingDatabase != null) { + String message = + String.format( + "Cannot create database '%s' because it conflicts with existing database '%s'. " + + "AWS Glue stores database names in lowercase, so '%s' and '%s' would both be stored as '%s'.", + databaseName, + conflictingDatabase, + databaseName, + conflictingDatabase, + databaseName.toLowerCase()); + throw new DatabaseAlreadyExistException( + getName(), databaseName, new CatalogException(message)); + } + + // Safe to create - no exact match and no case conflicts + glueDatabaseOperations.createDatabase(databaseName, catalogDatabase); + } + + /** + * Drops an existing database in Glue. + * + * @param databaseName the name of the database to drop + * @param ignoreIfNotExists flag to ignore the exception if the database doesn't exist + * @param cascade flag indicating whether to cascade the operation to drop related objects + * @throws DatabaseNotExistException if the database does not exist and ignoreIfNotExists is + * false + * @throws DatabaseNotEmptyException if the database contains objects and cascade is false + * @throws CatalogException if an error occurs while dropping the database + */ + @Override + public void dropDatabase(String databaseName, boolean ignoreIfNotExists, boolean cascade) + throws DatabaseNotExistException, DatabaseNotEmptyException, CatalogException { + Preconditions.checkArgument( + !StringUtils.isNullOrWhitespaceOnly(databaseName), + "databaseName cannot be null or empty"); + + if (!databaseExists(databaseName)) { + if (!ignoreIfNotExists) { + throw new DatabaseNotExistException(getName(), databaseName); + } + return; // Database doesn't exist and ignoreIfNotExists is true + } + + // Check if database is empty (contains no tables, views, or functions) + boolean isEmpty = isDatabaseEmpty(databaseName); + + if (!isEmpty && !cascade) { + throw new DatabaseNotEmptyException(getName(), databaseName); + } + + if (!isEmpty && cascade) { + // Drop all objects in the database before dropping the database + dropAllObjectsInDatabase(databaseName); + } + + // Drop the database + glueDatabaseOperations.dropGlueDatabase(databaseName); + } + + /** + * Checks if a database is empty (contains no tables, views, or functions). + * + * @param databaseName the name of the database to check + * @return true if the database is empty, false otherwise + * @throws CatalogException if an error occurs while checking the database contents + */ + private boolean isDatabaseEmpty(String databaseName) throws CatalogException { + try { + // Check for tables + List<String> tables = listTables(databaseName); + if (!tables.isEmpty()) { + return false; + } + + // Check for views + List<String> views = listViews(databaseName); + if (!views.isEmpty()) { + return false; + } + + // Check for functions + List<String> functions = listFunctions(databaseName); + if (!functions.isEmpty()) { + return false; + } + + return true; + } catch (DatabaseNotExistException e) { + // This shouldn't happen since we checked existence earlier, but handle it gracefully + throw new CatalogException("Database " + databaseName + " does not exist", e); + } + } + + /** + * Drops all objects (tables, views, functions) in a database. This is used when cascade=true in + * dropDatabase. + * + * @param databaseName the name of the database + * @throws CatalogException if an error occurs while dropping objects + */ + private void dropAllObjectsInDatabase(String databaseName) throws CatalogException { + try { + // Drop all tables + List<String> tables = listTables(databaseName); + for (String tableName : tables) { + ObjectPath tablePath = new ObjectPath(databaseName, tableName); + dropTable(tablePath, true); // Use ifExists=true to avoid exceptions + } + + // Drop all views (views are also stored as tables in Glue, so they should be handled by + // dropTable above) + // But let's be explicit and handle them separately if needed + List<String> views = listViews(databaseName); + for (String viewName : views) { + ObjectPath viewPath = new ObjectPath(databaseName, viewName); + // Views are handled as tables in Glue, so dropTable should work + dropTable(viewPath, true); + } + + // Drop all functions + List<String> functions = listFunctions(databaseName); + for (String functionName : functions) { + ObjectPath functionPath = new ObjectPath(databaseName, functionName); + dropFunction(functionPath, true); // Use ignoreIfNotExists=true to avoid exceptions + } + + LOG.info("Successfully dropped all objects in database: {}", databaseName); + } catch (DatabaseNotExistException e) { + throw new CatalogException("Database " + databaseName + " does not exist", e); + } catch (TableNotExistException | FunctionNotExistException e) { + // This could happen in concurrent scenarios, but we use ifExists/ignoreIfNotExists + // flags + LOG.warn( + "Object was already deleted while cascading drop for database: {}", + databaseName, + e); + } + } + + /** + * Lists all tables in a specified database. + * + * @param databaseName the name of the database + * @return a list of table names in the database + * @throws DatabaseNotExistException if the database does not exist + * @throws CatalogException if an error occurs while listing the tables + */ + @Override + public List<String> listTables(String databaseName) + throws DatabaseNotExistException, CatalogException { + Preconditions.checkArgument( + !StringUtils.isNullOrWhitespaceOnly(databaseName), + "databaseName cannot be null or empty"); + + validateDatabaseExists(databaseName); + + // Use the proper database name resolution + String glueDatabaseName = findGlueDatabaseName(databaseName); + if (glueDatabaseName == null) { + throw new DatabaseNotExistException(getName(), databaseName); + } + + // Return original table names with case preserved + return glueTableOperations.listTablesWithOriginalNames(glueDatabaseName); + } + + /** + * Retrieves a table from the catalog using its object path. + * + * @param objectPath the object path of the table to retrieve + * @return the corresponding CatalogBaseTable for the specified table + * @throws TableNotExistException if the table does not exist + * @throws CatalogException if an error occurs while retrieving the table + */ + @Override + public CatalogBaseTable getTable(ObjectPath objectPath) + throws TableNotExistException, CatalogException { + String originalDatabaseName = objectPath.getDatabaseName(); + String originalTableName = objectPath.getObjectName(); + + // Convert to Glue storage names - Use proper database resolution + String glueDatabaseName = findGlueDatabaseName(originalDatabaseName); + if (glueDatabaseName == null) { + throw new TableNotExistException(getName(), objectPath); + } + + // Use direct lowercase lookup first (like databases), then fall back to complex search + String glueTableName = findGlueTableName(glueDatabaseName, originalTableName); + if (glueTableName == null) { + throw new TableNotExistException(getName(), objectPath); + } + + // Get the table using Glue storage names + Table glueTable = glueTableOperations.getGlueTable(glueDatabaseName, glueTableName); + return getCatalogBaseTableFromGlueTable(glueTable); + } + + /** + * Checks if a table exists in the Glue catalog. + * + * @param objectPath the object path of the table to check + * @return true if the table exists, false otherwise + * @throws CatalogException if an error occurs while checking the table's existence + */ + @Override + public boolean tableExists(ObjectPath objectPath) throws CatalogException { + String originalDatabaseName = objectPath.getDatabaseName(); + String originalTableName = objectPath.getObjectName(); + + // Convert to Glue storage names - Use proper database resolution + String glueDatabaseName = findGlueDatabaseName(originalDatabaseName); + if (glueDatabaseName == null) { + return false; // Database doesn't exist, so table can't exist + } + + // Use efficient table name resolution + String glueTableName = findGlueTableName(glueDatabaseName, originalTableName); + return glueTableName != null; + } + + /** + * Drops a table from the Glue catalog. + * + * @param objectPath the object path of the table to drop + * @param ifExists flag indicating whether to ignore the exception if the table does not exist + * @throws TableNotExistException if the table does not exist and ifExists is false + * @throws CatalogException if an error occurs while dropping the table + */ + @Override + public void dropTable(ObjectPath objectPath, boolean ifExists) + throws TableNotExistException, CatalogException { + String originalDatabaseName = objectPath.getDatabaseName(); + String originalTableName = objectPath.getObjectName(); + + // Convert to Glue storage names - Use proper database resolution + String glueDatabaseName = findGlueDatabaseName(originalDatabaseName); + if (glueDatabaseName == null) { + if (!ifExists) { + throw new TableNotExistException(getName(), objectPath); + } + return; // Database doesn't exist, so table can't exist + } + + // Use efficient table name resolution + String glueTableName = findGlueTableName(glueDatabaseName, originalTableName); + if (glueTableName == null) { + if (!ifExists) { + throw new TableNotExistException(getName(), objectPath); + } + return; // Table doesn't exist, and IF EXISTS is true + } + + // Drop the table using Glue storage names + glueTableOperations.dropTable(glueDatabaseName, glueTableName); + } + + /** + * Creates a table in the Glue catalog. + * + * @param objectPath the object path of the table to create + * @param catalogBaseTable the table definition containing the schema and properties + * @param ifNotExists flag indicating whether to ignore the exception if the table already + * exists + * @throws NullPointerException if objectPath or catalogBaseTable is null + * @throws TableAlreadyExistException if the table already exists and ifNotExists is false + * @throws DatabaseNotExistException if the database does not exist + * @throws CatalogException if an error occurs while creating the table + */ + @Override + public void createTable( + ObjectPath objectPath, CatalogBaseTable catalogBaseTable, boolean ifNotExists) + throws TableAlreadyExistException, DatabaseNotExistException, CatalogException { + + // Validate that required parameters are not null + Preconditions.checkNotNull(objectPath, "ObjectPath cannot be null"); + Preconditions.checkNotNull(catalogBaseTable, "CatalogBaseTable cannot be null"); + + String originalDatabaseName = objectPath.getDatabaseName(); + String originalTableName = objectPath.getObjectName(); + + // Check if the database exists + validateDatabaseExists(originalDatabaseName); + + // Check for exact case match first + boolean exactExists = tableExists(objectPath); + if (exactExists && !ifNotExists) { + throw new TableAlreadyExistException(getName(), objectPath); + } + if (exactExists) { + return; // Table exists with exact case, and IF NOT EXISTS is true + } + + // Check for case-insensitive collision (Glue limitation) + String conflictingTable = findCaseInsensitiveTableConflict(objectPath); + if (conflictingTable != null) { + String message = + String.format( + "Cannot create table '%s.%s' because it conflicts with existing table '%s.%s'. " + + "AWS Glue stores table names in lowercase, so '%s' and '%s' would both be stored as '%s'.", + originalDatabaseName, + originalTableName, + originalDatabaseName, + conflictingTable, + originalTableName, + conflictingTable, + originalTableName.toLowerCase()); + throw new TableAlreadyExistException( + getName(), objectPath, new CatalogException(message)); + } + + // Get common properties + Map<String, String> tableProperties = new HashMap<>(catalogBaseTable.getOptions()); + + try { + // Process based on table type + if (catalogBaseTable.getTableKind() == CatalogBaseTable.TableKind.TABLE) { + createRegularTable(objectPath, (CatalogTable) catalogBaseTable, tableProperties); + } else if (catalogBaseTable.getTableKind() == CatalogBaseTable.TableKind.VIEW) { + createView(objectPath, (CatalogView) catalogBaseTable, tableProperties); + } else { + throw new CatalogException( + "Unsupported table kind: " + catalogBaseTable.getTableKind()); + } + LOG.info( + "Successfully created {}.{} of kind {}", + originalDatabaseName, + originalTableName, + catalogBaseTable.getTableKind()); + } catch (Exception e) { + throw new CatalogException( + String.format( + "Failed to create table %s.%s: %s", + originalDatabaseName, originalTableName, e.getMessage()), + e); + } + } + + /** + * Lists all views in a specified database. + * + * @param databaseName the name of the database + * @return a list of view names in the database + * @throws DatabaseNotExistException if the database does not exist + * @throws CatalogException if an error occurs while listing the views + */ + @Override + public List<String> listViews(String databaseName) Review Comment: we dont seem to be paginating in listViews. Is that intentional? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
