adds support for BLOB, improves support for LASTVAL(), uses connection instead of engine directly

dmalan · web-flow · commit 74e6dd0dfd4b · 2019-12-14T12:12:34.000-05:00
diff --git a/src/cs50/sql.py b/src/cs50/sql.py
@@ -51,21 +51,31 @@ def __init__(self, url, **kwargs):
             if not os.path.isfile(matches.group(1)):
                 raise RuntimeError("not a file: {}".format(matches.group(1)))
 
-            # Create engine, raising exception if back end's module not installed
-            self.engine = sqlalchemy.create_engine(url, **kwargs)
+            # Create engine, disabling SQLAlchemy's own autocommit mode, raising exception if back end's module not installed
+            engine = sqlalchemy.create_engine(url, **kwargs).execution_options(autocommit=False)
 
-            # Enable foreign key constraints
+            # Listener for connections
             def connect(dbapi_connection, connection_record):
+
+                # Disable underlying API's own emitting of BEGIN and COMMIT
+                dbapi_connection.isolation_level = None
+
+                # Enable foreign key constraints
                 if type(dbapi_connection) is sqlite3.Connection:  # If back end is sqlite
                     cursor = dbapi_connection.cursor()
                     cursor.execute("PRAGMA foreign_keys=ON")
                     cursor.close()
-            sqlalchemy.event.listen(self.engine, "connect", connect)
+
+            # Register listener
+            sqlalchemy.event.listen(engine, "connect", connect)
 
         else:
 
             # Create engine, raising exception if back end's module not installed
-            self.engine = sqlalchemy.create_engine(url, **kwargs)
+            engine = sqlalchemy.create_engine(url, **kwargs)
+
+        # Connect to database (for transactions' sake)
+        self._connection = engine.connect().execution_options(autocommit=False)
 
         # Log statements to standard error
         logging.basicConfig(level=logging.DEBUG)
@@ -260,8 +270,11 @@ def execute(self, sql, *args, **kwargs):
             # Prepare, execute statement
             try:
 
+                # Join tokens into statement, abbreviating binary data as <class 'bytes'>
+                _statement = "".join([str(bytes) if token.ttype == sqlparse.tokens.Other else str(token) for token in tokens])
+
                 # Execute statement
-                result = self.engine.execute(sqlalchemy.text(statement))
+                result = self._connection.execute(sqlalchemy.text(statement))
 
                 # Return value
                 ret = True
@@ -273,22 +286,33 @@ def execute(self, sql, *args, **kwargs):
                     # If SELECT, return result set as list of dict objects
                     if value == "SELECT":
 
-                        # Coerce any decimal.Decimal objects to float objects
-                        # https://groups.google.com/d/msg/sqlalchemy/0qXMYJvq8SA/oqtvMD9Uw-kJ
+                        # 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"]:
-                            result = self.engine.execute("SELECT LASTVAL()")
-                            ret = result.first()[0]
+                        if self._connection.engine.url.get_backend_name() in ["postgres", "postgresql"]:
+                            try:
+                                result = self._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.lastrowid > 0 else None
+                            ret = result.lastrowid if result.rowcount == 1 else None
 
                     # If DELETE or UPDATE, return number of rows matched
                     elif value in ["DELETE", "UPDATE"]:
@@ -310,7 +334,7 @@ def execute(self, sql, *args, **kwargs):
 
             # Return value
             else:
-                self._logger.debug(termcolor.colored(statement, "green"))
+                self._logger.debug(termcolor.colored(_statement, "green"))
                 return ret
 
     def _escape(self, value):
@@ -333,65 +357,68 @@ def __escape(value):
             if type(value) is bool:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Number,
-                    sqlalchemy.types.Boolean().literal_processor(self.engine.dialect)(value))
+                    sqlalchemy.types.Boolean().literal_processor(self._connection.engine.dialect)(value))
 
             # bytearray, bytes
             elif type(value) in [bytearray, bytes]:
-                raise RuntimeError("unsupported value")  # TODO
+                if self._connection.engine.url.get_backend_name() in ["mysql", "sqlite"]:
+                    return sqlparse.sql.Token(sqlparse.tokens.Other, f"x'{value.hex()}'")  # https://dev.mysql.com/doc/refman/8.0/en/hexadecimal-literals.html
+                elif self._connection.engine.url.get_backend_name() == "postgresql":
+                    return sqlparse.sql.Token(sqlparse.tokens.Other, f"'\\x{value.hex()}'")  # https://dba.stackexchange.com/a/203359
+                else:
+                    raise RuntimeError("unsupported value: {}".format(value))
 
             # datetime.date
             elif type(value) is datetime.date:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self.engine.dialect)(value.strftime("%Y-%m-%d")))
+                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value.strftime("%Y-%m-%d")))
 
             # datetime.datetime
             elif type(value) is datetime.datetime:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self.engine.dialect)(value.strftime("%Y-%m-%d %H:%M:%S")))
+                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value.strftime("%Y-%m-%d %H:%M:%S")))
 
             # datetime.time
             elif type(value) is datetime.time:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self.engine.dialect)(value.strftime("%H:%M:%S")))
+                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value.strftime("%H:%M:%S")))
 
             # float
             elif type(value) is float:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Number,
-                    sqlalchemy.types.Float().literal_processor(self.engine.dialect)(value))
+                    sqlalchemy.types.Float().literal_processor(self._connection.engine.dialect)(value))
 
             # int
             elif type(value) is int:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Number,
