From 5bcaa92a2171d5734bdfd74eb0d54b42f3c2208f Mon Sep 17 00:00:00 2001 From: steffnay Date: Sun, 10 Oct 2021 14:15:46 -0700 Subject: [PATCH 1/3] feat: accept TableListItem where TableReference is accepted --- google/cloud/bigquery/client.py | 62 +++++++++++++++++++++++---------- 1 file changed, 43 insertions(+), 19 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index a8a1c1e16..468792f58 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -805,13 +805,16 @@ 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 not isinstance(table, (TableReference)): + raise TypeError( + "table must be a Table, TableReference, or TableListItem, or string" + ) if requested_policy_version != 1: raise ValueError("only IAM policy version 1 is supported") @@ -834,14 +837,17 @@ 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(table, (TableReference)): + raise TypeError( + "table must be a Table, TableReference, or TableListItem, or string" + ) if not isinstance(policy, (Policy)): raise TypeError("policy must be a Policy") @@ -868,13 +874,16 @@ 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) + if not isinstance(table, (TableReference)): + raise TypeError( + "table must be a Table, TableReference, or TableListItem, or string" + ) body = {"permissions": permissions} @@ -981,7 +990,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: @@ -991,6 +1000,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. @@ -1756,7 +1766,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, @@ -1770,6 +1780,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, @@ -2256,7 +2267,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, @@ -2277,6 +2288,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 @@ -2338,7 +2350,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, @@ -2359,6 +2371,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 @@ -2698,7 +2711,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, @@ -2732,6 +2745,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 @@ -2979,9 +2993,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, @@ -2999,11 +3017,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, \ ] \ ], \ @@ -3012,6 +3032,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. @@ -3083,7 +3104,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, @@ -3103,6 +3124,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, \ @@ -3464,7 +3486,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, @@ -3482,6 +3504,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. @@ -3604,7 +3627,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]: @@ -3614,6 +3637,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 From 297f9a3f6f4a1fd83535f3f82aa22f2d9e218259 Mon Sep 17 00:00:00 2001 From: steffnay Date: Tue, 12 Oct 2021 13:29:16 -0700 Subject: [PATCH 2/3] update unit test --- tests/unit/test_client.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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): From 9467abcb04291ecca46df9d8492e2356fd970916 Mon Sep 17 00:00:00 2001 From: steffnay Date: Tue, 26 Oct 2021 12:21:27 -0700 Subject: [PATCH 3/3] remove raise --- google/cloud/bigquery/client.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/google/cloud/bigquery/client.py b/google/cloud/bigquery/client.py index 8e69e99d3..4bdd43e8f 100644 --- a/google/cloud/bigquery/client.py +++ b/google/cloud/bigquery/client.py @@ -812,10 +812,6 @@ def get_iam_policy( timeout: float = DEFAULT_TIMEOUT, ) -> Policy: table = _table_arg_to_table_ref(table, default_project=self.project) - if not isinstance(table, (TableReference)): - raise TypeError( - "table must be a Table, TableReference, or TableListItem, or string" - ) if requested_policy_version != 1: raise ValueError("only IAM policy version 1 is supported") @@ -845,10 +841,6 @@ def set_iam_policy( timeout: float = DEFAULT_TIMEOUT, ) -> Policy: table = _table_arg_to_table_ref(table, default_project=self.project) - if not isinstance(table, (TableReference)): - raise TypeError( - "table must be a Table, TableReference, or TableListItem, or string" - ) if not isinstance(policy, (Policy)): raise TypeError("policy must be a Policy") @@ -881,10 +873,6 @@ def test_iam_permissions( timeout: float = DEFAULT_TIMEOUT, ) -> Dict[str, Any]: table = _table_arg_to_table_ref(table, default_project=self.project) - if not isinstance(table, (TableReference)): - raise TypeError( - "table must be a Table, TableReference, or TableListItem, or string" - ) body = {"permissions": permissions}