<commit_before>
import os
import tarfile
import zipfile
import tempfile
import traceback

import synapse.exc as s_exc
from synapse.common import *

# 10 MB
read_chunk_sz = 1048576*10

'''
Provide a generic opener API for paths that cross into supported container files.

For example: /dir0/dir1/foo.zip/d0/bar
The above represents a file named bar located in the d0 directory inside the foo.zip zip archive
located on the filesystem in the /dir0/dir1 directory.
'''

class FilePath(object):
    def __init__(self, path, remainder, parent=None):
        '''
        base path object for regular filesystems
        '''
        self.parent = parent
        self.path = path
        self.remainder = remainder

        self.child = None
        self.tempfd = None

        if self.remainder:
            self._process_next()

    def _type(self):
        return 'fs.reg.file'

    def isfile(self):
        return True

    def _process_next(self):
        return

    def getChild(self):
        return self.child

class FilePathFile(FilePath):
    '''
    The base path object for filesystem files.
    '''
    def _open(self, mode='r'):
        return open(self.path, mode=mode)

class FilePathDir(FilePath):
    '''
    The base dir object for filesystem files.
    '''
    def _type(self):
        return 'fs.reg.dir'

    def isfile(self):
        return False

    def _process_next(self):
        child_path, child_remainder = _get_child_paths(self.path, self.remainder)

        cls = _get_path_class(child_path)
        self.child = cls(child_path, child_remainder, parent=self)

class FilePathTarDir(FilePath):
    def __init__(self, tarfd, path, remainder, parent=None):
        '''
        This is a directory in a tar archive.
        '''
        self.tarfd = tarfd
        super(FilePathTarDir, self).__init__(path, remainder, parent=parent)

    def _type(self):
        return 'inner.tar.dir'

    def isfile(self):
        return False

    def _process_next(self):
        '''
        This will only be called if the member from the remainder is a file or dir
        '''

        tar_dirs, tar_files = tar_enumerate(self.tarfd)

        child_path, child_remainder = _get_child_paths(self.path, self.remainder)

        child_tar_path = _container_pathnorm(child_path)
        if child_tar_path in tar_dirs:
            self.child = FilePathTarDir(self.tarfd, child_path, child_remainder, parent=self)

        elif child_tar_path in tar_files:
            if not child_remainder:
                self.child = FilePathTarFile(self.tarfd, child_path, child_remainder, parent=self)
                return

            # if the child has a remainder then extract otherwise, don't
            self.tempfd = tempfile.NamedTemporaryFile(prefix='syn_sfp_')

            tfd = self.tarfd.extractfile(child_path.strip('/'))
            tbuf = tfd.read(read_chunk_sz)
            while tbuf:
                self.tempfd.write(tbuf)
                tbuf = tfd.read(read_chunk_sz)
            self.tempfd.flush()
            child_path = self.tempfd.name
            self.child = _get_path_class(child_path)(child_path, child_remainder, parent=self)

        elif self.remainder:
            raise s_exc.NoSuchPath(path=os.path.join(self.path, self.remainder))

class FilePathTarFile(FilePathTarDir):
    def _type(self):
        return 'inner.tar.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return self.tarfd.extractfile(self.path.lstrip('/'))

class FilePathTar(FilePathTarDir):
    def __init__(self, path, remainder, parent=None):
        tarfd = tarfile.open(path)
        self.tarpath = path

        super(FilePathTar, self).__init__(tarfd, '/', remainder, parent=parent)

    def _type(self):
        return 'fs.tar.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return open(self.tarpath, mode=mode)

class FilePathZipDir(FilePath):
    def __init__(self, zipfd, path, remainder, parent=None):
        '''
        This is a directory in a zip.
        '''
        self.zipfd = zipfd
        super(FilePathZipDir, self).__init__(path, remainder, parent=parent)

    def _type(self):
        return 'inner.zip.dir'

    def isfile(self):
        return False

    def _process_next(self):
        '''
        This will only be called if the member from the remainder is a file or dir
        '''

        zip_dirs, zip_files = zip_enumerate(self.zipfd)

        child_path, child_remainder = _get_child_paths(self.path, self.remainder)

        child_zip_path = _container_pathnorm(child_path)
        if child_zip_path in zip_dirs:
            self.child = FilePathZipDir(self.zipfd, child_path, child_remainder, parent=self)

        elif child_zip_path in zip_files:
            if not child_remainder:
                self.child = FilePathZipFile(self.zipfd, child_path, child_remainder, parent=self)
                return

            # extract the file to a directory and parse it
            self.tempfd = tempfile.NamedTemporaryFile(prefix='syn_sfp_')

            zfd = self.zipfd.open(child_path.strip('/'), mode='r')
            zbuf = zfd.read(read_chunk_sz)
            while zbuf:
                self.tempfd.write(zbuf)
                zbuf = zfd.read(read_chunk_sz)

            self.tempfd.flush()
            child_path = self.tempfd.name
            self.child = _get_path_class(child_path)(child_path, child_remainder, parent=self)

        elif self.remainder:
            raise s_exc.NoSuchPath(path=os.path.join(self.path, self.remainder))

class FilePathZipFile(FilePathZipDir):
    def _type(self):
        return 'inner.zip.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return self.zipfd.open(self.path.lstrip('/'), mode='r')

class FilePathZip(FilePathZipDir):
    def __init__(self, path, remainder, parent=None):
        zipfd = zipfile.ZipFile(path)
        self.zippath = path

        super(FilePathZip, self).__init__(zipfd, '/', remainder, parent=parent)

    def _type(self):
        return 'fs.zip.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return open(self.zippath, mode=mode)

def zip_enumerate(zipfd):
    '''
    Return the files and dirs in a zip file
    '''
    files = set()
    dirs = set()

    for path in zipfd.namelist():

        for spath in _get_subpaths(path):
            try:
                info = zipfd.getinfo(spath)
                files.add(spath)
            except KeyError as e:
                dirs.add(spath)

    return dirs, files

def tar_enumerate(tarfd):
    '''
    Return the files and dirs in a tar archive
    '''
    files = set()
    dirs = set()

    for path in tarfd.getnames():

        for spath in _get_subpaths(path):
            info = tarfd.getmember(spath)

            if info.isdir():
                dirs.add(spath)
            if info.isfile():
                files.add(spath)

    return dirs, files

def _get_path_class(path):
    '''
    Returns the class to handle the type of item located at path.  This function
    only operates on regular os. accessible paths
    '''

    if not os.path.exists(path):
        raise s_exc.NoSuchPath(path=path)

    if os.path.isdir(path):
        return _path_ctors.get('fs.reg.dir')
    mime = _mime_file(path)
    return _path_ctors.get(mime)

def _mime_file(path):
    '''
    Assumes the path exists and is a file.
    returns reg.file unless it is a known container
    '''
    if zipfile.is_zipfile(path):
        return 'fs.zip.file'
    if tarfile.is_tarfile(path):
        return 'fs.tar.file'

    return 'fs.reg.file'

def _lsplit(path):
    '''
    Split a pathname.  Returns tuple "(head, tail)" where "tail" is everything after the
    FIRST slash that is not at position 0.  Either part may be empty.

    It's the inverse of os.path.split
    '''

    drive, rpath = os.path.splitdrive(path)
    parts = os.path.normpath(rpath).split(os.path.sep)
    if parts[0] == '':
        parts[0] = os.path.sep
    head = os.path.join(drive, parts[0])

    tail = ''
    if parts[1:]:
        tail = os.path.join(*parts[1:])
    return head, tail

def _get_child_paths(headpath, tailpath):
    '''
    Take a split head/tail path and create a new head tail that is one directory lower
    in the joined path
    '''
    child_part, child_remainder = _lsplit(tailpath)
    child_headpath = os.path.join(headpath, child_part)
    return child_headpath, child_remainder

def _container_pathnorm(path):
    return path.replace('\\', '/').strip('/')

def _get_subpaths(path):
    '''
    Returns a list of subpaths in a path, one for each level in the path
    This is an internal function used for ONLY for iterating over paths in a container
    As such it should not be used for filesystem paths since separators will vary across platforms
    '''
    path = path.replace('\\', '/')
    path = path.rstrip('/')

    path_parts = path.split('/')
    if path_parts[0] == '':
        path_parts[0] = '/'

    paths = []
    for i in range(len(path_parts)):

        ipath = os.path.join(*path_parts[:i+1])
        paths.append(ipath)

    return paths

def _parse_path(*paths):
    '''
    Internal function to parse the incoming path.
    lists of paths are joined prior to parsing
    '''

    if not paths or None in paths:
        return None

    path = genpath(*paths)

    if not path:
        return None

    fpobj = None
    child_path, child_remainder = _lsplit(path)

    try:
        cls = _get_path_class(child_path)
        fpobj = cls(child_path, child_remainder)

        while fpobj.getChild():
            fpobj = fpobj.getChild()
    except s_exc.NoSuchPath as e:
        return None

    return fpobj

def openfile(*paths, mode='r'):
    '''
    Returns a read-only file-like object even if the path terminates inside a container file.
    If the path is a regular os accessible path mode is passed through.  If the path terminates
    in a container file mode is ignored.

    If the path does not exist a NoSuchPath exception is raised.

    ex.
    openfile('/foo/bar/baz.egg/path/inside/zip/to/file')
    '''
    fpobj = _parse_path(*paths)
    if not fpobj:
        raise s_exc.NoSuchPath(path='%r' % (paths,))
    if not fpobj.isfile():
        path = genpath(*paths)
        raise s_exc.NoSuchPath(path=path)

    fd = fpobj._open(mode=mode)
    return fd

def isfile(*paths):
    '''
    Determines if the path is a file, even if the path terminates inside a container file.
    If a list of paths are provided, they are joined first.
    Returns a boolean.
    '''
    fpath = _parse_path(*paths)
    if fpath and fpath.isfile():
        return True
    return False

def isdir(*paths):
    '''
    Determines if the path is a directory, even if the path terminates inside a container file.
    If a list of paths are provided, they are joined first.
    Returns a boolean.
    '''
    fpath = _parse_path(*paths)
    if fpath and not fpath.isfile():
        return True
    return False

def exists(*paths):
    '''
    Determines if the path exists even if the path terminates inside a container file.
    If a list of paths are provided, they are joined first.
    Returns a boolean.
    '''
    if _parse_path(*paths) == None:
        return False
    return True

