Skip to content

Commit a1d53a3

Browse files
Chenglong-MSChenglong-MS
committed
some fixes around data loader
1 parent 16c599f commit a1d53a3

File tree

9 files changed

+19
-40
lines changed

9 files changed

+19
-40
lines changed

‎py-src/data_formulator/agents/__init__.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,5 +20,5 @@
2020
"SQLDataRecAgent",
2121
"DataLoadAgent",
2222
"SortDataAgent",
23-
"DataCleanAgent"
23+
"DataCleanAgent",
2424
]

‎py-src/data_formulator/agents/agent_query_completion.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,7 @@ def __init__(self, client):
5454

5555
def run(self, data_source_metadata, query):
5656

57-
user_query = f"[DATA SOURCE]\n\n{json.dumps(data_source_metadata, indent=2)}\n\n[USER INPUTS]\n\n{query}\n\n[REASONING]\n"
57+
user_query = f"[DATA SOURCE]\n\n{json.dumps(data_source_metadata, indent=2)}\n\n[USER INPUTS]\n\n{query}\n\n"
5858

5959
logger.info(user_query)
6060

@@ -63,11 +63,11 @@ def run(self, data_source_metadata, query):
6363

6464
###### the part that calls open_ai
6565
response = self.client.get_completion(messages = messages)
66-
response_content = '[REASONING]\n' + response.choices[0].message.content
66+
response_content = response.choices[0].message.content
6767

6868
logger.info(f"=== query completion output ===>\n{response_content}\n")
6969

70-
reasoning = extract_json_objects(response_content.split("[REASONING]")[1].split("[QUERY]")[0].strip())[0]
70+
reasoning = extract_json_objects(response_content.split("[QUERY]")[0].strip())[0]
7171
output_query = response_content.split("[QUERY]")[1].strip()
7272

7373
# Extract the query by removing the language markers

‎py-src/data_formulator/data_loader/azure_blob_data_loader.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -369,4 +369,4 @@ def ingest_data_from_query(self, query: str, name_as: str):
369369
# Execute the query and get results as a DataFrame
370370
df = self.duck_db_conn.execute(query).df()
371371
# Use the base class's method to ingest the DataFrame
372-
self.ingest_df_to_duckdb(df, name_as)
372+
self.ingest_df_to_duckdb(df, sanitize_table_name(name_as))

‎py-src/data_formulator/data_loader/mssql_data_loader.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -445,7 +445,7 @@ def ingest_data_from_query(self, query: str, name_as: str) -> pd.DataFrame:
445445
try:
446446
df = self._execute_query(query)
447447
# Use the base class's method to ingest the DataFrame
448-
self.ingest_df_to_duckdb(df, name_as)
448+
self.ingest_df_to_duckdb(df, sanitize_table_name(name_as))
449449
log.info(f"Successfully ingested {len(df)} rows from custom query to {name_as}")
450450
return df
451451
except Exception as e:

