MCPcopy Create free account
hub / github.com/CL-lau/SQL-GPT / SQLHelper

Class SQLHelper

sql/SQLOperator.py:30–89  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

28
29@singleton
30class SQLHelper:
31 def __init__(self):
32 self.generated_ids = set()
33
34 self.dbMap = {}
35
36 def operate(self, sql: Optional[str], dbID: Optional[str]) -> [int, Any]:
37 sqlType = get_db_operation_class(sql=sql)
38
39 db = self.dbMap[dbID]
40 cursor = db.cursor()
41
42 if sqlType == SQL_class.UPDATE:
43 cursor.execute(sql)
44 db.commit()
45 updated_rows = cursor.rowcount
46 return SQL_class.UPDATE, updated_rows
47 elif sqlType == SQL_class.DELETE:
48 cursor.execute(sql)
49 db.commit()
50 deleted_rows = cursor.rowcount
51 return SQL_class.DELETE, deleted_rows
52 elif sqlType == SQL_class.INSERT:
53 cursor.execute(sql)
54 db.commit()
55 new_row_id = cursor.lastrowid
56 return SQL_class.INSERT, new_row_id
57 elif sqlType == SQL_class.SELECT:
58 cursor.execute(sql)
59 results = cursor.fetchall()
60 cursor.close()
61 db.close()
62 dataframe = pd.DataFrame(results, columns=[i[0] for i in cursor.description]) # 根据你的列名进行修改
63 return SQL_class.SELECT, dataframe
64 else:
65 return -1, -1
66
67 def add_db(self, dbAddress: Optional[str], dbPort: Optional[str], userName: Optional[str], password: Optional[str], dbName: Optional[str], dbID: Optional[str]) -> str:
68 db = pymysql.connect(host=dbAddress, port=int(dbPort), user=userName, password=password, database=dbName)
69
70 for key, value in self.dbMap.items():
71 if value.host == dbAddress and value.port == int(dbPort) and value.user == userName and value.password == password and value.database == dbName:
72 return key
73
74 if dbID is None:
75 dbID = self.generated_id()
76 self.dbMap[dbID] = db
77 else:
78 self.dbMap[dbID] = db
79 return dbID
80
81 def generated_id(self) -> str:
82 characters = string.ascii_letters + string.digits
83
84 while True:
85 idt = ''.join(random.choice(characters) for _ in range(13))
86 if idt not in self.generated_ids:
87 self.generated_ids.add(idt)

Callers 1

__init__Method · 0.90

Calls

no outgoing calls

Tested by

no test coverage detected