_path_ctors = {
    'fs.reg.dir': FilePathDir,
    'fs.reg.file': FilePathFile,
    'fs.tar.file': FilePathTar,
    'fs.zip.file': FilePathZip,
    'inner.tar.dir': FilePathTarDir,
    'inner.zip.dir': FilePathZipDir,
    'inner.tar.file': FilePathTarFile,
    'inner.zip.file': FilePathZipFile,
}

<commit_msg>updated openfile to take generic keyword args so that it is syntax compliant with python2.7<commit_after>
import os
import tarfile
import zipfile
import tempfile
import traceback

import synapse.exc as s_exc
from synapse.common import *

# 10 MB
read_chunk_sz = 1048576*10

'''
Provide a generic opener API for paths that cross into supported container files.

For example: /dir0/dir1/foo.zip/d0/bar
The above represents a file named bar located in the d0 directory inside the foo.zip zip archive
located on the filesystem in the /dir0/dir1 directory.
'''

class FilePath(object):
    def __init__(self, path, remainder, parent=None):
        '''
        base path object for regular filesystems
        '''
        self.parent = parent
        self.path = path
        self.remainder = remainder

        self.child = None
        self.tempfd = None

        if self.remainder:
            self._process_next()

    def _type(self):
        return 'fs.reg.file'

    def isfile(self):
        return True

    def _process_next(self):
        return

    def getChild(self):
        return self.child

class FilePathFile(FilePath):
    '''
    The base path object for filesystem files.
    '''
    def _open(self, mode='r'):
        return open(self.path, mode=mode)

class FilePathDir(FilePath):
    '''
    The base dir object for filesystem files.
    '''
    def _type(self):
        return 'fs.reg.dir'

    def isfile(self):
        return False

    def _process_next(self):
        child_path, child_remainder = _get_child_paths(self.path, self.remainder)

        cls = _get_path_class(child_path)
        self.child = cls(child_path, child_remainder, parent=self)

class FilePathTarDir(FilePath):
    def __init__(self, tarfd, path, remainder, parent=None):
        '''
        This is a directory in a tar archive.
        '''
        self.tarfd = tarfd
        super(FilePathTarDir, self).__init__(path, remainder, parent=parent)

    def _type(self):
        return 'inner.tar.dir'

    def isfile(self):
        return False

    def _process_next(self):
        '''
        This will only be called if the member from the remainder is a file or dir
        '''

        tar_dirs, tar_files = tar_enumerate(self.tarfd)

        child_path, child_remainder = _get_child_paths(self.path, self.remainder)

        child_tar_path = _container_pathnorm(child_path)
        if child_tar_path in tar_dirs:
            self.child = FilePathTarDir(self.tarfd, child_path, child_remainder, parent=self)

        elif child_tar_path in tar_files:
            if not child_remainder:
                self.child = FilePathTarFile(self.tarfd, child_path, child_remainder, parent=self)
                return

            # if the child has a remainder then extract otherwise, don't
            self.tempfd = tempfile.NamedTemporaryFile(prefix='syn_sfp_')

            tfd = self.tarfd.extractfile(child_path.strip('/'))
            tbuf = tfd.read(read_chunk_sz)
            while tbuf:
                self.tempfd.write(tbuf)
                tbuf = tfd.read(read_chunk_sz)
            self.tempfd.flush()
            child_path = self.tempfd.name
            self.child = _get_path_class(child_path)(child_path, child_remainder, parent=self)

        elif self.remainder:
            raise s_exc.NoSuchPath(path=os.path.join(self.path, self.remainder))

class FilePathTarFile(FilePathTarDir):
    def _type(self):
        return 'inner.tar.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return self.tarfd.extractfile(self.path.lstrip('/'))

class FilePathTar(FilePathTarDir):
    def __init__(self, path, remainder, parent=None):
        tarfd = tarfile.open(path)
        self.tarpath = path

        super(FilePathTar, self).__init__(tarfd, '/', remainder, parent=parent)

    def _type(self):
        return 'fs.tar.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return open(self.tarpath, mode=mode)

class FilePathZipDir(FilePath):
    def __init__(self, zipfd, path, remainder, parent=None):
        '''
        This is a directory in a zip.
        '''
        self.zipfd = zipfd
        super(FilePathZipDir, self).__init__(path, remainder, parent=parent)

    def _type(self):
        return 'inner.zip.dir'

    def isfile(self):
        return False

    def _process_next(self):
        '''
        This will only be called if the member from the remainder is a file or dir
        '''

        zip_dirs, zip_files = zip_enumerate(self.zipfd)

        child_path, child_remainder = _get_child_paths(self.path, self.remainder)

        child_zip_path = _container_pathnorm(child_path)
        if child_zip_path in zip_dirs:
            self.child = FilePathZipDir(self.zipfd, child_path, child_remainder, parent=self)

        elif child_zip_path in zip_files:
            if not child_remainder:
                self.child = FilePathZipFile(self.zipfd, child_path, child_remainder, parent=self)
                return

            # extract the file to a directory and parse it
            self.tempfd = tempfile.NamedTemporaryFile(prefix='syn_sfp_')

            zfd = self.zipfd.open(child_path.strip('/'), mode='r')
            zbuf = zfd.read(read_chunk_sz)
            while zbuf:
                self.tempfd.write(zbuf)
                zbuf = zfd.read(read_chunk_sz)

            self.tempfd.flush()
            child_path = self.tempfd.name
            self.child = _get_path_class(child_path)(child_path, child_remainder, parent=self)

        elif self.remainder:
            raise s_exc.NoSuchPath(path=os.path.join(self.path, self.remainder))

class FilePathZipFile(FilePathZipDir):
    def _type(self):
        return 'inner.zip.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return self.zipfd.open(self.path.lstrip('/'), mode='r')

class FilePathZip(FilePathZipDir):
    def __init__(self, path, remainder, parent=None):
        zipfd = zipfile.ZipFile(path)
        self.zippath = path

        super(FilePathZip, self).__init__(zipfd, '/', remainder, parent=parent)

    def _type(self):
        return 'fs.zip.file'

    def isfile(self):
        return True

    def _open(self, mode='r'):
        return open(self.zippath, mode=mode)

def zip_enumerate(zipfd):
    '''
    Return the files and dirs in a zip file
    '''
    files = set()
    dirs = set()

    for path in zipfd.namelist():

        for spath in _get_subpaths(path):
            try:
                info = zipfd.getinfo(spath)
                files.add(spath)
            except KeyError as e:
                dirs.add(spath)

    return dirs, files

def tar_enumerate(tarfd):
    '''
    Return the files and dirs in a tar archive
    '''
    files = set()
    dirs = set()

    for path in tarfd.getnames():

        for spath in _get_subpaths(path):
            info = tarfd.getmember(spath)

            if info.isdir():
                dirs.add(spath)
            if info.isfile():
                files.add(spath)

    return dirs, files

def _get_path_class(path):
    '''
    Returns the class to handle the type of item located at path.  This function
    only operates on regular os. accessible paths
    '''

    if not os.path.exists(path):
        raise s_exc.NoSuchPath(path=path)

    if os.path.isdir(path):
        return _path_ctors.get('fs.reg.dir')
    mime = _mime_file(path)
    return _path_ctors.get(mime)

def _mime_file(path):
    '''
    Assumes the path exists and is a file.
    returns reg.file unless it is a known container
    '''
    if zipfile.is_zipfile(path):
        return 'fs.zip.file'
    if tarfile.is_tarfile(path):
        return 'fs.tar.file'

    return 'fs.reg.file'

def _lsplit(path):
    '''
    Split a pathname.  Returns tuple "(head, tail)" where "tail" is everything after the
    FIRST slash that is not at position 0.  Either part may be empty.

    It's the inverse of os.path.split
    '''

    drive, rpath = os.path.splitdrive(path)
    parts = os.path.normpath(rpath).split(os.path.sep)
    if parts[0] == '':
        parts[0] = os.path.sep
    head = os.path.join(drive, parts[0])

    tail = ''
    if parts[1:]:
        tail = os.path.join(*parts[1:])
    return head, tail

def _get_child_paths(headpath, tailpath):
    '''
    Take a split head/tail path and create a new head tail that is one directory lower
    in the joined path
    '''
    child_part, child_remainder = _lsplit(tailpath)
    child_headpath = os.path.join(headpath, child_part)
    return child_headpath, child_remainder

def _container_pathnorm(path):
    return path.replace('\\', '/').strip('/')

def _get_subpaths(path):
    '''
    Returns a list of subpaths in a path, one for each level in the path
    This is an internal function used for ONLY for iterating over paths in a container
    As such it should not be used for filesystem paths since separators will vary across platforms
    '''
    path = path.replace('\\', '/')
    path = path.rstrip('/')

    path_parts = path.split('/')
    if path_parts[0] == '':
        path_parts[0] = '/'

    paths = []
    for i in range(len(path_parts)):

        ipath = os.path.join(*path_parts[:i+1])
        paths.append(ipath)

    return paths

def _parse_path(*paths):
    '''
    Internal function to parse the incoming path.
    lists of paths are joined prior to parsing
    '''

    if not paths or None in paths:
        return None

    path = genpath(*paths)

    if not path:
        return None

    fpobj = None
    child_path, child_remainder = _lsplit(path)

    try:
        cls = _get_path_class(child_path)
        fpobj = cls(child_path, child_remainder)

        while fpobj.getChild():
            fpobj = fpobj.getChild()
    except s_exc.NoSuchPath as e:
        return None

    return fpobj

# KILL2.7 kwargs instead of mode='r' here because of python2.7
def openfile(*paths, **kwargs):
    '''
    Returns a read-only file-like object even if the path terminates inside a container file.

    If the path is a regular os accessible path mode may be passed through as a keyword argument.
    If the path terminates in a container file mode is ignored.

    If the path does not exist a NoSuchPath exception is raised.

    ex.
    openfile('/foo/bar/baz.egg/path/inside/zip/to/file')
    '''
    fpobj = _parse_path(*paths)
    if not fpobj:
        raise s_exc.NoSuchPath(path='%r' % (paths,))
    if not fpobj.isfile():
        path = genpath(*paths)
        raise s_exc.NoSuchPath(path=path)

    fd = fpobj._open(mode=kwargs.get('mode', 'r'))
    return fd

def isfile(*paths):
    '''
    Determines if the path is a file, even if the path terminates inside a container file.
    If a list of paths are provided, they are joined first.
    Returns a boolean.
    '''
    fpath = _parse_path(*paths)
    if fpath and fpath.isfile():
        return True
    return False

