Create Redshift Connection and Utilites for sink

Author: Samrat002

flink-connector-aws/flink-connector-redshift/src/main/java/org/apache/flink/connector/redshift/internal/connection/RedshiftConnectionProvider.java

@@ -1,2 +1,69 @@
-package org.apache.flink.connector.redshift.internal.connection;public interface JdbcConnectionProvider {
+/*
+ * 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.connector.redshift.internal.connection;
+
+import org.apache.flink.annotation.Internal;
+
+import com.amazon.redshift.core.BaseConnection;
+
+import javax.annotation.Nullable;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+
+/** Redshift connection provider. */
+@Internal
+public interface RedshiftConnectionProvider {
+    /**
+     * Get existing connection.
+     *
+     * @return existing connection
+     */
+    @Nullable
+    BaseConnection getConnection();
+
+    /**
+     * Check whether possible existing connection is valid or not through {@link
+     * Connection#isValid(int)}.
+     *
+     * @return true if existing connection is valid
+     * @throws SQLException sql exception throw from {@link Connection#isValid(int)}
+     */
+    boolean isConnectionValid() throws SQLException;
+
+    /**
+     * Get existing connection or establish a new one if there is none.
+     *
+     * @return existing connection or newly established connection
+     * @throws SQLException sql exception
+     * @throws ClassNotFoundException driver class not found
+     */
+    BaseConnection getOrEstablishConnection() throws SQLException, ClassNotFoundException;
+
+    /** Close possible existing connection. */
+    void closeConnection();
+
+    /**
+     * Close possible existing connection and establish an new one.
+     *
+     * @return newly established connection
+     * @throws SQLException sql exception
+     * @throws ClassNotFoundException driver class not found
+     */
+    BaseConnection reestablishConnection() throws SQLException, ClassNotFoundException;
 }

flink-connector-aws/flink-connector-redshift/src/main/java/org/apache/flink/connector/redshift/internal/connection/RedshiftJdbcConnectionProvider.java

@@ -1,2 +1,141 @@
-package org.apache.flink.connector.redshift.internal.connection;public class RedshiftJdbcConnectionProvider {
+/*
+ * 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.connector.redshift.internal.connection;
+
+import org.apache.flink.annotation.Internal;
+import org.apache.flink.configuration.Configuration;
+import org.apache.flink.connector.redshift.sink.config.RedshiftSinkConfigConstants;
+import org.apache.flink.util.FlinkRuntimeException;
+
+import com.amazon.redshift.core.BaseConnection;
+import com.amazon.redshift.jdbc.RedshiftConnectionImpl;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import javax.annotation.concurrent.NotThreadSafe;
+
+import java.io.Serializable;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+
+/** Redshift connection provider. */
+@Internal
+@NotThreadSafe
+public class RedshiftJdbcConnectionProvider implements RedshiftConnectionProvider, Serializable {
+    private static final long serialVersionUID = 1L;
+
+    private static final Logger LOG = LoggerFactory.getLogger(RedshiftJdbcConnectionProvider.class);
+
+    private static final String REDSHIFT_DRIVER_NAME = "com.amazon.redshift.Driver";
+
+    private static final String REDSHIFT_JDBC_IDENTIFIER = "jdbc:redshift://";
+
+    private transient BaseConnection connection;
+
+    private final Configuration options;
+
+    public RedshiftJdbcConnectionProvider(Configuration config) {
+        this.options = config;
+    }
+
+    @Override
+    public BaseConnection getConnection() {
+        if (connection == null) {
+            try {
+                connection =
+                        createConnection(
+                                options.getString(RedshiftSinkConfigConstants.HOSTNAME),
+                                options.getInteger(RedshiftSinkConfigConstants.PORT),
+                                options.getString(RedshiftSinkConfigConstants.DATABASE_NAME));
+            } catch (SQLException e) {
+                throw new FlinkRuntimeException(e);
+            }
+        }
+        return connection;
+    }
+
+    private BaseConnection createConnection(String hostname, int portNumber, String databaseName)
+            throws SQLException {
+        RedshiftConnectionImpl redshiftConnection;
+        String url = REDSHIFT_JDBC_IDENTIFIER + hostname + ":" + portNumber + "/" + databaseName;
+        LOG.info("connection to {}", url);
+
+        try {
+            Class.forName(REDSHIFT_DRIVER_NAME);
+        } catch (ClassNotFoundException e) {
+            throw new SQLException(e);
+        }
+
+        final String username = options.getString(RedshiftSinkConfigConstants.USERNAME);
+        final String password = options.getString(RedshiftSinkConfigConstants.PASSWORD);
+        if (username != null
+                && !username.matches("")
+                && password != null
+                && !password.matches("")) {
+            redshiftConnection =
+                    (RedshiftConnectionImpl) DriverManager.getConnection(url, username, password);
+        } else {
+            redshiftConnection = (RedshiftConnectionImpl) DriverManager.getConnection(url);
+        }
+
+        return redshiftConnection;
+    }
+
+    @Override
+    public boolean isConnectionValid() throws SQLException {
+        return connection.isValid(
+                (int) options.get(RedshiftSinkConfigConstants.TIMEOUT).getSeconds());
+    }
+
+    @Override
+    public BaseConnection getOrEstablishConnection() throws SQLException, ClassNotFoundException {
+        if (!connection.isClosed()) {
+            return connection;
+        }
+        connection =
+                createConnection(
+                        options.getString(RedshiftSinkConfigConstants.HOSTNAME),
+                        options.getInteger(RedshiftSinkConfigConstants.PORT),
+                        options.getString(RedshiftSinkConfigConstants.DATABASE_NAME));
+        return connection;
+    }
+
+    @Override
+    public void closeConnection() {
+        if (connection != null) {
+            try {
+                connection.close();
+            } catch (SQLException e) {
+                LOG.warn("Redshift connection close failed.", e);
+            } finally {
+                connection = null;
+            }
+        }
+    }
+
+    @Override
+    public BaseConnection reestablishConnection() throws SQLException, ClassNotFoundException {
+        closeConnection();
+        connection =
+                createConnection(
+                        options.getString(RedshiftSinkConfigConstants.HOSTNAME),
+                        options.getInteger(RedshiftSinkConfigConstants.PORT),
+                        options.getString(RedshiftSinkConfigConstants.DATABASE_NAME));
+        return connection;
+    }
 }

flink-connector-aws/flink-connector-redshift/src/main/java/org/apache/flink/connector/redshift/mode/jdbc/RedshiftJDBCModeRowConverterImpl.java

@@ -1,5 +1,26 @@
-package org.apache.flink.connector.redshift.jdbc.converter;
+/*
+ * 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.connector.redshift.mode.jdbc;
+
+import org.apache.flink.connector.redshift.converter.AbstractRedshiftRowConverter;
+import org.apache.flink.connector.redshift.converter.DeserializationConverter;
+import org.apache.flink.connector.redshift.converter.SerializationConverter;
 import org.apache.flink.table.data.GenericRowData;
 import org.apache.flink.table.data.RowData;
 import org.apache.flink.table.types.logical.LogicalType;
@@ -12,7 +33,7 @@
 import java.sql.SQLException;
 
 /** Redshift Converter Implementation. */
-public class RedshiftRowConverterImpl extends AbstractRedshiftRowConverter {
+public class RedshiftJDBCModeRowConverterImpl extends AbstractRedshiftRowConverter {
     private static final long serialVersionUID = 1L;
 
     private final RowType rowType;
@@ -21,7 +42,7 @@ public class RedshiftRowConverterImpl extends AbstractRedshiftRowConverter {
 
     private final SerializationConverter[] toExternalConverters;
 
-    public RedshiftRowConverterImpl(RowType rowType) {
+    public RedshiftJDBCModeRowConverterImpl(RowType rowType) {
         this.rowType = Preconditions.checkNotNull(rowType);
         LogicalType[] logicalTypes =
                 rowType.getFields().stream()
@@ -49,9 +70,11 @@ public RowData toInternal(ResultSet resultSet) throws SQLException {
         return genericRowData;
     }
 
+    RedshiftPreparedStatement insertStatement;
+
     @Override
-    public RedshiftPreparedStatement toExternal(
-            RowData rowData, RedshiftPreparedStatement insertStatement) throws SQLException {
+    public Object toExternal(RowData rowData) throws SQLException {
+
         for (int index = 0; index < rowData.getArity(); index++) {
             if (!rowData.isNullAt(index)) {
                 toExternalConverters[index].serialize(rowData, index, insertStatement);

flink-connector-aws/flink-connector-redshift/src/main/java/org/apache/flink/connector/redshift/sink/RedshiftSinkConfigConstants.java

@@ -0,0 +1,114 @@
+/*
+ * 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.connector.redshift.sink.config;
+
+import org.apache.flink.configuration.ConfigOption;
+import org.apache.flink.configuration.ConfigOptions;
+
+import java.time.Duration;
+
+import static java.time.temporal.ChronoUnit.SECONDS;
+
+/** Constants to be used with the RedshiftSink. */
+public class RedshiftSinkConfigConstants {
+
+    /** Enum to Define different Modes of Sink supported in Redshift. */
+    public enum SinkMode {
+        JDBC,
+        COPY
+    }
+
+    public static final ConfigOption<String> HOSTNAME =
+            ConfigOptions.key("hostname")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Hostname of Redshift Cluster");
+
+    public static final ConfigOption<Integer> PORT =
+            ConfigOptions.key("port")
+                    .intType()
+                    .defaultValue(5439)
+                    .withDescription("Port of the Redshift Cluster");
+
+    public static final ConfigOption<String> USERNAME =
+            ConfigOptions.key("username")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Username of Redshift Cluster");
+
+    public static final ConfigOption<String> PASSWORD =
+            ConfigOptions.key("password")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Password of Redshift Cluster");
+
+    public static final ConfigOption<String> DATABASE_NAME =
+            ConfigOptions.key("database-name")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Database to which it needs to be connected");
+
+    public static final ConfigOption<String> TABLE_NAME =
+            ConfigOptions.key("table-name")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Table to which it needs to be connected");
+
+    public static final ConfigOption<Integer> BATCH_SIZE =
+            ConfigOptions.key("sink.batch-size")
+                    .intType()
+                    .defaultValue(1000)
+                    .withDescription("Port of the Redshift Cluster");
+
+    public static final ConfigOption<SinkMode> SINK_MODE =
+            ConfigOptions.key("sink.write.mode")
+                    .enumType(SinkMode.class)
+                    .defaultValue(SinkMode.JDBC)
+                    .withDescription("Mode of sink");
+
+    public static final ConfigOption<String> IAM_ROLE_ARN =
+            ConfigOptions.key("aws.iam-role")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("IAM ROLE");
+
+    public static final ConfigOption<String> S3_URI =
+            ConfigOptions.key("sink.write.temp.s3-uri")
+                    .stringType()
+                    .noDefaultValue()
+                    .withDescription("Temporary S3 URI to store data");
+
+    public static final ConfigOption<Integer> MAX_RETIRES =
+            ConfigOptions.key("sink.max.retries")
+                    .intType()
+                    .defaultValue(2)
+                    .withDescription("Maximum number of Retries in case of Failure");
+
+    public static final ConfigOption<Duration> FLUSH_INTERVAL =
+            ConfigOptions.key("sink.flush.interval")
+                    .durationType()
+                    .defaultValue(Duration.of(10, SECONDS))
+                    .withDescription("Maximum number of Retries in case of Failure");
+
+    public static final ConfigOption<Duration> TIMEOUT =
+            ConfigOptions.key("sink.connection.timeout")
+                    .durationType()
+                    .defaultValue(Duration.ofMinutes(5))
+                    .withDescription("Duration of Connection Timeout (in minutes)");
+}