more fixes

Author: jrmccluskey

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

@@ -423,7 +423,7 @@ def chain_after(result):
   from google.cloud import bigquery as gcp_bigquery
   DatasetReference = gcp_bigquery.DatasetReference
   TableReference = gcp_bigquery.TableReference
-  JobReference = gcp_bigquery.JobReference
+  SchemaField = gcp_bigquery.SchemaField
 except ImportError:
 
   class DatasetReference(object):
@@ -432,9 +432,11 @@ class DatasetReference(object):
   class TableReference(object):
     pass
 
-  class JobReference(object):
+  class SchemaField(object):
     pass
 
+JobReference = bigquery_tools.BeamJobReference
+
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -2235,7 +2237,8 @@ def expand(self, pcoll):
         self.table_reference.project == bigquery_tools.FALLBACK_PROJECT):
       self.table_reference = TableReference(
           DatasetReference(
-              pcoll.pipeline.options.view_as(GoogleCloudOptions).project,
+              (pcoll.pipeline.options.view_as(GoogleCloudOptions).project or
+               bigquery_tools.FALLBACK_PROJECT),
               self.table_reference.dataset_id),
           self.table_reference.table_id)
 
@@ -2766,7 +2769,7 @@ def __init__(self, schema):
       self._value = schema
 
     def __enter__(self):
-      if not isinstance(self._value, (gcp_bigquery.SchemaField, )):
+      if not isinstance(self._value, (SchemaField, )):
         return bigquery_tools.get_bq_tableschema(self._value)
 
       return self._value
@@ -3047,8 +3050,8 @@ def _expand_direct_read(self, pcoll):
     project_id = None
     temp_table_ref = None
     if 'temp_dataset' in self._kwargs:
