Properly close database connections and transactions

This addresses the first item of issue #19.
blue-yonder · Jan 30, 2018 · bd4f1b8 · bd4f1b8
1 parent 20487fc
commit bd4f1b8
Show file tree

Hide file tree

Showing 3 changed files with 57 additions and 62 deletions.
diff --git a/postgraas_server/backends/postgres_cluster/postgres_cluster_driver.py b/postgraas_server/backends/postgres_cluster/postgres_cluster_driver.py
@@ -17,61 +17,61 @@ def _create_pg_connection(config):
 
 
 def check_db_or_user_exists(db_name, db_user, config):
-    con = _create_pg_connection(config)
-    cur = con.cursor()
-    cur.execute("SELECT 1 FROM pg_database WHERE datname='{}';".format(db_name))
-    db_exists = cur.fetchone() is not None
-    cur.execute("SELECT 1 FROM pg_roles WHERE rolname='{}';".format(db_user))
-    user = cur.fetchone()
-    user_exists = user is not None
-    return db_exists or user_exists
+    with _create_pg_connection(config) as con:
+        with con.cursor() as cur:
+            cur.execute("SELECT 1 FROM pg_database WHERE datname='{}';".format(db_name))
+            db_exists = cur.fetchone() is not None
+            cur.execute("SELECT 1 FROM pg_roles WHERE rolname='{}';".format(db_user))
+            user = cur.fetchone()
+            user_exists = user is not None
+            return db_exists or user_exists
 
 
 def create_postgres_db(connection_dict, config):
     if check_db_or_user_exists(connection_dict["db_name"], connection_dict["db_username"], config):
         raise ValueError("db or user already exists")
-    con = _create_pg_connection(config)
-    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-    cur = con.cursor()
-    create_role = "CREATE USER {db_username} WITH PASSWORD '{db_pwd}';".format(**connection_dict)
-    drop_role = "DROP ROLE {db_username};".format(**connection_dict)
-    grant_role = 'GRANT {db_username} TO "{postgraas_user}";'.format(
-        db_username=connection_dict['db_username'], postgraas_user=get_normalized_username(config['username'])
-    )
-    create_database = "CREATE DATABASE {db_name} OWNER {db_username};".format(**connection_dict)
-    try:
-        cur.execute(create_role)
-        cur.execute(grant_role)
-    except psycopg2.ProgrammingError as e:
-        raise ValueError(e.args[0])
-    # cleanup role in case database creation fails
-    # saidly 'CREATE DATABASE' cannot run inside a transaction block
-    try:
-        cur.execute(create_database)
-    except psycopg2.ProgrammingError as e:
-        cur.execute(drop_role)
-        raise ValueError(e.args[0])
+    with _create_pg_connection(config) as con:
+        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+        with con.cursor() as cur:
+            create_role = "CREATE USER {db_username} WITH PASSWORD '{db_pwd}';".format(**connection_dict)
+            drop_role = "DROP ROLE {db_username};".format(**connection_dict)
+            grant_role = 'GRANT {db_username} TO "{postgraas_user}";'.format(
+                db_username=connection_dict['db_username'], postgraas_user=get_normalized_username(config['username'])
+            )
+            create_database = "CREATE DATABASE {db_name} OWNER {db_username};".format(**connection_dict)
+            try:
+                cur.execute(create_role)
+                cur.execute(grant_role)
+            except psycopg2.ProgrammingError as e:
+                raise ValueError(e.args[0])
+            # cleanup role in case database creation fails
+            # saidly 'CREATE DATABASE' cannot run inside a transaction block
+            try:
+                cur.execute(create_database)
+            except psycopg2.ProgrammingError as e:
+                cur.execute(drop_role)
+                raise ValueError(e.args[0])
 
 
 def get_normalized_username(username):
     return username.split('@')[0]
 
 
 def delete_database(db_name, config):
-    con = _create_pg_connection(config)
-    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-    cur = con.cursor()
-    try:
-        cur.execute('''DROP DATABASE "{}";'''.format(db_name))
-    except psycopg2.ProgrammingError as e:
-        raise ValueError(e.args[0])
+    with _create_pg_connection(config) as con:
+        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+        with con.cursor() as cur:
+            try:
+                cur.execute('''DROP DATABASE "{}";'''.format(db_name))
+            except psycopg2.ProgrammingError as e:
+                raise ValueError(e.args[0])
 
 
 def delete_user(username, config):
-    con = _create_pg_connection(config)
-    con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
-    cur = con.cursor()
-    try:
-        cur.execute('''DROP USER "{}";'''.format(get_normalized_username(username)))
-    except psycopg2.ProgrammingError as e:
-        raise ValueError(e.args[0])
+    with _create_pg_connection(config) as con:
+        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
+        with con.cursor() as cur:
+            try:
+                cur.execute('''DROP USER "{}";'''.format(get_normalized_username(username)))
+            except psycopg2.ProgrammingError as e:
+                raise ValueError(e.args[0])
diff --git a/postgraas_server/management_resources.py b/postgraas_server/management_resources.py
@@ -77,23 +77,19 @@ def delete(self, id):
         if not entity:
             return {'status': 'failed', 'msg': 'Postgraas instance {} does not exist'.format(id)}
 
-        connection_error = None
         try:
-            conn = psycopg2.connect(
+            with psycopg2.connect(
                 user=entity.username,
                 password=args['db_pwd'],
                 host=current_app.postgraas_backend.master_hostname,
                 port=entity.port,
                 dbname=entity.db_name
-            )
-            conn.close()
+            ):
+                pass
         except Exception as ex:
-            connection_error = str(ex)
-
-        if connection_error is not None:
             return {
                 'status': 'failed',
-                'msg': 'Could not connect to postgres instance: {}'.format(connection_error)
+                'msg': 'Could not connect to postgres instance: {}'.format(str(ex))
             }
 
         if not current_app.postgraas_backend.exists(entity):

diff --git a/postgraas_server/prometheus_app.py b/postgraas_server/prometheus_app.py
@@ -20,26 +20,25 @@
 @contextmanager
 def db_connection(config):
     username = cfg.get_user(config)
-    connection = psycopg2.connect(
+    with psycopg2.connect(
         database=config['metadb']['db_name'],
         user=username,
         password=config['metadb']['db_pwd'],
         host=config['metadb']['host'],
         port=config['metadb']['port']
-    )
-    yield connection
-    connection.close()
+    ) as connection:
+        yield connection
 
 
 def do_count_query(config):
     with db_connection(config) as conn:
-        cursor = conn.cursor()
-        sql = (
-            "SELECT count(*) FROM pg_database "
-            "WHERE datistemplate = false;"
-        )
-        cursor.execute(sql)
-        (count,) = cursor.fetchone()
+        with conn.cursor() as cursor:
+            sql = (
+                "SELECT count(*) FROM pg_database "
+                "WHERE datistemplate = false;"
+            )
+            cursor.execute(sql)
+            (count,) = cursor.fetchone()
 
     return count