<commit_before>import os
import socket
import fnmatch
import json
import docker
import time
import re
import traceback
from threading import Thread
from colorama import Fore, Style

from toolset.utils.output_helper import log
from toolset.utils.metadata_helper import gather_tests
from toolset.utils.database_helper import test_database

def clean(benchmarker_config):
    '''
    Cleans all the docker images from the system
    '''
    # Clean the app server images
    client = docker.DockerClient(
        base_url=benchmarker_config.server_docker_host)

    client.images.prune()
    for image in client.images.list():
        if len(image.tags) > 0:
            # 'techempower/tfb.test.gemini:0.1' -> 'techempower/tfb.test.gemini'
            image_tag = image.tags[0].split(':')[0]
            if image_tag != 'techempower/tfb' and 'techempower' in image_tag:
                client.images.remove(image.id, force=True)
    client.images.prune()

    # Clean the database server images
    client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)

    client.images.prune()
    for image in client.images.list():
        if len(image.tags) > 0:
            # 'techempower/tfb.test.gemini:0.1' -> 'techempower/tfb.test.gemini'
            image_tag = image.tags[0].split(':')[0]
            if image_tag != 'techempower/tfb':
                client.images.remove(image.id, force=True)
    client.images.prune()

def build(benchmarker_config, test_names, build_log_dir=os.devnull):
    '''
    Builds the dependency chain as well as the test implementation docker images
    for the given tests.
    '''
    tests = gather_tests(
        include=test_names, benchmarker_config=benchmarker_config)

    for test in tests:
        log_prefix = "%s: " % test.name

        # Build the test image
        test_docker_file = "%s.dockerfile" % test.name
        build_log_file = build_log_dir
        if build_log_dir is not os.devnull:
            build_log_file = os.path.join(
                build_log_dir,
                "%s.log" % test_docker_file.replace(".dockerfile", "").lower())
        with open(build_log_file, 'w') as build_log:
            try:
                for line in docker.APIClient(
                        base_url=benchmarker_config.server_docker_host).build(
                            path=test.directory,
                            dockerfile=test_docker_file,
                            tag="techempower/tfb.test.%s" %
                            test_docker_file.replace(".dockerfile", ""),
                            forcerm=True):
                    if line.startswith('{"stream":'):
                        line = json.loads(line)
                        line = line[line.keys()[0]].encode('utf-8')
                        log(line,
                            prefix=log_prefix,
                            file=build_log,
                            color=Fore.WHITE + Style.BRIGHT \
                                if re.match(r'^Step \d+\/\d+', line) else '')
            except Exception:
                tb = traceback.format_exc()
                log("Docker build failed; terminating",
                    prefix=log_prefix,
                    file=build_log,
                    color=Fore.RED)
                log(tb, prefix=log_prefix, file=build_log)
                return 1

    return 0

def run(benchmarker_config, test, run_log_dir):
    '''
    Run the given Docker container(s)
    '''
    client = docker.DockerClient(
        base_url=benchmarker_config.server_docker_host)

    log_prefix = "%s: " % test.name
    container = None

    try:

        def watch_container(docker_container, docker_file):
            with open(
                    os.path.join(run_log_dir, "%s.log" % docker_file.replace(
                        ".dockerfile", "").lower()), 'w') as run_log:
                for line in docker_container.logs(stream=True):
                    log(line, prefix=log_prefix, file=run_log)

        extra_hosts = None
        name = "tfb-server"

        if benchmarker_config.network is None:
            extra_hosts = {
                socket.gethostname(): str(benchmarker_config.server_host),
                'tfb-server': str(benchmarker_config.server_host),
                'tfb-database': str(benchmarker_config.database_host)
            }
            name = None

        sysctl = {'net.core.somaxconn': 65535}

        ulimit = [{
            'name': 'nofile',
            'hard': 200000,
            'soft': 200000
        }, {
            'name': 'rtprio',
            'hard': 99,
            'soft': 99
        }]

        container = client.containers.run(
            "techempower/tfb.test.%s" % test.name,
            name=name,
            network=benchmarker_config.network,
            network_mode=benchmarker_config.network_mode,
            stderr=True,
            detach=True,
            init=True,
            extra_hosts=extra_hosts,
            privileged=True,
            ulimits=ulimit,
            sysctls=sysctl)

        watch_thread = Thread(
            target=watch_container,
            args=(
                container,
                "%s.dockerfile" % test.name,
            ))
        watch_thread.daemon = True
        watch_thread.start()

    except Exception:
        with open(
                os.path.join(run_log_dir, "%s.log" % test.name.lower()),
                'w') as run_log:
            tb = traceback.format_exc()
            log("Running docker cointainer: %s.dockerfile failed" % test.name,
                prefix=log_prefix,
                file=run_log)
            log(tb, prefix=log_prefix, file=run_log)

    return container

def stop(benchmarker_config=None,
         container=None,
         database_container=None,
         test=None):
    '''
    Attempts to stop the running test container.
    '''
    client = docker.DockerClient(
        base_url=benchmarker_config.server_docker_host)
    if container is None:
        for container in client.containers.list():
            if len(
                    container.image.tags
            ) > 0 and 'techempower' in container.image.tags[0] and 'tfb:latest' not in container.image.tags[0]:
                container.stop()
    else:
        # Stop the running container
        container.stop()

    database_client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)
    # Stop the database container
    if database_container is None:
        for container in database_client.containers.list():
            if len(
                    container.image.tags
            ) > 0 and 'techempower' in container.image.tags[0] and 'tfb:latest' not in container.image.tags[0]:
                container.stop()
    else:
        database_container.stop()

    client.containers.prune()

    if benchmarker_config.server_docker_host != benchmarker_config.database_docker_host:
        database_client.containers.prune()

def find(path, pattern):
    '''
    Finds and returns all the the files matching the given pattern recursively in
    the given path.
    '''
    for root, dirs, files in os.walk(path):
        for name in files:
            if fnmatch.fnmatch(name, pattern):
                return os.path.join(root, name)

def start_database(benchmarker_config, test, database):
    '''
    Sets up a container for the given database and port, and starts said docker
    container.
    '''
    image_name = "techempower/%s:latest" % database
    log_prefix = image_name + ": "

    database_dir = os.path.join(benchmarker_config.fwroot, "toolset",
                                "databases", database)
    docker_file = "%s.dockerfile" % database

    client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)
    try:
        # Don't pull if we have it
        client.images.get(image_name)
        log("Found published image; skipping build", prefix=log_prefix)
    except:
        # Build the database image
        for line in docker.APIClient(
                base_url=benchmarker_config.database_docker_host).build(
                    path=database_dir,
                    dockerfile=docker_file,
                    tag="techempower/%s" % database):
            if line.startswith('{"stream":'):
                line = json.loads(line)
                line = line[line.keys()[0]].encode('utf-8')
                log(line,
                    prefix=log_prefix,
                    color=Fore.WHITE + Style.BRIGHT \
                        if re.match(r'^Step \d+\/\d+', line) else '')

    client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)

    sysctl = {'net.core.somaxconn': 65535, 'kernel.sem': "250 32000 256 512"}

    ulimit = [{'name': 'nofile', 'hard': 65535, 'soft': 65535}]

    container = client.containers.run(
        "techempower/%s" % database,
        name="tfb-database",
        network=benchmarker_config.network,
        network_mode=benchmarker_config.network_mode,
        detach=True,
        ulimits=ulimit,
        sysctls=sysctl)

    # Sleep until the database accepts connections
    slept = 0
    max_sleep = 60
    database_ready = False
    while not database_ready and slept < max_sleep:
        time.sleep(1)
        slept += 1
        database_ready = test_database(benchmarker_config, database)

    if not database_ready:
        log("Database was not ready after startup", prefix=log_prefix)

    return container

def test_client_connection(benchmarker_config, url):
    '''
    Tests that the app server at the given url responds successfully to a
    request.
    '''
    client = docker.DockerClient(
        base_url=benchmarker_config.client_docker_host)

    try:
        client.images.get('techempower/tfb.wrk:latest')
    except:
        log("Attempting docker pull for image (this can take some time)",
            prefix="techempower/tfb.wrk:latest: ")
        client.images.pull('techempower/tfb.wrk:latest')

    try:
        client.containers.run(
            'techempower/tfb.wrk',
            'curl %s' % url,
            network=benchmarker_config.network,
            network_mode=benchmarker_config.network_mode)
    except:
        return False

    return True

