Skip to content

Commit

Permalink
Setup base class
Browse files Browse the repository at this point in the history
This commit does several things at once:

- Setup BaseCommand class which is basically the base script that was done in shell script.
- Add setup_logigng method so that we can log errors and communicate with the user.
- Remove unused code
- Use CliContext class to create a click object. We create a click object so we can save the state across different CLI commands.
  • Loading branch information
Charles Short authored and portante committed Sep 30, 2020
1 parent 20bfd2b commit 0e7f706
Show file tree
Hide file tree
Showing 6 changed files with 212 additions and 64 deletions.
106 changes: 106 additions & 0 deletions lib/pbench/agent/base.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
import abc
import datetime
import os
import pathlib
import socket
import sys

import click

from pbench.agent import PbenchAgentConfig, fs


class BaseCommand(metaclass=abc.ABCMeta):
"""A base class used to define the command interface."""

def __init__(self, context):
self.context = context

self.config = PbenchAgentConfig(self.context.config)

self.pbench_run = self.get_path(os.environ.get("pbench_run", None))
if self.pbench_run is None:
self.pbench_run = self.config.pbench_run
if not self.pbench_run.exists():
click.secho(
f"[ERROR] the provided pbench run directory, {self.pbench_run}, does not exist"
)
sys.exit(1)

# the pbench temporary directory is always relative to the $pbench_run
# directory
self.pbench_tmp = self.pbench_run / "tmp"
try:
fs.safe_mkdir(self.pbench_tmp)
except Exception:
click.secho(f"[ERROR] unable to create TMP dir, {self.pbench_tmp}")
sys.exit(1)

# log file - N.B. not a directory
self.pbench_log = self.get_path(os.environ.get("pbench_log", None))
if self.pbench_log is None:
self.pbench_log = self.config.pbench_log

self.pbench_install_dir = self.get_path(
os.environ.get("pbench_install_dir", None)
)
if self.pbench_install_dir is None:
self.pbench_install_dir = self.config.pbench_install_dir
if not self.pbench_install_dir.exists():
click.secho(
f"[ERROR] pbench installation directory, {self.pbench_install_dir}, does not exist"
)
sys.exit(1)

self.pbench_bspp_dir = self.pbench_install_dir / "bench-scripts" / "postprocess"
self.pbench_lib_dir = self.pbench_install_dir / "lib"

self.ssh_opts = self.config.ssh_opts
os.environ["ssh_opts"] = self.ssh_opts

self.scp_opts = self.config.scp_opts
os.environ["scp_opts"] = self.scp_opts

os.environ["_pbench_debug_mode"] = "0"
if os.environ.get("_PBENCH_UNIT_TESTS"):
self.date = "1900-01-01T00:00:00"
self.date_suffix = "1900.01.01T00.00.00"
self.hostname = "testhost"
self.full_hostname = "testhost.example.com"
else:
self.date = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H:%M:%s")
self.date_suffix = datetime.datetime.utcnow().strftime("%Y-%m-%dT%H.%M.%s")
self.hostname = socket.gethostname()
self.full_hostname = socket.getfqdn()

# Backwards compatibility and for toolmeister
pbench_env = {
"date": self.date,
"date_suffix": self.date_suffix,
"hostname": self.hostname,
"full_hostname": self.full_hostname,
"pbench_run": str(self.pbench_run),
"pbench_install_dir": str(self.pbench_install_dir),
"pbench_tmp": str(self.pbench_tmp),
"pbench_log": str(self.pbench_log),
"pbench_bspp_dir": str(self.pbench_bspp_dir),
"pbench_lib_dir": str(self.pbench_lib_dir),
}
for k, v in pbench_env.items():
os.environ[k] = v

@abc.abstractmethod
def execute(self):
"""
This is the main method of the application
"""
pass

def get_path(self, path):
"""Converts a string path into a pathlib object"""
if path is None:
return path
elif not isinstance(path, pathlib.PurePath):
return pathlib.Path(path)
else:
return path
23 changes: 23 additions & 0 deletions lib/pbench/agent/fs.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
import errno
import os
import shutil


def safe_rmtree(directory):
"""Delete a directory, if it's present otherwise no-op"""
if os.path.exists(directory):
shutil.rmtree(directory)


def safe_mkdir(directory, clean=False):
"""Safely create a directory.
Ensures a directory is present. If it's not there, it is created. If it is, it's a no-op. If
clean is True, ensures the directory is empty.
"""
if clean:
safe_rmtree(directory)
try:
os.makedirs(directory)
except OSError as e:
if e.errno != errno.EEXIST:
raise
81 changes: 47 additions & 34 deletions lib/pbench/agent/utils.py
Original file line number Diff line number Diff line change
@@ -1,41 +1,54 @@
import logging
import os
import sys
import pathlib

import six
import colorlog