-      temp_table_ref = gcp_bigquery.TableReference(
-          gcp_bigquery.DatasetReference(
+      temp_table_ref = TableReference(
+          DatasetReference(
               self._kwargs['temp_dataset'].project,
               self._kwargs['temp_dataset'].dataset_id),
           'beam_temp_table_' + uuid.uuid4().hex)

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

@@ -55,6 +55,7 @@
   from google.cloud import bigquery as gcp_bigquery
   DatasetReference = gcp_bigquery.DatasetReference
   TableReference = gcp_bigquery.TableReference
+  SchemaField = gcp_bigquery.SchemaField
 except ImportError:
 
   class DatasetReference(object):
@@ -63,6 +64,9 @@ class DatasetReference(object):
   class TableReference(object):
     pass
 
+  class SchemaField(object):
+    pass
+
 
 _LOGGER = logging.getLogger(__name__)
 

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

@@ -413,7 +413,7 @@ class DummyTable:
       class DummySchema:
         fields = []
 
-      numBytes = 5
+      num_bytes = 5
       schema = DummySchema()
 
     # TODO(https://github.com/apache/beam/issues/34549): This test relies on
@@ -528,7 +528,7 @@ class DummyTable:
       class DummySchema:
         fields = []
 
-      numBytes = 5
+      num_bytes = 5
       schema = DummySchema()
 
     with mock.patch('time.sleep'), \
@@ -683,7 +683,7 @@ def test_table_spec_display_data(self):
     sink = beam.io.BigQuerySink('dataset.table')
     dd = DisplayData.create_from(sink)
     expected_items = [
-        DisplayDataItemMatcher('table', 'apache-beam-testing:dataset.table'),
+        DisplayDataItemMatcher('table', 'dataset.table'),
         DisplayDataItemMatcher('validation', False)
     ]
     hc.assert_that(dd.items, hc.contains_inanyorder(*expected_items))

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

@@ -125,11 +125,19 @@ def __repr__(self):
 try:
   from google.cloud import bigquery as gcp_bigquery
   TableReference = gcp_bigquery.TableReference
+  DatasetReference = gcp_bigquery.DatasetReference
+  SchemaField = gcp_bigquery.SchemaField
 except ImportError:
 
   class TableReference(object):
     pass
 
+  class DatasetReference(object):
+    pass
+
+  class SchemaField(object):
+    pass
+
 
 # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
 
@@ -216,8 +224,7 @@ def get_hashable_destination(destination):
 
 
 def to_hashable_table_ref(
-    table_ref_elem_kv: tuple[Union[str, gcp_bigquery.TableReference], V]
-) -> tuple[str, V]:
+    table_ref_elem_kv: tuple[Union[str, TableReference], V]) -> tuple[str, V]:
   """Turns the key of the input tuple to its string representation. The key
   should be either a string or a TableReference.
 
@@ -243,7 +250,7 @@ def _parse_schema_field(field):
     mode = field.get('mode', 'NULLABLE')
     description = field.get('description')
     fields = tuple([_parse_schema_field(x) for x in field.get('fields', [])])
-    return gcp_bigquery.SchemaField(
+    return SchemaField(
         field['name'],
         field['type'],
         mode=mode,
@@ -281,10 +288,9 @@ def parse_table_reference(table, dataset=None, project=None):
       format.
   """
 
-  if isinstance(table, gcp_bigquery.TableReference):
-    return gcp_bigquery.TableReference(
-        gcp_bigquery.DatasetReference(table.project, table.dataset_id),
-        table.table_id)
+  if isinstance(table, TableReference):
+    return TableReference(
+        DatasetReference(table.project, table.dataset_id), table.table_id)
   elif callable(table):
     return table
   elif isinstance(table, value_provider.ValueProvider):
@@ -307,8 +313,7 @@ def parse_table_reference(table, dataset=None, project=None):
     # A dummy project is used. It's often overridden by the pipeline options
     project = FALLBACK_PROJECT
 
-  return gcp_bigquery.TableReference(
-      gcp_bigquery.DatasetReference(project, dataset), table)
+  return TableReference(DatasetReference(project, dataset), table)
 
 
 # -----------------------------------------------------------------------------
@@ -784,8 +789,8 @@ def _create_table(
 
     additional_parameters = additional_parameters or {}
     table = gcp_bigquery.Table(
-        table_ref=gcp_bigquery.TableReference(
-            gcp_bigquery.DatasetReference(project_id, dataset_id), table_id),
+        table_ref=TableReference(
+            DatasetReference(project_id, dataset_id), table_id),
         schema=schema,
         **additional_parameters)
     response = self.client.create_table(table)
@@ -1327,7 +1332,7 @@ def convert_row_to_dict(self, row, schema):
       if isinstance(schema, (tuple, list)):
         cell = row.f[index]
         value = from_json_value(cell.v) if cell.v is not None else None
-      elif isinstance(schema, gcp_bigquery.SchemaField):
+      elif isinstance(schema, SchemaField):
         cell = row['f'][index]
         value = cell['v'] if 'v' in cell else None
       if field.mode == 'REPEATED':
@@ -1577,11 +1582,11 @@ def beam_row_from_dict(row: dict, schema):
   Returns:
     ~apache_beam.pvalue.Row: The converted row.
   """
-  if not isinstance(schema, (tuple, list, gcp_bigquery.SchemaField)):
+  if not isinstance(schema, (tuple, list, SchemaField)):
     schema = get_bq_tableschema(schema)
   beam_row = {}
   fields_to_iterate = schema.fields if isinstance(
-      schema, gcp_bigquery.SchemaField) else schema
+      schema, SchemaField) else schema
   for field in fields_to_iterate:
     name = field.name
     mode = field.mode.upper()
@@ -1632,7 +1637,7 @@ def get_table_schema_from_string(schema):
   schema_list = [s.strip() for s in schema.split(',')]
   for field_and_type in schema_list:
     field_name, field_type = field_and_type.split(':')
-    field_schema = gcp_bigquery.SchemaField(
+    field_schema = SchemaField(
         name=field_name, field_type=field_type, mode='NULLABLE')
     table_schema.append(field_schema)
   return table_schema
@@ -1748,11 +1753,11 @@ def get_beam_typehints_from_tableschema(schema, type_overrides=None):
     Nested and repeated fields are supported.
   """
   effective_types = {**BIGQUERY_TYPE_TO_PYTHON_TYPE, **(type_overrides or {})}
-  if not isinstance(schema, (tuple, list, gcp_bigquery.SchemaField)):
+  if not isinstance(schema, (tuple, list, SchemaField)):
     schema = get_bq_tableschema(schema)
   typehints = []
   fields_to_iterate = schema.fields if isinstance(
-      schema, gcp_bigquery.SchemaField) else schema
+      schema, SchemaField) else schema
   for field in fields_to_iterate:
     name, field_type, mode = field.name, field.field_type.upper(), field.mode.upper()
 
@@ -1823,11 +1828,10 @@ def check_schema_equal(
   Returns:
     bool: True if the schemas are equivalent, False otherwise.
   """
-  if type(left) != type(right) or not isinstance(
-      left, (tuple, gcp_bigquery.SchemaField)):
+  if type(left) != type(right) or not isinstance(left, (tuple, SchemaField)):
     return False
 
-  if isinstance(left, gcp_bigquery.SchemaField):
+  if isinstance(left, SchemaField):
     if left.name != right.name:
       return False
 

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

@@ -56,8 +56,6 @@
 # Protect against environments where bigquery library is not available.
 # pylint: disable=wrong-import-order, wrong-import-position
 try:
-  from apitools.base.py.exceptions import GoogleAPICallError
-  from apitools.base.py.exceptions import HttpForbiddenError
   from google.api_core import exceptions as google_api_core_exceptions
   from google.api_core.exceptions import ClientError
   from google.api_core.exceptions import DeadlineExceeded
@@ -155,7 +153,7 @@ def test_calling_with_partially_qualified_table_ref(self):
     parsed_ref = parse_table_reference(partially_qualified_table)
     self.assertEqual(parsed_ref.dataset_id, datasetId)
     self.assertEqual(parsed_ref.table_id, tableId)
-    self.assertEqual(parsed_ref.project, 'apache-beam-testing')
+    self.assertEqual(parsed_ref.project, 'beam_fallback_project')
 
   def test_calling_with_insufficient_table_ref(self):
     table = 'test_table'
@@ -581,7 +579,7 @@ def test_start_query_job_priority_configuration(self):
         priority=beam.io.BigQueryQueryPriority.BATCH)
 
     self.assertEqual(
-        client.query.call_args[0][0].job.configuration.query.priority, 'BATCH')
+        client.query.call_args[1]['job_config'].priority, 'BATCH')
 
     wrapper._start_query_job(
         "my_project",
@@ -592,7 +590,7 @@ def test_start_query_job_priority_configuration(self):
         priority=beam.io.BigQueryQueryPriority.INTERACTIVE)
 
     self.assertEqual(
-        client.query.call_args[0][0].job.configuration.query.priority,
+        client.query.call_args[1]['job_config'].priority,
         'INTERACTIVE')
 
   def test_get_temp_table_project_with_temp_table_ref(self):