cevaris
diff --git a/‎functions/sql/mysql_sample.py
Lines changed: 16 additions & 8 deletions b/‎functions/sql/mysql_sample.py
Lines changed: 16 additions & 8 deletions
diff --git a/‎functions/sql/postgres_sample.py
Lines changed: 16 additions & 9 deletions b/‎functions/sql/postgres_sample.py
Lines changed: 16 additions & 9 deletions
@@ -16,8 +16,7 @@
 from os import getenv
 
 import pymysql
-
-is_production = getenv('SUPERVISOR_HOSTNAME') is not None
+from pymysql.err import OperationalError
 
 # TODO(developer): specify SQL connection details
 CONNECTION_NAME = getenv(
@@ -36,13 +35,9 @@
   'autocommit': True
 }
 
-if is_production:
-    mysql_config['unix_socket'] = \
-      '/cloudsql/' + CONNECTION_NAME
-
 # Create SQL connection globally to enable reuse
 # PyMySQL does not include support for connection pooling
-mysql_conn = pymysql.connect(**mysql_config)
+mysql_conn = None
 
 
 def __get_cursor():
@@ -53,12 +48,25 @@ def __get_cursor():
     """
     try:
         return mysql_conn.cursor()
-    except Exception:
+    except OperationalError:
         mysql_conn.ping(reconnect=True)
         return mysql_conn.cursor()
 
 
 def mysql_demo(request):
+    global mysql_conn
+
+    # Initialize connections lazily, in case SQL access isn't needed for this
+    # GCF instance. Doing so minimizes the number of active SQL connections,
+    # which helps keep your GCF instances under SQL connection limits.
+    if not mysql_conn:
+        try:
+            mysql_conn = pymysql.connect(**mysql_config)
+        except OperationalError:
+            # If production settings fail, use local development ones
+            mysql_config['unix_socket'] = f'/cloudsql/{CONNECTION_NAME}'
+            mysql_conn = pymysql.connect(**mysql_config)
+
     # Remember to close SQL resources declared while running this function.
     # Keep any declared in global scope (e.g. mysql_conn) for later reuse.
     with __get_cursor() as cursor:
 
@@ -15,10 +15,9 @@
 # [START functions_sql_postgres]
 from os import getenv
 
+from psycopg2 import OperationalError
 from psycopg2.pool import SimpleConnectionPool
 
-is_production = getenv('SUPERVISOR_HOSTNAME') is not None
-
 # TODO(developer): specify SQL connection details
 CONNECTION_NAME = getenv(
   'INSTANCE_CONNECTION_NAME',
@@ -30,27 +29,35 @@
 pg_config = {
   'user': DB_USER,
   'password': DB_PASSWORD,
-  'dbname': DB_NAME,
+  'dbname': DB_NAME
 }
 
-if is_production:
-    pg_config['host'] = '/cloudsql/' + CONNECTION_NAME
-else:
-    pg_config['host'] = 'localhost'
-
 # Connection pools reuse connections between invocations,
 # and handle dropped or expired connections automatically.
 pg_pool = None
 
 
+def __connect(host):
+    """
+    Helper function to connect to Postgres
+    """
+    global pg_pool
+    pg_config['host'] = host
+    pg_pool = SimpleConnectionPool(1, 1, **pg_config)
+
+
 def postgres_demo(request):
     global pg_pool
 
     # Initialize the pool lazily, in case SQL access isn't needed for this
     # GCF instance. Doing so minimizes the number of active SQL connections,
     # which helps keep your GCF instances under SQL connection limits.
     if not pg_pool:
-        pg_pool = SimpleConnectionPool(1, 1, **pg_config)
+        try:
+            __connect(f'/cloudsql/{CONNECTION_NAME}')
+        except OperationalError:
+            # If production settings fail, use local development ones
+            __connect('localhost')
 
     # Remember to close SQL resources declared while running this function.
     # Keep any declared in global scope (e.g. pg_pool) for later reuse.