
from tablestore import *
import traceback

def getRow_ots_primary(row):
    _dict = dict()
    if row is None:
        return _dict
    for part in row.attribute_columns:
        _dict[part[0]] = part[1]
    for part in row.primary_key:
        _dict[part[0]] = part[1]
    return _dict

def getRow_ots(rows):
    list_dict = []
    for row in rows:
        _dict = dict()
        for part in row:
            for v in part:
                _dict[v[0]] = v[1]
        list_dict.append(_dict)
    return list_dict

def getPartitionKey(docid):
    return int(docid)%500+1

def getDocument(_query,ots_client,columns,sort="page_time",table_name="document",table_index="document_index"):
    try:
        item = _query.get("query")
        _limit = _query.get("limit",100)
        keyword = _query.get("keyword")
        l_rows = []
        rows,next_token,total_count,is_all_succeed = ots_client.search(table_name,table_index,
                                                                       SearchQuery(item,sort=Sort(sorters=[FieldSort(sort,SortOrder.ASC)]),limit=min(_limit,100),get_total_count=True),
                                                                       ColumnsToGet(columns,return_type=ColumnReturnType.SPECIFIED))
        dict_row = getRow_ots(rows)
        for _row in dict_row:
            _row["keyword"] = keyword
        l_rows.extend(dict_row)
        _count = len(dict_row)
        while next_token:
            rows,next_token,total_count,is_all_succeed = ots_client.search(table_name,table_index,
                                                                           SearchQuery(item,next_token=next_token,limit=min(_limit,100),get_total_count=True),
                                                                           ColumnsToGet(columns,return_type=ColumnReturnType.SPECIFIED))
            dict_row = getRow_ots(rows)
            for _row in dict_row:
                _row["keyword"] = keyword
            l_rows.extend(dict_row)
            if _limit and len(l_rows)>=_limit:
                break
            _count += len(dict_row)
    except Exception as e:
        traceback.print_exc()
    return l_rows