diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 9cb6af8f0..4bdd43e8f 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -806,13 +806,12 @@ def get_dataset( def get_iam_policy( self, - table: Union[Table, TableReference], + table: Union[Table, TableReference, TableListItem, str], requested_policy_version: int = 1, retry: retries.Retry = DEFAULT_RETRY, timeout: float = DEFAULT_TIMEOUT, ) -> Policy: - if not isinstance(table, (Table, TableReference)): - raise TypeError("table must be a Table or TableReference") + table = _table_arg_to_table_ref(table, default_project=self.project) if requested_policy_version != 1: raise ValueError("only IAM policy version 1 is supported") @@ -835,14 +834,13 @@ def get_iam_policy( def set_iam_policy( self, - table: Union[Table, TableReference], + table: Union[Table, TableReference, TableListItem, str], policy: Policy, updateMask: str = None, retry: retries.Retry = DEFAULT_RETRY, timeout: float = DEFAULT_TIMEOUT, ) -> Policy: - if not isinstance(table, (Table, TableReference)): - raise TypeError("table must be a Table or TableReference") + table = _table_arg_to_table_ref(table, default_project=self.project) if not isinstance(policy, (Policy)): raise TypeError("policy must be a Policy") @@ -869,13 +867,12 @@ def set_iam_policy( def test_iam_permissions( self, - table: Union[Table, TableReference], + table: Union[Table, TableReference, TableListItem, str], permissions: Sequence[str], retry: retries.Retry = DEFAULT_RETRY, timeout: float = DEFAULT_TIMEOUT, ) -> Dict[str, Any]: - if not isinstance(table, (Table, TableReference)): - raise TypeError("table must be a Table or TableReference") + table = _table_arg_to_table_ref(table, default_project=self.project) body = {"permissions": permissions} @@ -982,7 +979,7 @@ def get_routine( def get_table( self, - table: Union[Table, TableReference, str], + table: Union[Table, TableReference, TableListItem, str], retry: retries.Retry = DEFAULT_RETRY, timeout: float = DEFAULT_TIMEOUT, ) -> Table: @@ -992,6 +989,7 @@ def get_table( table (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): A reference to the table to fetch from the BigQuery API. @@ -1757,7 +1755,7 @@ def delete_routine( def delete_table( self, - table: Union[Table, TableReference, str], + table: Union[Table, TableReference, TableListItem, str], retry: retries.Retry = DEFAULT_RETRY, timeout: float = DEFAULT_TIMEOUT, not_found_ok: bool = False, @@ -1771,6 +1769,7 @@ def delete_table( table (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): A reference to the table to delete. If a string is passed in, @@ -2257,7 +2256,7 @@ def api_request(*args, **kwargs): def load_table_from_uri( self, source_uris: Union[str, Sequence[str]], - destination: Union[Table, TableReference, str], + destination: Union[Table, TableReference, TableListItem, str], job_id: str = None, job_id_prefix: str = None, location: str = None, @@ -2278,6 +2277,7 @@ def load_table_from_uri( destination (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): Table into which data is to be loaded. If a string is passed @@ -2339,7 +2339,7 @@ def load_table_from_uri( def load_table_from_file( self, file_obj: BinaryIO, - destination: Union[Table, TableReference, str], + destination: Union[Table, TableReference, TableListItem, str], rewind: bool = False, size: int = None, num_retries: int = _DEFAULT_NUM_RETRIES, @@ -2360,6 +2360,7 @@ def load_table_from_file( destination (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): Table into which data is to be loaded. If a string is passed @@ -2699,7 +2700,7 @@ def load_table_from_dataframe( def load_table_from_json( self, json_rows: Iterable[Dict[str, Any]], - destination: Union[Table, TableReference, str], + destination: Union[Table, TableReference, TableListItem, str], num_retries: int = _DEFAULT_NUM_RETRIES, job_id: str = None, job_id_prefix: str = None, @@ -2733,6 +2734,7 @@ def load_table_from_json( destination (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): Table into which data is to be loaded. If a string is passed @@ -2980,9 +2982,13 @@ def _do_multipart_upload( def copy_table( self, sources: Union[ - Table, TableReference, str, Sequence[Union[Table, TableReference, str]] + Table, + TableReference, + TableListItem, + str, + Sequence[Union[Table, TableReference, TableListItem, str]], ], - destination: Union[Table, TableReference, str], + destination: Union[Table, TableReference, TableListItem, str], job_id: str = None, job_id_prefix: str = None, location: str = None, @@ -3000,11 +3006,13 @@ def copy_table( sources (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ Sequence[ \ Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ] \ ], \ @@ -3013,6 +3021,7 @@ def copy_table( destination (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): Table into which data is to be copied. @@ -3084,7 +3093,7 @@ def copy_table( def extract_table( self, - source: Union[Table, TableReference, Model, ModelReference, str], + source: Union[Table, TableReference, TableListItem, Model, ModelReference, str], destination_uris: Union[str, Sequence[str]], job_id: str = None, job_id_prefix: str = None, @@ -3104,6 +3113,7 @@ def extract_table( source (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ google.cloud.bigquery.model.Model, \ google.cloud.bigquery.model.ModelReference, \ src, \ @@ -3465,7 +3475,7 @@ def insert_rows_from_dataframe( def insert_rows_json( self, - table: Union[Table, TableReference, str], + table: Union[Table, TableReference, TableListItem, str], json_rows: Sequence[Dict], row_ids: Union[Iterable[str], AutoRowIDs, None] = AutoRowIDs.GENERATE_UUID, skip_invalid_rows: bool = None, @@ -3483,6 +3493,7 @@ def insert_rows_json( table (Union[ \ google.cloud.bigquery.table.Table \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str \ ]): The destination table for the row data, or a reference to it. @@ -3605,7 +3616,7 @@ def insert_rows_json( def list_partitions( self, - table: Union[Table, TableReference, str], + table: Union[Table, TableReference, TableListItem, str], retry: retries.Retry = DEFAULT_RETRY, timeout: float = DEFAULT_TIMEOUT, ) -> Sequence[str]: @@ -3615,6 +3626,7 @@ def list_partitions( table (Union[ \ google.cloud.bigquery.table.Table, \ google.cloud.bigquery.table.TableReference, \ + google.cloud.bigquery.table.TableListItem, \ str, \ ]): The table or reference from which to get partition info diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 48dacf7e2..11b336728 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -1554,7 +1554,7 @@ def test_get_iam_policy_w_invalid_table(self): self.PROJECT, self.DS_ID, self.TABLE_ID, ) - with self.assertRaises(TypeError): + with self.assertRaises(ValueError): client.get_iam_policy(table_resource_string) def test_get_iam_policy_w_invalid_version(self): @@ -1675,7 +1675,7 @@ def test_set_iam_policy_w_invalid_table(self): self.TABLE_ID, ) - with self.assertRaises(TypeError): + with self.assertRaises(ValueError): client.set_iam_policy(table_resource_string, policy) def test_test_iam_permissions(self): @@ -1717,7 +1717,7 @@ def test_test_iam_permissions_w_invalid_table(self): PERMISSIONS = ["bigquery.tables.get", "bigquery.tables.update"] - with self.assertRaises(TypeError): + with self.assertRaises(ValueError): client.test_iam_permissions(table_resource_string, PERMISSIONS) def test_update_dataset_w_invalid_field(self):