Content-Length: 352376 | pFad | https://github.com/apache/airflow/commit/f2ebc292fe63d2ddd0686d90c3acc0630f017a07

8C Provide missing project id and creds for TabularDataset (#31991) · apache/airflow@f2ebc29 · GitHub
Skip to content

Commit f2ebc29

Browse files
Provide missing project id and creds for TabularDataset (#31991)
1 parent 10aa704 commit f2ebc29

File tree

2 files changed

+16
-3
lines changed

2 files changed

+16
-3
lines changed

airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -352,11 +352,16 @@ def execute(self, context: Context):
352352
gcp_conn_id=self.gcp_conn_id,
353353
impersonation_chain=self.impersonation_chain,
354354
)
355+
credentials, _ = self.hook.get_credentials_and_project_id()
355356
model, training_id = self.hook.create_auto_ml_tabular_training_job(
356357
project_id=self.project_id,
357358
region=self.region,
358359
display_name=self.display_name,
359-
dataset=datasets.TabularDataset(dataset_name=self.dataset_id),
360+
dataset=datasets.TabularDataset(
361+
dataset_name=self.dataset_id,
362+
project=self.project_id,
363+
credentials=credentials,
364+
),
360365
target_column=self.target_column,
361366
optimization_prediction_type=self.optimization_prediction_type,
362367
optimization_objective=self.optimization_objective,

tests/providers/google/cloud/operators/test_vertex_ai.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
from __future__ import annotations
1818

1919
from unittest import mock
20+
from unittest.mock import MagicMock
2021

2122
from google.api_core.gapic_v1.method import DEFAULT
2223
from google.api_core.retry import Retry
@@ -783,7 +784,12 @@ class TestVertexAICreateAutoMLTabularTrainingJobOperator:
783784
@mock.patch("google.cloud.aiplatform.datasets.TabularDataset")
784785
@mock.patch(VERTEX_AI_PATH.format("auto_ml.AutoMLHook"))
785786
def test_execute(self, mock_hook, mock_dataset):
786-
mock_hook.return_value.create_auto_ml_tabular_training_job.return_value = (None, "training_id")
787+
mock_hook.return_value = MagicMock(
788+
**{
789+
"create_auto_ml_tabular_training_job.return_value": (None, "training_id"),
790+
"get_credentials_and_project_id.return_value": ("creds", "project_id"),
791+
}
792+
)
787793
op = CreateAutoMLTabularTrainingJobOperator(
788794
task_id=TASK_ID,
789795
gcp_conn_id=GCP_CONN_ID,
@@ -798,7 +804,9 @@ def test_execute(self, mock_hook, mock_dataset):
798804
)
799805
op.execute(context={"ti": mock.MagicMock()})
800806
mock_hook.assert_called_once_with(gcp_conn_id=GCP_CONN_ID, impersonation_chain=IMPERSONATION_CHAIN)
801-
mock_dataset.assert_called_once_with(dataset_name=TEST_DATASET_ID)
807+
mock_dataset.assert_called_once_with(
808+
dataset_name=TEST_DATASET_ID, project=GCP_PROJECT, credentials="creds"
809+
)
802810
mock_hook.return_value.create_auto_ml_tabular_training_job.assert_called_once_with(
803811
project_id=GCP_PROJECT,
804812
region=GCP_LOCATION,

0 commit comments

Comments
 (0)








ApplySandwichStrip

pFad - (p)hone/(F)rame/(a)nonymizer/(d)eclutterfier!      Saves Data!


--- a PPN by Garber Painting Akron. With Image Size Reduction included!

Fetched URL: https://github.com/apache/airflow/commit/f2ebc292fe63d2ddd0686d90c3acc0630f017a07

Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy