python之使用 apache-airflow (cloud-composer) 调度从 BigQuery 到 MongoDB 的负载

我正在尝试设置将数据从 GCS 移动到 BigQuery 的数据管道，执行某些任务/处理并将其加载到 MongoDB 集群(所有这些都使用 DAG 在 python 中设置)。在加载到 MongoDB 之前，我一直能够实现这一点。是否有任何现有的 Airflow 运营商可以做到这一点？如果没有，是否可以使用 airflow 中提供的 mongoDB Hook 创建自定义代码？

谢谢， GT

编辑 1

我使用了 MongoHook 和来自 BigQueryGetDataOperator 的源代码(下面的代码片段)。我现在的问题是，我需要为 10++ 百万行工作，当我增加 BigQueryGetDataOperator 中的 max_results='100' 默认值时，我收到一个错误:

sqlalchemy.exc.InvalidRequestError:由于刷新期间的先前异常，此 session 的事务已回滚。要使用此 session 开始新事务，首先发出 Session.rollback()。原始异常是:(_mysql_exceptions.DataError)(1406，“第 1 行的列‘值’的数据太长”)

我知道我应该将数据以 block 的形式推送到 XCom，但不确定是否真的可以做到。 Xcom 中是否有解析大量数据的标准方法？使用 Airflow 实现此目的的任何其他替代方法也会有所帮助。我唯一能想到的就是将数据写入 GCS，加载到 MongoDB，然后删除 GCS 文件。

#-------- COPY PASTED BigQueryGetDataOperator SECTION: START -------------- 
'''Source: https://airflow.readthedocs.io/en/stable/_modules/airflow/contrib/operators/bigquery_get_data.html#BigQueryGetDataOperator ''' 
 
from airflow.contrib.hooks.bigquery_hook import BigQueryHook 
from airflow.models import BaseOperator 
from airflow.utils.decorators import apply_defaults 
 
 
class BigQueryGetDataOperator(BaseOperator): 
    template_fields = ('dataset_id', 'table_id', 'max_results') 
    ui_color = '#e4f0e8' 
 
    @apply_defaults 
    def __init__(self, 
                 dataset_id, 
                 table_id, 
                 max_results='100', 
                 selected_fields=None, 
                 bigquery_conn_id='bigquery_default', 
                 delegate_to=None, 
                 *args, 
                 **kwargs): 
        super(BigQueryGetDataOperator, self).__init__(*args, **kwargs) 
        self.dataset_id = dataset_id 
        self.table_id = table_id 
        self.max_results = max_results 
        self.selected_fields = selected_fields 
        self.bigquery_conn_id = bigquery_conn_id 
        self.delegate_to = delegate_to 
 
    def execute(self, context): 
        self.log.info('Fetching Data from:') 
        self.log.info('Dataset: %s ; Table: %s ; Max Results: %s', 
                      self.dataset_id, self.table_id, self.max_results) 
 
        hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id, 
                            delegate_to=self.delegate_to) 
 
        conn = hook.get_conn() 
        cursor = conn.cursor() 
        response = cursor.get_tabledata(dataset_id=self.dataset_id, 
                                        table_id=self.table_id, 
                                        max_results=self.max_results, 
                                        selected_fields=self.selected_fields) 
 
        self.log.info('Total Extracted rows: %s', response['totalRows']) 
        rows = response['rows'] 
        return rows 
        # Below lines were commented as I did not want a list but a json 
        # table_data = [] 
        # for dict_row in rows: 
        #     single_row = [] 
        #     for fields in dict_row['f']: 
        #         single_row.append(fields['v']) 
        #     table_data.append(single_row) 
 
        # return table_data 
#----------------------- COPY PASTED SECTION: END ---------------- 
 
from airflow import models 
from airflow.operators.python_operator import PythonOperator 
from airflow.utils import trigger_rule 
from airflow.contrib.operators import gcs_to_bq 
from airflow.contrib.operators import bigquery_to_gcs 
from airflow.contrib.operators import bigquery_operator 
from airflow.contrib.operators import bigquery_get_data 
from airflow.contrib.operators import MongoHook 
def get_dlist(**kwargs): 
  import logging as log 
  #Import pymongo 
  from pymongo import MongoClient 
  #Pull the data saved in XCom 
  value = kwargs.get('task_instance').xcom_pull(task_ids='get_data_in_list_from_bq') 
 
  header = ['V1','V2'] 
  data=[] 
  for rows in value: 
    onerow={} 
    for i,f in zip(range(len(rows['f'])),rows['f']): 
      onerow[header[i]] = f['v'] 
    data.append(onerow) 
  log.info("Pulled...") 
  log.info(data) 
  log.info("Pushing into mongodb...") 
  client = MongoClient(localhost:27017) 
  db = client.test 
  collection = db.testingbq2mongo 
  collection.insert(data) 
  log.info("Written to mongoDB...") 
  client.close() 
 
