# Infinite push

#

# Copyright 2016 Facebook, Inc.

#

# This software may be used and distributed according to the terms of the

# GNU General Public License version 2 or any later version.

""" store some pushes in a remote blob store on the server (EXPERIMENTAL)

    [infinitepush]

    # Server-side and client-side option. Pattern of the infinitepush bookmark

    branchpattern = PATTERN

    # Server or client

    server = False

    # Server-side option. Possible values: 'disk' or 'sql'. Fails if not set

    indextype = disk

    # Server-side option. Used only if indextype=sql.

    # Format: 'IP:PORT:DB_NAME:USER:PASSWORD'

    sqlhost = IP:PORT:DB_NAME:USER:PASSWORD

    # Server-side option. Used only if indextype=disk.

    # Filesystem path to the index store

    indexpath = PATH

    # Server-side option. Possible values: 'disk' or 'external'

    # Fails if not set

    storetype = disk

    # Server-side option.

    # Path to the binary that will save bundle to the bundlestore

    # Formatted cmd line will be passed to it (see `put_args`)

    put_binary = put

    # Serser-side option. Used only if storetype=external.

    # Format cmd-line string for put binary. Placeholder: {filename}

    put_args = {filename}

    # Server-side option.

    # Path to the binary that get bundle from the bundlestore.

    # Formatted cmd line will be passed to it (see `get_args`)

    get_binary = get

    # Serser-side option. Used only if storetype=external.

    # Format cmd-line string for get binary. Placeholders: {filename} {handle}

    get_args = {filename} {handle}

    # Server-side option

    logfile = FIlE

    # Server-side option

    loglevel = DEBUG

    # Server-side option. Used only if indextype=sql.

    # Sets mysql wait_timeout option.

    waittimeout = 300

    # Server-side option. Used only if indextype=sql.

    # Sets mysql innodb_lock_wait_timeout option.

    locktimeout = 120

    # Server-side option. Used only if indextype=sql.

    # Name of the repository

    reponame = ''

    # Client-side option. Used by --list-remote option. List of remote scratch

    # patterns to list if no patterns are specified.

    defaultremotepatterns = ['*']

    # Instructs infinitepush to forward all received bundle2 parts to the

    # bundle for storage. Defaults to False.

    storeallparts = True

    # routes each incoming push to the bundlestore. defaults to False

    pushtobundlestore = True

    [remotenames]

    # Client-side option

    # This option should be set only if remotenames extension is enabled.

    # Whether remote bookmarks are tracked by remotenames extension.

    bookmarks = True

"""

from __future__ import absolute_import

import collections

import contextlib

import errno

import functools

import logging

import os

import random

import re

import socket

import subprocess

import tempfile

import time

from mercurial.node import (

    bin,

    hex,

)

from mercurial.i18n import _

from mercurial.utils import (

    procutil,

    stringutil,

)

from mercurial import (

    bundle2,

    changegroup,

    commands,

    discovery,

    encoding,

    error,

    exchange,

    extensions,

    hg,

    localrepo,

    phases,

    pushkey,

    pycompat,

    registrar,

    util,

    wireprototypes,

    wireprotov1peer,

    wireprotov1server,

)

from . import (

    bundleparts,

    common,

)

# Note for extension authors: ONLY specify testedwith = 'ships-with-hg-core' for

# extensions which SHIP WITH MERCURIAL. Non-mainline extensions should

# be specifying the version(s) of Mercurial they are tested with, or

# leave the attribute unspecified.

testedwith = 'ships-with-hg-core'

configtable = {}

configitem = registrar.configitem(configtable)

configitem('infinitepush', 'server',

    default=False,

)

configitem('infinitepush', 'storetype',

    default='',

)

configitem('infinitepush', 'indextype',

    default='',

)

configitem('infinitepush', 'indexpath',

    default='',

)

configitem('infinitepush', 'storeallparts',

    default=False,

)

configitem('infinitepush', 'reponame',

    default='',

)

configitem('scratchbranch', 'storepath',

    default='',

)

configitem('infinitepush', 'branchpattern',

    default='',

)

configitem('infinitepush', 'pushtobundlestore',

    default=False,

)

configitem('experimental', 'server-bundlestore-bookmark',

    default='',

)

configitem('experimental', 'infinitepush-scratchpush',

    default=False,

)

experimental = 'experimental'

configbookmark = 'server-bundlestore-bookmark'

configscratchpush = 'infinitepush-scratchpush'

scratchbranchparttype = bundleparts.scratchbranchparttype

