Merge pull request #157 from cs50/refactor/scoped-session-fix

v7.0.0

Author: Kareem Zidane

setup.py

@@ -10,11 +10,11 @@
         "Topic :: Software Development :: Libraries :: Python Modules"
     ],
     description="CS50 library for Python",
-    install_requires=["Flask>=1.0", "SQLAlchemy", "sqlparse", "termcolor"],
+    install_requires=["Flask>=1.0", "SQLAlchemy<2", "sqlparse", "termcolor"],
     keywords="cs50",
     name="cs50",
     package_dir={"": "src"},
     packages=["cs50"],
     url="https://github.com/cs50/python-cs50",
-    version="6.0.5"
+    version="7.0.0"
 )

src/cs50/__init__.py

@@ -1,20 +1,5 @@
-import logging
-import os
-import sys
-
-
-# Disable cs50 logger by default
-logging.getLogger("cs50").disabled = True
-
-# Import cs50_*
-from .cs50 import get_char, get_float, get_int, get_string
-try:
-    from .cs50 import get_long
-except ImportError:
-    pass
-
-# Hook into flask importing
-from . import flask
-
-# Wrap SQLAlchemy
+from .cs50 import get_float, get_int, get_string
 from .sql import SQL
+from ._logger import _setup_logger
+
+_setup_logger()

src/cs50/_engine.py

@@ -0,0 +1,66 @@
+import threading
+
+from ._engine_util import create_engine
+
+
+thread_local_data = threading.local()
+
+
+class Engine:
+    """Wraps a SQLAlchemy engine.
+    """
+
+    def __init__(self, url):
+        self._engine = create_engine(url)
+
+    def get_transaction_connection(self):
+        """
+        :returns: A new connection with autocommit disabled (to be used for transactions).
+        """
+
+        _thread_local_connections()[self] = self._engine.connect().execution_options(
+            autocommit=False)
+        return self.get_existing_transaction_connection()
+
+    def get_connection(self):
+        """
+        :returns: A new connection with autocommit enabled
+        """
+
+        return self._engine.connect().execution_options(autocommit=True)
+
+    def get_existing_transaction_connection(self):
+        """
+        :returns: The transaction connection bound to this Engine instance, if one exists, or None.
+        """
+
+        return _thread_local_connections().get(self)
+
+    def close_transaction_connection(self):
+        """Closes the transaction connection bound to this Engine instance, if one exists and
+        removes it.
+        """
+
+        connection = self.get_existing_transaction_connection()
+        if connection:
+            connection.close()
+            del _thread_local_connections()[self]
+
+    def is_postgres(self):
+        return self._engine.dialect.name in {"postgres", "postgresql"}
+
+    def __getattr__(self, attr):
+        return getattr(self._engine, attr)
+
+def _thread_local_connections():
+    """
+    :returns: A thread local dict to keep track of transaction connection. If one does not exist,
+    creates one.
+    """
+
+    try:
+        connections = thread_local_data.connections
+    except AttributeError:
+        connections = thread_local_data.connections = {}
+
+    return connections

src/cs50/_engine_util.py

@@ -0,0 +1,43 @@
+"""Utility functions used by _session.py.
+"""
+
+import os
+import sqlite3
+
+import sqlalchemy
+
+sqlite_url_prefix = "sqlite:///"
+
+
+def create_engine(url, **kwargs):
+    """Creates a new SQLAlchemy engine. If ``url`` is a URL for a SQLite database, makes sure that
+    the SQLite file exits and enables foreign key constraints.
+    """
+
+    try:
+        engine = sqlalchemy.create_engine(url, **kwargs)
+    except sqlalchemy.exc.ArgumentError:
+        raise RuntimeError(f"invalid URL: {url}") from None
+
+    if _is_sqlite_url(url):
+        _assert_sqlite_file_exists(url)
+        sqlalchemy.event.listen(engine, "connect", _enable_sqlite_foreign_key_constraints)
+
+    return engine
+
+def _is_sqlite_url(url):
+    return url.startswith(sqlite_url_prefix)
+
+
+def _assert_sqlite_file_exists(url):
+    path = url[len(sqlite_url_prefix):]
+    if not os.path.exists(path):
+        raise RuntimeError(f"does not exist: {path}")
+    if not os.path.isfile(path):
+        raise RuntimeError(f"not a file: {path}")
+
+
+def _enable_sqlite_foreign_key_constraints(dbapi_connection, _):
+    cursor = dbapi_connection.cursor()
+    cursor.execute("PRAGMA foreign_keys=ON")
+    cursor.close()

