Fix the bug of dataflow with ID #43

Author: zhanglongbin

data_celery/datasource/hive/tasks.py

@@ -12,6 +12,7 @@
                                ensure_directory_exists_remove, get_datasource_csg_hub_server_dir)
 from data_celery.mongo_tools.tools import insert_datasource_run_task_log_info,insert_datasource_run_task_log_error
 from data_server.datasource.services.datasource import get_datasource_connector
+from data_server.datasource.schemas import DataSourceCreate
 from data_engine.exporter.load import load_exporter
 from pathlib import Path
 import pandas as pd
@@ -103,13 +104,32 @@ def collection_hive_task(task_uid: str,user_name: str,user_token: str):
             max_line = extra_config["max_line_json"]
         if use_type == "sql":
             if use_sql:
-                connector = get_datasource_connector(collection_task.datasource)
-                if not connector.test_connection():
+                try:
+                    # 将数据库对象转换为 DataSourceCreate 对象
+                    datasource_create = DataSourceCreate(
+                        name=collection_task.datasource.name,
+                        des=collection_task.datasource.des,
+                        source_type=collection_task.datasource.source_type,
+                        host=collection_task.datasource.host,
+                        port=collection_task.datasource.port,
+                        username=collection_task.datasource.username,
+                        password=collection_task.datasource.password,
+                        database=collection_task.datasource.database,
+                        auth_type=collection_task.datasource.auth_type
+                    )
+                    connector = get_datasource_connector(datasource_create)
+                    test_result = connector.test_connection()
+                    if not test_result or not test_result.get("success", False):
+                        collection_task.task_status = DataSourceTaskStatusEnum.ERROR.value
+                        error_msg = test_result.get("message", "Connection failed") if test_result else "Connection test returned None"
+                        insert_datasource_run_task_log_error(task_uid, f"Task with UID {task_uid} failed to connect to the database: {error_msg}")
+                        return False
+                    get_table_dataset_by_sql(connector, task_uid, use_sql, db_session, collection_task,
+                                             datasource_temp_parquet_dir, max_line=max_line)
+                except Exception as e:
                     collection_task.task_status = DataSourceTaskStatusEnum.ERROR.value
-                    insert_datasource_run_task_log_error(task_uid, f"Task with UID {task_uid} failed to connect to the database.")
+                    insert_datasource_run_task_log_error(task_uid, f"Error occurred while executing the task: {str(e)}")
                     return False