def isdir(*paths):
    '''
    Determines if the path is a directory, even if the path terminates inside a container file.
    If a list of paths are provided, they are joined first.
    Returns a boolean.
    '''
    fpath = _parse_path(*paths)
    if fpath and not fpath.isfile():
        return True
    return False

def exists(*paths):
    '''
    Determines if the path exists even if the path terminates inside a container file.
    If a list of paths are provided, they are joined first.
    Returns a boolean.
    '''
    if _parse_path(*paths) == None:
        return False
    return True

_path_ctors = {
    'fs.reg.dir': FilePathDir,
    'fs.reg.file': FilePathFile,
    'fs.tar.file': FilePathTar,
    'fs.zip.file': FilePathZip,
    'inner.tar.dir': FilePathTarDir,
    'inner.zip.dir': FilePathZipDir,
    'inner.tar.file': FilePathTarFile,
    'inner.zip.file': FilePathZipFile,
}

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2014-2016 OpenMarket Ltd
#
# 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.

from twisted.internet import defer

from synapse.api.errors import StoreError

from ._base import SQLBaseStore
from .engines import PostgresEngine, Sqlite3Engine

import collections
import logging
import ujson as json

logger = logging.getLogger(__name__)

OpsLevel = collections.namedtuple(
    "OpsLevel",
    ("ban_level", "kick_level", "redact_level",)
)

class RoomStore(SQLBaseStore):

    @defer.inlineCallbacks
    def store_room(self, room_id, room_creator_user_id, is_public):
        """Stores a room.

        Args:
            room_id (str): The desired room ID, can be None.
            room_creator_user_id (str): The user ID of the room creator.
            is_public (bool): True to indicate that this room should appear in
            public room lists.
        Raises:
            StoreError if the room could not be stored.
        """
        try:
            def store_room_txn(txn, next_id):
                self._simple_insert_txn(
                    txn,
                    "rooms",
                    {
                        "room_id": room_id,
                        "creator": room_creator_user_id,
                        "is_public": is_public,
                    },
                )
                if is_public:
                    self._simple_insert_txn(
                        txn,
                        table="public_room_list_stream",
                        values={
                            "stream_id": next_id,
                            "room_id": room_id,
                            "visibility": is_public,
                        }
                    )
            with self._public_room_id_gen.get_next() as next_id:
                yield self.runInteraction(
                    "store_room_txn",
                    store_room_txn, next_id,
                )
        except Exception as e:
            logger.error("store_room with room_id=%s failed: %s", room_id, e)
            raise StoreError(500, "Problem creating room.")

    def get_room(self, room_id):
        """Retrieve a room.

        Args:
            room_id (str): The ID of the room to retrieve.
        Returns:
            A namedtuple containing the room information, or an empty list.
        """
        return self._simple_select_one(
            table="rooms",
            keyvalues={"room_id": room_id},
            retcols=("room_id", "is_public", "creator"),
            desc="get_room",
            allow_none=True,
        )

    @defer.inlineCallbacks
    def set_room_is_public(self, room_id, is_public):
        def set_room_is_public_txn(txn, next_id):
            self._simple_update_one_txn(
                txn,
                table="rooms",
                keyvalues={"room_id": room_id},
                updatevalues={"is_public": is_public},
            )

            entries = self._simple_select_list_txn(
                txn,
                table="public_room_list_stream",
                keyvalues={
                    "room_id": room_id,
                    "appservice_id": None,
                    "network_id": None,
                },
                retcols=("stream_id", "visibility"),
            )

            entries.sort(key=lambda r: r["stream_id"])

            add_to_stream = True
            if entries:
                add_to_stream = bool(entries[-1]["visibility"]) != is_public

            if add_to_stream:
                self._simple_insert_txn(
                    txn,
                    table="public_room_list_stream",
                    values={
                        "stream_id": next_id,
                        "room_id": room_id,
                        "visibility": is_public,
                        "appservice_id": None,
                        "network_id": None,
                    }
                )

        with self._public_room_id_gen.get_next() as next_id:
            yield self.runInteraction(
                "set_room_is_public",
                set_room_is_public_txn, next_id,
            )

    @defer.inlineCallbacks
    def set_room_is_public_appservice(self, room_id, appservice_id, network_id,
                                      is_public):
        """Edit the appservice/network specific public room list.

        Each appservice can have a number of published room lists associated
        with them, keyed off of an appservice defined `network_id`, which
        basically represents a single instance of a bridge to a third party
        network.

        Args:
            room_id (str)
            appservice_id (str)
            network_id (str)
            is_public (bool): Whether to publish or unpublish the room from the
                list.
        """
        def set_room_is_public_appservice_txn(txn, next_id):
            if is_public:
                try:
                    self._simple_insert_txn(
                        txn,
                        table="appservice_room_list",
                        values={
                            "appservice_id": appservice_id,
                            "network_id": "network_id",
                            "room_id": room_id
                        },
                    )
                except self.database_engine.module.IntegrityError:
                    # We've already inserted, nothing to do.
                    return
            else:
                self._simple_delete_txn(
                    txn,
                    table="appservice_room_list",
                    keyvalues={
                        "appservice_id": appservice_id,
                        "network_id": network_id,
                        "room_id": room_id
                    },
                )

            entries = self._simple_select_list_txn(
                txn,
                table="public_room_list_stream",
                keyvalues={
                    "room_id": room_id,
                    "appservice_id": appservice_id,
                    "network_id": network_id,
                },
                retcols=("stream_id", "visibility"),
            )

            entries.sort(key=lambda r: r["stream_id"])

            add_to_stream = True
            if entries:
                add_to_stream = bool(entries[-1]["visibility"]) != is_public

            if add_to_stream:
                self._simple_insert_txn(
                    txn,
                    table="public_room_list_stream",
                    values={
                        "stream_id": next_id,
                        "room_id": room_id,
                        "visibility": is_public,
                        "appservice_id": appservice_id,
                        "network_id": network_id,
                    }
                )

        with self._public_room_id_gen.get_next() as next_id:
            yield self.runInteraction(
                "set_room_is_public_appservice",
                set_room_is_public_appservice_txn, next_id,
            )

    def get_public_room_ids(self):
        return self._simple_select_onecol(
            table="rooms",
            keyvalues={
                "is_public": True,
            },
            retcol="room_id",
            desc="get_public_room_ids",
        )

    def get_room_count(self):
        """Retrieve a list of all rooms
        """

        def f(txn):
            sql = "SELECT count(*)  FROM rooms"
            txn.execute(sql)
            row = txn.fetchone()
            return row[0] or 0

        return self.runInteraction(
            "get_rooms", f
        )

    def _store_room_topic_txn(self, txn, event):
        if hasattr(event, "content") and "topic" in event.content:
            self._simple_insert_txn(
                txn,
                "topics",
                {
                    "event_id": event.event_id,
                    "room_id": event.room_id,
                    "topic": event.content["topic"],
                },
            )

            self._store_event_search_txn(
                txn, event, "content.topic", event.content["topic"]
            )

    def _store_room_name_txn(self, txn, event):
        if hasattr(event, "content") and "name" in event.content:
            self._simple_insert_txn(
                txn,
                "room_names",
                {
                    "event_id": event.event_id,
                    "room_id": event.room_id,
                    "name": event.content["name"],
                }
            )

            self._store_event_search_txn(
                txn, event, "content.name", event.content["name"]
            )

    def _store_room_message_txn(self, txn, event):
        if hasattr(event, "content") and "body" in event.content:
            self._store_event_search_txn(
                txn, event, "content.body", event.content["body"]
            )

    def _store_history_visibility_txn(self, txn, event):
        self._store_content_index_txn(txn, event, "history_visibility")

    def _store_guest_access_txn(self, txn, event):
        self._store_content_index_txn(txn, event, "guest_access")

    def _store_content_index_txn(self, txn, event, key):
        if hasattr(event, "content") and key in event.content:
            sql = (
                "INSERT INTO %(key)s"
                " (event_id, room_id, %(key)s)"
                " VALUES (?, ?, ?)" % {"key": key}
            )
            txn.execute(sql, (
                event.event_id,
                event.room_id,
                event.content[key]
            ))

    def _store_event_search_txn(self, txn, event, key, value):
        if isinstance(self.database_engine, PostgresEngine):
            sql = (
                "INSERT INTO event_search"
                " (event_id, room_id, key, vector, stream_ordering, origin_server_ts)"
                " VALUES (?,?,?,to_tsvector('english', ?),?,?)"
            )
            txn.execute(
                sql,
                (
                    event.event_id, event.room_id, key, value,
                    event.internal_metadata.stream_ordering,
                    event.origin_server_ts,
                )
            )
        elif isinstance(self.database_engine, Sqlite3Engine):
            sql = (
                "INSERT INTO event_search (event_id, room_id, key, value)"
                " VALUES (?,?,?,?)"
            )
            txn.execute(sql, (event.event_id, event.room_id, key, value,))
        else:
            # This should be unreachable.
            raise Exception("Unrecognized database engine")

    def add_event_report(self, room_id, event_id, user_id, reason, content,
                         received_ts):
        next_id = self._event_reports_id_gen.get_next()
        return self._simple_insert(
            table="event_reports",
            values={
                "id": next_id,
                "received_ts": received_ts,
                "room_id": room_id,
                "event_id": event_id,
                "user_id": user_id,
                "reason": reason,
                "content": json.dumps(content),
            },
            desc="add_event_report"
        )

    def get_current_public_room_stream_id(self):
        return self._public_room_id_gen.get_current_token()

    def get_public_room_ids_at_stream_id(self, stream_id, network_tuple):
        """Get pulbic rooms for a particular list, or across all lists.

        Args:
            stream_id (int)
            network_tuple (ThirdPartyInstanceID): The list to use (None, None)
                means the main list, None means all lsits.
        """
        return self.runInteraction(
            "get_public_room_ids_at_stream_id",
            self.get_public_room_ids_at_stream_id_txn,
            stream_id, network_tuple=network_tuple
        )

    def get_public_room_ids_at_stream_id_txn(self, txn, stream_id,
                                             network_tuple):
        return {
            rm
            for rm, vis in self.get_published_at_stream_id_txn(
                txn, stream_id, network_tuple=network_tuple
            ).items()
            if vis
        }

    def get_published_at_stream_id_txn(self, txn, stream_id, network_tuple):
        if network_tuple:
            # We want to get from a particular list. No aggregation required.

            sql = ("""
                SELECT room_id, visibility FROM public_room_list_stream
                INNER JOIN (
                    SELECT room_id, max(stream_id) AS stream_id
                    FROM public_room_list_stream
                    WHERE stream_id <= ? %s
                    GROUP BY room_id
                ) grouped USING (room_id, stream_id)
            """)

            if network_tuple.appservice_id is not None:
                txn.execute(
                    sql % ("AND appservice_id = ? AND network_id = ?",),
                    (stream_id, network_tuple.appservice_id, network_tuple.network_id,)
                )
            else:
                txn.execute(
                    sql % ("AND appservice_id IS NULL",),
                    (stream_id,)
                )
            return dict(txn.fetchall())
        else:
            # We want to get from all lists, so we need to aggregate the results

            logger.info("Executing full list")

            sql = ("""
                SELECT room_id, visibility
                FROM public_room_list_stream
                INNER JOIN (
                    SELECT
                        room_id, max(stream_id) AS stream_id, appservice_id,
                        network_id
                    FROM public_room_list_stream
                    WHERE stream_id <= ?
                    GROUP BY room_id, appservice_id, network_id
                ) grouped USING (room_id, stream_id)
            """)

            txn.execute(
                sql,
                (stream_id,)
            )

            results = {}
            # A room is visible if its visible on any list.
            for room_id, visibility in txn.fetchall():
                results[room_id] = bool(visibility) or results.get(room_id, False)

            return results

    def get_public_room_changes(self, prev_stream_id, new_stream_id,
                                network_tuple):
        def get_public_room_changes_txn(txn):
            then_rooms = self.get_public_room_ids_at_stream_id_txn(
                txn, prev_stream_id, network_tuple
            )

            now_rooms_dict = self.get_published_at_stream_id_txn(
                txn, new_stream_id, network_tuple
            )

            now_rooms_visible = set(
                rm for rm, vis in now_rooms_dict.items() if vis
            )
            now_rooms_not_visible = set(
                rm for rm, vis in now_rooms_dict.items() if not vis
            )

            newly_visible = now_rooms_visible - then_rooms
            newly_unpublished = now_rooms_not_visible & then_rooms

            return newly_visible, newly_unpublished

        return self.runInteraction(
            "get_public_room_changes", get_public_room_changes_txn
        )

    def get_all_new_public_rooms(self, prev_id, current_id, limit):
        def get_all_new_public_rooms(txn):
            sql = ("""
                SELECT stream_id, room_id, visibility, appservice_id, network_id
                FROM public_room_list_stream
                WHERE stream_id > ? AND stream_id <= ?
                ORDER BY stream_id ASC
                LIMIT ?
            """)

            txn.execute(sql, (prev_id, current_id, limit,))
            return txn.fetchall()

        if prev_id == current_id:
            return defer.succeed([])

        return self.runInteraction(
            "get_all_new_public_rooms", get_all_new_public_rooms
        )
