Beispiel #1
0
    def test_execute(self, mock_hook):
        filter_ = "filter"
        page_size = 42

        op = AutoMLTablesListTableSpecsOperator(
            dataset_id=DATASET_ID,
            location=GCP_LOCATION,
            project_id=GCP_PROJECT_ID,
            filter_=filter_,
            page_size=page_size,
            task_id=TASK_ID,
        )
        op.execute(context=None)
        mock_hook.return_value.list_table_specs.assert_called_once_with(
            dataset_id=DATASET_ID,
            filter_=filter_,
            location=GCP_LOCATION,
            metadata=None,
            page_size=page_size,
            project_id=GCP_PROJECT_ID,
            retry=None,
            timeout=None,
        )
Beispiel #2
0
    MODEL["dataset_id"] = dataset_id

    # [START howto_operator_automl_import_data]
    import_dataset_task = AutoMLImportDataOperator(
        task_id="import_dataset_task",
        dataset_id=dataset_id,
        location=GCP_AUTOML_LOCATION,
        input_config=IMPORT_INPUT_CONFIG,
    )
    # [END howto_operator_automl_import_data]

    # [START howto_operator_automl_specs]
    list_tables_spec_task = AutoMLTablesListTableSpecsOperator(
        task_id="list_tables_spec_task",
        dataset_id=dataset_id,
        location=GCP_AUTOML_LOCATION,
        project_id=GCP_PROJECT_ID,
    )
    # [END howto_operator_automl_specs]

    # [START howto_operator_automl_column_specs]
    list_columns_spec_task = AutoMLTablesListColumnSpecsOperator(
        task_id="list_columns_spec_task",
        dataset_id=dataset_id,
        table_spec_id=
        "{{ extract_object_id(task_instance.xcom_pull('list_tables_spec_task')[0]) }}",
        location=GCP_AUTOML_LOCATION,
        project_id=GCP_PROJECT_ID,
    )
    # [END howto_operator_automl_column_specs]