from pbench.agent.logger import logger

tools_group_prefix = "tools-v1"
def setup_logging(name=None, debug=False, logfile=None):
"""Setup logging for client
:param None: name of the python object
:param debug: Turn on debug logging
:param logfile: Logfile to write to
"""
if not name:
log = logging.getLogger() # root logger
else:
log = logging.getLogger(name)

# Make sh logging a bit less verbose
logging.getLogger("sh").setLevel(logging.WARNING)

def init_wrapper(config):
"""Initialize agent envrionment before running a command
# FIXME: since we dont do debugging level yet
log.setLevel(logging.INFO)

:param config: a configparser object
"""
if six.PY2:
logger.error("python3 is required, either directly or through SCL")
sys.exit(1)

pbench_run = pathlib.Path(config.rundir)
if pbench_run.exists():
# Its possible to run pbench without root
# but check to make sure that the rundir is writable
# before we do anything else
if os.access(pbench_run, os.W_OK) is not True:
logger.error("%s is not writable", pbench_run)
sys.exit(1)
pbench_tmp = pathlib.Path(pbench_run, "tmp")
if not pbench_run.exists():
# the pbench temporary directory is always relative to pbench run
pbench_tmp.mkdir(parents=True, exists_ok=True)
else:
logger.error("the provided pbench run directory %s does not exist.", pbench_run)
sys.exit(1)
pbench_install_dir = pathlib.Path(config.installdir)
if not pbench_install_dir.exists():
logger.error(
"pbench installation directory %s does not exist", pbench_install_dir
)
sys.exit(1)
format_str = "%(message)s"
date_format = "%Y-%m-%d %H:%M:%S"
cformat = "%(log_color)s" + format_str
colors = {
"DEBUG": "green",
"INFO": "cyan",
"WARNING": "bold_yellow",
"ERROR": "bold_red",
"CRITICAL": "bold_purple",
}
# Setup console
formatter = colorlog.ColoredFormatter(cformat, date_format, log_colors=colors)
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(formatter)

# Setup log file
if logfile is not None:
if not os.environ.get("_PBENCH_UNIT_TESTS"):
format_str = "[%(levelname)-1s][%(asctime)s.%(msecs)d] %(message)s"
else:
format_str = "[%(levelname)-1s][1900-01-01T00:00:00.000000] %(message)s"

_formatter = logging.Formatter(format_str)
log_file = logging.FileHandler(logfile)
log_file.setLevel(logging.DEBUG)
log_file.setFormatter(_formatter)
log.addHandler(log_file)

log.addHandler(stream_handler)

return log
21 changes: 21 additions & 0 deletions lib/pbench/cli/agent/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,21 @@
"""
Create a click context object that holds the state of the agent
invocation. The CliContext will keep track of passed parameters,
what command created it, which resources need to be cleaned up,
and etc.
We create an empty object at the begining and populate the object
with configuration, group names, at the beginning of the agent
execution.
"""

import click


class CliContext:
"""Inialize an empty click object"""

pass


pass_cli_context = click.make_pass_decorator(CliContext, ensure=True)
44 changes: 14 additions & 30 deletions lib/pbench/cli/agent/options.py
Original file line number Diff line number Diff line change
@@ -1,47 +1,31 @@
import os
import click


#
# Agent options
#
def pbench_upload_user(f):
return click.option(
"-u", "--user", "user", default="", help="Specify username for server upload"
)(f)
from pbench.cli.agent import CliContext


def pbench_server_prefix(f):
return click.option(
"-p", "--prefix", default="", help="Specify a prefix for server upload"
)(f)
def common_options(f):
f = _pbench_agent_config(f)
return f


def pbench_show_server(f):
return click.option("-S", "--show-server", required=False, help="Show server",)(f)
def _pbench_agent_config(f):
"""Option for agent configuration"""

def callback(ctx, param, value):
clictx = ctx.ensure_object(CliContext)
clictx.config = value
return value

#
# Default options
#
def pbench_agent_config(f):
"""Option for agent configuration"""
return click.option(
"-C",
"--config",
default=os.environ.get("_PBENCH_AGENT_CONFIG"),
envvar="_PBENCH_AGENT_CONFIG",
type=click.Path(exists=True),
callback=callback,
expose_value=False,
help=(
"Path to a pbench-agent config. If provided pbench will load "
"this config file first. By default is looking for config in "
"'_PBENCH_AGENT_CONFIG' envrionment variable."
),
)(f)


def pbench_agent_debug(f):
"""Turn on/off debug"""
return click.option(
"--debug",
default=False,
help="Enable or disable debug mode. Default is disabled",
)(f)
1 change: 1 addition & 0 deletions requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -13,3 +13,4 @@ python-pidfile
python-daemon
bottle
pyesbulk
colorlog

0 comments on commit 0e7f706

Please sign in to comment.