diff --git a/setup.py b/setup.py index 6a8bf06..3b32066 100644 --- a/setup.py +++ b/setup.py @@ -16,5 +16,5 @@ package_dir={"": "src"}, packages=["cs50"], url="https://github.com/cs50/python-cs50", - version="4.0.3" + version="4.0.4" ) diff --git a/src/cs50/sql.py b/src/cs50/sql.py index 6898b75..3bc52d6 100644 --- a/src/cs50/sql.py +++ b/src/cs50/sql.py @@ -75,9 +75,11 @@ def execute(self, sql, *args, **kwargs): import termcolor import warnings + # Parse statement, stripping comments and then leading/trailing whitespace + statements = sqlparse.parse(sqlparse.format(sql, strip_comments=True).strip()) + # Allow only one statement at a time, since SQLite doesn't support multiple # https://docs.python.org/3/library/sqlite3.html#sqlite3.Cursor.execute - statements = sqlparse.parse(sql) if len(statements) > 1: raise RuntimeError("too many statements at once") elif len(statements) == 0: diff --git a/tests/sql.py b/tests/sql.py index 9b28930..ceee9dc 100644 --- a/tests/sql.py +++ b/tests/sql.py @@ -63,6 +63,12 @@ def test_select_where(self): self.assertEqual(self.db.execute("SELECT * FROM cs50 WHERE id = :id OR val = :val", id=rows[1]["id"], val=rows[2]["val"]), rows[1:3]) + def test_select_with_comments(self): + self.assertEqual(self.db.execute("--comment\nSELECT * FROM cs50;\n--comment"), []) + + def test_select_with_semicolon(self): + self.assertEqual(self.db.execute("SELECT * FROM cs50;\n--comment"), []) + def test_update_returns_affected_rows(self): rows = [ {"id": 1, "val": "foo"},