Refactor and make getting at SQL connections easy from python (#102)

Refactor, expose "noteable.sql.get_sqla_connection() and get_sqla_engine()

Author: James Robinson

noteable/__init__.py

@@ -2,25 +2,22 @@
 
 __version__ = pkg_resources.get_distribution("noteable").version
 
-from .data_loader import LOCAL_DB_CONN_HANDLE, NoteableDataLoaderMagic, get_db_connection
+from .data_loader import NoteableDataLoaderMagic
 from .datasources import bootstrap_datasources
 from .logging import configure_logging
 from .ntbl import NTBLMagic
+from .sql.connection import bootstrap_duckdb
 from .sql.magic import SqlMagic
-from .sql.run import add_commit_blacklist_dialect
 
 
 def load_ipython_extension(ipython):
-    # Initialize any remote datasource connections
-    bootstrap_datasources()
-
-    # Always prevent sql-magic from trying to autocommit bigquery,
-    # for the legacy datasource support for Expel and whomever.
-    add_commit_blacklist_dialect('bigquery')
+    configure_logging(False, "INFO", "DEBUG")
 
-    # Initialize the noteable local (duck_db) database connection
-    get_db_connection(LOCAL_DB_CONN_HANDLE)
+    # Initialize any remote datasource connections.
+    bootstrap_datasources()
 
-    configure_logging(False, "INFO", "DEBUG")
+    # Initialize the noteable local (duck_db) database connection.
+    bootstrap_duckdb()
 
+    # Register all of our magics.
     ipython.register_magics(NoteableDataLoaderMagic, NTBLMagic, SqlMagic)

noteable/data_loader.py

@@ -1,5 +1,4 @@
 import mimetypes
-from typing import Optional
 
 import pandas as pd
 from IPython.core.magic import Magics, line_cell_magic, magics_class
@@ -8,46 +7,12 @@
 from traitlets import Bool, Int
 from traitlets.config import Configurable
 
-from noteable.sql.connection import Connection
+from noteable.sql.connection import LOCAL_DB_CONN_HANDLE, get_db_connection
 
 EXCEL_MIMETYPES = {
     "application/vnd.ms-excel",  # .xls
     "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",  # .xlsx
 }
-LOCAL_DB_CONN_HANDLE = "@noteable"
-duckdb_location = "duckdb:///:memory:"
-
-
-def get_db_connection(sql_cell_handle_or_human_name: str) -> Optional['Connection']:
-    """Return the sql.connection.Connection corresponding to the requested
-        datasource sql_cell_handle or human name.
-
-    If the cell handle happens to correspond to the 'local database' DuckDB database,
-    then we will bootstrap it upon demand. Otherwise, try to find and return
-    the connection.
-
-    Will return None if the given handle isn't @noteable and isn't present in
-    the connections dict already (created after this kernel was launched?)
-    """
-    if (
-        sql_cell_handle_or_human_name == LOCAL_DB_CONN_HANDLE
-        and sql_cell_handle_or_human_name not in Connection.connections
-    ):
-        # Bootstrap the DuckDB database if asked and needed.
-        return Connection.set(
-            duckdb_location,
-            human_name="Local Database",
-            name=LOCAL_DB_CONN_HANDLE,
-        )
-    else:
-        # If, say, they created the datasource *after* this kernel was launched, then
-        # this will come up empty and the caller should handle gracefully.
-        for conn in Connection.connections.values():
-            if (
-                conn.name == sql_cell_handle_or_human_name
-                or conn.human_name == sql_cell_handle_or_human_name
-            ):
-                return conn
 
 
 @magics_class
@@ -56,7 +21,6 @@ class NoteableDataLoaderMagic(Magics, Configurable):
         True, config=True, help="Return the first N rows from the loaded pandas dataframe"
     )
     display_example = Bool(True, config=True, help="Show example SQL query")
-    display_connection_str = Bool(False, config=True, help="Show connection string after execute")
     pandas_limit = Int(10, config=True, help="The limit of rows to returns in the pandas dataframe")
 
     @line_cell_magic("create_or_replace_data_view")
@@ -116,19 +80,13 @@ def execute(self, line="", cell=""):
                 f"Could not find datasource identified by {args.connection!r}. Perhaps restart the kernel?"
             )
 