src/cs50/_logger.py

@@ -0,0 +1,98 @@
+"""Sets up logging for the library.
+"""
+
+import logging
+import os.path
+import re
+import sys
+import traceback
+
+import termcolor
+
+
+def green(msg):
+    return _colored(msg, "green")
+
+
+def red(msg):
+    return _colored(msg, "red")
+
+
+def yellow(msg):
+    return _colored(msg, "yellow")
+
+
+def _colored(msg, color):
+    return termcolor.colored(str(msg), color)
+
+
+def _setup_logger():
+    _configure_default_logger()
+    _patch_root_handler_format_exception()
+    _configure_cs50_logger()
+    _patch_excepthook()
+
+
+def _configure_default_logger():
+    """Configures a default handler and formatter to prevent flask and werkzeug from adding theirs.
+    """
+
+    logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.DEBUG)
+
+
+def _patch_root_handler_format_exception():
+    """Patches formatException for the root handler to use ``_format_exception``.
+    """
+
+    try:
+        formatter = logging.root.handlers[0].formatter
+        formatter.formatException = lambda exc_info: _format_exception(*exc_info)
+    except IndexError:
+        pass
+
+
+def _configure_cs50_logger():
+    """Disables the cs50 logger by default. Disables logging propagation to prevent messages from
+    being logged more than once. Sets the logging handler and formatter.
+    """
+
+    _logger = logging.getLogger("cs50")
+    _logger.disabled = True
+    _logger.setLevel(logging.DEBUG)
+
+    # Log messages once
+    _logger.propagate = False
+
+    handler = logging.StreamHandler()
+    handler.setLevel(logging.DEBUG)
+
+    formatter = logging.Formatter("%(levelname)s: %(message)s")
+    formatter.formatException = lambda exc_info: _format_exception(*exc_info)
+    handler.setFormatter(formatter)
+    _logger.addHandler(handler)
+
+
+def _patch_excepthook():
+    sys.excepthook = lambda type_, value, exc_tb: print(
+        _format_exception(type_, value, exc_tb), file=sys.stderr)
+
+
+def _format_exception(type_, value, exc_tb):
+    """Formats traceback, darkening entries from global site-packages directories and user-specific
+    site-packages directory.
+    https://stackoverflow.com/a/46071447/5156190
+    """
+
+    # Absolute paths to site-packages
+    packages = tuple(os.path.join(os.path.abspath(p), "") for p in sys.path[1:])
+
+    # Highlight lines not referring to files in site-packages
+    lines = []
+    for line in traceback.format_exception(type_, value, exc_tb):
+        matches = re.search(r"^  File \"([^\"]+)\", line \d+, in .+", line)
+        if matches and matches.group(1).startswith(packages):
+            lines += line
+        else:
+            matches = re.search(r"^(\s*)(.*?)(\s*)$", line, re.DOTALL)
+            lines.append(matches.group(1) + yellow(matches.group(2)) + matches.group(3))
+    return "".join(lines).rstrip()

src/cs50/_sql_sanitizer.py

