fix eating an exception

use dispose with sqlalchemy

Author: JohnMount

build/lib/data_algebra/BigQuery.py

@@ -131,13 +131,13 @@ def insert_table(
             self.drop_table(conn, table_name)
         else:
             table_exists = True
+            # noinspection PyBroadException
             try:
                 self.read_query(
                     conn,
                     "SELECT * FROM " + self.quote_table_name(table_name) + " LIMIT 1",
                 )
-                table_exists = True
-            except Exception as e:
+            except Exception:
                 table_exists = False
             if table_exists:
                 raise ValueError("table " + prepped_table_name + " already exists")
@@ -214,6 +214,7 @@ def example_handle():
     # assert os.path.isfile(credential_file)
     os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = credential_file
     # os.environ["GOOGLE_APPLICATION_CREDENTIALS"]  # trigger key error if not present
+    # noinspection PyBroadException
     try:
         data_catalog = "data-algebra-test"
         data_schema = "test_1"
@@ -222,5 +223,5 @@ def example_handle():
         ).db_handle(google.cloud.bigquery.Client())
         db_handle.db_model.prepare_connection(db_handle.conn)
         return db_handle
-    except Exception as e:
+    except Exception:
         return None

build/lib/data_algebra/MySQL.py

@@ -1,3 +1,8 @@
+"""
+Partial adapter of data algebra for MySQL. Not all data algebra operations are supported on this database at this time.
+"""
+
+
 import data_algebra.data_ops
 import data_algebra.db_model
 
@@ -54,7 +59,7 @@ class MySQLModel(data_algebra.db_model.DBModel):
        Assuming we are using a sqlalchemy engine as our connection.
     """
 
-    def __init__(self, *, supports_with=False):
+    def __init__(self):
         data_algebra.db_model.DBModel.__init__(
             self,
             string_type="CHAR",
@@ -80,8 +85,7 @@ def example_handle():
     """
     # TODO: parameterize this
     assert have_sqlalchemy
-    db_handle = MySQLModel().db_handle(
-        sqlalchemy.engine.create_engine("mysql+pymysql://jmount@localhost/jmount")
-    )
+    db_engine = sqlalchemy.engine.create_engine("mysql+pymysql://jmount@localhost/jmount")
+    db_handle = MySQLModel().db_handle(conn=db_engine, db_engine=db_engine)
     db_handle.db_model.prepare_connection(db_handle.conn)
     return db_handle

build/lib/data_algebra/PostgreSQL.py

