<commit_before>"""Conversion functions for pre- and postprocessing of bib(la)tex fields."""

import bibpy.compat
import bibpy.date
import bibpy.error
import bibpy.grammar
import bibpy.lexers
import bibpy.name
import bibpy.parser
import calendar
import re

_MONTH_ABBREVIATIONS = [
    'jan', 'feb', 'mar', 'apr', 'may', 'jun',
    'jul', 'aug', 'sep', 'oct', 'nov', 'dec'
]

# TODO: Let users split on default names and custom ones
_SPLIT_NAMES = frozenset([
    'author', 'afterword', 'bookauthor', 'commentator',
    'editor', 'editora', 'editorb', 'editorc', 'foreword',
    'holder', 'introduction', 'language', 'origpublisher',
    'publisher', 'shortauthor', 'shorteditor', 'translator'
])

def postprocess_braces(value, **options):
    """Remove any braces from a string value."""
    if bibpy.compat.is_string(value):
        return "".join([e for e in bibpy.parser.parse_braced_expr(value)
                        if e not in '{}'])

    return value

def postprocess_namelist(field, names, **options):
    """Convert a string of authors to a list."""
    if not names:
        return []

    # First, split on zero brace-level 'and'
    names = bibpy.lexers.lex_namelist(names)

    # Second, if requested, parse each name
    if field in options.get('split_names', []):
        return [postprocess_name(field, n) for n in names]

    # Remove any leftover curly braces after splitting
    if not options.get('split_names', False):
        names = [re.sub('\{(.+?)\}', '\\1', name) for name in names]

    return names

def postprocess_name(field, author, **options):
    """Attempts to split an author name into first, middle and last name."""
    if author and bibpy.compat.is_string(author):
        return bibpy.name.Name.fromstring(author)
    else:
        return author

def postprocess_keywords(field, keywords, **options):
    """Convert a string of keywords to a list."""
    if not keywords:
        return []

    return [keyword.strip() for keyword in keywords.split(';')
            if keyword.strip()]

def postprocess_int(field, value, **options):
    """Convert a string to an integer."""
    try:
        return int(value)
    except (ValueError, TypeError):
        return value

def postprocess_date(field, datestring, **options):
    """Convert a string to a bibpy.date.DateRange."""
    if not datestring:
        return bibpy.date.DateRange((None, None, None), (None, None, None),
                                    False)

    return bibpy.date.DateRange.fromstring(datestring)

def get_month_name(i):
    """Return the name of a month from its one-based index ."""
    return calendar.month_name[i]

def postprocess_month(field, month, **options):
    """Convert a month number to its name."""
    try:
        i = int(month)

        if i not in range(1, 13):
            return month

        return get_month_name(i)
    except (ValueError, IndexError):
        month_name = month.lower()

        if month_name in _MONTH_ABBREVIATIONS:
            # The first element of calendar.month_name is an empty string
            return get_month_name(_MONTH_ABBREVIATIONS.index(month_name) + 1)

    return month

def postprocess_keylist(field, keylist, **options):
    """Convert a comma-separated string of keys to a list."""
    if not keylist:
        return []

    return list(filter(None, [key.strip() for key in keylist.split(',')]))

def postprocess_pages(field, pages, **options):
    """Convert a page range to a 2-element tuple."""
    values = re.split('\-+', pages)

    if len(values) == 2:
        try:
            return (int(values[0]), int(values[1]))
        except (ValueError, TypeError):
            return pages
    else:
        return pages

# A dictionary of fields as keys and the functions that postprocess them as
# values, e.g. 'year' should be converted to an integer etc.
postprocess_functions = {'address':       postprocess_namelist,
                         'afterword':     postprocess_namelist,
                         'author':        postprocess_namelist,
                         'bookauthor':    postprocess_namelist,
                         'commentator':   postprocess_namelist,
                         'date':          postprocess_date,
                         'editor':        postprocess_namelist,
                         'editora':       postprocess_namelist,
                         'editorb':       postprocess_namelist,
                         'editorc':       postprocess_namelist,
                         'eventdate':     postprocess_date,
                         'foreword':      postprocess_namelist,
                         'holder':        postprocess_namelist,
                         'institution':   postprocess_namelist,
                         'introduction':  postprocess_namelist,
                         'keywords':      postprocess_keywords,
                         'language':      postprocess_namelist,
                         'location':      postprocess_namelist,
                         'month':         postprocess_month,
                         'organization':  postprocess_namelist,
                         'origdate':      postprocess_date,
                         'origlocation':  postprocess_namelist,
                         'origpublisher': postprocess_namelist,
                         'pages':         postprocess_pages,
                         'publisher':     postprocess_namelist,
                         'related':       postprocess_keylist,
                         'school':        postprocess_namelist,
                         'shortauthor':   postprocess_namelist,
                         'shorteditor':   postprocess_namelist,
                         'translator':    postprocess_namelist,
                         'urldate':       postprocess_date,
                         'xdata':         postprocess_keylist,
                         'volume':        postprocess_int,
                         'year':          postprocess_int}

def postprocess(entry, fields, **options):
    """Postprocess a subset of fields in a list of parsed entries."""
    if type(fields) not in (bool, list):
        raise ValueError("postprocess takes either a bool or a list for fields"
                         ", not '{0}'".format(type(fields)))

    split_names = options.get('split_names', [])

    if type(split_names) not in (bool, list):
        raise ValueError("split_names takes either a bool or a list for fields"
                         ", not '{0}'".format(type(split_names)))

    remove_braces = options.get('remove_braces', False)
    _fields = []

    if fields:
        if type(fields) is bool:
            _fields = entry.fields
        elif type(fields) is list:
            _fields = fields
    elif remove_braces:
        _fields = entry.fields

    if type(split_names) is bool:
        split_names = _SPLIT_NAMES if split_names else []

    options['split_names'] = split_names

    for field in _fields:
        value = getattr(entry, field)

        if remove_braces:
            value = postprocess_braces(value, **options)

        if field in postprocess_functions:
            yield field, postprocess_functions[field](field, value, **options)
        else:
            yield field, value

# TODO: split_names should also work like postprocess and take the fields to
# postprocess as allowed
def postprocess_entry(entry, **options):
    """Convenience function for postprocessing the fields in an entry."""
    processed_fields =\
        bibpy.postprocess.postprocess(entry, options.get('postprocess', []),
                                      **options)

    for field, value in processed_fields:
        setattr(entry, field, value)
<commit_msg>Add missing fields to postprocess mapping<commit_after>"""Conversion functions for pre- and postprocessing of bib(la)tex fields."""

import bibpy.compat
import bibpy.date
import bibpy.error
import bibpy.grammar
import bibpy.lexers
import bibpy.name
import bibpy.parser
import calendar
import re

_MONTH_ABBREVIATIONS = [
    'jan', 'feb', 'mar', 'apr', 'may', 'jun',
    'jul', 'aug', 'sep', 'oct', 'nov', 'dec'
]

# TODO: Let users split on default names and custom ones
_SPLIT_NAMES = frozenset([
    'author', 'afterword', 'bookauthor', 'commentator',
    'editor', 'editora', 'editorb', 'editorc', 'foreword',
    'holder', 'introduction', 'language', 'origpublisher',
    'publisher', 'shortauthor', 'shorteditor', 'translator'
])

def postprocess_braces(value, **options):
    """Remove any braces from a string value."""
    if bibpy.compat.is_string(value):
        return "".join([e for e in bibpy.parser.parse_braced_expr(value)
                        if e not in '{}'])

    return value

def postprocess_namelist(field, names, **options):
    """Convert a string of authors to a list."""
    if not names:
        return []

    # First, split on zero brace-level 'and'
    names = bibpy.lexers.lex_namelist(names)

    # Second, if requested, parse each name
    if field in options.get('split_names', []):
        return [postprocess_name(field, n) for n in names]

    # Remove any leftover curly braces after splitting
    if not options.get('split_names', False):
        names = [re.sub('\{(.+?)\}', '\\1', name) for name in names]

    return names

def postprocess_name(field, author, **options):
    """Attempts to split an author name into first, middle and last name."""
    if author and bibpy.compat.is_string(author):
        return bibpy.name.Name.fromstring(author)
    else:
        return author

def postprocess_keywords(field, keywords, **options):
    """Convert a string of keywords to a list."""
    if not keywords:
        return []

    return [keyword.strip() for keyword in keywords.split(';')
            if keyword.strip()]

def postprocess_int(field, value, **options):
    """Convert a string to an integer."""
    try:
        return int(value)
    except (ValueError, TypeError):
        return value

def postprocess_date(field, datestring, **options):
    """Convert a string to a bibpy.date.DateRange."""
    if not datestring:
        return bibpy.date.DateRange((None, None, None), (None, None, None),
                                    False)

    return bibpy.date.DateRange.fromstring(datestring)

def get_month_name(i):
    """Return the name of a month from its one-based index ."""
    return calendar.month_name[i]