‎py-src/data_formulator/data_loader/mysql_data_loader.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,9 @@ def __init__(self, params: Dict[str, Any], duck_db_conn: duckdb.DuckDBPyConnecti
6363
try:
6464
self.duck_db_conn.execute("DETACH mysqldb;")
6565
except:
66-
pass # Ignore if mysqldb doesn't exist # Register MySQL connection
66+
pass # Ignore if mysqldb doesn't exist
67+
68+
# Register MySQL connection
6769
self.duck_db_conn.execute(f"ATTACH '{attach_string}' AS mysqldb (TYPE mysql);")
6870

6971
def list_tables(self, table_filter: str = None):
@@ -129,4 +131,4 @@ def ingest_data_from_query(self, query: str, name_as: str) -> pd.DataFrame:
129131
# Execute the query and get results as a DataFrame
130132
df = self.duck_db_conn.execute(query).df()
131133
# Use the base class's method to ingest the DataFrame
132-
self.ingest_df_to_duckdb(df, name_as)
134+
self.ingest_df_to_duckdb(df, sanitize_table_name(name_as))

‎py-src/data_formulator/data_loader/postgresql_data_loader.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,5 +128,5 @@ def ingest_data_from_query(self, query: str, name_as: str) -> pd.DataFrame:
128128
# Execute the query and get results as a DataFrame
129129
df = self.duck_db_conn.execute(query).df()
130130
# Use the base class's method to ingest the DataFrame
131-
self.ingest_df_to_duckdb(df, name_as)
131+
self.ingest_df_to_duckdb(df, sanitize_table_name(name_as))
132132
return df

‎py-src/data_formulator/data_loader/s3_data_loader.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,4 +203,4 @@ def ingest_data_from_query(self, query: str, name_as: str):
203203
# Execute the query and get results as a DataFrame
204204
df = self.duck_db_conn.execute(query).df()
205205
# Use the base class's method to ingest the DataFrame
206-
self.ingest_df_to_duckdb(df, name_as)
206+
self.ingest_df_to_duckdb(df, sanitize_table_name(name_as))

‎src/views/DBTableManager.tsx

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1396,21 +1396,21 @@ export const DataQueryForm: React.FC<{
13961396
}}
13971397
/>
13981398
</Box>
1399-
{queryResult?.status === "error" && <Box sx={{display: "flex", flexDirection: "row", gap: 1, alignItems: "center"}}>
1399+
{queryResult?.status === "error" && <Box sx={{display: "flex", flexDirection: "row", gap: 1, alignItems: "center", overflow: "auto"}}>
14001400
<Typography variant="body2" sx={{color: "text.secondary", fontSize: 11, backgroundColor: "rgba(255, 0, 0, 0.1)", p: 0.5, borderRadius: 0.5}}>
14011401
{queryResult?.message}
14021402
</Typography>
1403-
<Button variant="outlined" color="primary" size="small" sx={{textTransform: "none", height: 24, ml: 1, minWidth: 120}}
1404-
startIcon={<PrecisionManufacturingIcon />} onClick={() => aiCompleteQuery(queryResult.code + "\n error:" + queryResult.message)}>
1405-
help me fix it
1406-
</Button>
14071403
</Box>}
1408-
<Box sx={{display: "flex", flexDirection: "row", gap: 1, justifyContent: "space-between"}}>
1404+
<Box sx={{display: "flex", flexDirection: "row", gap: 1, justifyContent: "flex-end"}}>
14091405
<Button variant="outlined" color="primary" size="small" sx={{textTransform: "none"}} disabled={queryResult?.status === "error"}
14101406
startIcon={<PrecisionManufacturingIcon />} onClick={() => aiCompleteQuery(query)}>
14111407
help me complete the query from selected tables
14121408
</Button>
1413-
<Button variant="contained" color="primary" size="small" sx={{textTransform: "none", width: 80}}
1409+
{queryResult?.status === "error" && <Button variant="contained" color="primary" size="small" sx={{textTransform: "none",minWidth: 120}}
1410+
startIcon={<PrecisionManufacturingIcon />} onClick={() => aiCompleteQuery(queryResult.code + "\n error:" + queryResult.message)}>
1411+
help me fix the error
1412+
</Button>}
1413+
<Button variant="contained" color="primary" size="small" sx={{textTransform: "none", ml: 'auto', width: 80}}
14141414
onClick={() => handleViewQuerySample(query)}>
14151415
run query
14161416
</Button>

‎src/views/TableSelectionView.tsx

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -68,25 +68,6 @@ function a11yProps(index: number) {
6868
};
6969
}
7070

71-
function syncTableWithServerDB(table: DictTable) {
72-
console.log(`syncing table ${table.id} with server db`)
73-
const formData = new FormData();
74-
formData.append('table_name', table.id);
75-
formData.append('raw_data', JSON.stringify(table.rows));
76-
77-
fetch(`${getUrls().CREATE_TABLE}`, {
78-
method: 'POST',
79-
body: formData
80-
}).then((response) => {
81-
return response.json();
82-
}).then((data) => {
83-
console.log(data);
84-
}).catch((error) => {
85-
console.log(error);
86-
});
87-
}
88-
89-
9071
export interface TableChallenges {
9172
name: string;
9273
challenges: { text: string; difficulty: 'easy' | 'medium' | 'hard'; }[];
@@ -237,7 +218,6 @@ export const TableSelectionDialog: React.FC<{ buttonElement: any }> = function T
237218
.then((text) => {
238219
let fullTable = createTableFromFromObjectArray(tableChallenges.table.id, JSON.parse(text), true);
239220
if (fullTable) {
240-
syncTableWithServerDB(fullTable);
241221
dispatch(dfActions.loadTable(fullTable));
242222
dispatch(fetchFieldSemanticType(fullTable));
243223
dispatch(dfActions.addChallenges({
@@ -332,7 +312,6 @@ export const TableUploadDialog: React.FC<TableUploadDialogProps> = ({ buttonElem
332312
if (arrayBuffer) {
333313
let tables = loadBinaryDataWrapper(uniqueName, arrayBuffer);
334314
for (let table of tables) {
335-
syncTableWithServerDB(table);
336315
dispatch(dfActions.loadTable(table));
337316
dispatch(fetchFieldSemanticType(table));
338317
}
@@ -426,7 +405,6 @@ export const TableURLDialog: React.FC<TableURLDialogProps> = ({ buttonElement, d
426405
}
427406

428407
if (table) {
429-
syncTableWithServerDB(table);
430408
dispatch(dfActions.loadTable(table));
431409
dispatch(fetchFieldSemanticType(table));
432410
}
@@ -523,7 +501,6 @@ export const TableCopyDialogV2: React.FC<TableCopyDialogProps> = ({ buttonElemen
523501
table = createTableFromText(uniqueName, tableStr);
524502
}
525503
if (table) {
526-
syncTableWithServerDB(table);
527504
dispatch(dfActions.loadTable(table));
528505
dispatch(fetchFieldSemanticType(table));
529506
}

0 commit comments

Comments
 (0)