feat: add timeout paramter to load_table_from_file and it dependent methods (#327)

Author: HemangChothani

google/cloud/bigquery/client.py

@@ -1591,14 +1591,17 @@ def job_from_resource(self, resource):
  return job.QueryJob.from_api_repr(resource, self)
  return job.UnknownJob.from_api_repr(resource, self)
 
- def create_job(self, job_config, retry=DEFAULT_RETRY):
+ def create_job(self, job_config, retry=DEFAULT_RETRY, timeout=None):
  """Create a new job.
  Args:
  job_config (dict): configuration job representation returned from the API.
 
  Keyword Arguments:
  retry (Optional[google.api_core.retry.Retry]):
  How to retry the RPC.
+ timeout (Optional[float]):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
 
  Returns:
  Union[ \
@@ -1617,7 +1620,11 @@ def create_job(self, job_config, retry=DEFAULT_RETRY):
  destination = _get_sub_prop(job_config, ["load", "destinationTable"])
  source_uris = _get_sub_prop(job_config, ["load", "sourceUris"])
  return self.load_table_from_uri(
- source_uris, destination, job_config=load_job_config, retry=retry
+ source_uris,
+ destination,
+ job_config=load_job_config,
+ retry=retry,
+ timeout=timeout,
  )
  elif "copy" in job_config:
  copy_job_config = google.cloud.bigquery.job.CopyJobConfig.from_api_repr(
@@ -1633,7 +1640,11 @@ def create_job(self, job_config, retry=DEFAULT_RETRY):
  table_ref = TableReference.from_api_repr(source_config)
  sources.append(table_ref)
  return self.copy_table(
- sources, destination, job_config=copy_job_config, retry=retry
+ sources,
+ destination,
+ job_config=copy_job_config,
+ retry=retry,
+ timeout=timeout,
  )
  elif "extract" in job_config:
  extract_job_config = google.cloud.bigquery.job.ExtractJobConfig.from_api_repr(
@@ -1650,6 +1661,7 @@ def create_job(self, job_config, retry=DEFAULT_RETRY):
  destination_uris,
  job_config=extract_job_config,
  retry=retry,
+ timeout=timeout,
  source_type=source_type,
  )
  elif "query" in job_config:
@@ -1659,7 +1671,9 @@ def create_job(self, job_config, retry=DEFAULT_RETRY):
  copy_config
  )
  query = _get_sub_prop(copy_config, ["query", "query"])
- return self.query(query, job_config=query_job_config, retry=retry)
+ return self.query(
+ query, job_config=query_job_config, retry=retry, timeout=timeout
+ )
  else:
  raise TypeError("Invalid job configuration received.")
 
@@ -1981,6 +1995,7 @@ def load_table_from_file(
  location=None,
  project=None,
  job_config=None,
+ timeout=None,
  ):
  """Upload the contents of this table from a file-like object.
 
@@ -2020,6 +2035,9 @@ def load_table_from_file(
  to the client's project.
  job_config (Optional[google.cloud.bigquery.job.LoadJobConfig]):
  Extra configuration options for the job.
+ timeout (Optional[float]):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
 
  Returns:
  google.cloud.bigquery.job.LoadJob: A new load job.
@@ -2058,11 +2076,11 @@ def load_table_from_file(
  try:
  if size is None or size >= _MAX_MULTIPART_SIZE:
  response = self._do_resumable_upload(
- file_obj, job_resource, num_retries
+ file_obj, job_resource, num_retries, timeout
  )
  else:
  response = self._do_multipart_upload(
- file_obj, job_resource, size, num_retries
+ file_obj, job_resource, size, num_retries, timeout
  )
  except resumable_media.InvalidResponse as exc:
  raise exceptions.from_http_response(exc.response)
@@ -2080,6 +2098,7 @@ def load_table_from_dataframe(
  project=None,
  job_config=None,
  parquet_compression="snappy",
+ timeout=None,
  ):
  """Upload the contents of a table from a pandas DataFrame.
 
@@ -2143,6 +2162,9 @@ def load_table_from_dataframe(
  passed as the ``compression`` argument to the underlying
  ``DataFrame.to_parquet()`` method.
  https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_parquet.html#pandas.DataFrame.to_parquet
+ timeout (Optional[float]):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
 
  Returns:
  google.cloud.bigquery.job.LoadJob: A new load job.
@@ -2249,6 +2271,7 @@ def load_table_from_dataframe(
  location=location,
  project=project,
  job_config=job_config,
+ timeout=timeout,
  )
 
  finally:
@@ -2264,6 +2287,7 @@ def load_table_from_json(
  location=None,
  project=None,
  job_config=None,
+ timeout=None,
  ):
  """Upload the contents of a table from a JSON string or dict.
 
@@ -2313,6 +2337,9 @@ def load_table_from_json(
  Extra configuration options for the job. The ``source_format``
  setting is always set to
  :attr:`~google.cloud.bigquery.job.SourceFormat.NEWLINE_DELIMITED_JSON`.
+ timeout (Optional[float]):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
 
  Returns:
  google.cloud.bigquery.job.LoadJob: A new load job.
@@ -2357,9 +2384,10 @@ def load_table_from_json(
  location=location,
  project=project,
  job_config=job_config,
+ timeout=timeout,
  )
 
- def _do_resumable_upload(self, stream, metadata, num_retries):
+ def _do_resumable_upload(self, stream, metadata, num_retries, timeout):
  """Perform a resumable upload.
 
  Args:
@@ -2371,21 +2399,25 @@ def _do_resumable_upload(self, stream, metadata, num_retries):
  Number of upload retries. (Deprecated: This
  argument will be removed in a future release.)
 
+ timeout (float):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
+
  Returns:
  requests.Response:
  The "200 OK" response object returned after the final chunk
  is uploaded.
  """
  upload, transport = self._initiate_resumable_upload(
- stream, metadata, num_retries
+ stream, metadata, num_retries, timeout
  )
 
  while not upload.finished:
  response = upload.transmit_next_chunk(transport)
 
  return response
 
- def _initiate_resumable_upload(self, stream, metadata, num_retries):
+ def _initiate_resumable_upload(self, stream, metadata, num_retries, timeout):
  """Initiate a resumable upload.
 
  Args:
@@ -2397,6 +2429,10 @@ def _initiate_resumable_upload(self, stream, metadata, num_retries):
  Number of upload retries. (Deprecated: This
  argument will be removed in a future release.)
 
+ timeout (float):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
+
  Returns:
  Tuple:
  Pair of
@@ -2419,12 +2455,17 @@ def _initiate_resumable_upload(self, stream, metadata, num_retries):
  )
 
  upload.initiate(
- transport, stream, metadata, _GENERIC_CONTENT_TYPE, stream_final=False
+ transport,
+ stream,
+ metadata,
+ _GENERIC_CONTENT_TYPE,
+ stream_final=False,
+ timeout=timeout,
  )
 
  return upload, transport
 
- def _do_multipart_upload(self, stream, metadata, size, num_retries):
+ def _do_multipart_upload(self, stream, metadata, size, num_retries, timeout):
  """Perform a multipart upload.
 
  Args:
@@ -2441,6 +2482,10 @@ def _do_multipart_upload(self, stream, metadata, size, num_retries):
  Number of upload retries. (Deprecated: This
  argument will be removed in a future release.)
 
+ timeout (float):
+ The number of seconds to wait for the underlying HTTP transport
+ before using ``retry``.
+
  Returns:
  requests.Response:
  The "200 OK" response object returned after the multipart
@@ -2466,7 +2511,9 @@ def _do_multipart_upload(self, stream, metadata, size, num_retries):
  max_retries=num_retries
  )
 
- response = upload.transmit(self._http, data, metadata, _GENERIC_CONTENT_TYPE)
+ response = upload.transmit(
+ self._http, data, metadata, _GENERIC_CONTENT_TYPE, timeout=timeout
+ )
 
  return response