@@ -0,0 +1,95 @@
+import datetime
+import re
+
+import sqlalchemy
+import sqlparse
+
+
+class SQLSanitizer:
+    """Sanitizes SQL values.
+    """
+
+    def __init__(self, dialect):
+        self._dialect = dialect
+
+    def escape(self, value):
+        """Escapes value using engine's conversion function.
+        https://docs.sqlalchemy.org/en/latest/core/type_api.html#sqlalchemy.types.TypeEngine.literal_processor
+
+        :param value: The value to be sanitized
+
+        :returns: The sanitized value
+        """
+        # pylint: disable=too-many-return-statements
+        if isinstance(value, (list, tuple)):
+            return self.escape_iterable(value)
+
+        if isinstance(value, bool):
+            return sqlparse.sql.Token(
+                sqlparse.tokens.Number,
+                sqlalchemy.types.Boolean().literal_processor(self._dialect)(value))
+
+        if isinstance(value, bytes):
+            if self._dialect.name in {"mysql", "sqlite"}:
+                # https://dev.mysql.com/doc/refman/8.0/en/hexadecimal-literals.html
+                return sqlparse.sql.Token(sqlparse.tokens.Other, f"x'{value.hex()}'")
+            if self._dialect.name in {"postgres", "postgresql"}:
+                # https://dba.stackexchange.com/a/203359
+                return sqlparse.sql.Token(sqlparse.tokens.Other, f"'\\x{value.hex()}'")
+
+            raise RuntimeError(f"unsupported value: {value}")
+
+        string_processor = sqlalchemy.types.String().literal_processor(self._dialect)
+        if isinstance(value, datetime.date):
+            return sqlparse.sql.Token(
+                sqlparse.tokens.String, string_processor(value.strftime("%Y-%m-%d")))
+
+        if isinstance(value, datetime.datetime):
+            return sqlparse.sql.Token(
+                sqlparse.tokens.String, string_processor(value.strftime("%Y-%m-%d %H:%M:%S")))
+
+        if isinstance(value, datetime.time):
+            return sqlparse.sql.Token(
+                sqlparse.tokens.String, string_processor(value.strftime("%H:%M:%S")))
+
+        if isinstance(value, float):
+            return sqlparse.sql.Token(
+                sqlparse.tokens.Number,
+                sqlalchemy.types.Float().literal_processor(self._dialect)(value))
+
+        if isinstance(value, int):
+            return sqlparse.sql.Token(
+                sqlparse.tokens.Number,
+                sqlalchemy.types.Integer().literal_processor(self._dialect)(value))
+
+        if isinstance(value, str):
+            return sqlparse.sql.Token(sqlparse.tokens.String, string_processor(value))
+
+        if value is None:
+            return sqlparse.sql.Token(
+                sqlparse.tokens.Keyword,
+                sqlalchemy.types.NullType().literal_processor(self._dialect)(value))
+
+        raise RuntimeError(f"unsupported value: {value}")
+
+    def escape_iterable(self, iterable):
+        """Escapes each value in iterable and joins all the escaped values with ", ", formatted for
+        SQL's ``IN`` operator.
+
+        :param: An iterable of values to be escaped
+
+        :returns: A comma-separated list of escaped values from ``iterable``
+        :rtype: :class:`sqlparse.sql.TokenList`
+        """
+
+        return sqlparse.sql.TokenList(
+            sqlparse.parse(", ".join([str(self.escape(v)) for v in iterable])))
+
+
+def escape_verbatim_colon(value):
+    """Escapes verbatim colon from a value so as it is not confused with a parameter marker.
+    """
+
+    # E.g., ':foo, ":foo,   :foo will be replaced with
+    # '\:foo, "\:foo,   \:foo respectively
+    return re.sub(r"(^(?:'|\")|\s+):", r"\1\:", value)

src/cs50/_sql_util.py

@@ -0,0 +1,51 @@
+"""Utility functions used by sql.py.
+"""
+
+import contextlib
+import decimal
+import warnings
+
+import sqlalchemy
+
+
+def process_select_result(result):
+    """Converts a SQLAlchemy result to a ``list`` of ``dict`` objects, each of which represents a
+    row in the result set.
+
+    :param result: A SQLAlchemy result
+    :type result: :class:`sqlalchemy.engine.Result`
+    """
+    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 isinstance(row[column], decimal.Decimal):
+                row[column] = float(row[column])
+
+            # Coerce memoryview objects (as from PostgreSQL's bytea columns) to bytes
+            elif isinstance(row[column], memoryview):
+                row[column] = bytes(row[column])
+
+    return rows
+
+
+@contextlib.contextmanager
+def raise_errors_for_warnings():
+    """Catches warnings and raises errors instead.
+    """
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        yield
+
+
+def postgres_lastval(connection):
+    """
+    :returns: The ID of the last inserted row, if defined in this session, or None
+    """
+
+    try:
+        return connection.execute("SELECT LASTVAL()").first()[0]
+    except sqlalchemy.exc.OperationalError:
+        return None