<commit_msg>Notify replication. Use correct network_id<commit_after># -*- coding: utf-8 -*-
# Copyright 2014-2016 OpenMarket Ltd
#
# 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.

from twisted.internet import defer

from synapse.api.errors import StoreError

from ._base import SQLBaseStore
from .engines import PostgresEngine, Sqlite3Engine

import collections
import logging
import ujson as json

logger = logging.getLogger(__name__)

OpsLevel = collections.namedtuple(
    "OpsLevel",
    ("ban_level", "kick_level", "redact_level",)
)

class RoomStore(SQLBaseStore):

    @defer.inlineCallbacks
    def store_room(self, room_id, room_creator_user_id, is_public):
        """Stores a room.

        Args:
            room_id (str): The desired room ID, can be None.
            room_creator_user_id (str): The user ID of the room creator.
            is_public (bool): True to indicate that this room should appear in
            public room lists.
        Raises:
            StoreError if the room could not be stored.
        """
        try:
            def store_room_txn(txn, next_id):
                self._simple_insert_txn(
                    txn,
                    "rooms",
                    {
                        "room_id": room_id,
                        "creator": room_creator_user_id,
                        "is_public": is_public,
                    },
                )
                if is_public:
                    self._simple_insert_txn(
                        txn,
                        table="public_room_list_stream",
                        values={
                            "stream_id": next_id,
                            "room_id": room_id,
                            "visibility": is_public,
                        }
                    )
            with self._public_room_id_gen.get_next() as next_id:
                yield self.runInteraction(
                    "store_room_txn",
                    store_room_txn, next_id,
                )
        except Exception as e:
            logger.error("store_room with room_id=%s failed: %s", room_id, e)
            raise StoreError(500, "Problem creating room.")

    def get_room(self, room_id):
        """Retrieve a room.

        Args:
            room_id (str): The ID of the room to retrieve.
        Returns:
            A namedtuple containing the room information, or an empty list.
        """
        return self._simple_select_one(
            table="rooms",
            keyvalues={"room_id": room_id},
            retcols=("room_id", "is_public", "creator"),
            desc="get_room",
            allow_none=True,
        )

    @defer.inlineCallbacks
    def set_room_is_public(self, room_id, is_public):
        def set_room_is_public_txn(txn, next_id):
            self._simple_update_one_txn(
                txn,
                table="rooms",
                keyvalues={"room_id": room_id},
                updatevalues={"is_public": is_public},
            )

            entries = self._simple_select_list_txn(
                txn,
                table="public_room_list_stream",
                keyvalues={
                    "room_id": room_id,
                    "appservice_id": None,
                    "network_id": None,
                },
                retcols=("stream_id", "visibility"),
            )

            entries.sort(key=lambda r: r["stream_id"])

            add_to_stream = True
            if entries:
                add_to_stream = bool(entries[-1]["visibility"]) != is_public

            if add_to_stream:
                self._simple_insert_txn(
                    txn,
                    table="public_room_list_stream",
                    values={
                        "stream_id": next_id,
                        "room_id": room_id,
                        "visibility": is_public,
                        "appservice_id": None,
                        "network_id": None,
                    }
                )

        with self._public_room_id_gen.get_next() as next_id:
            yield self.runInteraction(
                "set_room_is_public",
                set_room_is_public_txn, next_id,
            )
        self.hs.get_notifier().on_new_replication_data()

    @defer.inlineCallbacks
    def set_room_is_public_appservice(self, room_id, appservice_id, network_id,
                                      is_public):
        """Edit the appservice/network specific public room list.

        Each appservice can have a number of published room lists associated
        with them, keyed off of an appservice defined `network_id`, which
        basically represents a single instance of a bridge to a third party
        network.

        Args:
            room_id (str)
            appservice_id (str)
            network_id (str)
            is_public (bool): Whether to publish or unpublish the room from the
                list.
        """
        def set_room_is_public_appservice_txn(txn, next_id):
            if is_public:
                try:
                    self._simple_insert_txn(
                        txn,
                        table="appservice_room_list",
                        values={
                            "appservice_id": appservice_id,
                            "network_id": network_id,
                            "room_id": room_id
                        },
                    )
                except self.database_engine.module.IntegrityError:
                    # We've already inserted, nothing to do.
                    return
            else:
                self._simple_delete_txn(
                    txn,
                    table="appservice_room_list",
                    keyvalues={
                        "appservice_id": appservice_id,
                        "network_id": network_id,
                        "room_id": room_id
                    },
                )

            entries = self._simple_select_list_txn(
                txn,
                table="public_room_list_stream",
                keyvalues={
                    "room_id": room_id,
                    "appservice_id": appservice_id,
                    "network_id": network_id,
                },
                retcols=("stream_id", "visibility"),
            )

            entries.sort(key=lambda r: r["stream_id"])

            add_to_stream = True
            if entries:
                add_to_stream = bool(entries[-1]["visibility"]) != is_public

            if add_to_stream:
                self._simple_insert_txn(
                    txn,
                    table="public_room_list_stream",
                    values={
                        "stream_id": next_id,
                        "room_id": room_id,
                        "visibility": is_public,
                        "appservice_id": appservice_id,
                        "network_id": network_id,
                    }
                )

        with self._public_room_id_gen.get_next() as next_id:
            yield self.runInteraction(
                "set_room_is_public_appservice",
                set_room_is_public_appservice_txn, next_id,
            )
        self.hs.get_notifier().on_new_replication_data()

    def get_public_room_ids(self):
        return self._simple_select_onecol(
            table="rooms",
            keyvalues={
                "is_public": True,
            },
            retcol="room_id",
            desc="get_public_room_ids",
        )

    def get_room_count(self):
        """Retrieve a list of all rooms
        """

        def f(txn):
            sql = "SELECT count(*)  FROM rooms"
            txn.execute(sql)
            row = txn.fetchone()
            return row[0] or 0

        return self.runInteraction(
            "get_rooms", f
        )

    def _store_room_topic_txn(self, txn, event):
        if hasattr(event, "content") and "topic" in event.content:
            self._simple_insert_txn(
                txn,
                "topics",
                {
                    "event_id": event.event_id,
                    "room_id": event.room_id,
                    "topic": event.content["topic"],
                },
            )

            self._store_event_search_txn(
                txn, event, "content.topic", event.content["topic"]
            )

    def _store_room_name_txn(self, txn, event):
        if hasattr(event, "content") and "name" in event.content:
            self._simple_insert_txn(
                txn,
                "room_names",
                {
                    "event_id": event.event_id,
                    "room_id": event.room_id,
                    "name": event.content["name"],
                }
            )

            self._store_event_search_txn(
                txn, event, "content.name", event.content["name"]
            )

    def _store_room_message_txn(self, txn, event):
        if hasattr(event, "content") and "body" in event.content:
            self._store_event_search_txn(
                txn, event, "content.body", event.content["body"]
            )

    def _store_history_visibility_txn(self, txn, event):
        self._store_content_index_txn(txn, event, "history_visibility")

    def _store_guest_access_txn(self, txn, event):
        self._store_content_index_txn(txn, event, "guest_access")

    def _store_content_index_txn(self, txn, event, key):
        if hasattr(event, "content") and key in event.content:
            sql = (
                "INSERT INTO %(key)s"
                " (event_id, room_id, %(key)s)"
                " VALUES (?, ?, ?)" % {"key": key}
            )
            txn.execute(sql, (
                event.event_id,
                event.room_id,
                event.content[key]
            ))

    def _store_event_search_txn(self, txn, event, key, value):
        if isinstance(self.database_engine, PostgresEngine):
            sql = (
                "INSERT INTO event_search"
                " (event_id, room_id, key, vector, stream_ordering, origin_server_ts)"
                " VALUES (?,?,?,to_tsvector('english', ?),?,?)"
            )
            txn.execute(
                sql,
                (
                    event.event_id, event.room_id, key, value,
                    event.internal_metadata.stream_ordering,
                    event.origin_server_ts,
                )
            )
        elif isinstance(self.database_engine, Sqlite3Engine):
            sql = (
                "INSERT INTO event_search (event_id, room_id, key, value)"
                " VALUES (?,?,?,?)"
            )
            txn.execute(sql, (event.event_id, event.room_id, key, value,))
        else:
            # This should be unreachable.
            raise Exception("Unrecognized database engine")

    def add_event_report(self, room_id, event_id, user_id, reason, content,
                         received_ts):
        next_id = self._event_reports_id_gen.get_next()
        return self._simple_insert(
            table="event_reports",
            values={
                "id": next_id,
                "received_ts": received_ts,
                "room_id": room_id,
                "event_id": event_id,
                "user_id": user_id,
                "reason": reason,
                "content": json.dumps(content),
            },
            desc="add_event_report"
        )

    def get_current_public_room_stream_id(self):
        return self._public_room_id_gen.get_current_token()

    def get_public_room_ids_at_stream_id(self, stream_id, network_tuple):
        """Get pulbic rooms for a particular list, or across all lists.

        Args:
            stream_id (int)
            network_tuple (ThirdPartyInstanceID): The list to use (None, None)
                means the main list, None means all lsits.
        """
        return self.runInteraction(
            "get_public_room_ids_at_stream_id",
            self.get_public_room_ids_at_stream_id_txn,
            stream_id, network_tuple=network_tuple
        )

    def get_public_room_ids_at_stream_id_txn(self, txn, stream_id,
                                             network_tuple):
        return {
            rm
            for rm, vis in self.get_published_at_stream_id_txn(
                txn, stream_id, network_tuple=network_tuple
            ).items()
            if vis
        }

    def get_published_at_stream_id_txn(self, txn, stream_id, network_tuple):
        if network_tuple:
            # We want to get from a particular list. No aggregation required.

            sql = ("""
                SELECT room_id, visibility FROM public_room_list_stream
                INNER JOIN (
                    SELECT room_id, max(stream_id) AS stream_id
                    FROM public_room_list_stream
                    WHERE stream_id <= ? %s
                    GROUP BY room_id
                ) grouped USING (room_id, stream_id)
            """)

            if network_tuple.appservice_id is not None:
                txn.execute(
                    sql % ("AND appservice_id = ? AND network_id = ?",),
                    (stream_id, network_tuple.appservice_id, network_tuple.network_id,)
                )
            else:
                txn.execute(
                    sql % ("AND appservice_id IS NULL",),
                    (stream_id,)
                )
            return dict(txn.fetchall())
        else:
            # We want to get from all lists, so we need to aggregate the results

            logger.info("Executing full list")

            sql = ("""
                SELECT room_id, visibility
                FROM public_room_list_stream
                INNER JOIN (
                    SELECT
                        room_id, max(stream_id) AS stream_id, appservice_id,
                        network_id
                    FROM public_room_list_stream
                    WHERE stream_id <= ?
                    GROUP BY room_id, appservice_id, network_id
                ) grouped USING (room_id, stream_id)
            """)

            txn.execute(
                sql,
                (stream_id,)
            )

            results = {}
            # A room is visible if its visible on any list.
            for room_id, visibility in txn.fetchall():
                results[room_id] = bool(visibility) or results.get(room_id, False)

            return results

    def get_public_room_changes(self, prev_stream_id, new_stream_id,
                                network_tuple):
        def get_public_room_changes_txn(txn):
            then_rooms = self.get_public_room_ids_at_stream_id_txn(
                txn, prev_stream_id, network_tuple
            )

            now_rooms_dict = self.get_published_at_stream_id_txn(
                txn, new_stream_id, network_tuple
            )

            now_rooms_visible = set(
                rm for rm, vis in now_rooms_dict.items() if vis
            )
            now_rooms_not_visible = set(
                rm for rm, vis in now_rooms_dict.items() if not vis
            )

            newly_visible = now_rooms_visible - then_rooms
            newly_unpublished = now_rooms_not_visible & then_rooms

            return newly_visible, newly_unpublished

        return self.runInteraction(
            "get_public_room_changes", get_public_room_changes_txn
        )

    def get_all_new_public_rooms(self, prev_id, current_id, limit):
        def get_all_new_public_rooms(txn):
            sql = ("""
                SELECT stream_id, room_id, visibility, appservice_id, network_id
                FROM public_room_list_stream
                WHERE stream_id > ? AND stream_id <= ?
                ORDER BY stream_id ASC
                LIMIT ?
            """)

            txn.execute(sql, (prev_id, current_id, limit,))
            return txn.fetchall()

        if prev_id == current_id:
            return defer.succeed([])

        return self.runInteraction(
            "get_all_new_public_rooms", get_all_new_public_rooms
        )