@@ -43,8 +43,7 @@ def example_handle():
     """
     # TODO: parameterize this
     assert have_sqlalchemy
-    db_handle = PostgreSQLModel().db_handle(
-        sqlalchemy.engine.create_engine(r"postgresql://johnmount@localhost/johnmount")
-    )
+    db_engine = sqlalchemy.engine.create_engine(r"postgresql://johnmount@localhost/johnmount")
+    db_handle = PostgreSQLModel().db_handle(conn=db_engine, db_engine=db_engine)
     db_handle.db_model.prepare_connection(db_handle.conn)
     return db_handle

build/lib/data_algebra/SQLite.py

@@ -1,3 +1,8 @@
+
+"""
+Adapt data_algebra to SQLite database.
+"""
+
 import math
 import copy
 import numpy
@@ -19,6 +24,10 @@
 # at least capture to is_bad, which appears to not be implemented
 # unless we call prepare connection
 def _sqlite_is_bad_expr(dbmodel, expression):
+    """
+    Return SQL to check for bad values.
+    """
+
     return (
         "is_bad("
         + dbmodel.expr_to_sql(expression.args[0], want_inline_parens=False)
@@ -32,6 +41,10 @@ def _sqlite_is_bad_expr(dbmodel, expression):
 
 
 def _check_scalar_bad(x):
+    """
+    Return 1 if scalar value is none or nan, else 0.
+    """
+
     if x is None:
         return 1
     if not isinstance(x, numbers.Number):
@@ -42,13 +55,23 @@ def _check_scalar_bad(x):
 
 
 class MedianAgg:
+    """
+    Aggregate as median. SQLite user class.
+    """
+
     def __init__(self):
         self.collection = []
 
     def step(self, value):
+        """
+        Observe value
+        """
         self.collection.append(value)
 
     def finalize(self):
+        """
+        Return result.
+        """
         return numpy.median(self.collection)
 
 
@@ -75,6 +98,9 @@ def _unquote_identifier(self, s: str) -> str:
         return res
 
     def prepare_connection(self, conn):
+        """
+        Insert user functions into db.
+        """
         # https://docs.python.org/3/library/sqlite3.html#sqlite3.Connection.create_function
         conn.create_function("is_bad", 1, _check_scalar_bad)
         saw = set()
@@ -262,6 +288,9 @@ def natural_join_to_near_sql(
         sql_format_options=None,
         left_is_first=True
     ):
+        """
+        Translate a join into SQL, converting right and full joins to replacement code (as SQLite doesn't have these).
+        """
         if join_node.node_name != "NaturalJoinNode":
             raise TypeError(
                 "Expected join_node to be a data_algebra.data_ops.NaturalJoinNode)"

build/lib/data_algebra/SparkSQL.py

@@ -1,3 +1,8 @@
+
+"""
+SparkSQL adapter for the data algebra.
+"""
+
 import data_algebra.data_ops
 import data_algebra.db_model
 
@@ -36,8 +41,14 @@ def _sparksql_is_bad_expr(dbmodel, expression):
 
 
 # treat NaN as NULL, as Pandas has a hard time distinguishing the two
-def _sparksql_coalesce_expr(dbmodel, expression):
-    def coalesce_step(x):
+def _sparksql_coalesce_expr(dbmodel, expression) -> str:
+    """
+    Return coalesce expression.
+    """
+    def coalesce_step(x: str) -> str:
+        """
+        Return one caes of coalesce.
+        """
         assert isinstance(x, str)
         return f" WHEN ({x} IS NOT NULL) AND (NOT isNaN({x})) THEN {x} "
 
@@ -89,6 +100,9 @@ def _sparksql_db_mapv(dbmodel, expression):
 
 
 class SparkConnection:
+    """
+    Holder for spark conext and session as a connection (defines close).
+    """
     def __init__(self, *, spark_context, spark_session):
         assert have_Spark
         assert isinstance(spark_context, pyspark.context.SparkContext)
@@ -97,6 +111,9 @@ def __init__(self, *, spark_context, spark_session):
         self.spark_session = spark_session
 
     def close(self):
+        """
+        Stop context and release reference to context and session.
+        """
         if self.spark_context is not None:
             self.spark_context.stop()  # probably only for local demos
             self.spark_context = None

build/lib/data_algebra/db_model.py

@@ -1,3 +1,7 @@
+"""
+Base clase for SQL adapters for data algebra.
+"""
+
 import math
 import re
 from collections import OrderedDict
@@ -53,6 +57,7 @@ def __init__(
     def __str__(self):
         return self.__repr__()
 
+    # noinspection PyUnusedLocal
     def _repr_pretty_(self, p, cycle):
         """
         IPython pretty print, used at implicit print time
@@ -103,6 +108,7 @@ def _db_mean_expr(dbmodel, expression):
     )
 
 
+# noinspection PyUnusedLocal
 def _db_size_expr(dbmodel, expression):
     return "SUM(1)"
 
@@ -306,9 +312,6 @@ def _db_nunique_expr(dbmodel, expression):
     )
 
 
-# fns that had been in bigquery_user_fns
-
-
 def _as_int64(dbmodel, expression):
     return (
         "CAST("
@@ -487,7 +490,7 @@ def _base_Sunday(dbmodel, expression):
     "**": _db_pow_expr,
     "nunique": _db_nunique_expr,
     "mapv": _db_mapv,
-    # fns that had been in bigquery_user_fns
+    # additional fns
     "as_int64": _as_int64,
     "as_str": _as_str,
     "trimstr": _trimstr,
@@ -588,8 +591,13 @@ def __init__(
         self.union_all_term_start = union_all_term_start
         self.union_all_term_end = union_all_term_end
 
-    def db_handle(self, conn):
-        return DBHandle(db_model=self, conn=conn)
+    def db_handle(self, conn, db_engine=None):
+        """
+
+        :param conn: database connection
+        :param db_engine: optional sqlalchemy style engine (for closing)
+        """
+        return DBHandle(db_model=self, conn=conn, db_engine=db_engine)
 
     def prepare_connection(self, conn):
         pass
@@ -630,8 +638,8 @@ def read_query(self, conn, q):
     def table_exists(self, conn, table_name: str) -> bool:
         assert isinstance(table_name, str)
         q_table_name = self.quote_table_name(table_name)
-        # noinspection PyBroadException
         table_exists = True
+        # noinspection PyBroadException
         try:
             self.read_query(conn, "SELECT * FROM " + q_table_name + " LIMIT 1")
         except Exception:
@@ -1856,9 +1864,19 @@ def __repr__(self):
 
 
 class DBHandle:
-    def __init__(self, *, db_model: DBModel, conn):
+    """
+    Container for database connection handles.
+    """
+    def __init__(self, *, db_model: DBModel, conn, db_engine=None):
+        """
+
+        :param db_model: associated database model
+        :param conn: database connection
+        :param db_engine: optional sqlalchemy style engine (for closing)
+        """
         assert isinstance(db_model, DBModel)
         self.db_model = db_model
+        self.db_engine = db_engine
         self.conn = conn
 
     def __enter__(self):
@@ -1924,8 +1942,21 @@ def __repr__(self):
 
     def close(self) -> None:
         if self.conn is not None:
-            try:
-                self.conn.close()
-            except Exception:
-                pass
+            caught = None
+            if self.db_engine is not None:
+                # sqlalchemy style handle
+                # noinspection PyBroadException
+                try:
+                    self.db_engine.dispose()
+                except Exception as ex:
+                    caught = ex
+            else:
+                # noinspection PyBroadException
+                try:
+                    self.conn.close()
+                except Exception as ex:
+                    caught = ex
+            self.db_engine = None
             self.conn = None
+            if caught is not None:
+                raise ValueError('close caught: ' + str(caught))

build/lib/data_algebra/expr_rep.py

@@ -612,7 +612,7 @@ def mapv(self, value_map, default_value):
             "mapv", x=value_map, y=default_value, inline=False, method=True
         )
 
-    # fns that had been in bigquery_user_fns
+    # additional fns
 
     def as_int64(self):
         return self.__uop_expr__("as_int64")

build/lib/data_algebra/pandas_base.py

@@ -1,5 +1,5 @@
 """
-Base clas for adapters for Pandas-like APIs
+Base class for adapters for Pandas-like APIs
 """
 
 from abc import ABC
@@ -206,7 +206,7 @@ def populate_impl_map(data_model):
             a, b
         ),
         "mapv": _map_v,
