Skip to content

Commit 9acd002

Browse files
cgivrecgivre
committed
Fix state
1 parent efc2d12 commit 9acd002

29 files changed

Lines changed: 2093 additions & 91 deletions

‎exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/SqlTranspiler.java‎

Lines changed: 55 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,16 @@
1717
*/
1818
package org.apache.drill.exec.server.rest;
1919

20-
import java.io.IOException;
21-
import java.io.InputStream;
22-
import java.nio.charset.StandardCharsets;
23-
2420
import org.graalvm.polyglot.Context;
2521
import org.graalvm.polyglot.Source;
2622
import org.graalvm.polyglot.Value;
2723
import org.slf4j.Logger;
2824
import org.slf4j.LoggerFactory;
2925

26+
import java.io.IOException;
27+
import java.io.InputStream;
28+
import java.nio.charset.StandardCharsets;
29+
3030
/**
3131
* Singleton service that wraps a GraalPy context with sqlglot
3232
* for SQL dialect transpilation. Thread-safe via synchronized access
@@ -49,6 +49,8 @@ public final class SqlTranspiler {
4949
private Context pythonContext;
5050
private Value transpileFunc;
5151
private Value jsonParseFunc;
52+
private Value convertDataTypeFunc;
53+
private Value formatSqlFunc;
5254
private boolean available;
5355

5456
/**
@@ -77,6 +79,8 @@ private SqlTranspiler() {
7779

7880
transpileFunc = pythonContext.getBindings("python").getMember("transpile_sql");
7981
jsonParseFunc = pythonContext.getBindings("python").getMember("json_parse");
82+
convertDataTypeFunc = pythonContext.getBindings("python").getMember("convert_data_type_raw");
83+
formatSqlFunc = pythonContext.getBindings("python").getMember("format_sql");
8084
available = true;
8185
logger.info("SqlTranspiler initialized successfully with GraalPy + sqlglot");
8286
} catch (Exception e) {
@@ -104,29 +108,66 @@ public synchronized String transpile(String sql, String sourceDialect, String ta
104108
*/
105109
public synchronized String transpile(String sql, String sourceDialect,
106110
String targetDialect, String schemasJson) {
107-
System.out.println("=== TRANSPILE INPUT ===");
108-
System.out.println("Available: " + available);
109-
System.out.println("SQL: " + sql);
110-
System.out.println("Source dialect: " + sourceDialect);
111-
System.out.println("Target dialect: " + targetDialect);
112-
System.out.println("Schemas JSON: " + schemasJson);
113-
System.out.println("======================");
114111
if (!available) {
115112
return sql;
116113
}
117114
try {
118115
Value schemas = jsonParseFunc.execute(schemasJson != null ? schemasJson : "[]");
119116
String result = transpileFunc.execute(sql, sourceDialect, targetDialect, schemas).asString();
120-
System.out.println("=== TRANSPILE OUTPUT ===");
121-
System.out.println("Result: " + result);
122-
System.out.println("========================");
123117
return result;
124118
} catch (Exception e) {
125119
logger.debug("SQL transpilation failed, returning original SQL: {}", e.getMessage());
126120
return sql;
127121
}
128122
}
129123

124+
/**
125+
* Convert a column's data type in a SQL query using sqlglot AST manipulation.
126+
*
127+
* @param sql the SQL query
128+
* @param columnName the column to convert
129+
* @param dataType the target SQL data type (e.g. "INTEGER", "VARCHAR")
130+
* @param columnsJson optional JSON object mapping column names to types (for star queries)
131+
* @return the transformed SQL, or null if conversion fails
132+
*/
133+
public synchronized String convertDataType(String sql, String columnName,
134+
String dataType, String columnsJson) {
135+
if (!available) {
136+
logger.warn("convertDataType called but transpiler is not available");
137+
return null;
138+
}
139+
try {
140+
String result;
141+
if (columnsJson != null && !columnsJson.isEmpty()) {
142+
result = convertDataTypeFunc.execute(sql, columnName, dataType, columnsJson).asString();
143+
} else {
144+
result = convertDataTypeFunc.execute(sql, columnName, dataType).asString();
145+
}
146+
return result;
147+
} catch (Exception e) {
148+
logger.warn("Data type conversion failed: {}", e.getMessage(), e);
149+
return null;
150+
}
151+
}
152+
153+
/**
154+
* Pretty-print a SQL string using sqlglot.
155+
*
156+
* @param sql the SQL string to format
157+
* @return the formatted SQL, or the original if formatting fails
158+
*/
159+
public synchronized String formatSql(String sql) {
160+
if (!available) {
161+
return sql;
162+
}
163+
try {
164+
return formatSqlFunc.execute(sql).asString();
165+
} catch (Exception e) {
166+
logger.debug("SQL formatting failed, returning original: {}", e.getMessage());
167+
return sql;
168+
}
169+
}
170+
130171
/**
131172
* Returns whether the transpiler is available (GraalPy + sqlglot initialized successfully).
132173
*/