<|endoftext|>


<commit_before>"""
Utilities to manage AWS Elastic Block Store volumes and snapshots.

To delete EBS volumes or snapshots, use ``aegea rm``.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, argparse, getpass
from datetime import datetime

from . import register_parser
from .ls import add_name, filter_collection, filter_and_tabulate, register_filtering_parser
from .util import Timestamp, paginate
from .util.printing import format_table, page_output, get_field, get_cell, tabulate
from .util.aws import ARN, resources, clients, ensure_vpc, ensure_subnet, resolve_instance_id, add_tags
from .util.compat import lru_cache

def ebs(args):
    ebs_parser.print_help()

ebs_parser = register_parser(ebs, help="Manage Elastic Block Store resources", description=__doc__,
                             formatter_class=argparse.RawTextHelpFormatter)

def ls(args):
    @lru_cache()
    def instance_id_to_name(i):
        return add_name(resources.ec2.Instance(i)).name
    table = [{f: get_cell(i, f) for f in args.columns} for i in filter_collection(resources.ec2.volumes, args)]
    if "attachments" in args.columns:
        for row in table:
            row["attachments"] = ", ".join(instance_id_to_name(a["InstanceId"]) for a in row["attachments"])
    page_output(tabulate(table, args))

parser = register_filtering_parser(ls, parent=ebs_parser, help="List EC2 EBS volumes")

def snapshots(args):
    page_output(filter_and_tabulate(resources.ec2.snapshots.filter(OwnerIds=[ARN.get_account_id()]), args))

parser = register_filtering_parser(snapshots, parent=ebs_parser, help="List EC2 EBS snapshots")

def create(args):
    tags = dict(tag.split("=", 1) for tag in args.tags)
    create_args = dict(Size=args.size)
    for arg in "dry_run snapshot_id availability_zone volume_type iops encrypted kms_key_id".split():
        if getattr(args, arg) is not None:
            create_args["".join(x.capitalize() for x in arg.split("_"))] = getattr(args, arg)
    if "AvailabilityZone" not in create_args:
        create_args["AvailabilityZone"] = ensure_subnet(ensure_vpc()).availability_zone
    res = clients.ec2.create_volume(**create_args)
    clients.ec2.get_waiter('volume_available').wait(VolumeIds=[res["VolumeId"]])
    if tags:
        add_tags(resources.ec2.Volume(res["VolumeId"]), **tags)
    return res

parser = register_parser(create, parent=ebs_parser, help="Create an EBS volume")
parser.add_argument("--size-gb", dest="size", type=int, help="Volume size in gigabytes", required=True)
parser.add_argument("--dry-run", action="store_true")
parser.add_argument("--snapshot-id")
parser.add_argument("--availability-zone")
parser.add_argument("--volume-type", choices={"standard", "io1", "gp2", "sc1", "st1"},
                    help="io1, PIOPS SSD; gp2, general purpose SSD; sc1, cold HDD; st1, throughput optimized HDD")
parser.add_argument("--storage", type=int)
parser.add_argument("--iops", type=int)
parser.add_argument("--encrypted", action="store_true")
parser.add_argument("--kms-key-id")
parser.add_argument("--tags", nargs="+", default=[], metavar="TAG_NAME=VALUE")

def snapshot(args):
    return clients.ec2.create_snapshot(DryRun=args.dry_run, VolumeId=args.volume_id)
parser_snapshot = register_parser(snapshot, parent=ebs_parser, help="Create an EBS snapshot")

def attach(args):
    res = clients.ec2.attach_volume(DryRun=args.dry_run,
                                    VolumeId=args.volume_id,
                                    InstanceId=args.instance,
                                    Device=args.device)
    clients.ec2.get_waiter('volume_in_use').wait(VolumeIds=[res["VolumeId"]])
    return res
parser_attach = register_parser(attach, parent=ebs_parser, help="Attach an EBS volume to an EC2 instance")

def detach(args):
    res = clients.ec2.detach_volume(DryRun=args.dry_run,
                                    VolumeId=args.volume_id,
                                    InstanceId=args.instance,
                                    Device=args.device,
                                    Force=args.force)
    clients.ec2.get_waiter('volume_available').wait(VolumeIds=[res["VolumeId"]])
    return res
parser_detach = register_parser(detach, parent=ebs_parser, help="Detach an EBS volume from an EC2 instance")

def complete_volume_id(**kwargs):
    return [i["VolumeId"] for i in clients.ec2.describe_volumes()["Volumes"]]

for parser in parser_snapshot, parser_attach, parser_detach:
    parser.add_argument("volume_id").completer = complete_volume_id
    parser.add_argument("--dry-run", action="store_true")
    if parser in (parser_attach, parser_detach):
        parser.add_argument("instance", type=resolve_instance_id)
        parser.add_argument("device", choices=["xvd" + chr(i + 1) for i in range(ord("a"), ord("z"))])

parser_detach.add_argument("--force", action="store_true")
<commit_msg>Add aegea ebs modify<commit_after>"""
Utilities to manage AWS Elastic Block Store volumes and snapshots.

