Merge pull request #97 from cs50/file

fixes bug whereby comments before query yielded wrong return value

Author: dmalan

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"
 )

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:

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"},