-        tmp_df.to_sql(tablename, conn.session, if_exists="replace", index=args.include_index)
-
-        if self.display_connection_str:
-            print(f"Connect with: %sql {conn.name}")
+        tmp_df.to_sql(
+            tablename, conn.sqla_connection, if_exists="replace", index=args.include_index
+        )
 
         if self.display_example:
-            if conn.human_name:
-                noun = f'{conn.human_name!r}'
-            else:
-                # Hmm. "Legacy" created datasource. Err on the engine's dialect name?
-                noun = conn._engine.dialect.name
             print(
-                f"""Create a {noun} SQL cell and then input query. """
+                f"""Create a {conn.human_name!r} SQL cell and then input query. """
                 f"Example: 'SELECT * FROM \"{tablename}\" LIMIT 10'"
             )
 

noteable/sql/__init__.py

@@ -1 +1 @@
-from .magic import *
+from .connection import get_sqla_connection, get_sqla_engine  # noqa

noteable/sql/connection.py

@@ -8,26 +8,16 @@
 logger = structlog.get_logger(__name__)
 
 
-class UnknownConnectionError(Exception):
-    pass
-
-
-def rough_dict_get(dct, sought, default=None):
-    """
-    Like dct.get(sought), but any key containing sought will do.
+LOCAL_DB_CONN_HANDLE = "@noteable"
+LOCAL_DB_CONN_NAME = "Local Database"
+DUCKDB_LOCATION = "duckdb:///:memory:"
 
-    If there is a `@` in sought, seek each piece separately.
-    This lets `me@server` match `me:***@myserver/db`
-    """
 
-    sought = sought.split("@")
-    for key, val in dct.items():
-        if not any(s.lower() not in key.lower() for s in sought):
-            return val
-    return default
+class UnknownConnectionError(Exception):
+    pass
 
 
-class Connection(object):
+class Connection:
     current = None
     connections: Dict[str, 'Connection'] = {}
     bootstrapping_failures: Dict[str, str] = {}
@@ -107,23 +97,30 @@ def __init__(self, connect_str=None, name=None, human_name=None, **create_engine
         self.metadata = sqlalchemy.MetaData(bind=self._engine)
         self.name = name or self.assign_name(self._engine)
         self.human_name = human_name
-        self._session = None
+        self._sqla_connection = None
         self.connections[name or repr(self.metadata.bind.url)] = self
 
         Connection.current = self
 
     @property
-    def session(self) -> sqlalchemy.engine.base.Connection:
-        """Lazily connect to the database.
+    def engine(self) -> sqlalchemy.engine.base.Engine:
+        return self._engine
 
-        Despite the name, this is a SQLA Connection, not a Session. And 'Connection'
-        is highly overused term around here.
-        """
+    @property
+    def sqla_connection(self) -> sqlalchemy.engine.base.Connection:
+        """Lazily connect to the database. Return a SQLA Connection object, or die trying."""
+
+        if not self._sqla_connection:
+            self._sqla_connection = self._engine.connect()
 
-        if not self._session:
-            self._session = self._engine.connect()
+        return self._sqla_connection
 
-        return self._session
+    def reset_connection_pool(self):
+        """Reset the SQLA connection pool, such as after an exception suspected to indicate
+        a broken connection has been raised.
+        """
+        self._engine.dispose()
+        self._sqla_connection = None
 
     @classmethod
     def set(
@@ -164,14 +161,23 @@ def connection_list(cls):
             result.append(template.format(engine_url.__repr__()))
         return "\n".join(result)
 
+    @classmethod
+    def find(cls, name: str) -> Optional['Connection']:
+        """Find a connection by SQL cell handle or by human assigned name"""
+        # TODO: Capt. Obvious says to double-register the instance by both of these keys
+        # to then be able to do lookups properly in this dict?
+        for c in cls.connections.values():
+            if c.name == name or c.human_name == name:
+                return c
+
     @classmethod
     def get_engine(cls, name: str) -> Optional[Engine]:
         """Return the SQLAlchemy Engine given either the sql_cell_handle or
         end-user assigned name for the connection.
         """
-        for c in cls.connections.values():
-            if c.name == name or c.human_name == name:
-                return c._engine
+        maybe_conn = cls.find(name)
+        if maybe_conn:
+            return maybe_conn.engine
 
     @classmethod
     def add_bootstrapping_failure(cls, name: str, human_name: Optional[str], error_message: str):
@@ -204,7 +210,56 @@ def _close(cls, descriptor):
             )
         cls.connections.pop(conn.name, None)
         cls.connections.pop(str(conn.metadata.bind.url), None)
-        conn.session.close()
+        conn.sqla_connection.close()
 
     def close(self):
         self.__class__._close(self)
+
+
+def rough_dict_get(dct, sought, default=None):
+    """
+    Like dct.get(sought), but any key containing sought will do.
+
+    If there is a `@` in sought, seek each piece separately.
+    This lets `me@server` match `me:***@myserver/db`
+    """
+
+    sought = sought.split("@")
+    for key, val in dct.items():
+        if not any(s.lower() not in key.lower() for s in sought):
+            return val
+    return default
+
+
+def get_db_connection(name_or_handle: str) -> Optional[Connection]:
+    """Return the noteable.sql.connection.Connection corresponding to the requested
+        datasource a name or handle.
+
+    Will return None if the given handle isn't present in
+    the connections dict already (created after this kernel was launched?)
+    """
+    return Connection.find(name_or_handle)
+
+
+def get_sqla_connection(name_or_handle: str) -> Optional[sqlalchemy.engine.base.Connection]:
+    """Return a SQLAlchemy connection given a name or handle
+    Returns None if cannot find by this string.
+    """
+    nconn = get_db_connection(name_or_handle)
+    if nconn:
+        return nconn.sqla_connection
+
+
+def get_sqla_engine(name_or_handle: str) -> Optional[Engine]:
+    """Return a SQLAlchemy Engine given a name or handle.
+    Returns None if cannot find by this string.
+    """
+    return Connection.get_engine(name_or_handle)
+
+
+def bootstrap_duckdb():
+    Connection.set(
+        DUCKDB_LOCATION,
+        human_name=LOCAL_DB_CONN_NAME,
+        name=LOCAL_DB_CONN_HANDLE,
+    )

noteable/sql/magic.py

@@ -149,7 +149,6 @@ def execute(self, line="", cell="", local_ns={}):  # noqa: C901
             OperationalError,
             MetaCommandException,
         ) as e:
-
             # Normal syntax errors, missing table, etc. should come back as
             # ProgrammingError. And the rest indicate something fundamentally
             # broken at the DBAPI layer.
@@ -185,8 +184,8 @@ def execute(self, line="", cell="", local_ns={}):  # noqa: C901
                 #
                 # "Restart Kernel" is too big of a hammer here.
                 #
-                conn._engine.dispose()
-                conn._session = None
+                conn.reset_connection_pool()
+
                 eprint(
                     "Encoutered the following unexpected exception while trying to run the statement."
                     " Closed the connection just to be safe. Re-run the cell to try again!\n\n"

noteable/sql/run.py

@@ -107,7 +107,7 @@ def _commit(conn, config):
 
     if _should_commit:
         try:
-            conn.session.execute("commit")
+            conn.sqla_connection.execute("commit")
         except sqlalchemy.exc.OperationalError:
             pass  # not all engines can commit
 
@@ -140,7 +140,6 @@ def _commit(conn, config):
 
 
 def run(conn, sql, config, user_namespace, skip_boxing_scalar_result: bool):
-
     if sql.strip():
         for statement in sqlparse.split(sql):
             first_word = sql.strip().split()[0].lower()
@@ -154,7 +153,7 @@ def run(conn, sql, config, user_namespace, skip_boxing_scalar_result: bool):
             bind_dict = {str(idx + 1): elem for (idx, elem) in enumerate(bind_list)}
 
             txt = sqlalchemy.sql.text(query)
-            result = conn.session.execute(txt, bind_dict)
+            result = conn.sqla_connection.execute(txt, bind_dict)
 
             _commit(conn=conn, config=config)
 

tests/conftest.py

@@ -18,7 +18,7 @@
     FileProgressUpdateContent,
     FileProgressUpdateMessage,
 )
-from noteable.sql.connection import Connection
+from noteable.sql.connection import Connection, bootstrap_duckdb
 from noteable.sql.magic import SqlMagic
 from noteable.sql.run import add_commit_blacklist_dialect
 
@@ -117,6 +117,14 @@ def with_empty_connections() -> None:
     Connection.connections = preexisting_connections
 
 
+@pytest.fixture
+def with_duckdb_bootstrapped(with_empty_connections) -> None:
+    # Normal magics bootstrapping will leave us with DuckDB connection populated.
+    bootstrap_duckdb()
+
+    yield
+
+
 @pytest.fixture
 def ipython_shell() -> InteractiveShell:
     return InteractiveShell()