def benchmark(benchmarker_config, script, variables, raw_file):
    '''
    Runs the given remote_script on the wrk container on the client machine.
    '''

    def watch_container(container, raw_file):
        with open(raw_file, 'w') as benchmark_file:
            for line in container.logs(stream=True):
                log(line, file=benchmark_file)

    client = docker.DockerClient(
        base_url=benchmarker_config.client_docker_host)

    sysctl = {'net.core.somaxconn': 65535}

    ulimit = [{'name': 'nofile', 'hard': 65535, 'soft': 65535}]

    watch_container(
        client.containers.run(
            "techempower/tfb.wrk",
            "/bin/bash /%s" % script,
            environment=variables,
            network=benchmarker_config.network,
            network_mode=benchmarker_config.network_mode,
            detach=True,
            stderr=True,
            ulimits=ulimit,
            sysctls=sysctl), raw_file)
<commit_msg>check dockerhub for latest during build (#3585)<commit_after>import os
import socket
import fnmatch
import json
import docker
import time
import re
import traceback
from threading import Thread
from colorama import Fore, Style

from toolset.utils.output_helper import log
from toolset.utils.metadata_helper import gather_tests
from toolset.utils.database_helper import test_database

def clean(benchmarker_config):
    '''
    Cleans all the docker images from the system
    '''
    # Clean the app server images
    client = docker.DockerClient(
        base_url=benchmarker_config.server_docker_host)

    client.images.prune()
    for image in client.images.list():
        if len(image.tags) > 0:
            # 'techempower/tfb.test.gemini:0.1' -> 'techempower/tfb.test.gemini'
            image_tag = image.tags[0].split(':')[0]
            if image_tag != 'techempower/tfb' and 'techempower' in image_tag:
                client.images.remove(image.id, force=True)
    client.images.prune()

    # Clean the database server images
    client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)

    client.images.prune()
    for image in client.images.list():
        if len(image.tags) > 0:
            # 'techempower/tfb.test.gemini:0.1' -> 'techempower/tfb.test.gemini'
            image_tag = image.tags[0].split(':')[0]
            if image_tag != 'techempower/tfb':
                client.images.remove(image.id, force=True)
    client.images.prune()

def build(benchmarker_config, test_names, build_log_dir=os.devnull):
    '''
    Builds the dependency chain as well as the test implementation docker images
    for the given tests.
    '''
    tests = gather_tests(
        include=test_names, benchmarker_config=benchmarker_config)

    for test in tests:
        log_prefix = "%s: " % test.name

        # Build the test image
        test_docker_file = "%s.dockerfile" % test.name
        build_log_file = build_log_dir
        if build_log_dir is not os.devnull:
            build_log_file = os.path.join(
                build_log_dir,
                "%s.log" % test_docker_file.replace(".dockerfile", "").lower())
        with open(build_log_file, 'w') as build_log:
            try:
                for line in docker.APIClient(
                        base_url=benchmarker_config.server_docker_host).build(
                            path=test.directory,
                            dockerfile=test_docker_file,
                            tag="techempower/tfb.test.%s" %
                            test_docker_file.replace(".dockerfile", ""),
                            forcerm=True,
                            pull=True):
                    if line.startswith('{"stream":'):
                        line = json.loads(line)
                        line = line[line.keys()[0]].encode('utf-8')
                        log(line,
                            prefix=log_prefix,
                            file=build_log,
                            color=Fore.WHITE + Style.BRIGHT \
                                if re.match(r'^Step \d+\/\d+', line) else '')
            except Exception:
                tb = traceback.format_exc()
                log("Docker build failed; terminating",
                    prefix=log_prefix,
                    file=build_log,
                    color=Fore.RED)
                log(tb, prefix=log_prefix, file=build_log)
                return 1

    return 0

def run(benchmarker_config, test, run_log_dir):
    '''
    Run the given Docker container(s)
    '''
    client = docker.DockerClient(
        base_url=benchmarker_config.server_docker_host)

    log_prefix = "%s: " % test.name
    container = None

    try:

        def watch_container(docker_container, docker_file):
            with open(
                    os.path.join(run_log_dir, "%s.log" % docker_file.replace(
                        ".dockerfile", "").lower()), 'w') as run_log:
                for line in docker_container.logs(stream=True):
                    log(line, prefix=log_prefix, file=run_log)

        extra_hosts = None
        name = "tfb-server"

        if benchmarker_config.network is None:
            extra_hosts = {
                socket.gethostname(): str(benchmarker_config.server_host),
                'tfb-server': str(benchmarker_config.server_host),
                'tfb-database': str(benchmarker_config.database_host)
            }
            name = None

        sysctl = {'net.core.somaxconn': 65535}

        ulimit = [{
            'name': 'nofile',
            'hard': 200000,
            'soft': 200000
        }, {
            'name': 'rtprio',
            'hard': 99,
            'soft': 99
        }]

        container = client.containers.run(
            "techempower/tfb.test.%s" % test.name,
            name=name,
            network=benchmarker_config.network,
            network_mode=benchmarker_config.network_mode,
            stderr=True,
            detach=True,
            init=True,
            extra_hosts=extra_hosts,
            privileged=True,
            ulimits=ulimit,
            sysctls=sysctl)

        watch_thread = Thread(
            target=watch_container,
            args=(
                container,
                "%s.dockerfile" % test.name,
            ))
        watch_thread.daemon = True
        watch_thread.start()

    except Exception:
        with open(
                os.path.join(run_log_dir, "%s.log" % test.name.lower()),
                'w') as run_log:
            tb = traceback.format_exc()
            log("Running docker cointainer: %s.dockerfile failed" % test.name,
                prefix=log_prefix,
                file=run_log)
            log(tb, prefix=log_prefix, file=run_log)

    return container

def stop(benchmarker_config=None,
         container=None,
         database_container=None,
         test=None):
    '''
    Attempts to stop the running test container.
    '''
    client = docker.DockerClient(
        base_url=benchmarker_config.server_docker_host)
    if container is None:
        for container in client.containers.list():
            if len(
                    container.image.tags
            ) > 0 and 'techempower' in container.image.tags[0] and 'tfb:latest' not in container.image.tags[0]:
                container.stop()
    else:
        # Stop the running container
        container.stop()

    database_client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)
    # Stop the database container
    if database_container is None:
        for container in database_client.containers.list():
            if len(
                    container.image.tags
            ) > 0 and 'techempower' in container.image.tags[0] and 'tfb:latest' not in container.image.tags[0]:
                container.stop()
    else:
        database_container.stop()

    client.containers.prune()

    if benchmarker_config.server_docker_host != benchmarker_config.database_docker_host:
        database_client.containers.prune()

def find(path, pattern):
    '''
    Finds and returns all the the files matching the given pattern recursively in
    the given path.
    '''
    for root, dirs, files in os.walk(path):
        for name in files:
            if fnmatch.fnmatch(name, pattern):
                return os.path.join(root, name)

def start_database(benchmarker_config, test, database):
    '''
    Sets up a container for the given database and port, and starts said docker
    container.
    '''
    image_name = "techempower/%s:latest" % database
    log_prefix = image_name + ": "

    database_dir = os.path.join(benchmarker_config.fwroot, "toolset",
                                "databases", database)
    docker_file = "%s.dockerfile" % database

    client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)
    try:
        # Don't pull if we have it
        client.images.get(image_name)
        log("Found published image; skipping build", prefix=log_prefix)
    except:
        # Build the database image
        for line in docker.APIClient(
                base_url=benchmarker_config.database_docker_host).build(
                    path=database_dir,
                    dockerfile=docker_file,
                    tag="techempower/%s" % database):
            if line.startswith('{"stream":'):
                line = json.loads(line)
                line = line[line.keys()[0]].encode('utf-8')
                log(line,
                    prefix=log_prefix,
                    color=Fore.WHITE + Style.BRIGHT \
                        if re.match(r'^Step \d+\/\d+', line) else '')

    client = docker.DockerClient(
        base_url=benchmarker_config.database_docker_host)

    sysctl = {'net.core.somaxconn': 65535, 'kernel.sem': "250 32000 256 512"}

    ulimit = [{'name': 'nofile', 'hard': 65535, 'soft': 65535}]

    container = client.containers.run(
        "techempower/%s" % database,
        name="tfb-database",
        network=benchmarker_config.network,
        network_mode=benchmarker_config.network_mode,
        detach=True,
        ulimits=ulimit,
        sysctls=sysctl)

    # Sleep until the database accepts connections
    slept = 0
    max_sleep = 60
    database_ready = False
    while not database_ready and slept < max_sleep:
        time.sleep(1)
        slept += 1
        database_ready = test_database(benchmarker_config, database)

    if not database_ready:
        log("Database was not ready after startup", prefix=log_prefix)

    return container

