Fix db connection leaks

Author: DraTeots

python/rcdb/cli/db.py

@@ -24,6 +24,7 @@ def db_command(ctx):
             print("ERROR connection string is missing.")
             exit(1)
         provider.connect(connection_str, check_version=False)
+        ctx.call_on_close(provider.disconnect)
         query = select(SchemaVersion).order_by(SchemaVersion.version.desc())
         schema_version, = provider.session.execute(query).first()
         print("Schema version: {} - '{}'".format(schema_version.version, schema_version.comment))
@@ -88,6 +89,7 @@ def _print_sqlite_table_sizes(engine):
 @pass_rcdb_context
 def update(context):
     provider = RCDBProvider(context.connection_str, check_version=False)
+    click.get_current_context().call_on_close(provider.disconnect)
 
     # Check something exists
     if not sqlalchemy.inspect(provider.engine).has_table(SchemaVersion.__tablename__):
@@ -117,6 +119,7 @@ def update(context):
         return
 
     provider = RCDBProvider(context.connection_str, check_version=False)
+    click.get_current_context().call_on_close(provider.disconnect)
 
     # Create alias table
     Alias.__table__.create(provider.engine)
@@ -181,6 +184,7 @@ def init(context, drop_all, no_defaults, confirm, add_tests, add_cpp_tests):
     # That we will need for DB
     metadata = rcdb.model.Base.metadata
     provider = RCDBProvider(context.connection_str, check_version=False)
+    click.get_current_context().call_on_close(provider.disconnect)
 
     # Drop all if needed
     if drop_all:

python/rcdb/cli/rp.py

@@ -55,6 +55,7 @@ def add(context, name, description, run_min, run_max, start_date, end_date):
 
     provider = RCDBProvider()
     provider.connect(connection_str, check_version=False)
+    click.get_current_context().call_on_close(provider.disconnect)
     session = provider.session
 
     # Convert date strings to Date objects if provided
@@ -104,6 +105,7 @@ def rm(context, period_id, yes):
 
     provider = RCDBProvider()
     provider.connect(connection_str, check_version=False)
+    click.get_current_context().call_on_close(provider.disconnect)
     session = provider.session
 
     rp_item = session.query(RunPeriod).filter(RunPeriod.id == period_id).one_or_none()
@@ -142,6 +144,7 @@ def update(context, period_id, name, description, run_min, run_max, start_date,
 
     provider = RCDBProvider()
     provider.connect(connection_str, check_version=False)
+    click.get_current_context().call_on_close(provider.disconnect)
     session = provider.session
 
     rp_item = session.query(RunPeriod).filter(RunPeriod.id == period_id).one_or_none()

python/rcdb/provider.py

@@ -133,18 +133,28 @@ def connect(self, connection_string="mysql+pymysql://rcdb@127.0.0.1/rcdb", check
         self._connection_string = connection_string
 
         if check_version:
-            db_version = self.get_schema_version()
-            if db_version != rcdb.SQL_SCHEMA_VERSION:
-                message = "SQL schema version doesn't match. " \
-                          "Retrieved DB version is {0}, required version is {1}. " \
-                    .format(db_version, rcdb.SQL_SCHEMA_VERSION)
-                if db_version is not None and db_version < rcdb.SQL_SCHEMA_VERSION:
-                    message += "The database schema is older than this RCDB version. " \
-                               "Run 'rcdb db update' to migrate the database to the current schema."
-                else:
-                    message += "The database schema is newer than this RCDB version. " \
-                               "Update your RCDB installation to match the database."
-                raise rcdb.errors.SqlSchemaVersionError(message)
+            # This version query is the first statement that actually opens a
+            # DB connection. If it fails (unreachable host, access denied, wrong
+            # schema, ...) the engine/session we just built must be disposed
+            # here: the exception propagates out of connect()/__init__, so the
+            # caller never receives this provider to disconnect() it themselves,
+            # and the engine's open connection would otherwise leak until GC.
+            try:
+                db_version = self.get_schema_version()
+                if db_version != rcdb.SQL_SCHEMA_VERSION:
+                    message = "SQL schema version doesn't match. " \
+                              "Retrieved DB version is {0}, required version is {1}. " \
+                        .format(db_version, rcdb.SQL_SCHEMA_VERSION)
+                    if db_version is not None and db_version < rcdb.SQL_SCHEMA_VERSION:
+                        message += "The database schema is older than this RCDB version. " \
+                                   "Run 'rcdb db update' to migrate the database to the current schema."
+                    else:
+                        message += "The database schema is newer than this RCDB version. " \
+                                   "Update your RCDB installation to match the database."
+                    raise rcdb.errors.SqlSchemaVersionError(message)
+            except Exception:
+                self.disconnect()
+                raise
 
     # ------------------------------------------------
     # Closes connection to data
@@ -164,6 +174,11 @@ def disconnect(self):
         if self.engine is not None:
             self.engine.dispose()
 
+    # ``close`` is an alias for ``disconnect``: a lot of calling code (and other
+    # resource-holding objects) expects a ``close()`` method, while existing
+    # RCDB code and users rely on the historical ``disconnect()`` name. Keep both.
+    close = disconnect
+
     # -------------------------------------------------
     # indicates ether the connection is open or not
     # -------------------------------------------------

python/rcdb/web/modules/runs.py

@@ -157,8 +157,11 @@ def elog(run_number):
         from urllib.request import urlopen
         from urllib.error import HTTPError
     try:
-        elog_json = urlopen('https://logbooks.jlab.org/api/elog/entries?book=hdrun&title=Run_{}&limit=1'
-                            .format(run_number)).read()
+        # Context-manage the response so the underlying socket is closed
+        # deterministically instead of leaking until GC.
+        with urlopen('https://logbooks.jlab.org/api/elog/entries?book=hdrun&title=Run_{}&limit=1'
+                     .format(run_number)) as response:
+            elog_json = response.read()
     except HTTPError as e:
         return jsonify(stat=str(e.code))