Skip to content

Commit

Permalink
feat: accept TableListItem where TableReference is accepted (#1016)
Browse files Browse the repository at this point in the history 
Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:
- [ ] Make sure to open an issue as a [bug/issue](https://github.com/googleapis/python-bigquery/issues/new/choose) before writing your code!  That way we can discuss the change, evaluate designs, and agree on the general idea
- [ ] Ensure the tests and linter pass
- [ ] Code coverage does not decrease (if any source code was changed)
- [ ] Appropriate docs were updated (if necessary)

🦕
  • Loading branch information
@steffnay
steffnay authored Nov 1, 2021
1 parent 4ecd64b commit fe16adc
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 22 deletions.
50 changes: 31 additions & 19 deletions google/cloud/bigquery/client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand All @@ -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")
Expand All @@ -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}

Expand Down Expand Up @@ -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:
Expand All @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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,
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -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, \
] \
], \
Expand All @@ -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.
Expand Down Expand Up @@ -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,
Expand All @@ -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, \
Expand Down Expand Up @@ -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,
Expand All @@ -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.
Expand Down Expand Up @@ -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]:
Expand All @@ -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
Expand Down
6 changes: 3 additions & 3 deletions tests/unit/test_client.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down

0 comments on commit fe16adc

Please sign in to comment.