feat(bigquery): add support of model for extract job

google/cloud/bigquery/client.py

@@ -1365,6 +1365,8 @@ def create_job(self, job_config, retry=DEFAULT_RETRY):
                 job_config
             )
             source = _get_sub_prop(job_config, ["extract", "sourceTable"])
+            if not source:
+                source = _get_sub_prop(job_config, ["extract", "sourceModel"])
             destination_uris = _get_sub_prop(job_config, ["extract", "destinationUris"])
             return self.extract_table(
                 source, destination_uris, job_config=extract_job_config, retry=retry
@@ -2345,10 +2347,10 @@ def extract_table(
             location = self.location
 
         job_ref = job._JobReference(job_id, project=project, location=location)
-
-        if source_type.lower() == "table":
+        src = source_type.lower()
+        if src == "table":
             source = _table_arg_to_table_ref(source, default_project=self.project)
-        elif source_type.lower() == "model":
+        elif src == "model":
             source = _model_arg_to_model_ref(source, default_project=self.project)
         else:
             raise ValueError(

google/cloud/bigquery/job.py

@@ -1990,8 +1990,11 @@ class ExtractJob(_AsyncJob):
     Args:
         job_id (str): the job's ID.
 
-        source (google.cloud.bigquery.table.TableReference):
-            Table into which data is to be loaded.
+        source (Union[ \
+            google.cloud.bigquery.table.TableReference, \
+            google.cloud.bigquery.model.ModelReference \
+        ]):
+            Table or Model into which data is to be loaded.
 
         destination_uris (List[str]):
             URIs describing where the extracted data will be written in Cloud
@@ -2073,9 +2076,9 @@ def to_api_repr(self):
             "datasetId": self.source.dataset_id,
         }
 
+        source = "sourceTable"
         if isinstance(self.source, TableReference):
             source_ref["tableId"] = self.source.table_id
-            source = "sourceTable"
         else:
             source_ref["modelId"] = self.source.model_id
             source = "sourceModel"

google/cloud/bigquery/model.py

@@ -441,7 +441,7 @@ def _model_arg_to_model_ref(value, default_project=None):
     This function keeps ModelReference and other kinds of objects unchanged.
     """
     if isinstance(value, six.string_types):
-        value = ModelReference.from_string(value, default_project=default_project)
+        return ModelReference.from_string(value, default_project=default_project)
     if isinstance(value, Model):
-        value = value.reference
+        return value.reference
     return value

tests/unit/test_client.py

@@ -2884,6 +2884,21 @@ def test_create_job_extract_config(self):
             configuration, "google.cloud.bigquery.client.Client.extract_table",
         )
 
+    def test_create_job_extract_config_for_model(self):
+        configuration = {
+            "extract": {
+                "sourceModel": {
+                    "projectId": self.PROJECT,
+                    "datasetId": self.DS_ID,
+                    "modelId": "source_model",
+                },
+                "destinationUris": ["gs://test_bucket/dst_object*"],
+            }
+        }
+        self._create_job_helper(
+            configuration, "google.cloud.bigquery.client.Client.extract_table",
+        )
+
     def test_create_job_query_config(self):
         configuration = {
             "query": {"query": "query", "destinationTable": {"tableId": "table_id"}}