To delete EBS volumes or snapshots, use ``aegea rm``.
"""

from __future__ import absolute_import, division, print_function, unicode_literals

import os, sys, argparse, getpass
from datetime import datetime

from . import register_parser
from .ls import add_name, filter_collection, filter_and_tabulate, register_filtering_parser
from .util import Timestamp, paginate
from .util.printing import format_table, page_output, get_field, get_cell, tabulate
from .util.aws import ARN, resources, clients, ensure_vpc, ensure_subnet, resolve_instance_id, add_tags
from .util.compat import lru_cache

def ebs(args):
    ebs_parser.print_help()

ebs_parser = register_parser(ebs, help="Manage Elastic Block Store resources", description=__doc__,
                             formatter_class=argparse.RawTextHelpFormatter)

def ls(args):
    @lru_cache()
    def instance_id_to_name(i):
        return add_name(resources.ec2.Instance(i)).name
    table = [{f: get_cell(i, f) for f in args.columns} for i in filter_collection(resources.ec2.volumes, args)]
    if "attachments" in args.columns:
        for row in table:
            row["attachments"] = ", ".join(instance_id_to_name(a["InstanceId"]) for a in row["attachments"])
    page_output(tabulate(table, args))

parser = register_filtering_parser(ls, parent=ebs_parser, help="List EC2 EBS volumes")

def snapshots(args):
    page_output(filter_and_tabulate(resources.ec2.snapshots.filter(OwnerIds=[ARN.get_account_id()]), args))

parser = register_filtering_parser(snapshots, parent=ebs_parser, help="List EC2 EBS snapshots")

def create(args):
    tags = dict(tag.split("=", 1) for tag in args.tags)
    create_args = dict(Size=args.size)
    for arg in "dry_run snapshot_id availability_zone volume_type iops encrypted kms_key_id".split():
        if getattr(args, arg) is not None:
            create_args["".join(x.capitalize() for x in arg.split("_"))] = getattr(args, arg)
    if "AvailabilityZone" not in create_args:
        create_args["AvailabilityZone"] = ensure_subnet(ensure_vpc()).availability_zone
    res = clients.ec2.create_volume(**create_args)
    clients.ec2.get_waiter('volume_available').wait(VolumeIds=[res["VolumeId"]])
    if tags:
        add_tags(resources.ec2.Volume(res["VolumeId"]), **tags)
    return res

parser_create = register_parser(create, parent=ebs_parser, help="Create an EBS volume")
parser_create.add_argument("--dry-run", action="store_true")
parser_create.add_argument("--snapshot-id")
parser_create.add_argument("--availability-zone")
parser_create.add_argument("--encrypted", action="store_true")
parser_create.add_argument("--kms-key-id")
parser_create.add_argument("--tags", nargs="+", default=[], metavar="TAG_NAME=VALUE")

def snapshot(args):
    return clients.ec2.create_snapshot(DryRun=args.dry_run, VolumeId=args.volume_id)
parser_snapshot = register_parser(snapshot, parent=ebs_parser, help="Create an EBS snapshot")

def attach(args):
    res = clients.ec2.attach_volume(DryRun=args.dry_run,
                                    VolumeId=args.volume_id,
                                    InstanceId=args.instance,
                                    Device=args.device)
    clients.ec2.get_waiter('volume_in_use').wait(VolumeIds=[res["VolumeId"]])
    return res
parser_attach = register_parser(attach, parent=ebs_parser, help="Attach an EBS volume to an EC2 instance")

def detach(args):
    res = clients.ec2.detach_volume(DryRun=args.dry_run,
                                    VolumeId=args.volume_id,
                                    InstanceId=args.instance,
                                    Device=args.device,
                                    Force=args.force)
    clients.ec2.get_waiter('volume_available').wait(VolumeIds=[res["VolumeId"]])
    return res
parser_detach = register_parser(detach, parent=ebs_parser, help="Detach an EBS volume from an EC2 instance")

def modify(args):
    modify_args = dict(VolumeId=args.volume_id, DryRun=args.dry_run)
    if args.size:
        modify_args.update(Size=args.size)
    if args.volume_type:
        modify_args.update(VolumeType=args.volume_type)
    if args.iops:
        modify_args.update(Iops=args.iops)
    res = clients.ec2.modify_volume(**modify_args)["VolumeModification"]
    #if args.wait:
    #    waiter = make_waiter(clients.ec2.describe_volumes_modifications, "VolumesModifications[].ModificationState",
    #                         "optimizing", "pathAny")
    #    waiter.wait(VolumeIds=[args.volume_id])
    return res
parser_modify = register_parser(modify, parent=ebs_parser, help="Change the size, type, or IOPS of an EBS volume")

for parser in parser_create, parser_modify:
    parser.add_argument("--size-gb", dest="size", type=int, help="Volume size in gigabytes")
    parser.add_argument("--volume-type", choices={"standard", "io1", "gp2", "sc1", "st1"},
                        help="io1, PIOPS SSD; gp2, general purpose SSD; sc1, cold HDD; st1, throughput optimized HDD")
    parser.add_argument("--iops", type=int)

def complete_volume_id(**kwargs):
    return [i["VolumeId"] for i in clients.ec2.describe_volumes()["Volumes"]]

for parser in parser_snapshot, parser_attach, parser_detach, parser_modify:
    parser.add_argument("volume_id").completer = complete_volume_id
    parser.add_argument("--dry-run", action="store_true")
    if parser in (parser_attach, parser_detach):
        parser.add_argument("instance", type=resolve_instance_id)
        parser.add_argument("device", choices=["xvd" + chr(i + 1) for i in range(ord("a"), ord("z"))])

parser_detach.add_argument("--force", action="store_true")
<|endoftext|>


<commit_before>import os
import shutil
import logging
import tempfile
import unittest
import threading
import contextlib

logging.basicConfig(level=logging.WARNING)

import synapse.cortex as s_cortex

from synapse.eventbus import Waiter
import synapse.lib.output as s_output
import synapse.lib.thishost as s_thishost

from synapse.common import *

class TooFewEvents(Exception):pass

class TestEnv:

    def __init__(self):
        self.items = {}
        self.tofini = []

    def __getattr__(self, prop):
        item = self.items.get(prop)
        if item == None:
            raise AttributeError(prop)
        return item

    @contextlib.contextmanager
    def getWithBlock(self):
        yield self
        self.fini()

    def add(self, name, item, fini=False):
        self.items[name] = item
        if fini:
            self.tofini.append(item)

    def fini(self):
        for bus in self.tofini:
            bus.fini()

class TestOutPut(s_output.OutPutStr):

    def expect(self, substr):
        outs = str(self)
        if outs.find(substr) == -1:
            raise Exception('TestOutPut.expect(%s) not in %s' % (substr,outs))

class SynTest(unittest.TestCase):

    def getTestWait(self, bus, size, *evts):
        return Waiter(bus, size, *evts)

    def getPgCore(self):
        url = os.getenv('SYN_TEST_PG_URL')
        if url != None:
            return s_cortex.openurl(url)

        db = os.getenv('SYN_TEST_PG_DB')
        if db == None:
            raise unittest.SkipTest('no SYN_TEST_PG_DB or SYN_TEST_PG_URL')

        table = 'syn_test_%s' % guid()

        core = s_cortex.openurl('postgres:///%s/%s' % (db,table))

        def droptable():
            with core.cursor() as c:
                c.execute('DROP TABLE %s' % (table,))

        core.onfini(droptable)
        return core

    def getTestOutp(self):
        return TestOutPut()

    def thisHostMust(self, **props):
        for k,v in props.items():
            if s_thishost.get(k) != v:
                raise unittest.SkipTest('skip thishost: %s!=%r' % (k,v))

    def thisHostMustNot(self, **props):
        for k,v in props.items():
            if s_thishost.get(k) == v:
                raise unittest.SkipTest('skip thishost: %s==%r' % (k,v))

    @contextlib.contextmanager
    def getTestDir(self):
        tempdir = tempfile.mkdtemp()
        yield tempdir
        shutil.rmtree(tempdir)

    def eq(self, x, y):
        self.assertEqual(x,y)

    def true(self, x):
        self.assertTrue(x)

    def false(self, x):
        self.assertFalse(x)

    def nn(self, x):
        self.assertIsNotNone(x)

    def sorteq(self, x, y):
        return self.eq( sorted(x), sorted(y) )

testdir = os.path.dirname(__file__)
def getTestPath(*paths):
    return os.path.join(testdir,*paths)
<commit_msg>ignore errors when removing temporary files from unit tests<commit_after>import os
import shutil
import logging
import tempfile
import unittest
import threading
import contextlib

logging.basicConfig(level=logging.WARNING)

import synapse.cortex as s_cortex
import synapse.eventbus as s_eventbus

import synapse.lib.output as s_output
import synapse.lib.thishost as s_thishost

from synapse.common import *

class TooFewEvents(Exception):pass

class TestEnv:

    def __init__(self):
        self.items = {}
        self.tofini = []

    def __getattr__(self, prop):
        item = self.items.get(prop)
        if item == None:
            raise AttributeError(prop)
        return item

    def __enter__(self):
        return self

    def __exit__(self, cls, exc, tb):
        self.fini()

    def add(self, name, item, fini=False):
        self.items[name] = item
        if fini:
            self.tofini.append(item)

    def fini(self):
        for bus in self.tofini:
            bus.fini()

class TestOutPut(s_output.OutPutStr):

    def expect(self, substr):
        outs = str(self)
        if outs.find(substr) == -1:
            raise Exception('TestOutPut.expect(%s) not in %s' % (substr,outs))

class SynTest(unittest.TestCase):

    def getTestWait(self, bus, size, *evts):
        return s_eventbus.Waiter(bus, size, *evts)

    def getPgCore(self):
        url = os.getenv('SYN_TEST_PG_URL')
        if url != None:
            return s_cortex.openurl(url)

        db = os.getenv('SYN_TEST_PG_DB')
        if db == None:
            raise unittest.SkipTest('no SYN_TEST_PG_DB or SYN_TEST_PG_URL')

        table = 'syn_test_%s' % guid()

        core = s_cortex.openurl('postgres:///%s/%s' % (db,table))

        def droptable():
            with core.cursor() as c:
                c.execute('DROP TABLE %s' % (table,))

        core.onfini(droptable)
        return core

    def getTestOutp(self):
        return TestOutPut()

    def thisHostMust(self, **props):
        for k,v in props.items():
            if s_thishost.get(k) != v:
                raise unittest.SkipTest('skip thishost: %s!=%r' % (k,v))

    def thisHostMustNot(self, **props):
        for k,v in props.items():
            if s_thishost.get(k) == v:
                raise unittest.SkipTest('skip thishost: %s==%r' % (k,v))

    @contextlib.contextmanager
    def getTestDir(self):
        tempdir = tempfile.mkdtemp()
        yield tempdir
        shutil.rmtree(tempdir, ignore_errors=True)

    def eq(self, x, y):
        self.assertEqual(x,y)

    def true(self, x):
        self.assertTrue(x)

    def false(self, x):
        self.assertFalse(x)

    def nn(self, x):
        self.assertIsNotNone(x)

    def sorteq(self, x, y):
        return self.eq( sorted(x), sorted(y) )

testdir = os.path.dirname(__file__)
def getTestPath(*paths):
    return os.path.join(testdir,*paths)
<|endoftext|>


<commit_before>import os
import logging
from itertools import product,chain

import pandas as pd
import numpy as np

import util, data
from util import merge_dicts, hash_obj

# the series can be:
#    - a callable (taking the DataFrame), e.g. lambda df: df.column**2
#    - one of the columns in the frame, e.g. 'column'
#    - a non-string value, e.g. 1
def get_series(series, df, astype=None):
    if hasattr(series, '__call__'):
        r = series(df)
    elif series in df.columns:
        r = df[series]
    elif not isinstance(series, basestring):
        r = pd.Series(series, index=df.index)
    else:
        raise ValueError('Invalid series: %s' % series)

    if astype is not None:
        r = r.astype(astype)
    return r

# AggregateSeries consist of a series and a function
class AggregateSeries(object):
    def __init__(self, series, function, astype=None):
        self.series = series
        self.function = function

        # default for pandas keywords is float32
        # e.g. 'mean', 'min', 'sum', etc.
        # much faster!
        if astype is None and isinstance(function, basestring):
            astype = np.float32
        self.astype = astype

    def apply_series(self, df):
        return get_series(self.series, df, self.astype)

    def __hash__(self):
        s = hash_obj(self.series)
        f = hash_obj(self.function)

        return hash((s,f, self.astype))

class AggregateBase(object):
    def __init__(self, columns, aggregate_series):
        self.columns = columns
        self.aggregate_series = aggregate_series

    # default is that series and columns are one-to-one
    def apply(self, series):
        return series

    def __div__(self, other):
        f = Fraction(numerator=self, denominator=other)
        return f

class Fraction(AggregateBase):
    def __init__(self, numerator, denominator, name='{numerator}_per_{denominator}',
            include_numerator=False, include_denominator=False, include_fraction=True):
        self.numerator = numerator
        self.denominator = denominator

        self.include_numerator=include_numerator
        self.include_denominator=include_denominator
        self.include_fraction=include_fraction

        columns = []
        if include_fraction:
            if hasattr(name, '__iter__') and len(name) == \
                    len(numerator.columns)*len(denominator.columns):
                columns.extend(name)
            elif isinstance(name, basestring):
                columns.extend([name.format(numerator=n, denominator=d)
                    for n,d in product(numerator.columns, denominator.columns)])
            else:
                raise ValueError('Name must either be a list of names of the correct length or a format string')
        if include_numerator:
            columns.extend(numerator.columns)
        if include_denominator:
            columns.extend(denominator.columns)

        aggregate_series = []
        if include_fraction:
            aggregate_series += numerator.aggregate_series + denominator.aggregate_series
        else:
            if include_numerator:
                aggregate_series += numerator.aggregate_series
            if include_denominator:
                aggregate_series += denominator.aggregate_series

        AggregateBase.__init__(self, columns=columns, aggregate_series=aggregate_series)

    def apply(self, series):
        if self.include_fraction or (self.include_numerator and self.include_denominator):
            ncolumns = self.numerator.apply(series[:len(self.numerator.aggregate_series)])
            dcolumns = self.denominator.apply(series[len(self.numerator.aggregate_series):])
        elif self.include_numerator:
            ncolumns = self.numerator.apply(series)
        elif self.include_denominator:
            dcolumns = self.denominator.apply(series)

        columns = []
        if self.include_fraction:
            columns = [n / d for n,d in product(ncolumns, dcolumns)]
        if self.include_numerator:
            columns.extend(ncolumns)
        if self.include_denominator:
            columns.extend(dcolumns)

        return columns

class Aggregate(AggregateBase):
    def __init__(self, series, f, name=None, fname=None):
        """
        series can be single or iterable
        functions can be a single function or iterable
        total aggregates are the products of series and functions
        default name is str(series), default function name is str(f)
        column names are {name}_{fname}
        unless single function and function_names=False, then column name is just name
        """
        series = util.make_list(series)
        f = util.make_list(f)

        name = series if name is None else util.make_list(name)
        fname = f if fname is None else util.make_list(fname)

        if not (len(series) == len(name)):
            raise ValueError('series and name must have same length or name must be None')
        if not (len(series) == len(name)):
            raise ValueError('f and fname must have same length or fname must be None')

        if fname == [False]:
            if len(f) > 1:
                raise ValueError('Must use function names for multiple functions')
            columns = name
        else:
            columns = ['%s_%s' % n for n in product(name, fname)]

        AggregateBase.__init__(self, columns,
                [AggregateSeries(*sf) for sf in product(series, f)])

class Count(Fraction):
    """
    Aggregate a given series (as accepted by get_series above) by summing
    By default that series is 1, resulting in a count (hence the name).
    Also useful when series is a boolean.
    If prop=True then also divide that series by a specified parent series, also summed.
    """
    def __init__(self, series=None, name=None, parent=None, parent_name=None, prop=False, prop_only=False):
        fname = 'count'
        if series is None:
            series = 1
            if name is None:
                name = 'count'
                fname = False # don't use fname when series is 'count'

        if parent is not None or prop_only:
            prop=True

        numerator = Aggregate(series, 'sum', name, fname=fname)
        if not prop:
            Fraction.__init__(self, numerator=numerator,
                    denominator=None, include_fraction=False,
                    include_numerator=True)
        else:
            if parent is None:
                parent = 1

            denominator = Aggregate(parent, 'sum', parent_name)
            Fraction.__init__(self, numerator=numerator,
                    denominator=denominator,
                    include_numerator=not prop_only,
                    # fraction names are numerator names with
                    # 'count' replaced by 'prop'
                    name=[n[:-5]+'prop' for n in numerator.columns])

class Proportion(Count):
    def __init__(self, series=None, name=None, parent=None, parent_name=None):
        Count.__init__(self, series=series, name=name, parent=parent, parent_name=parent_name, prop_only=True)

def _collect_columns(aggregates):
    columns = set()
    for a in aggregates:
        intersection = columns.intersection(a.columns)
        if len(intersection) > 0: raise ValueError('Repeated columns: %s' % intersection)
        columns.update(a.columns)

    if len(columns) == 0:
        raise ValueError('Aggregator needs at least one output columns')

    return columns

class Aggregator(object):
    def __init__(self, df, aggregates):
        self.df = df
        self.aggregates = aggregates

        self.columns = _collect_columns(aggregates)

        # collect all the series
        series = {} # unique series
        aggregate_series = {} # references to those series, one for each (series, function)
        series_refs = [] # series_refs[i] contains references to aggregate_series for aggregates[i]
        for a in aggregates:
            sr = []
            for aseries in a.aggregate_series:
                h = hash_obj(aseries.series)
                i = hash(aseries)
                if h not in series: # only apply each series once
                    series[h] = aseries.apply_series(df)
                    aggregate_series[h] = {i: aseries.function}
                if i not in aggregate_series[h]: # collect each aseries only once
                    aggregate_series[h][i] = aseries.function

                sr.append(i)
            series_refs.append(sr)

        self.adf = pd.DataFrame(series)
        self.aggregate_series = aggregate_series
        self.series_refs = series_refs

    # all the aggregates for a single series
    # returns a hash(aseries): series dict
    def get_series_aggregates(self, groupby, h):
        a = groupby[h].agg(self.aggregate_series[h])
        return {i:a[i] for i in a.columns}

    def aggregate(self, index):
        if isinstance(index, basestring):
            groupby = self.adf.groupby(self.df[index])
        else:
            # groupby does not accept a dataframe so we have to set the index then put it back
            self.adf.index = pd.MultiIndex.from_arrays([self.df[i] for i in index])
            groupby = self.adf.groupby(level=index)
            self.adf.index = self.df.index

        aggregated = merge_dicts(*[self.get_series_aggregates(groupby, h) for h in self.aggregate_series])

        all_series = []
        for a,r in zip(self.aggregates, self.series_refs):
            series = a.apply([aggregated[i].copy() for i in r])
            # rename series according to aggregate.columns
            for s,c in zip(series, a.columns): s.name = c
            all_series.extend(series)

        return pd.concat(all_series, axis=1)

def aggregate_list(l):
    return list(np.concatenate(l.values))

def aggregate_set(l):
    return set(np.concatenate(l.values))

def aggregate_counts(l):
    lists = [list(i) for i in l.values if len(i) > 0]
    if len(lists) == 0:
        return None
    else:
        l = np.concatenate(lists)
        return np.unique(l, return_counts=True)

def days(date1, date2):
    """
    returns a lambda that determines the number of days between the two dates
    the dates can be strings (column names) or actual dates
    e.g. Aggregate(days('date', today), ['min','max'], 'days_since')
    TODO: should there be a Days(AggregateBase) for this? handle naming
    """
    if isinstance(date1, basestring) and isinstance(date2, basestring):
        return lambda d: (d[date2] - d[date1])/util.day
    elif isinstance(date1, basestring):
        return lambda d: (date2 - d[date1])/util.day
    elif isinstance(date2, basestring):
        return lambda d: (d[date2] - date1)/util.day
    else:
        return (date2 - date1)/util.day

def date_min(d):
    """
    groupby()['date_colum'].aggregate('min') returns a float?
    convert it back to a timestamp
    """
    return pd.to_datetime(d.min())

def date_max(d):
    return pd.to_datetime(d.max())
<commit_msg>fix upcast<commit_after>import os
import logging
from itertools import product,chain

import pandas as pd
import numpy as np

import util, data
from util import merge_dicts, hash_obj

# the series can be:
#    - a callable (taking the DataFrame), e.g. lambda df: df.column**2
#    - one of the columns in the frame, e.g. 'column'
#    - a non-string value, e.g. 1
def get_series(series, df, astype=None):
    if hasattr(series, '__call__'):
        r = series(df)
    elif series in df.columns:
        r = df[series]
    elif not isinstance(series, basestring):
        r = pd.Series(series, index=df.index)
    else:
        raise ValueError('Invalid series: %s' % series)

    if astype is not None:
        r = r.astype(astype)
    return r

# AggregateSeries consist of a series and a function
class AggregateSeries(object):
    def __init__(self, series, function, astype=None):
        self.series = series
        self.function = function

        # default for pandas keywords is float32
        # e.g. 'mean', 'min', 'sum', etc.
        # much faster!
        if astype is None and isinstance(function, basestring):
            astype = np.float32
        self.astype = astype

    def apply_series(self, df):
        return get_series(self.series, df, self.astype)

    def __hash__(self):
        s = hash_obj(self.series)
        f = hash_obj(self.function)

        return hash((s,f, self.astype))

class AggregateBase(object):
    def __init__(self, columns, aggregate_series):
        self.columns = columns
        self.aggregate_series = aggregate_series

    # default is that series and columns are one-to-one
    def apply(self, series):
        return series

    def __div__(self, other):
        f = Fraction(numerator=self, denominator=other)
        return f

class Fraction(AggregateBase):
    def __init__(self, numerator, denominator, name='{numerator}_per_{denominator}',
            include_numerator=False, include_denominator=False, include_fraction=True):
        self.numerator = numerator
        self.denominator = denominator

        self.include_numerator=include_numerator
        self.include_denominator=include_denominator
        self.include_fraction=include_fraction

        columns = []
        if include_fraction:
            if hasattr(name, '__iter__') and len(name) == \
                    len(numerator.columns)*len(denominator.columns):
                columns.extend(name)
            elif isinstance(name, basestring):
                columns.extend([name.format(numerator=n, denominator=d)
                    for n,d in product(numerator.columns, denominator.columns)])
            else:
                raise ValueError('Name must either be a list of names of the correct length or a format string')
        if include_numerator:
            columns.extend(numerator.columns)
        if include_denominator:
            columns.extend(denominator.columns)

        aggregate_series = []
        if include_fraction:
            aggregate_series += numerator.aggregate_series + denominator.aggregate_series
        else:
            if include_numerator:
                aggregate_series += numerator.aggregate_series
            if include_denominator:
                aggregate_series += denominator.aggregate_series

        AggregateBase.__init__(self, columns=columns, aggregate_series=aggregate_series)

    def apply(self, series):
        if self.include_fraction or (self.include_numerator and self.include_denominator):
            ncolumns = self.numerator.apply(series[:len(self.numerator.aggregate_series)])
            dcolumns = self.denominator.apply(series[len(self.numerator.aggregate_series):])
        elif self.include_numerator:
            ncolumns = self.numerator.apply(series)
        elif self.include_denominator:
            dcolumns = self.denominator.apply(series)

        columns = []
        if self.include_fraction:
            # sometimes n / d will upcast to float64 so downcast it back to float32
            # this should really be handled by Aggregate(dtype=np.float32) in the future
            columns = [(n / d).astype(np.float32) for n,d in product(ncolumns, dcolumns)]
        if self.include_numerator:
            columns.extend(ncolumns)
        if self.include_denominator:
            columns.extend(dcolumns)

        return columns

class Aggregate(AggregateBase):
    def __init__(self, series, f, name=None, fname=None):
        """
        series can be single or iterable
        functions can be a single function or iterable
        total aggregates are the products of series and functions
        default name is str(series), default function name is str(f)
        column names are {name}_{fname}
        unless single function and function_names=False, then column name is just name
        """
        series = util.make_list(series)
        f = util.make_list(f)

        name = series if name is None else util.make_list(name)
        fname = f if fname is None else util.make_list(fname)

        if not (len(series) == len(name)):
            raise ValueError('series and name must have same length or name must be None')
        if not (len(series) == len(name)):
            raise ValueError('f and fname must have same length or fname must be None')

        if fname == [False]:
            if len(f) > 1:
                raise ValueError('Must use function names for multiple functions')
            columns = name
        else:
            columns = ['%s_%s' % n for n in product(name, fname)]

        AggregateBase.__init__(self, columns,
                [AggregateSeries(*sf) for sf in product(series, f)])

class Count(Fraction):
    """
    Aggregate a given series (as accepted by get_series above) by summing
    By default that series is 1, resulting in a count (hence the name).
    Also useful when series is a boolean.
    If prop=True then also divide that series by a specified parent series, also summed.
    """
    def __init__(self, series=None, name=None, parent=None, parent_name=None, prop=False, prop_only=False):
        fname = 'count'
        if series is None:
            series = 1
            if name is None:
                name = 'count'
                fname = False # don't use fname when series is 'count'

        if parent is not None or prop_only:
            prop=True

        numerator = Aggregate(series, 'sum', name, fname=fname)
        if not prop:
            Fraction.__init__(self, numerator=numerator,
                    denominator=None, include_fraction=False,
                    include_numerator=True)
        else:
            if parent is None:
                parent = 1

            denominator = Aggregate(parent, 'sum', parent_name)
            Fraction.__init__(self, numerator=numerator,
                    denominator=denominator,
                    include_numerator=not prop_only,
                    # fraction names are numerator names with
                    # 'count' replaced by 'prop'
                    name=[n[:-5]+'prop' for n in numerator.columns])

class Proportion(Count):
    def __init__(self, series=None, name=None, parent=None, parent_name=None):
        Count.__init__(self, series=series, name=name, parent=parent, parent_name=parent_name, prop_only=True)

def _collect_columns(aggregates):
    columns = set()
    for a in aggregates:
        intersection = columns.intersection(a.columns)
        if len(intersection) > 0: raise ValueError('Repeated columns: %s' % intersection)
        columns.update(a.columns)

    if len(columns) == 0:
        raise ValueError('Aggregator needs at least one output columns')

    return columns

class Aggregator(object):
    def __init__(self, df, aggregates):
        self.df = df
        self.aggregates = aggregates

        self.columns = _collect_columns(aggregates)

        # collect all the series
        series = {} # unique series
        aggregate_series = {} # references to those series, one for each (series, function)
        series_refs = [] # series_refs[i] contains references to aggregate_series for aggregates[i]
        for a in aggregates:
            sr = []
            for aseries in a.aggregate_series:
                h = hash_obj(aseries.series)
                i = hash(aseries)
                if h not in series: # only apply each series once
                    series[h] = aseries.apply_series(df)
                    aggregate_series[h] = {i: aseries.function}
                if i not in aggregate_series[h]: # collect each aseries only once
                    aggregate_series[h][i] = aseries.function

                sr.append(i)
            series_refs.append(sr)

        self.adf = pd.DataFrame(series)
        self.aggregate_series = aggregate_series
        self.series_refs = series_refs

    # all the aggregates for a single series
    # returns a hash(aseries): series dict
    def get_series_aggregates(self, groupby, h):
        a = groupby[h].agg(self.aggregate_series[h])
        return {i:a[i] for i in a.columns}

    def aggregate(self, index):
        if isinstance(index, basestring):
            groupby = self.adf.groupby(self.df[index])
        else:
            # groupby does not accept a dataframe so we have to set the index then put it back
            self.adf.index = pd.MultiIndex.from_arrays([self.df[i] for i in index])
            groupby = self.adf.groupby(level=index)
            self.adf.index = self.df.index

        aggregated = merge_dicts(*[self.get_series_aggregates(groupby, h) for h in self.aggregate_series])

        all_series = []
        for a,r in zip(self.aggregates, self.series_refs):
            series = a.apply([aggregated[i].copy() for i in r])
            # rename series according to aggregate.columns
            for s,c in zip(series, a.columns): s.name = c
            all_series.extend(series)

        return pd.concat(all_series, axis=1)

def aggregate_list(l):
    return list(np.concatenate(l.values))

def aggregate_set(l):
    return set(np.concatenate(l.values))

def aggregate_counts(l):
    lists = [list(i) for i in l.values if len(i) > 0]
    if len(lists) == 0:
        return None
    else:
        l = np.concatenate(lists)
        return np.unique(l, return_counts=True)

def days(date1, date2):
    """
    returns a lambda that determines the number of days between the two dates
    the dates can be strings (column names) or actual dates
    e.g. Aggregate(days('date', today), ['min','max'], 'days_since')
    TODO: should there be a Days(AggregateBase) for this? handle naming
    """
    if isinstance(date1, basestring) and isinstance(date2, basestring):
        return lambda d: (d[date2] - d[date1])/util.day
    elif isinstance(date1, basestring):
        return lambda d: (date2 - d[date1])/util.day
    elif isinstance(date2, basestring):
        return lambda d: (d[date2] - date1)/util.day
    else:
        return (date2 - date1)/util.day

def date_min(d):
    """
    groupby()['date_colum'].aggregate('min') returns a float?
    convert it back to a timestamp
    """
    return pd.to_datetime(d.min())

def date_max(d):
    return pd.to_datetime(d.max())
<|endoftext|>


<commit_before>#!/usr/bin/python

""" Downloads questions from dropbox to be integrated into srvy """

import dropbox

access_token = "I4gEG90KcJAAAAAAAAACORywbA7QoPUArtQGf0TstZFTmkIsDTrLksBSIGBHKD8I"

dbx = dropbox.Dropbox(access_token)

local = '/home/pi/srvy/synch/questions.csv'

with open (local, 'w') as f:
    metadata, res = dbx.files_download('/questions.csv')
    f.write(res.content)
<commit_msg>Add push chart and csv to pull_questions<commit_after>#!/usr/bin/python

""" Downloads questions from dropbox to be integrated into srvy """
""" Uploads charts and csv's"""

