autotest: tko: connect via cloudsql_proxy if possible
We have setup cloudsql proxy on all autotest servers. In this change, we
try to connect the cloudsql instance by using the proxy socket. When
failed, it fallbacks to connect using username/password/host.
BUG=chromium:868052
TEST=Tested on my test autotest server:
1. setup shadow_config.ini properly with global_db_proxy_socket
information.
2. Ran below code and ensure no exception.
import common
from autotest_lib.tko import db
sql = db.db()
3. Removed global_db_proxy_socket from shadow_config.ini
4. Re-ran above code and ensure connection failed because my testing
server IP isn't in the whitelist.
Change-Id: I39738d39b4c3bb10471cd705abddba041f423b61
Reviewed-on: https://chromium-review.googlesource.com/1164064
Commit-Ready: Congbin Guo <guocb@chromium.org>
Tested-by: Congbin Guo <guocb@chromium.org>
Reviewed-by: Congbin Guo <guocb@chromium.org>
diff --git a/tko/db.py b/tko/db.py
index d8e2720..2b3b58b 100644
--- a/tko/db.py
+++ b/tko/db.py
@@ -67,10 +67,10 @@
"""Data access."""
def __init__(self, debug=False, autocommit=True, host=None,
- database=None, user=None, password=None):
+ database=None, user=None, password=None, proxy_socket=None):
self.debug = debug
self.autocommit = autocommit
- self._load_config(host, database, user, password)
+ self._load_config(host, database, user, password, proxy_socket)
self.con = None
self._init_db()
@@ -92,7 +92,7 @@
self.machine_group = {}
- def _load_config(self, host, database, user, password):
+ def _load_config(self, host, database, user, password, proxy_socket):
"""Loads configuration settings required to connect to the database.
This will try to connect to use the settings prefixed with global_db_.
@@ -101,6 +101,8 @@
If parameters are supplied, these will be taken instead of the values
in global_config.
+ When proxy_socket is set, 'host' will be bypassed.
+
@param host: If set, this host will be used, if not, the host will be
retrieved from global_config.
@param database: If set, this database will be used, if not, the
@@ -109,6 +111,8 @@
user will be retrieved from global_config.
@param password: If set, this password will be used, if not, the
password will be retrieved from global_config.
+ @param proxy_socket: If set, this proxy_socket will be used, if not, the
+ proxy_socket will be retrieved from global_config.
"""
database_settings = database_settings_helper.get_global_db_config()
@@ -116,9 +120,10 @@
self.host = host or database_settings['HOST']
self.database = database or database_settings['NAME']
- # grab the user and password
+ # grab authentication information
self.user = user or database_settings['USER']
self.password = password or database_settings['PASSWORD']
+ self.proxy_socket = proxy_socket or database_settings['PROXY_SOCKET']
# grab the timeout configuration
self.query_timeout =(
@@ -157,7 +162,8 @@
# create the db connection and cursor
self.con = self.connect(self.host, self.database,
- self.user, self.password, self.port)
+ self.user, self.password, self.port,
+ self.proxy_socket)
self.cur = self.con.cursor()
@@ -168,18 +174,27 @@
@retry.retry(driver.OperationalError, timeout_min=10,
delay_sec=5, callback=_connection_retry_callback)
- def connect(self, host, database, user, password, port):
+ def connect(self, host, database, user, password, port, proxy_socket):
"""Open and return a connection to mysql database."""
connection_args = {
- 'host': host,
- 'user': user,
'db': database,
+ 'user': user,
'passwd': password,
'connect_timeout': 20,
}
if port:
connection_args['port'] = int(port)
- return driver.connect(**connection_args)
+ # Connect using proxy socket if possible.
+ if proxy_socket:
+ try:
+ return driver.connect(unix_socket=proxy_socket,
+ **connection_args)
+ # pylint: disable=catching-non-exception
+ except driver.OperationalError:
+ # Fallback to connect using user/host/password.
+ pass
+
+ return driver.connect(host=host, **connection_args)
def run_with_retry(self, function, *args, **dargs):