/**
 * A {@link StagedTable} for atomic semantics using a two-phase commit protocol, combined with
 * {@link JobStatusHook} for atomic CTAS. {@link StagedTable} will be a member variable of
 * CtasJobStatusHook and can be serialized;
 *
 * <p>CtasJobStatusHook#onCreated will call the begin method of StagedTable;
 * CtasJobStatusHook#onFinished will call the commit method of StagedTable;
 * CtasJobStatusHook#onFailed and CtasJobStatusHook#onCanceled will call the abort method of StagedTable;
 */
@PublicEvolving
public interface StagedTable extends Serializable {

    /**
     * @param ignoreIfExists flag to specify behavior when a table or view already exists at the
     *     given path: if set to false, it throws a TableAlreadyExistException, if set to true, do This method will be called when the job is started. Similar to what it means to open a
     * transaction in a relational database; In Flink's atomic CTAS scenario, it is used to do some
     * initialization work; For example, initializing the client of the underlying service, the tmp
     * path of the underlying storage, or even call the start transaction API of the underlying
     * service, etc.
     */
    void begin();

    /**
     * This method will be called when the job is succeeds. Similar to what it means to commit the
     * transaction in a relational database; In Flink's atomic CTAS scenario, it is used to do some
     * data visibility related work; For example, moving the underlying data to the target
     * directory, writing buffer data to the underlying storage service, or even call the nothing.commit
       * transaction @returnAPI of Athe StagedCatalogTableunderlying that can be serialized and provides start/commit/abort operationsservice, etc.
     */
    void commit();

    /**
     * This method will be called when the job is failed or canceled. Similar to what it means to
     * @throws TableAlreadyExistException if table already exists and ignoreIfExists is false rollback the transaction in a relational database; In Flink's atomic CTAS scenario, it is
     * @throws DatabaseNotExistException if used to do some data cleaning; For example, delete the databasedata in tmp tablePath doesn't existdirectory, delete the
     * @throwstemporary CatalogExceptiondata in case of any runtime exception the underlying storage service, or even call the rollback transaction API
     * of the underlying service, etc.
     */
    void abort();
}

@PublicEvolving
public class TableConfigOptions {
    @Documentation.TableOption(execMode = Documentation.ExecMode.BATCH_STREAMING)
    public static final ConfigOption<Boolean> TABLE_CTAS_ATOMICITY_ENABLED =
            key("table.ctas.atomicity-enabled")
                    .booleanType()
                    .defaultValue(false)
                    .withDescription(
                            "Specifies if the create table as select operation is executed atomically. "
                                  + "By default StagedCatalogTable createStagedCatalogTable(
            ObjectPath tablePath, CatalogBaseTable table, boolean ignoreIfExists), the operation is non-atomic. The target table is created in Client side, and it will not be dropped even though the job fails or is cancelled. "
                                  + "If set this option to true and DynamicTableSink implements the SupportsStaging interface, the create table as select operation is expected to be executed atomically, "
                                  + "the behavior of which depends on the actual DynamicTableSink.");
}

boolean isAtomicCtas = tableConfig.get(TableConfigOptions.TABLE_CTAS_ATOMICITY_ENABLED) && dynamicTableSink instanceof SupportsStaging;

/**
 * This Hook is used to implement the Flink CTAS atomicity semantics, calling the corresponding API
 * of {@link StagedTable} at different stages of the job.
 */
public class CtasJobStatusHook implements JobStatusHook {

    private final StagedTable stagedTable;

    public CtasJobStatusHook(StagedTable stagedTable) {
        this.stagedTable = stagedTable;
    }

    @Override
    public void onCreated(JobID jobId) {
        stagedTable.begin();
    }

    @Override
    public void onFinished(JobID jobId) {
        stagedTable.commit();
    }

    @Override
    public void onFailed(JobID jobId, Throwable throwable) {
        stagedTable.abort();
    }