-                get_table_dataset_by_sql(connector, task_uid, use_sql, db_session, collection_task,
-                                         datasource_temp_parquet_dir, max_line=max_line)
                 upload_path = datasource_temp_parquet_dir.join('run_sql')
                 upload_to_csg_hub_server(csg_hub_dataset_id,
                                          csg_hub_dataset_name,
@@ -125,14 +145,34 @@ def collection_hive_task(task_uid: str,user_name: str,user_token: str):
                 source = hive_config["source"]
                 total_count = 0
                 records_count = 0
-                connector = get_datasource_connector(collection_task.datasource)
-                if not connector.test_connection():
+                try:
+                    # 将数据库对象转换为 DataSourceCreate 对象
+                    datasource_create = DataSourceCreate(
+                        name=collection_task.datasource.name,
+                        des=collection_task.datasource.des,
+                        source_type=collection_task.datasource.source_type,
+                        host=collection_task.datasource.host,
+                        port=collection_task.datasource.port,
+                        username=collection_task.datasource.username,
+                        password=collection_task.datasource.password,
+                        database=collection_task.datasource.database,
+                        auth_type=collection_task.datasource.auth_type
+                    )
+                    connector = get_datasource_connector(datasource_create)
+                    test_result = connector.test_connection()
+                    if not test_result or not test_result.get("success", False):
+                        collection_task.task_status = DataSourceTaskStatusEnum.ERROR.value
+                        error_msg = test_result.get("message", "Connection failed") if test_result else "Connection test returned None"
+                        insert_datasource_run_task_log_error(task_uid, f"Task with UID {task_uid} failed to connect to the database: {error_msg}")
+                        return False
+                    for table_name in source.keys():
+                        table_total = connector.get_table_total_count_hive(table_name)
+                        total_count += table_total
+                except Exception as e:
                     collection_task.task_status = DataSourceTaskStatusEnum.ERROR.value
-                    insert_datasource_run_task_log_error(task_uid, f"Task with UID {task_uid} failed to connect to the database.")
+                    insert_datasource_run_task_log_error(task_uid, f"Error occurred while executing the task: {str(e)}")
                     return False
-                for table_name in source.keys():
-                    table_total = connector.get_table_total_count_hive(table_name)
-                    total_count += table_total
+                
                 collection_task.total_count = total_count
                 collection_task.records_count = records_count
                 db_session.commit()
@@ -165,8 +205,14 @@ def collection_hive_task(task_uid: str,user_name: str,user_token: str):
     except Exception as e:
         if collection_task:
             collection_task.task_status = DataSourceTaskStatusEnum.ERROR.value
+        error_type = type(e).__name__
+        error_msg = str(e)
+        error_traceback = traceback.format_exc()
+        logger.error(f"Task {task_uid} error: {error_type}: {error_msg}")
+        logger.error(f"Full traceback:\n{error_traceback}")
         traceback.print_exc()
-        insert_datasource_run_task_log_error(task_uid, f"Error occurred while executing the task: {e}")
+        insert_datasource_run_task_log_error(task_uid, f"Error occurred while executing the task: {error_type}: {error_msg}")
+        insert_datasource_run_task_log_error(task_uid, f"Traceback: {error_traceback}")
         return False
     finally:
         if collection_task:

data_server/api/endpoints/datasource.py

@@ -1,4 +1,5 @@
 import datetime
+import asyncio
 
 from fastapi import FastAPI, APIRouter, HTTPException, status, Header, Depends,Body
 from sqlalchemy import func
@@ -93,9 +94,25 @@ async def create_datasource(datasource: DataSourceCreate, db: Session = Depends(
         if datasource.source_type not in [item.value for item in DataSourceTypeEnum]:
             return response_fail(msg="不支持的数据源类型")
         # user_id = 54
+        
+        # 处理分支信息：修正前端传递的分支信息
+        # 前端可能将用户填写的分支名错误地放在了 csg_hub_dataset_name 中
+        if datasource.extra_config is None:
+            datasource.extra_config = {}
+        
+        current_branch = datasource.extra_config.get("csg_hub_dataset_default_branch", "")
+        dataset_name = datasource.extra_config.get("csg_hub_dataset_name", "")
+        dataset_id = datasource.extra_config.get("csg_hub_dataset_id", "")
+        
+        # 如果用户选择了数据流向，且分支是 main，但 dataset_name 有值，使用 dataset_name 作为分支
+        if dataset_id and current_branch == "main" and dataset_name and dataset_name != "main" and dataset_name.strip():
+            datasource.extra_config["csg_hub_dataset_default_branch"] = dataset_name
 
         connector = get_datasource_connector(datasource)
-        if not connector.test_connection():
+        # Run the synchronized test_connection method in the thread pool
+        loop = asyncio.get_event_loop()
+        test_result = await loop.run_in_executor(None, connector.test_connection)
+        if not test_result.get('success', False):
             datasource.source_status = DataSourceStatusEnum.INACTIVE.value
         else:
             if datasource.is_run:
@@ -104,7 +121,7 @@ async def create_datasource(datasource: DataSourceCreate, db: Session = Depends(
                 datasource.source_status = DataSourceStatusEnum.WAITING.value
         if not user_id:
             return response_fail(msg="用户ID不能为空")
-        data_source_id = create_data_source(connector.test_connection(), db, datasource, int(user_id), user_name,
+        data_source_id = create_data_source(test_result, db, datasource, int(user_id), user_name,
                                             user_token)
         return response_success(data=data_source_id)
     except Exception as e:
@@ -145,7 +162,10 @@ async def test_datasource_connection(datasource: DataSourceCreate):
 
     try:
         connector = get_datasource_connector(datasource)
-        return response_success(data=connector.test_connection())
+        # Run the synchronized test_connection method in the thread pool to avoid blocking the event loop
+        loop = asyncio.get_event_loop()
+        result = await loop.run_in_executor(None, connector.test_connection)
+        return response_success(data=result)
     except Exception as e:
         logger.error(f"test_datasource_connection: {str(e)}")
         return response_fail(msg=f"测试连接失败:{str(e)}")
@@ -154,6 +174,16 @@ async def test_datasource_connection(datasource: DataSourceCreate):
 @router.put("/datasource/edit/{datasource_id}", response_model=dict)
 async def update_datasource(datasource_id: int, datasource: DataSourceUpdate, db: Session = Depends(get_sync_session)):
     try:
+        # 处理分支信息：修正前端传递的分支信息（与创建接口相同的逻辑）
+        if datasource.extra_config is not None:
+            current_branch = datasource.extra_config.get("csg_hub_dataset_default_branch", "")
+            dataset_name = datasource.extra_config.get("csg_hub_dataset_name", "")
+            dataset_id = datasource.extra_config.get("csg_hub_dataset_id", "")
+            
+            # 如果用户选择了数据流向，且分支是 main，但 dataset_name 有值，使用 dataset_name 作为分支
+            if dataset_id and current_branch == "main" and dataset_name and dataset_name != "main" and dataset_name.strip():
+                datasource.extra_config["csg_hub_dataset_default_branch"] = dataset_name
+        
         data_source = update_data_source(db, datasource_id, datasource)
         if not data_source:
             return response_fail(msg="更新失败")
@@ -213,12 +243,14 @@ async def get_datasource_tables(datasource: DataSourceCreate):
     try:
         # if datasource.source_type == DataSourceTypeEnum.MONGODB.value:
 
-
         connector = get_datasource_connector(datasource)
-        if not connector.test_connection():
+        # test_the_connection_in_the_thread_pool
+        loop = asyncio.get_event_loop()
+        test_result = await loop.run_in_executor(None, connector.test_connection)
+        if not test_result.get('success', False):
             return response_fail(msg="数据源连接失败")
 
-        tables = connector.get_tables()
+        tables = await loop.run_in_executor(None, connector.get_tables)
         return response_success(data=tables)
     except Exception as e:
         logger.error(f"获取表列表失败: {str(e)}")
@@ -233,17 +265,16 @@ async def get_datasource_table_columns(datasource: DataSourceCreate, table_name:
             return response_fail(msg="MongoDB不支持获取表和字段列表")
 
         connector = get_datasource_connector(datasource)
-        if not connector.test_connection():
+        loop = asyncio.get_event_loop()
+        test_result = await loop.run_in_executor(None, connector.test_connection)
+        if not test_result.get('success', False):
             return response_fail(msg="数据源连接失败")
 
-        columns = connector.get_table_columns(table_name)
+        columns = await loop.run_in_executor(None, connector.get_table_columns, table_name)
         return response_success(data=columns)
     except Exception as e:
         logger.error(f"获取表字段失败: {str(e)}")
         return response_fail(msg=f"获取表字段失败: {str(e)}")
-    except Exception as e:
-        logger.error(f"获取表字段失败: {str(e)}")
-        return response_fail(msg=f"获取字段失败: {str(e)}")
 
 
 @router.get("/datasource/info", response_model=dict)
@@ -266,8 +297,10 @@ async def get_datasource_tables_and_columns(datasource: DataSourceCreate):
             return response_fail(msg="MongoDB不支持获取表和字段列表")
 
         connector = get_datasource_connector(datasource)
-        if not connector.test_connection():
-            return response_fail(msg="数据源连接失败")
+        test_result = connector.test_connection()
+        if not test_result or not test_result.get("success", False):
+            error_msg = test_result.get("message", "Connection failed") if test_result else "Connection test returned None"
+            return response_fail(msg=f"数据源连接失败: {error_msg}")
 
         tables_and_columns = connector.get_tables_and_columns()
         return response_success(data=tables_and_columns)

data_server/datasource/DatasourceManager.py

@@ -36,6 +36,7 @@ def create_data_source(is_connection: bool, db_session: Session, datasource: Dat
         int: ID of the created data source
     """
     # create db model
+    extra_config = datasource.extra_config or {}
     data_source_db = DataSource(name=datasource.name,
                                 des=datasource.des,
                                 source_type=datasource.source_type,
@@ -44,8 +45,9 @@ def create_data_source(is_connection: bool, db_session: Session, datasource: Dat
                                 username=datasource.username,
                                 password=datasource.password,
                                 database=datasource.database,
+                                auth_type=datasource.auth_type,
                                 task_run_time=datasource.task_run_time,
-                                extra_config=json.dumps(datasource.extra_config, ensure_ascii=False, indent=4))
+                                extra_config=json.dumps(extra_config, ensure_ascii=False, indent=4))
     data_source_db.source_status = datasource.source_status
     data_source_db.owner_id = user_id
     db_session.add(data_source_db)
@@ -140,6 +142,8 @@ def update_data_source(db_session: Session, data_source_id: int, update_data: Da
         data_source.password = update_data.password
     if update_data.database is not None:
         data_source.database = update_data.database
+    if update_data.auth_type is not None:
+        data_source.auth_type = update_data.auth_type
     if update_data.extra_config is not None:
         data_source.extra_config = json.dumps(update_data.extra_config, ensure_ascii=False, indent=4)
     db_session.commit()