def test_client_connection(benchmarker_config, url):
    '''
    Tests that the app server at the given url responds successfully to a
    request.
    '''
    client = docker.DockerClient(
        base_url=benchmarker_config.client_docker_host)

    try:
        client.images.get('techempower/tfb.wrk:latest')
    except:
        log("Attempting docker pull for image (this can take some time)",
            prefix="techempower/tfb.wrk:latest: ")
        client.images.pull('techempower/tfb.wrk:latest')

    try:
        client.containers.run(
            'techempower/tfb.wrk',
            'curl %s' % url,
            network=benchmarker_config.network,
            network_mode=benchmarker_config.network_mode)
    except:
        return False

    return True

def benchmark(benchmarker_config, script, variables, raw_file):
    '''
    Runs the given remote_script on the wrk container on the client machine.
    '''

    def watch_container(container, raw_file):
        with open(raw_file, 'w') as benchmark_file:
            for line in container.logs(stream=True):
                log(line, file=benchmark_file)

    client = docker.DockerClient(
        base_url=benchmarker_config.client_docker_host)

    sysctl = {'net.core.somaxconn': 65535}

    ulimit = [{'name': 'nofile', 'hard': 65535, 'soft': 65535}]

    watch_container(
        client.containers.run(
            "techempower/tfb.wrk",
            "/bin/bash /%s" % script,
            environment=variables,
            network=benchmarker_config.network,
            network_mode=benchmarker_config.network_mode,
            detach=True,
            stderr=True,
            ulimits=ulimit,
            sysctls=sysctl), raw_file)
<|endoftext|>


<commit_before>
QUICKRELEASE_CONSTANTS = {
   'MAKE' : 'make',
   'MV' : 'mv',
   'PERL' : 'perl',
   'S3CURL' : 's3curl.pl',
   'SVN' : 'svn',
   'RSYNC' : 'rsync',
   'UNZIP' : 'unzip',
   'WGET' : 'wget',

   'BUILD_PLATFORMS': ('windows-i686-msvc8',
                       'macosx-i686',
                       'linux-i686',
                       'linux-x86_64'),

   'BUILD_PLATFORM_EXTENSIONS': { 'windows-i686-msvc8': 'exe',
                                  'macosx-i686': 'dmg',
                                  'linux-i686': 'tar.gz',
                                  'linux-x86_64': 'tar.gz' },

   # in seconds, so five minutes
   'RUN_SHELL_COMMAND_DEFAULT_TIMEOUT': 60 * 5,

   # in seconds, so 10 mintues.
   'S3_PUSH_TIMEOUT': 60 * 10,

   'S3_MIME_TYPES': { 'asc' : 'text/plain',
                      'bz2' : 'application/x-bzip2',
                      'dmg' : 'application/x-apple-diskimage',
                      'exe' : 'application/octet-stream',
                      'mar' : 'application/octet-stream',
                      'md5' : 'text/plain',
                      'tar.gz' : 'application/x-gzip',
                      'txt': 'text/plain',
                      'zip': 'application/zip',
                    },
}

CONSTANTS_FROM_ENV_HANDLERS = {
   'BUILD_PLATFORMS': lambda val: tuple(val.split()),
   'BUILD_PLATFORM_EXTENSIONS': lambda val: NotImplementedError("Need to turn BUILD_PLATFORM_EXTENSIONS overloads into a dict!"),
   'S3_MIME_TYPES': lambda val: NotImplementedError("Need to turn S3_MIME_TYPES overloads into a dict!"),
}

<commit_msg>Add constants for tools required for tarball generation.<commit_after>
QUICKRELEASE_CONSTANTS = {
   'BZIP_PROG': 'bzip2',
   'MAKE' : 'make',
   'MD5SUM' : 'md5sum',
   'MV' : 'mv',
   'PERL' : 'perl',
   'S3CURL' : 's3curl.pl',
   'SVN' : 'svn',
   'RSYNC' : 'rsync',
   'TAR' : 'tar',
   'UNZIP' : 'unzip',
   'WGET' : 'wget',

   'BUILD_PLATFORMS': ('windows-i686-msvc8',
                       'macosx-i686',
                       'linux-i686',
                       'linux-x86_64'),

   'BUILD_PLATFORM_EXTENSIONS': { 'windows-i686-msvc8': 'exe',
                                  'macosx-i686': 'dmg',
                                  'linux-i686': 'tar.gz',
                                  'linux-x86_64': 'tar.gz' },

   # in seconds, so five minutes
   'RUN_SHELL_COMMAND_DEFAULT_TIMEOUT': 60 * 5,

   # in seconds, so 10 mintues.
   'S3_PUSH_TIMEOUT': 60 * 10,

   'S3_MIME_TYPES': { 'asc' : 'text/plain',
                      'bz2' : 'application/x-bzip2',
                      'dmg' : 'application/x-apple-diskimage',
                      'exe' : 'application/octet-stream',
                      'mar' : 'application/octet-stream',
                      'md5' : 'text/plain',
                      'tar.gz' : 'application/x-gzip',
                      'txt': 'text/plain',
                      'zip': 'application/zip',
                    },
}

CONSTANTS_FROM_ENV_HANDLERS = {
   'BUILD_PLATFORMS': lambda val: tuple(val.split()),
   'BUILD_PLATFORM_EXTENSIONS': lambda val: NotImplementedError("Need to turn BUILD_PLATFORM_EXTENSIONS overloads into a dict!"),
   'S3_MIME_TYPES': lambda val: NotImplementedError("Need to turn S3_MIME_TYPES overloads into a dict!"),
}

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2014, 2015 Cyan, Inc.
# Copyright 2018, 2019 Ciena Corporation
#
# 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 logging
import os
import os.path
import re
import shutil
import subprocess
import tempfile
import uuid
from datetime import datetime

from six.moves.urllib.parse import urlparse

from .service import ExternalService, SpawnedService
from .testutil import get_open_port, random_string

class KafkaHarness(object):
    """
    `KafkaHarness` starts a Kafka cluster.

    :ivar str chroot:
        ZooKeeper chroot the Kafka brokers are configured to use.
    :ivar zk: `_ZookeeperFixture` instance
    :ivar brokers: `list` of `_KafakBroker` instances
    """
    @classmethod
    def start(cls, replicas, **kw):
        """
        Create and run a `KafkaHarness`.

        :param int replicas:
            Number of Kafka brokers to start. Must be an odd number.
        :param kw:
            Keyword arguments passed through to `_KafkaFixture`. The
            *replicas*, *zk_host*, *zk_port*, and *zk_chroot* are set
            automatically. *partitions* defaults to 8.
        """
        kw['zk_chroot'] = chroot = random_string(10)
        zk = _ZookeeperFixture.instance(chroot)
        kw['replicas'] = replicas
        kw.setdefault('partitions', 8)
        kw['zk_host'] = zk.host
        kw['zk_port'] = zk.port

        brokers = []
        for broker_id in range(replicas):
            brokers.append(_KafkaFixture.instance(broker_id, **kw))

        return cls(chroot, zk, brokers)

    def __init__(self, chroot, zk, brokers):
        """
        Private constructor.
        """
        self.chroot = chroot
        self.zk = zk
        self.brokers = brokers

    def halt(self):
        for broker in self.brokers:
            broker.close()
        self.zk.close()

    @property
    def bootstrap_hosts(self):
        """
        List of (host, port) tuples for the brokers.
        """
        return ["{}:{}".format(b.host, b.port) for b in self.brokers]

class _Fixture(object):
    kafka_version = os.environ.get('KAFKA_VERSION', '0.9.0.1')
    scala_version = os.environ.get("SCALA_VERSION", '2.10.0')
    project_root = os.environ.get(
        'PROJECT_ROOT', os.path.abspath(
            os.path.join(os.path.dirname(__file__), "../..")))
    kafka_root = os.environ.get(
        "KAFKA_ROOT", os.path.join(project_root, 'servers',
                                   kafka_version, "kafka-bin"))

    @classmethod
    def test_resource(cls, filename):
        return os.path.join(cls.project_root, "servers", cls.kafka_version,
                            "resources", filename)

    @classmethod
    def kafka_run_class_args(cls, *args):
        result = [os.path.join(cls.kafka_root, 'bin', 'kafka-run-class.sh')]
        result.extend(args)
        return result

    @classmethod
    def kafka_run_class_env(cls):
        env = os.environ.copy()
        env['KAFKA_LOG4J_OPTS'] = "-Dlog4j.configuration=file:{}".format(
            cls.test_resource("log4j.properties"))
        return env

    @classmethod
    def render_template(cls, source_file, target_file, binding):
        with open(source_file, "r") as handle:
            template = handle.read()
        with open(target_file, "w") as handle:
            handle.write(template.format(**binding))