default_dag_args = { 
    # Setting start date as yesterday starts the DAG immediately when it is 
    # detected in the Cloud Storage bucket. 
    'start_date':yesterday, 
    # To email on failure or retry set 'email' arg to your email and enable 
    # emailing here. 
    'email_on_failure': False, 
    'email_on_retry': False, 
    # If a task fails, retry it once after waiting at least 5 minutes 
    'retries': 0, 
    #'retry_delay': datetime.timedelta(minutes=5), 
    'project_id': 'data-rubrics' 
} 
 
try: 
  # [START composer_quickstart_schedule] 
  with models.DAG( 
        'composer_testing00', 
        # Continue to run DAG once per day 
        schedule_interval=datetime.timedelta(days=1), 
        default_args=default_dag_args) as dag: 
    # [END composer_quickstart_schedule] 
 
  data_list = bigquery_get_data.BigQueryGetDataOperator(\ 
task_id='get_data_in_list_from_bq',\ 
dataset_id='testcomposer',\ # Name of the dataset which contains the table ( a BQ terminology) 
table_id='summarized_sample_T1' # Name of the BQ table you want to push into MongoDB 
) 
 
  op_push2mongo = PythonOperator(task_id='Push_to_MongoDB', python_callable=get_dlist, provide_context=True) 
  data_list >> op_push2mongo 
except Exception as e: 
  raise(e) 

编辑 2

    #-------- COPY PASTED BigQueryGetDataOperator SECTION: START -------------- 
    '''Source: https://airflow.readthedocs.io/en/stable/_modules/airflow/contrib/operators/bigquery_get_data.html#BigQueryGetDataOperator ''' 
 
    from airflow.contrib.hooks.bigquery_hook import BigQueryHook 
    from airflow.models import BaseOperator 
    from airflow.utils.decorators import apply_defaults 
 
 
    class BigQueryGetDataOperator(BaseOperator): 
        template_fields = ('dataset_id', 'table_id', 'max_results') 
        ui_color = '#e4f0e8' 
 
        @apply_defaults 
        def __init__(self, 
                     dataset_id, 
                     table_id, 
                     max_results='100', 
                     selected_fields=None, 
                     bigquery_conn_id='bigquery_default', 
                     delegate_to=None, 
                     *args, 
                     **kwargs): 
            super(BigQueryGetDataOperator, self).__init__(*args, **kwargs) 
            self.dataset_id = dataset_id 
            self.table_id = table_id 
            self.max_results = max_results 
            self.selected_fields = selected_fields 
            self.bigquery_conn_id = bigquery_conn_id 
            self.delegate_to = delegate_to 
 
        def execute(self, context): 
            self.log.info('Fetching Data from:') 
            self.log.info('Dataset: %s ; Table: %s ; Max Results: %s', 
                          self.dataset_id, self.table_id, self.max_results) 
 
            hook = BigQueryHook(bigquery_conn_id=self.bigquery_conn_id, 
                                delegate_to=self.delegate_to) 
 
            conn = hook.get_conn() 
            cursor = conn.cursor() 
 
    #----------------------- COPY PASTED SECTION: END ---------------- 
 
        # Trying to add to a MongoDB here itself - coed by GT 
        from pymongo import MongoClient 
        header = ['day', 'ticker','app_id','area', 'store_types', 'devices_in_store', 'devices_in_store_or_plot', 'matched_devices', \ 
'all_devices'] 
        client = MongoClient('35.237.46.25:27017') 
        db = client.test03 
        collection = db.advan_t1_sample_mongo00 
 
        response = cursor.get_tabledata(dataset_id=self.dataset_id, start_index=0, 
                                        table_id=self.table_id, 
                                        max_results='2', 
                                        selected_fields=self.selected_fields) 
        total_rows=int(response['totalRows']) 
        chunksize=100000 
        for chunk in range(0,total_rows,chunksize): 
          rows=[] 
          if chunk+chunksize<total_rows: 
            self.log.info("Extracting chunk %d to %d"%(chunk,chunk+chunksize)) 
            response = cursor.get_tabledata(dataset_id=self.dataset_id, start_index=chunk, 
                                        table_id=self.table_id, 
                                        max_results=str(chunksize), 
                                        selected_fields=self.selected_fields) 
 
            rows = response['rows'] 
 
            for row in rows: 
              onerow={} 
              for i,f in zip(range(len(row['f'])),row['f']): 
                onerow[header[i]] = f['v'] 
              collection.insert_one(onerow) 
            self.log.info("------------------------- Document size: %d --------------------"%(collection.find().count())) 
          else: 
            self.log.info("Extracting chunk %d to %d"%(chunk,total_rows)) 
            response = cursor.get_tabledata(dataset_id=self.dataset_id, start_index=chunk, 
                                        table_id=self.table_id, 
                                        max_results=total_rows, 
                                        selected_fields=self.selected_fields) 
 
            rows = response['rows'] 
 
            for row in rows: 
              onerow={} 
              for i,f in zip(range(len(row['f'])),row['f']): 
                onerow[header[i]] = f['v'] 
              collection.insert_one(onerow) 
            self.log.info("------------------------- Document size: %d --------------------"%(collection.find().count())) 
 
 
        self.log.info("Pushed into %s"%collection.name) 
 
        if total_rows == collection.find().count(): 
          self.log.info("Successfully pushed %d records into %s"%(total_rows,collection.name)) 
          return(1) 
        else: 
          self.log.warning("Push Failed! Total Rows: %d Document Size: %d"%(total_rows,collection.find().count())) 
          return(0) 
