ottertune/client/driver/utils.py

170 lines
5.4 KiB
Python
Raw Normal View History

import importlib
import os
from fabric.api import hide, local, settings, task
from fabric.api import get as _get, put as _put, run as _run, sudo as _sudo
2019-11-08 03:42:32 -08:00
dconf = None # pylint: disable=invalid-name
def load_driver_conf():
# The default config file is 'driver_config.py' but you can use
# set the env 'DRIVER_CONFIG' to the path of a different config
# file to override it.
2019-11-08 03:42:32 -08:00
global dconf # pylint: disable=global-statement,invalid-name
if not dconf:
driver_conf = os.environ.get('DRIVER_CONFIG', 'driver_config')
if driver_conf.endswith('.py'):
driver_conf = driver_conf[:-len('.py')]
mod = importlib.import_module(driver_conf)
dconf = mod
# Generate the login string of the host connection
if dconf.HOST_CONN == 'local':
login_str = 'localhost'
elif dconf.HOST_CONN == 'remote':
if not dconf.LOGIN_HOST:
raise ValueError("LOGIN_HOST must be set if HOST_CONN=remote")
login_str = dconf.LOGIN_HOST
if dconf.LOGIN_NAME:
login_str = '{}@{}'.format(dconf.LOGIN_NAME, login_str)
if dconf.LOGIN_PORT:
login_str += ':{}'.format(dconf.LOGIN_PORT)
elif dconf.HOST_CONN == 'docker':
if not dconf.CONTAINER_NAME:
raise ValueError("CONTAINER_NAME must be set if HOST_CONN=docker")
login_str = 'localhost'
else:
raise ValueError(("Invalid HOST_CONN: {}. Valid values are "
"'local', 'remote', or 'docker'.").format(dconf.HOST_CONN))
dconf.LOGIN = login_str
return dconf
def parse_bool(value):
if not isinstance(value, bool):
value = str(value).lower() == 'true'
return value
@task
2019-11-18 08:23:15 -08:00
def run(cmd, capture=True, **kwargs):
capture = parse_bool(capture)
try:
if dconf.HOST_CONN == 'remote':
res = _run(cmd, **kwargs)
elif dconf.HOST_CONN == 'local':
2019-11-18 08:23:15 -08:00
res = local(cmd, capture=capture, **kwargs)
else: # docker
opts = ''
cmdd = cmd
if cmd.endswith('&'):
cmdd = cmd[:-1].strip()
opts = '-d '
res = local('docker exec {} -ti {} /bin/bash -c "{}"'.format(
2019-11-18 08:23:15 -08:00
opts, dconf.CONTAINER_NAME, cmdd), capture=capture, **kwargs)
except TypeError as e:
err = str(e).strip()
if 'unexpected keyword argument' in err:
offender = err.rsplit(' ', 1)[-1][1:-1]
kwargs.pop(offender)
res = run(cmd, **kwargs)
else:
raise e
return res
@task
2019-11-18 08:23:15 -08:00
def sudo(cmd, user=None, capture=True, **kwargs):
capture = parse_bool(capture)
if dconf.HOST_CONN == 'remote':
res = _sudo(cmd, user=user, **kwargs)
elif dconf.HOST_CONN == 'local':
pre_cmd = 'sudo '
if user:
pre_cmd += '-u {} '.format(user)
2019-11-18 08:23:15 -08:00
res = local(pre_cmd + cmd, capture=capture, **kwargs)
else: # docker
user = user or 'root'
opts = '-ti -u {}'.format(user or 'root')
if user == 'root':
opts += ' -w /'
res = local('docker exec {} {} /bin/bash -c "{}"'.format(
2019-11-18 08:23:15 -08:00
opts, dconf.CONTAINER_NAME, cmd), capture=capture)
return res
@task
def get(remote_path, local_path, use_sudo=False):
use_sudo = parse_bool(use_sudo)
if dconf.HOST_CONN == 'remote':
res = _get(remote_path, local_path, use_sudo=use_sudo)
elif dconf.HOST_CONN == 'local':
pre_cmd = 'sudo ' if use_sudo else ''
opts = '-r' if os.path.isdir(remote_path) else ''
res = local('{}cp {} {} {}'.format(pre_cmd, opts, remote_path, local_path))
else: # docker
res = local('docker cp {}:{} {}'.format(dconf.CONTAINER_NAME, remote_path, local_path))
return res
@task
def put(local_path, remote_path, use_sudo=False):
use_sudo = parse_bool(use_sudo)
if dconf.HOST_CONN == 'remote':
res = _put(local_path, remote_path, use_sudo=use_sudo)
elif dconf.HOST_CONN == 'local':
pre_cmd = 'sudo ' if use_sudo else ''
opts = '-r' if os.path.isdir(local_path) else ''
res = local('{}cp {} {} {}'.format(pre_cmd, opts, local_path, remote_path))
else: # docker
res = local('docker cp {} {}:{}'.format(local_path, dconf.CONTAINER_NAME, remote_path))
return res
@task
def run_sql_script(scriptfile, *args):
if dconf.DB_TYPE == 'oracle':
if dconf.HOST_CONN != 'local':
scriptdir = '/home/oracle/oracleScripts'
remote_path = os.path.join(scriptdir, scriptfile)
if not file_exists(remote_path):
run('mkdir -p {}'.format(scriptdir))
put(os.path.join('./oracleScripts', scriptfile), remote_path)
sudo('chown -R oracle:oinstall /home/oracle/oracleScripts')
res = run('sh {} {}'.format(remote_path, ' '.join(args)))
else:
raise Exception("Database Type {} Not Implemented !".format(dconf.DB_TYPE))
return res
@task
def file_exists(filename):
2019-11-08 03:42:32 -08:00
with settings(warn_only=True), hide('warnings'): # pylint: disable=not-context-manager
res = run('[ -f {} ]'.format(filename))
return res.return_code == 0
@task
def dir_exists(dirname):
2019-11-08 03:42:32 -08:00
with settings(warn_only=True), hide('warnings'): # pylint: disable=not-context-manager
res = run('[ -d {} ]'.format(dirname))
return res.return_code == 0
class FabricException(Exception):
pass