cs50 · kzidane · Feb 3, 2020 · Jan 31, 2020 · Jan 31, 2020 · Jan 31, 2020
diff --git a/setup.py b/setup.py
@@ -16,5 +16,5 @@
     package_dir={"": "src"},
     packages=["cs50"],
     url="https://github.com/cs50/python-cs50",
-    version="5.0.2"
+    version="5.0.3"
 )
diff --git a/src/cs50/sql.py b/src/cs50/sql.py
@@ -120,6 +120,14 @@ def execute(self, sql, *args, **kwargs):
         if len(args) > 0 and len(kwargs) > 0:
             raise RuntimeError("cannot pass both named and positional parameters")
 
+        # Infer command from (unflattened) statement
+        for token in statements[0]:
+            if token.ttype in [sqlparse.tokens.Keyword.DDL, sqlparse.tokens.Keyword.DML]:
+                command = token.value.upper()
+                break
+        else:
+            command = None
+
         # Flatten statement
         tokens = list(statements[0].flatten())
 
@@ -313,45 +321,41 @@ def shutdown_session(exception=None):
 
                 # Return value
                 ret = True
-                if tokens[0].ttype == sqlparse.tokens.Keyword.DML:
-
-                    # Uppercase token's value
-                    value = tokens[0].value.upper()
-
-                    # If SELECT, return result set as list of dict objects
-                    if value == "SELECT":
-
-                        # Coerce types
-                        rows = [dict(row) for row in result.fetchall()]
-                        for row in rows:
-                            for column in row:
-
-                                # Coerce decimal.Decimal objects to float objects
-                                # https://groups.google.com/d/msg/sqlalchemy/0qXMYJvq8SA/oqtvMD9Uw-kJ
-                                if type(row[column]) is decimal.Decimal:
-                                    row[column] = float(row[column])
-
-                                # Coerce memoryview objects (as from PostgreSQL's bytea columns) to bytes
-                                elif type(row[column]) is memoryview:
-                                    row[column] = bytes(row[column])
-
-                        # Rows to be returned
-                        ret = rows
-
-                    # If INSERT, return primary key value for a newly inserted row (or None if none)
-                    elif value == "INSERT":
-                        if self._engine.url.get_backend_name() in ["postgres", "postgresql"]:
-                            try:
-                                result = connection.execute("SELECT LASTVAL()")
-                                ret = result.first()[0]
-                            except sqlalchemy.exc.OperationalError:  # If lastval is not yet defined in this session
-                                ret = None
-                        else:
-                            ret = result.lastrowid if result.rowcount == 1 else None
-
-                    # If DELETE or UPDATE, return number of rows matched
-                    elif value in ["DELETE", "UPDATE"]:
-                        ret = result.rowcount
+
+                # If SELECT, return result set as list of dict objects
+                if command == "SELECT":
+
+                    # Coerce types
+                    rows = [dict(row) for row in result.fetchall()]
+                    for row in rows:
+                        for column in row:
+
+                            # Coerce decimal.Decimal objects to float objects
+                            # https://groups.google.com/d/msg/sqlalchemy/0qXMYJvq8SA/oqtvMD9Uw-kJ
+                            if type(row[column]) is decimal.Decimal:
+                                row[column] = float(row[column])
+
+                            # Coerce memoryview objects (as from PostgreSQL's bytea columns) to bytes
+                            elif type(row[column]) is memoryview:
+                                row[column] = bytes(row[column])
+
+                    # Rows to be returned
+                    ret = rows
+
+                # If INSERT, return primary key value for a newly inserted row (or None if none)
+                elif command == "INSERT":
+                    if self._engine.url.get_backend_name() in ["postgres", "postgresql"]:
+                        try:
+                            result = connection.execute("SELECT LASTVAL()")
+                            ret = result.first()[0]
+                        except sqlalchemy.exc.OperationalError:  # If lastval is not yet defined in this session
+                            ret = None
+                    else:
+                        ret = result.lastrowid if result.rowcount == 1 else None
+
+                # If DELETE or UPDATE, return number of rows matched
+                elif command in ["DELETE", "UPDATE"]:
+                    ret = result.rowcount
 
             # If constraint violated, return None
             except sqlalchemy.exc.IntegrityError as e:

diff --git a/tests/sql.py b/tests/sql.py
@@ -158,6 +158,9 @@ def setUpClass(self):
     def setUp(self):
         self.db.execute("CREATE TABLE cs50 (id SERIAL PRIMARY KEY, val VARCHAR(16), bin BYTEA)")
 
+    def test_cte(self):
+        self.assertEqual(self.db.execute("WITH foo AS ( SELECT 1 AS bar ) SELECT bar FROM foo"), [{"bar": 1}])
+
 class SQLiteTests(SQLTests):
     @classmethod
     def setUpClass(self):
@@ -306,6 +309,9 @@ def test_numeric(self):
         self.assertRaises(RuntimeError, self.db.execute, "INSERT INTO foo VALUES (:1, :2)", 'bar', 'baz', 'qux')
         self.assertRaises(RuntimeError, self.db.execute, "INSERT INTO foo VALUES (:1, :2)", 'bar', baz='baz')
 
+    def test_cte(self):
+        self.assertEqual(self.db.execute("WITH foo AS ( SELECT 1 AS bar ) SELECT bar FROM foo"), [{"bar": 1}])
+
 
 if __name__ == "__main__":
     suite = unittest.TestSuite([