revsetpredicate = registrar.revsetpredicate()

templatekeyword = registrar.templatekeyword()

_scratchbranchmatcher = lambda x: False

_maybehash = re.compile(r'^[a-f0-9]+$').search

def _buildexternalbundlestore(ui):

    put_args = ui.configlist('infinitepush', 'put_args', [])

    put_binary = ui.config('infinitepush', 'put_binary')

    if not put_binary:

        raise error.Abort('put binary is not specified')

    get_args = ui.configlist('infinitepush', 'get_args', [])

    get_binary = ui.config('infinitepush', 'get_binary')

    if not get_binary:

        raise error.Abort('get binary is not specified')

    from . import store

    return store.externalbundlestore(put_binary, put_args, get_binary, get_args)

def _buildsqlindex(ui):

    sqlhost = ui.config('infinitepush', 'sqlhost')

    if not sqlhost:

        raise error.Abort(_('please set infinitepush.sqlhost'))

    host, port, db, user, password = sqlhost.split(':')

    reponame = ui.config('infinitepush', 'reponame')

    if not reponame:

        raise error.Abort(_('please set infinitepush.reponame'))

    logfile = ui.config('infinitepush', 'logfile', '')

    waittimeout = ui.configint('infinitepush', 'waittimeout', 300)

    locktimeout = ui.configint('infinitepush', 'locktimeout', 120)

    from . import sqlindexapi

    return sqlindexapi.sqlindexapi(

        reponame, host, port, db, user, password,

        logfile, _getloglevel(ui), waittimeout=waittimeout,

        locktimeout=locktimeout)

def _getloglevel(ui):

    loglevel = ui.config('infinitepush', 'loglevel', 'DEBUG')

    numeric_loglevel = getattr(logging, loglevel.upper(), None)

    if not isinstance(numeric_loglevel, int):

        raise error.Abort(_('invalid log level %s') % loglevel)

    return numeric_loglevel

def _tryhoist(ui, remotebookmark):

    '''returns a bookmarks with hoisted part removed

    Remotenames extension has a 'hoist' config that allows to use remote

    bookmarks without specifying remote path. For example, 'hg update master'

    works as well as 'hg update remote/master'. We want to allow the same in

    infinitepush.

    '''

    if common.isremotebooksenabled(ui):

        hoist = ui.config('remotenames', 'hoistedpeer') + '/'

        if remotebookmark.startswith(hoist):

            return remotebookmark[len(hoist):]

    return remotebookmark

class bundlestore(object):

    def __init__(self, repo):

        self._repo = repo

        storetype = self._repo.ui.config('infinitepush', 'storetype')

        if storetype == 'disk':

            from . import store

            self.store = store.filebundlestore(self._repo.ui, self._repo)

        elif storetype == 'external':

            self.store = _buildexternalbundlestore(self._repo.ui)

        else:

            raise error.Abort(

                _('unknown infinitepush store type specified %s') % storetype)

        indextype = self._repo.ui.config('infinitepush', 'indextype')

        if indextype == 'disk':

            from . import fileindexapi

            self.index = fileindexapi.fileindexapi(self._repo)

        elif indextype == 'sql':

            self.index = _buildsqlindex(self._repo.ui)

        else:

            raise error.Abort(

                _('unknown infinitepush index type specified %s') % indextype)

def _isserver(ui):

    return ui.configbool('infinitepush', 'server')

def reposetup(ui, repo):

    if _isserver(ui) and repo.local():

        repo.bundlestore = bundlestore(repo)

def extsetup(ui):

    commonsetup(ui)

    if _isserver(ui):

        serverextsetup(ui)

    else:

        clientextsetup(ui)

def commonsetup(ui):

    wireprotov1server.commands['listkeyspatterns'] = (

        wireprotolistkeyspatterns, 'namespace patterns')

    scratchbranchpat = ui.config('infinitepush', 'branchpattern')

    if scratchbranchpat:

        global _scratchbranchmatcher

        kind, pat, _scratchbranchmatcher = \

                stringutil.stringmatcher(scratchbranchpat)