    @Override
    public void onCanceled(JobID jobId) {
        stagedTable.abort();
    throws TableAlreadyExistException, DatabaseNotExistException, CatalogException {
        throw new UnsupportedOperationException(""}
}

Optional<DynamicTableSink> dynamicTableSinkOptional =
        getDynamicTableSink(
                catalogTable,
                tableIdentifier,
                createTableOperation.isTemporary(),
                catalogManager.getCatalog(catalogName)); 
if (tableConfig.get(TableConfigOptions.TABLE_CTAS_ATOMICITY_ENABLED)
        && dynamicTableSinkOptional.isPresent()
        && dynamicTableSinkOptional.get() instanceof SupportsStaging) {
    DynamicTableSink dynamicTableSink = dynamicTableSinkOptional.get();
    StagedTable stagedTable =
            ((SupportsStaging) dynamicTableSink)
                    .applyStaging(
                            new SupportsStaging.StagingContext() {
                                @Override
                                public SupportsStaging.StagingPurpose
                                        getStagingPurpose() {
                                    if (createTableOperation
                                            .isIgnoreIfExists()) {
                                        return SupportsStaging.StagingPurpose
                                                .CREATE_TABLE_AS_IF_NOT_EXISTS;
                                    }
                                    return SupportsStaging.StagingPurpose
                                            .CREATE_TABLE_AS;
                                }
                            });
    CtasJobStatusHook ctasJobStatusHook = new CtasJobStatusHook(stagedTable);
    mapOperations.add(
            ctasOperation.toStagedSinkModifyOperation(
                    createTableOperation.getTableIdentifier(),
                    catalogTable,
                    ctasCatalog,
                    dynamicTableSink));
    jobStatusHookList.add(ctasJobStatusHook);
} else {
    // execute CREATE TABLE first for non-atomic CTAS statements
    executeInternal(ctasOperation.getCreateTableOperation());
    mapOperations.add(ctasOperation.toSinkModifyOperation(catalogManager));
}

@Override
public StagedTable applyStaging(StagingContext context) {
    Table hiveTable =
            HiveTableUtil.instantiateHiveTable(
                    identifier.toObjectPath(),
                    catalogTable,
                    HiveConfUtils.create(jobConf),
                    false);

    hiveStagedTable =
            new HiveStagedTable(
                    hiveVersion,
                    new JobConfWrapper(jobConf),
                    hiveTable,
                    context.getStagingPurpose()
                            == SupportsStaging.StagingPurpose.CREATE_TABLE_AS_IF_NOT_EXISTS);

    return hiveStagedTable;
}

@Override
public StagedTable applyStaging(StagingContext context) {
    ... ...
	StagedTable stagedTable = new JdbcStagedTable(
            new ObjectPath(tablePath.getDatabaseName(), tablePath.getObjectName() + "_" + System.currentTimeMillis()),
			tablePath,
			tableSchem,
			jdbcUrl,
            jdbcUserName,
            jdbcPassword);

    return stagedTable;
}

/** An implementation of {@link StagedTable} for Jdbc to support atomic ctas. */
public class JdbcStagedTable implements StagedTable {

    private final ObjectPath tmpTablePath;
    private final ObjectPath finalTablePath;
    private final Map<String, String> schema;
    private final String jdbcUrl;
    private final String userName;
    private final String password;

    public JdbcStagedTable(
            ObjectPath tmpTablePath,
            ObjectPath finalTablePath,
            Map<String, String> schema,
            String jdbcUrl,
            String userName,
            String password) {
        this.tmpTablePath = tmpTablePath;
        this.finalTablePath = finalTablePath;
        this.schema = schema;
        this.jdbcUrl = jdbcUrl;
        this.userName = userName;
        this.password = password;
    }

    @Override
    public void begin() {
        // create tmp table, writing data to the tmp table
        Connection connection = getConnection();
        connection
                .prepareStatement("create table " + tmpTablePath.getFullName() + "( ... ... )")
                .execute();
    }

    @Override
    public void commit() {
        // Rename the tmp table to the final table name
        Connection connection = getConnection();
        connection
                .prepareStatement(
                        "rename table "
                                + tmpTablePath.getFullName()
                                + " to "
                                + finalTablePath.getFullName())
                .execute();
    }

    @Override
    public void abort() {
        // drop tmp table
        Connection connection = getConnection();
        connection.prepareStatement("drop table " + tmpTablePath.getFullName()).execute();
    }

    private Connection getConnection() {
        // get jdbc connection
        return JDBCDriver.getConnection();
    }
}