‎exec/java-exec/src/main/java/org/apache/drill/exec/server/rest/TranspileResources.java‎

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,10 +83,19 @@ public static class TranspileResponse {
8383
@JsonProperty
8484
public boolean success;
8585

86+
@JsonProperty
87+
public String formattedOriginal;
88+
8689
public TranspileResponse(String sql, boolean success) {
8790
this.sql = sql;
8891
this.success = success;
8992
}
93+
94+
public TranspileResponse(String sql, boolean success, String formattedOriginal) {
95+
this.sql = sql;
96+
this.success = success;
97+
this.formattedOriginal = formattedOriginal;
98+
}
9099
}
91100

92101
// ==================== Endpoints ====================
@@ -118,6 +127,100 @@ public Response transpile(TranspileRequest request) {
118127
return Response.ok(new TranspileResponse(result, true)).build();
119128
}
120129

130+
// ==================== Convert Data Type ====================
131+
132+
public static class ConvertDataTypeRequest {
133+
@JsonProperty
134+
public String sql;
135+
136+
@JsonProperty
137+
public String columnName;
138+
139+
@JsonProperty
140+
public String dataType;
141+
142+
@JsonProperty
143+
public Map<String, Object> columns;
144+
145+
public ConvertDataTypeRequest() {
146+
}
147+
148+
@JsonCreator
149+
public ConvertDataTypeRequest(
150+
@JsonProperty("sql") String sql,
151+
@JsonProperty("columnName") String columnName,
152+
@JsonProperty("dataType") String dataType,
153+
@JsonProperty("columns") Map<String, Object> columns) {
154+
this.sql = sql;
155+
this.columnName = columnName;
156+
this.dataType = dataType;
157+
this.columns = columns;
158+
}
159+
}
160+
161+
@POST
162+
@Path("/convert-type")
163+
@Consumes(MediaType.APPLICATION_JSON)
164+
@Produces(MediaType.APPLICATION_JSON)
165+
@Operation(summary = "Convert column data type",
166+
description = "Wraps a column in a CAST expression using sqlglot AST manipulation")
167+
public Response convertDataType(ConvertDataTypeRequest request) {
168+
if (request.sql == null || request.sql.trim().isEmpty()) {
169+
return Response.ok(new TranspileResponse("", true)).build();
170+
}
171+
if (request.columnName == null || request.dataType == null) {
172+
return Response.status(Response.Status.BAD_REQUEST)
173+
.entity(new TranspileResponse(request.sql, false)).build();
174+
}
175+
176+
SqlTranspiler transpiler = SqlTranspiler.getInstance();
177+
String columnsJson = null;
178+
if (request.columns != null && !request.columns.isEmpty()) {
179+
try {
180+
columnsJson = new ObjectMapper().writeValueAsString(request.columns);
181+
} catch (Exception e) {
182+
columnsJson = null;
183+
}
184+
}
185+
186+
String result = transpiler.convertDataType(
187+
request.sql, request.columnName, request.dataType, columnsJson);
188+
if (result == null) {
189+
return Response.ok(new TranspileResponse(request.sql, false)).build();
190+
}
191+
String formattedOriginal = transpiler.formatSql(request.sql);
192+
return Response.ok(new TranspileResponse(result, true, formattedOriginal)).build();
193+
}
194+
195+
// ==================== Format SQL ====================
196+
197+
public static class FormatRequest {
198+
@JsonProperty
199+
public String sql;
200+
201+
public FormatRequest() {
202+
}
203+
204+
@JsonCreator
205+
public FormatRequest(@JsonProperty("sql") String sql) {
206+
this.sql = sql;
207+
}
208+
}
209+
210+
@POST
211+
@Path("/format")
212+
@Consumes(MediaType.APPLICATION_JSON)
213+
@Produces(MediaType.APPLICATION_JSON)
214+
@Operation(summary = "Format SQL",
215+
description = "Pretty-prints a SQL string using sqlglot")
216+
public Response formatSql(FormatRequest request) {
217+
if (request.sql == null || request.sql.trim().isEmpty()) {
218+
return Response.ok(new TranspileResponse("", true)).build();
219+
}
220+
String result = SqlTranspiler.getInstance().formatSql(request.sql);
221+
return Response.ok(new TranspileResponse(result, true)).build();
222+
}
223+
121224
@GET
122225
@Path("/status")
123226
@Produces(MediaType.APPLICATION_JSON)

0 commit comments

Comments
 (0)