diff --git a/samples/model-builder/create_and_import_dataset_tabular_bigquery_sample.py b/samples/model-builder/create_and_import_dataset_tabular_bigquery_sample.py new file mode 100644 index 0000000000..ea92a5a6dc --- /dev/null +++ b/samples/model-builder/create_and_import_dataset_tabular_bigquery_sample.py @@ -0,0 +1,44 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Union + +from google.cloud import aiplatform + + +# [START aiplatform_sdk_create_and_import_dataset_tabular_bigquery_sample] +def create_and_import_dataset_tabular_bigquery_sample( + project: str, + location: str, + display_name: str, + src_uris: Union[str, List[str]], + sync: bool, +): + aiplatform.init(project=project, location=location) + + ds = aiplatform.TabularDataset.create( + display_name=display_name, + bq_source=src_uris, + sync=sync, + ) + + if not sync: + ds.wait() + + print(ds.display_name) + print(ds.resource_name) + return ds + + +# [END aiplatform_sdk_create_and_import_dataset_tabular_bigquery_sample] diff --git a/samples/model-builder/create_and_import_dataset_tabular_bigquery_test.py b/samples/model-builder/create_and_import_dataset_tabular_bigquery_test.py new file mode 100644 index 0000000000..82b9fc2eb6 --- /dev/null +++ b/samples/model-builder/create_and_import_dataset_tabular_bigquery_test.py @@ -0,0 +1,39 @@ +# Copyright 2021 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +from google.cloud.aiplatform import schema + +import create_and_import_dataset_tabular_bigquery_sample +import test_constants as constants + + +def test_create_and_import_dataset_tabular_bigquery_sample(mock_sdk_init, mock_create_tabular_dataset): + + create_and_import_dataset_tabular_bigquery_sample.create_and_import_dataset_tabular_bigquery_sample( + project=constants.PROJECT, + location=constants.LOCATION, + src_uris=constants.BIGQUERY_SOURCE, + display_name=constants.DISPLAY_NAME, + sync=True, + ) + + mock_sdk_init.assert_called_once_with( + project=constants.PROJECT, location=constants.LOCATION + ) + mock_create_tabular_dataset.assert_called_once_with( + display_name=constants.DISPLAY_NAME, + bq_source=constants.BIGQUERY_SOURCE, + sync=True, + )