1919
2020
2121import com .google .cloud .bigquery .BigQuery ;
22+ import com .google .cloud .bigquery .TableId ;
2223import com .google .cloud .bigquery .BigQueryError ;
2324import com .google .cloud .bigquery .BigQueryException ;
2425import com .google .cloud .bigquery .InsertAllRequest ;
3738import java .util .Map ;
3839
3940/**
40- * A {@link BigQueryWriter} capable of updating BigQuery table schemas.
41+ * A {@link BigQueryWriter} capable of updating BigQuery table schemas and creating non-existed tables automatically .
4142 */
4243public class AdaptiveBigQueryWriter extends BigQueryWriter {
4344 private static final Logger logger = LoggerFactory .getLogger (AdaptiveBigQueryWriter .class );
4445
45- // The maximum number of retries we will attempt to write rows after updating a BQ table schema.
46- private static final int AFTER_UPDATE_RETY_LIMIT = 5 ;
46+ // The maximum number of retries we will attempt to write rows after creating a table or updating a BQ table schema.
47+ private static final int RETRY_LIMIT = 5 ;
48+ // Wait for about 30s between each retry since both creating table and updating schema take up to 2~3 minutes to take effect.
49+ private static final int RETRY_WAIT_TIME = 30000 ;
4750
4851 private final BigQuery bigQuery ;
4952 private final SchemaManager schemaManager ;
53+ private final boolean autoUpdateSchemas ;
54+ private final boolean autoCreateTables ;
5055
5156 /**
5257 * @param bigQuery Used to send write requests to BigQuery.
@@ -57,10 +62,14 @@ public class AdaptiveBigQueryWriter extends BigQueryWriter {
5762 public AdaptiveBigQueryWriter (BigQuery bigQuery ,
5863 SchemaManager schemaManager ,
5964 int retry ,
60- long retryWait ) {
65+ long retryWait ,
66+ boolean autoUpdateSchemas ,
67+ boolean autoCreateTables ) {
6168 super (retry , retryWait );
6269 this .bigQuery = bigQuery ;
6370 this .schemaManager = schemaManager ;
71+ this .autoUpdateSchemas = autoUpdateSchemas ;
72+ this .autoCreateTables = autoCreateTables ;
6473 }
6574
6675 private boolean isTableMissingSchema (BigQueryException exception ) {
@@ -69,6 +78,12 @@ private boolean isTableMissingSchema(BigQueryException exception) {
6978 return exception .getReason () != null && exception .getReason ().equalsIgnoreCase ("invalid" );
7079 }
7180
81+ private boolean isTableNotExistedException (BigQueryException exception ) {
82+ // If a table does not exist, it will raise a BigQueryException that the input is notFound
83+ // Referring to Google Cloud Error Codes Doc: https://cloud.google.com/bigquery/docs/error-messages?hl=en
84+ return exception .getCode () == 404 ;
85+ }
86+
7287 /**
7388 * Sends the request to BigQuery, then checks the response to see if any errors have occurred. If
7489 * any have, and all errors can be blamed upon invalid columns in the rows sent, attempts to
@@ -86,21 +101,24 @@ public Map<Long, List<BigQueryError>> performWriteRequest(
86101 try {
87102 request = createInsertAllRequest (tableId , rows );
88103 writeResponse = bigQuery .insertAll (request );
89- // Should only perform one schema update attempt; may have to continue insert attempts due to
90- // BigQuery schema updates taking up to two minutes to take effect
104+ // Should only perform one schema update attempt.
91105 if (writeResponse .hasErrors ()
92- && onlyContainsInvalidSchemaErrors (writeResponse .getInsertErrors ())) {
106+ && onlyContainsInvalidSchemaErrors (writeResponse .getInsertErrors ()) && autoUpdateSchemas ) {
93107 attemptSchemaUpdate (tableId , topic );
94108 }
95109 } catch (BigQueryException exception ) {
96- if (isTableMissingSchema (exception )) {
110+ // Should only perform one table creation attempt.
111+ if (isTableNotExistedException (exception ) && autoCreateTables && bigQuery .getTable (tableId .getBaseTableId ()) == null ) {
112+ attemptTableCreate (tableId .getBaseTableId (), topic );
113+ } else if (isTableMissingSchema (exception ) && autoUpdateSchemas ) {
97114 attemptSchemaUpdate (tableId , topic );
98115 } else {
99116 throw exception ;
100117 }
101118 }
102119
103- // Schema update might be delayed, so multiple insertion attempts may be necessary
120+ // Creating tables or updating table schemas in BigQuery takes up to 2~3 minutes to take affect,
121+ // so multiple insertion attempts may be necessary.
104122 int attemptCount = 0 ;
105123 while (writeResponse == null || writeResponse .hasErrors ()) {
106124 logger .trace ("insertion failed" );
@@ -117,10 +135,15 @@ && onlyContainsInvalidSchemaErrors(writeResponse.getInsertErrors())) {
117135 return writeResponse .getInsertErrors ();
118136 }
119137 attemptCount ++;
120- if (attemptCount >= AFTER_UPDATE_RETY_LIMIT ) {
138+ if (attemptCount >= RETRY_LIMIT ) {
121139 throw new BigQueryConnectException (
122140 "Failed to write rows after BQ schema update within "
123- + AFTER_UPDATE_RETY_LIMIT + " attempts for: " + tableId .getBaseTableId ());
141+ + RETRY_LIMIT + " attempts for: " + tableId .getBaseTableId ());
142+ }
143+ try {
144+ Thread .sleep (RETRY_WAIT_TIME );
145+ } catch (InterruptedException e ) {
146+ // no-op, we want to keep retrying the insert
124147 }
125148 }
126149 logger .debug ("table insertion completed successfully" );
@@ -136,6 +159,16 @@ private void attemptSchemaUpdate(PartitionedTableId tableId, String topic) {
136159 }
137160 }
138161
162+ private void attemptTableCreate (TableId tableId , String topic ) {
163+ try {
164+ schemaManager .createTable (tableId , topic );
165+ logger .info ("Table {} does not exist, auto-created table for topic {}" , tableId , topic );
166+ } catch (BigQueryException exception ) {
167+ throw new BigQueryConnectException (
168+ "Failed to create table " + tableId , exception );
169+ }
170+ }
171+
139172 /*
140173 * Currently, the only way to determine the cause of an insert all failure is by examining the map
141174 * object returned by the insertErrors() method of an insert all response. The only way to
0 commit comments