DB - Add method to use a database session with a context block.

Co-authored-by: aaarendt <[email protected]>

snowexsql/db.py

@@ -108,6 +108,18 @@ def get_db(credentials_path=None, return_metadata=False):
         return engine, session
 
 
+@contextmanager
+def db_session_with_credentials(credentials_path=None):
+    """
+    Helper method to allow database session with a context block.
+
+    Args:
+        credentials_path (string): Full path to credentials file (Optional)
+
+    """
+    yield get_db(credentials_path)
+
+
 def get_table_attributes(DataCls):
     """
     Returns a list of all the table columns to be used for each entry

tests/test_db.py

@@ -1,8 +1,8 @@
 import pytest
 import snowexsql
-from snowexsql.db import (
-    DB_CONNECTION_PROTOCOL, db_connection_string, get_db, load_credentials
-)
+from snowexsql.db import (DB_CONNECTION_PROTOCOL, db_connection_string,
+                          db_session_with_credentials, get_db,
+                          load_credentials)
 from sqlalchemy import Engine, MetaData
 from sqlalchemy.orm import Session
 
@@ -52,6 +52,16 @@ def test_returns_metadata(self):
             MetaData
         )
 
+    @pytest.mark.usefixtures('db_connection_string_patch')
+    def test_db_session_with_credentials(self, monkeypatch, test_db_info):
+        engine, session = None, None
+        with db_session_with_credentials() as (test_engine, test_session):
+            engine = test_engine
+            session = test_session
+
+        assert isinstance(engine, Engine)
+        assert isinstance(session, Session)
+
     @pytest.mark.usefixtures('db_connection_string_patch')
     @pytest.mark.parametrize("return_metadata, expected_objs", [
         (False, 2),