def postprocess_month(field, month, **options):
    """Convert a month number to its name."""
    try:
        i = int(month)

        if i not in range(1, 13):
            return month

        return get_month_name(i)
    except (ValueError, IndexError):
        month_name = month.lower()

        if month_name in _MONTH_ABBREVIATIONS:
            # The first element of calendar.month_name is an empty string
            return get_month_name(_MONTH_ABBREVIATIONS.index(month_name) + 1)

    return month

def postprocess_keylist(field, keylist, **options):
    """Convert a comma-separated string of keys to a list."""
    if not keylist:
        return []

    return list(filter(None, [key.strip() for key in keylist.split(',')]))

def postprocess_pages(field, pages, **options):
    """Convert a page range to a 2-element tuple."""
    values = re.split('\-+', pages)

    if len(values) == 2:
        try:
            return (int(values[0]), int(values[1]))
        except (ValueError, TypeError):
            return pages
    else:
        return pages

# A dictionary of fields as keys and the functions that postprocess them as
# values, e.g. 'year' should be converted to an integer etc.
postprocess_functions = {'address':       postprocess_namelist,
                         'afterword':     postprocess_namelist,
                         'annotator':     postprocess_namelist,
                         'author':        postprocess_namelist,
                         'bookauthor':    postprocess_namelist,
                         'commentator':   postprocess_namelist,
                         'date':          postprocess_date,
                         'edition':       postprocess_int,
                         'editor':        postprocess_namelist,
                         'editora':       postprocess_namelist,
                         'editorb':       postprocess_namelist,
                         'editorc':       postprocess_namelist,
                         'eventdate':     postprocess_date,
                         'foreword':      postprocess_namelist,
                         'holder':        postprocess_namelist,
                         'institution':   postprocess_namelist,
                         'introduction':  postprocess_namelist,
                         'keywords':      postprocess_keywords,
                         'language':      postprocess_namelist,
                         'location':      postprocess_namelist,
                         'month':         postprocess_month,
                         'number':        postprocess_int,
                         'organization':  postprocess_namelist,
                         'origdate':      postprocess_date,
                         'origlocation':  postprocess_namelist,
                         'origpublisher': postprocess_namelist,
                         'pages':         postprocess_pages,
                         'pagetotal':     postprocess_int,
                         'publisher':     postprocess_namelist,
                         'related':       postprocess_keylist,
                         'school':        postprocess_namelist,
                         'shortauthor':   postprocess_namelist,
                         'shorteditor':   postprocess_namelist,
                         'translator':    postprocess_namelist,
                         'urldate':       postprocess_date,
                         'xdata':         postprocess_keylist,
                         'volume':        postprocess_int,
                         'year':          postprocess_int}

def postprocess(entry, fields, **options):
    """Postprocess a subset of fields in a list of parsed entries."""
    if type(fields) not in (bool, list):
        raise ValueError("postprocess takes either a bool or a list for fields"
                         ", not '{0}'".format(type(fields)))

    split_names = options.get('split_names', [])

    if type(split_names) not in (bool, list):
        raise ValueError("split_names takes either a bool or a list for fields"
                         ", not '{0}'".format(type(split_names)))

    remove_braces = options.get('remove_braces', False)
    _fields = []

    if fields:
        if type(fields) is bool:
            _fields = entry.fields
        elif type(fields) is list:
            _fields = fields
    elif remove_braces:
        _fields = entry.fields

    if type(split_names) is bool:
        split_names = _SPLIT_NAMES if split_names else []

    options['split_names'] = split_names

    for field in _fields:
        value = getattr(entry, field)

        if remove_braces:
            value = postprocess_braces(value, **options)

        if field in postprocess_functions:
            yield field, postprocess_functions[field](field, value, **options)
        else:
            yield field, value

# TODO: split_names should also work like postprocess and take the fields to
# postprocess as allowed
def postprocess_entry(entry, **options):
    """Convenience function for postprocessing the fields in an entry."""
    processed_fields =\
        bibpy.postprocess.postprocess(entry, options.get('postprocess', []),
                                      **options)

    for field, value in processed_fields:
        setattr(entry, field, value)
<|endoftext|>


<commit_before>from __future__ import division, print_function

import numpy as np, gus_utils as gu
from scipy.optimize import brentq

def Phi(r,params):

    """
    SPL potential.

    Arguments
    ---------

    r: array_like
        radii at which to evaluate the potential

    params: list
        [vesc(Rsun), alpha] the potential parameters
    """

    vesc_Rsun,alpha = params
    return -.5*vesc_Rsun**2.*(r/8.5)**-alpha

def rperi_rapo_eq(r,E,L,params):

    """equation for finding rp, ra

    Arguments
    ---------

    r: array_like
        radii at which to evalute the equation

    E: float
        the energy of the orbit

    L: float
        the angular momentum of the orbit

    params: list
        the potential parameters

    Returns
    -------

    eq: array_like
        values of the LHS of the equation at r

    """

    return E  - .5*L**2. / r**2. - Phi(r,params)

def rp_ra_getstart(r,E,L,params,apo=False):
    """Find suitable place to start the root finding

    Arguments
    ---------

    r: array_like
        radii at which to evalute the equation

    E: float
        the energy of the orbit

    L: float
        the angular momentum of the orbit

    params: list
        the potential parameters

    Returns
    -------

    rtry: float
        a good place to start using a root finder

    """
    if apo:
        rtry = 1.1*r
    else:
        rtry = 0.9*r
    while rperi_rapo_eq(rtry,E,L,params)>=0. and rtry>1e-9:
        if apo:
            if rtry>1e10:
                raise Exception("This orbit looks unbound.")
            rtry*=1.1
        else:
            rtry*=0.9
    if rtry<1e-9:
        return 0.
    return rtry

def get_rp_ra(r,vr,vT,params):

    """Find the limiting radii of the orbit

    Arguments
    ---------

    r: float
        the current radius of the dwarf

    vr: float
        the current radial velocity of the dwarf

    vT: float
        the current tangential velocity of the dwarf

    params: list
        the potential parameters
    """

    E = .5*(vr**2. + vT**2.) + Phi(r,params)
    L = vT*r
    eps=1e-8
    if np.abs(rperi_rapo_eq(r,E,L,params))<1e-7: #we are at peri or apo
        peps = rperi_rapo_eq(r+eps,E,L,params)
        meps = rperi_rapo_eq(r-eps,E,L,params)
        if peps<0. and meps>0.: #we are at apo
            ra = r
            rstart = rp_ra_getstart(r-eps,E,L,params)
            if rstart==0.: rp = 0.
            else:
                rp = brentq(rperi_rapo_eq,rstart,r-eps,(E,L,params),maxiter=200)
        elif peps>0. and meps<0.: #we are at peri
            rp = r
            rend = rp_ra_getstart(r,E,L,params,apo=True)
            ra = brentq(rperi_rapo_eq,r+eps,rend,rend)
        else: #circular orbit
            rp = r
            ra = r
    else:
        rstart = rp_ra_getstart(r,E,L,params)
        if rstart==0.: rp=0.
        else:
            rp = brentq(rperi_rapo_eq,rstart,rstart/0.9,(E,L,params),maxiter=200)
        rend = rp_ra_getstart(r,E,L,params,apo=True)
        ra = brentq(rperi_rapo_eq,rend/1.1,rend,(E,L,params),maxiter=200)
    return (rp,ra)

def sample_rp_ra_e_distributions(r,vr,chain,thin_by=100,burnin=200):

    """
    sample rp,ra,epsilon distribution given chain and measurement

    Arguments
    ---------

    r: float
        galactocentric radius of the dwarf

    vr: float
        the radial velocity of the dwarf

    chain: array_like
        mcmc chain of the potential parameters and k

    thin_by: (=100) int
        the thinning factor for the mcmc chain

    burnin: (=200) int
        the number of steps per walker to discard for burn in

    Returns
    -------

    rp: array_like
        posterior samples on rp

    ra: array_like
        posterior samples on ra

    e: array_like
        posterior samples on the eccentricity
    """

    c = gu.reshape_chain(chain)[:,burnin::thin_by,:]
    c = np.reshape(c, (c.shape[0]*c.shape[1],c.shape[2]))
    k = c[:,2]
    samples = c[:,-2:]
    rp,ra,e = np.ones(len(samples))*np.nan,np.ones(len(samples))*np.nan,np.ones(len(samples))*np.nan

    for i,params in enumerate(samples):
        try:
            vesc_Rsun,alpha = params
            vesc = vesc_Rsun*(r/8.5)**(-alpha/2.)
            v = vr + (vesc - vr)*(1. - np.random.power(k[i]+1) )
            vT = np.sqrt(v**2. - vr**2.)
            rp[i],ra[i] = get_rp_ra(r,vr,vT,params)
            e[i] = (ra[i] - rp[i])/(ra[i] + rp[i])
        except:
            pass

    return rp[~np.isnan(rp)],ra[~np.isnan(rp)],e[~np.isnan(rp)]