class _ZookeeperFixture(_Fixture):
    @classmethod
    def instance(cls, kafka_chroot):
        if "ZOOKEEPER_URI" in os.environ:  # pragma: no cover
            parse = urlparse(os.environ["ZOOKEEPER_URI"])
            (host, port) = (parse.hostname, parse.port)
            fixture = ExternalService(host, port)
        else:
            (host, port) = ("127.0.0.1", get_open_port())
            fixture = cls(host, port)

        fixture.open(kafka_chroot)
        return fixture

    def __init__(self, host, port):
        self.host = host
        self.port = port

        self.tmp_dir = None
        self.child = None
        self._log = logging.getLogger('fixtures.ZK')

    def __repr__(self):
        return "{}<{}:{}>".format(
            self.__class__.__name__,
            self.host,
            self.port,
        )

    def open(self, kafka_chroot):
        self.tmp_dir = tempfile.mkdtemp()
        properties_file = os.path.join(self.tmp_dir, "zookeeper.properties")

        properties = (
            "dataDir={tmp_dir}\n"
            "clientPort={port}\n"
            "clientPortAddress={host}\n"
            "maxClientCnxns=0\n"
            # Use reduced timeouts to speed up failovers.
            "tickTime=2000\n"
            "maxSessionTimeout=6000\n"
        ).format(
            tmp_dir=self.tmp_dir,
            host=self.host,
            port=self.port,
        )
        with open(properties_file, 'w') as f:
            f.write(properties)
        self._log.info("Running local instance with config:\n%s", properties)

        # Configure Zookeeper child process
        args = self.kafka_run_class_args("org.apache.zookeeper.server.quorum.QuorumPeerMain", properties_file)
        env = self.kafka_run_class_env()
        start_re = re.compile(
            "binding to port /{host}:|Starting server.*ZooKeeperServerMain"
            .format(host=re.escape(self.host)),
        )
        self._child = SpawnedService('zookeeper', self._log, args, env, start_re)
        self._child.start()

        self._log.debug("Creating ZooKeeper chroot node %r...", kafka_chroot)
        args = self.kafka_run_class_args(
            "org.apache.zookeeper.ZooKeeperMain",
            "-server", "%s:%d" % (self.host, self.port),
            "create",
            "/%s" % kafka_chroot,
            "afkak",
        )
        env = self.kafka_run_class_env()
        proc = subprocess.Popen(args, env=env, stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT)
        stdout, _stderr = proc.communicate()
        if proc.returncode != 0:  # pragma: no cover
            self._log.error("Failed to create Zookeeper chroot node. Process exited %d and output:\n%s",
                            proc.returncode, stdout.decode('utf-8', 'replace'))
            raise RuntimeError("Failed to create Zookeeper chroot node")
        self._log.debug("Done!")

    def close(self):
        self._log.debug("Stopping...")
        self._child.stop()
        self._child = None
        self._log.debug("Done!")
        shutil.rmtree(self.tmp_dir)

class _KafkaFixture(_Fixture):

    @classmethod
    def instance(cls, broker_id, zk_host, zk_port, zk_chroot, replicas,
                 partitions, message_max_bytes=1000000):
        if zk_chroot is None:
            zk_chroot = "afkak_" + str(uuid.uuid4()).replace("-", "_")
        if "KAFKA_URI" in os.environ:  # pragma: no cover
            parse = urlparse(os.environ["KAFKA_URI"])
            (host, port) = (parse.hostname, parse.port)
            fixture = ExternalService(host, port)
        else:
            (host, port) = ("127.0.0.1", get_open_port())
            fixture = cls(
                host=host, port=port, broker_id=broker_id, zk_host=zk_host,
                zk_port=zk_port, zk_chroot=zk_chroot, replicas=replicas,
                partitions=partitions, message_max_bytes=message_max_bytes,
            )
            fixture.open()
        return fixture

    def __init__(self, host, port, broker_id, zk_host, zk_port, zk_chroot,
                 replicas, partitions, message_max_bytes):
        self.host = host
        self.port = port
        self._log = logging.getLogger('fixtures.K{}'.format(broker_id))

        self.broker_id = broker_id

        self.zk_host = zk_host
        self.zk_port = zk_port
        self.zk_chroot = zk_chroot

        self.replicas = replicas
        self.partitions = partitions
        assert replicas % 2 == 1, "replica count must be odd, not {}".format(replicas)
        self.min_insync_replicas = replicas // 2 + 1

        self.message_max_bytes = message_max_bytes

        self.tmp_dir = None
        self.child = None
        self.running = False
        self.restartable = False  # Only restartable after stop() call

    def __repr__(self):
        return '{}<node_id={} {}:{} {}>'.format(
            self.__class__.__name__,
            self.broker_id,
            self.host,
            self.port,
            'running' if self.running else 'stopped',
        )

    def open(self):
        if self.running:  # pragma: no cover
            self._log.debug("Instance already running")
            return

        self.tmp_dir = tempfile.mkdtemp()
        self._properties_file = os.path.join(self.tmp_dir, "kafka.properties")
        self._log.info("Running local instance...")
        self._log.info("  host       = %s", self.host)
        self._log.info("  port       = %s", self.port)
        self._log.info("  broker_id  = %s", self.broker_id)
        self._log.info("  zk_host    = %s", self.zk_host)
        self._log.info("  zk_port    = %s", self.zk_port)
        self._log.info("  zk_chroot  = %s", self.zk_chroot)
        self._log.info("  replicas   = %s", self.replicas)
        self._log.info("  partitions = %s", self.partitions)
        self._log.info("  msg_max_sz = %s", self.message_max_bytes)
        self._log.info("  tmp_dir    = %s", self.tmp_dir)

        # Create directories
        os.mkdir(os.path.join(self.tmp_dir, "logs"))
        os.mkdir(os.path.join(self.tmp_dir, "data"))

        # Generate configs
        template = self.test_resource("kafka.properties")
        self.render_template(template, self._properties_file, vars(self))

        self._log.debug("Starting...")
        self.child = self._make_child()
        self.child.start()
        self._log.debug("Done!")
        self.running = True

    def _make_child(self):
        """
        Configure Kafka child process
        """
        name = 'kafka{}'.format(self.broker_id)
        args = self.kafka_run_class_args("kafka.Kafka", self._properties_file)
        env = self.kafka_run_class_env()
        # Match a message like (0.9.0.1 and earlier):
        #
        #     [2018-07-17 18:06:00,915] INFO [Kafka Server 0], started (kafka.server.KafkaServer)
        #
        # Or like this (1.1.1):
        #
        #     [2018-09-27 17:23:46,818] INFO [KafkaServer id=0] started (kafka.server.KafkaServer)
        start_re = re.compile((
            r"("
            r"\[Kafka Server {broker_id}\], [Ss]tarted"
            r"|"
            r"\[KafkaServer id={broker_id}\] started"
            r")"
        ).format(broker_id=self.broker_id))
        return SpawnedService(name, self._log, args, env, start_re)

    def close(self):
        if not self.running:  # pragma: no cover
            self._log.warning("Instance already stopped")
            return

        self._log.info("Stopping... %s at %s",
                       self.tmp_dir, datetime.utcnow().isoformat())
        self.child.stop()
        self.child = None
        self._log.info("Done!")
        shutil.rmtree(self.tmp_dir)
        self.running = False

    def stop(self):
        """Stop/cleanup the child SpawnedService"""
        self.child.stop()
        self.child = None
        self._log.info("Child stopped.")
        self.running = False
        self.restartable = True

    def restart(self):
        """Start a new child SpawnedService with same settings and tmpdir"""
        if not self.restartable:  # pragma: no cover
            self._log.error("*** Kafka [%s:%d]: Restart attempted when not stopped.",
                            self.host, self.port)
            return
        self.child = self._make_child()
        self._log.debug("Starting...")
        self.child.start()
        self._log.debug("Done!")
        self.running = True
        self.restartable = False
<commit_msg>Remove unused scala_version ivar<commit_after># -*- coding: utf-8 -*-
# Copyright 2014, 2015 Cyan, Inc.
# Copyright 2018, 2019 Ciena Corporation
#
# 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 logging
import os
import os.path
import re
import shutil
import subprocess
import tempfile
import uuid
from datetime import datetime

from six.moves.urllib.parse import urlparse

from .service import ExternalService, SpawnedService
from .testutil import get_open_port, random_string

