<commit_before>from trivup import trivup
from trivup.apps.KerberosKdcApp import KerberosKdcApp

import contextlib
import os
import socket
import time

class KafkaBrokerApp (trivup.App):
    """ Kafka broker app
        Depends on ZookeeperApp """
    def __init__(self, cluster, conf=None, on=None):
        """
        @param cluster     Current cluster
        @param conf        Configuration dict, see below.
        @param on          Node name to run on

        Supported conf keys:
           * version - Kafka version to use, will build 'trunk' from kafka_path,
                       otherwise the version is taken to be a formal release which
                       will be downloaded and deployed.
           * listeners - CSV list of listener types: PLAINTEXT,SSL,SASL,SASL_SSL
           * listener_host - alternative listener host instead of node name (e.g., '*')
           * advertised_hostname - hostname to use for advertised.listeners (defaults to 'on' node)
           * sasl_mechanisms - CSV list of SASL mechanisms to enable: GSSAPI,PLAIN,SCRAM-SHA-n
                               SASL listeners will be added automatically.
                               KerberosKdcApp is required for GSSAPI.
           * sasl_users - CSV list of SASL PLAIN/SCRAM of user=pass for authenticating clients
           * ssl_client_auth - ssl.client.auth broker property (def: required)
           * num_partitions - Topic auto-create partition count (3)
           * replication_Factor - Topic auto-create replication factor (1)
           * port_base - Low TCP port base to start allocating from (random)
           * kafka_path - Path to Kafka build tree (for trunk usage)
           * conf - arbitary server.properties config as a list of strings.
        """
        super(KafkaBrokerApp, self).__init__(cluster, conf=conf, on=on)

        self.zk = cluster.find_app('ZookeeperApp')
        if self.zk is None:
            raise Exception('ZookeeperApp required')

        # Kafka repo uses SVN nomenclature
        if self.conf['version'] == 'master':
            self.conf['version'] = 'trunk'

        listener_host = self.conf.get('listener_host', self.conf.get('nodename'))

        # Arbitrary (non-template) configuration statements
        conf_blob = self.conf.get('conf', list())
        jaas_blob = list()

        #
        # Configure listeners, SSL and SASL
        #
        listeners = self.conf.get('listeners', 'PLAINTEXT').split(',')

        # SASL support
        sasl_mechs = [x for x in self.conf.get('sasl_mechanisms', '').replace(' ', '').split(',') if len(x) > 0]
        if len(sasl_mechs) > 0:
            listeners.append('SASL_PLAINTEXT')

        # SSL support
        if getattr(cluster, 'ssl', None) is not None:
            # Add SSL listeners
            listeners.append('SSL')
            if len(sasl_mechs) > 0:
                listeners.append('SASL_SSL')

        # Create listeners
        ports = [(x, trivup.TcpPortAllocator(self.cluster).next(self, self.conf.get('port_base', None))) for x in sorted(set(listeners))]
        self.conf['port'] = ports[0][1] # "Default" port
        self.conf['address'] = '%s:%d' % (listener_host, self.conf['port'])
        self.conf['listeners'] = ','.join(['%s://%s:%d' % (x[0], listener_host, x[1]) for x in ports])
        if 'advertised_hostname' not in self.conf:
            self.conf['advertised_hostname'] = self.conf['nodename']
        self.conf['advertised.listeners'] =','.join(['%s://%s:%d' % (x[0], self.conf['advertised_hostname'], x[1]) for x in ports])
        self.conf['auto_create_topics'] = self.conf.get('auto_create_topics', 'true')
        self.dbg('Listeners: %s' % self.conf['listeners'])
        self.dbg('Advertised Listeners: %s' % self.conf['advertised.listeners'])

        if len(sasl_mechs) > 0:
            self.dbg('SASL mechanisms: %s' % sasl_mechs)
            jaas_blob.append('KafkaServer {')

            conf_blob.append('sasl.enabled.mechanisms=%s' % ','.join(sasl_mechs))
            # Handle PLAIN and SCRAM-.. the same way
            for mech in sasl_mechs:
                if mech.find('SCRAM') != -1:
                    plugin = 'scram.Scram'
                elif mech == 'PLAIN':
                    plugin = 'plain.Plain'
                else:
                    continue

                sasl_users = self.conf.get('sasl_users', '')
                if len(sasl_users) == 0:
                    self.log('WARNING: No sasl_users configured for %s, expected CSV of user=pass,..' % plugin)
                else:
                    jaas_blob.append('org.apache.kafka.common.security.%sLoginModule required debug=true' % plugin)
                    for up in sasl_users.split(','):
                        u,p = up.split('=')
                        if plugin is 'plain.Plain':
                            jaas_blob.append('  user_%s="%s"' % (u, p))
                        elif plugin is 'scram.Scram':
                            jaas_blob.append('  username="%s" password="%s"' % (u, p))
                            # SCRAM users are set up in ZK with kafka-configs.sh
                            self.post_start_cmds.append('%s --zookeeper %s --alter --add-config \'%s=[iterations=4096,password=%s]\' --entity-type users --entity-name \'%s\'' % \
                                                        (kafka_configs_sh, self.conf['zk_connect'], mech, p, u))

                    jaas_blob[-1] += ';'

            if 'GSSAPI' in sasl_mechs:
                conf_blob.append('sasl.kerberos.service.name=%s' % 'kafka')
                kdc = self.cluster.find_app(KerberosKdcApp)
                self.env_add('KRB5_CONFIG', kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Djava.security.krb5.conf=%s' % kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Dsun.security.krb5.debug=true')
                self.kerberos_principal,self.kerberos_keytab = kdc.add_principal('kafka', self.conf['advertised_hostname'])
                jaas_blob.append('com.sun.security.auth.module.Krb5LoginModule required')
                jaas_blob.append('useKeyTab=true storeKey=true doNotPrompt=true')
                jaas_blob.append('keyTab="%s"' % self.kerberos_keytab)
                jaas_blob.append('debug=true')
                jaas_blob.append('principal="%s";' % self.kerberos_principal)

            jaas_blob.append('};\n')
            self.conf['jaas_file'] = self.create_file('jaas_broker.conf', data='\n'.join(jaas_blob))
            self.env_add('KAFKA_OPTS', '-Djava.security.auth.login.config=%s' % self.conf['jaas_file'])
            self.env_add('KAFKA_OPTS', '-Djava.security.debug=all')

        # SSL config and keys (et.al.)
        if getattr(cluster, 'ssl', None) is not None:
            ssl = cluster.ssl
            keystore,truststore,_,_ = ssl.create_keystore('broker%s' % self.appid)
            conf_blob.append('ssl.protocol=TLS')
            conf_blob.append('ssl.enabled.protocols=TLSv1.2,TLSv1.1,TLSv1')
            conf_blob.append('ssl.keystore.type = JKS')
            conf_blob.append('ssl.keystore.location = %s' % keystore)
            conf_blob.append('ssl.keystore.password = %s ' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.key.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.truststore.type = JKS')
            conf_blob.append('ssl.truststore.location = %s' % truststore)
            conf_blob.append('ssl.truststore.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.client.auth = %s' % self.conf.get('ssl_client_auth', 'required'))

        # Kafka Configuration properties
        self.conf['log_dirs'] = self.create_dir('logs')
        if 'num_partitions' not in self.conf:
            self.conf['num_partitions'] = 3
        self.conf['zk_connect'] = self.zk.get('address', None)
        if 'replication_factor' not in self.conf:
            self.conf['replication_factor'] = 1

        # Generate config file
        self.conf['conf_file'] = self.create_file_from_template('server.properties',
                                                                self.conf,
                                                                append_data='\n'.join(conf_blob))

        # Generate LOG4J file (if app debug is enabled)
        if self.debug:
            self.conf['log4j_file'] = self.create_file_from_template('log4j.properties', self.conf, subst=False)
            self.env_add('KAFKA_LOG4J_OPTS', '-Dlog4j.configuration=file:%s' % self.conf['log4j_file'])

        # Runs in foreground, stopped by Ctrl-C
        # This is the default for no-deploy use: will be overwritten by deploy() if enabled.
        if self.conf.get('kafka_path', None) is not None:
            self.conf['destdir'] = self.conf['kafka_path']
            self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
            start_sh = os.path.join(self.conf['bindir'], 'kafka-server-start.sh')
        else:
            start_sh = 'kafka-server-start.sh'

        self.conf['start_cmd'] = '%s %s' % (start_sh, self.conf['conf_file'])
        self.conf['stop_cmd'] = None # Ctrl-C

    def operational (self):
        self.dbg('Checking if operational')
        addr, port = self.get('address').split(':')
        with contextlib.closing(socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
            return s.connect_ex((addr, int(port))) == 0

    def deploy (self):
        destdir = os.path.join(self.cluster.mkpath(self.__class__.__name__), 'kafka', self.get('version'))
        self.dbg('Deploy %s version %s on %s to %s' %
                 (self.name, self.get('version'), self.node.name, destdir))
        deploy_exec = self.resource_path('deploy.sh')
        if not os.path.exists(deploy_exec):
            raise NotImplementedError('Kafka deploy.sh script missing in %s' %
                                      deploy_exec)
        t_start = time.time()
        cmd = '%s %s "%s" "%s"' % \
              (deploy_exec, self.get('version'), self.get('kafka_path', destdir), destdir)
        r = os.system(cmd)
        if r != 0:
            raise Exception('Deploy "%s" returned exit code %d' % (cmd, r))
        self.dbg('Deployed version %s in %ds' %
                 (self.get('version'), time.time() - t_start))

        self.conf['destdir'] = destdir
        self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
        # Override start command with updated path.
        self.conf['start_cmd'] = '%s/bin/kafka-server-start.sh %s' % (destdir, self.conf['conf_file'])
        self.dbg('Updated start_cmd to %s' % self.conf['start_cmd'])
<commit_msg>Kafka: add 'num_partitions' and 'replication_factor' config support<commit_after>from trivup import trivup
from trivup.apps.KerberosKdcApp import KerberosKdcApp

import contextlib
import os
import socket
import time

class KafkaBrokerApp (trivup.App):
    """ Kafka broker app
        Depends on ZookeeperApp """
    def __init__(self, cluster, conf=None, on=None):
        """
        @param cluster     Current cluster
        @param conf        Configuration dict, see below.
        @param on          Node name to run on

        Supported conf keys:
           * version - Kafka version to use, will build 'trunk' from kafka_path,
                       otherwise the version is taken to be a formal release which
                       will be downloaded and deployed.
           * listeners - CSV list of listener types: PLAINTEXT,SSL,SASL,SASL_SSL
           * listener_host - alternative listener host instead of node name (e.g., '*')
           * advertised_hostname - hostname to use for advertised.listeners (defaults to 'on' node)
           * sasl_mechanisms - CSV list of SASL mechanisms to enable: GSSAPI,PLAIN,SCRAM-SHA-n
                               SASL listeners will be added automatically.
                               KerberosKdcApp is required for GSSAPI.
           * sasl_users - CSV list of SASL PLAIN/SCRAM of user=pass for authenticating clients
           * ssl_client_auth - ssl.client.auth broker property (def: required)
           * num_partitions - Topic auto-create partition count (3)
           * replication_Factor - Topic auto-create replication factor (1)
           * port_base - Low TCP port base to start allocating from (random)
           * kafka_path - Path to Kafka build tree (for trunk usage)
           * conf - arbitary server.properties config as a list of strings.
        """
        super(KafkaBrokerApp, self).__init__(cluster, conf=conf, on=on)

        self.zk = cluster.find_app('ZookeeperApp')
        if self.zk is None:
            raise Exception('ZookeeperApp required')

        # Kafka repo uses SVN nomenclature
        if self.conf['version'] == 'master':
            self.conf['version'] = 'trunk'

        listener_host = self.conf.get('listener_host', self.conf.get('nodename'))
        # Kafka Configuration properties
        self.conf['log_dirs'] = self.create_dir('logs')
        if 'num_partitions' not in self.conf:
            self.conf['num_partitions'] = 3
        self.conf['zk_connect'] = self.zk.get('address', None)
        if 'replication_factor' not in self.conf:
            self.conf['replication_factor'] = 1

        # Kafka paths
        if self.conf.get('kafka_path', None) is not None:
            self.conf['destdir'] = self.conf['kafka_path']
            self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
            start_sh = os.path.join(self.conf['bindir'], 'kafka-server-start.sh')
            kafka_configs_sh = os.path.join(self.conf['bindir'], 'kafka-configs.sh')
        else:
            start_sh = 'kafka-server-start.sh'
            kafka_configs_sh = 'kafka-configs.sh'

        # Arbitrary (non-template) configuration statements
        conf_blob = self.conf.get('conf', list())
        jaas_blob = list()

        #
        # Configure listeners, SSL and SASL
        #
        listeners = self.conf.get('listeners', 'PLAINTEXT').split(',')

        # SASL support
        sasl_mechs = [x for x in self.conf.get('sasl_mechanisms', '').replace(' ', '').split(',') if len(x) > 0]
        if len(sasl_mechs) > 0:
            listeners.append('SASL_PLAINTEXT')

        # SSL support
        if getattr(cluster, 'ssl', None) is not None:
            # Add SSL listeners
            listeners.append('SSL')
            if len(sasl_mechs) > 0:
                listeners.append('SASL_SSL')

        # Create listeners
        ports = [(x, trivup.TcpPortAllocator(self.cluster).next(self, self.conf.get('port_base', None))) for x in sorted(set(listeners))]
        self.conf['port'] = ports[0][1] # "Default" port
        self.conf['address'] = '%s:%d' % (listener_host, self.conf['port'])
        self.conf['listeners'] = ','.join(['%s://%s:%d' % (x[0], listener_host, x[1]) for x in ports])
        if 'advertised_hostname' not in self.conf:
            self.conf['advertised_hostname'] = self.conf['nodename']
        self.conf['advertised.listeners'] =','.join(['%s://%s:%d' % (x[0], self.conf['advertised_hostname'], x[1]) for x in ports])
        self.conf['auto_create_topics'] = self.conf.get('auto_create_topics', 'true')
        self.dbg('Listeners: %s' % self.conf['listeners'])
        self.dbg('Advertised Listeners: %s' % self.conf['advertised.listeners'])

        if len(sasl_mechs) > 0:
            self.dbg('SASL mechanisms: %s' % sasl_mechs)
            jaas_blob.append('KafkaServer {')

            conf_blob.append('sasl.enabled.mechanisms=%s' % ','.join(sasl_mechs))
            # Handle PLAIN and SCRAM-.. the same way
            for mech in sasl_mechs:
                if mech.find('SCRAM') != -1:
                    plugin = 'scram.Scram'
                elif mech == 'PLAIN':
                    plugin = 'plain.Plain'
                else:
                    continue

                sasl_users = self.conf.get('sasl_users', '')
                if len(sasl_users) == 0:
                    self.log('WARNING: No sasl_users configured for %s, expected CSV of user=pass,..' % plugin)
                else:
                    jaas_blob.append('org.apache.kafka.common.security.%sLoginModule required debug=true' % plugin)
                    for up in sasl_users.split(','):
                        u,p = up.split('=')
                        if plugin is 'plain.Plain':
                            jaas_blob.append('  user_%s="%s"' % (u, p))
                        elif plugin is 'scram.Scram':
                            jaas_blob.append('  username="%s" password="%s"' % (u, p))
                            # SCRAM users are set up in ZK with kafka-configs.sh
                            self.post_start_cmds.append('%s --zookeeper %s --alter --add-config \'%s=[iterations=4096,password=%s]\' --entity-type users --entity-name \'%s\'' % \
                                                        (kafka_configs_sh, self.conf['zk_connect'], mech, p, u))

                    jaas_blob[-1] += ';'

            if 'GSSAPI' in sasl_mechs:
                conf_blob.append('sasl.kerberos.service.name=%s' % 'kafka')
                kdc = self.cluster.find_app(KerberosKdcApp)
                self.env_add('KRB5_CONFIG', kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Djava.security.krb5.conf=%s' % kdc.conf['krb5_conf'])
                self.env_add('KAFKA_OPTS', '-Dsun.security.krb5.debug=true')
                self.kerberos_principal,self.kerberos_keytab = kdc.add_principal('kafka', self.conf['advertised_hostname'])
                jaas_blob.append('com.sun.security.auth.module.Krb5LoginModule required')
                jaas_blob.append('useKeyTab=true storeKey=true doNotPrompt=true')
                jaas_blob.append('keyTab="%s"' % self.kerberos_keytab)
                jaas_blob.append('debug=true')
                jaas_blob.append('principal="%s";' % self.kerberos_principal)

            jaas_blob.append('};\n')
            self.conf['jaas_file'] = self.create_file('jaas_broker.conf', data='\n'.join(jaas_blob))
            self.env_add('KAFKA_OPTS', '-Djava.security.auth.login.config=%s' % self.conf['jaas_file'])
            self.env_add('KAFKA_OPTS', '-Djava.security.debug=all')

        # SSL config and keys (et.al.)
        if getattr(cluster, 'ssl', None) is not None:
            ssl = cluster.ssl
            keystore,truststore,_,_ = ssl.create_keystore('broker%s' % self.appid)
            conf_blob.append('ssl.protocol=TLS')
            conf_blob.append('ssl.enabled.protocols=TLSv1.2,TLSv1.1,TLSv1')
            conf_blob.append('ssl.keystore.type = JKS')
            conf_blob.append('ssl.keystore.location = %s' % keystore)
            conf_blob.append('ssl.keystore.password = %s ' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.key.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.truststore.type = JKS')
            conf_blob.append('ssl.truststore.location = %s' % truststore)
            conf_blob.append('ssl.truststore.password = %s' % ssl.conf.get('ssl_key_pass'))
            conf_blob.append('ssl.client.auth = %s' % self.conf.get('ssl_client_auth', 'required'))

        # Generate config file
        self.conf['conf_file'] = self.create_file_from_template('server.properties',
                                                                self.conf,
                                                                append_data='\n'.join(conf_blob))

        # Generate LOG4J file (if app debug is enabled)
        if self.debug:
            self.conf['log4j_file'] = self.create_file_from_template('log4j.properties', self.conf, subst=False)
            self.env_add('KAFKA_LOG4J_OPTS', '-Dlog4j.configuration=file:%s' % self.conf['log4j_file'])

        # Runs in foreground, stopped by Ctrl-C
        # This is the default for no-deploy use: will be overwritten by deploy() if enabled.

        self.conf['start_cmd'] = '%s %s' % (start_sh, self.conf['conf_file'])
        self.conf['stop_cmd'] = None # Ctrl-C

    def operational (self):
        self.dbg('Checking if operational')
        addr, port = self.get('address').split(':')
        with contextlib.closing(socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
            return s.connect_ex((addr, int(port))) == 0

    def deploy (self):
        destdir = os.path.join(self.cluster.mkpath(self.__class__.__name__), 'kafka', self.get('version'))
        self.dbg('Deploy %s version %s on %s to %s' %
                 (self.name, self.get('version'), self.node.name, destdir))
        deploy_exec = self.resource_path('deploy.sh')
        if not os.path.exists(deploy_exec):
            raise NotImplementedError('Kafka deploy.sh script missing in %s' %
                                      deploy_exec)
        t_start = time.time()
        cmd = '%s %s "%s" "%s"' % \
              (deploy_exec, self.get('version'), self.get('kafka_path', destdir), destdir)
        r = os.system(cmd)
        if r != 0:
            raise Exception('Deploy "%s" returned exit code %d' % (cmd, r))
        self.dbg('Deployed version %s in %ds' %
                 (self.get('version'), time.time() - t_start))

        self.conf['destdir'] = destdir
        self.conf['bindir'] = os.path.join(self.conf['destdir'], 'bin')
        # Override start command with updated path.
        self.conf['start_cmd'] = '%s/bin/kafka-server-start.sh %s' % (destdir, self.conf['conf_file'])
        self.dbg('Updated start_cmd to %s' % self.conf['start_cmd'])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Tue Jan 10 00:15:50 2017

@author: sakurai
"""

import numpy as np
from sklearn.preprocessing import LabelEncoder
from fuel.streams import DataStream
from fuel.schemes import IterationScheme, BatchSizeScheme, SequentialScheme

from cars196_dataset import Cars196Dataset
from cub200_2011_dataset import Cub200_2011Dataset
from online_products_dataset import OnlineProductsDataset
from random_fixed_size_crop_mod import RandomFixedSizeCrop

def get_streams(batch_size=50, dataset='cars196', method='n_pairs_mc',
                load_in_memory=False):
    '''
    args:
       batch_size (int):
           number of examples per batch
       dataset (str):
           specify the dataset from 'cars196', 'cub200_2011', 'products'.
       method (str):
           batch construction method. Select from 'n_pairs_mc', 'clustering'.
    '''

    if dataset == 'cars196':
        dataset_class = Cars196Dataset
    elif dataset == 'cub200_2011':
        dataset_class = Cub200_2011Dataset
    elif dataset == 'products':
        dataset_class = OnlineProductsDataset
    else:
        raise ValueError(
            "`dataset` must be 'cars196', 'cub200_2011' or 'products'.")

    dataset_train = dataset_class(['train'], load_in_memory=load_in_memory)
    dataset_test = dataset_class(['test'], load_in_memory=load_in_memory)

    if method == 'n_pairs_mc':
        labels = dataset_class(
            ['train'], sources=['targets'], load_in_memory=True).data_sources
        scheme = NPairLossScheme(labels, batch_size)
    elif method == 'clustering':
        scheme = EpochwiseShuffledInfiniteScheme(
            dataset_train.num_examples, batch_size)
    elif issubclass(method, IterationScheme):
        scheme = method(batch_size, dataset=dataset_train)
    else:
        raise ValueError("`method` must be 'n_pairs_mc' or 'clustering' "
                         "or subclass of IterationScheme.")
    stream = DataStream(dataset_train, iteration_scheme=scheme)
    stream_train = RandomFixedSizeCrop(stream, which_sources=('images',),
                                       random_lr_flip=True)

    stream_train_eval = RandomFixedSizeCrop(DataStream(
        dataset_train, iteration_scheme=SequentialScheme(
            dataset_train.num_examples, batch_size)),
        which_sources=('images',), random_lr_flip=True)
    stream_test = RandomFixedSizeCrop(DataStream(
        dataset_test, iteration_scheme=SequentialScheme(
            dataset_test.num_examples, batch_size)),
        which_sources=('images',), random_lr_flip=True)

    return stream_train, stream_train_eval, stream_test

class NPairLossScheme(BatchSizeScheme):
    def __init__(self, labels, batch_size):
        self._labels = np.array(labels).flatten()
        self._label_encoder = LabelEncoder().fit(self._labels)
        self._classes = self._label_encoder.classes_
        self.num_classes = len(self._classes)
        assert batch_size % 2 == 0, ("batch_size must be even number.")
        assert batch_size <= self.num_classes * 2, (
               "batch_size must not exceed twice the number of classes"
               "(i.e. set batch_size <= {}).".format(self.num_classes * 2))
        self.batch_size = batch_size

        self._class_to_indexes = []
        for c in self._classes:
            self._class_to_indexes.append(
                np.argwhere(self._labels == c).ravel())

    def __iter__(self):
        return self

    def next(self):
        anchor_indexes, positive_indexes = self._generate_indexes()
        indexes = anchor_indexes + positive_indexes
        return indexes

    def _generate_indexes(self):
        random_classes = np.random.choice(
            self.num_classes, self.batch_size / 2, False)
        anchor_indexes = []
        positive_indexes = []
        for c in random_classes:
            a, p = np.random.choice(self._class_to_indexes[c], 2, False)
            anchor_indexes.append(a)
            positive_indexes.append(p)
        return anchor_indexes, positive_indexes

    def get_request_iterator(self):
        return self

class EpochwiseShuffledInfiniteScheme(BatchSizeScheme):
    def __init__(self, indexes, batch_size):
        if isinstance(indexes, int):
            indexes = range(indexes)
        if batch_size > len(indexes):
            raise ValueError('batch_size must not be larger than the indexes.')
        if len(indexes) != len(np.unique(indexes)):
            raise ValueError('Items in indexes must be unique.')
        self._original_indexes = np.array(indexes, dtype=np.int).flatten()
        self.batch_size = batch_size
        self._shuffled_indexes = np.array([], dtype=np.int)

    def __iter__(self):
        return self

    def next(self):
        batch_size = self.batch_size

        # if remaining indexes are shorter than batch_size then new shuffled
        # indexes are appended to the remains.
        num_remains = len(self._shuffled_indexes)
        if num_remains < batch_size:
            num_overrun = batch_size - num_remains
            new_shuffled_indexes = self._original_indexes.copy()

            # ensure the batch of indexes from the joint part does not contain
            # duplicate index.
            np.random.shuffle(new_shuffled_indexes)
            overrun = new_shuffled_indexes[:num_overrun]
            next_indexes = np.concatenate((self._shuffled_indexes, overrun))
            while len(next_indexes) != len(np.unique(next_indexes)):
                np.random.shuffle(new_shuffled_indexes)
                overrun = new_shuffled_indexes[:num_overrun]
                next_indexes = np.concatenate(
                    (self._shuffled_indexes, overrun))
            self._shuffled_indexes = np.concatenate(
                (self._shuffled_indexes, new_shuffled_indexes))
        next_indexes = self._shuffled_indexes[:batch_size]
        self._shuffled_indexes = self._shuffled_indexes[batch_size:]
        return next_indexes

    def get_request_iterator(self):
        return self
<commit_msg>Fix EpochwiseShuffledInfiniteScheme to iterate lists instead of ndarrays<commit_after># -*- coding: utf-8 -*-
"""
Created on Tue Jan 10 00:15:50 2017

@author: sakurai
"""

import collections
import numpy as np
from sklearn.preprocessing import LabelEncoder
from fuel.streams import DataStream
from fuel.schemes import IterationScheme, BatchSizeScheme, SequentialScheme

from cars196_dataset import Cars196Dataset
from cub200_2011_dataset import Cub200_2011Dataset
from online_products_dataset import OnlineProductsDataset
from random_fixed_size_crop_mod import RandomFixedSizeCrop

def get_streams(batch_size=50, dataset='cars196', method='n_pairs_mc',
                load_in_memory=False):
    '''
    args:
       batch_size (int):
           number of examples per batch
       dataset (str):
           specify the dataset from 'cars196', 'cub200_2011', 'products'.
       method (str):
           batch construction method. Select from 'n_pairs_mc', 'clustering'.
    '''

    if dataset == 'cars196':
        dataset_class = Cars196Dataset
    elif dataset == 'cub200_2011':
        dataset_class = Cub200_2011Dataset
    elif dataset == 'products':
        dataset_class = OnlineProductsDataset
    else:
        raise ValueError(
            "`dataset` must be 'cars196', 'cub200_2011' or 'products'.")

    dataset_train = dataset_class(['train'], load_in_memory=load_in_memory)
    dataset_test = dataset_class(['test'], load_in_memory=load_in_memory)

    if method == 'n_pairs_mc':
        labels = dataset_class(
            ['train'], sources=['targets'], load_in_memory=True).data_sources
        scheme = NPairLossScheme(labels, batch_size)
    elif method == 'clustering':
        scheme = EpochwiseShuffledInfiniteScheme(
            dataset_train.num_examples, batch_size)
    elif issubclass(method, IterationScheme):
        scheme = method(batch_size, dataset=dataset_train)
    else:
        raise ValueError("`method` must be 'n_pairs_mc' or 'clustering' "
                         "or subclass of IterationScheme.")
    stream = DataStream(dataset_train, iteration_scheme=scheme)
    stream_train = RandomFixedSizeCrop(stream, which_sources=('images',),
                                       random_lr_flip=True)

    stream_train_eval = RandomFixedSizeCrop(DataStream(
        dataset_train, iteration_scheme=SequentialScheme(
            dataset_train.num_examples, batch_size)),
        which_sources=('images',), random_lr_flip=True)
    stream_test = RandomFixedSizeCrop(DataStream(
        dataset_test, iteration_scheme=SequentialScheme(
            dataset_test.num_examples, batch_size)),
        which_sources=('images',), random_lr_flip=True)

    return stream_train, stream_train_eval, stream_test

class NPairLossScheme(BatchSizeScheme):
    def __init__(self, labels, batch_size):
        self._labels = np.array(labels).flatten()
        self._label_encoder = LabelEncoder().fit(self._labels)
        self._classes = self._label_encoder.classes_
        self.num_classes = len(self._classes)
        assert batch_size % 2 == 0, ("batch_size must be even number.")
        assert batch_size <= self.num_classes * 2, (
               "batch_size must not exceed twice the number of classes"
               "(i.e. set batch_size <= {}).".format(self.num_classes * 2))
        self.batch_size = batch_size

        self._class_to_indexes = []
        for c in self._classes:
            self._class_to_indexes.append(
                np.argwhere(self._labels == c).ravel())

    def __iter__(self):
        return self

    def next(self):
        anchor_indexes, positive_indexes = self._generate_indexes()
        indexes = anchor_indexes + positive_indexes
        return indexes

    def _generate_indexes(self):
        random_classes = np.random.choice(
            self.num_classes, self.batch_size / 2, False)
        anchor_indexes = []
        positive_indexes = []
        for c in random_classes:
            a, p = np.random.choice(self._class_to_indexes[c], 2, False)
            anchor_indexes.append(a)
            positive_indexes.append(p)
        return anchor_indexes, positive_indexes

    def get_request_iterator(self):
        return self

class EpochwiseShuffledInfiniteScheme(BatchSizeScheme):
    def __init__(self, indexes, batch_size):
        if not isinstance(indexes, collections.Iterable):
            indexes = range(indexes)
        if batch_size > len(indexes):
            raise ValueError('batch_size must not be larger than the indexes.')
        if len(indexes) != len(np.unique(indexes)):
            raise ValueError('Items in indexes must be unique.')
        self._original_indexes = np.array(indexes, dtype=np.int).flatten()
        self.batch_size = batch_size
        self._shuffled_indexes = np.array([], dtype=np.int)

    def __iter__(self):
        return self

    def next(self):
        batch_size = self.batch_size

        # if remaining indexes are shorter than batch_size then new shuffled
        # indexes are appended to the remains.
        num_remains = len(self._shuffled_indexes)
        if num_remains < batch_size:
            num_overrun = batch_size - num_remains
            new_shuffled_indexes = self._original_indexes.copy()

            # ensure the batch of indexes from the joint part does not contain
            # duplicate index.
            np.random.shuffle(new_shuffled_indexes)
            overrun = new_shuffled_indexes[:num_overrun]
            next_indexes = np.concatenate((self._shuffled_indexes, overrun))
            while len(next_indexes) != len(np.unique(next_indexes)):
                np.random.shuffle(new_shuffled_indexes)
                overrun = new_shuffled_indexes[:num_overrun]
                next_indexes = np.concatenate(
                    (self._shuffled_indexes, overrun))
            self._shuffled_indexes = np.concatenate(
                (self._shuffled_indexes, new_shuffled_indexes))
        next_indexes = self._shuffled_indexes[:batch_size]
        self._shuffled_indexes = self._shuffled_indexes[batch_size:]
        return next_indexes.tolist()

    def get_request_iterator(self):
        return self
<|endoftext|>


<commit_before>#!/usr/bin/python
# vim: set fileencoding=UTF-8 :
"""
callsp

Call a stored procedure in database.

Depends on dboperator which takes care of connection and other db stuff.
"""
import sys,os,getopt
from time import localtime, strftime

import dboperator_debug as dboperator

def show(message):
  print(strftime("%Y-%m-%d %H:%M:%S", localtime())+" "+message)

def load(database,schema,procedure,verbose=False):
  show("begin with "+database+" "+schema+" "+procedure)

  sql = "execute "+database+"."+schema+"."+procedure
  try:
    dboperator.execute(sql)
except sqlerror as e from dboperator:
    show("Executing procedure failed.")
    print(e)
    dboperator.close()
    exit(2) # lopeta virheeseen

  dboperator.close()

  show("ready")

def usage():
  print("""
usage: callsp.py [-e|--schema <schema>] -p|--procedure <procedure> [-v|--verbose]

schema defaults to $SCHEMA then to "" (for database default if set)
procedure is mandatory argument. Name of the procedure to execute.
""")

def main(argv):
  # variables that are given as arguments with possible default values
  database = os.getenv("DATABASE_NAME")
  schema = os.getenv("SCHEMA") or ""
  procedure = ""
  verbose = False

  try:
    opts, args = getopt.getopt(argv,"d:e:p:v",["database=","schema=","procedure=","verbose"])
  except getopt.GetoptError as err:
    print(err)
    usage()
    sys.exit(2)
  for opt, arg in opts:
    if opt in ("-d", "--database"): database = arg
    elif opt in ("-e", "--schema"): schema = arg
    elif opt in ("-p", "--procedure"): procedure = arg
    elif opt in ("-v", "--verbose"): verbose = True
  if not schema or not procedure:
    usage()
    sys.exit(2)

  load(database,schema,procedure,verbose)

if __name__ == "__main__":
  main(sys.argv[1:])
<commit_msg>CSCVIPUNEN-2386 tuning dboperator<commit_after>#!/usr/bin/python
# vim: set fileencoding=UTF-8 :
"""
callsp

Call a stored procedure in database.

Depends on dboperator which takes care of connection and other db stuff.
"""
import sys,os,getopt
from time import localtime, strftime

import dboperator_debug as dboperator

def show(message):
  print(strftime("%Y-%m-%d %H:%M:%S", localtime())+" "+message)

def load(database,schema,procedure,verbose=False):
  show("begin with "+database+" "+schema+" "+procedure)

  sql = "execute "+database+"."+schema+"."+procedure
  try:
    dboperator.execute(sql)
  except sqlerror as e from dboperator:
    show("Executing procedure failed.")
    print(e)
    dboperator.close()
    exit(2) # lopeta virheeseen

  dboperator.close()

  show("ready")

def usage():
  print("""
usage: callsp.py [-e|--schema <schema>] -p|--procedure <procedure> [-v|--verbose]

schema defaults to $SCHEMA then to "" (for database default if set)
procedure is mandatory argument. Name of the procedure to execute.
""")

def main(argv):
  # variables that are given as arguments with possible default values
  database = os.getenv("DATABASE_NAME")
  schema = os.getenv("SCHEMA") or ""
  procedure = ""
  verbose = False

  try:
    opts, args = getopt.getopt(argv,"d:e:p:v",["database=","schema=","procedure=","verbose"])
  except getopt.GetoptError as err:
    print(err)
    usage()
    sys.exit(2)
  for opt, arg in opts:
    if opt in ("-d", "--database"): database = arg
    elif opt in ("-e", "--schema"): schema = arg
    elif opt in ("-p", "--procedure"): procedure = arg
    elif opt in ("-v", "--verbose"): verbose = True
  if not schema or not procedure:
    usage()
    sys.exit(2)

  load(database,schema,procedure,verbose)

if __name__ == "__main__":
  main(sys.argv[1:])
<|endoftext|>


<commit_before>import os
from unittest import mock

from tornado.testing import AsyncTestCase, gen_test

from remoteappmanager.docker.container import Container
from remoteappmanager.docker.docker_labels import SIMPHONY_NS_RUNINFO
from remoteappmanager.docker.container_manager import ContainerManager, \
    OperationInProgress
from remoteappmanager.docker.image import Image
from remoteappmanager.tests import utils
from remoteappmanager.tests.mocking.virtual.docker_client import (
    create_docker_client)

class TestContainerManager(AsyncTestCase):
    def setUp(self):
        super().setUp()
        self.manager = ContainerManager(docker_config={},
                                        realm="myrealm")
        self.mock_docker_client = create_docker_client()
        self.manager._docker_client._sync_client = self.mock_docker_client

    def test_instantiation(self):
        self.assertIsNotNone(self.manager._docker_client)

    @gen_test
    def test_start_stop(self):
        mock_client = self.mock_docker_client

        result = yield self.manager.start_container(
            "username",
            'image_id1',
            "new_mapping_id",
            "/base/url",
            None,
            None)
        self.assertTrue(mock_client.start.called)
        self.assertTrue(mock_client.create_container.called)

        runinfo_labels = mock_client.create_container.call_args[1]["labels"]

        self.assertEqual(runinfo_labels[SIMPHONY_NS_RUNINFO.user], "username")
        self.assertEqual(runinfo_labels[SIMPHONY_NS_RUNINFO.realm], "myrealm")
        self.assertIn(SIMPHONY_NS_RUNINFO.url_id, runinfo_labels)
        self.assertEqual(runinfo_labels[SIMPHONY_NS_RUNINFO.mapping_id],
                         "new_mapping_id")

        self.assertIsInstance(result, Container)
        self.assertFalse(mock_client.stop.called)
        self.assertFalse(mock_client.remove_container.called)

        yield self.manager.stop_and_remove_container(result.docker_id)

        self.assertTrue(mock_client.stop.called)
        self.assertTrue(mock_client.remove_container.called)

    @gen_test
    def test_containers_from_mapping_id(self):
        """ Test containers_for_mapping_id returns a list of Container """
        result = yield self.manager.containers_from_mapping_id("user_name",
                                                               "mapping_id")
        expected = Container(docker_id='container_id1',
                             mapping_id="mapping_id",
                             name='/myrealm-username-mapping_5Fid',
                             image_name='image_name1',  # noqa
                             user="user_name",
                             image_id='image_id1',
                             ip='127.0.0.1',
                             port=666,
                             url_id='url_id',
                             realm="myrealm",
                             urlpath="/user/username/containers/url_id")

        self.assertEqual(len(result), 1)
        utils.assert_containers_equal(self, result[0], expected)

    @gen_test
    def test_containers_from_url_id(self):
        """ Test containers_for_mapping_id returns a list of Container """
        result = yield self.manager.container_from_url_id("url_id")
        expected = Container(docker_id='container_id1',
                             mapping_id="mapping_id",
                             name='/myrealm-username-mapping_5Fid',
                             image_name='image_name1',  # noqa
                             user="user_name",
                             image_id='image_id1',
                             ip='127.0.0.1',
                             port=666,
                             url_id='url_id',
                             realm="myrealm",
                             urlpath="/user/username/containers/url_id",
                             )

        utils.assert_containers_equal(self, result, expected)

    @gen_test
    def test_containers_from_url_id_exceptions(self):
        """ Test containers_for_mapping_id returns a list of Container """
        docker_client = self.mock_docker_client
        docker_client.port = mock.Mock(side_effect=Exception("Boom!"))

        result = yield self.manager.container_from_url_id("url_id")
        self.assertEqual(result, None)

        # Making it so that no valid dictionary is returned.
        docker_client.port = mock.Mock(return_value=1234)
        self.mock_docker_client = docker_client
        self.manager._docker_client._sync_client = docker_client

        result = yield self.manager.container_from_url_id("url_id")
        self.assertEqual(result, None)

    @gen_test
    def test_running_containers(self):
        result = yield self.manager.running_containers()
        self.assertEqual(len(result), 1)

    @gen_test
    def test_race_condition_spawning(self):
        # Start the operations, and retrieve the future.
        # they will stop at the first yield and not go further until
        # we yield them
        f1 = self.manager.start_container("username",
                                          "image_id1",
                                          "mapping_id",
                                          "/foo/bar",
                                          None,
                                          )

        f2 = self.manager.start_container("username",
                                          "image_id1",
                                          "mapping_id",
                                          "/foo/baz",
                                          None,
                                          )

        # If these yielding raise a KeyError, it is because the second
        # one tries to remove the same key from the list, but it has been
        # already removed by the first one. Race condition.
        yield f1

        with self.assertRaises(OperationInProgress):
            yield f2

        self.assertEqual(self.mock_docker_client.start.call_count, 1)

    @gen_test
    def test_race_condition_stopping(self):
        f1 = self.manager.stop_and_remove_container("container_id1")
        f2 = self.manager.stop_and_remove_container("container_id1")

        yield f1

        with self.assertRaises(OperationInProgress):
            yield f2

        self.assertEqual(self.mock_docker_client.stop.call_count, 1)

    @gen_test
    def test_start_already_present_container(self):
        mock_client = self.mock_docker_client

        result = yield self.manager.start_container(
            "username",
            "image_name1",
            "mapping_id",
            "/base/url",
            None,
            None)
        self.assertTrue(mock_client.start.called)
        self.assertIsInstance(result, Container)

        # Stop should have been called and the container removed
        self.assertTrue(mock_client.stop.called)
        self.assertTrue(mock_client.remove_container.called)

    @gen_test
    def test_image(self):
        image = yield self.manager.image('image_name1')
        self.assertIsInstance(image, Image)
        self.assertEqual(image.description,
                         'Ubuntu machine with mayavi preinstalled')
        self.assertEqual(image.icon_128, "")
        self.assertEqual(image.ui_name, "Mayavi 4.4.4")
        self.assertEqual(image.docker_id, 'image_id1')

        image = yield self.manager.image("whatev")
        self.assertIsNone(image)

    @gen_test
    def test_start_container_with_nonexisting_volume_source(self):
        # These volume sources are invalid
        volumes = {'~no_way_this_be_valid': {'bind': '/target_vol1',
                                             'mode': 'ro'},
                   '/no_way_this_be_valid': {'bind': '/target_vol2',
                                             'mode': 'ro'}}

        # This volume source is valid
        good_path = os.path.abspath('.')
        volumes[good_path] = {'bind': '/target_vol3',
                              'mode': 'ro'}

        yield self.manager.start_container("username",
                                           "image_id1",
                                           "mapping_id",
                                           "/foo/bar",
                                           volumes,
                                           )

        # Call args and keyword args that create_container receives
        docker_client = self.manager._docker_client
        args = docker_client._sync_client.create_container.call_args
        actual_volume_targets = args[1]['volumes']

        # Invalid volume paths should have been filtered away
        self.assertNotIn('/target_vol1', actual_volume_targets)
        self.assertNotIn('/target_vol2', actual_volume_targets)

        # The current directory is valid, should stay
        self.assertIn('/target_vol3', actual_volume_targets)

    @gen_test
    def test_start_container_exception_cleanup(self):
        self.assertFalse(self.mock_docker_client.stop.called)
        self.assertFalse(self.mock_docker_client.remove_container.called)

        def raiser(*args, **kwargs):
            raise Exception("Boom!")

        self.manager._docker_client.start = mock.Mock(side_effect=raiser)

        with self.assertRaisesRegex(Exception, 'Boom!'):
            yield self.manager.start_container("username",
                                               "image_id1",
                                               "mapping_id",
                                               "/base/url",
                                               None,
                                               None)

        self.assertTrue(self.mock_docker_client.stop.called)
        self.assertTrue(self.mock_docker_client.remove_container.called)

    @gen_test
    def test_start_container_exception_cleanup_2(self):
        # Same test as above, but checks after the start (at ip and port)
        self.assertFalse(self.mock_docker_client.stop.called)
        self.assertFalse(self.mock_docker_client.remove_container.called)

        def raiser(*args, **kwargs):
            raise Exception("Boom!")

        self.manager._docker_client.port = mock.Mock(side_effect=raiser)

        with self.assertRaisesRegex(Exception, 'Boom!'):
            yield self.manager.start_container("username",
                                               "image_id1",
                                               "mapping_id",
                                               "/base/url",
                                               None,
                                               None)

        self.assertTrue(self.mock_docker_client.stop.called)
        self.assertTrue(self.mock_docker_client.remove_container.called)

    @gen_test
    def test_start_container_with_environment(self):
        mock_client = self.mock_docker_client

        environment = {
            "FOO": "bar"
        }

        yield self.manager.start_container(
            "username",
            "image_name1",
            "mapping_id",
            "/base/url",
            None,
            environment)

        self.assertEqual(
            mock_client.create_container.call_args[1]["environment"]["FOO"],
            "bar")

    @gen_test
    def test_different_realm(self):
        manager = ContainerManager(docker_config={},
                                   realm="anotherrealm")
        manager._docker_client._sync_client = create_docker_client()

        result = yield manager.containers_from_mapping_id("user_name",
                                                          "mapping_id")
        self.assertEqual(result, [])
<commit_msg>Coverage for not stopping when realm is different<commit_after>import os
from unittest import mock

from tornado.testing import AsyncTestCase, gen_test

from remoteappmanager.docker.container import Container
from remoteappmanager.docker.docker_labels import SIMPHONY_NS_RUNINFO
from remoteappmanager.docker.container_manager import ContainerManager, \
    OperationInProgress
from remoteappmanager.docker.image import Image
from remoteappmanager.tests import utils
from remoteappmanager.tests.mocking.virtual.docker_client import (
    create_docker_client)

class TestContainerManager(AsyncTestCase):
    def setUp(self):
        super().setUp()
        self.manager = ContainerManager(docker_config={},
                                        realm="myrealm")
        self.mock_docker_client = create_docker_client()
        self.manager._docker_client._sync_client = self.mock_docker_client

    def test_instantiation(self):
        self.assertIsNotNone(self.manager._docker_client)

    @gen_test
    def test_start_stop(self):
        mock_client = self.mock_docker_client

        result = yield self.manager.start_container(
            "username",
            'image_id1',
            "new_mapping_id",
            "/base/url",
            None,
            None)
        self.assertTrue(mock_client.start.called)
        self.assertTrue(mock_client.create_container.called)

        runinfo_labels = mock_client.create_container.call_args[1]["labels"]

        self.assertEqual(runinfo_labels[SIMPHONY_NS_RUNINFO.user], "username")
        self.assertEqual(runinfo_labels[SIMPHONY_NS_RUNINFO.realm], "myrealm")
        self.assertIn(SIMPHONY_NS_RUNINFO.url_id, runinfo_labels)
        self.assertEqual(runinfo_labels[SIMPHONY_NS_RUNINFO.mapping_id],
                         "new_mapping_id")

        self.assertIsInstance(result, Container)
        self.assertFalse(mock_client.stop.called)
        self.assertFalse(mock_client.remove_container.called)

        yield self.manager.stop_and_remove_container(result.docker_id)

        self.assertTrue(mock_client.stop.called)
        self.assertTrue(mock_client.remove_container.called)

    @gen_test
    def test_containers_from_mapping_id(self):
        """ Test containers_for_mapping_id returns a list of Container """
        result = yield self.manager.containers_from_mapping_id("user_name",
                                                               "mapping_id")
        expected = Container(docker_id='container_id1',
                             mapping_id="mapping_id",
                             name='/myrealm-username-mapping_5Fid',
                             image_name='image_name1',  # noqa
                             user="user_name",
                             image_id='image_id1',
                             ip='127.0.0.1',
                             port=666,
                             url_id='url_id',
                             realm="myrealm",
                             urlpath="/user/username/containers/url_id")

        self.assertEqual(len(result), 1)
        utils.assert_containers_equal(self, result[0], expected)

    @gen_test
    def test_containers_from_url_id(self):
        """ Test containers_for_mapping_id returns a list of Container """
        result = yield self.manager.container_from_url_id("url_id")
        expected = Container(docker_id='container_id1',
                             mapping_id="mapping_id",
                             name='/myrealm-username-mapping_5Fid',
                             image_name='image_name1',  # noqa
                             user="user_name",
                             image_id='image_id1',
                             ip='127.0.0.1',
                             port=666,
                             url_id='url_id',
                             realm="myrealm",
                             urlpath="/user/username/containers/url_id",
                             )

        utils.assert_containers_equal(self, result, expected)

    @gen_test
    def test_containers_from_url_id_exceptions(self):
        """ Test containers_for_mapping_id returns a list of Container """
        docker_client = self.mock_docker_client
        docker_client.port = mock.Mock(side_effect=Exception("Boom!"))

        result = yield self.manager.container_from_url_id("url_id")
        self.assertEqual(result, None)

        # Making it so that no valid dictionary is returned.
        docker_client.port = mock.Mock(return_value=1234)
        self.mock_docker_client = docker_client
        self.manager._docker_client._sync_client = docker_client

        result = yield self.manager.container_from_url_id("url_id")
        self.assertEqual(result, None)

    @gen_test
    def test_running_containers(self):
        result = yield self.manager.running_containers()
        self.assertEqual(len(result), 1)

    @gen_test
    def test_race_condition_spawning(self):
        # Start the operations, and retrieve the future.
        # they will stop at the first yield and not go further until
        # we yield them
        f1 = self.manager.start_container("username",
                                          "image_id1",
                                          "mapping_id",
                                          "/foo/bar",
                                          None,
                                          )

        f2 = self.manager.start_container("username",
                                          "image_id1",
                                          "mapping_id",
                                          "/foo/baz",
                                          None,
                                          )

        # If these yielding raise a KeyError, it is because the second
        # one tries to remove the same key from the list, but it has been
        # already removed by the first one. Race condition.
        yield f1

        with self.assertRaises(OperationInProgress):
            yield f2

        self.assertEqual(self.mock_docker_client.start.call_count, 1)

    @gen_test
    def test_race_condition_stopping(self):
        f1 = self.manager.stop_and_remove_container("container_id1")
        f2 = self.manager.stop_and_remove_container("container_id1")

        yield f1

        with self.assertRaises(OperationInProgress):
            yield f2

        self.assertEqual(self.mock_docker_client.stop.call_count, 1)

    @gen_test
    def test_start_already_present_container(self):
        mock_client = self.mock_docker_client

        result = yield self.manager.start_container(
            "username",
            "image_name1",
            "mapping_id",
            "/base/url",
            None,
            None)
        self.assertTrue(mock_client.start.called)
        self.assertIsInstance(result, Container)

        # Stop should have been called and the container removed
        self.assertTrue(mock_client.stop.called)
        self.assertTrue(mock_client.remove_container.called)

    @gen_test
    def test_image(self):
        image = yield self.manager.image('image_name1')
        self.assertIsInstance(image, Image)
        self.assertEqual(image.description,
                         'Ubuntu machine with mayavi preinstalled')
        self.assertEqual(image.icon_128, "")
        self.assertEqual(image.ui_name, "Mayavi 4.4.4")
        self.assertEqual(image.docker_id, 'image_id1')

        image = yield self.manager.image("whatev")
        self.assertIsNone(image)

    @gen_test
    def test_start_container_with_nonexisting_volume_source(self):
        # These volume sources are invalid
        volumes = {'~no_way_this_be_valid': {'bind': '/target_vol1',
                                             'mode': 'ro'},
                   '/no_way_this_be_valid': {'bind': '/target_vol2',
                                             'mode': 'ro'}}

        # This volume source is valid
        good_path = os.path.abspath('.')
        volumes[good_path] = {'bind': '/target_vol3',
                              'mode': 'ro'}

        yield self.manager.start_container("username",
                                           "image_id1",
                                           "mapping_id",
                                           "/foo/bar",
                                           volumes,
                                           )

        # Call args and keyword args that create_container receives
        docker_client = self.manager._docker_client
        args = docker_client._sync_client.create_container.call_args
        actual_volume_targets = args[1]['volumes']

        # Invalid volume paths should have been filtered away
        self.assertNotIn('/target_vol1', actual_volume_targets)
        self.assertNotIn('/target_vol2', actual_volume_targets)

        # The current directory is valid, should stay
        self.assertIn('/target_vol3', actual_volume_targets)

    @gen_test
    def test_start_container_exception_cleanup(self):
        self.assertFalse(self.mock_docker_client.stop.called)
        self.assertFalse(self.mock_docker_client.remove_container.called)

        def raiser(*args, **kwargs):
            raise Exception("Boom!")

        self.manager._docker_client.start = mock.Mock(side_effect=raiser)

        with self.assertRaisesRegex(Exception, 'Boom!'):
            yield self.manager.start_container("username",
                                               "image_id1",
                                               "mapping_id",
                                               "/base/url",
                                               None,
                                               None)

        self.assertTrue(self.mock_docker_client.stop.called)
        self.assertTrue(self.mock_docker_client.remove_container.called)

    @gen_test
    def test_start_container_exception_cleanup_2(self):
        # Same test as above, but checks after the start (at ip and port)
        self.assertFalse(self.mock_docker_client.stop.called)
        self.assertFalse(self.mock_docker_client.remove_container.called)

        def raiser(*args, **kwargs):
            raise Exception("Boom!")

        self.manager._docker_client.port = mock.Mock(side_effect=raiser)

        with self.assertRaisesRegex(Exception, 'Boom!'):
            yield self.manager.start_container("username",
                                               "image_id1",
                                               "mapping_id",
                                               "/base/url",
                                               None,
                                               None)

        self.assertTrue(self.mock_docker_client.stop.called)
        self.assertTrue(self.mock_docker_client.remove_container.called)

    @gen_test
    def test_start_container_with_environment(self):
        mock_client = self.mock_docker_client

        environment = {
            "FOO": "bar"
        }

        yield self.manager.start_container(
            "username",
            "image_name1",
            "mapping_id",
            "/base/url",
            None,
            environment)

        self.assertEqual(
            mock_client.create_container.call_args[1]["environment"]["FOO"],
            "bar")

    @gen_test
    def test_different_realm(self):
        manager = ContainerManager(docker_config={},
                                   realm="anotherrealm")
        manager._docker_client._sync_client = create_docker_client()

        result = yield manager.containers_from_mapping_id("user_name",
                                                          "mapping_id")
        self.assertEqual(result, [])

        result = yield manager.running_containers()
        self.assertEqual(result, [])

    @gen_test
    def test_not_stopping_if_different_realm(self):
        manager = ContainerManager(docker_config={},
                                   realm="anotherrealm")
        manager._docker_client._sync_client = create_docker_client()
        yield manager.stop_and_remove_container("container_id1")

        self.assertFalse(self.mock_docker_client.stop.called)
        self.assertFalse(self.mock_docker_client.remove_container.called)
<|endoftext|>


<commit_before><commit_msg>Delete __init__.py<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
#   repository.py — Class to handle the repository
#
#   This file is part of debexpo - http://debexpo.workaround.org
#
#   Copyright © 2008 Jonny Lamb <jonnylamb@jonnylamb.com
#
#   Permission is hereby granted, free of charge, to any person
#   obtaining a copy of this software and associated documentation
#   files (the "Software"), to deal in the Software without
#   restriction, including without limitation the rights to use,
#   copy, modify, merge, publish, distribute, sublicense, and/or sell
#   copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following
#   conditions:
#
#   The above copyright notice and this permission notice shall be
#   included in all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#   OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
#   HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
#   WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
#   OTHER DEALINGS IN THE SOFTWARE.

"""
This module holds the Repository class to handle the repository.
"""

__author__ = 'Jonny Lamb'
__copyright__ = 'Copyright © 2008 Jonny Lamb'
__license__ = 'MIT'

import os
from debian_bundle import deb822
from sqlalchemy import select
import gzip
import bz2

from debexpo.lib.utils import get_package_dir
from debexpo.model import meta
from debexpo.model.package_files import PackageFile
from debexpo.model.source_packages import SourcePackage
from debexpo.model.package_versions import PackageVersion

class Repository(object):
    """
    Class to handle the repository.
    """
    def __init__(self, repository):
        """
        Class constructor. Sets the repository base directory variable and other misc variables.
        """
        self.repository = repository
        self.compression = [(gzip.GzipFile, 'gz'), (bz2.BZ2File, 'bz2')]

    def _dsc_to_sources(self, package_file):
        """
        Reads the contents of a dsc file and converts it to a Sources file entry.

        ``file``
            Filename of the dsc to read.
        """
        filename = os.path.join(self.repository, package_file.filename)
        package_version = package_file.source_package.package_version
        package = package_version.package

        # Read the dsc file.
        dsc = deb822.Dsc(file(filename))

        # There are a few differences between a dsc file and a Sources entry, listed and acted
        # upon below:

        # Firstly, the "Source" field in the dsc is simply renamed to "Package".
        dsc['Package'] = dsc.pop('Source')

        # There needs to be a "Directory" field to tell the package manager where to download the
        # package from. This is in the format (for the test package in the component "main"):
        #   pool/main/t/test
        dsc['Directory'] = 'pool/%s/%s' % (package_version.component, get_package_dir(package.name))

        # The dsc file, its size, and its md5sum needs to be added to the "Files" field. This is
        # unsurprisingly not in the original dsc file!
        dsc['Files'].append({'md5sum' : package_file.md5sum, 'size': package_file.size, 'name' : package_file.filename.split('/')[-1]})

        # Get a nice rfc822 output of this dsc, now Sources, entry.
        return dsc.dump()

    def _get_sources_file(self, distribution, component):
        """
        Does a query to find all packages that fit the criteria of distribution
        and component and returns the contents of a Sources file.

        ``distribution``
            Name of the distribution to look at.

        ``component``
            Name of the component to look at.
        """
        # Get all PackageFile instances...
        dscfiles = meta.session.query(PackageFile)

        # ...include only *.dsc files...
        dscfiles = dscfiles.filter(PackageFile.filename.like('%.dsc'))

        # ...where there is a SourcePackage instance...
        dscfiles = dscfiles.filter(PackageFile.source_package_id == SourcePackage.id)

        # ...where there is a PackageVersion instance...
        dscfiles = dscfiles.filter(SourcePackage.package_version_id == PackageVersion.id)

        # ...include only packages in the specified component...
        dscfiles = dscfiles.filter(PackageVersion.component == component)

        # ...include only package in the specified distribution...
        dscfiles = dscfiles.filter(PackageVersion.distribution == distribution)

        # ...and finally create a list of PackageFile instances.
        dscfiles = dscfiles.all()

        entries = []

        # Loop through dsc files.
        for file in dscfiles:
            entries.append(self._dsc_to_sources(file))

        # Each entry is simply joined by a blank newline, so do just that to create
        # the finished Sources file.
        return '\n'.join(entries)

    def _get_archs_list(self, list):
        """
        Takes a list of directory names and returns an altered list where only items
        with "binary-" prefix stay. This "binary-" prefix is also removed.

        For example, an input of ['binary-i386', 'binary-amd64', 'foo'] would produce
        an output of ['i386', 'amd64'].

        ``list``
            List of directory names to look at.
        """
        for item in list:
            if not item.startswith('binary-'):
                list.remove(item)

        return [z[7:] for z in list]

    def _append_current_distributions(self, distributions):
        """
        Take a look at the current directory layout and add distribution - component
        entries to the ``distributions`` dict. This is useful as if the contents of the
        Sources file is empty then the file is deleted. If a package has no longer present
        but still exists in the current Sources file, this process will remove it.
        """
        distsdir = os.path.join(self.repository, 'dists')

        # The first time the repository dists files are generated, the dists directory
        # isn't present.
        if not os.path.isdir(distsdir):
            return distributions

        dists = os.listdir(distsdir)

        # Components are subdirectories in distribution directories.
        for dist in dists:
            components = os.listdir(os.path.join(distsdir, dist))

            comps = {}

            # Loop through each component.
            for component in components:
                # Get arch list.
                archs = self._get_archs_list(os.listdir(os.path.join(distsdir, dist, component)))

                if dist in distributions.keys():
                    # The distribution already exists in the big dictionary. Therefore we need to
                    # look at each component to see whether it already exists.

                    if component in distributions[dist].keys():
                        # The component already exists in the big dictionary under the current
                        # distribution. Therefore we need to look at each arch to see whether it
                        # already exists.

                        for arch in archs:
                            if arch not in distributions[dist][component]:
                                # The arch doesn't already exist. Add it.
                                distributions[dist][component].append(arch)
                    else:
                        # The component doesn't already exist. It can easily be added now then.
                        distributions[dist][component] = archs
                else:
                    # The distribution doesn't already exist in the big dictionary. Therefore
                    # we can use the temporary comps dictionary to store { component : [archs] }
                    # information.
                    comps[component] = archs

            # Test whether the comps dictionary has been used. It will only have been used when
            # the distribution didn't already exist in the big dictionary.
            if len(comps) != 0:
                distributions[dist] = comps

        return distributions

    def _get_dists_comps_archs(self):
        """
        Return a dictionary of distribution, components and dists in the distribution::

            { distribution1 : { component1 : [ arch1, arch2, ... ],
                                component2 : [ arch3, ... ] },
                                ...
              distribution2 : ...
            }

        I'm not particularly happy with the implementation of this. Perhaps one day when I'm
        much more skilled at SQLAlchemy, I'll take a look at this and see if I can perform
        all the operations on the db level. That would be nice.
        """
        # Get distinct distributions from package_versions.
        dists = meta.engine.execute(select([PackageVersion.c.distribution]).distinct())

        distributions = {}

        # Loop through each distribution.
        for dist in dists.fetchall():
            distributions[dist[0]] = {}

            components = meta.engine.execute(select([PackageVersion.c.component], PackageVersion.c.distribution == dist[0]).distinct())

            comps = {}

            # Loop through each component within the current distribution.
            for comp in components.fetchall():

                comps[comp[0]] = []

                # Get all distinct archs...
                archs_query = select([BinaryPackage.arch]).distinct()

                # ...where an PackageVersion instance exists for it...
                archs_query = archs_query.where(BinaryPackage.package_version_id == PackageVersion.id)

                # ...where its distribution is the current one in iteration...
                archs_query = archs_query.where(PackageVersion.distribution == dist[0])

                # ...where its component is the current one in iteration...
                archs_query = archs_query.where(PackageVersion.component == comp[0])

                # ...execute that.
                archs = meta.engine.execute(archs_query)

                # Loop through each arch.
                for arch in archs.fetchall():
                    comps[comp[0]].append(arch[0])

            # Add distribution to dictionary.
            distributions[dist[0]] = comps

        return self._append_current_distributions(distributions)

    def _check_directories(self, dist, component, arch=None):
        """
        Checks whether the directories needed for a dists file are present, and if not
        it creates them.

        ``dist``
            Name of the distribution.

        ``component``
            Name of the component within the distribution.

        ``arch``
            Architecture of the dists file. This defaults to *source*.
        """
        dir = self.repository

        # Default to source, otherwise binary-%(arch)s.
        if arch is None:
            arch = 'source'
        else:
            arch = 'binary-%s' % arch

        # Check each subdirectory of the dists directory.
        for dirname in ['dists', dist, component, arch]:
            dir = os.path.join(dir, dirname)

            if not os.path.isdir(dir):
                os.mkdir(dir)

    def update_sources(self):
        """
        Updates all the Sources.{gz,bz2} files for all distributions and components
        by looking at all source packages.
        """
        # Get distributions and components.
        dists = self._get_dists_comps_archs()

        for dist, components in dists.iteritems():
            for component in components.keys():
                # Make sure all directories are present.
                self._check_directories(dist, component)

                # Create Sources file content.
                sources = self._get_sources_file(dist, component)
                filename = os.path.join(self.repository, 'dists', dist, component, 'source', 'Sources')

                # If the Sources file is empty, remove Sources.{gz,bz2} files. This way the
                # directory will be removed on a clean-up as it is empty.
                if sources == '':
                    for format, extension in self.compression:
                        if os.path.isfile('%s.%s' % (filename, extension)):
                            os.remove('%s.%s' % (filename, extension))
                else:
                    for format, extension in self.compression:
                        # Create the Sources files.
                        f = format('%s.%s' % (filename, extension), 'w')
                        f.write(sources)
                        f.close()
<commit_msg>Added methods to debexpo.lib.repository to update Packages files.<commit_after># -*- coding: utf-8 -*-
#
#   repository.py — Class to handle the repository
#
#   This file is part of debexpo - http://debexpo.workaround.org
#
#   Copyright © 2008 Jonny Lamb <jonnylamb@jonnylamb.com
#
#   Permission is hereby granted, free of charge, to any person
#   obtaining a copy of this software and associated documentation
#   files (the "Software"), to deal in the Software without
#   restriction, including without limitation the rights to use,
#   copy, modify, merge, publish, distribute, sublicense, and/or sell
#   copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following
#   conditions:
#
#   The above copyright notice and this permission notice shall be
#   included in all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
#   EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
#   OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
#   HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
#   WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
#   OTHER DEALINGS IN THE SOFTWARE.

"""
This module holds the Repository class to handle the repository.
"""

__author__ = 'Jonny Lamb'
__copyright__ = 'Copyright © 2008 Jonny Lamb'
__license__ = 'MIT'

import os
from debian_bundle import deb822, debfile
from sqlalchemy import select
import gzip
import bz2

from debexpo.lib.utils import get_package_dir
from debexpo.model import meta
from debexpo.model.package_files import PackageFile
from debexpo.model.source_packages import SourcePackage
from debexpo.model.binary_packages import BinaryPackage
from debexpo.model.package_versions import PackageVersion

class Repository(object):
    """
    Class to handle the repository.
    """
    def __init__(self, repository):
        """
        Class constructor. Sets the repository base directory variable and other misc variables.
        """
        self.repository = repository
        self.compression = [(gzip.GzipFile, 'gz'), (bz2.BZ2File, 'bz2')]

    def _dsc_to_sources(self, package_file):
        """
        Reads the contents of a dsc file and converts it to a Sources file entry.

        ``file``
            Filename of the dsc to read.
        """
        filename = os.path.join(self.repository, package_file.filename)
        package_version = package_file.source_package.package_version
        package = package_version.package

        # Read the dsc file.
        dsc = deb822.Dsc(file(filename))

        # There are a few differences between a dsc file and a Sources entry, listed and acted
        # upon below:

        # Firstly, the "Source" field in the dsc is simply renamed to "Package".
        dsc['Package'] = dsc.pop('Source')

        # There needs to be a "Directory" field to tell the package manager where to download the
        # package from. This is in the format (for the test package in the component "main"):
        #   pool/main/t/test
        dsc['Directory'] = 'pool/%s/%s' % (package_version.component, get_package_dir(package.name))

        # The dsc file, its size, and its md5sum needs to be added to the "Files" field. This is
        # unsurprisingly not in the original dsc file!
        dsc['Files'].append({'md5sum' : package_file.md5sum, 'size': package_file.size, 'name' : package_file.filename.split('/')[-1]})

        # Get a nice rfc822 output of this dsc, now Sources, entry.
        return dsc.dump()

    def _deb_to_packages(self, package_file):
        """
        Reads a binary package and outputs a Packages file entry.

        ``file``
            Filename of the deb to read.
        """
        filename = os.path.join(self.repository, package_file.filename)
        package_version = package_file.binary_package.package_version
        package = package_version.package

        # Read the deb file.
        deb = debfile.DebFile(filename).debcontrol()

        # There are a few additions toa debian/control file to make a Packages entry, listed
        # and acted upon below:

        deb['Filename'] = package_file.filename
        deb['Size'] = str(package_file.size)
        deb['MD5sum'] = package_file.md5sum

        return deb.dump()

    def _get_sources_file(self, distribution, component):
        """
        Does a query to find all packages that fit the criteria of distribution
        and component and returns the contents of a Sources file.

        ``distribution``
            Name of the distribution to look at.

        ``component``
            Name of the component to look at.
        """
        # Get all PackageFile instances...
        dscfiles = meta.session.query(PackageFile)

        # ...include only *.dsc files...
        dscfiles = dscfiles.filter(PackageFile.filename.like('%.dsc'))

        # ...where there is a SourcePackage instance...
        dscfiles = dscfiles.filter(PackageFile.source_package_id == SourcePackage.id)

        # ...where there is a PackageVersion instance...
        dscfiles = dscfiles.filter(SourcePackage.package_version_id == PackageVersion.id)

        # ...include only packages in the specified component...
        dscfiles = dscfiles.filter(PackageVersion.component == component)

        # ...include only package in the specified distribution...
        dscfiles = dscfiles.filter(PackageVersion.distribution == distribution)

        # ...and finally create a list of PackageFile instances.
        dscfiles = dscfiles.all()

        entries = []

        # Loop through dsc files.
        for file in dscfiles:
            entries.append(self._dsc_to_sources(file))

        # Each entry is simply joined by a blank newline, so do just that to create
        # the finished Sources file.
        return '\n'.join(entries)

    def _get_packages_file(self, distribution, component, arch):
        """
        Does a query to find all packages that fit the criteria of distribution,
        component and architecture and returns the contents of a Packages file.

        ``distribution``
            Name of the distribution to look at.

        ``component``
            Name of the component to look at.

        ``arch``
            Name of the architecture to look at.
        """
        # Get all PackageFile instances...
        debfiles = meta.session.query(PackageFile)

        # ...include only *.deb files...
        debfiles = debfiles.filter(PackageFile.filename.like('%.deb'))

        # ...where there is a BinaryPackage instance...
        debfiles = debfiles.filter(PackageFile.binary_package_id == BinaryPackage.id)

        # ...where the BinaryPackage has Arch: %(arch)s or Arch: all...
        debfiles = debfiles.filter(BinaryPackage.arch == arch or BinaryPackage.arch == 'all')

        # ...where there is a PackageVersion instance...
        debfiles = debfiles.filter(BinaryPackage.package_version_id == PackageVersion.id)

        # ...include only packages in the specified component...
        debfiles = debfiles.filter(PackageVersion.component == component)

        # ...include only package in the specified distribution...
        debfiles = debfiles.filter(PackageVersion.distribution == distribution)

        # ...and finally create a list of PackageFile instances.
        debfiles = debfiles.all()

        entries = []

        # Loop through deb files.
        for file in debfiles:
            entries.append(self._deb_to_packages(file))

        # Each entry is simply joined by a blank newline, so do just that to create
        # the finished Packages file.
        return '\n'.join(entries)

    def _get_archs_list(self, list):
        """
        Takes a list of directory names and returns an altered list where only items
        with "binary-" prefix stay. This "binary-" prefix is also removed.

        For example, an input of ['binary-i386', 'binary-amd64', 'foo'] would produce
        an output of ['i386', 'amd64'].

        ``list``
            List of directory names to look at.
        """
        for item in list:
            if not item.startswith('binary-'):
                list.remove(item)

        return [z[7:] for z in list]

    def _append_current_distributions(self, distributions):
        """
        Take a look at the current directory layout and add distribution - component
        entries to the ``distributions`` dict. This is useful as if the contents of the
        Sources file is empty then the file is deleted. If a package has no longer present
        but still exists in the current Sources file, this process will remove it.
        """
        distsdir = os.path.join(self.repository, 'dists')

        # The first time the repository dists files are generated, the dists directory
        # isn't present.
        if not os.path.isdir(distsdir):
            return distributions

        dists = os.listdir(distsdir)

        # Components are subdirectories in distribution directories.
        for dist in dists:
            components = os.listdir(os.path.join(distsdir, dist))

            comps = {}

            # Loop through each component.
            for component in components:
                # Get arch list.
                archs = self._get_archs_list(os.listdir(os.path.join(distsdir, dist, component)))

                if dist in distributions.keys():
                    # The distribution already exists in the big dictionary. Therefore we need to
                    # look at each component to see whether it already exists.

                    if component in distributions[dist].keys():
                        # The component already exists in the big dictionary under the current
                        # distribution. Therefore we need to look at each arch to see whether it
                        # already exists.

                        for arch in archs:
                            if arch not in distributions[dist][component]:
                                # The arch doesn't already exist. Add it.
                                distributions[dist][component].append(arch)
                    else:
                        # The component doesn't already exist. It can easily be added now then.
                        distributions[dist][component] = archs
                else:
                    # The distribution doesn't already exist in the big dictionary. Therefore
                    # we can use the temporary comps dictionary to store { component : [archs] }
                    # information.
                    comps[component] = archs

            # Test whether the comps dictionary has been used. It will only have been used when
            # the distribution didn't already exist in the big dictionary.
            if len(comps) != 0:
                distributions[dist] = comps

        return distributions

    def _get_dists_comps_archs(self):
        """
        Return a dictionary of distribution, components and dists in the distribution::

            { distribution1 : { component1 : [ arch1, arch2, ... ],
                                component2 : [ arch3, ... ] },
                                ...
              distribution2 : ...
            }

        I'm not particularly happy with the implementation of this. Perhaps one day when I'm
        much more skilled at SQLAlchemy, I'll take a look at this and see if I can perform
        all the operations on the db level. That would be nice.
        """
        # Get distinct distributions from package_versions.
        dists = meta.engine.execute(select([PackageVersion.c.distribution]).distinct())

        distributions = {}

        # Loop through each distribution.
        for dist in dists.fetchall():
            distributions[dist[0]] = {}

            components = meta.engine.execute(select([PackageVersion.c.component], PackageVersion.c.distribution == dist[0]).distinct())

            comps = {}

            # Loop through each component within the current distribution.
            for comp in components.fetchall():

                comps[comp[0]] = []

                # Get all distinct archs...
                archs_query = select([BinaryPackage.arch]).distinct()

                # ...where an PackageVersion instance exists for it...
                archs_query = archs_query.where(BinaryPackage.package_version_id == PackageVersion.id)

                # ...where its distribution is the current one in iteration...
                archs_query = archs_query.where(PackageVersion.distribution == dist[0])

                # ...where its component is the current one in iteration...
                archs_query = archs_query.where(PackageVersion.component == comp[0])

                # ...execute that.
                archs = meta.engine.execute(archs_query)

                # Loop through each arch.
                for arch in archs.fetchall():
                    comps[comp[0]].append(arch[0])

            # Add distribution to dictionary.
            distributions[dist[0]] = comps

        return self._append_current_distributions(distributions)

    def _check_directories(self, dist, component, arch=None):
        """
        Checks whether the directories needed for a dists file are present, and if not
        it creates them.

        ``dist``
            Name of the distribution.

        ``component``
            Name of the component within the distribution.

        ``arch``
            Architecture of the dists file. This defaults to *source*.
        """
        dir = self.repository

        # Default to source, otherwise binary-%(arch)s.
        if arch is None:
            arch = 'source'
        else:
            arch = 'binary-%s' % arch

        # Check each subdirectory of the dists directory.
        for dirname in ['dists', dist, component, arch]:
            dir = os.path.join(dir, dirname)

            if not os.path.isdir(dir):
                os.mkdir(dir)

    def update_sources(self):
        """
        Updates all the Sources.{gz,bz2} files for all distributions and components
        by looking at all source packages.
        """
        # Get distributions and components.
        dists = self._get_dists_comps_archs()

        for dist, components in dists.iteritems():
            for component in components.keys():
                # Make sure all directories are present.
                self._check_directories(dist, component)

                # Create Sources file content.
                sources = self._get_sources_file(dist, component)
                filename = os.path.join(self.repository, 'dists', dist, component, 'source', 'Sources')

                # If the Sources file is empty, remove Sources.{gz,bz2} files. This way the
                # directory will be removed on a clean-up as it is empty.
                if sources == '':
                    for format, extension in self.compression:
                        if os.path.isfile('%s.%s' % (filename, extension)):
                            os.remove('%s.%s' % (filename, extension))
                else:
                    for format, extension in self.compression:
                        # Create the Sources files.
                        f = format('%s.%s' % (filename, extension), 'w')
                        f.write(sources)
                        f.close()

    def update_packages(self):
        """
        Updates all the Packages.{gz,bz2} files for all distributions and components
        by looking at all binary packages.
        """
        # Get distributions and components.
        dists = self._get_dists_comps_archs()

        for dist, components in dists.iteritems():
            for component, archs in components.iteritems():
                for arch in archs:
                    # Make sure all directories are present.
                    self._check_directories(dist, component, arch)

                    # Create Packages file content.
                    packages = self._get_packages_file(dist, component, arch)
                    filename = os.path.join(self.repository, 'dists', dist, component, 'binary-%s' % arch, 'Packages')

                    # If the Packages file is empty, remove Packages.{gz,bz2} files. This way
                    # the directory will be removed on a clean-up as it is empty.
                    if packages == '':
                        for format, extension in self.compression:
                            if os.path.isfile('%s.%s' % (filename, extension)):
                                os.remove('%s.%s' % (filename, extension))
                    else:
                        for format, extension in self.compression:
                            # Create the Packages files.
                            f = format('%s.%s' % (filename, extension), 'w')
                            f.write(packages)
                            f.close()

    def update(self):
        """
        Updates both Sources and Packages files in the repository.
        """

        self.update_sources()
        self.update_packages()
<|endoftext|>