def test_rpra_solver(params):

    """
    Given some potential parameters, test the peri and apo solver
    by passing it orbits with known rp, ra.

    Arguments
    ---------

    params: list
        potential parameters [vesc_Rsun, alphas]

    Returns
    -------

    delta_rp: array_like
        the differences between the true and solver pericentres

    delta_ra: array_like
        the differences between the true and solver apocentres
    """

    rp = np.random.uniform(1.,30.,100)
    ra = np.random.uniform(30.,1000.,100)
    E = -(rp**2.*Phi(rp,params) - ra**2.*Phi(ra,params))/(ra**2. - rp**2.)
    L = np.sqrt(2.*(Phi(ra,params) - Phi(rp,params))/(rp**-2. - ra**-2.))
    dr = ra-rp
    f = np.random.uniform(size=100)
    rs = rp + f*dr
    vT = L/rs
    vr = np.sqrt(2.*(E - Phi(rs,params)) - vT**2.)
    rp_s,ra_s = np.zeros(100), np.zeros(100)
    for i in np.arange(100):
        rp_s[i],ra_s[i] = get_rp_ra(rs[i],vr[i],vT[i],params)

    return rp - rp_s, ra - ra_s

<commit_msg>test in a less stupid way<commit_after>from __future__ import division, print_function

import numpy as np, gus_utils as gu
from scipy.optimize import brentq

def Phi(r,params):

    """
    SPL potential.

    Arguments
    ---------

    r: array_like
        radii at which to evaluate the potential

    params: list
        [vesc(Rsun), alpha] the potential parameters
    """

    vesc_Rsun,alpha = params
    return -.5*vesc_Rsun**2.*(r/8.5)**-alpha

def rperi_rapo_eq(r,E,L,params):

    """equation for finding rp, ra

    Arguments
    ---------

    r: array_like
        radii at which to evalute the equation

    E: float
        the energy of the orbit

    L: float
        the angular momentum of the orbit

    params: list
        the potential parameters

    Returns
    -------

    eq: array_like
        values of the LHS of the equation at r

    """

    return E  - .5*L**2. / r**2. - Phi(r,params)

def rp_ra_getstart(r,E,L,params,apo=False):
    """Find suitable place to start the root finding

    Arguments
    ---------

    r: array_like
        radii at which to evalute the equation

    E: float
        the energy of the orbit

    L: float
        the angular momentum of the orbit

    params: list
        the potential parameters

    Returns
    -------

    rtry: float
        a good place to start using a root finder

    """
    if apo:
        rtry = 1.1*r
    else:
        rtry = 0.9*r
    while rperi_rapo_eq(rtry,E,L,params)>=0. and rtry>1e-9:
        if apo:
            if rtry>1e10:
                raise Exception("This orbit looks unbound.")
            rtry*=1.1
        else:
            rtry*=0.9
    if rtry<1e-9:
        return 0.
    return rtry

def get_rp_ra(r,vr,vT,params):

    """Find the limiting radii of the orbit

    Arguments
    ---------

    r: float
        the current radius of the dwarf

    vr: float
        the current radial velocity of the dwarf

    vT: float
        the current tangential velocity of the dwarf

    params: list
        the potential parameters
    """

    E = .5*(vr**2. + vT**2.) + Phi(r,params)
    L = vT*r
    eps=1e-8
    if np.abs(rperi_rapo_eq(r,E,L,params))<1e-7: #we are at peri or apo
        peps = rperi_rapo_eq(r+eps,E,L,params)
        meps = rperi_rapo_eq(r-eps,E,L,params)
        if peps<0. and meps>0.: #we are at apo
            ra = r
            rstart = rp_ra_getstart(r-eps,E,L,params)
            if rstart==0.: rp = 0.
            else:
                rp = brentq(rperi_rapo_eq,rstart,r-eps,(E,L,params),maxiter=200)
        elif peps>0. and meps<0.: #we are at peri
            rp = r
            rend = rp_ra_getstart(r,E,L,params,apo=True)
            ra = brentq(rperi_rapo_eq,r+eps,rend,rend)
        else: #circular orbit
            rp = r
            ra = r
    else:
        rstart = rp_ra_getstart(r,E,L,params)
        if rstart==0.: rp=0.
        else:
            rp = brentq(rperi_rapo_eq,rstart,rstart/0.9,(E,L,params),maxiter=200)
        rend = rp_ra_getstart(r,E,L,params,apo=True)
        ra = brentq(rperi_rapo_eq,rend/1.1,rend,(E,L,params),maxiter=200)
    return (rp,ra)

def sample_rp_ra_e_distributions(r,vr,chain,thin_by=100,burnin=200):

    """
    sample rp,ra,epsilon distribution given chain and measurement

    Arguments
    ---------

    r: float
        galactocentric radius of the dwarf

    vr: float
        the radial velocity of the dwarf

    chain: array_like
        mcmc chain of the potential parameters and k

    thin_by: (=100) int
        the thinning factor for the mcmc chain

    burnin: (=200) int
        the number of steps per walker to discard for burn in

    Returns
    -------

    rp: array_like
        posterior samples on rp

    ra: array_like
        posterior samples on ra

    e: array_like
        posterior samples on the eccentricity
    """

    c = gu.reshape_chain(chain)[:,burnin::thin_by,:]
    c = np.reshape(c, (c.shape[0]*c.shape[1],c.shape[2]))
    k = c[:,2]
    samples = c[:,-2:]
    rp,ra,e = np.ones(len(samples))*np.nan,np.ones(len(samples))*np.nan,np.ones(len(samples))*np.nan

    for i,params in enumerate(samples):
        try:
            vesc_Rsun,alpha = params
            vesc = vesc_Rsun*(r/8.5)**(-alpha/2.)
            v = vr + (vesc - vr)*(1. - np.random.power(k[i]+1) )
            vT = np.sqrt(v**2. - vr**2.)
            rp[i],ra[i] = get_rp_ra(r,vr,vT,params)
            e[i] = (ra[i] - rp[i])/(ra[i] + rp[i])
        except:
            pass

    return rp[~np.isnan(rp)],ra[~np.isnan(rp)],e[~np.isnan(rp)]

def test_rpra_solver(params):

    """
    Given some potential parameters, test the peri and apo solver
    by passing it orbits with known rp, ra.

    Arguments
    ---------

    params: list
        potential parameters [vesc_Rsun, alphas]

    Returns
    -------

    delta_rp: array_like
        the differences between the true and solver pericentres

    delta_ra: array_like
        the differences between the true and solver apocentres
    """

    ra = np.random.uniform(1.,1000.,size=100)
    rp_frac = np.random.uniform(size=100)
    rp = rp_frac*ra
    E = -(rp**2.*Phi(rp,params) - ra**2.*Phi(ra,params))/(ra**2. - rp**2.)
    L = np.sqrt(2.*(Phi(ra,params) - Phi(rp,params))/(rp**-2. - ra**-2.))
    dr = ra-rp
    f = np.random.uniform(size=100)
    rs = rp + f*dr
    vT = L/rs
    vr = np.sqrt(2.*(E - Phi(rs,params)) - vT**2.)
    rp_s,ra_s = np.zeros(100), np.zeros(100)
    for i in np.arange(100):
        rp_s[i],ra_s[i] = get_rp_ra(rs[i],vr[i],vT[i],params)

    return rp - rp_s, ra - ra_s

<|endoftext|>


<commit_before>from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
import argparse
import contextlib
import logging
import yaml
import sys

from radosgw_agent import client
from radosgw_agent import sync

def check_positive_int(string):
    value = int(string)
    if value < 1:
        msg = '%r is not a positive integer' % string
        raise argparse.ArgumentTypeError(msg)
    return value

def check_endpoint(endpoint):
    try:
        return client.parse_endpoint(endpoint)
    except client.InvalidProtocol as e:
        raise argparse.ArgumentTypeError(str(e))
    except client.InvalidHost as e:
        raise argparse.ArgumentTypeError(str(e))

