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)