如何让Celery多Worker共享数据库游标实现并行跨库数据迁移？

问题背景 #

目前使用单个Celery Worker结合SQLAlchemy实现跨库数据迁移：从源数据库指定表读取数据，写入目标数据库指定表。希望改为多Worker并行处理，将总数据量拆分给多个Worker处理。已知可通过拆分数据起止行实现，但想了解是否能让多Worker共享proxy.fetchmany()调用直至数据耗尽。

以下是单Worker运行的示例代码：

logging.info("Connecting to Database")
connection_url = URL.create("mssql+pyodbc", query={"odbc_connect": self.SOURCE_CONNECT_STRING})
source_eng = sqlalchemy.create_engine(connection_url)
source_connection = source_eng.connect()

# Select all rows
query = f"SELECT * FROM {sourceTable} WITH (NOLOCK)"

stmt = sqlalchemy.text(query)

proxy = source_connection.execution_options(stream_results=True).execute(stmt)

column_names = list(proxy.keys())
logging.info(column_names)

insert_stmt = f"""INSERT INTO {destTable} ({','.join(column_names)}) VALUES ({','.join(['?']*len(column_names))});"""

if ( hasIdentity ):
  insert_stmt = f"""
    SET IDENTITY_INSERT {destTable} ON;
    SET NOCOUNT ON;
    {insert_stmt}
    SET NOCOUNT OFF;
    SET IDENTITY_INSERT {destTable} OFF"""

completed = 0
while 'batch not empty':  # equivalent of 'while True'
  batch = proxy.fetchmany(self.INITIAL_BATCH_SIZE)
  if not batch:
    break

  # Strip whitespace from batch
  batch = [self._strip_tuple(i) for i in batch]
  
  logging.info(f"Inserting {str(len(batch))} rows.")
  with pyodbc.connect(self.DEST_CONNECT_STRING) as conn:
    with conn.cursor() as cursor:
      cursor.fast_executemany = True
      # Send pyodbc call
      logging.info(insert_stmt)
      cursor.executemany(insert_stmt, list(batch))
      conn.commit()
      completed += len(batch)
  logging.info(f"Completed inserting {completed} rows into {destTable}.")

proxy.close()

核心结论 #

无法让多Worker直接共享同一个proxy对象。原因在于：

• proxy是SQLAlchemy绑定到单个数据库连接和线程的结果集对象，Celery Worker为独立进程/线程，跨进程无法共享该对象的状态。
• 数据库游标本身是会话绑定的，跨进程访问会引发连接异常、数据不一致等问题。

可行的并行处理方案 #

方案1：按数据范围拆分任务（Worker独立查询） #

通过预计算数据范围，给每个Worker分配独立的查询区间，Worker自行从源库读取对应区间的数据并写入目标库。

实现思路：

1. 主进程先查询源表总条数：SELECT COUNT(*) FROM {sourceTable} WITH (NOLOCK)。
2. 根据Worker数量，将总数据量拆分为N个连续区间（比如总条数1000，3个Worker则分配0-333、334-666、667-1000）。
3. 为每个区间生成Celery任务，传递起止位置给Worker。
4. Worker使用有序字段（如主键ID）执行分页查询，避免数据重复或遗漏：

   query = f"""
   SELECT * FROM {sourceTable} WITH (NOLOCK)
   ORDER BY id
   OFFSET ? ROWS FETCH NEXT ? ROWS ONLY
   """
   

优缺点：

• 优点：Worker独立访问源库，无单点瓶颈，适合超大数据量迁移；故障隔离性好，单个Worker失败不影响其他任务。
• 缺点：需要依赖有序字段实现分页；需预计算总条数，若源表数据实时变化可能导致区间偏差。

方案2：主进程分发批次数据（Worker仅负责写入） #

主进程作为生产者，持续通过proxy.fetchmany()读取批量数据，将每个批次封装为Celery任务发送到队列，Worker从队列中取出批次数据并写入目标库。

实现修改：

1. 主进程保留源库连接和proxy读取逻辑，移除写入代码，改为发送任务：

   # 主进程代码片段
   completed = 0
   while True:
     batch = proxy.fetchmany(self.INITIAL_BATCH_SIZE)
     if not batch:
       break
     # 发送迁移任务到Celery
     migrate_batch.delay(
         batch=batch,
         column_names=column_names,
         destTable=destTable,
         hasIdentity=hasIdentity,
         dest_connect_string=self.DEST_CONNECT_STRING
     )
     completed += len(batch)
     logging.info(f"Dispatched {completed} rows to workers.")
   proxy.close()
   

2. 定义Celery Worker任务：

   @app.task
   def migrate_batch(batch, column_names, destTable, hasIdentity, dest_connect_string):
       # 处理数据清洗
       cleaned_batch = [_strip_tuple(row) for row in batch]
       # 构造插入语句
       insert_stmt = f"""INSERT INTO {destTable} ({','.join(column_names)}) VALUES ({','.join(['?']*len(column_names))});"""
       if hasIdentity:
           insert_stmt = f"""
             SET IDENTITY_INSERT {destTable} ON;
             SET NOCOUNT ON;
             {insert_stmt}
             SET NOCOUNT OFF;
             SET IDENTITY_INSERT {destTable} OFF"""
       # 写入目标库
       with pyodbc.connect(dest_connect_string) as conn:
           with conn.cursor() as cursor:
               cursor.fast_executemany = True
               cursor.executemany(insert_stmt, cleaned_batch)
               conn.commit()
       return len(cleaned_batch)
   

优缺点：

• 优点：逻辑简单，无需预计算数据范围；Worker仅负责写入，无需访问源库，降低源库连接压力。
• 缺点：主进程成为单点瓶颈，若主进程故障会导致整个迁移中断；需控制任务发送速度，避免队列积压。

内容的提问来源于stack exchange，提问作者Optimus