def test_execute(self, mock_hook): max_results = '100' selected_fields = 'DATE' operator = BigQueryGetDataOperator(task_id=TASK_ID, dataset_id=TEST_DATASET, table_id=TEST_TABLE_ID, max_results=max_results, selected_fields=selected_fields, ) operator.execute(None) mock_hook.return_value \ .get_conn.return_value \ .cursor.return_value \ .get_tabledata \ .assert_called_once_with( dataset_id=TEST_DATASET, table_id=TEST_TABLE_ID, max_results=max_results, selected_fields=selected_fields, )
def test_execute(self, mock_hook): max_results = '100' selected_fields = 'DATE' operator = BigQueryGetDataOperator(task_id=TASK_ID, dataset_id=TEST_DATASET, table_id=TEST_TABLE_ID, max_results=max_results, selected_fields=selected_fields, ) operator.execute(None) mock_hook.return_value \ .get_conn.return_value \ .cursor.return_value \ .get_tabledata \ .assert_called_once_with( dataset_id=TEST_DATASET, table_id=TEST_TABLE_ID, max_results=max_results, selected_fields=selected_fields, )
destination_dataset_table="{}.save_query_result".format(DATASET_NAME), query_params=[{ "name": "to_address", "parameterType": { "type": "STRING" }, "parameterValue": { "value": WALLET_ADDRESS }, }], ) get_data = BigQueryGetDataOperator( task_id="get-data", dataset_id=DATASET_NAME, table_id="save_query_result", max_results="10", selected_fields="value,to_address", ) get_data_result = BashOperator( task_id="get-data-result", bash_command="echo \"{{ task_instance.xcom_pull('get-data') }}\"") create_external_table = BigQueryCreateExternalTableOperator( task_id="create-external-table", bucket=DATA_SAMPLE_GCS_BUCKET_NAME, source_objects=[DATA_SAMPLE_GCS_OBJECT_NAME], destination_project_dataset_table="{}.external_table".format( DATASET_NAME), skip_leading_rows=1,
# Load data from Bigquery, store result in a temp table load_test_data_task = BigQueryOperator( task_id='Load_test_data', dag=dag, sql=bigquerySQL, destination_dataset_table="%s.%s" % (dag_config['destination_dataset'], dag_config['destination_table']), use_legacy_sql=False, write_disposition='WRITE_TRUNCATE' ) # Load Bigquery result data from temp table load_query_results_task = BigQueryGetDataOperator( task_id='bq_get_data', dataset_id=dag_config['destination_dataset'], table_id=dag_config['destination_table'], dag=dag, max_results=100000 ) # Load data into Pandas and apply test logic def load_bq_result_intoPandas(**kwargs): print(kwargs) ti = kwargs['ti'] resultData = ti.xcom_pull(task_ids='bq_get_data') df = pd.DataFrame(resultData) df.columns=['your columns', 'your columns'] df['timestamp'] = pd.to_datetime(df['timestamp'], unit='s', utc=True).dt.tz_convert('US/Pacific') print("Retrieved ", len(df), " rows") print(df.head())