From 1b3adae0915ab9238ffb7b4b815a8c4bdc756f2e Mon Sep 17 00:00:00 2001 From: ek Date: Mon, 10 Jun 2024 16:54:54 +0530 Subject: [PATCH] Context manager support with connection object --- ibm_db_ctx.py | 25 ++++++++ ibm_db_tests/test_context_ConnActive.py | 61 +++++++++++++++++++ ibm_db_tests/test_context_ConnDb.py | 34 +++++++++++ .../test_context_ConnDbUncatalogedConn.py | 36 +++++++++++ setup.py | 2 +- 5 files changed, 157 insertions(+), 1 deletion(-) create mode 100644 ibm_db_ctx.py create mode 100644 ibm_db_tests/test_context_ConnActive.py create mode 100644 ibm_db_tests/test_context_ConnDb.py create mode 100644 ibm_db_tests/test_context_ConnDbUncatalogedConn.py diff --git a/ibm_db_ctx.py b/ibm_db_ctx.py new file mode 100644 index 00000000..18c99d22 --- /dev/null +++ b/ibm_db_ctx.py @@ -0,0 +1,25 @@ +import ibm_db + +class Db2connect: + """Context manager to handle connections to DB2.""" + #__cxn: Optional['IBM_DBConnection'] + + def __init__(self, dsn: str, username: str, password: str) -> None: + """Instantiate a DB2 connection.""" + #print("init method called") + self.__dsn = dsn + self.__username = username + self.__password = password + self.__cxn = None + + def __enter__(self) -> 'IBM_DBConnection': + """Connect to DB2.""" + self.__cxn = ibm_db.connect(self.__dsn, '', '') + #print("enter method called") + return self.__cxn + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + """Disconnect from DB2.""" + #print("exit method called") + ibm_db.close(self.__cxn) + self.__cxn = None diff --git a/ibm_db_tests/test_context_ConnActive.py b/ibm_db_tests/test_context_ConnActive.py new file mode 100644 index 00000000..d6bd85ca --- /dev/null +++ b/ibm_db_tests/test_context_ConnActive.py @@ -0,0 +1,61 @@ +# +# Licensed Materials - Property of IBM +# +# (c) Copyright IBM Corp. 2007-2008 +# + +from __future__ import print_function +import sys +import unittest +import ibm_db +import config +import platform +import ibm_db_ctx + +from testfunctions import IbmDbTestFunctions + +class IbmDbTestCase(unittest.TestCase): + @unittest.skipIf(platform.system() == 'z/OS',"Test fails with z/OS ODBC driver") + def test_context_ConnActive(self): + obj = IbmDbTestFunctions() + obj.assert_expect(self.run_test_context_ConnActive) + + def run_test_context_ConnActive(self): + conn = None + is_alive = ibm_db.active(conn) + if is_alive: + print("Is active") + else: + print("Is not active") + + with ibm_db_ctx.Db2connect(config.database, config.user, config.password) as conn: + is_alive = ibm_db.active(conn) + if is_alive: + print("Is active") + else: + print("Is not active") + + is_alive = ibm_db.active(conn) + if is_alive: + print("Is active") + else: + print("Is not active") + + +#__END__ +#__LUW_EXPECTED__ +#Is not active +#Is active +#Is not active +#__ZOS_EXPECTED__ +#Is not active +#Is active +#Is not active +#__SYSTEMI_EXPECTED__ +#Is not active +#Is active +#Is not active +#__IDS_EXPECTED__ +#Is not active +#Is active +#Is not active \ No newline at end of file diff --git a/ibm_db_tests/test_context_ConnDb.py b/ibm_db_tests/test_context_ConnDb.py new file mode 100644 index 00000000..ae635323 --- /dev/null +++ b/ibm_db_tests/test_context_ConnDb.py @@ -0,0 +1,34 @@ +# +# Licensed Materials - Property of IBM +# +# (c) Copyright IBM Corp. 2007-2008 +# + +from __future__ import print_function +import sys +import unittest +import ibm_db +import config +from testfunctions import IbmDbTestFunctions +import ibm_db_ctx + +class IbmDbTestCase(unittest.TestCase): + + def test_context_ConnDb(self): + obj = IbmDbTestFunctions() + obj.assert_expect(self.run_test_context_ConnDb) + + def run_test_context_ConnDb(self): + with ibm_db_ctx.Db2connect(config.database, config.user, config.password) as conn: + print("Connection succeeded.") + + +#__END__ +#__LUW_EXPECTED__ +#Connection succeeded. +#__ZOS_EXPECTED__ +#Connection succeeded. +#__SYSTEMI_EXPECTED__ +#Connection succeeded. +#__IDS_EXPECTED__ +#Connection succeeded. \ No newline at end of file diff --git a/ibm_db_tests/test_context_ConnDbUncatalogedConn.py b/ibm_db_tests/test_context_ConnDbUncatalogedConn.py new file mode 100644 index 00000000..44e3c515 --- /dev/null +++ b/ibm_db_tests/test_context_ConnDbUncatalogedConn.py @@ -0,0 +1,36 @@ +# +# Licensed Materials - Property of IBM +# +# (c) Copyright IBM Corp. 2007-2008 +# + +from __future__ import print_function +import sys +import unittest +import ibm_db +import config +from testfunctions import IbmDbTestFunctions +import ibm_db_ctx + +class IbmDbTestCase(unittest.TestCase): + + def test_context_ConnDbUncatalogedConn(self): + obj = IbmDbTestFunctions() + obj.assert_expect(self.run_test_context_ConnDbUncatalogedConn) + + def run_test_context_ConnDbUncatalogedConn(self): + conn_str = "DATABASE=%s;HOSTNAME=%s;PORT=%d;PROTOCOL=TCPIP;UID=%s;PWD=%s;" % (config.database, config.hostname, config.port, config.user, config.password) + with ibm_db_ctx.Db2connect(conn_str,'','') as conn: + print("Connection succeeded.") + + + +#__END__ +#__LUW_EXPECTED__ +#Connection succeeded. +#__ZOS_EXPECTED__ +#Connection succeeded. +#__SYSTEMI_EXPECTED__ +#Connection succeeded. +#__IDS_EXPECTED__ +#Connection succeeded. \ No newline at end of file diff --git a/setup.py b/setup.py index f4b0414a..e7325338 100644 --- a/setup.py +++ b/setup.py @@ -486,7 +486,7 @@ def print_exception( e, url): (get_python_lib(), ['./LICENSE']), (get_python_lib(), ['./config.py.sample'])] -modules = ['ibm_db_dbi', 'testfunctions', 'ibmdb_tests'] +modules = ['ibm_db_dbi', 'testfunctions', 'ibmdb_tests', 'ibm_db_ctx'] if 'zos' == sys.platform: ext_modules = _ext_modules(os.path.join(os.getcwd(), include_dir), library, ibm_db_lib, ibm_db_lib_runtime)