-        # fns that had been in bigquery_user_fns
+        # additonal fns
         # x is a pandas Series
         "as_int64": lambda x: x.astype("int64").copy(),
         "as_str": lambda x: x.astype("str").copy(),

build/lib/data_algebra/test_util.py

@@ -153,7 +153,7 @@ def _run_handle_experiments(
     *,
     db_handle,
     data: Dict,
-    ops: ValueError,
+    ops: ViewRepresentation,
     sql_statements: Iterable[str],
     expect,
     float_tol: float = 1e-8,
@@ -164,14 +164,18 @@ def _run_handle_experiments(
     alter_cache: bool = True,
     test_direct_ops_path=False,
 ):
+    """
+    Run ops and sql_statements on db_handle, checking if result matches expect.
+    """
     assert isinstance(db_handle, data_algebra.db_model.DBHandle)
     assert isinstance(db_handle.db_model, data_algebra.db_model.DBModel)
+    assert isinstance(ops, ViewRepresentation)
     assert db_handle.conn is not None
     if isinstance(db_handle.db_model, data_algebra.SQLite.SQLiteModel):
         test_direct_ops_path = True
     db_handle_key = str(db_handle.db_model)
     sql_statements = list(sql_statements)
-    res_db_sql = [None] * len(sql_statements)
+    res_db_sql = list([None] * len(sql_statements))  # extra list() wrapper for PyCharm's type checker
     res_db_ops = None
     need_to_run = True
     dict_keys = list(data.keys())
@@ -192,7 +196,7 @@ def mk_key(ii):
             except KeyError:
                 pass
         need_to_run = test_direct_ops_path or numpy.any(
-            [resi is None for resi in res_db_sql]
+            [result_i is None for result_i in res_db_sql]
         )
     # generate any new needed results
     if need_to_run:
@@ -222,6 +226,7 @@ def mk_key(ii):
             raise ValueError(f"{db_handle} error in test " + str(caught))
     # check results
     for res in res_db_sql:
+        assert res is not None
         if not equivalent_frames(
             res,
             expect,

build/lib/data_algebra/util.py

@@ -40,6 +40,7 @@ def pandas_to_example_str(obj, *, local_data_model=None) -> str:
 
 def table_is_keyed_by_columns(table, column_names: Iterable[str]) -> bool:
     """
+    Check if a table is keyed by a given list of column names.
 
     :param table: pandas DataFrame
     :param column_names: list of column names
@@ -64,6 +65,11 @@ def table_is_keyed_by_columns(table, column_names: Iterable[str]) -> bool:
 
 # noinspection PyBroadException
 def _mk_type_conversion_table():
+    """
+    Build up conversion from type aliases we do not want into standard types. Eat any errors or warnings during table
+    construction.
+    """
+
     type_conversions_table = dict()
     # DeprecationWarning: `np.bool` is a deprecated alias for the builtin `bool`.
     # To silence this warning, use `bool` by itself. Doing this will not modify any behavior and is safe.