plumb pipeline options

Author: jrmccluskey

sdks/python/apache_beam/io/gcp/bigquery.py

@@ -1398,7 +1398,8 @@ def __init__(
       with_batched_input=False,
       ignore_unknown_columns=False,
       max_retries=MAX_INSERT_RETRIES,
-      max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE):
+      max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE,
+      options=None):
     """Initialize a WriteToBigQuery transform.
 
     Args:
@@ -1489,6 +1490,7 @@ def __init__(
     self.ignore_unknown_columns = ignore_unknown_columns
     self._max_retries = max_retries
     self._max_insert_payload_size = max_insert_payload_size
+    self.options = options
 
   def display_data(self):
     return {
@@ -1524,6 +1526,10 @@ def get_table_schema(schema):
       return bigquery_tools.parse_table_schema_from_json(schema)
     elif isinstance(schema, dict):
       return bigquery_tools.parse_table_schema_from_json(json.dumps(schema))
+    elif isinstance(schema, (tuple, list)):
+      return tuple(schema)
+    elif hasattr(schema, 'fields'):
+      return BigQueryWriteFn.get_table_schema(schema.fields)
     else:
       raise TypeError('Unexpected schema argument: %s.' % schema)
 
@@ -1532,7 +1538,7 @@ def start_bundle(self):
 
     if not self.bigquery_wrapper:
       self.bigquery_wrapper = bigquery_tools.BigQueryWrapper(
-          client=self.test_client)
+          client=self.test_client, pipeline_options=self.options)
 
     (
         bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER.
@@ -1876,7 +1882,8 @@ def expand(self, input):
         ignore_unknown_columns=self.ignore_unknown_columns,
         with_batched_input=self.with_auto_sharding,
         max_retries=self._max_retries,
-        max_insert_payload_size=self._max_insert_payload_size)
+        max_insert_payload_size=self._max_insert_payload_size,
+        options=input.pipeline.options)
 
     def _add_random_shard(element):
       key = element[0]

sdks/python/apache_beam/io/gcp/bigquery_tools.py

@@ -369,8 +369,14 @@ class BigQueryWrapper(object):
 
   HISTOGRAM_METRIC_LOGGER = MetricLogger()
 
-  def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
-    self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions())
+  def __init__(
+      self,
+      client=None,
+      temp_dataset_id=None,
+      temp_table_ref=None,
+      pipeline_options=None):
+    self.client = client or BigQueryWrapper._bigquery_client(
+        pipeline_options or PipelineOptions())
     self.gcp_bq_client = self.client
 
     self._unique_row_id = 0
@@ -1362,7 +1368,8 @@ def convert_row_to_dict(self, row, schema):
   @staticmethod
   def from_pipeline_options(pipeline_options: PipelineOptions):
     return BigQueryWrapper(
-        client=BigQueryWrapper._bigquery_client(pipeline_options))
+        client=BigQueryWrapper._bigquery_client(pipeline_options),
+        pipeline_options=pipeline_options)
 
   @staticmethod
   def _bigquery_client(pipeline_options: PipelineOptions):
@@ -1650,6 +1657,8 @@ def table_schema_to_dict(table_schema):
   def get_table_field(field):
     """Create a dictionary representation of a table field
     """
+    if isinstance(field, dict):
+      return field
     result = {}
     result['name'] = field.name
     result['type'] = getattr(field, 'field_type', getattr(field, 'type', None))
@@ -1705,14 +1714,16 @@ def get_bq_tableschema(schema):
   Returns:
     Sequence[``google.cloud.bigquery.schema.SchemaField``]: The schema as a TableSchema object.
   """
-  if (isinstance(schema, (tuple, value_provider.ValueProvider)) or
+  if (isinstance(schema, (tuple, list, value_provider.ValueProvider)) or
       callable(schema) or schema is None):
-    return schema
+    return tuple(schema) if isinstance(schema, list) else schema
   elif isinstance(schema, str):
     return get_table_schema_from_string(schema)
   elif isinstance(schema, dict):
     schema_string = json.dumps(schema)
     return parse_table_schema_from_json(schema_string)
+  elif hasattr(schema, 'fields'):
+    return get_bq_tableschema(schema.fields)
   else:
     raise TypeError('Unexpected schema argument: %s.' % schema)