1818import static org .apache .beam .it .truthmatchers .PipelineAsserts .assertThatPipeline ;
1919
2020import com .google .cloud .teleport .v2 .spanner .migrations .shard .Shard ;
21+ import com .google .cloud .teleport .v2 .spanner .migrations .source .config .JdbcShardConfig ;
2122import com .google .gson .Gson ;
2223import com .google .gson .JsonArray ;
2324import com .google .gson .JsonObject ;
2930import java .io .InputStream ;
3031import java .io .InputStreamReader ;
3132import java .nio .charset .StandardCharsets ;
33+ import java .util .ArrayList ;
3234import java .util .Collections ;
3335import java .util .List ;
3436import java .util .Map ;
@@ -102,11 +104,10 @@ protected void createAndUploadReverseShardConfigToGcs(
102104 shard .setPassword (cloudSqlResourceManager .getPassword ());
103105 shard .setPort (String .valueOf (cloudSqlResourceManager .getPort ()));
104106 shard .setDbName (cloudSqlResourceManager .getDatabaseName ());
105- JsonObject jsObj = new Gson ().toJsonTree (shard ).getAsJsonObject ();
106- jsObj .remove ("secretManagerUri" ); // remove field secretManagerUri
107- JsonArray ja = new JsonArray ();
108- ja .add (jsObj );
109- String shardFileContents = ja .toString ();
107+ JdbcShardConfig jdbcShardConfig = new JdbcShardConfig ();
108+ jdbcShardConfig .setShardConfigs (Collections .singletonList (shard ));
109+ JsonObject jsObj = new Gson ().toJsonTree (jdbcShardConfig ).getAsJsonObject ();
110+ String shardFileContents = jsObj .toString ();
110111 LOG .info ("Shard file contents: {}" , shardFileContents );
111112 gcsResourceManager .createArtifact ("input/shard.json" , shardFileContents );
112113 }
@@ -116,7 +117,7 @@ protected void createAndUploadReverseShardConfigToGcs(
116117 List <CloudSqlResourceManager > cloudSqlResourceManagers ,
117118 String privateHost ,
118119 List <String > logicalShardIds ) {
119- JsonArray ja = new JsonArray ();
120+ List < Shard > shardConfigs = new ArrayList <> ();
120121 for (int i = 0 ; i < cloudSqlResourceManagers .size (); ++i ) {
121122 Shard shard = new Shard ();
122123 shard .setLogicalShardId (logicalShardIds .get (i ));
@@ -126,11 +127,12 @@ protected void createAndUploadReverseShardConfigToGcs(
126127 shard .setPassword (cloudSqlResourceManager .getPassword ());
127128 shard .setPort (String .valueOf (cloudSqlResourceManager .getPort ()));
128129 shard .setDbName (cloudSqlResourceManager .getDatabaseName ());
129- JsonObject jsObj = new Gson ().toJsonTree (shard ).getAsJsonObject ();
130- jsObj .remove ("secretManagerUri" ); // remove field secretManagerUri
131- ja .add (jsObj );
130+ shardConfigs .add (shard );
132131 }
133- String shardFileContents = ja .toString ();
132+ JdbcShardConfig jdbcShardConfig = new JdbcShardConfig ();
133+ jdbcShardConfig .setShardConfigs (shardConfigs );
134+ JsonObject jsObj = new Gson ().toJsonTree (jdbcShardConfig ).getAsJsonObject ();
135+ String shardFileContents = jsObj .toString ();
134136 LOG .info ("Shard file contents: {}" , shardFileContents );
135137 gcsResourceManager .createArtifact ("input/shard.json" , shardFileContents );
136138 }
0 commit comments