class KafkaHarness(object):
    """
    `KafkaHarness` starts a Kafka cluster.

    :ivar str chroot:
        ZooKeeper chroot the Kafka brokers are configured to use.
    :ivar zk: `_ZookeeperFixture` instance
    :ivar brokers: `list` of `_KafakBroker` instances
    """
    @classmethod
    def start(cls, replicas, **kw):
        """
        Create and run a `KafkaHarness`.

        :param int replicas:
            Number of Kafka brokers to start. Must be an odd number.
        :param kw:
            Keyword arguments passed through to `_KafkaFixture`. The
            *replicas*, *zk_host*, *zk_port*, and *zk_chroot* are set
            automatically. *partitions* defaults to 8.
        """
        kw['zk_chroot'] = chroot = random_string(10)
        zk = _ZookeeperFixture.instance(chroot)
        kw['replicas'] = replicas
        kw.setdefault('partitions', 8)
        kw['zk_host'] = zk.host
        kw['zk_port'] = zk.port

        brokers = []
        for broker_id in range(replicas):
            brokers.append(_KafkaFixture.instance(broker_id, **kw))

        return cls(chroot, zk, brokers)

    def __init__(self, chroot, zk, brokers):
        """
        Private constructor.
        """
        self.chroot = chroot
        self.zk = zk
        self.brokers = brokers

    def halt(self):
        for broker in self.brokers:
            broker.close()
        self.zk.close()

    @property
    def bootstrap_hosts(self):
        """
        List of (host, port) tuples for the brokers.
        """
        return ["{}:{}".format(b.host, b.port) for b in self.brokers]

class _Fixture(object):
    kafka_version = os.environ.get('KAFKA_VERSION', '0.9.0.1')
    project_root = os.environ.get(
        'PROJECT_ROOT', os.path.abspath(
            os.path.join(os.path.dirname(__file__), "../..")))
    kafka_root = os.environ.get(
        "KAFKA_ROOT", os.path.join(project_root, 'servers',
                                   kafka_version, "kafka-bin"))

    @classmethod
    def test_resource(cls, filename):
        return os.path.join(cls.project_root, "servers", cls.kafka_version,
                            "resources", filename)

    @classmethod
    def kafka_run_class_args(cls, *args):
        result = [os.path.join(cls.kafka_root, 'bin', 'kafka-run-class.sh')]
        result.extend(args)
        return result

    @classmethod
    def kafka_run_class_env(cls):
        env = os.environ.copy()
        env['KAFKA_LOG4J_OPTS'] = "-Dlog4j.configuration=file:{}".format(
            cls.test_resource("log4j.properties"))
        return env

    @classmethod
    def render_template(cls, source_file, target_file, binding):
        with open(source_file, "r") as handle:
            template = handle.read()
        with open(target_file, "w") as handle:
            handle.write(template.format(**binding))

class _ZookeeperFixture(_Fixture):
    @classmethod
    def instance(cls, kafka_chroot):
        if "ZOOKEEPER_URI" in os.environ:  # pragma: no cover
            parse = urlparse(os.environ["ZOOKEEPER_URI"])
            (host, port) = (parse.hostname, parse.port)
            fixture = ExternalService(host, port)
        else:
            (host, port) = ("127.0.0.1", get_open_port())
            fixture = cls(host, port)

        fixture.open(kafka_chroot)
        return fixture

    def __init__(self, host, port):
        self.host = host
        self.port = port

        self.tmp_dir = None
        self.child = None
        self._log = logging.getLogger('fixtures.ZK')

    def __repr__(self):
        return "{}<{}:{}>".format(
            self.__class__.__name__,
            self.host,
            self.port,
        )

    def open(self, kafka_chroot):
        self.tmp_dir = tempfile.mkdtemp()
        properties_file = os.path.join(self.tmp_dir, "zookeeper.properties")

        properties = (
            "dataDir={tmp_dir}\n"
            "clientPort={port}\n"
            "clientPortAddress={host}\n"
            "maxClientCnxns=0\n"
            # Use reduced timeouts to speed up failovers.
            "tickTime=2000\n"
            "maxSessionTimeout=6000\n"
        ).format(
            tmp_dir=self.tmp_dir,
            host=self.host,
            port=self.port,
        )
        with open(properties_file, 'w') as f:
            f.write(properties)
        self._log.info("Running local instance with config:\n%s", properties)

        # Configure Zookeeper child process
        args = self.kafka_run_class_args("org.apache.zookeeper.server.quorum.QuorumPeerMain", properties_file)
        env = self.kafka_run_class_env()
        start_re = re.compile(
            "binding to port /{host}:|Starting server.*ZooKeeperServerMain"
            .format(host=re.escape(self.host)),
        )
        self._child = SpawnedService('zookeeper', self._log, args, env, start_re)
        self._child.start()

        self._log.debug("Creating ZooKeeper chroot node %r...", kafka_chroot)
        args = self.kafka_run_class_args(
            "org.apache.zookeeper.ZooKeeperMain",
            "-server", "%s:%d" % (self.host, self.port),
            "create",
            "/%s" % kafka_chroot,
            "afkak",
        )
        env = self.kafka_run_class_env()
        proc = subprocess.Popen(args, env=env, stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT)
        stdout, _stderr = proc.communicate()
        if proc.returncode != 0:  # pragma: no cover
            self._log.error("Failed to create Zookeeper chroot node. Process exited %d and output:\n%s",
                            proc.returncode, stdout.decode('utf-8', 'replace'))
            raise RuntimeError("Failed to create Zookeeper chroot node")
        self._log.debug("Done!")

    def close(self):
        self._log.debug("Stopping...")
        self._child.stop()
        self._child = None
        self._log.debug("Done!")
        shutil.rmtree(self.tmp_dir)

class _KafkaFixture(_Fixture):

    @classmethod
    def instance(cls, broker_id, zk_host, zk_port, zk_chroot, replicas,
                 partitions, message_max_bytes=1000000):
        if zk_chroot is None:
            zk_chroot = "afkak_" + str(uuid.uuid4()).replace("-", "_")
        if "KAFKA_URI" in os.environ:  # pragma: no cover
            parse = urlparse(os.environ["KAFKA_URI"])
            (host, port) = (parse.hostname, parse.port)
            fixture = ExternalService(host, port)
        else:
            (host, port) = ("127.0.0.1", get_open_port())
            fixture = cls(
                host=host, port=port, broker_id=broker_id, zk_host=zk_host,
                zk_port=zk_port, zk_chroot=zk_chroot, replicas=replicas,
                partitions=partitions, message_max_bytes=message_max_bytes,
            )
            fixture.open()
        return fixture

    def __init__(self, host, port, broker_id, zk_host, zk_port, zk_chroot,
                 replicas, partitions, message_max_bytes):
        self.host = host
        self.port = port
        self._log = logging.getLogger('fixtures.K{}'.format(broker_id))

        self.broker_id = broker_id

        self.zk_host = zk_host
        self.zk_port = zk_port
        self.zk_chroot = zk_chroot

        self.replicas = replicas
        self.partitions = partitions
        assert replicas % 2 == 1, "replica count must be odd, not {}".format(replicas)
        self.min_insync_replicas = replicas // 2 + 1

        self.message_max_bytes = message_max_bytes

        self.tmp_dir = None
        self.child = None
        self.running = False
        self.restartable = False  # Only restartable after stop() call

    def __repr__(self):
        return '{}<node_id={} {}:{} {}>'.format(
            self.__class__.__name__,
            self.broker_id,
            self.host,
            self.port,
            'running' if self.running else 'stopped',
        )

    def open(self):
        if self.running:  # pragma: no cover
            self._log.debug("Instance already running")
            return

        self.tmp_dir = tempfile.mkdtemp()
        self._properties_file = os.path.join(self.tmp_dir, "kafka.properties")
        self._log.info("Running local instance...")
        self._log.info("  host       = %s", self.host)
        self._log.info("  port       = %s", self.port)
        self._log.info("  broker_id  = %s", self.broker_id)
        self._log.info("  zk_host    = %s", self.zk_host)
        self._log.info("  zk_port    = %s", self.zk_port)
        self._log.info("  zk_chroot  = %s", self.zk_chroot)
        self._log.info("  replicas   = %s", self.replicas)
        self._log.info("  partitions = %s", self.partitions)
        self._log.info("  msg_max_sz = %s", self.message_max_bytes)
        self._log.info("  tmp_dir    = %s", self.tmp_dir)

        # Create directories
        os.mkdir(os.path.join(self.tmp_dir, "logs"))
        os.mkdir(os.path.join(self.tmp_dir, "data"))

        # Generate configs
        template = self.test_resource("kafka.properties")
        self.render_template(template, self._properties_file, vars(self))

        self._log.debug("Starting...")
        self.child = self._make_child()
        self.child.start()
        self._log.debug("Done!")
        self.running = True

    def _make_child(self):
        """
        Configure Kafka child process
        """
        name = 'kafka{}'.format(self.broker_id)
        args = self.kafka_run_class_args("kafka.Kafka", self._properties_file)
        env = self.kafka_run_class_env()
        # Match a message like (0.9.0.1 and earlier):
        #
        #     [2018-07-17 18:06:00,915] INFO [Kafka Server 0], started (kafka.server.KafkaServer)
        #
        # Or like this (1.1.1):
        #
        #     [2018-09-27 17:23:46,818] INFO [KafkaServer id=0] started (kafka.server.KafkaServer)
        start_re = re.compile((
            r"("
            r"\[Kafka Server {broker_id}\], [Ss]tarted"
            r"|"
            r"\[KafkaServer id={broker_id}\] started"
            r")"
        ).format(broker_id=self.broker_id))
        return SpawnedService(name, self._log, args, env, start_re)

    def close(self):
        if not self.running:  # pragma: no cover
            self._log.warning("Instance already stopped")
            return

        self._log.info("Stopping... %s at %s",
                       self.tmp_dir, datetime.utcnow().isoformat())
        self.child.stop()
        self.child = None
        self._log.info("Done!")
        shutil.rmtree(self.tmp_dir)
        self.running = False

    def stop(self):
        """Stop/cleanup the child SpawnedService"""
        self.child.stop()
        self.child = None
        self._log.info("Child stopped.")
        self.running = False
        self.restartable = True

    def restart(self):
        """Start a new child SpawnedService with same settings and tmpdir"""
        if not self.restartable:  # pragma: no cover
            self._log.error("*** Kafka [%s:%d]: Restart attempted when not stopped.",
                            self.host, self.port)
            return
        self.child = self._make_child()
        self._log.debug("Starting...")
        self.child.start()
        self._log.debug("Done!")
        self.running = True
        self.restartable = False