import dropbox
import pathlib
import re
import datetime

""" S E T U P _ V A R I A B L E S """

today = str(datetime.now().strftime('%Y-%m-%d'))
yesterday = str((datetime.now() - timedelta(days=1)).strftime('%Y-%m-%d'))
#You will need to generate your own access token for dropbox
access_token = "I4gEG90KcJAAAAAAAAACORywbA7QoPUArtQGf0TstZFTmkIsDTrLksBSIGBHKD8I"
#Creates a dropbox object
dbx = dropbox.Dropbox(access_token)

""" D O W N L O A D I N G _ FUNCTIONS """

def download_questions():
    local = '/home/pi/srvy/synch/questions.csv'#The folder/file locally that we want the files to download to.
    with open (local, 'w') as f:
        metadata, res = dbx.files_download('/questions.csv')
        f.write(res.content)

""" U P L O A D I N G _ FUNCTIONS"""

# the source files to upload
chart_to_upload= yesterday + ".html","srvy"
csv_to_upload = yesterday + ".csv"
files_to_upload = chart_to_upload, csv_to_upload

def foopa(files_to_upload):
    folder = pathlib.Path("/home/pi/srvy/export/")    # located in this folder
    for f in files_to_upload:
        filepath = folder + f
        upload_files(filepath, f)

def upload_files(filepath, filename):
    # open the file and upload it
    target = "responses/"
    with filepath.open("rb") as f:
        # upload gives you metadata about the file
        # we want to overwite any previous version of the file
        targetfile = target + filename
        meta = dbx.files_upload(f.read(), targetfile, mode=dropbox.files.WriteMode("overwrite"))

""" M A I N """
download_questions()
foopa(files_upload)
<|endoftext|>