18
18
19
19
package org .apache .flink .connector .redshift .internal .executor ;
20
20
21
- import org .apache .flink .configuration .Configuration ;
22
- import org .apache .flink .connector .redshift .config .RedshiftSinkConfigConstants ;
23
21
import org .apache .flink .connector .redshift .converter .RedshiftRowConverter ;
24
22
import org .apache .flink .connector .redshift .internal .connection .RedshiftConnectionProvider ;
25
23
import org .apache .flink .connector .redshift .internal .statement .RedshiftStatement ;
26
24
import org .apache .flink .connector .redshift .mode .copy .RedshiftCopyModeRowConverterImpl ;
25
+ import org .apache .flink .connector .redshift .options .InternalRedshiftConfigOptions ;
27
26
import org .apache .flink .connector .redshift .util .S3Util ;
28
27
import org .apache .flink .table .data .RowData ;
29
28
import org .apache .flink .table .types .logical .LogicalType ;
30
- import org .apache .flink .util .FlinkRuntimeException ;
31
29
32
- import com .amazon .redshift .core .BaseConnection ;
33
30
import com .amazon .redshift .jdbc .RedshiftConnectionImpl ;
34
31
import com .amazon .redshift .jdbc .RedshiftPreparedStatement ;
35
32
import org .slf4j .Logger ;
43
40
import java .util .function .Function ;
44
41
45
42
/** Executor for Redshift Upload Upsert Operation. */
46
- public class RedshiftUploadUpsertExecutor implements RedshiftExecutor {
43
+ public class RedshiftCopyModeUpsertExecutor implements RedshiftExecutor {
47
44
48
45
private static final long serialVersionUID = 1L ;
49
46
50
- private static final Logger LOG = LoggerFactory .getLogger (RedshiftUploadUpsertExecutor .class );
47
+ private static final Logger LOG = LoggerFactory .getLogger (RedshiftCopyModeUpsertExecutor .class );
51
48
52
49
private final int maxRetries ;
53
50
@@ -87,43 +84,43 @@ public class RedshiftUploadUpsertExecutor implements RedshiftExecutor {
87
84
88
85
private final transient S3Client s3Client ;
89
86
90
- private RedshiftConnectionProvider connectionProvider ;
91
-
92
- public RedshiftUploadUpsertExecutor (
87
+ public RedshiftCopyModeUpsertExecutor (
93
88
String [] fieldNames ,
94
89
String [] keyFields ,
95
90
LogicalType [] fieldTypes ,
96
91
RedshiftRowConverter deleteConverter ,
97
92
Function <RowData , RowData > deleteExtractor ,
98
- Configuration options ) {
93
+ InternalRedshiftConfigOptions options ) {
99
94
100
- this .maxRetries = options .getInteger ( RedshiftSinkConfigConstants . MAX_RETIRES );
95
+ this .maxRetries = options .getMaxRetries ( );
101
96
this .fieldNames = fieldNames ;
102
97
this .keyFields = keyFields ;
103
98
this .deleteConverter = deleteConverter ;
104
99
this .deleteExtractor = deleteExtractor ;
105
100
this .csvInsertData = new ArrayList <>();
106
101
this .csvUpdateData = new ArrayList <>();
107
102
108
- this .tableName = options .getString ( RedshiftSinkConfigConstants . TABLE_NAME );
109
- this .iamRoleArn = options .getString ( RedshiftSinkConfigConstants . IAM_ROLE_ARN );
103
+ this .tableName = options .getTableName ( );
104
+ this .iamRoleArn = options .getIamRoleArn ( );
110
105
this .s3Client = S3Client .create ();
111
106
this .copyRowConverter = new RedshiftCopyModeRowConverterImpl (fieldTypes );
112
107
this .stageTableName = "_" + tableName + "_stage" ;
113
- this .tempS3Uri =
114
- S3Util .getS3UriWithFileName (options .getString (RedshiftSinkConfigConstants .S3_URI ));
108
+ this .tempS3Uri = S3Util .getS3UriWithFileName (options .getS3Uri ());
115
109
}
116
110
117
111
@ Override
118
- public void prepareStatement (BaseConnection connection ) throws SQLException {
112
+ public void prepareStatements (RedshiftConnectionProvider connectionProvider )
113
+ throws SQLException {
119
114
final String createTableSql =
120
115
RedshiftStatement .getCreateTempTableAsStatement (tableName , stageTableName );
121
116
final String insertSql =
122
- RedshiftStatement .getInsertFromStageTable (tableName , stageTableName , fieldNames );
117
+ RedshiftStatement .getInsertFromStageTableStatement (
118
+ tableName , stageTableName , fieldNames );
123
119
deleteSql = RedshiftStatement .getDeleteStatement (tableName , keyFields );
124
120
final String deleteFromStageSql =
125
- RedshiftStatement .getDeleteFromStageTable (tableName , stageTableName , keyFields );
126
- final String truncateSql = RedshiftStatement .getTruncateTable (stageTableName );
121
+ RedshiftStatement .getDeleteFromStageTableStatement (
122
+ tableName , stageTableName , keyFields );
123
+ final String truncateSql = RedshiftStatement .getTruncateTableStatement (stageTableName );
127
124
copyInsertSql =
128
125
RedshiftStatement .getTableCopyStatement (
129
126
tableName , tempS3Uri , fieldNames , iamRoleArn );
@@ -141,27 +138,16 @@ public void prepareStatement(BaseConnection connection) throws SQLException {
141
138
deleteFromStageSql ,
142
139
insertSql ,
143
140
"END;" );
144
- if (connection instanceof RedshiftConnectionImpl ) {
145
- RedshiftConnectionImpl redshiftConnection = (RedshiftConnectionImpl ) connection ;
146
- insertStatement =
147
- (RedshiftPreparedStatement ) redshiftConnection .prepareStatement (copyInsertSql );
148
- updateTrxStatement =
149
- (RedshiftPreparedStatement )
150
- redshiftConnection .prepareStatement (updateTransactionSql );
151
- deleteStatement =
152
- (RedshiftPreparedStatement ) redshiftConnection .prepareStatement (deleteSql );
153
- }
154
- }
155
141
156
- @ Override
157
- public void prepareStatement ( RedshiftConnectionProvider connectionProvider )
158
- throws SQLException {
159
- this . connectionProvider = connectionProvider ;
160
- try {
161
- prepareStatement ( connectionProvider . getOrEstablishConnection ());
162
- } catch ( ClassNotFoundException e ) {
163
- throw new FlinkRuntimeException ( e );
164
- }
142
+ RedshiftConnectionImpl redshiftConnection =
143
+ ( RedshiftConnectionImpl ) connectionProvider . getConnection ();
144
+ insertStatement =
145
+ ( RedshiftPreparedStatement ) redshiftConnection . prepareStatement ( copyInsertSql ) ;
146
+ updateTrxStatement =
147
+ ( RedshiftPreparedStatement )
148
+ redshiftConnection . prepareStatement ( updateTransactionSql );
149
+ deleteStatement =
150
+ ( RedshiftPreparedStatement ) redshiftConnection . prepareStatement ( deleteSql );
165
151
}
166
152
167
153
@ Override
@@ -230,7 +216,7 @@ public void closeStatement() {
230
216
231
217
@ Override
232
218
public String getName () {
233
- return "redshift-upload -upsert-executor" ;
219
+ return "redshift-copy-mode -upsert-executor" ;
234
220
}
235
221
236
222
@ Override
@@ -247,8 +233,6 @@ public String toString() {
247
233
+ '\''
248
234
+ ", maxRetries="
249
235
+ maxRetries
250
- + ", connectionProvider="
251
- + connectionProvider
252
236
+ '}' ;
253
237
}
254
238
}
0 commit comments