tests/common/test_sql_composite_key.py

# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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 collections import OrderedDict
from enum import Enum

import pytest

from pyignite import GenericObjectMeta, AioClient
from pyignite.datatypes import IntObject, String


class StudentKey(
    metaclass=GenericObjectMeta,
    type_name='test.model.StudentKey',
    schema=OrderedDict([
        ('ID', IntObject),
        ('DEPT', String)
    ])
):
    pass


class Student(
    metaclass=GenericObjectMeta,
    type_name='test.model.Student',
    schema=OrderedDict([
        ('NAME', String),
    ])
):
    pass


create_query = '''CREATE TABLE StudentTable (
    id INT(11),
    dept VARCHAR,
    name CHAR(24),
    PRIMARY KEY (id, dept))
    WITH "CACHE_NAME=StudentCache, KEY_TYPE=test.model.StudentKey, VALUE_TYPE=test.model.Student"'''

insert_query = '''INSERT INTO StudentTable (id, dept, name) VALUES (?, ?, ?)'''

select_query = 'SELECT id, dept, name FROM StudentTable'

select_kv_query = 'SELECT _key, _val FROM StudentTable'

drop_query = 'DROP TABLE StudentTable IF EXISTS'


@pytest.fixture
def student_table_fixture(client):
    yield from __create_student_table_fixture(client)


@pytest.fixture
async def async_student_table_fixture(async_client):
    async for _ in __create_student_table_fixture(async_client):
        yield


def __create_student_table_fixture(client):
    def inner():
        client.sql(drop_query)
        client.sql(create_query)
        yield None
        client.sql(drop_query)

    async def inner_async():
        await client.sql(drop_query)
        await client.sql(create_query)
        yield None
        await client.sql(drop_query)

    return inner_async() if isinstance(client, AioClient) else inner()


class InsertMode(Enum):
    SQL = 1
    CACHE = 2


@pytest.mark.parametrize('insert_mode', [InsertMode.SQL, InsertMode.CACHE])
def test_sql_composite_key(client, insert_mode, student_table_fixture):
    __perform_test(client, insert_mode)


@pytest.mark.asyncio
@pytest.mark.parametrize('insert_mode', [InsertMode.SQL, InsertMode.CACHE])
async def test_sql_composite_key_async(async_client, insert_mode, async_student_table_fixture):
    await __perform_test(async_client, insert_mode)


def __perform_test(client, insert=InsertMode.SQL):
    student_key = StudentKey(2, 'Business')
    student_val = Student('Abe')

    def validate_query_result(key, val, query_result):
        """
        Compare query result with expected key and value.
        """
        assert len(query_result) == 2
        sql_row = dict(zip(query_result[0], query_result[1]))

        assert sql_row['ID'] == key.ID
        assert sql_row['DEPT'] == key.DEPT
        assert sql_row['NAME'] == val.NAME

    def validate_kv_query_result(key, val, query_result):
        """
        Compare query result with expected key and value.
        """
        assert len(query_result) == 2
        sql_row = dict(zip(query_result[0], query_result[1]))

        sql_key, sql_val = sql_row['_KEY'], sql_row['_VAL']
        assert sql_key.ID == key.ID
        assert sql_key.DEPT == key.DEPT
        assert sql_val.NAME == val.NAME

    def inner():
        if insert == InsertMode.SQL:
            result = client.sql(insert_query, query_args=[student_key.ID, student_key.DEPT, student_val.NAME])
            assert next(result)[0] == 1
        else:
            studentCache = client.get_cache('StudentCache')
            studentCache.put(student_key, student_val)
            val = studentCache.get(student_key)
            assert val is not None
            assert val.NAME == student_val.NAME

        query_result = list(client.sql(select_query, include_field_names=True))
        validate_query_result(student_key, student_val, query_result)

        query_result = list(client.sql(select_kv_query, include_field_names=True))
        validate_kv_query_result(student_key, student_val, query_result)

    async def inner_async():
        if insert == InsertMode.SQL:
            result = await client.sql(insert_query, query_args=[student_key.ID, student_key.DEPT, student_val.NAME])
            assert (await result.__anext__())[0] == 1
        else:
            studentCache = await client.get_cache('StudentCache')
            await studentCache.put(student_key, student_val)
            val = await studentCache.get(student_key)
            assert val is not None
            assert val.NAME == student_val.NAME

        async with client.sql(select_query, include_field_names=True) as cursor:
            query_result = [r async for r in cursor]
            validate_query_result(student_key, student_val, query_result)

        async with client.sql(select_kv_query, include_field_names=True) as cursor:
            query_result = [r async for r in cursor]
            validate_kv_query_result(student_key, student_val, query_result)

    return inner_async() if isinstance(client, AioClient) else inner()