add reload() method

fiftyone/__public__.py

@@ -5,15 +5,39 @@
 | `voxel51.com <https://voxel51.com/>`_
 |
 """
-import fiftyone.core.config as _foc
-import fiftyone.core.odm as _foo
+config = None
+annotation_config = None
+evaluation_config = None
+app_config = None
 
-config = _foc.load_config()
-annotation_config = _foc.load_annotation_config()
-evaluation_config = _foc.load_evaluation_config()
-app_config = _foc.load_app_config()
 
-_foo.establish_db_conn(config)
+def reload(hard=False):
+    """Reloads the current database connection.
+
+    Args:
+        hard (False): whether to reconnect using the current in-memory config
+            values (False) or reload configs from environment variables (True)
+    """
+    global config
+    global annotation_config
+    global evaluation_config
+    global app_config
+
+    import fiftyone.core.config as foc
+    import fiftyone.core.odm as foo
+
+    if hard:
+        config = foc.load_config()
+        annotation_config = foc.load_annotation_config()
+        evaluation_config = foc.load_evaluation_config()
+        app_config = foc.load_app_config()
+
+    foo.disconnect_db()
+    foo.establish_db_conn(config)
+
+
+reload(hard=True)
+
 
 from .core.aggregations import (
     Aggregation,

fiftyone/core/odm/__init__.py

@@ -11,6 +11,7 @@
     aggregate,
     get_db_config,
     establish_db_conn,
+    disconnect_db,
     get_db_client,
     get_db_conn,
     get_async_db_client,

fiftyone/core/odm/database.py

@@ -14,7 +14,7 @@
 import asyncio
 from bson import json_util, ObjectId
 from bson.codec_options import CodecOptions
-from mongoengine import connect
+from mongoengine import connect, disconnect
 import mongoengine.errors as moe
 import motor.motor_asyncio as mtr
 
@@ -210,15 +210,7 @@ def establish_db_conn(config):
                 "own MongoDB instance or cluster "
             )
 
-    _client = pymongo.MongoClient(
-        **_connection_kwargs, appname=foc.DATABASE_APPNAME
-    )
-    _validate_db_version(config, _client)
-
-    # Register cleanup method
-    atexit.register(_delete_non_persistent_datasets_if_allowed)
-
-    connect(config.database_name, **_connection_kwargs)
+    _connect(config)
 
     config = get_db_config()
     if foc.CLIENT_TYPE != config.type:
@@ -227,20 +219,35 @@ def establish_db_conn(config):
             % (config.type, foc.CLIENT_TYPE)
         )
 
+    # Register cleanup method
+    atexit.register(_delete_non_persistent_datasets_if_allowed)
+
+
+def disconnect_db():
+    """Disconnects from the database, if necessary."""
+    _disconnect()
+
 
-def _connect():
+def _connect(config=None):
     global _client
+
     if _client is None:
         global _connection_kwargs
 
+        if config is None:
+            config = fo.config
+
         _client = pymongo.MongoClient(
             **_connection_kwargs, appname=foc.DATABASE_APPNAME
         )
-        connect(fo.config.database_name, **_connection_kwargs)
+        _validate_db_version(config, _client)
+
+        connect(config.database_name, **_connection_kwargs)
 
 
 def _async_connect(use_global=False):
     global _async_client
+
     if not use_global or _async_client is None:
         global _connection_kwargs
         client = mtr.AsyncIOMotorClient(
@@ -255,6 +262,16 @@ def _async_connect(use_global=False):
     return client
 
 
+def _disconnect():
+    global _client
+
+    if _client is not None:
+        _client.close()
+        _client = None
+
+        disconnect()
+
+
 def _delete_non_persistent_datasets_if_allowed():
     """Deletes all non-persistent datasets if and only if we are the only
     client currently connected to the database.