def parse_args():
    conf_parser = argparse.ArgumentParser(add_help=False)
    conf_parser.add_argument(
        '-c', '--conf',
        type=file,
        help='configuration file'
        )
    args, remaining = conf_parser.parse_known_args()
    defaults = dict(
        sync_scope='incremental',
        log_lock_time=20,
        )
    if args.conf is not None:
        with contextlib.closing(args.conf):
            config = yaml.safe_load_all(args.conf)
            for new in config:
                defaults.update(new)

    parser = argparse.ArgumentParser(
        parents=[conf_parser],
        description='Synchronize radosgw installations',
        )
    parser.set_defaults(**defaults)
    verbosity = parser.add_mutually_exclusive_group(required=False)
    verbosity.add_argument(
        '-v', '--verbose',
        action='store_true', dest='verbose',
        help='be more verbose',
        )
    verbosity.add_argument(
        '-q', '--quiet',
        action='store_true', dest='quiet',
        help='be less verbose',
        )
    parser.add_argument(
        '--src-access-key',
        required='src_access_key' not in defaults,
        help='access key for source zone system user',
        )
    parser.add_argument(
        '--src-secret-key',
        required='src_secret_key' not in defaults,
        help='secret key for source zone system user',
        )
    parser.add_argument(
        '--dest-access-key',
        required='dest_access_key' not in defaults,
        help='access key for destination zone system user',
        )
    parser.add_argument(
        '--dest-secret-key',
        required='dest_secret_key' not in defaults,
        help='secret key for destination zone system user',
        )
    parser.add_argument(
        'destination',
        type=check_endpoint,
        nargs=None if 'destination' not in defaults else '?',
        help='radosgw endpoint to which to sync '
        '(e.g. http://zone2.example.org:8080)',
        )
    src_options = parser.add_mutually_exclusive_group(required=False)
    src_options.add_argument(
        '--source',
        type=check_endpoint,
        help='radosgw endpoint from which to sync '
        '(e.g. http://zone1.example.org:8080)',
        )
    src_options.add_argument(
        '--src-zone',
        help='radosgw zone from which to sync',
        )
    parser.add_argument(
        '--metadata-only',
        action='store_true',
        help='sync bucket and user metadata, but not bucket contents',
        )
    parser.add_argument(
        '--num-workers',
        default=1,
        type=check_positive_int,
        help='number of items to sync at once',
        )
    parser.add_argument(
        '--sync-scope',
        choices=['full', 'incremental'],
        default='incremental',
        help='synchronize everything (for a new region) or only things that '
             'have changed since the last run',
        )
    parser.add_argument(
        '--lock-timeout',
        type=check_positive_int,
        default=60,
        help='timeout in seconds after which a log segment lock will expire if '
             'not refreshed',
        )
    parser.add_argument(
        '--log-file',
        help='where to store log output',
        )
    parser.add_argument(
        '--max-entries',
        type=check_positive_int,
        default=1000,
        help='maximum number of log entries to process at once during '
        'continuous sync',
        )
    parser.add_argument(
        '--incremental-sync-delay',
        type=check_positive_int,
        default=30,
        help='seconds to wait between syncs',
        )
    parser.add_argument(
        '--object-sync-timeout',
        type=check_positive_int,
        default=60 * 60 * 60,
        help='seconds to wait for an individual object to sync before '
        'assuming failure',
        )
    parser.add_argument(
        '--prepare-error-delay',
        type=check_positive_int,
        default=10,
        help='seconds to wait before retrying when preparing '
        'an incremental sync fails',
        )
    parser.add_argument(
        '--rgw-data-log-window',
        type=check_positive_int,
        default=30,
        help='period until a data log entry is valid - '
        'must match radosgw configuration',
        )
    parser.add_argument(
        '--test-server-host',
        # host to run a simple http server for testing the sync agent on,
        help=argparse.SUPPRESS,
        )
    parser.add_argument(
        '--test-server-port',
        # port to run a simple http server for testing the sync agent on,
        type=check_positive_int,
        default=8080,
        help=argparse.SUPPRESS,
        )
    return parser.parse_args(remaining)

class TestHandler(BaseHTTPRequestHandler):
    """HTTP handler for testing radosgw-agent.

    This should never be used outside of testing.
    """
    num_workers = None
    lock_timeout = None
    max_entries = None
    rgw_data_log_window = 30
    src = None
    dest = None

    def do_POST(self):
        log = logging.getLogger(__name__)
        status = 200
        resp = ''
        sync_cls = None
        if self.path.startswith('/metadata/full'):
            sync_cls = sync.MetaSyncerFull
        elif self.path.startswith('/metadata/incremental'):
            sync_cls = sync.MetaSyncerInc
        elif self.path.startswith('/data/full'):
            sync_cls = sync.DataSyncerFull
        elif self.path.startswith('/data/incremental'):
            sync_cls = sync.DataSyncerInc
        else:
            log.warn('invalid request, ignoring')
            status = 400
            resp = 'bad path'

        try:
            if sync_cls is not None:
                syncer = sync_cls(TestHandler.src, TestHandler.dest,
                                  TestHandler.max_entries,
                                  rgw_data_log_window=TestHandler.rgw_data_log_window,
                                  object_sync_timeout=TestHandler.object_sync_timeout)
                syncer.prepare()
                syncer.sync(
                    TestHandler.num_workers,
                    TestHandler.lock_timeout,
                    )
        except Exception as e:
            log.exception('error during sync')
            status = 500
            resp = str(e)

        self.log_request(status, len(resp))
        if status >= 400:
            self.send_error(status, resp)
        else:
            self.send_response(status)
            self.end_headers()

def main():
    args = parse_args()
    log = logging.getLogger()
    log_level = logging.INFO
    lib_log_level = logging.WARN
    if args.verbose:
        log_level = logging.DEBUG
        lib_log_level = logging.DEBUG
    elif args.quiet:
        log_level = logging.WARN
    logging.basicConfig(level=log_level)
    logging.getLogger('boto').setLevel(lib_log_level)
    logging.getLogger('requests').setLevel(lib_log_level)

    if args.log_file is not None:
        handler = logging.handlers.WatchedFileHandler(
            filename=args.log_file,
            )
        formatter = logging.Formatter(
            fmt='%(asctime)s.%(msecs)03d %(process)d:%(levelname)s:%(name)s:%(message)s',
            datefmt='%Y-%m-%dT%H:%M:%S',
            )
        handler.setFormatter(formatter)
        logging.getLogger().addHandler(handler)

    dest = args.destination
    dest.access_key = args.dest_access_key
    dest.secret_key = args.dest_secret_key
    src = args.source or client.Endpoint(None, None, None)
    if args.src_zone:
        src.zone = args.src_zone
    dest_conn = client.connection(dest)

    try:
        region_map = client.get_region_map(dest_conn)
    except Exception:
        log.exception('Could not retrieve region map from destination')
        sys.exit(1)

    try:
        client.configure_endpoints(region_map, dest, src, args.metadata_only)
    except client.ClientException as e:
        log.error(e)
        sys.exit(1)

    src.access_key = args.src_access_key
    src.secret_key = args.src_secret_key

    if args.test_server_host:
        log.warn('TEST MODE - do not run unless you are testing this program')
        TestHandler.src = src
        TestHandler.dest = dest
        TestHandler.num_workers = args.num_workers
        TestHandler.lock_timeout = args.lock_timeout
        TestHandler.max_entries = args.max_entries
        TestHandler.rgw_data_log_window = args.rgw_data_log_window
        TestHandler.object_sync_timeout = args.object_sync_timeout
        server = HTTPServer((args.test_server_host, args.test_server_port),
                            TestHandler)
        server.serve_forever()
        sys.exit()

    if args.sync_scope == 'full':
        meta_cls = sync.MetaSyncerFull
        data_cls = sync.DataSyncerFull
    else:
        meta_cls = sync.MetaSyncerInc
        data_cls = sync.DataSyncerInc

    meta_syncer = meta_cls(src, dest, args.max_entries)
    data_syncer = data_cls(src, dest, args.max_entries,
                           rgw_data_log_window=args.rgw_data_log_window,
                           object_sync_timeout=args.object_sync_timeout)

    # fetch logs first since data logs need to wait before becoming usable
    # due to rgw's window of data log updates during which the bucket index
    # log may still be updated without the data log getting a new entry for
    # the bucket
    sync.prepare_sync(meta_syncer, args.prepare_error_delay)
    if not args.metadata_only:
        sync.prepare_sync(data_syncer, args.prepare_error_delay)

    if args.sync_scope == 'full':
        log.info('syncing all metadata')
        meta_syncer.sync(args.num_workers, args.lock_timeout)
        if not args.metadata_only:
            log.info('syncing all data')
            data_syncer.sync(args.num_workers, args.lock_timeout)
        log.info('Finished full sync. Check logs to see any issues that '
                 'incremental sync will retry.')
    else:
        sync.incremental_sync(meta_syncer, data_syncer,
                              args.num_workers,
                              args.lock_timeout,
                              args.incremental_sync_delay,
                              args.metadata_only,
                              args.prepare_error_delay)
<commit_msg>cli: need to import submodule of logging to use it<commit_after>from BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
import argparse
import contextlib
import logging
import logging.handlers
import yaml
import sys

from radosgw_agent import client
from radosgw_agent import sync

def check_positive_int(string):
    value = int(string)
    if value < 1:
        msg = '%r is not a positive integer' % string
        raise argparse.ArgumentTypeError(msg)
    return value

def check_endpoint(endpoint):
    try:
        return client.parse_endpoint(endpoint)
    except client.InvalidProtocol as e:
        raise argparse.ArgumentTypeError(str(e))
    except client.InvalidHost as e:
        raise argparse.ArgumentTypeError(str(e))