#        return rows 
 
    from airflow import models 
    from airflow.operators.python_operator import PythonOperator 
    from airflow.utils import trigger_rule 
    from airflow.contrib.operators import gcs_to_bq 
    from airflow.contrib.operators import bigquery_to_gcs 
    from airflow.contrib.operators import bigquery_operator 
    from airflow.contrib.operators import bigquery_get_data 
    from airflow.contrib.operators import MongoHook 
    def get_dlist(**kwargs): 
      import logging as log 
      #Import pymongo 
      from pymongo import MongoClient 
      #Pull the data saved in XCom 
      value = kwargs.get('task_instance').xcom_pull(task_ids='get_data_in_list_from_bq') 
 
      header = ['V1','V2'] 
      data=[] 
      for rows in value: 
        onerow={} 
        for i,f in zip(range(len(rows['f'])),rows['f']): 
          onerow[header[i]] = f['v'] 
        data.append(onerow) 
      log.info("Pulled...") 
      log.info(data) 
      log.info("Pushing into mongodb...") 
      client = MongoClient(localhost:27017) 
      db = client.test 
      collection = db.testingbq2mongo 
      collection.insert(data) 
      log.info("Written to mongoDB...") 
      client.close() 
 
    default_dag_args = { 
        # Setting start date as yesterday starts the DAG immediately when it is 
        # detected in the Cloud Storage bucket. 
        'start_date':yesterday, 
        # To email on failure or retry set 'email' arg to your email and enable 
        # emailing here. 
        'email_on_failure': False, 
        'email_on_retry': False, 
        # If a task fails, retry it once after waiting at least 5 minutes 
        'retries': 0, 
        #'retry_delay': datetime.timedelta(minutes=5), 
        'project_id': 'data-rubrics' 
    } 
 
    try: 
      # [START composer_quickstart_schedule] 
      with models.DAG( 
            'composer_testing00', 
            # Continue to run DAG once per day 
            schedule_interval=datetime.timedelta(days=1), 
            default_args=default_dag_args) as dag: 
        # [END composer_quickstart_schedule] 
 
      data_list = bigquery_get_data.BigQueryGetDataOperator(\ 
    task_id='get_data_in_list_from_bq',\ 
    dataset_id='testcomposer',\ # Name of the dataset which contains the table ( a BQ terminology) 
    table_id='summarized_sample_T1' # Name of the BQ table you want to push into MongoDB 
    ) 
 
      op_push2mongo = PythonOperator(task_id='Push_to_MongoDB', python_callable=get_dlist, provide_context=True) 
      data_list >> op_push2mongo 
    except Exception as e: 
      raise(e) 

请您参考如下方法:

在 XCOM 中传递大量数据被认为是一种反模式。我建议将数据从 BigQuery 写入持久存储服务(如 Cloud Storage)，然后从那里加载到 MongoDB。