Skip to content

Commit

Permalink
feat: add support for unrecognized model types (#401)
Browse files Browse the repository at this point in the history
* feat: add support for unrecognized model types

* refactor

Co-authored-by: Tim Swast <swast@google.com>
  • Loading branch information
steffnay and tswast authored Nov 30, 2020
1 parent 5a422eb commit 168f035
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 3 deletions.
12 changes: 9 additions & 3 deletions google/cloud/bigquery/model.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,9 +305,15 @@ def from_api_repr(cls, resource):
start_time = datetime_helpers.from_microseconds(1e3 * float(start_time))
training_run["startTime"] = datetime_helpers.to_rfc3339(start_time)

this._proto = json_format.ParseDict(
resource, types.Model()._pb, ignore_unknown_fields=True
)
try:
this._proto = json_format.ParseDict(
resource, types.Model()._pb, ignore_unknown_fields=True
)
except json_format.ParseError:
resource["modelType"] = "MODEL_TYPE_UNSPECIFIED"
this._proto = json_format.ParseDict(
resource, types.Model()._pb, ignore_unknown_fields=True
)
return this

def _build_resource(self, filter_fields):
Expand Down
17 changes: 17 additions & 0 deletions tests/unit/model/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -186,6 +186,23 @@ def test_from_api_repr_w_unknown_fields(target_class):
assert got._properties is resource


def test_from_api_repr_w_unknown_type(target_class):
from google.cloud.bigquery import ModelReference

resource = {
"modelReference": {
"projectId": "my-project",
"datasetId": "my_dataset",
"modelId": "my_model",
},
"modelType": "BE_A_GOOD_ROLE_MODEL",
}
got = target_class.from_api_repr(resource)
assert got.reference == ModelReference.from_string("my-project.my_dataset.my_model")
assert got.model_type == 0
assert got._properties is resource


@pytest.mark.parametrize(
"resource,filter_fields,expected",
[
Expand Down

0 comments on commit 168f035

Please sign in to comment.