<|endoftext|>


<commit_before>"""
byceps.blueprints.snippet.templating
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

:Copyright: 2006-2019 Jochen Kupperschmidt
:License: Modified BSD, see LICENSE for details.
"""

import sys
import traceback

from flask import abort, g, render_template, url_for
from jinja2 import TemplateNotFound

from ...services.snippet import service as snippet_service
from ...services.snippet.service import SnippetNotFound
from ...services.snippet.transfer.models import Scope
from ...util.templating import get_variable_value, load_template

def render_snippet_as_page(version):
    """Render the given version of the snippet, or an error page if
    that fails.
    """
    try:
        context = get_snippet_context(version)
        return render_template('snippet/view.html', **context)
    except TemplateNotFound:
        abort(404)
    except Exception as e:
        print('Error in snippet markup:', e, file=sys.stderr)
        traceback.print_exc()
        context = {
            'message': str(e),
        }
        return render_template('snippet/error.html', **context), 500

def get_snippet_context(version):
    """Return the snippet context to insert into the outer template."""
    template = _load_template_with_globals(version.body)

    current_page = get_variable_value(template, 'current_page')
    title = version.title
    head = _render_template(version.head) if version.head else None
    body = template.render()

    return {
        'title': title,
        'current_page': current_page,
        'head': head,
        'body': body,
    }

def render_snippet_as_partial(
    name, *, scope=None, ignore_if_unknown=False, context=None
):
    """Render the latest version of the snippet with the given name and
    return the result.
    """
    if scope is None:
        scope = Scope.for_site(g.site_id)

    current_version = snippet_service.find_current_version_of_snippet_with_name(
        scope, name
    )

    if current_version is None:
        if ignore_if_unknown:
            return ''
        else:
            raise SnippetNotFound(scope, name)

    if context is None:
        context = {}

    return _render_template(current_version.body, context=context)

def _render_template(source, *, context=None):
    template = _load_template_with_globals(source)

    if context is None:
        context = {}

    return template.render(**context)

def _load_template_with_globals(source):
    template_globals = {
        'render_snippet': render_snippet_as_partial,
        'url_for': url_for,
        'url_for_snippet': url_for_snippet,
    }

    return load_template(source, template_globals=template_globals)

def url_for_snippet(name, **kwargs):
    """Render an URL pointing to the snippet document's mountpoint."""
    return url_for(f'snippet.{name}', **kwargs)
<commit_msg>Clarify argument name<commit_after>"""
byceps.blueprints.snippet.templating
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

:Copyright: 2006-2019 Jochen Kupperschmidt
:License: Modified BSD, see LICENSE for details.
"""

import sys
import traceback

from flask import abort, g, render_template, url_for
from jinja2 import TemplateNotFound

from ...services.snippet import service as snippet_service
from ...services.snippet.service import SnippetNotFound
from ...services.snippet.transfer.models import Scope
from ...util.templating import get_variable_value, load_template

def render_snippet_as_page(version):
    """Render the given version of the snippet, or an error page if
    that fails.
    """
    try:
        context = get_snippet_context(version)
        return render_template('snippet/view.html', **context)
    except TemplateNotFound:
        abort(404)
    except Exception as e:
        print('Error in snippet markup:', e, file=sys.stderr)
        traceback.print_exc()
        context = {
            'message': str(e),
        }
        return render_template('snippet/error.html', **context), 500

def get_snippet_context(version):
    """Return the snippet context to insert into the outer template."""
    template = _load_template_with_globals(version.body)

    current_page = get_variable_value(template, 'current_page')
    title = version.title
    head = _render_template(version.head) if version.head else None
    body = template.render()

    return {
        'title': title,
        'current_page': current_page,
        'head': head,
        'body': body,
    }

def render_snippet_as_partial(
    name, *, scope=None, ignore_if_unknown=False, context=None
):
    """Render the latest version of the snippet with the given name and
    return the result.
    """
    if scope is None:
        scope = Scope.for_site(g.site_id)

    current_version = snippet_service.find_current_version_of_snippet_with_name(
        scope, name
    )

    if current_version is None:
        if ignore_if_unknown:
            return ''
        else:
            raise SnippetNotFound(scope, name)

    if context is None:
        context = {}

    return _render_template(current_version.body, context=context)

def _render_template(source, *, context=None):
    template = _load_template_with_globals(source)

    if context is None:
        context = {}

    return template.render(**context)

def _load_template_with_globals(source):
    template_globals = {
        'render_snippet': render_snippet_as_partial,
        'url_for': url_for,
        'url_for_snippet': url_for_snippet,
    }

    return load_template(source, template_globals=template_globals)

def url_for_snippet(endpoint_suffix, **kwargs):
    """Render an URL pointing to the snippet document's mountpoint."""
    return url_for(f'snippet.{endpoint_suffix}', **kwargs)
<|endoftext|>


<commit_before>import getopt
import tornado.ioloop
import tornado.web
import sys
from api_resources import Resources
from sentiment_analysis import SentimentAnalysis

import time

