class EtcdbFunction(object):
    """EtcdbFunction represents an SQL function.

    :param function_name: python function that implements SQL function.
    :type function_name: callable
    :param group: True if the functions is aggregate function
    :type group: bool
    :param args: Arguments to pass to function_name
    :param kwargs: Keyword arguments"""
    def __init__(self, *args, **kwargs):

        self._function = args[0]
        self._group = kwargs.get('group', False)
        self._args = args[1:]

    @property
    def function(self):
        """Return function name"""
        return self._function

    @property
    def group(self):
        """Return whether the function is aggregate"""
        return self._group

    def __call__(self, *args, **kwargs):
        return self._function(*args, **kwargs)

    def __eq__(self, other):
        return all(
            (
                isinstance(other, EtcdbFunction),
                self.function == other.function,
                self.group == other.group
            )
        )

    def __ne__(self, other):
        return not self.__eq__(other) def etcdb_count(result_set):
    """
    Count rows in result set

    :param result_set: ResultSet instance
    :type result_set: ResultSet
    :return: number of rows in ResultSet
    :rtype: int
    """
    return len(result_set.rows) def eval_row(table_columns, table_row, tree):
    """Find values of a row. table_columns are fields in the table.
    The result columns is taken from tree.expressions.

    :param table_columns: Columns in the table row.
    :type table_columns: ColumnSet
    :param table_row: Input row.
    :type table_row: Row
    :param tree: Parsing tree.
    :type tree: SQLTree
    """
    result_row = ()

    for select_item in tree.expressions:
        expr = select_item[0]
        expr_value = eval_expr((table_columns, table_row),
                               tree=expr)[1]
        if isinstance(expr_value, EtcdbFunction) and not expr_value.group:
            expr_value = expr_value()

        result_row += (expr_value, )

    return result_row def prepare_columns(tree):
    """
    Generate ColumnSet for query result. ColumnsSet doesn't include
    a grouping function.

    :return: Columns of the query result.
    :rtype: ColumnSet
    """
    columns = ColumnSet()
    for select_item in tree.expressions:

        expr, alias = select_item

        colname, _ = eval_expr(None, tree=expr)
        if alias:
            colname = alias

        columns.add(Column(colname))

    return columns def group_result_set(func, result_set, table_row, tree, pos):
    """Apply a group function to result set and return an aggregated row.

    :param func: Aggregation function.
    :type func: callable
    :param result_set: Result set to aggregate.
    :type result_set: ResultSet
    :param table_row: Table row to base aggregated row on.
    :type table_row: Row
    :param tree: Parsing tree.
    :type tree: SQLTree
    :param pos: Aggregate function position in the resulting row.
    :type pos: int
    :return: Result set with aggregated row.
    :rtype: ResultSet"""
    group_value = func(result_set)
    values = list(eval_row(result_set.columns, table_row, tree))
    values[pos] = group_value
    row = Row(tuple(values))
    return ResultSet(prepare_columns(tree), [row]) def get_row_by_primary_key(etcd_client, db, table, primary_key, **kwargs):
    """
    Read row from etcd by its primary key value.

    :param etcd_client:
    :type etcd_client: Client
    :param db:
    :param table:
    :param primary_key: Primary key value.
    :param kwargs: See below.
    :return: Row
    :rtype: Row

    :Keyword Arguments:
            * **wait** (bool) - If True it will wait for a change in the key.
            * **wait_index** (int) - When waiting you can specify index to
                wait for.
    """
    key = "/{db}/{tbl}/{pk}".format(
        db=db,
        tbl=table,
        pk=primary_key
    )
    client_kwargs = {}
    if 'wait' in kwargs:
        client_kwargs['wait'] = kwargs.get('wait')
        if 'wait_index' in kwargs:
            client_kwargs['waitIndex'] = kwargs.get('wait_index')
    etcd_response = None
    for i in xrange(30):
        try:
            etcd_response = etcd_client.read(key, **client_kwargs)
            break
        except EtcdEmptyResponse as err:
            LOG.warning("Retry #%d after error: %s", i, err)
            time.sleep(1)
    if not etcd_response:
        raise InternalError('Failed to get response from etcd')

    row = ()
    field_values = json.loads(etcd_response.node['value'])
    table_columns = get_table_columns(etcd_client, db, table)
    for col in table_columns:
        row += (field_values[str(col)], )

    try:
        etcd_index = etcd_response.x_etcd_index
    except AttributeError:
        etcd_index = 0
    return Row(row, etcd_index=etcd_index,
               modified_index=etcd_response.node['modifiedIndex']) class ColumnSet(object):
    """
    Instantiate a Column set

    :param columns: Optional dictionary with column definitions
    :type columns: dict
    """
    def __init__(self, columns=None):

        self._columns = []
        self._column_position = 0
        if columns:
            for colname, value in columns.iteritems():
                self._columns.append(Column(colname,
                                            coltype=value['type'],
                                            options=ColumnOptions(
                                                value['options']
                                            )))
        else:
            self._columns = []
            self._column_position = 0

    def add(self, column):
        """Add column to ColumnSet

        :param column: Column instance
        :type column: Column
        :return: Updated CoulmnSet instance
        :rtype: ColumnSet
        """
        if isinstance(column, Column):
            self._columns.append(column)
        else:
            raise ProgrammingError('%s must be Column type' % column)
        return self

    @property
    def empty(self):
        """True if there are no columns in the ColumnSet"""
        return not self._columns

    @property
    def columns(self):
        """Returns list of Columns"""
        return self._columns

    @property
    def primary(self):
        """Return primary key column"""
        for col in self._columns:
            if col.primary:
                return col
        return None

    def __contains__(self, item):
        return item in self._columns

    def __eq__(self, other):
        return self.columns == other.columns

    def __ne__(self, other):
        return not self.__eq__(other)

    def __str__(self):
        return "[" + ', '.join(['"%s"' % str(x) for x in self._columns]) + "]"

    def __len__(self):
        return len(self._columns)

    def __iter__(self):
        return self

    def next(self):
        """Return next Column"""
        self._column_position += 1
        try:
            return self._columns[self._column_position - 1]
        except IndexError:
            self._column_position = 0
            raise StopIteration()

    def __getitem__(self,  # pylint: disable=inconsistent-return-statements
                    key):
        if isinstance(key, int):
            return self._columns[key]
        else:
            for col in self._columns:
                if col == Column(key):
                    return col

    def index(self, column):
        """Find position of the given column in the set."""
        return self._columns.index(column) class Column(object):
    """
    Instantiate a Column

    :param colname: Column name.
    :type colname: str
    :param coltype: Column type
    :type coltype: str
    :param options: Column options
    :type options: ColumnOptions
    """
    def __init__(self, colname, coltype=None, options=None):
        self._colname = colname
        self._print_width = len(self._colname)
        self._coltype = coltype
        self._options = options

    @property
    def name(self):
        """Column name"""
        return self._colname

    @property
    def type(self):
        """Column type e.g. INT, VARCHAR, etc."""
        return self._coltype

    @property
    def auto_increment(self):
        """True if column is auto_incrementing."""
        return self._options.auto_increment

    @property
    def primary(self):
        """True if column is primary key."""
        return self._options.primary

    @property
    def nullable(self):
        """True if column is NULL-able."""
        return self._options.nullable

    @property
    def default(self):
        """Column default value."""
        return self._options.default

    @property
    def unique(self):
        """True if column is unique key."""
        return self._options.unique

    @property
    def print_width(self):
        """How many symbols client has to spare to print column value.
         A column name can be short, but its values may be longer.
         To align column and its values print_width is number of characters
         a client should allocate for the column name so it will be as lager
         as the largest columns length value."""
        return self._print_width

    def set_print_width(self, width):
        """Sets print_width."""
        self._print_width = width

    def __eq__(self, other):
        return self.name == other.name

    def __ne__(self, other):
        return not self.__eq__(other)

    def __str__(self):
        return self.name class Row(object):
    """
    Row class

    :param row: Row values
    :type row: tuple
    """
    def __init__(self, row, etcd_index=0, modified_index=0):
        if not isinstance(row, tuple):
            raise ProgrammingError('%s must be tuple')
        self._row = row
        self._field_position = 0
        self._etcd_index = etcd_index
        self._modified_index = modified_index

    @property
    def row(self):
        """
        :return: Return tuple with row values..
        :rtype: tuple
        """
        return self._row

    @property
    def etcd_index(self):
        """A row in etcdb is a key. Etcd index corresponds to X-Etcd-Index
        in etcd response header.

        :return: Etcd index.
        :rtype: int
        """
        return self._etcd_index

    @property
    def modified_index(self):
        """modifiedIndex of a key in etcd"""
        return self._modified_index

    def __str__(self):
        return json.dumps(self._row)

    def __eq__(self, other):
        try:
            return self._row == other.row
        except AttributeError:
            return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __iter__(self):
        return self

    def next(self):
        """Return next field in the row."""
        self._field_position += 1
        try:
            return self._row[self._field_position - 1]
        except IndexError:
            self._field_position = 0
            raise StopIteration()

    def __getitem__(self, key):
        return self._row[key]

    def __repr__(self):
        return self.__str__() class ResultSet(object):
    """
    Represents query result

    :param columns: Column set instance.
    :type columns: ColumnSet
    :param rows: List of Rows
    :type rows: list(Row)
    """
    def __init__(self, columns, rows=None):

        if not isinstance(columns, ColumnSet):
            raise ProgrammingError('%s must be ColumnSet' % columns)

        if columns.empty:
            raise ProgrammingError('columns must not be empty')
        self.columns = columns
        if rows:
            self.rows = rows
            for row in self.rows:
                self._update_print_width(row)
        else:
            self.rows = []

        self._pos = 0

    def __eq__(self, other):
        try:
            return all((self.columns == other.columns,
                        self.rows == other.rows))
        except AttributeError:
            return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __str__(self):
        return "Columns: %s\nRows: %s" \
               % (
                   self.columns,
                   "[" + ', '.join(['%s' % str(x) for x in self.rows]) + "]"
               )

    def __iter__(self):
        return self

    def next(self):
        """Return next row in the result set."""
        self._pos += 1
        try:
            return self.rows[self._pos - 1]
        except IndexError:
            raise StopIteration()

    def __getitem__(self, key):
        return self.rows[key]

    def __len__(self):
        return len(self.rows)

    def rewind(self):
        """Move internal records pointer to the beginning of the result set.
        After this call .fetchone() will start returning rows again."""
        self._pos = 0

    @property
    def n_rows(self):
        """Return number of rows in the result set."""
        return len(self.rows)

    @property
    def n_cols(self):
        """Return number of columns in the result set."""
        return len(self.columns)

    def add_row(self, row):
        """
        Add row to result set

        :param row: Row instance
        :type row: Row
        :return: Updated result set
        :rtype: ResultSet
        :raise InternalError: if row is not a Row class instance.
        """
        if not isinstance(row, Row):
            raise InternalError('%r must be of type Row' % row)
        self.rows.append(row)
        self._update_print_width(row)
        # return self

    def _update_print_width(self, row):
        i = 0
        for field in row:
            if len(str(field)) > self.columns[i].print_width:
                self.columns[i].set_print_width(len(str(field)))
            i += 1 class SQLTree(object):
    def __init__(self):
        self.query = None
        self.success = False
        self.query_type = None
        self.expressions = []
        self.db = None
        self.table = None
        self.fields = {}
        self.options = {}
        self.where = None
        self.limit = None
        self.order = {
            'by': None,
            'direction': 'ASC'
        }
        self.wait = False
        self.lock = None

    def reset(self):
        self.__init__()

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return json.dumps({
            'query': self.query,
            'success': self.success,
            'query_type': self.query_type,
            'expressions': self.expressions,
            'db': self.db,
            'table': self.table,
            'fields': self.fields,
            'options': self.options,
            'where': self.where,
            'limit': self.limit,
            'order': self.order
        })
