chore(wren-ai-service): add correlation ID to langfuse

Canner · Oct 30, 2024 · 2b5e265 · 2b5e265
1 parent 1860cf8
commit 2b5e265
Show file tree

Hide file tree

Showing 4 changed files with 26 additions and 14 deletions.
diff --git a/wren-ai-service/src/pipelines/common.py b/wren-ai-service/src/pipelines/common.py
@@ -88,14 +88,14 @@ async def _check_if_sql_executable(
         project_id: str | None = None,
     ):
         async with aiohttp.ClientSession() as session:
-            status, _, error = await self._engine.execute_sql(
+            status, _, addition = await self._engine.execute_sql(
                 sql,
                 session,
                 project_id=project_id,
             )
 
         if not status:
-            logger.exception(f"SQL is not executable: {error}")
+            logger.exception(f"SQL is not executable: {addition["error_message"]}")
 
         return status
 
@@ -151,22 +151,24 @@ async def _task(result: Dict[str, str]):
             quoted_sql, no_error = add_quotes(result["sql"])
 
             if no_error:
-                status, _, error = await self._engine.execute_sql(
+                status, _, addition = await self._engine.execute_sql(
                     quoted_sql, session, project_id=project_id
                 )
 
                 if status:
                     valid_generation_results.append(
                         {
                             "sql": quoted_sql,
+                            "correlation_id": addition["correlation_id"],
                         }
                     )
                 else:
                     invalid_generation_results.append(
                         {
                             "sql": quoted_sql,
                             "type": "DRY_RUN",
-                            "error": error,
+                            "error": addition["error_message"],
+                            "correlation_id": addition["correlation_id"],
                         }
                     )
             else:

diff --git a/wren-ai-service/src/providers/engine/wren.py b/wren-ai-service/src/providers/engine/wren.py
@@ -49,11 +49,16 @@ async def execute_sql(
         ) as response:
             res = await response.json()
             if data := res.get("data"):
-                return True, data, None
+                return True, data, {
+                    "correlation_id": res.get("correlationId"),
+                }
             return (
                 False,
                 None,
-                res.get("errors", [{}])[0].get("message", "Unknown error"),
+                {
+                    "error_message": res.get("errors", [{}])[0].get("message", "Unknown error"),
+                    "correlation_id": res.get("extensions", {}).get("other", {}).get("correlationId"),
+                }
             )
 
 

diff --git a/wren-ui/src/apollo/server/resolvers/modelResolver.ts b/wren-ui/src/apollo/server/resolvers/modelResolver.ts
@@ -933,14 +933,13 @@ export class ModelResolver {
       ? await ctx.projectService.getProjectById(projectId)
       : await ctx.projectService.getCurrentProject();
     const { manifest } = await ctx.deployService.getLastDeployment(project.id);
-    const previewRes = await ctx.queryService.preview(sql, {
+    return await ctx.queryService.preview(sql, {
       project,
       limit: limit,
       modelingOnly: false,
       manifest,
       dryRun,
     });
-    return dryRun ? { dryRun: 'success' } : previewRes;
   }
 
   public async getNativeSql(

diff --git a/wren-ui/src/apollo/server/services/queryService.ts b/wren-ui/src/apollo/server/services/queryService.ts
@@ -28,8 +28,8 @@ export interface ColumnMetadata {
 export interface PreviewDataResponse {
   correlationId?: string;
   processTime?: string;
-  columns: ColumnMetadata[];
-  data: any[][];
+  columns?: ColumnMetadata[];
+  data?: any[][];
 }
 
 export interface DescribeStatementResponse {
@@ -117,8 +117,7 @@ export class QueryService implements IQueryService {
       this.checkDataSourceIsSupported(dataSource);
       logger.debug('Use ibis adaptor to preview');
       if (dryRun) {
-        await this.ibisDryRun(sql, dataSource, connectionInfo, mdl);
-        return true;
+        return await this.ibisDryRun(sql, dataSource, connectionInfo, mdl);
       } else {
         return await this.ibisQuery(sql, dataSource, connectionInfo, mdl, limit);
       }
@@ -178,7 +177,7 @@ export class QueryService implements IQueryService {
     dataSource: DataSourceName,
     connectionInfo: any,
     mdl: Manifest
-  ) {
+  ): Promise<PreviewDataResponse> {
     const event = TelemetryEvent.IBIS_DRY_RUN;
     try {
       const res = await this.ibisAdaptor.dryRun(sql, {
@@ -187,6 +186,9 @@ export class QueryService implements IQueryService {
         mdl,
       });
       this.sendIbisEvent(event, res, { dataSource, sql });
+      return {
+        correlationId: res.correlationId,
+      }
     } catch (err: any) {
       this.sendIbisFailedEvent(event, err, { dataSource, sql });
       throw err;
@@ -209,7 +211,11 @@ export class QueryService implements IQueryService {
         limit,
       });
       this.sendIbisEvent(event, res, { dataSource, sql });
-      return this.transformDataType(res);
+      const data = this.transformDataType(res);
+      return { 
+        correlationId: res.correlationId,
+        ...data
+      };
     } catch (err: any) {
       this.sendIbisFailedEvent(event, err, { dataSource, sql });
       throw err;