def execute(self, context: 'Context'):
hook = CloudAutoMLHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Updating AutoML dataset %s.", self.dataset["name"])
result = hook.update_dataset(
dataset=self.dataset,
update_mask=self.update_mask,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
self.log.info("Dataset updated.")
return Dataset.to_dict(result)
def execute(self, context: 'Context'):
hook = CloudAutoMLHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Creating dataset")
result = hook.create_dataset(
dataset=self.dataset,
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
result = Dataset.to_dict(result)
dataset_id = hook.extract_object_id(result)
self.log.info("Creating completed. Dataset id: %s", dataset_id)
self.xcom_push(context, key="dataset_id", value=dataset_id)
return result
def test_execute(self, mock_hook):
mock_hook.return_value.update_dataset.return_value = Dataset(
name=DATASET_PATH)
dataset = copy.deepcopy(DATASET)
dataset["name"] = DATASET_ID
op = AutoMLTablesUpdateDatasetOperator(
dataset=dataset,
update_mask=MASK,
location=GCP_LOCATION,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.update_dataset.assert_called_once_with(
dataset=dataset,
metadata=None,
retry=None,
timeout=None,
update_mask=MASK,
)
def execute(self, context: 'Context'):
hook = CloudAutoMLHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
)
self.log.info("Requesting datasets")
page_iterator = hook.list_datasets(
location=self.location,
project_id=self.project_id,
retry=self.retry,
timeout=self.timeout,
metadata=self.metadata,
)
result = [Dataset.to_dict(dataset) for dataset in page_iterator]
self.log.info("Datasets obtained.")
self.xcom_push(
context,
key="dataset_id_list",
value=[hook.extract_object_id(d) for d in result],
)
return result
def test_execute(self, mock_hook, mock_xcom):
mock_hook.return_value.create_dataset.return_value = Dataset(
name=DATASET_PATH)
mock_hook.return_value.extract_object_id = extract_object_id
op = AutoMLCreateDatasetOperator(
dataset=DATASET,
location=GCP_LOCATION,
project_id=GCP_PROJECT_ID,
task_id=TASK_ID,
)
op.execute(context=None)
mock_hook.return_value.create_dataset.assert_called_once_with(
dataset=DATASET,
location=GCP_LOCATION,
metadata=None,
project_id=GCP_PROJECT_ID,
retry=None,
timeout=None,
)
mock_xcom.assert_called_once_with(None,
key="dataset_id",
value=DATASET_ID)