improves postgres compatibility + more granular feedback config (#798)

* improved support for postgres

* fix

* test updates

* update changelog

* typo

* deletes old comment

* adds a generic way to access the magic instance, the feedback config can hide the resultset footer

* shorter displaylimit footer

* feedback config controls displaying switching connections and current connection

* update changelog

* fix version

* test fix

* documentation updates

* cleanup

* moves data frame persist logic to connection, and creates method to handle sqlalchemy errors

* fix

* `--persist/--persist-replace` perform `ROLLBACK` automatically when needed

* changelog

* testing error when using --persist/--persist-replace with dbapi connections

* testing --persist uses error handling method

* adds some missing comments

* fix

* clean up __init__.py

* adds missing docstrings

* feedback doc update

CHANGELOG.md

@@ -1,8 +1,16 @@
 # CHANGELOG
 
-## 0.9.2dev
+## 0.10.0dev
 
+* [Fix] Perform `ROLLBACK` when SQLAlchemy raises `PendingRollbackError`
+* [Fix] Perform `ROLLBACK` when `psycopg2` raises `current transaction is aborted, commands ignored until end of transaction block`
+* [Fix] Perform `ROLLBACK` when `psycopg2` raises `server closed the connection unexpectedly` (#677)
 * [Fix] Fix a bug that caused a cell with a CTE to fail if it referenced a table/view with the same name as an existing snippet (#753)
+* [Fix] Shorter `displaylimit` footer
+* [API Change] `%config SqlMagic.feedback` now takes values `0` (disabled), `1` (normal), `2` (verbose)
+* [Fix] `ResultSet` footer only displayed when `feedback=2`
+* [Fix] Current connection and switching connections message only displayed when `feedback>=1`
+* [Fix] `--persist/--persist-replace` perform `ROLLBACK` automatically when needed
 
 ## 0.9.1 (2023-08-10)
 

doc/api/configuration.md

@@ -66,7 +66,7 @@ If you have autopandas set to true, the displaylimit option will not apply. You
 ## Changing configuration
 
 ```{code-cell} ipython3
-%config SqlMagic.feedback = False
+%config SqlMagic.feedback = 0
 ```
 
 ## `displaycon`
@@ -198,31 +198,22 @@ To unset:
 
 ## `feedback`
 
-Default: `True`
-
-Print number of rows affected by DML.
-
-```{code-cell} ipython3
-%config SqlMagic.feedback = True
+```{versionchanged} 0.10
+`feedback` takes values `0`, `1`, and `2` instead of `True`/`False`
 ```
 
-```{code-cell} ipython3
-%%sql
-CREATE TABLE my_points (x, y);
-INSERT INTO my_points VALUES (0, 0);
-INSERT INTO my_points VALUES (1, 1);
-```
+Default: `1`
 
-```{code-cell} ipython3
-%config SqlMagic.feedback = False
-```
+Control the quantity of messages displayed when performing certain operations. Each
+value enables the ones from previous values plus new ones:
 
-```{code-cell} ipython3
-%%sql
-CREATE TABLE more_points (x, y);
-INSERT INTO more_points VALUES (0, 0);
-INSERT INTO more_points VALUES (1, 1);
-```
+- `0`: Minimal feedback
+- `1`: Normal feedback (default)
+  - Connection name when switching
+  - Connection name when running a query
+  - Number of rows afffected by DML (e.g., `INSERT`, `UPDATE`, `DELETE`)
+- `2`: All feedback
+  - Footer to distinguish pandas/polars data frames from JupySQL's result sets
 
 ## `style`
 
@@ -244,6 +235,9 @@ print(res)
 
 ## `named_parameters`
 
+```{versionadded} 0.9
+```
+
 Default: `False`
 
 If True, it enables named parameters `:variable`. Learn more in the [tutorial.](../user-guide/template.md)
@@ -263,7 +257,7 @@ FROM languages
 WHERE rating > :rating
 ```
 
-## Loading configuration from a `pyproject.toml` file
+## Loading from `pyproject.toml`
 
 ```{versionadded} 0.9
 ```

src/sql/__init__.py

@@ -1,12 +1,7 @@
-from sql.magic import RenderMagic, SqlMagic, load_ipython_extension
-from sql.connection import PLOOMBER_DOCS_LINK_STR
+from sql.magic import load_ipython_extension
 
-__version__ = "0.9.2dev"
 
+__version__ = "0.10.0dev"
 
-__all__ = [
-    "RenderMagic",
-    "SqlMagic",
-    "load_ipython_extension",
-    "PLOOMBER_DOCS_LINK_STR",
-]
+
+__all__ = ["load_ipython_extension"]

src/sql/_current.py

@@ -0,0 +1,27 @@
+"""Get/set the current SqlMagic instance."""
+
+__sql_magic = None
+
+
+def _get_sql_magic():
+    """Returns the current SqlMagic instance."""
+    if __sql_magic is None:
+        raise RuntimeError("%sql has not been loaded yet. Run %load_ext sql")
+
+    return __sql_magic
+
+
+def _set_sql_magic(sql_magic):
+    """Sets the current SqlMagic instance."""
+    global __sql_magic
+    __sql_magic = sql_magic
+
+
+def _config_feedback_all():
+    """Returns True if the current feedback level is >=2"""
+    return _get_sql_magic().feedback >= 2
+
+
+def _config_feedback_normal_or_more():
+    """Returns True if the current feedback level is >=1"""
+    return _get_sql_magic().feedback >= 1

src/sql/connection/connection.py

@@ -4,10 +4,17 @@
 import os
 from difflib import get_close_matches
 import atexit
+from functools import partial
 
 import sqlalchemy
 from sqlalchemy.engine import Engine
-from sqlalchemy.exc import NoSuchModuleError, OperationalError, StatementError
+from sqlalchemy.exc import (
+    NoSuchModuleError,
+    OperationalError,
+    StatementError,
+    PendingRollbackError,
+    InternalError,
+)
 from IPython.core.error import UsageError
 import sqlglot
 import sqlparse
@@ -19,7 +26,8 @@
 from sql import exceptions, display
 from sql.error_message import detail
 from sql.parse import escape_string_literals_with_colon_prefix, find_named_parameters
-from sql.warnings import JupySQLQuotedNamedParametersWarning
+from sql.warnings import JupySQLQuotedNamedParametersWarning, JupySQLRollbackPerformed
+from sql import _current
 
 
 PLOOMBER_DOCS_LINK_STR = (
@@ -208,7 +216,10 @@ def set(
                 # passing an existing descriptor and not alias: use existing connection
                 elif existing and alias is None:
                     cls.current = existing
-                    display.message(f"Switching to connection {descriptor}")
+
+                    if _current._config_feedback_normal_or_more():
+                        display.message(f"Switching to connection {descriptor}")
+
                 # passing the same URL but different alias: create a new connection
                 elif existing is None or existing.alias != alias:
                     cls.current = cls.from_connect_str(
@@ -221,7 +232,7 @@ def set(
 
         else:
             if cls.connections:
-                if displaycon:
+                if displaycon and _current._config_feedback_normal_or_more():
                     cls.display_current_connection()
             elif os.getenv("DATABASE_URL"):
                 cls.current = cls.from_connect_str(
@@ -397,6 +408,11 @@ def _get_database_information(self):
         """
         pass
 
+    @abc.abstractmethod
+    def to_table(self, table_name, data_frame, if_exists, index):
+        """Create a table from a pandas DataFrame"""
+        pass
+
     def close(self):
         """Close the connection"""
         for rs in self._result_sets:
@@ -636,12 +652,8 @@ def _connection_execute(self, query, parameters=None):
         # not be a SELECT
         is_select = first_word_statement in {"select", "with", "from"}
 
-        if IS_SQLALCHEMY_ONE:
-            out = self._connection.execute(sqlalchemy.text(query), **parameters)
-        else:
-            out = self._connection.execute(
-                sqlalchemy.text(query), parameters=parameters
-            )
+        operation = partial(self._execute_with_parameters, query, parameters)
+        out = self._execute_with_error_handling(operation)
 
         if self._requires_manual_commit:
             # calling connection.commit() when using duckdb-engine will yield
@@ -666,6 +678,17 @@ def _connection_execute(self, query, parameters=None):
 
         return out
 
+    def _execute_with_parameters(self, query, parameters):
+        """Execute the query with the given parameters"""
+        if IS_SQLALCHEMY_ONE:
+            out = self._connection.execute(sqlalchemy.text(query), **parameters)
+        else:
+            out = self._connection.execute(
+                sqlalchemy.text(query), parameters=parameters
+            )
+
+        return out
+
     def raw_execute(self, query, parameters=None, with_=None):
         """Run the query without any preprocessing
 
@@ -733,6 +756,69 @@ def raw_execute(self, query, parameters=None, with_=None):
                         )
                 raise
 
+    def _execute_with_error_handling(self, operation):
+        """Execute a database operation and handle errors
+
+        Parameters
+        ----------
+        operation : callable
+            A callable that takes no parameters to execute a database operation
+        """
+        rollback_needed = False
+
+        try:
+            out = operation()
+
+        # this is a generic error but we've seen it in postgres. it helps recover
+        # from a idle session timeout (happens in psycopg 2 and psycopg 3)
+        except PendingRollbackError:
+            warnings.warn(
+                "Found invalid transaction. JupySQL executed a ROLLBACK operation.",
+                category=JupySQLRollbackPerformed,
+            )
+            rollback_needed = True
+
+        # postgres error
+        except InternalError as e:
+            # message from psycopg 2 and psycopg 3
+            message = (
+                "current transaction is aborted, "
+                "commands ignored until end of transaction block"
+            )
+            if type(e.orig).__name__ == "InFailedSqlTransaction" and message in str(
+                e.orig
+            ):
+                warnings.warn(
+                    (
+                        "Current transaction is aborted. "
+                        "JupySQL executed a ROLLBACK operation."
+                    ),
+                    category=JupySQLRollbackPerformed,
+                )
+                rollback_needed = True
+            else:
+                raise
+
+        # postgres error
+        except OperationalError as e:
+            # message from psycopg 2 and psycopg 3
+            message = "server closed the connection unexpectedly"
+
+            if type(e.orig).__name__ == "OperationalError" and message in str(e.orig):
+                warnings.warn(
+                    "Server closed connection. JupySQL executed a ROLLBACK operation.",
+                    category=JupySQLRollbackPerformed,
+                )
+                rollback_needed = True
+            else:
+                raise
+
+        if rollback_needed:
+            self._connection.rollback()
+            out = operation()
+
+        return out
+
     def _get_database_information(self):
         dialect = self._connection_sqlalchemy.dialect
 
@@ -789,6 +875,27 @@ def _error_invalid_connection_info(cls, e, connect_str):
         err.modify_exception = True
         return err
 
+    def to_table(self, table_name, data_frame, if_exists, index):
+        """Create a table from a pandas DataFrame"""
+        operation = partial(
+            data_frame.to_sql,
+            table_name,
+            self.connection_sqlalchemy,
+            if_exists=if_exists,
+            index=index,
+        )
+
+        try:
+            self._execute_with_error_handling(operation)
+        except ValueError:
+            raise exceptions.ValueError(
+                f"Table {table_name!r} already exists. Consider using "
+                "--persist-replace to drop the table before "
+                "persisting the data frame"
+            )
+
+        display.message_success(f"Success! Persisted {table_name} to the database.")
+
 
 class DBAPIConnection(AbstractConnection):
     """A connection object for generic DBAPI connections"""
@@ -879,6 +986,12 @@ def connection_sqlalchemy(self):
             "This feature is only available for SQLAlchemy connections"
         )
 
+    def to_table(self, table_name, data_frame, if_exists, index):
+        raise exceptions.NotImplementedError(
+            "--persist/--persist-replace is not available for DBAPI connections"
+            " (only available for SQLAlchemy connections)"
+        )
+
 
 def _check_if_duckdb_dbapi_connection(conn):
     """Check if the connection is a native duckdb connection"""

src/sql/exceptions.py

@@ -39,6 +39,7 @@ def _error(message):
 RuntimeError = exception_factory("RuntimeError")
 ValueError = exception_factory("ValueError")
 FileNotFoundError = exception_factory("FileNotFoundError")
+NotImplementedError = exception_factory("NotImplementedError")
 
 # The following are internal exceptions that should not be raised directly