class GetWordInformation(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos", ['v', 'n', 'a', 'r'])
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}
            if isinstance(pos, list):
                response['synsets'] = {}
                for postagging in pos:
                    pos_result = sentiwordnet.get_info_word(word, postagging, language)
                    if pos_result:
                        response['synsets'].update(pos_result)
                response["elapsed_time"] = time.time() - start
            else:
                response["synsets"] = sentiwordnet.get_info_word(word, pos, language)
                response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class Synonym(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["sentiment"] = sentiwordnet.get_synonym(word, pos, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetAffect(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["affects"] = sentiwordnet.get_affects(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetDommain(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["dommains"] = sentiwordnet.get_domains(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetSentiment(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["sentiment"] = sentiwordnet.get_sentiment(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetSentimentEmotion(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            sentiment = SentimentAnalysis(text, language, self.sentiwordnet)
            sentiment.analyze()
            response = sentiment.tostring()
            self.write(response)
        else:
            # Create DS to save result
            response = {}

            response.update(sentiwordnet.get_sentiment_and_emotion(text, language))
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class PosTagging(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["postagging"] = sentiwordnet.get_postagging(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetInformation(tornado.web.RequestHandler):
    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos")
        language = self.get_argument("language")

        score = sentiwordnet.get_info_first_word(word, pos, language)

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}
            response["score"] = score

            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class Translator(tornado.web.RequestHandler):
    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos")
        from_language = self.get_argument("from_language")
        to_language = self.get_argument("to_language")

        translation = sentiwordnet.get_translation(word, pos, from_language, to_language)

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}
            response["translation"] = translation

            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

# Check if server por is valid
try:
    opts, args = getopt.getopt(sys.argv[1:], "p:", ["port="])
    assert (len(opts) == 1)
except:
    print 'api_resources_service.py -p <port>'
    sys.exit()

for o, a in opts:
    if o == "-p":
        port = a
    else:
        pass

if __name__ == "__main__":
    # Init treetagger with a specific language
    sentiwordnet = Resources()

    # Init Tornado web server
    application = tornado.web.Application([
        (r"/get_information", GetInformation, dict(sentiwordnet=sentiwordnet)),
        (r"/translate", Translator, dict(sentiwordnet=sentiwordnet)),
        (r"/get_synonym", Synonym, dict(sentiwordnet=sentiwordnet)),
        (r"/get_postagging", PosTagging, dict(sentiwordnet=sentiwordnet)),
        (r"/get_sentiment", GetSentiment, dict(sentiwordnet=sentiwordnet)),
        (r"/get_affect", GetAffect, dict(sentiwordnet=sentiwordnet)),
        (r"/get_dommain", GetDommain, dict(sentiwordnet=sentiwordnet)),
        (r"/get_word_information", GetWordInformation, dict(sentiwordnet=sentiwordnet)),
        (r"/get_sentiment_emotion", GetSentimentEmotion, dict(sentiwordnet=sentiwordnet)),
    ])

    # Listen on specific port and start server
    application.listen(port)
    tornado.ioloop.IOLoop.instance().start()
<commit_msg>Fixed encoding error.<commit_after>import getopt
import tornado.ioloop
import tornado.web
import sys
from api_resources import Resources
from sentiment_analysis import SentimentAnalysis

import time

class GetWordInformation(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos", ['v', 'n', 'a', 'r'])
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}
            if isinstance(pos, list):
                response['synsets'] = {}
                for postagging in pos:
                    pos_result = sentiwordnet.get_info_word(word, postagging, language)
                    if pos_result:
                        response['synsets'].update(pos_result)
                response["elapsed_time"] = time.time() - start
            else:
                response["synsets"] = sentiwordnet.get_info_word(word, pos, language)
                response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class Synonym(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["sentiment"] = sentiwordnet.get_synonym(word, pos, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetAffect(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language").encode('ascii', errors='ignore')

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["affects"] = sentiwordnet.get_affects(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetDommain(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["dommains"] = sentiwordnet.get_domains(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetSentiment(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text").encode('ascii', errors='ignore')
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["sentiment"] = sentiwordnet.get_sentiment(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetSentimentEmotion(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            sentiment = SentimentAnalysis(text, language, self.sentiwordnet)
            sentiment.analyze()
            response = sentiment.tostring()
            self.write(response)
        else:
            # Create DS to save result
            response = {}

            response.update(sentiwordnet.get_sentiment_and_emotion(text, language))
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class PosTagging(tornado.web.RequestHandler):

    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        text = self.get_argument("text")
        language = self.get_argument("language")

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}

            response["postagging"] = sentiwordnet.get_postagging(text, language)
            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class GetInformation(tornado.web.RequestHandler):
    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos")
        language = self.get_argument("language")

        score = sentiwordnet.get_info_first_word(word, pos, language)

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}
            response["score"] = score

            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

class Translator(tornado.web.RequestHandler):
    def initialize(self, sentiwordnet):
        self.sentiwordnet = sentiwordnet

    def get(self):
        start = time.time()
        # Check input text
        word = self.get_argument("word")
        pos = self.get_argument("pos")
        from_language = self.get_argument("from_language")
        to_language = self.get_argument("to_language")

        translation = sentiwordnet.get_translation(word, pos, from_language, to_language)

        if self.get_argument("format", None) == "onyx":
            self.clear()
            self.set_status(400)
            self.finish("ONYX response not implemented")
        else:
            # Create DS to save result
            response = {}
            response["translation"] = translation

            response["elapsed_time"] = time.time() - start
            # Return result
            self.write(response)

# Check if server por is valid
try:
    opts, args = getopt.getopt(sys.argv[1:], "p:", ["port="])
    assert (len(opts) == 1)
except:
    print 'api_resources_service.py -p <port>'
    sys.exit()

for o, a in opts:
    if o == "-p":
        port = a
    else:
        pass

if __name__ == "__main__":
    # Init treetagger with a specific language
    sentiwordnet = Resources()

    # Init Tornado web server
    application = tornado.web.Application([
        (r"/get_information", GetInformation, dict(sentiwordnet=sentiwordnet)),
        (r"/translate", Translator, dict(sentiwordnet=sentiwordnet)),
        (r"/get_synonym", Synonym, dict(sentiwordnet=sentiwordnet)),
        (r"/get_postagging", PosTagging, dict(sentiwordnet=sentiwordnet)),
        (r"/get_sentiment", GetSentiment, dict(sentiwordnet=sentiwordnet)),
        (r"/get_affect", GetAffect, dict(sentiwordnet=sentiwordnet)),
        (r"/get_dommain", GetDommain, dict(sentiwordnet=sentiwordnet)),
        (r"/get_word_information", GetWordInformation, dict(sentiwordnet=sentiwordnet)),
        (r"/get_sentiment_emotion", GetSentimentEmotion, dict(sentiwordnet=sentiwordnet)),
    ])

    # Listen on specific port and start server
    application.listen(port)
    tornado.ioloop.IOLoop.instance().start()
<|endoftext|>


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

#-----------------------------------------------------------------------------
# Copyright (c) 2013--, bipy development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from StringIO import StringIO

from bipy.util.misc import safe_md5
from bipy.util.unit_test import TestCase, main

class UtilTests(TestCase):
    """Test object for the miscellaneous utility functions"""
    def test_safe_md5(self):
        """Make sure we have the expected md5"""
        exp = 'ab07acbb1e496801937adfa772424bf7'

        fd = StringIO('foo bar baz')
        obs = safe_md5(fd)
        self.assertEqual(obs.hexdigest(),exp)

        fd.close()

if __name__ == '__main__':
    main()
<commit_msg>Remove whitespace<commit_after>#!/usr/bin/env python

#-----------------------------------------------------------------------------
# Copyright (c) 2013--, bipy development team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

from StringIO import StringIO

from bipy.util.misc import safe_md5
from bipy.util.unit_test import TestCase, main

class UtilTests(TestCase):
    """Test object for the miscellaneous utility functions"""
    def test_safe_md5(self):
        """Make sure we have the expected md5"""
        exp = 'ab07acbb1e496801937adfa772424bf7'

        fd = StringIO('foo bar baz')
        obs = safe_md5(fd)
        self.assertEqual(obs.hexdigest(),exp)

        fd.close()

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Copy files to latest/ in the project's default file storage.
"""
import os
import re
import boto3
import random
import logging
from django.conf import settings
from calaccess_raw.models import RawDataVersion
from django.core.management.base import BaseCommand
logger = logging.getLogger(__name__)

class Command(BaseCommand):
    """
    Copy files to latest/ in the project's default file storage.
    """
    help = "Copy files to latest/ in the project's default file storage."

    def handle(self, *args, **options):
        # set up boto session
        self.session = boto3.Session(
            aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
            aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
            region_name=settings.AWS_S3_REGION_NAME
        )

        # and clients
        self.s3 = self.session.client('s3')
        self.cloudfront = self.session.client('cloudfront')

        # Delete existing latest files
        latest_key_list = self.s3.list_objects_v2(
            Bucket=settings.AWS_STORAGE_BUCKET_NAME,
            Prefix='latest/',
        )
        if latest_key_list['KeyCount'] > 0:
            self.delete_keys(latest_key_list)

        # get the version of last update that finished
        v = RawDataVersion.objects.exclude(
            update_finish_datetime=None
        ).latest('update_finish_datetime')

        logger.debug(
            'Copying files for {:%m-%d-%Y %H:%M:%S} version to latest/'.format(
                v.release_datetime
            )
        )

        # save downloaded zip to the latest dir
        if v.download_zip_archive:
            # strip the datetime from the zip name
            zip_name = self.strip_datetime(
                os.path.basename(v.download_zip_archive.name),
            )
            self.copy(
                v.download_zip_archive.name,
                self.get_latest_path(zip_name),
            )
        # save cleaned zip to the latest dir
        if v.clean_zip_archive:
            # strip the datetime from the zip name
            zip_name = self.strip_datetime(
                os.path.basename(v.clean_zip_archive.name),
            )
            self.copy(
                v.clean_zip_archive.name,
                self.get_latest_path(zip_name),
            )

        # loop through all of the raw data files
        for f in v.files.all():
            # save downloaded file to the latest directory
            self.copy(
                f.download_file_archive.name,
                self.get_latest_path(f.download_file_archive.name)
            )

            if f.clean_file_archive:
                # save cleaned file to the latest directory
                self.copy(
                    f.clean_file_archive.name,
                    self.get_latest_path(f.clean_file_archive.name)
                )

            if f.error_log_archive:
                # save error log file to the latest directory
                self.copy(
                    f.error_log_archive.name,
                    self.get_latest_path(f.error_log_archive.name)
                )

        # save processed zip to the latest dir
        if v.processed_version:
            for zf in v.processed_version.zips.all():
                # strip the datetime from the zip name
                zip_name = self.strip_datetime(
                    os.path.basename(zf.zip_archive.name),
                )
                self.copy(
                    zf.zip_archive.name,
                    self.get_latest_path(zip_name),
                )
            # same for processed files
            for f in v.processed_version.files.all():
                # save processed file to the latest directory
                if f.file_archive:
                    # save cleaned file to the latest directory
                    self.copy(
                        f.file_archive.name,
                        self.get_latest_path(f.file_archive.name)
                    )

        # Clear the cloudfront cache by sending an invalidation request
        if latest_key_list['KeyCount'] > 0:
            self.invalidate_keys(latest_key_list)

    def strip_datetime(self, filename):
        """
        Removes the datetime portion from filename.
        """
        return re.sub(
            r'_\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}',
            '',
            filename,
        )

    def get_latest_path(self, old_path):
        """
        Convert the file path to a latest file path
        """
        base_name = os.path.basename(old_path)
        return os.path.join("latest", base_name)

    def copy(self, source, target):
        """
        Copies the provided source key to the provided target key
        """
        logger.debug('Saving copy of %s' % os.path.basename(source))
        copy_source = {
            'Bucket': settings.AWS_STORAGE_BUCKET_NAME,
            'Key': source
        }
        self.s3.copy(
            copy_source,
            settings.AWS_STORAGE_BUCKET_NAME,
            target,
        )

    def delete_keys(self, key_list):
        """
        Delete all the provided s3 keys.
        """
        logger.debug(
            'Deleting %s keys currently under latest/' % key_list['KeyCount']
        )
        # format
        objects = [{'Key': o['Key']} for o in key_list['Contents']]
        # delete
        self.s3.delete_objects(
            Bucket=settings.AWS_STORAGE_BUCKET_NAME,
            Delete={
                'Objects': objects,
            }
        )

    def invalidate_keys(self, key_list):
        """
        Send CloudFront an invalidation request to clear.
        """
        logger.debug(
            "Sending invalidation request for %s keys under latest/" % key_list['KeyCount']
        )
        items = ["/{}".format(o['Key']) for o in key_list['Contents']]
        self.cloudfront.create_invalidation(
            DistributionId=settings.CLOUDFRONT_ARCHIVED_DATA_DISTRIBUTION,
            InvalidationBatch={
                # What to invalidate
                'Paths': {
                    'Quantity': key_list['KeyCount'],
                    'Items': items
                },
                # A random name for the request
                'CallerReference': str(random.getrandbits(128))
            }
        )
<commit_msg>Expand logging in createlatestlinks<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Copy files to latest/ in the project's default file storage.
"""
import os
import re
import boto3
import random
import logging
from django.conf import settings
from calaccess_raw.models import RawDataVersion
from django.core.management.base import BaseCommand
logger = logging.getLogger(__name__)

class Command(BaseCommand):
    """
    Copy files to latest/ in the project's default file storage.
    """
    help = "Copy files to latest/ in the project's default file storage."

    def handle(self, *args, **options):
        # set up boto session
        self.session = boto3.Session(
            aws_access_key_id=settings.AWS_ACCESS_KEY_ID,
            aws_secret_access_key=settings.AWS_SECRET_ACCESS_KEY,
            region_name=settings.AWS_S3_REGION_NAME
        )

        # and clients
        self.s3 = self.session.client('s3')
        self.cloudfront = self.session.client('cloudfront')

        # Delete existing latest files
        latest_key_list = self.s3.list_objects_v2(
            Bucket=settings.AWS_STORAGE_BUCKET_NAME,
            Prefix='latest/',
        )
        if latest_key_list['KeyCount'] > 0:
            self.delete_keys(latest_key_list)

        # get the version of last update that finished
        v = RawDataVersion.objects.exclude(
            update_finish_datetime=None
        ).latest('update_finish_datetime')

        logger.debug(
            'Copying files for {:%m-%d-%Y %H:%M:%S} version to latest/'.format(
                v.release_datetime
            )
        )

        # save downloaded zip to the latest dir
        if v.download_zip_archive:
            # strip the datetime from the zip name
            zip_name = self.strip_datetime(
                os.path.basename(v.download_zip_archive.name),
            )
            self.copy(
                v.download_zip_archive.name,
                self.get_latest_path(zip_name),
            )
        # save cleaned zip to the latest dir
        if v.clean_zip_archive:
            # strip the datetime from the zip name
            zip_name = self.strip_datetime(
                os.path.basename(v.clean_zip_archive.name),
            )
            self.copy(
                v.clean_zip_archive.name,
                self.get_latest_path(zip_name),
            )

        # loop through all of the raw data files
        for f in v.files.all():
            # save downloaded file to the latest directory
            self.copy(
                f.download_file_archive.name,
                self.get_latest_path(f.download_file_archive.name)
            )

            if f.clean_file_archive:
                # save cleaned file to the latest directory
                self.copy(
                    f.clean_file_archive.name,
                    self.get_latest_path(f.clean_file_archive.name)
                )

            if f.error_log_archive:
                # save error log file to the latest directory
                self.copy(
                    f.error_log_archive.name,
                    self.get_latest_path(f.error_log_archive.name)
                )

        # save processed zip to the latest dir
        if v.processed_version:
            for zf in v.processed_version.zips.all():
                # strip the datetime from the zip name
                zip_name = self.strip_datetime(
                    os.path.basename(zf.zip_archive.name),
                )
                self.copy(
                    zf.zip_archive.name,
                    self.get_latest_path(zip_name),
                )
            # same for processed files
            for f in v.processed_version.files.all():
                # save processed file to the latest directory
                if f.file_archive:
                    # save cleaned file to the latest directory
                    self.copy(
                        f.file_archive.name,
                        self.get_latest_path(f.file_archive.name)
                    )

        # Clear the cloudfront cache by sending an invalidation request
        if latest_key_list['KeyCount'] > 0:
            self.invalidate_keys(latest_key_list)

    def strip_datetime(self, filename):
        """
        Removes the datetime portion from filename.
        """
        return re.sub(
            r'_\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2}',
            '',
            filename,
        )

    def get_latest_path(self, old_path):
        """
        Convert the file path to a latest file path
        """
        base_name = os.path.basename(old_path)
        return os.path.join("latest", base_name)

    def copy(self, source, target):
        """
        Copies the provided source key to the provided target key
        """
        logger.debug('Saving copy of {} to {}'.format(os.path.basename(source), target))
        copy_source = {
            'Bucket': settings.AWS_STORAGE_BUCKET_NAME,
            'Key': source
        }
        self.s3.copy(
            copy_source,
            settings.AWS_STORAGE_BUCKET_NAME,
            target,
        )

    def delete_keys(self, key_list):
        """
        Delete all the provided s3 keys.
        """
        logger.debug(
            'Deleting %s keys currently under latest/' % key_list['KeyCount']
        )
        # format
        objects = [{'Key': o['Key']} for o in key_list['Contents']]
        # delete
        self.s3.delete_objects(
            Bucket=settings.AWS_STORAGE_BUCKET_NAME,
            Delete={
                'Objects': objects,
            }
        )

    def invalidate_keys(self, key_list):
        """
        Send CloudFront an invalidation request to clear.
        """
        logger.debug(
            "Sending invalidation request for %s keys under latest/" % key_list['KeyCount']
        )
        items = ["/{}".format(o['Key']) for o in key_list['Contents']]
        self.cloudfront.create_invalidation(
            DistributionId=settings.CLOUDFRONT_ARCHIVED_DATA_DISTRIBUTION,
            InvalidationBatch={
                # What to invalidate
                'Paths': {
                    'Quantity': key_list['KeyCount'],
                    'Items': items
                },
                # A random name for the request
                'CallerReference': str(random.getrandbits(128))
            }
        )
<|endoftext|>