-                    sqlalchemy.types.Integer().literal_processor(self.engine.dialect)(value))
+                    sqlalchemy.types.Integer().literal_processor(self._connection.engine.dialect)(value))
 
             # str
             elif type(value) is str:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.String,
-                    sqlalchemy.types.String().literal_processor(self.engine.dialect)(value))
+                    sqlalchemy.types.String().literal_processor(self._connection.engine.dialect)(value))
 
             # None
             elif value is None:
                 return sqlparse.sql.Token(
                     sqlparse.tokens.Keyword,
-                    sqlalchemy.types.NullType().literal_processor(self.engine.dialect)(value))
+                    sqlalchemy.types.NullType().literal_processor(self._connection.engine.dialect)(value))
 
             # Unsupported value
             else:
                 raise RuntimeError("unsupported value: {}".format(value))
 
         # Escape value(s), separating with commas as needed
         if type(value) in [list, tuple]:
-            return sqlparse.sql.TokenList(sqlparse.parse(", ".join([str(__escape(v)) for v in value])))
+            return sqlparse.sql.TokenList([__escape(v) for v in value])
         else:
-            return sqlparse.sql.Token(
-                sqlparse.tokens.String,
-                __escape(value))
+            return __escape(value)
 
 
 def _parse_exception(e):
diff --git a/tests/sql.py b/tests/sql.py
@@ -32,9 +32,9 @@ def test_select_all(self):
         self.assertEqual(self.db.execute("SELECT * FROM cs50"), [])
 
         rows = [
-            {"id": 1, "val": "foo"},
-            {"id": 2, "val": "bar"},
-            {"id": 3, "val": "baz"}
+            {"id": 1, "val": "foo", "bin": None},
+            {"id": 2, "val": "bar", "bin": None},
+            {"id": 3, "val": "baz", "bin": None}
         ]
         for row in rows:
             self.db.execute("INSERT INTO cs50(val) VALUES(:val)", val=row["val"])
@@ -61,7 +61,7 @@ def test_select_where(self):
         for row in rows:
             self.db.execute("INSERT INTO cs50(val) VALUES(:val)", val=row["val"])
 
-        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])
+        self.assertEqual(self.db.execute("SELECT id, val 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"), [])
@@ -99,6 +99,29 @@ def test_string_literal_with_colon(self):
         self.assertEqual(self.db.execute("SELECT val FROM cs50 WHERE val = ':bar :baz'"), [{"val": ":bar :baz"}])
         self.assertEqual(self.db.execute("SELECT val FROM cs50 WHERE val = '  :bar :baz'"), [{"val": "  :bar :baz"}])
 
+    def test_blob(self):
+        rows = [
+            {"id": 1, "bin": b"\0"},
+            {"id": 2, "bin": b"\1"},
+            {"id": 3, "bin": b"\2"}
+        ]
+        for row in rows:
+            self.db.execute("INSERT INTO cs50(bin) VALUES(:bin)", bin=row["bin"])
+        self.assertEqual(self.db.execute("SELECT id, bin FROM cs50"), rows)
+
+    def test_commit(self):
+        self.db.execute("BEGIN")
+        self.db.execute("INSERT INTO cs50 (val) VALUES('foo')")
+        self.db.execute("COMMIT")
+        self.assertEqual(self.db.execute("SELECT val FROM cs50"), [{"val": "foo"}])
+
+    def test_rollback(self):
+        self.db.execute("BEGIN")
+        self.db.execute("INSERT INTO cs50 (val) VALUES('foo')")
+        self.db.execute("INSERT INTO cs50 (val) VALUES('bar')")
+        self.db.execute("ROLLBACK")
+        self.assertEqual(self.db.execute("SELECT val FROM cs50"), [])
+
     def tearDown(self):
         self.db.execute("DROP TABLE cs50")
         self.db.execute("DROP TABLE IF EXISTS foo")
@@ -119,15 +142,15 @@ def setUpClass(self):
         self.db = SQL("mysql://root@localhost/test")
 
     def setUp(self):
-        self.db.execute("CREATE TABLE cs50 (id INTEGER NOT NULL AUTO_INCREMENT, val VARCHAR(16), PRIMARY KEY (id))")
+        self.db.execute("CREATE TABLE cs50 (id INTEGER NOT NULL AUTO_INCREMENT, val VARCHAR(16), bin BLOB, PRIMARY KEY (id))")
 
 class PostgresTests(SQLTests):
     @classmethod
     def setUpClass(self):
         self.db = SQL("postgresql://postgres@localhost/test")
 
     def setUp(self):
-        self.db.execute("CREATE TABLE cs50 (id SERIAL PRIMARY KEY, val VARCHAR(16))")
+        self.db.execute("CREATE TABLE cs50 (id SERIAL PRIMARY KEY, val VARCHAR(16), bin BYTEA)")
 
 class SQLiteTests(SQLTests):
     @classmethod
@@ -136,7 +159,7 @@ def setUpClass(self):
         self.db = SQL("sqlite:///test.db")
 
     def setUp(self):
-        self.db.execute("CREATE TABLE cs50(id INTEGER PRIMARY KEY, val TEXT)")
+        self.db.execute("CREATE TABLE cs50(id INTEGER PRIMARY KEY, val TEXT, bin BLOB)")
 
     def test_lastrowid(self):
         self.db.execute("CREATE TABLE foo(id INTEGER PRIMARY KEY AUTOINCREMENT, firstname TEXT, lastname TEXT)")
