Skip to content

Commit

Permalink
fix(custom-pewee): Keep same signature of execute_sql function
Browse files Browse the repository at this point in the history
  • Loading branch information
tanmoysrt committed Feb 17, 2025
1 parent 5c10214 commit 2c294f7
Showing 1 changed file with 12 additions and 8 deletions.
20 changes: 12 additions & 8 deletions agent/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,9 @@ def default(self, obj):
return str(obj)


SENTINEL = object()


class CustomPeeweeDB(peewee.MySQLDatabase):
"""
Override peewee.MySQLDatabase to modify `execute_sql` method
Expand Down Expand Up @@ -674,18 +677,19 @@ class CustomPeeweeDB(peewee.MySQLDatabase):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)

def execute_sql(self, sql):
if self.in_transaction():
commit = False
elif self.commit_select:
commit = True
else:
commit = not sql[:6].lower().startswith("select")
def execute_sql(self, sql, params=None, commit=SENTINEL):
if commit is SENTINEL:
if self.in_transaction():
commit = False
elif self.commit_select:
commit = True
else:
commit = not sql[:6].lower().startswith("select")

with self.__exception_wrapper__:
cursor = self.cursor(commit)
try:
cursor.execute(sql, None) # params passed as none
cursor.execute(sql, params)
except Exception:
if self.autorollback and not self.in_transaction():
self.rollback()
Expand Down

0 comments on commit 2c294f7

Please sign in to comment.