def serverextsetup(ui):

    origpushkeyhandler = bundle2.parthandlermapping['pushkey']

    def newpushkeyhandler(*args, **kwargs):

        bundle2pushkey(origpushkeyhandler, *args, **kwargs)

    newpushkeyhandler.params = origpushkeyhandler.params

    bundle2.parthandlermapping['pushkey'] = newpushkeyhandler

    orighandlephasehandler = bundle2.parthandlermapping['phase-heads']

    newphaseheadshandler = lambda *args, **kwargs: \

        bundle2handlephases(orighandlephasehandler, *args, **kwargs)

    newphaseheadshandler.params = orighandlephasehandler.params

    bundle2.parthandlermapping['phase-heads'] = newphaseheadshandler

    extensions.wrapfunction(localrepo.localrepository, 'listkeys',

                            localrepolistkeys)

    wireprotov1server.commands['lookup'] = (

        _lookupwrap(wireprotov1server.commands['lookup'][0]), 'key')

    extensions.wrapfunction(exchange, 'getbundlechunks', getbundlechunks)

    extensions.wrapfunction(bundle2, 'processparts', processparts)

def clientextsetup(ui):

    entry = extensions.wrapcommand(commands.table, 'push', _push)

    entry[1].append(

        ('', 'bundle-store', None,

         _('force push to go to bundle store (EXPERIMENTAL)')))

    extensions.wrapcommand(commands.table, 'pull', _pull)

    extensions.wrapfunction(discovery, 'checkheads', _checkheads)

    wireprotov1peer.wirepeer.listkeyspatterns = listkeyspatterns

    partorder = exchange.b2partsgenorder

    index = partorder.index('changeset')

    partorder.insert(

        index, partorder.pop(partorder.index(scratchbranchparttype)))

def _checkheads(orig, pushop):

    if pushop.ui.configbool(experimental, configscratchpush, False):

        return

    return orig(pushop)

def wireprotolistkeyspatterns(repo, proto, namespace, patterns):

    patterns = wireprototypes.decodelist(patterns)

    d = repo.listkeys(encoding.tolocal(namespace), patterns).iteritems()

    return pushkey.encodekeys(d)

def localrepolistkeys(orig, self, namespace, patterns=None):

    if namespace == 'bookmarks' and patterns:

        index = self.bundlestore.index

        results = {}

        bookmarks = orig(self, namespace)

        for pattern in patterns:

            results.update(index.getbookmarks(pattern))

            if pattern.endswith('*'):

                pattern = 're:^' + pattern[:-1] + '.*'

            kind, pat, matcher = stringutil.stringmatcher(pattern)

            for bookmark, node in bookmarks.iteritems():

                if matcher(bookmark):

                    results[bookmark] = node

        return results

    else:

        return orig(self, namespace)

@wireprotov1peer.batchable

def listkeyspatterns(self, namespace, patterns):

    if not self.capable('pushkey'):

        yield {}, None

    f = wireprotov1peer.future()

    self.ui.debug('preparing listkeys for "%s" with pattern "%s"\n' %

                  (namespace, patterns))

    yield {

        'namespace': encoding.fromlocal(namespace),

        'patterns': wireprototypes.encodelist(patterns)

    }, f

    d = f.value

    self.ui.debug('received listkey for "%s": %i bytes\n'

                  % (namespace, len(d)))

    yield pushkey.decodekeys(d)

def _readbundlerevs(bundlerepo):

    return list(bundlerepo.revs('bundle()'))

def _includefilelogstobundle(bundlecaps, bundlerepo, bundlerevs, ui):

    '''Tells remotefilelog to include all changed files to the changegroup

    By default remotefilelog doesn't include file content to the changegroup.

    But we need to include it if we are fetching from bundlestore.

    '''

    changedfiles = set()

    cl = bundlerepo.changelog

    for r in bundlerevs:

        # [3] means changed files

        changedfiles.update(cl.read(r)[3])

    if not changedfiles:

        return bundlecaps

    changedfiles = '\0'.join(changedfiles)

    newcaps = []

    appended = False

    for cap in (bundlecaps or []):

        if cap.startswith('excludepattern='):

            newcaps.append('\0'.join((cap, changedfiles)))

            appended = True

        else:

            newcaps.append(cap)

    if not appended:

        # Not found excludepattern cap. Just append it

        newcaps.append('excludepattern=' + changedfiles)

    return newcaps

def _rebundle(bundlerepo, bundleroots, unknownhead):

    '''

    Bundle may include more revision then user requested. For example,

    if user asks for revision but bundle also consists its descendants.

    This function will filter out all revision that user is not requested.

    '''

    parts = []

    version = '02'

    outgoing = discovery.outgoing(bundlerepo, commonheads=bundleroots,

                                  missingheads=[unknownhead])

    cgstream = changegroup.makestream(bundlerepo, outgoing, version, 'pull')