Upload app.py

app.py

@@ -13,8 +13,6 @@ class QueryRequest(BaseModel):
 @app.get("/")
 def home():
  return {"message": "SQL Generation Server is running"}
-
-@app.post("/generate")
 def generate(request: QueryRequest):
  try:
  text = request.text
@@ -24,14 +22,30 @@ def generate(request: QueryRequest):
  generated_text = output[0]['generated_text']
  sql_query = generated_text.split("SQL query:")[-1].strip()
  
- # Basic validation
- if not sql_query.lower().startswith(('select', 'show', 'describe')):
+
+ # Basic validation to ensure it's a valid SQL query
+ if not sql_query.lower().startswith(('select', 'show', 'describe', 'insert', 'update', 'delete')):
  raise ValueError("Generated text is not a valid SQL query")
  
+ # Further validation to ensure no additional text
+ sql_query = sql_query.split(';')[0].strip()
+
+ # Comprehensive list of SQL keywords
+ allowed_keywords = {
+ 'select', 'insert', 'update', 'delete', 'show', 'describe', 'from', 'where', 'and', 'or', 'like', 'limit', 'order by', 'group by', 'join', 'inner join', 'left join', 'right join', 'full join', 'on', 'using', 'union', 'union all', 'distinct', 'having', 'into', 'values', 'set', 'create', 'alter', 'drop', 'table', 'database', 'index', 'view', 'trigger', 'procedure', 'function', 'if', 'exists', 'primary key', 'foreign key', 'references', 'check', 'constraint', 'default', 'auto_increment', 'null', 'not null', 'in', 'is', 'is not', 'between', 'case', 'when', 'then', 'else', 'end', 'asc', 'desc', 'count', 'sum', 'avg', 'min', 'max', 'timestamp', 'date', 'time', 'varchar', 'char', 'int', 'integer', 'smallint', 'bigint', 'decimal', 'numeric', 'float', 'real', 'double', 'boolean', 'enum', 'text', 'blob', 'clob'
+ }
+ # Ensure the query only contains allowed keywords
+ tokens = sql_query.lower().split()
+ for token in tokens:
+ if not any(token.startswith(keyword) for keyword in allowed_keywords):
+ raise ValueError("Generated text contains invalid SQL syntax")
+
  return {"output": sql_query}
  except Exception as e:
  raise HTTPException(status_code=500, detail=str(e))
 
+
+
 if __name__ == "__main__":
  import uvicorn
- uvicorn.run(app, host="0.0.0.0", port=7860)
+ uvicorn.run(app, host="0.0.0.0", port=7860)@app.post("/generate")