Skip to content

Commit 168f035

Browse files
steffnaytswast
andauthored
feat: add support for unrecognized model types (#401)
* feat: add support for unrecognized model types * refactor Co-authored-by: Tim Swast <[email protected]>
1 parent 5a422eb commit 168f035

File tree

2 files changed

+26
-3
lines changed

2 files changed

+26
-3
lines changed

google/cloud/bigquery/model.py

+9-3
Original file line numberDiff line numberDiff line change
@@ -305,9 +305,15 @@ def from_api_repr(cls, resource):
305305
start_time = datetime_helpers.from_microseconds(1e3 * float(start_time))
306306
training_run["startTime"] = datetime_helpers.to_rfc3339(start_time)
307307

308-
this._proto = json_format.ParseDict(
309-
resource, types.Model()._pb, ignore_unknown_fields=True
310-
)
308+
try:
309+
this._proto = json_format.ParseDict(
310+
resource, types.Model()._pb, ignore_unknown_fields=True
311+
)
312+
except json_format.ParseError:
313+
resource["modelType"] = "MODEL_TYPE_UNSPECIFIED"
314+
this._proto = json_format.ParseDict(
315+
resource, types.Model()._pb, ignore_unknown_fields=True
316+
)
311317
return this
312318

313319
def _build_resource(self, filter_fields):

tests/unit/model/test_model.py

+17
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,23 @@ def test_from_api_repr_w_unknown_fields(target_class):
186186
assert got._properties is resource
187187

188188

189+
def test_from_api_repr_w_unknown_type(target_class):
190+
from google.cloud.bigquery import ModelReference
191+
192+
resource = {
193+
"modelReference": {
194+
"projectId": "my-project",
195+
"datasetId": "my_dataset",
196+
"modelId": "my_model",
197+
},
198+
"modelType": "BE_A_GOOD_ROLE_MODEL",
199+
}
200+
got = target_class.from_api_repr(resource)
201+
assert got.reference == ModelReference.from_string("my-project.my_dataset.my_model")
202+
assert got.model_type == 0
203+
assert got._properties is resource
204+
205+
189206
@pytest.mark.parametrize(
190207
"resource,filter_fields,expected",
191208
[

0 commit comments

Comments
 (0)