import mock
import pytest
from pyetcd import EtcdResult
from etcdb.eval_expr import EtcdbFunction, etcdb_count
from etcdb.execute.dml.select import eval_row, prepare_columns, \
    group_result_set, get_row_by_primary_key
from etcdb.resultset import ColumnSet, Column, Row, ResultSet
from etcdb.sqlparser.sql_tree import SQLTree
    assert actual_cs == cs


@pytest.mark.parametrize('cs, row, expressions, result', [
    (
        ColumnSet().add(Column('id')).add(Column('name')),
        Row((1, 'aaa')),
        [
            (
                ('bool_primary',
                 ('predicate',
                  ('bit_expr',
                   ('simple_expr',
                    ('IDENTIFIER', 'id')
                    )
                   )
                  )
                 ),
                None),
            (
                ('bool_primary',
                 ('predicate',
                  ('bit_expr',
                   ('simple_expr',
                    ('IDENTIFIER', 'name')
                    )
                   )
                  )
                 ),
                None
            )
        ],
        (1, 'aaa')

    ),
    (
        ColumnSet().add(Column('id')).add(Column('COUNT(*)')),
        Row((1, 'aaa')),
        [
            (
                ('bool_primary',
                 ('predicate',
                  ('bit_expr',
                   ('simple_expr',
                    ('IDENTIFIER', 'id')
                    )
                   )
                  )
                 ),
                None),
            (
                ('bool_primary',
                 ('predicate',
                  ('bit_expr',
                   ('simple_expr',
                    ('function_call', 'COUNT')
                    )
                   )
                  )
                 ),
                '__count'
            )
        ],
        (1, EtcdbFunction(etcdb_count, group=True))

    ),
    (
        ColumnSet().add(Column('COUNT(*)')),
        None,
        [
            (
                ('bool_primary',
                 ('predicate',
                  ('bit_expr',
                   ('simple_expr',
                    ('function_call', 'COUNT')
                    )
                   )
                  )
                 ),
                '__count'
            )
        ],
        (EtcdbFunction(etcdb_count, group=True),)

    )

])
def test_eval_row(cs, row, expressions, result):
    tree = SQLTree()
    tree.expressions = expressions

    assert eval_row(cs, row, tree) == result


@pytest.mark.parametrize('rs, row, expressions, result', [
    (
        ResultSet(ColumnSet().add(Column(str(
            EtcdbFunction(etcdb_count, group=True)))), []),
        None,
        [
            (
                ('bool_primary',
                 ('predicate',
                  ('bit_expr',
                   ('simple_expr',
                    ('function_call', 'COUNT')
                    )
                   )
                  )
                 ),
                '__count'
            )
        ],
        ResultSet(
            ColumnSet().add(Column('__count')),
            [Row((0,))]
        )
    ),
    (

        ResultSet(ColumnSet().add(Column(str(