1111import com .fasterxml .jackson .databind .node .ObjectNode ;
1212import com .google .api .gax .rpc .HeaderProvider ;
1313import com .google .cloud .bigquery .BigQuery ;
14+ import com .google .cloud .bigquery .BigQueryError ;
1415import com .google .cloud .bigquery .BigQueryException ;
1516import com .google .cloud .bigquery .Clustering ;
1617import com .google .cloud .bigquery .Dataset ;
1718import com .google .cloud .bigquery .DatasetInfo ;
1819import com .google .cloud .bigquery .Field ;
1920import com .google .cloud .bigquery .FieldList ;
21+ import com .google .cloud .bigquery .InsertAllRequest ;
22+ import com .google .cloud .bigquery .InsertAllResponse ;
2023import com .google .cloud .bigquery .Job ;
2124import com .google .cloud .bigquery .JobId ;
2225import com .google .cloud .bigquery .JobInfo ;
2528import com .google .cloud .bigquery .Schema ;
2629import com .google .cloud .bigquery .StandardSQLTypeName ;
2730import com .google .cloud .bigquery .StandardTableDefinition ;
31+ import com .google .cloud .bigquery .Table ;
2832import com .google .cloud .bigquery .TableDefinition ;
2933import com .google .cloud .bigquery .TableId ;
3034import com .google .cloud .bigquery .TableInfo ;
3135import com .google .cloud .bigquery .TimePartitioning ;
3236import com .google .common .collect .ImmutableList ;
3337import com .google .common .collect .ImmutableMap ;
38+ import io .airbyte .commons .exceptions .ConfigErrorException ;
3439import io .airbyte .commons .json .Jsons ;
3540import io .airbyte .config .WorkerEnvConstants ;
3641import io .airbyte .integrations .base .JavaBaseConstants ;
4550import java .time .format .DateTimeParseException ;
4651import java .util .ArrayList ;
4752import java .util .List ;
53+ import java .util .Map ;
4854import java .util .Optional ;
4955import java .util .Set ;
5056import java .util .UUID ;
@@ -64,7 +70,8 @@ public class BigQueryUtils {
6470 DateTimeFormatter .ofPattern ("[yyyy][yy]['-']['/']['.'][' '][MMM][MM][M]['-']['/']['.'][' '][dd][d]" +
6571 "[[' ']['T']HH:mm[':'ss[.][SSSSSS][SSSSS][SSSS][SSS][' '][z][zzz][Z][O][x][XXX][XX][X]]]" );
6672 private static final String USER_AGENT_FORMAT = "%s (GPN: Airbyte)" ;
67- private static final String CHECK_TEST_DATASET_SUFFIX = "_airbyte_check_stage_tmp" ;
73+ private static final String CHECK_TEST_DATASET_SUFFIX = "_airbyte_check_stage_tmp_" + System .currentTimeMillis ();
74+ private static final String CHECK_TEST_TMP_TABLE_NAME = "test_connection_table_name" ;
6875
6976 public static ImmutablePair <Job , String > executeQuery (final BigQuery bigquery , final QueryJobConfiguration queryConfig ) {
7077 final JobId jobId = JobId .of (UUID .randomUUID ().toString ());
@@ -119,16 +126,67 @@ public static Dataset getOrCreateDataset(final BigQuery bigquery, final String d
119126
120127 public static void checkHasCreateAndDeleteDatasetRole (final BigQuery bigquery , final String datasetId , final String datasetLocation ) {
121128 final String tmpTestDatasetId = datasetId + CHECK_TEST_DATASET_SUFFIX ;
122- final Dataset dataset = bigquery .getDataset (tmpTestDatasetId );
129+ final DatasetInfo datasetInfo = DatasetInfo .newBuilder (tmpTestDatasetId ).setLocation (datasetLocation ).build ();
130+
131+ bigquery .create (datasetInfo );
123132
124- // remove possible tmp datasets from previous execution
125- if (dataset != null && dataset .exists ()) {
133+ try {
134+ attemptCreateTableAndTestInsert (bigquery , tmpTestDatasetId );
135+ } finally {
126136 bigquery .delete (tmpTestDatasetId );
127137 }
138+ }
128139
129- final DatasetInfo datasetInfo = DatasetInfo .newBuilder (tmpTestDatasetId ).setLocation (datasetLocation ).build ();
130- bigquery .create (datasetInfo );
131- bigquery .delete (tmpTestDatasetId );
140+ /**
141+ * Method is used to create tmp table and make dummy record insert. It's used in Check() connection
142+ * method to make sure that user has all required roles for upcoming data sync/migration. It also
143+ * verifies if BigQuery project is billable, if not - later sync will fail as non-billable project
144+ * has limitations with stream uploading and DML queries. More details may be found there:
145+ * https://cloud.google.com/bigquery/docs/streaming-data-into-bigquery
146+ * https://cloud.google.com/bigquery/docs/reference/standard-sql/data-manipulation-language
147+ *
148+ * @param bigquery - initialized bigquery client
149+ * @param tmpTestDatasetId - dataset name where tmp table will be created
150+ */
151+ private static void attemptCreateTableAndTestInsert (final BigQuery bigquery , final String tmpTestDatasetId ) {
152+ // Create dummy schema that will be used for tmp table creation
153+ final Schema testTableSchema = Schema .of (
154+ Field .of ("id" , StandardSQLTypeName .INT64 ),
155+ Field .of ("name" , StandardSQLTypeName .STRING ));
156+
157+ // Create tmp table to verify if user has a create table permission. Also below we will do test
158+ // records insert in it
159+ final Table test_connection_table_name = createTable (bigquery , tmpTestDatasetId ,
160+ CHECK_TEST_TMP_TABLE_NAME , testTableSchema );
161+
162+ // Try to make test (dummy records) insert to make sure that user has required permissions
163+ try {
164+ final InsertAllResponse response =
165+ bigquery .insertAll (InsertAllRequest
166+ .newBuilder (test_connection_table_name )
167+ .addRow (Map .of ("id" , 1 , "name" , "James" ))
168+ .addRow (Map .of ("id" , 2 , "name" , "Eugene" ))
169+ .addRow (Map .of ("id" , 3 , "name" , "Angelina" ))
170+ .build ());
171+
172+ if (response .hasErrors ()) {
173+ // If any of the insertions failed, this lets you inspect the errors
174+ for (Map .Entry <Long , List <BigQueryError >> entry : response .getInsertErrors ().entrySet ()) {
175+ throw new ConfigErrorException ("Failed to check connection: \n " + entry .getValue ());
176+ }
177+ }
178+ } catch (final BigQueryException e ) {
179+ throw new ConfigErrorException ("Failed to check connection: \n " + e .getMessage ());
180+ } finally {
181+ test_connection_table_name .delete ();
182+ }
183+ }
184+
185+ public static Table createTable (final BigQuery bigquery , String datasetName , String tableName , Schema schema ) {
186+ final TableId tableId = TableId .of (datasetName , tableName );
187+ final TableDefinition tableDefinition = StandardTableDefinition .of (schema );
188+ final TableInfo tableInfo = TableInfo .newBuilder (tableId , tableDefinition ).build ();
189+ return bigquery .create (tableInfo );
132190 }
133191
134192 // https://cloud.google.com/bigquery/docs/creating-partitioned-tables#java
0 commit comments