def parse_args():
    conf_parser = argparse.ArgumentParser(add_help=False)
    conf_parser.add_argument(
        '-c', '--conf',
        type=file,
        help='configuration file'
        )
    args, remaining = conf_parser.parse_known_args()
    defaults = dict(
        sync_scope='incremental',
        log_lock_time=20,
        )
    if args.conf is not None:
        with contextlib.closing(args.conf):
            config = yaml.safe_load_all(args.conf)
            for new in config:
                defaults.update(new)

    parser = argparse.ArgumentParser(
        parents=[conf_parser],
        description='Synchronize radosgw installations',
        )
    parser.set_defaults(**defaults)
    verbosity = parser.add_mutually_exclusive_group(required=False)
    verbosity.add_argument(
        '-v', '--verbose',
        action='store_true', dest='verbose',
        help='be more verbose',
        )
    verbosity.add_argument(
        '-q', '--quiet',
        action='store_true', dest='quiet',
        help='be less verbose',
        )
    parser.add_argument(
        '--src-access-key',
        required='src_access_key' not in defaults,
        help='access key for source zone system user',
        )
    parser.add_argument(
        '--src-secret-key',
        required='src_secret_key' not in defaults,
        help='secret key for source zone system user',
        )
    parser.add_argument(
        '--dest-access-key',
        required='dest_access_key' not in defaults,
        help='access key for destination zone system user',
        )
    parser.add_argument(
        '--dest-secret-key',
        required='dest_secret_key' not in defaults,
        help='secret key for destination zone system user',
        )
    parser.add_argument(
        'destination',
        type=check_endpoint,
        nargs=None if 'destination' not in defaults else '?',
        help='radosgw endpoint to which to sync '
        '(e.g. http://zone2.example.org:8080)',
        )
    src_options = parser.add_mutually_exclusive_group(required=False)
    src_options.add_argument(
        '--source',
        type=check_endpoint,
        help='radosgw endpoint from which to sync '
        '(e.g. http://zone1.example.org:8080)',
        )
    src_options.add_argument(
        '--src-zone',
        help='radosgw zone from which to sync',
        )
    parser.add_argument(
        '--metadata-only',
        action='store_true',
        help='sync bucket and user metadata, but not bucket contents',
        )
    parser.add_argument(
        '--num-workers',
        default=1,
        type=check_positive_int,
        help='number of items to sync at once',
        )
    parser.add_argument(
        '--sync-scope',
        choices=['full', 'incremental'],
        default='incremental',
        help='synchronize everything (for a new region) or only things that '
             'have changed since the last run',
        )
    parser.add_argument(
        '--lock-timeout',
        type=check_positive_int,
        default=60,
        help='timeout in seconds after which a log segment lock will expire if '
             'not refreshed',
        )
    parser.add_argument(
        '--log-file',
        help='where to store log output',
        )
    parser.add_argument(
        '--max-entries',
        type=check_positive_int,
        default=1000,
        help='maximum number of log entries to process at once during '
        'continuous sync',
        )
    parser.add_argument(
        '--incremental-sync-delay',
        type=check_positive_int,
        default=30,
        help='seconds to wait between syncs',
        )
    parser.add_argument(
        '--object-sync-timeout',
        type=check_positive_int,
        default=60 * 60 * 60,
        help='seconds to wait for an individual object to sync before '
        'assuming failure',
        )
    parser.add_argument(
        '--prepare-error-delay',
        type=check_positive_int,
        default=10,
        help='seconds to wait before retrying when preparing '
        'an incremental sync fails',
        )
    parser.add_argument(
        '--rgw-data-log-window',
        type=check_positive_int,
        default=30,
        help='period until a data log entry is valid - '
        'must match radosgw configuration',
        )
    parser.add_argument(
        '--test-server-host',
        # host to run a simple http server for testing the sync agent on,
        help=argparse.SUPPRESS,
        )
    parser.add_argument(
        '--test-server-port',
        # port to run a simple http server for testing the sync agent on,
        type=check_positive_int,
        default=8080,
        help=argparse.SUPPRESS,
        )
    return parser.parse_args(remaining)

class TestHandler(BaseHTTPRequestHandler):
    """HTTP handler for testing radosgw-agent.

    This should never be used outside of testing.
    """
    num_workers = None
    lock_timeout = None
    max_entries = None
    rgw_data_log_window = 30
    src = None
    dest = None

    def do_POST(self):
        log = logging.getLogger(__name__)
        status = 200
        resp = ''
        sync_cls = None
        if self.path.startswith('/metadata/full'):
            sync_cls = sync.MetaSyncerFull
        elif self.path.startswith('/metadata/incremental'):
            sync_cls = sync.MetaSyncerInc
        elif self.path.startswith('/data/full'):
            sync_cls = sync.DataSyncerFull
        elif self.path.startswith('/data/incremental'):
            sync_cls = sync.DataSyncerInc
        else:
            log.warn('invalid request, ignoring')
            status = 400
            resp = 'bad path'

        try:
            if sync_cls is not None:
                syncer = sync_cls(TestHandler.src, TestHandler.dest,
                                  TestHandler.max_entries,
                                  rgw_data_log_window=TestHandler.rgw_data_log_window,
                                  object_sync_timeout=TestHandler.object_sync_timeout)
                syncer.prepare()
                syncer.sync(
                    TestHandler.num_workers,
                    TestHandler.lock_timeout,
                    )
        except Exception as e:
            log.exception('error during sync')
            status = 500
            resp = str(e)

        self.log_request(status, len(resp))
        if status >= 400:
            self.send_error(status, resp)
        else:
            self.send_response(status)
            self.end_headers()

def main():
    args = parse_args()
    log = logging.getLogger()
    log_level = logging.INFO
    lib_log_level = logging.WARN
    if args.verbose:
        log_level = logging.DEBUG
        lib_log_level = logging.DEBUG
    elif args.quiet:
        log_level = logging.WARN
    logging.basicConfig(level=log_level)
    logging.getLogger('boto').setLevel(lib_log_level)
    logging.getLogger('requests').setLevel(lib_log_level)

    if args.log_file is not None:
        handler = logging.handlers.WatchedFileHandler(
            filename=args.log_file,
            )
        formatter = logging.Formatter(
            fmt='%(asctime)s.%(msecs)03d %(process)d:%(levelname)s:%(name)s:%(message)s',
            datefmt='%Y-%m-%dT%H:%M:%S',
            )
        handler.setFormatter(formatter)
        logging.getLogger().addHandler(handler)

    dest = args.destination
    dest.access_key = args.dest_access_key
    dest.secret_key = args.dest_secret_key
    src = args.source or client.Endpoint(None, None, None)
    if args.src_zone:
        src.zone = args.src_zone
    dest_conn = client.connection(dest)

    try:
        region_map = client.get_region_map(dest_conn)
    except Exception:
        log.exception('Could not retrieve region map from destination')
        sys.exit(1)

    try:
        client.configure_endpoints(region_map, dest, src, args.metadata_only)
    except client.ClientException as e:
        log.error(e)
        sys.exit(1)

    src.access_key = args.src_access_key
    src.secret_key = args.src_secret_key

    if args.test_server_host:
        log.warn('TEST MODE - do not run unless you are testing this program')
        TestHandler.src = src
        TestHandler.dest = dest
        TestHandler.num_workers = args.num_workers
        TestHandler.lock_timeout = args.lock_timeout
        TestHandler.max_entries = args.max_entries
        TestHandler.rgw_data_log_window = args.rgw_data_log_window
        TestHandler.object_sync_timeout = args.object_sync_timeout
        server = HTTPServer((args.test_server_host, args.test_server_port),
                            TestHandler)
        server.serve_forever()
        sys.exit()

    if args.sync_scope == 'full':
        meta_cls = sync.MetaSyncerFull
        data_cls = sync.DataSyncerFull
    else:
        meta_cls = sync.MetaSyncerInc
        data_cls = sync.DataSyncerInc

    meta_syncer = meta_cls(src, dest, args.max_entries)
    data_syncer = data_cls(src, dest, args.max_entries,
                           rgw_data_log_window=args.rgw_data_log_window,
                           object_sync_timeout=args.object_sync_timeout)

    # fetch logs first since data logs need to wait before becoming usable
    # due to rgw's window of data log updates during which the bucket index
    # log may still be updated without the data log getting a new entry for
    # the bucket
    sync.prepare_sync(meta_syncer, args.prepare_error_delay)
    if not args.metadata_only:
        sync.prepare_sync(data_syncer, args.prepare_error_delay)

    if args.sync_scope == 'full':
        log.info('syncing all metadata')
        meta_syncer.sync(args.num_workers, args.lock_timeout)
        if not args.metadata_only:
            log.info('syncing all data')
            data_syncer.sync(args.num_workers, args.lock_timeout)
        log.info('Finished full sync. Check logs to see any issues that '
                 'incremental sync will retry.')
    else:
        sync.incremental_sync(meta_syncer, data_syncer,
                              args.num_workers,
                              args.lock_timeout,
                              args.incremental_sync_delay,
                              args.metadata_only,
                              args.prepare_error_delay)
<|endoftext|>


<commit_before>#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
import os

from docutils import nodes
from docutils.parsers import rst
from docutils.statemachine import ViewList
from sphinx.util.nodes import nested_parse_with_titles
import yaml

