Fix cleanup of connections.

Author: thatnerdjosh

shared/wait-for-psql.py

@@ -16,18 +16,25 @@
 
     args = arg_parser.parse_args()
 
+    conn = None
     start_time = time.time()
+    error = ''
+    
     while (time.time() - start_time) < args.timeout:
         try:
+            if conn is not None:
+                conn.close()
+                conn = None  # Reset conn to None
+
             conn = psycopg2.connect(user=args.db_user, host=args.db_host, port=args.db_port, password=args.db_password, dbname=args.db_name)
-            error = ''
             break
         except psycopg2.OperationalError as e:
             error = e
-        else:
-            conn.close()
+
         time.sleep(1)
 
-    if error:
+    if conn:
+        conn.close()
+    else:
         print("Database connection failure: %s" % error, file=sys.stderr)
         sys.exit(1)