WORKFLOW_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '../../../workbooks', ))

def _title(name):
    return name.replace('_', ' ').capitalize() + " Workbook"

def _workbook_to_rst(name, workbook):

    title = _title(name)

    yield '.. _workbook-%s:' % name
    yield ''
    yield '=' * len(title)
    yield title
    yield '=' * len(title)
    yield ''
    yield ':Workbook name: {}'.format(workbook['name'])
    yield ''
    if 'description' in workbook:
        yield workbook['description']
        yield ''
    yield 'Workflows in the {} Workbook'.format(title)
    yield ''

    for wf_name, workflow in sorted(workbook['workflows'].items()):

        yield '.. object:: ' + workbook['name'] + '.' + wf_name
        yield ''

        if 'type' in workflow:
            yield '   :type: {}'.format(workflow['type'])
            yield ''

        if 'description' in workflow:
            yield '   {}'.format(workflow['description'])
            yield ''

        if 'input' in workflow:
            yield '   Workflow inputs:'
            yield ''
            for input_param in workflow['input']:
                try:
                    yield '   :input {}: Default: {}'.format(
                        *input_param.items()[0])
                except:
                    yield '   :input {}: Required.'.format(input_param)
            yield ''

def get_workbooks():

    all_workbooks = {}

    for root, dirs, files in os.walk(WORKFLOW_PATH):
        for file in files:
            with open(os.path.join(root, file)) as f:
                all_workbooks[file.split('.')[0]] = yaml.safe_load(f)

    return all_workbooks

def _write_workbook_pages(app):
    all_workbooks = get_workbooks()
    files = []

    for name, workbook in all_workbooks.items():
        filename = 'doc/source/reference/workbooks/%s.rst' % name
        app.info('generating workbook page for %s' % name)
        with open(filename, 'w') as f:
            f.write('\n'.join(_workbook_to_rst(name, workbook)))
        files.append(filename)
    return files

class WorkflowListDirective(rst.Directive):

    has_content = False

    def run(self):
        all_workbooks = get_workbooks()

        # Build the view of the data to be parsed for rendering.
        result = ViewList()
        for workbook_name in sorted(all_workbooks.keys()):
            workbook = all_workbooks[workbook_name]
            for line in _workbook_to_rst(workbook_name, workbook):
                result.append(line, '<' + __name__ + '>')

        # Parse what we have into a new section.
        node = nodes.section()
        node.document = self.state.document
        nested_parse_with_titles(self.state, result, node)

        return node.children

def setup(app):
    app.info('loading workbooks extension')
    app.add_directive('workbooklist', WorkflowListDirective)
    _write_workbook_pages(app)
<commit_msg>Remove the duplication of Workbook in the description.<commit_after>#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.
import os

from docutils import nodes
from docutils.parsers import rst
from docutils.statemachine import ViewList
from sphinx.util.nodes import nested_parse_with_titles
import yaml

WORKFLOW_PATH = os.path.abspath(
    os.path.join(os.path.dirname(__file__), '../../../workbooks', ))

def _title(name):
    return name.replace('_', ' ').capitalize() + " Workbook"

def _workbook_to_rst(name, workbook):

    title = _title(name)

    yield '.. _workbook-%s:' % name
    yield ''
    yield '=' * len(title)
    yield title
    yield '=' * len(title)
    yield ''
    yield ':Workbook name: {}'.format(workbook['name'])
    yield ''
    if 'description' in workbook:
        yield workbook['description']
        yield ''
    yield 'Workflows in the {}'.format(title)
    yield ''

    for wf_name, workflow in sorted(workbook['workflows'].items()):

        yield '.. object:: ' + workbook['name'] + '.' + wf_name
        yield ''

        if 'type' in workflow:
            yield '   :type: {}'.format(workflow['type'])
            yield ''

        if 'description' in workflow:
            yield '   {}'.format(workflow['description'])
            yield ''

        if 'input' in workflow:
            yield '   Workflow inputs:'
            yield ''
            for input_param in workflow['input']:
                try:
                    yield '   :input {}: Default: {}'.format(
                        *input_param.items()[0])
                except:
                    yield '   :input {}: Required.'.format(input_param)
            yield ''

def get_workbooks():

    all_workbooks = {}

    for root, dirs, files in os.walk(WORKFLOW_PATH):
        for file in files:
            with open(os.path.join(root, file)) as f:
                all_workbooks[file.split('.')[0]] = yaml.safe_load(f)

    return all_workbooks

def _write_workbook_pages(app):
    all_workbooks = get_workbooks()
    files = []

    for name, workbook in all_workbooks.items():
        filename = 'doc/source/reference/workbooks/%s.rst' % name
        app.info('generating workbook page for %s' % name)
        with open(filename, 'w') as f:
            f.write('\n'.join(_workbook_to_rst(name, workbook)))
        files.append(filename)
    return files

class WorkflowListDirective(rst.Directive):

    has_content = False

    def run(self):
        all_workbooks = get_workbooks()

        # Build the view of the data to be parsed for rendering.
        result = ViewList()
        for workbook_name in sorted(all_workbooks.keys()):
            workbook = all_workbooks[workbook_name]
            for line in _workbook_to_rst(workbook_name, workbook):
                result.append(line, '<' + __name__ + '>')

        # Parse what we have into a new section.
        node = nodes.section()
        node.document = self.state.document
        nested_parse_with_titles(self.state, result, node)

        return node.children

def setup(app):
    app.info('loading workbooks extension')
    app.add_directive('workbooklist', WorkflowListDirective)
    _write_workbook_pages(app)
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import socket
import time
import traceback
import logging

from resource_management.core import global_lock
from resource_management.libraries.functions import format
from resource_management.libraries.functions import get_kinit_path
from resource_management.core.resources import Execute
from resource_management.core.signal_utils import TerminateStrategy
from ambari_commons.os_check import OSConst
from ambari_commons.os_family_impl import OsFamilyFuncImpl, OsFamilyImpl

OK_MESSAGE = "Metastore OK - Hive command took {0:.3f}s"
CRITICAL_MESSAGE = "Metastore on {0} failed ({1})"
SECURITY_ENABLED_KEY = '{{cluster-env/security_enabled}}'
SMOKEUSER_KEYTAB_KEY = '{{cluster-env/smokeuser_keytab}}'
SMOKEUSER_PRINCIPAL_KEY = '{{cluster-env/smokeuser_principal_name}}'
SMOKEUSER_KEY = '{{cluster-env/smokeuser}}'
HIVE_METASTORE_URIS_KEY = '{{hive-site/hive.metastore.uris}}'

# The configured Kerberos executable search paths, if any
KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY = '{{kerberos-env/executable_search_paths}}'

# default keytab location
SMOKEUSER_KEYTAB_SCRIPT_PARAM_KEY = 'default.smoke.keytab'
SMOKEUSER_KEYTAB_DEFAULT = '/etc/security/keytabs/smokeuser.headless.keytab'

# default smoke principal
SMOKEUSER_PRINCIPAL_SCRIPT_PARAM_KEY = 'default.smoke.principal'
SMOKEUSER_PRINCIPAL_DEFAULT = 'ambari-qa@EXAMPLE.COM'

# default smoke user
SMOKEUSER_SCRIPT_PARAM_KEY = 'default.smoke.user'
SMOKEUSER_DEFAULT = 'ambari-qa'

STACK_ROOT = '{{cluster-env/stack_root}}'

HIVE_CONF_DIR_LEGACY = '/etc/hive/conf.server'

HIVE_BIN_DIR_LEGACY = '/usr/lib/hive/bin'

CHECK_COMMAND_TIMEOUT_KEY = 'check.command.timeout'
CHECK_COMMAND_TIMEOUT_DEFAULT = 60.0

HADOOPUSER_KEY = '{{cluster-env/hadoop.user.name}}'
HADOOPUSER_DEFAULT = 'hadoop'

logger = logging.getLogger('ambari_alerts')

@OsFamilyFuncImpl(os_family=OsFamilyImpl.DEFAULT)
def get_tokens():
  """
  Returns a tuple of tokens in the format {{site/property}} that will be used
  to build the dictionary passed into execute
  """
  return (SECURITY_ENABLED_KEY,SMOKEUSER_KEYTAB_KEY,SMOKEUSER_PRINCIPAL_KEY,
    HIVE_METASTORE_URIS_KEY, SMOKEUSER_KEY, KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY,
    STACK_ROOT)

@OsFamilyFuncImpl(os_family=OSConst.WINSRV_FAMILY)
def get_tokens():
  """
  Returns a tuple of tokens in the format {{site/property}} that will be used
  to build the dictionary passed into execute
  """
  return (HIVE_METASTORE_URIS_KEY, HADOOPUSER_KEY)

@OsFamilyFuncImpl(os_family=OsFamilyImpl.DEFAULT)
def execute(configurations={}, parameters={}, host_name=None):
  """
  Returns a tuple containing the result code and a pre-formatted result label

  Keyword arguments:
  configurations (dictionary): a mapping of configuration key to value
  parameters (dictionary): a mapping of script parameter key to value
  host_name (string): the name of this host where the alert is running
  """

  if configurations is None:
    return (('UNKNOWN', ['There were no configurations supplied to the script.']))

  if not HIVE_METASTORE_URIS_KEY in configurations:
    return (('UNKNOWN', ['Hive metastore uris were not supplied to the script.']))

  metastore_uris = configurations[HIVE_METASTORE_URIS_KEY].split(',')

  security_enabled = False
  if SECURITY_ENABLED_KEY in configurations:
    security_enabled = str(configurations[SECURITY_ENABLED_KEY]).upper() == 'TRUE'

  check_command_timeout = CHECK_COMMAND_TIMEOUT_DEFAULT
  if CHECK_COMMAND_TIMEOUT_KEY in parameters:
    check_command_timeout = float(parameters[CHECK_COMMAND_TIMEOUT_KEY])

  # defaults
  smokeuser_keytab = SMOKEUSER_KEYTAB_DEFAULT
  smokeuser_principal = SMOKEUSER_PRINCIPAL_DEFAULT
  smokeuser = SMOKEUSER_DEFAULT

  # check script params
  if SMOKEUSER_PRINCIPAL_SCRIPT_PARAM_KEY in parameters:
    smokeuser_principal = parameters[SMOKEUSER_PRINCIPAL_SCRIPT_PARAM_KEY]

  if SMOKEUSER_SCRIPT_PARAM_KEY in parameters:
    smokeuser = parameters[SMOKEUSER_SCRIPT_PARAM_KEY]

  if SMOKEUSER_KEYTAB_SCRIPT_PARAM_KEY in parameters:
    smokeuser_keytab = parameters[SMOKEUSER_KEYTAB_SCRIPT_PARAM_KEY]

  # check configurations last as they should always take precedence
  if SMOKEUSER_PRINCIPAL_KEY in configurations:
    smokeuser_principal = configurations[SMOKEUSER_PRINCIPAL_KEY]

  if SMOKEUSER_KEY in configurations:
    smokeuser = configurations[SMOKEUSER_KEY]

  result_code = None

  try:
    if security_enabled:
      if SMOKEUSER_KEYTAB_KEY in configurations:
        smokeuser_keytab = configurations[SMOKEUSER_KEYTAB_KEY]

      # Get the configured Kerberos executable search paths, if any
      if KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY in configurations:
        kerberos_executable_search_paths = configurations[KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY]
      else:
        kerberos_executable_search_paths = None

      kinit_path_local = get_kinit_path(kerberos_executable_search_paths)
      kinitcmd=format("{kinit_path_local} -kt {smokeuser_keytab} {smokeuser_principal}; ")

      # prevent concurrent kinit
      kinit_lock = global_lock.get_lock(global_lock.LOCK_TYPE_KERBEROS)
      kinit_lock.acquire()
      try:
        Execute(kinitcmd, user=smokeuser,
          path=["/bin/", "/usr/bin/", "/usr/lib/hive/bin/", "/usr/sbin/"],
          timeout=10)
      finally:
        kinit_lock.release()

    if host_name is None:
      host_name = socket.getfqdn()

    for uri in metastore_uris:
      if host_name in uri:
        metastore_uri = uri

    conf_dir = HIVE_CONF_DIR_LEGACY
    bin_dir = HIVE_BIN_DIR_LEGACY

    if STACK_ROOT in configurations:
      hive_conf_dir = configurations[STACK_ROOT] + format("/current/hive-metastore/conf/conf.server")
      hive_bin_dir = configurations[STACK_ROOT] + format("/current/hive-metastore/bin")

      if os.path.exists(hive_conf_dir):
        conf_dir = hive_conf_dir
        bin_dir = hive_bin_dir

    cmd = format("export HIVE_CONF_DIR='{conf_dir}' ; "
                 "hive --hiveconf hive.metastore.uris={metastore_uri}\
                 --hiveconf hive.metastore.client.connect.retry.delay=1\
                 --hiveconf hive.metastore.failure.retries=1\
                 --hiveconf hive.metastore.connect.retries=1\
                 --hiveconf hive.metastore.client.socket.timeout=14\
                 --hiveconf hive.execution.engine=mr -e 'show databases;'")

    start_time = time.time()

    try:
      Execute(cmd, user=smokeuser,
        path=["/bin/", "/usr/bin/", "/usr/sbin/", bin_dir],
        timeout=int(check_command_timeout),
        timeout_kill_strategy=TerminateStrategy.KILL_PROCESS_TREE,
      )

      total_time = time.time() - start_time

      result_code = 'OK'
      label = OK_MESSAGE.format(total_time)
    except:
      result_code = 'CRITICAL'
      label = CRITICAL_MESSAGE.format(host_name, traceback.format_exc())

  except:
    label = traceback.format_exc()
    result_code = 'UNKNOWN'

  return ((result_code, [label]))

@OsFamilyFuncImpl(os_family=OSConst.WINSRV_FAMILY)
def execute(configurations={}, parameters={}, host_name=None):
  """
  Returns a tuple containing the result code and a pre-formatted result label

  Keyword arguments:
  configurations (dictionary): a mapping of configuration key to value
  parameters (dictionary): a mapping of script parameter key to value
  host_name (string): the name of this host where the alert is running
  """

  from resource_management.libraries.functions import reload_windows_env
  reload_windows_env()
  hive_home = os.environ['HIVE_HOME']

  if configurations is None:
    return (('UNKNOWN', ['There were no configurations supplied to the script.']))
  if not HIVE_METASTORE_URIS_KEY in configurations:
    return (('UNKNOWN', ['Hive metastore uris were not supplied to the script.']))

  metastore_uris = configurations[HIVE_METASTORE_URIS_KEY].split(',')

  # defaults
  hiveuser = HADOOPUSER_DEFAULT

  if HADOOPUSER_KEY in configurations:
    hiveuser = configurations[HADOOPUSER_KEY]

  result_code = None
  try:
    if host_name is None:
      host_name = socket.getfqdn()
    for uri in metastore_uris:
      if host_name in uri:
        metastore_uri = uri

    hive_cmd = os.path.join(hive_home, "bin", "hive.cmd")
    cmd = format("cmd /c {hive_cmd} --hiveconf hive.metastore.uris={metastore_uri}\
                 --hiveconf hive.metastore.client.connect.retry.delay=1\
                 --hiveconf hive.metastore.failure.retries=1\
                 --hiveconf hive.metastore.connect.retries=1\
                 --hiveconf hive.metastore.client.socket.timeout=14\
                 --hiveconf hive.execution.engine=mr -e 'show databases;'")
    start_time = time.time()
    try:
      Execute(cmd, user=hiveuser, timeout=30)
      total_time = time.time() - start_time
      result_code = 'OK'
      label = OK_MESSAGE.format(total_time)
    except:
      result_code = 'CRITICAL'
      label = CRITICAL_MESSAGE.format(host_name, traceback.format_exc())
  except:
    label = traceback.format_exc()
    result_code = 'UNKNOWN'

  return ((result_code, [label]))
<commit_msg>AMBARI-19534. Hive Metastore process cannot be determined to be up and listening on the network. (swagle)<commit_after>#!/usr/bin/env python

"""
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements.  See the NOTICE file
distributed with this work for additional information
regarding copyright ownership.  The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License.  You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import socket
import time
import traceback
import logging

from resource_management.core import global_lock
from resource_management.libraries.functions import format
from resource_management.libraries.functions import get_kinit_path
from resource_management.core.resources import Execute
from resource_management.core.signal_utils import TerminateStrategy
from ambari_commons.os_check import OSConst
from ambari_commons.os_family_impl import OsFamilyFuncImpl, OsFamilyImpl

OK_MESSAGE = "Metastore OK - Hive command took {0:.3f}s"
CRITICAL_MESSAGE = "Metastore on {0} failed ({1})"
SECURITY_ENABLED_KEY = '{{cluster-env/security_enabled}}'
SMOKEUSER_KEYTAB_KEY = '{{cluster-env/smokeuser_keytab}}'
SMOKEUSER_PRINCIPAL_KEY = '{{cluster-env/smokeuser_principal_name}}'
SMOKEUSER_KEY = '{{cluster-env/smokeuser}}'
HIVE_METASTORE_URIS_KEY = '{{hive-site/hive.metastore.uris}}'

# The configured Kerberos executable search paths, if any
KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY = '{{kerberos-env/executable_search_paths}}'

# default keytab location
SMOKEUSER_KEYTAB_SCRIPT_PARAM_KEY = 'default.smoke.keytab'
SMOKEUSER_KEYTAB_DEFAULT = '/etc/security/keytabs/smokeuser.headless.keytab'

# default smoke principal
SMOKEUSER_PRINCIPAL_SCRIPT_PARAM_KEY = 'default.smoke.principal'
SMOKEUSER_PRINCIPAL_DEFAULT = 'ambari-qa@EXAMPLE.COM'

# default smoke user
SMOKEUSER_SCRIPT_PARAM_KEY = 'default.smoke.user'
SMOKEUSER_DEFAULT = 'ambari-qa'

STACK_ROOT = '{{cluster-env/stack_root}}'

HIVE_CONF_DIR_LEGACY = '/etc/hive/conf.server'

HIVE_BIN_DIR_LEGACY = '/usr/lib/hive/bin'

CHECK_COMMAND_TIMEOUT_KEY = 'check.command.timeout'
CHECK_COMMAND_TIMEOUT_DEFAULT = 60.0

HADOOPUSER_KEY = '{{cluster-env/hadoop.user.name}}'
HADOOPUSER_DEFAULT = 'hadoop'

logger = logging.getLogger('ambari_alerts')

@OsFamilyFuncImpl(os_family=OsFamilyImpl.DEFAULT)
def get_tokens():
  """
  Returns a tuple of tokens in the format {{site/property}} that will be used
  to build the dictionary passed into execute
  """
  return (SECURITY_ENABLED_KEY,SMOKEUSER_KEYTAB_KEY,SMOKEUSER_PRINCIPAL_KEY,
    HIVE_METASTORE_URIS_KEY, SMOKEUSER_KEY, KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY,
    STACK_ROOT)

@OsFamilyFuncImpl(os_family=OSConst.WINSRV_FAMILY)
def get_tokens():
  """
  Returns a tuple of tokens in the format {{site/property}} that will be used
  to build the dictionary passed into execute
  """
  return (HIVE_METASTORE_URIS_KEY, HADOOPUSER_KEY)

@OsFamilyFuncImpl(os_family=OsFamilyImpl.DEFAULT)
def execute(configurations={}, parameters={}, host_name=None):
  """
  Returns a tuple containing the result code and a pre-formatted result label

  Keyword arguments:
  configurations (dictionary): a mapping of configuration key to value
  parameters (dictionary): a mapping of script parameter key to value
  host_name (string): the name of this host where the alert is running
  """

  if configurations is None:
    return (('UNKNOWN', ['There were no configurations supplied to the script.']))

  if not HIVE_METASTORE_URIS_KEY in configurations:
    return (('UNKNOWN', ['Hive metastore uris were not supplied to the script.']))

  metastore_uris = configurations[HIVE_METASTORE_URIS_KEY].split(',')

  security_enabled = False
  if SECURITY_ENABLED_KEY in configurations:
    security_enabled = str(configurations[SECURITY_ENABLED_KEY]).upper() == 'TRUE'

  check_command_timeout = CHECK_COMMAND_TIMEOUT_DEFAULT
  if CHECK_COMMAND_TIMEOUT_KEY in parameters:
    check_command_timeout = float(parameters[CHECK_COMMAND_TIMEOUT_KEY])

  # defaults
  smokeuser_keytab = SMOKEUSER_KEYTAB_DEFAULT
  smokeuser_principal = SMOKEUSER_PRINCIPAL_DEFAULT
  smokeuser = SMOKEUSER_DEFAULT

  # check script params
  if SMOKEUSER_PRINCIPAL_SCRIPT_PARAM_KEY in parameters:
    smokeuser_principal = parameters[SMOKEUSER_PRINCIPAL_SCRIPT_PARAM_KEY]

  if SMOKEUSER_SCRIPT_PARAM_KEY in parameters:
    smokeuser = parameters[SMOKEUSER_SCRIPT_PARAM_KEY]

  if SMOKEUSER_KEYTAB_SCRIPT_PARAM_KEY in parameters:
    smokeuser_keytab = parameters[SMOKEUSER_KEYTAB_SCRIPT_PARAM_KEY]

  # check configurations last as they should always take precedence
  if SMOKEUSER_PRINCIPAL_KEY in configurations:
    smokeuser_principal = configurations[SMOKEUSER_PRINCIPAL_KEY]

  if SMOKEUSER_KEY in configurations:
    smokeuser = configurations[SMOKEUSER_KEY]

  result_code = None

  try:
    if security_enabled:
      if SMOKEUSER_KEYTAB_KEY in configurations:
        smokeuser_keytab = configurations[SMOKEUSER_KEYTAB_KEY]

      # Get the configured Kerberos executable search paths, if any
      if KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY in configurations:
        kerberos_executable_search_paths = configurations[KERBEROS_EXECUTABLE_SEARCH_PATHS_KEY]
      else:
        kerberos_executable_search_paths = None

      kinit_path_local = get_kinit_path(kerberos_executable_search_paths)
      kinitcmd=format("{kinit_path_local} -kt {smokeuser_keytab} {smokeuser_principal}; ")

      # prevent concurrent kinit
      kinit_lock = global_lock.get_lock(global_lock.LOCK_TYPE_KERBEROS)
      kinit_lock.acquire()
      try:
        Execute(kinitcmd, user=smokeuser,
          path=["/bin/", "/usr/bin/", "/usr/lib/hive/bin/", "/usr/sbin/"],
          timeout=10)
      finally:
        kinit_lock.release()

    if host_name is None:
      host_name = socket.getfqdn()

    for uri in metastore_uris:
      if host_name in uri:
        metastore_uri = uri

    conf_dir = HIVE_CONF_DIR_LEGACY
    bin_dir = HIVE_BIN_DIR_LEGACY

    if STACK_ROOT in configurations:
      hive_conf_dir = configurations[STACK_ROOT] + format("/current/hive-metastore/conf")
      hive_bin_dir = configurations[STACK_ROOT] + format("/current/hive-metastore/bin")

      if os.path.exists(hive_conf_dir):
        conf_dir = hive_conf_dir
        bin_dir = hive_bin_dir

    cmd = format("export HIVE_CONF_DIR='{conf_dir}' ; "
                 "hive --hiveconf hive.metastore.uris={metastore_uri}\
                 --hiveconf hive.metastore.client.connect.retry.delay=1\
                 --hiveconf hive.metastore.failure.retries=1\
                 --hiveconf hive.metastore.connect.retries=1\
                 --hiveconf hive.metastore.client.socket.timeout=14\
                 --hiveconf hive.execution.engine=mr -e 'show databases;'")

    start_time = time.time()

    try:
      Execute(cmd, user=smokeuser,
        path=["/bin/", "/usr/bin/", "/usr/sbin/", bin_dir],
        timeout=int(check_command_timeout),
        timeout_kill_strategy=TerminateStrategy.KILL_PROCESS_TREE,
      )

      total_time = time.time() - start_time

      result_code = 'OK'
      label = OK_MESSAGE.format(total_time)
    except:
      result_code = 'CRITICAL'
      label = CRITICAL_MESSAGE.format(host_name, traceback.format_exc())

  except:
    label = traceback.format_exc()
    result_code = 'UNKNOWN'

  return ((result_code, [label]))

@OsFamilyFuncImpl(os_family=OSConst.WINSRV_FAMILY)
def execute(configurations={}, parameters={}, host_name=None):
  """
  Returns a tuple containing the result code and a pre-formatted result label

  Keyword arguments:
  configurations (dictionary): a mapping of configuration key to value
  parameters (dictionary): a mapping of script parameter key to value
  host_name (string): the name of this host where the alert is running
  """

  from resource_management.libraries.functions import reload_windows_env
  reload_windows_env()
  hive_home = os.environ['HIVE_HOME']

  if configurations is None:
    return (('UNKNOWN', ['There were no configurations supplied to the script.']))
  if not HIVE_METASTORE_URIS_KEY in configurations:
    return (('UNKNOWN', ['Hive metastore uris were not supplied to the script.']))

  metastore_uris = configurations[HIVE_METASTORE_URIS_KEY].split(',')

  # defaults
  hiveuser = HADOOPUSER_DEFAULT

  if HADOOPUSER_KEY in configurations:
    hiveuser = configurations[HADOOPUSER_KEY]

  result_code = None
  try:
    if host_name is None:
      host_name = socket.getfqdn()
    for uri in metastore_uris:
      if host_name in uri:
        metastore_uri = uri

    hive_cmd = os.path.join(hive_home, "bin", "hive.cmd")
    cmd = format("cmd /c {hive_cmd} --hiveconf hive.metastore.uris={metastore_uri}\
                 --hiveconf hive.metastore.client.connect.retry.delay=1\
                 --hiveconf hive.metastore.failure.retries=1\
                 --hiveconf hive.metastore.connect.retries=1\
                 --hiveconf hive.metastore.client.socket.timeout=14\
                 --hiveconf hive.execution.engine=mr -e 'show databases;'")
    start_time = time.time()
    try:
      Execute(cmd, user=hiveuser, timeout=30)
      total_time = time.time() - start_time
      result_code = 'OK'
      label = OK_MESSAGE.format(total_time)
    except:
      result_code = 'CRITICAL'
      label = CRITICAL_MESSAGE.format(host_name, traceback.format_exc())
  except:
    label = traceback.format_exc()
    result_code = 'UNKNOWN'

  return ((result_code, [label]))
<|endoftext|>