<commit_before>#coding=utf8
from __future__ import unicode_literals
import sys
import os
import shutil
import re
import xlrd
import yaml
import alias
import typedef
import lseri

reload(sys)
sys.setdefaultencoding("utf8")
import codecs
def cp65001(name):
    if name.lower() == "cp65001":
        return codecs.lookup("utf8")
codecs.register(cp65001)

g_error_l = []
g_alias_d = None
g_alias_deps = None
g_struct_deps = {}
g_sheet_struct_deps = {}

fp_log = None
def fprint(msg):
    global fp_log
    if not fp_log:
        fp_log = open("log.txt", "w")
    fp_log.write(msg)
    fp_log.write(os.linesep)
    fp_log.flush()

def output(msg):
    sys.stdout.write(msg)
    sys.stdout.write(os.linesep)
    sys.stdout.flush()

def error(msg):
    output(lseri.tolua({"error":msg}))
    sys.exit(1)

def print_exc():
    import traceback
    traceback.print_exc()

# 基础数据类型：int, float, string
basictype_convert_tbl = {
        "int":int,
        "float":float,
        "string":lambda s:str(s).encode("utf8"),
        }
# 自定义struct
g_struct_d = None

# 容器数据类型: list<T>, dict<T1, T2>
CONTAINER_RE = re.compile("^(list|dict)<(.+?)(?:,\s*(.+?))??>$")
#m = CONTAINER_RE.match(u"list<int>")
#print(m.groups())

def find_struct_deps(type_s):
    if type_s in basictype_convert_tbl:
        return None
    if type_s in g_struct_d:
        if type_s in g_struct_deps:
            return ["s", g_struct_deps[type_s]]
        else:
            return None
    m = CONTAINER_RE.match(type_s)
    sg = m.groups()
    if sg[0] == "list":
        if sg[1] in g_struct_deps:
            return ["l", g_struct_deps[sg[1]]]
    elif sg[0] == "dict":
        if sg[2] in g_struct_deps:
            return ["d", g_struct_deps[sg[2]]]
    return None

def get_basic_or_struct_cf(s):
    if s in basictype_convert_tbl:
        return basictype_convert_tbl[s]
    if s not in g_struct_d:
        return None
    cfg = g_struct_d[s]
    def cf(cont):
        l = cont.split("|")
        assert len(l)==len(cfg), cont
        ret = {}
        for idx, v in enumerate(l):
            ret[cfg[idx].keys()[0]] = basictype_convert_tbl[cfg[idx].values()[0]](v)
        return ret
    return cf

def make_convert_func(type_s):
    cf = get_basic_or_struct_cf(type_s)
    if cf:
        return cf
    m = CONTAINER_RE.match(type_s)
    sg = m.groups()
    assert len(sg) == 3, "类型解析错误"
    if sg[0] == "list":
        assert not sg[2], "list定义有误"
        f = get_basic_or_struct_cf(sg[1])
        assert f, "list元素类型定义有误:%s"%type_s
        def cf(s):
            return [f(i) for i in s.split()]
        return cf
    elif sg[0] == "dict":
        k_f = basictype_convert_tbl[sg[1]]
        assert k_f, "dict key类型定义有误%s"%type_s
        v_f = get_basic_or_struct_cf(sg[2])
        assert v_f, "dict value类型定义有误:%s"%type_s
        def cf(s):
            d = {}
            for i in s.split():
                l = i.split(":")
                d[k_f(l[0])] = v_f(l[1])
            return d
        return cf
    raise Exception("未定义类型:%s"%type_s)

TYPE_DEFAULT_RE = re.compile("^default(?:\((.*)\))?$")
#m = TYPE_DEFAULT_RE.match("default")
#print m.group(2)
type_default_tbl = {
        "int":0,
        "float":0.0,
        "string":"",
        }
def parse_type_tag(ncol, tag_sl, type_s, conv_f):
    ret = {}
    def _key_f():
        assert ncol == 0, "key必须是第一列"
        assert type_s == "int" or type_s == "string", "类型:<%s>不能做key"%type_s
        ret["key"] = True
    def _ignore_f():
        ret["ignore"] = True
    def _raw_f():
        ret["raw"] = True
    def _key_alias_f():
        assert ncol == 1, "key_alias必须是第二列"
        ret["key_alias"] = True
    def _index_f():
        ret["index"] = True
    tag_fs = {
        "key":_key_f,
        "ignore":_ignore_f,
        "raw":_raw_f,
        "key_alias":_key_alias_f,
        "index":_index_f,
    }

    for tag_s in tag_sl:
        if not tag_s.startswith("default"):
            assert tag_s in tag_fs, "标签填写错误:<%s>"%tag_s
            tag_fs[tag_s]()
            continue
        # defautl 处理
        assert type_s not in g_struct_d, "自定义类型:<%s>不能设置默认"%type_s
        m = TYPE_DEFAULT_RE.match(tag_s)
        if m:
            default_val = m.group(1)
            if not default_val and type_s in type_default_tbl:
                default_val = type_default_tbl[type_s]
            else:
                default_val = conv_f(default_val if default_val else "")
        else:
            assert False, tag_s
        ret["default"] = default_val
    assert not ("key" in ret and "default" in ret), "key类型不能设置default"
    assert not ("key" in ret and "ignore" in ret), "key类型不能设置ignore"
    return ret

def open_xls(filename):
    return xlrd.open_workbook(filename)

# 返回的是下标
def _find_dup_in_list(l):
    d = {}
    for n, i in enumerate(l):
        if i in d:
            return n
        d[i] = 1
    return -1

def _num2colname(n):
    def _n2chr(i):
        return chr(65+i)
    if n < 26:
        return _n2chr(n)
    elif n < 26*27:
        return _n2chr(n/26-1) + _n2chr(n%26)
    return str(n)

def sheet_to_dict(sheet, alias_d):
    conv_funcs = []
    tags = []
    struct_deps_l = []
    struct_deps_d = {}
    alias_deps = alias_d["deps"] if alias_d and "deps" in alias_d else {}
    alias_d = alias_d["alias"] if alias_d and "alias" in alias_d else None
    try:
        # 处理第一行，类型
        end_col = None
        for n, i in enumerate(sheet.row_values(0)):
            end_col = n + 1
            # 允许类型列填空，意味该列忽略
            if n > 0 and i == "" :
                end_col = n
                break
            type_sl = i.split("#")
            conv_f = make_convert_func(type_sl[0])
            conv_funcs.append(conv_f)
            if len(type_sl) > 1:
                tags.append(parse_type_tag(n, type_sl[1:], type_sl[0], conv_f))
            else:
                tags.append({})
            struct_deps_l.append(find_struct_deps(type_sl[0]))
    except Exception, e:
        raise Exception("sheet:<%s>类型行%s列填写错误, 内容:<%s>, msg:%s"%(sheet.name, _num2colname(n), i, e))

    name_row = sheet.row_values(1,end_colx=end_col)
    dup_idx = _find_dup_in_list(name_row)
    if dup_idx != -1:
        raise Exception("sheet:<%s>列名重复:<%s>"%(sheet.name, name_row[dup_idx]))
    col_names = []
    if alias_d:
        for name in name_row:
            col_names.append(alias_d[name] if name in alias_d else None)
    else:
        col_names = name_row

    for n, i in enumerate(col_names):
        if i and struct_deps_l[n]:
            struct_deps_d[i] = struct_deps_l[n]

    check_tag_f = lambda x,s: True if s in x else False
    raw_flag = check_tag_f(tags[0],"raw")
    raw_keys = {}
    key_flag = check_tag_f(tags[0],"key")
    key_alias_flag = check_tag_f(tags[1],"key_alias") if len(tags) > 1 else False
    ret = {} if key_flag and not raw_flag else []
    key_alias_d = {} if key_alias_flag else None
    for nrow in xrange(2, sheet.nrows):
        row = sheet.row_values(nrow, end_colx=end_col)
        row_d = {}
        try:
            # 注释行忽略
            if isinstance(row[0], unicode) and row[0].startswith("//"):
                continue
            row_key = None
            for ncol, value in enumerate(row):
                col_name = col_names[ncol]
                if not col_name:
                    if not key_alias_flag or "key_alias" not in tags[ncol]:
                        continue
                tag = tags[ncol]
                cv = None
                if "ignore" in tag:
                    continue
                # 如果该格子不填，获取的是空串
                if value == "" and "default" in tag:
                    cv = tag["default"]
                else:
                    # raw key 列可以为空
                    if not raw_flag:
                        assert value != "", "表项为空"
                        cv = conv_funcs[ncol](value)
                    else:
                        if value != "":
                            cv = conv_funcs[ncol](value)
                if ncol == 0 and "key" in tags[ncol]:
                    row_key = cv
                if ncol == 1 and "key_alias" in tags[ncol]:
                    row_key_alias = cv
                    if not col_name:
                        continue
                if "index" in tags[ncol]:
                    # 普通index
                    if col_name not in alias_deps:
                        # print col_name
                        continue
                row_d[col_name] = cv
            if not raw_flag and key_flag and row_key == None:
                raise Exception("本行key列没有导出！")
            def _check_key(check_d):
                # 检查key是否重复
                if raw_flag and row_key == None:
                    return
                assert row_key not in check_d, "key列内容重复, 行:%s,值:%s"%(nrow+1, row_key)
                check_d[row_key] = row_d
                if key_alias_flag:
                    assert row_key_alias not in key_alias_d, "key_alias列内容重复, 行:%s,值:%s"%(nrow+1, row_key_alias)
                    key_alias_d[row_key_alias] = row_key
            if isinstance(ret, dict):
                _check_key(ret)
            else:
                if raw_flag and key_flag:
                    _check_key(raw_keys)
                ret.append(row_d)
        except Exception, e:
            raise Exception("sheet:%s, cell:<行%s-列%s>, %s"%(sheet.name, nrow+1, _num2colname(ncol), e))
    return ret, struct_deps_d, key_alias_d

def get_alias_conf(wbname, shname):
    if wbname in g_alias_d:
        return g_alias_d[wbname]
    key = alias.make_key(wbname, shname)
    if key in g_alias_d:
        return g_alias_d[key]
    return None

def convert_xls(filename, out_dir=None):
    try:
        wb = open_xls(filename)
        wbname = os.path.basename(filename)
        ret = {}
        ext = {}
        for sheet in wb.sheets():
            if sheet.name.startswith("_"):
                continue
            if sheet.nrows <= 2:
                continue
            data, deps_d, key_alias_d = sheet_to_dict(sheet, get_alias_conf(wbname, sheet.name))
            ret[sheet.name] = data
            d = {}
            if len(deps_d) > 0:
                d["deps"] = deps_d
            d["typ"] = "l" if type(data) == type([]) else "d"
            d["key_alias"] = key_alias_d
            if d["typ"] == "l" and key_alias_d:
                tmp = {}
                for k, v in key_alias_d.iteritems():
                    tmp[v] = k
                d["key_check"] = tmp
            ext[sheet.name] = d
        return ret, ext
    except Exception, e:
        # import traceback
        # traceback.print_exc()
        error("file:%s, error: %s"%(filename, e))

def run_dir(path):
    files = []
    def visit(arg, dirname, names):
        for name in names:
            if name.endswith(".xls") and not name.startswith("_"):
                files.append(os.path.join(dirname, name))
    os.path.walk(path, visit, None)
    for f in files:
        fn = f[len(path)+1:]
        data, ext = convert_xls(f)
        out = {}
        out["filename"] = fn
        out["data"] = data
        out["ext"] = ext
        #if len(deps_d) > 0:
            #out["struct_deps"] = deps_d
        output(lseri.tolua(out))

if __name__ == "__main__":
    fpath = sys.argv[1]
    import platform
    if platform.system() == "Windows":
        fpath = fpath.decode("gbk")
    else:
        fpath = fpath.decode("utf8")
    try:
        alias_raw, g_alias_d, g_alias_deps = \
                alias.parse(sys.argv[2])
        if len(alias_raw) > 0:
            output(lseri.tolua({"alias_fields":alias_raw}))
        if len(g_alias_deps) > 0:
            output(lseri.tolua({"alias_deps":g_alias_deps}))
        g_struct_d, g_struct_deps = typedef.parse(os.path.join(fpath, "struct.yaml"))
        if len(g_struct_deps) > 0:
            output(lseri.tolua({"struct_deps":g_struct_deps}))
    except Exception, e:
        error(str(e))
    run_dir(fpath)
<commit_msg>支持导出空白表<commit_after>#coding=utf8
from __future__ import unicode_literals
import sys
import os
import shutil
import re
import xlrd
import yaml
import alias
import typedef
import lseri

reload(sys)
sys.setdefaultencoding("utf8")
import codecs
def cp65001(name):
    if name.lower() == "cp65001":
        return codecs.lookup("utf8")
codecs.register(cp65001)

g_error_l = []
g_alias_d = None
g_alias_deps = None
g_struct_deps = {}
g_sheet_struct_deps = {}

fp_log = None
def fprint(msg):
    global fp_log
    if not fp_log:
        fp_log = open("log.txt", "w")
    fp_log.write(msg)
    fp_log.write(os.linesep)
    fp_log.flush()

def output(msg):
    sys.stdout.write(msg)
    sys.stdout.write(os.linesep)
    sys.stdout.flush()

def error(msg):
    output(lseri.tolua({"error":msg}))
    sys.exit(1)

def print_exc():
    import traceback
    traceback.print_exc()

# 基础数据类型：int, float, string
basictype_convert_tbl = {
        "int":int,
        "float":float,
        "string":lambda s:str(s).encode("utf8"),
        }
# 自定义struct
g_struct_d = None

# 容器数据类型: list<T>, dict<T1, T2>
CONTAINER_RE = re.compile("^(list|dict)<(.+?)(?:,\s*(.+?))??>$")
#m = CONTAINER_RE.match(u"list<int>")
#print(m.groups())

def find_struct_deps(type_s):
    if type_s in basictype_convert_tbl:
        return None
    if type_s in g_struct_d:
        if type_s in g_struct_deps:
            return ["s", g_struct_deps[type_s]]
        else:
            return None
    m = CONTAINER_RE.match(type_s)
    sg = m.groups()
    if sg[0] == "list":
        if sg[1] in g_struct_deps:
            return ["l", g_struct_deps[sg[1]]]
    elif sg[0] == "dict":
        if sg[2] in g_struct_deps:
            return ["d", g_struct_deps[sg[2]]]
    return None

def get_basic_or_struct_cf(s):
    if s in basictype_convert_tbl:
        return basictype_convert_tbl[s]
    if s not in g_struct_d:
        return None
    cfg = g_struct_d[s]
    def cf(cont):
        l = cont.split("|")
        assert len(l)==len(cfg), cont
        ret = {}
        for idx, v in enumerate(l):
            ret[cfg[idx].keys()[0]] = basictype_convert_tbl[cfg[idx].values()[0]](v)
        return ret
    return cf

def make_convert_func(type_s):
    cf = get_basic_or_struct_cf(type_s)
    if cf:
        return cf
    m = CONTAINER_RE.match(type_s)
    sg = m.groups()
    assert len(sg) == 3, "类型解析错误"
    if sg[0] == "list":
        assert not sg[2], "list定义有误"
        f = get_basic_or_struct_cf(sg[1])
        assert f, "list元素类型定义有误:%s"%type_s
        def cf(s):
            return [f(i) for i in s.split()]
        return cf
    elif sg[0] == "dict":
        k_f = basictype_convert_tbl[sg[1]]
        assert k_f, "dict key类型定义有误%s"%type_s
        v_f = get_basic_or_struct_cf(sg[2])
        assert v_f, "dict value类型定义有误:%s"%type_s
        def cf(s):
            d = {}
            for i in s.split():
                l = i.split(":")
                d[k_f(l[0])] = v_f(l[1])
            return d
        return cf
    raise Exception("未定义类型:%s"%type_s)

TYPE_DEFAULT_RE = re.compile("^default(?:\((.*)\))?$")
#m = TYPE_DEFAULT_RE.match("default")
#print m.group(2)
type_default_tbl = {
        "int":0,
        "float":0.0,
        "string":"",
        }
def parse_type_tag(ncol, tag_sl, type_s, conv_f):
    ret = {}
    def _key_f():
        assert ncol == 0, "key必须是第一列"
        assert type_s == "int" or type_s == "string", "类型:<%s>不能做key"%type_s
        ret["key"] = True
    def _ignore_f():
        ret["ignore"] = True
    def _raw_f():
        ret["raw"] = True
    def _key_alias_f():
        assert ncol == 1, "key_alias必须是第二列"
        ret["key_alias"] = True
    def _index_f():
        ret["index"] = True
    tag_fs = {
        "key":_key_f,
        "ignore":_ignore_f,
        "raw":_raw_f,
        "key_alias":_key_alias_f,
        "index":_index_f,
    }

    for tag_s in tag_sl:
        if not tag_s.startswith("default"):
            assert tag_s in tag_fs, "标签填写错误:<%s>"%tag_s
            tag_fs[tag_s]()
            continue
        # defautl 处理
        assert type_s not in g_struct_d, "自定义类型:<%s>不能设置默认"%type_s
        m = TYPE_DEFAULT_RE.match(tag_s)
        if m:
            default_val = m.group(1)
            if not default_val and type_s in type_default_tbl:
                default_val = type_default_tbl[type_s]
            else:
                default_val = conv_f(default_val if default_val else "")
        else:
            assert False, tag_s
        ret["default"] = default_val
    assert not ("key" in ret and "default" in ret), "key类型不能设置default"
    assert not ("key" in ret and "ignore" in ret), "key类型不能设置ignore"
    return ret

def open_xls(filename):
    return xlrd.open_workbook(filename)

# 返回的是下标
def _find_dup_in_list(l):
    d = {}
    for n, i in enumerate(l):
        if i in d:
            return n
        d[i] = 1
    return -1

def _num2colname(n):
    def _n2chr(i):
        return chr(65+i)
    if n < 26:
        return _n2chr(n)
    elif n < 26*27:
        return _n2chr(n/26-1) + _n2chr(n%26)
    return str(n)

def sheet_to_dict(sheet, alias_d):
    conv_funcs = []
    tags = []
    struct_deps_l = []
    struct_deps_d = {}
    alias_deps = alias_d["deps"] if alias_d and "deps" in alias_d else {}
    alias_d = alias_d["alias"] if alias_d and "alias" in alias_d else None
    try:
        # 处理第一行，类型
        end_col = None
        for n, i in enumerate(sheet.row_values(0)):
            end_col = n + 1
            # 允许类型列填空，意味该列忽略
            if n > 0 and i == "" :
                end_col = n
                break
            type_sl = i.split("#")
            conv_f = make_convert_func(type_sl[0])
            conv_funcs.append(conv_f)
            if len(type_sl) > 1:
                tags.append(parse_type_tag(n, type_sl[1:], type_sl[0], conv_f))
            else:
                tags.append({})
            struct_deps_l.append(find_struct_deps(type_sl[0]))
    except Exception, e:
        raise Exception("sheet:<%s>类型行%s列填写错误, 内容:<%s>, msg:%s"%(sheet.name, _num2colname(n), i, e))

    name_row = sheet.row_values(1,end_colx=end_col)
    dup_idx = _find_dup_in_list(name_row)
    if dup_idx != -1:
        raise Exception("sheet:<%s>列名重复:<%s>"%(sheet.name, name_row[dup_idx]))
    col_names = []
    if alias_d:
        for name in name_row:
            col_names.append(alias_d[name] if name in alias_d else None)
    else:
        col_names = name_row

    for n, i in enumerate(col_names):
        if i and struct_deps_l[n]:
            struct_deps_d[i] = struct_deps_l[n]

    check_tag_f = lambda x,s: True if s in x else False
    raw_flag = check_tag_f(tags[0],"raw")
    raw_keys = {}
    key_flag = check_tag_f(tags[0],"key")
    key_alias_flag = check_tag_f(tags[1],"key_alias") if len(tags) > 1 else False
    ret = {} if key_flag and not raw_flag else []
    key_alias_d = {} if key_alias_flag else None
    for nrow in xrange(2, sheet.nrows):
        row = sheet.row_values(nrow, end_colx=end_col)
        row_d = {}
        try:
            # 注释行忽略
            if isinstance(row[0], unicode) and row[0].startswith("//"):
                continue
            row_key = None
            for ncol, value in enumerate(row):
                col_name = col_names[ncol]
                if not col_name:
                    if not key_alias_flag or "key_alias" not in tags[ncol]:
                        continue
                tag = tags[ncol]
                cv = None
                if "ignore" in tag:
                    continue
                # 如果该格子不填，获取的是空串
                if value == "" and "default" in tag:
                    cv = tag["default"]
                else:
                    # raw key 列可以为空
                    if not raw_flag:
                        assert value != "", "表项为空"
                        cv = conv_funcs[ncol](value)
                    else:
                        if value != "":
                            cv = conv_funcs[ncol](value)
                if ncol == 0 and "key" in tags[ncol]:
                    row_key = cv
                if ncol == 1 and "key_alias" in tags[ncol]:
                    row_key_alias = cv
                    if not col_name:
                        continue
                if "index" in tags[ncol]:
                    # 普通index
                    if col_name not in alias_deps:
                        # print col_name
                        continue
                row_d[col_name] = cv
            if not raw_flag and key_flag and row_key == None:
                raise Exception("本行key列没有导出！")
            def _check_key(check_d):
                # 检查key是否重复
                if raw_flag and row_key == None:
                    return
                assert row_key not in check_d, "key列内容重复, 行:%s,值:%s"%(nrow+1, row_key)
                check_d[row_key] = row_d
                if key_alias_flag:
                    assert row_key_alias not in key_alias_d, "key_alias列内容重复, 行:%s,值:%s"%(nrow+1, row_key_alias)
                    key_alias_d[row_key_alias] = row_key
            if isinstance(ret, dict):
                _check_key(ret)
            else:
                if raw_flag and key_flag:
                    _check_key(raw_keys)
                ret.append(row_d)
        except Exception, e:
            raise Exception("sheet:%s, cell:<行%s-列%s>, %s"%(sheet.name, nrow+1, _num2colname(ncol), e))
    return ret, struct_deps_d, key_alias_d

def get_alias_conf(wbname, shname):
    if wbname in g_alias_d:
        return g_alias_d[wbname]
    key = alias.make_key(wbname, shname)
    if key in g_alias_d:
        return g_alias_d[key]
    return None

def convert_xls(filename, out_dir=None):
    try:
        wb = open_xls(filename)
        wbname = os.path.basename(filename)
        ret = {}
        ext = {}
        for sheet in wb.sheets():
            if sheet.name.startswith("_"):
                continue
            if sheet.nrows < 2:
                continue
            data, deps_d, key_alias_d = sheet_to_dict(sheet, get_alias_conf(wbname, sheet.name))
            ret[sheet.name] = data
            d = {}
            if len(deps_d) > 0:
                d["deps"] = deps_d
            d["typ"] = "l" if type(data) == type([]) else "d"
            d["key_alias"] = key_alias_d
            if d["typ"] == "l" and key_alias_d:
                tmp = {}
                for k, v in key_alias_d.iteritems():
                    tmp[v] = k
                d["key_check"] = tmp
            ext[sheet.name] = d
        return ret, ext
    except Exception, e:
        # import traceback
        # traceback.print_exc()
        error("file:%s, error: %s"%(filename, e))

def run_dir(path):
    files = []
    def visit(arg, dirname, names):
        for name in names:
            if name.endswith(".xls") and not name.startswith("_"):
                files.append(os.path.join(dirname, name))
    os.path.walk(path, visit, None)
    for f in files:
        fn = f[len(path)+1:]
        data, ext = convert_xls(f)
        out = {}
        out["filename"] = fn
        out["data"] = data
        out["ext"] = ext
        #if len(deps_d) > 0:
            #out["struct_deps"] = deps_d
        output(lseri.tolua(out))

if __name__ == "__main__":
    fpath = sys.argv[1]
    import platform
    if platform.system() == "Windows":
        fpath = fpath.decode("gbk")
    else:
        fpath = fpath.decode("utf8")
    try:
        alias_raw, g_alias_d, g_alias_deps = \
                alias.parse(sys.argv[2])
        if len(alias_raw) > 0:
            output(lseri.tolua({"alias_fields":alias_raw}))
        if len(g_alias_deps) > 0:
            output(lseri.tolua({"alias_deps":g_alias_deps}))
        g_struct_d, g_struct_deps = typedef.parse(os.path.join(fpath, "struct.yaml"))
        if len(g_struct_deps) > 0:
            output(lseri.tolua({"struct_deps":g_struct_deps}))
    except Exception, e:
        error(str(e))
    run_dir(fpath)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from models import item_container, item
import core.command
from core.localisation import _

class open(core.command.command):
	def run(self):
		"""
		c.run()

		Open an item container in the area where the player is.
		The result of the command is a list of the items of the container
		"""
		if len(self._args) == 0:
			raise core.command.exception(_('ERROR_OPEN_NO_CONTAINER_PROVIDED'))

		containers = item_container.factory.getAllFromIdAreaAndType(
			self._player.getAreaId(),
			self._args[0]
		)

		if len(containers) == 0:
			raise core.command.exception(_('ERROR_OPEN_CONTAINER_NOT_AVAILABLE'))
		elif len(containers) > 1:
			raise core.command.exception(_('ERROR_OPEN_MULTIPLE_CONTAINERS_AVAILABLE'))

		items = item.inventory.fromStr(containers[0]['items'])

		result = {'container_type': self._args[0], 'items': list()}
		for i in items:
			it = item.model.loadById(i)
			result['items'].append({
				'name': it['name'],
				'quantity': items[i]['quantity']
			})

		return result

	def render(self, data):
		output = list()
		output.append(_('ITEMS_IN_CONTAINER_%s') % data['container_type'])
		for i in data['items']:
			output.append(str(i['quantity']).rjust(3) + ' ' + i['name'])
		return '\n'.join(output)
<commit_msg>possibility to add the index of the container in the open command<commit_after># -*- coding: utf-8 -*-

from models import item_container, item
import core.command
from core.localisation import _

class open(core.command.command):
	def run(self):
		"""
		c.run()

		Open an item container in the area where the player is.
		The result of the command is a list of the items of the container
		"""
		index = None
		if len(self._args) == 0:
			raise core.command.exception(_('ERROR_OPEN_NO_CONTAINER_PROVIDED'))
		elif len(self._args) == 2:
			index = self._args[1]

		containers = item_container.factory.getAllFromIdAreaAndType(
			self._player.getAreaId(),
			self._args[0]
		)

		if len(containers) == 0:
			raise core.command.exception(_('ERROR_OPEN_CONTAINER_NOT_AVAILABLE'))
		elif len(containers) > 1 and index is None:
			raise core.command.exception(_('ERROR_OPEN_MULTIPLE_CONTAINERS_AVAILABLE'))

		items = item.inventory.fromStr(containers[int(index or 0)]['items'])

		result = {'container_type': self._args[0], 'items': list()}
		for i in items:
			it = item.model.loadById(i)
			result['items'].append({
				'name': it['name'],
				'quantity': items[i]['quantity']
			})

		return result

	def render(self, data):
		output = list()
		output.append(_('ITEMS_IN_CONTAINER_%s') % data['container_type'])
		for i in data['items']:
			output.append(str(i['quantity']).rjust(3) + ' ' + i['name'])
		return '\n'.join(output)
<|endoftext|>


<commit_before>from __future__ import division

from .. import bar, xcbq, window
from . import base

import xcffib
from xcffib.xproto import EventMask, SetMode
import atexit
import struct

class Icon(window._Window):
    _windowMask = EventMask.StructureNotify | \
        EventMask.Exposure

    def __init__(self, win, qtile, systray):
        window._Window.__init__(self, win, qtile)
        self.systray = systray
        self.width = systray.icon_size
        self.height = systray.icon_size

    def handle_ConfigureNotify(self, event):
        icon_size = self.systray.icon_size
        self.updateHints()

        try:
            width = self.hints["min_width"]
            height = self.hints["min_height"]
        except KeyError:
            width = icon_size
            height = icon_size

        if height > icon_size:
            width = width * icon_size // height
            height = icon_size
        if height <= 0:
            width = icon_size
            height = icon_size

        self.width = width
        self.height = height
        self.window.set_attribute(backpixmap=self.systray.drawer.pixmap)
        return False

    def handle_DestroyNotify(self, event):
        wid = event.window
        del(self.qtile.windowMap[wid])
        del(self.systray.icons[wid])
        self.systray.draw()
        return False

    handle_UnmapNotify = handle_DestroyNotify

class TrayWindow(window._Window):
    _windowMask = EventMask.StructureNotify | \
        EventMask.Exposure

    def __init__(self, win, qtile, systray):
        window._Window.__init__(self, win, qtile)
        self.systray = systray

    def handle_ClientMessage(self, event):
        atoms = self.qtile.conn.atoms

        opcode = event.type
        wid = event.window
        data = event.data.data32
        message = data[1]

        conn = self.qtile.conn.conn
        parent = self.systray.bar.window.window

        # message == 0 corresponds to SYSTEM_TRAY_REQUEST_DOCK
        # TODO: handle system tray messages http://standards.freedesktop.org/systemtray-spec/systemtray-spec-latest.html
        # TODO: re-add 'and message == 0' constraint which /should/ work but doesn't with xcffib
        if opcode == atoms['_NET_SYSTEM_TRAY_OPCODE']:
            try:
                w = xcbq.Window(self.qtile.conn, wid)
                icon = Icon(w, self.qtile, self.systray)
                self.systray.icons[wid] = icon
                self.qtile.windowMap[wid] = icon

                # add icon window to the save-set, so it gets reparented
                # to the root window when qtile dies
                # TODO: figure out why this gives a BadMatch with xcffib.
                # conn.core.ChangeSaveSet(SetMode.Insert, wid)

                conn.core.ReparentWindow(wid, parent.wid, 0, 0)
                conn.flush()
                w.map()
            except xcffib.xproto.DrawableError:
                # The icon wasn't ready to be drawn yet... (NetworkManager does
                # this sometimes), so we just forget about it and wait for the
                # next event.
                pass
        return False

class Systray(base._Widget):
    """
        A widget that manage system tray
    """
    defaults = [
        ('icon_size', 20, 'Icon width'),
        ('padding', 5, 'Padding between icons'),
    ]

    def __init__(self, **config):
        base._Widget.__init__(self, bar.CALCULATED, **config)
        self.add_defaults(Systray.defaults)
        self.traywin = None
        self.icons = {}

    def button_press(self, x, y, button):
        pass

    def calculate_width(self):
        width = sum([i.width for i in self.icons.values()])
        width += self.padding * len(self.icons)
        return width

    def _configure(self, qtile, bar):
        base._Widget._configure(self, qtile, bar)
        self.qtile = qtile
        self.bar = bar
        atoms = qtile.conn.atoms
        win = qtile.conn.create_window(-1, -1, 1, 1)
        self.traywin = TrayWindow(win, self.qtile, self)
        qtile.windowMap[win.wid] = self.traywin
        qtile.conn.conn.core.SetSelectionOwner(
            win.wid,
            atoms['_NET_SYSTEM_TRAY_S0'],
            xcffib.CurrentTime
        )
        event = struct.pack(
            'BBHII5I', 33, 32, 0, qtile.root.wid,
            atoms['MANAGER'],
            xcffib.CurrentTime, atoms['_NET_SYSTEM_TRAY_S0'],
            win.wid, 0, 0
        )
        qtile.root.send_event(event, mask=EventMask.StructureNotify)

        # cleanup before exit
        atexit.register(self.cleanup)

    def draw(self):
        self.drawer.clear(self.background or self.bar.background)
        self.drawer.draw(self.offset, self.calculate_width())
        xoffset = self.padding
        for pos, icon in enumerate(self.icons.values()):
            icon.place(
                self.offset + xoffset,
                self.bar.height / 2 - self.icon_size / 2,
                icon.width, self.icon_size,
                0,
                None
            )
            xoffset += icon.width + self.padding

    def cleanup(self):
        atoms = self.qtile.conn.atoms
        self.qtile.conn.conn.core.SetSelectionOwner(
            0,
            atoms['_NET_SYSTEM_TRAY_S0'],
            xcffib.CurrentTime,
        )
        self.traywin.hide()
<commit_msg>fix systray integer division<commit_after>from __future__ import division

from .. import bar, xcbq, window
from . import base

import xcffib
from xcffib.xproto import EventMask, SetMode
import atexit
import struct

class Icon(window._Window):
    _windowMask = EventMask.StructureNotify | \
        EventMask.Exposure

    def __init__(self, win, qtile, systray):
        window._Window.__init__(self, win, qtile)
        self.systray = systray
        self.width = systray.icon_size
        self.height = systray.icon_size

    def handle_ConfigureNotify(self, event):
        icon_size = self.systray.icon_size
        self.updateHints()

        try:
            width = self.hints["min_width"]
            height = self.hints["min_height"]
        except KeyError:
            width = icon_size
            height = icon_size

        if height > icon_size:
            width = width * icon_size // height
            height = icon_size
        if height <= 0:
            width = icon_size
            height = icon_size

        self.width = width
        self.height = height
        self.window.set_attribute(backpixmap=self.systray.drawer.pixmap)
        return False

    def handle_DestroyNotify(self, event):
        wid = event.window
        del(self.qtile.windowMap[wid])
        del(self.systray.icons[wid])
        self.systray.draw()
        return False

    handle_UnmapNotify = handle_DestroyNotify

class TrayWindow(window._Window):
    _windowMask = EventMask.StructureNotify | \
        EventMask.Exposure

    def __init__(self, win, qtile, systray):
        window._Window.__init__(self, win, qtile)
        self.systray = systray

    def handle_ClientMessage(self, event):
        atoms = self.qtile.conn.atoms

        opcode = event.type
        wid = event.window
        data = event.data.data32
        message = data[1]

        conn = self.qtile.conn.conn
        parent = self.systray.bar.window.window

        # message == 0 corresponds to SYSTEM_TRAY_REQUEST_DOCK
        # TODO: handle system tray messages http://standards.freedesktop.org/systemtray-spec/systemtray-spec-latest.html
        # TODO: re-add 'and message == 0' constraint which /should/ work but doesn't with xcffib
        if opcode == atoms['_NET_SYSTEM_TRAY_OPCODE']:
            try:
                w = xcbq.Window(self.qtile.conn, wid)
                icon = Icon(w, self.qtile, self.systray)
                self.systray.icons[wid] = icon
                self.qtile.windowMap[wid] = icon

                # add icon window to the save-set, so it gets reparented
                # to the root window when qtile dies
                # TODO: figure out why this gives a BadMatch with xcffib.
                # conn.core.ChangeSaveSet(SetMode.Insert, wid)

                conn.core.ReparentWindow(wid, parent.wid, 0, 0)
                conn.flush()
                w.map()
            except xcffib.xproto.DrawableError:
                # The icon wasn't ready to be drawn yet... (NetworkManager does
                # this sometimes), so we just forget about it and wait for the
                # next event.
                pass
        return False

class Systray(base._Widget):
    """
        A widget that manage system tray
    """
    defaults = [
        ('icon_size', 20, 'Icon width'),
        ('padding', 5, 'Padding between icons'),
    ]

    def __init__(self, **config):
        base._Widget.__init__(self, bar.CALCULATED, **config)
        self.add_defaults(Systray.defaults)
        self.traywin = None
        self.icons = {}

    def button_press(self, x, y, button):
        pass

    def calculate_width(self):
        width = sum([i.width for i in self.icons.values()])
        width += self.padding * len(self.icons)
        return width

    def _configure(self, qtile, bar):
        base._Widget._configure(self, qtile, bar)
        self.qtile = qtile
        self.bar = bar
        atoms = qtile.conn.atoms
        win = qtile.conn.create_window(-1, -1, 1, 1)
        self.traywin = TrayWindow(win, self.qtile, self)
        qtile.windowMap[win.wid] = self.traywin
        qtile.conn.conn.core.SetSelectionOwner(
            win.wid,
            atoms['_NET_SYSTEM_TRAY_S0'],
            xcffib.CurrentTime
        )
        event = struct.pack(
            'BBHII5I', 33, 32, 0, qtile.root.wid,
            atoms['MANAGER'],
            xcffib.CurrentTime, atoms['_NET_SYSTEM_TRAY_S0'],
            win.wid, 0, 0
        )
        qtile.root.send_event(event, mask=EventMask.StructureNotify)

        # cleanup before exit
        atexit.register(self.cleanup)

    def draw(self):
        self.drawer.clear(self.background or self.bar.background)
        self.drawer.draw(self.offset, self.calculate_width())
        xoffset = self.padding
        for pos, icon in enumerate(self.icons.values()):
            icon.place(
                self.offset + xoffset,
                self.bar.height // 2 - self.icon_size // 2,
                icon.width, self.icon_size,
                0,
                None
            )
            xoffset += icon.width + self.padding

    def cleanup(self):
        atoms = self.qtile.conn.atoms
        self.qtile.conn.conn.core.SetSelectionOwner(
            0,
            atoms['_NET_SYSTEM_TRAY_S0'],
            xcffib.CurrentTime,
        )
        self.traywin.hide()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

from os import path

from ..platformio.project_recognition import ProjectRecognition
from ..libraries.quick_menu import QuickMenu

class ProjectCheck(QuickMenu):
    """
    ProjectCheck handles the actions between sublime text and platformio.
    Before run a platformio command like initilize, compile or upload, this
    class check if the project meets the requirements to proceed with the
    command, for example if the current file has been saved, or if it's saved
    is in the src folder when the platformio sutrcture options is marked
    """
    def __init__(self):
        super(ProjectCheck, self).__init__()

        self.board_id = None
        self.port_id = None

    def is_iot(self):
        """IOT

        Checks if the file in the current view is in the list
        of the IOT types (accepted) or not

        Returns:
            bool -- true if is in the list false if not
        """
        ext = self.get_file_extension()
        accepted = ['ino', 'pde', 'cpp', 'c', '.S']

        if(ext not in accepted):
            return False
        return True

    def is_empty(self):
        """Empty File

        Checks if the file is empty or not

        Returns:
            bool -- true is if empty
        """
        size = self.view.size()

        if(size > 0):
            return False
        return True

    def is_unsaved(self):
        """Unsaved View

        Check if the view has unsaved changes

        Returns:
            bool -- True if it's unsaved
        """
        return self.view.is_dirty()

    def structurize_project(self):
        """Structure Files

        If a project isn't initialized, it need to be checked
        if the open file is inside of the src folder, if it isn't
        the file need to be moved to the src folder
        """
        pio_structure = self.get_structure_option()

        if(pio_structure):
            file_path = self.get_file_path()
            if('src' not in file_path):
                from shutil import move

                self.close_file()

                dst = add_folder_to_filepath(file_path, 'src')
                move(file_path, dst)

                self.window.open_file(dst)
                return

        self.override_src()

    def override_src(self):
        """Adds src_dir

        When you don't want to keep the platformio file structure, you need to add
        the 'src_dir' flag in the platformio.ini with the path of your sketch/project.
        Here we add that option when platformio structure is not enabled
        """
        from ..libraries.configobj.configobj import ConfigObj

        source = {'src_dir': self.get_project_path()}
        ini_path = self.get_ini_path()
        config = ConfigObj(ini_path)

        config['platformio'] = source

        config.write()

    def close_file(self):
        """Close File Window

        Close the current focused windows in sublime text
        """
        self.window.run_command('close_file')

    def check_board_selected(self):
        """Checks Board Selection

        If an environment is stored in the preferences file, it will
        be loaded in the board_id object, if not, it will show the
        quick panel to select the board
        """
        self.board_id = self.get_environment()

        if(not self.board_id):
            selected_boards = self.get_selected_boards()

            if(not selected_boards):
                QuickMenu().quick_boards()
                return

            QuickMenu().quick_environments()
            return

    def check_port_selected(self):
        """Checks Serial Port Selection

        If the serial port is stored in the preferences file, it will
        be loaded in the port_id object, if not, it will show the
        quick panel to select the port
        """
        self.port_id = self.get_serial_port()

        if(not self.port_id):
            QuickMenu().quick_serial_ports()

def add_folder_to_filepath(src_path, new_folder):
    """Add folder

    Add a new folder at the end of the given specific path

    Arguments:
        src_path {str} -- initial path including the filename
        new_folder {str} -- string to add after the last folder in the path

    Returns:
        str -- file path with the new folder added
    """

    folder = path.dirname(src_path)
    file_name = path.basename(src_path)
    new_path = path.join(folder, new_folder, file_name)

    return new_path<commit_msg>Improved serial ports check when a port is not available anymore<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division
from __future__ import unicode_literals

from os import path

from ..platformio.project_recognition import ProjectRecognition
from .serial import serial_port_list
from .quick_menu import QuickMenu

class ProjectCheck(QuickMenu):
    """
    ProjectCheck handles the actions between sublime text and platformio.
    Before run a platformio command like initilize, compile or upload, this
    class check if the project meets the requirements to proceed with the
    command, for example if the current file has been saved, or if it's saved
    is in the src folder when the platformio sutrcture options is marked
    """
    def __init__(self):
        super(ProjectCheck, self).__init__()

        self.board_id = None
        self.port_id = None

    def is_iot(self):
        """IOT

        Checks if the file in the current view is in the list
        of the IOT types (accepted) or not

        Returns:
            bool -- true if is in the list false if not
        """
        ext = self.get_file_extension()
        accepted = ['ino', 'pde', 'cpp', 'c', '.S']

        if(ext not in accepted):
            return False
        return True

    def is_empty(self):
        """Empty File

        Checks if the file is empty or not

        Returns:
            bool -- true is if empty
        """
        size = self.view.size()

        if(size > 0):
            return False
        return True

    def is_unsaved(self):
        """Unsaved View

        Check if the view has unsaved changes

        Returns:
            bool -- True if it's unsaved
        """
        return self.view.is_dirty()

    def structurize_project(self):
        """Structure Files

        If a project isn't initialized, it need to be checked
        if the open file is inside of the src folder, if it isn't
        the file need to be moved to the src folder
        """
        pio_structure = self.get_structure_option()

        if(pio_structure):
            file_path = self.get_file_path()
            if('src' not in file_path):
                from shutil import move

                self.close_file()

                dst = add_folder_to_filepath(file_path, 'src')
                move(file_path, dst)

                self.window.open_file(dst)
                return

        self.override_src()

    def override_src(self):
        """Adds src_dir

        When you don't want to keep the platformio file structure, you need to add
        the 'src_dir' flag in the platformio.ini with the path of your sketch/project.
        Here we add that option when platformio structure is not enabled
        """
        from ..libraries.configobj.configobj import ConfigObj

        source = {'src_dir': self.get_project_path()}
        ini_path = self.get_ini_path()
        config = ConfigObj(ini_path)

        config['platformio'] = source

        config.write()

    def close_file(self):
        """Close File Window

        Close the current focused windows in sublime text
        """
        self.window.run_command('close_file')

    def check_board_selected(self):
        """Checks Board Selection

        If an environment is stored in the preferences file, it will
        be loaded in the board_id object, if not, it will show the
        quick panel to select the board
        """
        self.board_id = self.get_environment()

        if(not self.board_id):
            selected_boards = self.get_selected_boards()

            if(not selected_boards):
                QuickMenu().quick_boards()
                return

            QuickMenu().quick_environments()
            return

    def check_port_selected(self):
        """Checks Serial Port Selection

        If the serial port is stored in the preferences file, it will
        be loaded in the port_id object, if not, it will show the
        quick panel to select the port
        """
        self.port_id = self.get_serial_port()
        ports_list = serial_port_list()

        if(not self.port_id or self.port_id not in ports_list):
            QuickMenu().quick_serial_ports()
            self.port_id = None

def add_folder_to_filepath(src_path, new_folder):
    """Add folder

    Add a new folder at the end of the given specific path

    Arguments:
        src_path {str} -- initial path including the filename
        new_folder {str} -- string to add after the last folder in the path

    Returns:
        str -- file path with the new folder added
    """

    folder = path.dirname(src_path)
    file_name = path.basename(src_path)
    new_path = path.join(folder, new_folder, file_name)

    return new_path<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import uuid
from collections import defaultdict

from django.db import OperationalError

try:
    from random import choices
except ImportError:
    from corehq.sql_db.backports import choices

from django.conf import settings
from django import db
from django.db.utils import InterfaceError as DjangoInterfaceError
from functools import wraps
from psycopg2._psycopg import InterfaceError as Psycopg2InterfaceError
import six
from memoized import memoized

from corehq.sql_db.config import partition_config
from corehq.util.quickcache import quickcache

ACCEPTABLE_STANDBY_DELAY_SECONDS = 3
STALE_CHECK_FREQUENCY = 30

def paginate_query_across_partitioned_databases(model_class, q_expression, annotate=None, query_size=5000,
                                                values=None):
    """
    Runs a query across all partitioned databases in small chunks and produces a generator
    with the results.

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param annotate: (optional) If specified, should be a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :param values: (optional) If specified, should be a list of values to retrieve rather
    than retrieving entire objects.

    :return: A generator with the results
    """
    db_names = get_db_aliases_for_partitioned_query()
    for db_name in db_names:
        for row in paginate_query(db_name, model_class, q_expression, annotate, query_size, values):
            yield row

def paginate_query(db_name, model_class, q_expression, annotate=None, query_size=5000, values=None):
    """
    Runs a query on the given database in small chunks and produces a generator
    with the results.

    Iteration logic adopted from https://djangosnippets.org/snippets/1949/

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param annotate: (optional) If specified, should be a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :param values: (optional) If specified, should be a list of values to retrieve rather
    than retrieving entire objects.

    :return: A generator with the results
    """
    sort_col = 'pk'

    return_values = None
    if values:
        return_values = [sort_col] + values

    qs = model_class.objects.using(db_name)
    if annotate:
        qs = qs.annotate(**annotate)

    qs = qs.filter(q_expression).order_by(sort_col)

    if return_values:
        qs = qs.values_list(*return_values)
    value = 0

    filter_expression = {}
    while True:
        results = qs.filter(**filter_expression)[:query_size]
        if len(results) == 0:
            break

        for row in results:
            if return_values:
                value = row[0]
                yield row[1:]
            else:
                value = row.pk
                yield row
        filter_expression = {'{}__gt'.format(sort_col): value}

def split_list_by_db_partition(partition_values):
    """
    :param partition_values: Iterable of partition values (e.g. case IDs)
    :return: list of tuples (db_name, list(partition_values))
    """
    mapping = defaultdict(list)
    for value in partition_values:
        db_name = get_db_alias_for_partitioned_doc(value)
        mapping[db_name].append(value)
    return list(mapping.items())

def get_db_alias_for_partitioned_doc(partition_value):
    if settings.USE_PARTITIONED_DATABASE:
        from corehq.form_processor.backends.sql.dbaccessors import ShardAccessor
        db_name = ShardAccessor.get_database_for_doc(partition_value)
    else:
        db_name = 'default'
    return db_name

def get_db_aliases_for_partitioned_query():
    if settings.USE_PARTITIONED_DATABASE:
        db_names = partition_config.get_form_processing_dbs()
    else:
        db_names = ['default']
    return db_names

def get_default_db_aliases():
    return ['default']

def get_all_db_aliases():
    return list(settings.DATABASES)

def get_default_and_partitioned_db_aliases():
    return list(set(get_db_aliases_for_partitioned_query() + get_default_db_aliases()))

def new_id_in_same_dbalias(partition_value):
    """
    Returns a new partition value that belongs to the same db alias as
        the given partition value does
    """
    old_db_name = get_db_alias_for_partitioned_doc(partition_value)
    new_db_name = None
    while old_db_name != new_db_name:
        # todo; guard against infinite recursion
        new_partition_value = six.text_type(uuid.uuid4())
        new_db_name = get_db_alias_for_partitioned_doc(new_partition_value)
    return new_partition_value

def handle_connection_failure(get_db_aliases=get_default_db_aliases):
    def _inner2(fn):
        @wraps(fn)
        def _inner(*args, **kwargs):
            try:
                return fn(*args, **kwargs)
            except db.utils.DatabaseError:
                # we have to do this manually to avoid issues with
                # open transactions and already closed connections
                for db_name in get_db_aliases():
                    db.transaction.rollback(using=db_name)

                # re raise the exception for additional error handling
                raise
            except (Psycopg2InterfaceError, DjangoInterfaceError):
                # force closing the connection to prevent Django from trying to reuse it.
                # http://www.tryolabs.com/Blog/2014/02/12/long-time-running-process-and-django-orm/
                for db_name in get_db_aliases():
                    db.connections[db_name].close()

                # re raise the exception for additional error handling
                raise

        return _inner

    return _inner2

def get_all_sharded_models():
    from corehq.sql_db.models import PartitionedModel
    for subclass in _get_all_nested_subclasses(PartitionedModel):
        if not subclass._meta.abstract:
            yield subclass

def _get_all_nested_subclasses(cls):
    seen = set()
    for subclass in cls.__subclasses__():
        for sub_subclass in _get_all_nested_subclasses(subclass):
            # in case of multiple inheritance
            if sub_subclass not in seen:
                seen.add(sub_subclass)
                yield sub_subclass
        if subclass not in seen:
            seen.add(subclass)
            yield subclass

@memoized
def get_standby_databases():
    standby_dbs = []
    for db_alias in settings.DATABASES:
        with db.connections[db_alias].cursor() as cursor:
            cursor.execute("SELECT pg_is_in_recovery()")
            [(is_standby, )] = cursor.fetchall()
            if is_standby:
                standby_dbs.append(db_alias)
    return standby_dbs

def get_replication_delay_for_standby(db_alias):
    """
    Finds the replication delay for given database by running a SQL query on standby database.
        See https://www.postgresql.org/message-id/CADKbJJWz9M0swPT3oqe8f9+tfD4-F54uE6Xtkh4nERpVsQnjnw@mail.gmail.com

    If the given database is not a standby database, zero delay is returned
    If standby process (wal_receiver) is not running on standby a `VERY_LARGE_DELAY` is returned
    """
    if db_alias not in get_standby_databases():
        return 0
    # used to indicate that the wal_receiver process on standby is not running
    VERY_LARGE_DELAY = 100000
    try:
        sql = """
        SELECT
        CASE
            WHEN NOT EXISTS (SELECT 1 FROM pg_stat_wal_receiver) THEN %(delay)s
            WHEN pg_last_xlog_receive_location() = pg_last_xlog_replay_location() THEN 0
            ELSE EXTRACT (EPOCH FROM now() - pg_last_xact_replay_timestamp())::INTEGER
        END
        AS replication_lag;
        """
        with db.connections[db_alias].cursor() as cursor:
            cursor.execute(sql, {'delay': VERY_LARGE_DELAY})
            [(delay, )] = cursor.fetchall()
            return delay
    except OperationalError:
        return VERY_LARGE_DELAY

@memoized
def get_standby_delays_by_db():
    ret = {}
    for _db, config in six.iteritems(settings.DATABASES):
        delay = config.get('HQ_ACCEPTABLE_STANDBY_DELAY')
        if delay:
            ret[_db] = delay
    return ret

@quickcache(['dbs'], timeout=STALE_CHECK_FREQUENCY, skip_arg=lambda *args: settings.UNIT_TESTING)
def filter_out_stale_standbys(dbs):
    # from given list of databases filters out those with more than
    #   acceptable standby delay, if that database is a standby
    delays_by_db = get_standby_delays_by_db()
    return [
        db
        for db in dbs
        if get_replication_delay_for_standby(db) <= delays_by_db.get(db, ACCEPTABLE_STANDBY_DELAY_SECONDS)
    ]

def select_db_for_read(weighted_dbs):
    """
    Returns a randomly selected database per the weights assigned from
        a list of databases. If any database is standby and its replication has
        more than acceptable delay, that db is dropped from selection

    Args:
        weighted_dbs: a list of tuple of db and the weight.
            [
                ("pgmain", 5),
                ("pgmainstandby", 5)
            ]

    """
    # convert to a db to weight dictionary
    weights_by_db = {_db: weight for _db, weight in weighted_dbs}

    # filter out stale standby dbs
    fresh_dbs = filter_out_stale_standbys(weights_by_db)
    dbs = []
    weights = []
    for _db, weight in six.iteritems(weights_by_db):
        if _db in fresh_dbs:
            dbs.append(_db)
            weights.append(weight)

    if dbs:
        return choices(dbs, weights=weights)[0]
<commit_msg>Remove unnecessary value<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import uuid
from collections import defaultdict

from django.db import OperationalError

try:
    from random import choices
except ImportError:
    from corehq.sql_db.backports import choices

from django.conf import settings
from django import db
from django.db.utils import InterfaceError as DjangoInterfaceError
from functools import wraps
from psycopg2._psycopg import InterfaceError as Psycopg2InterfaceError
import six
from memoized import memoized

from corehq.sql_db.config import partition_config
from corehq.util.quickcache import quickcache

ACCEPTABLE_STANDBY_DELAY_SECONDS = 3
STALE_CHECK_FREQUENCY = 30

def paginate_query_across_partitioned_databases(model_class, q_expression, annotate=None, query_size=5000,
                                                values=None):
    """
    Runs a query across all partitioned databases in small chunks and produces a generator
    with the results.

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param annotate: (optional) If specified, should be a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :param values: (optional) If specified, should be a list of values to retrieve rather
    than retrieving entire objects.

    :return: A generator with the results
    """
    db_names = get_db_aliases_for_partitioned_query()
    for db_name in db_names:
        for row in paginate_query(db_name, model_class, q_expression, annotate, query_size, values):
            yield row

def paginate_query(db_name, model_class, q_expression, annotate=None, query_size=5000, values=None):
    """
    Runs a query on the given database in small chunks and produces a generator
    with the results.

    Iteration logic adopted from https://djangosnippets.org/snippets/1949/

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param annotate: (optional) If specified, should be a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :param values: (optional) If specified, should be a list of values to retrieve rather
    than retrieving entire objects.

    :return: A generator with the results
    """
    sort_col = 'pk'

    return_values = None
    if values:
        return_values = [sort_col] + values

    qs = model_class.objects.using(db_name)
    if annotate:
        qs = qs.annotate(**annotate)

    qs = qs.filter(q_expression).order_by(sort_col)

    if return_values:
        qs = qs.values_list(*return_values)

    filter_expression = {}
    while True:
        results = qs.filter(**filter_expression)[:query_size]
        if len(results) == 0:
            break

        for row in results:
            if return_values:
                value = row[0]
                yield row[1:]
            else:
                value = row.pk
                yield row
        filter_expression = {'{}__gt'.format(sort_col): value}

def split_list_by_db_partition(partition_values):
    """
    :param partition_values: Iterable of partition values (e.g. case IDs)
    :return: list of tuples (db_name, list(partition_values))
    """
    mapping = defaultdict(list)
    for value in partition_values:
        db_name = get_db_alias_for_partitioned_doc(value)
        mapping[db_name].append(value)
    return list(mapping.items())

def get_db_alias_for_partitioned_doc(partition_value):
    if settings.USE_PARTITIONED_DATABASE:
        from corehq.form_processor.backends.sql.dbaccessors import ShardAccessor
        db_name = ShardAccessor.get_database_for_doc(partition_value)
    else:
        db_name = 'default'
    return db_name

def get_db_aliases_for_partitioned_query():
    if settings.USE_PARTITIONED_DATABASE:
        db_names = partition_config.get_form_processing_dbs()
    else:
        db_names = ['default']
    return db_names

def get_default_db_aliases():
    return ['default']

def get_all_db_aliases():
    return list(settings.DATABASES)

def get_default_and_partitioned_db_aliases():
    return list(set(get_db_aliases_for_partitioned_query() + get_default_db_aliases()))

def new_id_in_same_dbalias(partition_value):
    """
    Returns a new partition value that belongs to the same db alias as
        the given partition value does
    """
    old_db_name = get_db_alias_for_partitioned_doc(partition_value)
    new_db_name = None
    while old_db_name != new_db_name:
        # todo; guard against infinite recursion
        new_partition_value = six.text_type(uuid.uuid4())
        new_db_name = get_db_alias_for_partitioned_doc(new_partition_value)
    return new_partition_value

def handle_connection_failure(get_db_aliases=get_default_db_aliases):
    def _inner2(fn):
        @wraps(fn)
        def _inner(*args, **kwargs):
            try:
                return fn(*args, **kwargs)
            except db.utils.DatabaseError:
                # we have to do this manually to avoid issues with
                # open transactions and already closed connections
                for db_name in get_db_aliases():
                    db.transaction.rollback(using=db_name)

                # re raise the exception for additional error handling
                raise
            except (Psycopg2InterfaceError, DjangoInterfaceError):
                # force closing the connection to prevent Django from trying to reuse it.
                # http://www.tryolabs.com/Blog/2014/02/12/long-time-running-process-and-django-orm/
                for db_name in get_db_aliases():
                    db.connections[db_name].close()

                # re raise the exception for additional error handling
                raise

        return _inner

    return _inner2

def get_all_sharded_models():
    from corehq.sql_db.models import PartitionedModel
    for subclass in _get_all_nested_subclasses(PartitionedModel):
        if not subclass._meta.abstract:
            yield subclass

def _get_all_nested_subclasses(cls):
    seen = set()
    for subclass in cls.__subclasses__():
        for sub_subclass in _get_all_nested_subclasses(subclass):
            # in case of multiple inheritance
            if sub_subclass not in seen:
                seen.add(sub_subclass)
                yield sub_subclass
        if subclass not in seen:
            seen.add(subclass)
            yield subclass

@memoized
def get_standby_databases():
    standby_dbs = []
    for db_alias in settings.DATABASES:
        with db.connections[db_alias].cursor() as cursor:
            cursor.execute("SELECT pg_is_in_recovery()")
            [(is_standby, )] = cursor.fetchall()
            if is_standby:
                standby_dbs.append(db_alias)
    return standby_dbs

def get_replication_delay_for_standby(db_alias):
    """
    Finds the replication delay for given database by running a SQL query on standby database.
        See https://www.postgresql.org/message-id/CADKbJJWz9M0swPT3oqe8f9+tfD4-F54uE6Xtkh4nERpVsQnjnw@mail.gmail.com

    If the given database is not a standby database, zero delay is returned
    If standby process (wal_receiver) is not running on standby a `VERY_LARGE_DELAY` is returned
    """
    if db_alias not in get_standby_databases():
        return 0
    # used to indicate that the wal_receiver process on standby is not running
    VERY_LARGE_DELAY = 100000
    try:
        sql = """
        SELECT
        CASE
            WHEN NOT EXISTS (SELECT 1 FROM pg_stat_wal_receiver) THEN %(delay)s
            WHEN pg_last_xlog_receive_location() = pg_last_xlog_replay_location() THEN 0
            ELSE EXTRACT (EPOCH FROM now() - pg_last_xact_replay_timestamp())::INTEGER
        END
        AS replication_lag;
        """
        with db.connections[db_alias].cursor() as cursor:
            cursor.execute(sql, {'delay': VERY_LARGE_DELAY})
            [(delay, )] = cursor.fetchall()
            return delay
    except OperationalError:
        return VERY_LARGE_DELAY

@memoized
def get_standby_delays_by_db():
    ret = {}
    for _db, config in six.iteritems(settings.DATABASES):
        delay = config.get('HQ_ACCEPTABLE_STANDBY_DELAY')
        if delay:
            ret[_db] = delay
    return ret

@quickcache(['dbs'], timeout=STALE_CHECK_FREQUENCY, skip_arg=lambda *args: settings.UNIT_TESTING)
def filter_out_stale_standbys(dbs):
    # from given list of databases filters out those with more than
    #   acceptable standby delay, if that database is a standby
    delays_by_db = get_standby_delays_by_db()
    return [
        db
        for db in dbs
        if get_replication_delay_for_standby(db) <= delays_by_db.get(db, ACCEPTABLE_STANDBY_DELAY_SECONDS)
    ]

def select_db_for_read(weighted_dbs):
    """
    Returns a randomly selected database per the weights assigned from
        a list of databases. If any database is standby and its replication has
        more than acceptable delay, that db is dropped from selection

    Args:
        weighted_dbs: a list of tuple of db and the weight.
            [
                ("pgmain", 5),
                ("pgmainstandby", 5)
            ]

    """
    # convert to a db to weight dictionary
    weights_by_db = {_db: weight for _db, weight in weighted_dbs}

    # filter out stale standby dbs
    fresh_dbs = filter_out_stale_standbys(weights_by_db)
    dbs = []
    weights = []
    for _db, weight in six.iteritems(weights_by_db):
        if _db in fresh_dbs:
            dbs.append(_db)
            weights.append(weight)

    if dbs:
        return choices(dbs, weights=weights)[0]
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import uuid
from collections import defaultdict
try:
    from random import choices
except ImportError:
    from corehq.sql_db.backports import choices

from django.conf import settings
from django import db
from django.db.utils import InterfaceError as DjangoInterfaceError
from functools import wraps
from psycopg2._psycopg import InterfaceError as Psycopg2InterfaceError
import six
from memoized import memoized

from corehq.sql_db.config import partition_config
from corehq.util.quickcache import quickcache

ACCEPTABLE_STANDBY_DELAY_SECONDS = 3
STALE_CHECK_FREQUENCY = 30

def run_query_across_partitioned_databases(model_class, q_expression, values=None, annotate=None):
    """
    Runs a query across all partitioned databases and produces a generator
    with the results.

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param values: (optional) If specified, should be a list of values to retrieve rather
    than retrieving entire objects. If a list with a single value is given, the result will
    be a generator of single values. If a list with multiple values is given, the result
    will be a generator of tuples.

    :param annotate: (optional) If specified, should by a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :return: A generator with the results
    """
    db_names = get_db_aliases_for_partitioned_query()

    if values and not isinstance(values, (list, tuple)):
        raise ValueError("Expected a list or tuple")

    for db_name in db_names:
        qs = model_class.objects.using(db_name)
        if annotate:
            qs = qs.annotate(**annotate)

        qs = qs.filter(q_expression)
        if values:
            if len(values) == 1:
                qs = qs.values_list(*values, flat=True)
            else:
                qs = qs.values_list(*values)

        for result in qs.iterator():
            yield result

def paginate_query_across_partitioned_databases(model_class, q_expression, annotate=None, query_size=5000):
    """
    Runs a query across all partitioned databases in small chunks and produces a generator
    with the results.

    Iteration logic adopted from https://djangosnippets.org/snippets/1949/

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param annotate: (optional) If specified, should by a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :return: A generator with the results
    """
    db_names = get_db_aliases_for_partitioned_query()

    for db_name in db_names:
        qs = model_class.objects.using(db_name)
        if annotate:
            qs = qs.annotate(**annotate)

        qs = qs.filter(q_expression)
        sort_col = 'pk'
        value = 0
        last_value = qs.order_by('-{}'.format(sort_col)).values_list(sort_col, flat=True).first()
        if last_value is not None:
            qs = qs.order_by(sort_col)
            while value < last_value:
                filter_expression = {'{}__gt'.format(sort_col): value}
                for row in qs.filter(**filter_expression)[:query_size]:
                    value = row.pk
                    yield row

def split_list_by_db_partition(partition_values):
    """
    :param partition_values: Iterable of partition values (e.g. case IDs)
    :return: list of tuples (db_name, list(partition_values))
    """
    mapping = defaultdict(list)
    for value in partition_values:
        db_name = get_db_alias_for_partitioned_doc(value)
        mapping[db_name].append(value)
    return list(mapping.items())

def get_db_alias_for_partitioned_doc(partition_value):
    if settings.USE_PARTITIONED_DATABASE:
        from corehq.form_processor.backends.sql.dbaccessors import ShardAccessor
        db_name = ShardAccessor.get_database_for_doc(partition_value)
    else:
        db_name = 'default'
    return db_name

def get_db_aliases_for_partitioned_query():
    if settings.USE_PARTITIONED_DATABASE:
        db_names = partition_config.get_form_processing_dbs()
    else:
        db_names = ['default']
    return db_names

def get_default_db_aliases():
    return ['default']

def get_all_db_aliases():
    return list(settings.DATABASES)

def get_default_and_partitioned_db_aliases():
    return list(set(get_db_aliases_for_partitioned_query() + get_default_db_aliases()))

def new_id_in_same_dbalias(partition_value):
    """
    Returns a new partition value that belongs to the same db alias as
        the given partition value does
    """
    old_db_name = get_db_alias_for_partitioned_doc(partition_value)
    new_db_name = None
    while old_db_name != new_db_name:
        # todo; guard against infinite recursion
        new_partition_value = six.text_type(uuid.uuid4())
        new_db_name = get_db_alias_for_partitioned_doc(new_partition_value)
    return new_partition_value

def handle_connection_failure(get_db_aliases=get_default_db_aliases):
    def _inner2(fn):
        @wraps(fn)
        def _inner(*args, **kwargs):
            try:
                return fn(*args, **kwargs)
            except db.utils.DatabaseError:
                # we have to do this manually to avoid issues with
                # open transactions and already closed connections
                for db_name in get_db_aliases():
                    db.transaction.rollback(using=db_name)

                # re raise the exception for additional error handling
                raise
            except (Psycopg2InterfaceError, DjangoInterfaceError):
                # force closing the connection to prevent Django from trying to reuse it.
                # http://www.tryolabs.com/Blog/2014/02/12/long-time-running-process-and-django-orm/
                for db_name in get_db_aliases():
                    db.connections[db_name].close()

                # re raise the exception for additional error handling
                raise

        return _inner

    return _inner2

def get_all_sharded_models():
    from corehq.sql_db.models import PartitionedModel
    for subclass in _get_all_nested_subclasses(PartitionedModel):
        if not subclass._meta.abstract:
            yield subclass

def _get_all_nested_subclasses(cls):
    seen = set()
    for subclass in cls.__subclasses__():
        for sub_subclass in _get_all_nested_subclasses(subclass):
            # in case of multiple inheritance
            if sub_subclass not in seen:
                seen.add(sub_subclass)
                yield sub_subclass
        if subclass not in seen:
            seen.add(subclass)
            yield subclass

@memoized
def get_standby_databases():
    standby_dbs = []
    for db_alias in settings.DATABASES:
        with db.connections[db_alias].cursor() as cursor:
            cursor.execute("SELECT pg_is_in_recovery()")
            [(is_standby, )] = cursor.fetchall()
            if is_standby:
                standby_dbs.append(db_alias)
    return standby_dbs

def get_replication_delay_for_standby(db_alias):
    """
    Finds the replication delay for given database by running a SQL query on standby database.
        See https://www.postgresql.org/message-id/CADKbJJWz9M0swPT3oqe8f9+tfD4-F54uE6Xtkh4nERpVsQnjnw@mail.gmail.com

    If the given database is not a standby database, zero delay is returned
    If standby process (wal_receiver) is not running on standby a `VERY_LARGE_DELAY` is returned
    """
    if db_alias not in get_standby_databases():
        return 0
    # used to indicate that the wal_receiver process on standby is not running
    VERY_LARGE_DELAY = 100000
    sql = """
    SELECT
    CASE
        WHEN NOT EXISTS (SELECT 1 FROM pg_stat_wal_receiver) THEN {delay}
        WHEN pg_last_xlog_receive_location() = pg_last_xlog_replay_location() THEN 0
        ELSE EXTRACT (EPOCH FROM now() - pg_last_xact_replay_timestamp())::INTEGER
    END
    AS replication_lag;
    """.format(delay=VERY_LARGE_DELAY)
    with db.connections[db_alias].cursor() as cursor:
        cursor.execute(sql)
        [(delay, )] = cursor.fetchall()
        return delay

@memoized
def get_standby_delays_by_db():
    ret = {}
    for _db, config in six.iteritems(settings.DATABASES):
        delay = config.get('HQ_ACCEPTABLE_STANDBY_DELAY')
        if delay:
            ret[_db] = delay
    return ret

@quickcache(['dbs'], timeout=STALE_CHECK_FREQUENCY, skip_arg=lambda *args: settings.UNIT_TESTING)
def filter_out_stale_standbys(dbs):
    # from given list of databases filters out those with more than
    #   acceptable standby delay, if that database is a standby
    delays_by_db = get_standby_delays_by_db()
    return [
        db
        for db in dbs
        if get_replication_delay_for_standby(db) <= delays_by_db.get(db, ACCEPTABLE_STANDBY_DELAY_SECONDS)
    ]

def select_db_for_read(weighted_dbs):
    """
    Returns a randomly selected database per the weights assigned from
        a list of databases. If any database is standby and its replication has
        more than accesptable delay, that db is dropped from selection

    Args:
        weighted_dbs: a list of tuple of db and the weight.
            [
                ("pgmain", 5),
                ("pgmainstandby", 5)
            ]

    """
    # convert to a db to weight dictionary
    weights_by_db = {_db: weight for _db, weight in weighted_dbs}

    # filter out stale standby dbs
    fresh_dbs = filter_out_stale_standbys(weights_by_db)
    dbs = []
    weights = []
    for _db, weight in six.iteritems(weights_by_db):
        if _db in fresh_dbs:
            dbs.append(_db)
            weights.append(weight)

    if dbs:
        return choices(dbs, weights=weights)[0]
<commit_msg>fix typo<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import unicode_literals

import uuid
from collections import defaultdict
try:
    from random import choices
except ImportError:
    from corehq.sql_db.backports import choices

from django.conf import settings
from django import db
from django.db.utils import InterfaceError as DjangoInterfaceError
from functools import wraps
from psycopg2._psycopg import InterfaceError as Psycopg2InterfaceError
import six
from memoized import memoized

from corehq.sql_db.config import partition_config
from corehq.util.quickcache import quickcache

ACCEPTABLE_STANDBY_DELAY_SECONDS = 3
STALE_CHECK_FREQUENCY = 30

def run_query_across_partitioned_databases(model_class, q_expression, values=None, annotate=None):
    """
    Runs a query across all partitioned databases and produces a generator
    with the results.

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param values: (optional) If specified, should be a list of values to retrieve rather
    than retrieving entire objects. If a list with a single value is given, the result will
    be a generator of single values. If a list with multiple values is given, the result
    will be a generator of tuples.

    :param annotate: (optional) If specified, should by a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :return: A generator with the results
    """
    db_names = get_db_aliases_for_partitioned_query()

    if values and not isinstance(values, (list, tuple)):
        raise ValueError("Expected a list or tuple")

    for db_name in db_names:
        qs = model_class.objects.using(db_name)
        if annotate:
            qs = qs.annotate(**annotate)

        qs = qs.filter(q_expression)
        if values:
            if len(values) == 1:
                qs = qs.values_list(*values, flat=True)
            else:
                qs = qs.values_list(*values)

        for result in qs.iterator():
            yield result

def paginate_query_across_partitioned_databases(model_class, q_expression, annotate=None, query_size=5000):
    """
    Runs a query across all partitioned databases in small chunks and produces a generator
    with the results.

    Iteration logic adopted from https://djangosnippets.org/snippets/1949/

    :param model_class: A Django model class

    :param q_expression: An instance of django.db.models.Q representing the
    filter to apply

    :param annotate: (optional) If specified, should by a dictionary of annotated fields
    and their calculations. The dictionary will be splatted into the `.annotate` function

    :return: A generator with the results
    """
    db_names = get_db_aliases_for_partitioned_query()

    for db_name in db_names:
        qs = model_class.objects.using(db_name)
        if annotate:
            qs = qs.annotate(**annotate)

        qs = qs.filter(q_expression)
        sort_col = 'pk'
        value = 0
        last_value = qs.order_by('-{}'.format(sort_col)).values_list(sort_col, flat=True).first()
        if last_value is not None:
            qs = qs.order_by(sort_col)
            while value < last_value:
                filter_expression = {'{}__gt'.format(sort_col): value}
                for row in qs.filter(**filter_expression)[:query_size]:
                    value = row.pk
                    yield row

def split_list_by_db_partition(partition_values):
    """
    :param partition_values: Iterable of partition values (e.g. case IDs)
    :return: list of tuples (db_name, list(partition_values))
    """
    mapping = defaultdict(list)
    for value in partition_values:
        db_name = get_db_alias_for_partitioned_doc(value)
        mapping[db_name].append(value)
    return list(mapping.items())

def get_db_alias_for_partitioned_doc(partition_value):
    if settings.USE_PARTITIONED_DATABASE:
        from corehq.form_processor.backends.sql.dbaccessors import ShardAccessor
        db_name = ShardAccessor.get_database_for_doc(partition_value)
    else:
        db_name = 'default'
    return db_name

def get_db_aliases_for_partitioned_query():
    if settings.USE_PARTITIONED_DATABASE:
        db_names = partition_config.get_form_processing_dbs()
    else:
        db_names = ['default']
    return db_names

def get_default_db_aliases():
    return ['default']

def get_all_db_aliases():
    return list(settings.DATABASES)

def get_default_and_partitioned_db_aliases():
    return list(set(get_db_aliases_for_partitioned_query() + get_default_db_aliases()))

def new_id_in_same_dbalias(partition_value):
    """
    Returns a new partition value that belongs to the same db alias as
        the given partition value does
    """
    old_db_name = get_db_alias_for_partitioned_doc(partition_value)
    new_db_name = None
    while old_db_name != new_db_name:
        # todo; guard against infinite recursion
        new_partition_value = six.text_type(uuid.uuid4())
        new_db_name = get_db_alias_for_partitioned_doc(new_partition_value)
    return new_partition_value

def handle_connection_failure(get_db_aliases=get_default_db_aliases):
    def _inner2(fn):
        @wraps(fn)
        def _inner(*args, **kwargs):
            try:
                return fn(*args, **kwargs)
            except db.utils.DatabaseError:
                # we have to do this manually to avoid issues with
                # open transactions and already closed connections
                for db_name in get_db_aliases():
                    db.transaction.rollback(using=db_name)

                # re raise the exception for additional error handling
                raise
            except (Psycopg2InterfaceError, DjangoInterfaceError):
                # force closing the connection to prevent Django from trying to reuse it.
                # http://www.tryolabs.com/Blog/2014/02/12/long-time-running-process-and-django-orm/
                for db_name in get_db_aliases():
                    db.connections[db_name].close()

                # re raise the exception for additional error handling
                raise

        return _inner

    return _inner2

def get_all_sharded_models():
    from corehq.sql_db.models import PartitionedModel
    for subclass in _get_all_nested_subclasses(PartitionedModel):
        if not subclass._meta.abstract:
            yield subclass

def _get_all_nested_subclasses(cls):
    seen = set()
    for subclass in cls.__subclasses__():
        for sub_subclass in _get_all_nested_subclasses(subclass):
            # in case of multiple inheritance
            if sub_subclass not in seen:
                seen.add(sub_subclass)
                yield sub_subclass
        if subclass not in seen:
            seen.add(subclass)
            yield subclass

@memoized
def get_standby_databases():
    standby_dbs = []
    for db_alias in settings.DATABASES:
        with db.connections[db_alias].cursor() as cursor:
            cursor.execute("SELECT pg_is_in_recovery()")
            [(is_standby, )] = cursor.fetchall()
            if is_standby:
                standby_dbs.append(db_alias)
    return standby_dbs

def get_replication_delay_for_standby(db_alias):
    """
    Finds the replication delay for given database by running a SQL query on standby database.
        See https://www.postgresql.org/message-id/CADKbJJWz9M0swPT3oqe8f9+tfD4-F54uE6Xtkh4nERpVsQnjnw@mail.gmail.com

    If the given database is not a standby database, zero delay is returned
    If standby process (wal_receiver) is not running on standby a `VERY_LARGE_DELAY` is returned
    """
    if db_alias not in get_standby_databases():
        return 0
    # used to indicate that the wal_receiver process on standby is not running
    VERY_LARGE_DELAY = 100000
    sql = """
    SELECT
    CASE
        WHEN NOT EXISTS (SELECT 1 FROM pg_stat_wal_receiver) THEN {delay}
        WHEN pg_last_xlog_receive_location() = pg_last_xlog_replay_location() THEN 0
        ELSE EXTRACT (EPOCH FROM now() - pg_last_xact_replay_timestamp())::INTEGER
    END
    AS replication_lag;
    """.format(delay=VERY_LARGE_DELAY)
    with db.connections[db_alias].cursor() as cursor:
        cursor.execute(sql)
        [(delay, )] = cursor.fetchall()
        return delay

@memoized
def get_standby_delays_by_db():
    ret = {}
    for _db, config in six.iteritems(settings.DATABASES):
        delay = config.get('HQ_ACCEPTABLE_STANDBY_DELAY')
        if delay:
            ret[_db] = delay
    return ret

@quickcache(['dbs'], timeout=STALE_CHECK_FREQUENCY, skip_arg=lambda *args: settings.UNIT_TESTING)
def filter_out_stale_standbys(dbs):
    # from given list of databases filters out those with more than
    #   acceptable standby delay, if that database is a standby
    delays_by_db = get_standby_delays_by_db()
    return [
        db
        for db in dbs
        if get_replication_delay_for_standby(db) <= delays_by_db.get(db, ACCEPTABLE_STANDBY_DELAY_SECONDS)
    ]

def select_db_for_read(weighted_dbs):
    """
    Returns a randomly selected database per the weights assigned from
        a list of databases. If any database is standby and its replication has
        more than acceptable delay, that db is dropped from selection

    Args:
        weighted_dbs: a list of tuple of db and the weight.
            [
                ("pgmain", 5),
                ("pgmainstandby", 5)
            ]

    """
    # convert to a db to weight dictionary
    weights_by_db = {_db: weight for _db, weight in weighted_dbs}

    # filter out stale standby dbs
    fresh_dbs = filter_out_stale_standbys(weights_by_db)
    dbs = []
    weights = []
    for _db, weight in six.iteritems(weights_by_db):
        if _db in fresh_dbs:
            dbs.append(_db)
            weights.append(weight)

    if dbs:
        return choices(dbs, weights=weights)[0]
<|endoftext|>


<commit_before>import os
import shutil

from django.core.management.base import CommandError
from rest_framework import status
from rest_framework.test import APITestCase
from django.test import TestCase
from django.utils.six import StringIO

from countylimits.models import CountyLimit, County, State
from countylimits.management.commands.load_county_limits import Command

class CountyLimitTest(APITestCase):

    fixtures = ['states.json']

    def populate_db(self):
        """ Prepopulate DB with dummy data. """
        sDC = State.objects.get(state_abbr='DC')
        sVA = State.objects.get(state_abbr='VA')
        cDC1 = County(
            county_name='DC County 1',
            county_fips=333,
            state_id=sDC.id)
        cDC1.save()
        cDC2 = County(
            county_name='DC County 2',
            county_fips=444,
            state_id=sDC.id)
        cDC2.save()
        cVA1 = County(
            county_name='VA County 1',
            county_fips=555,
            state_id=sVA.id)
        cVA1.save()
        cl1 = CountyLimit(
            county=cDC1,
            fha_limit=1,
            gse_limit=1,
            va_limit=1)
        cl1.save()
        cl2 = CountyLimit(
            county=cDC2,
            fha_limit=10,
            gse_limit=10,
            va_limit=10)
        cl2.save()
        cl3 = CountyLimit(
            county=cVA1,
            fha_limit=100,
            gse_limit=100,
            va_limit=100)
        cl3.save()

    def setUp(self):
        self.url = '/oah-api/county/'
        self.populate_db()

    def test_county_limits_by_state__no_args(self):
        """ ... when state is blank """
        response = self.client.get(self.url, {})
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertEqual(
            response.data,
            {'detail': 'Required parameter state is missing'})

    def test_county_limit_by_state__invalid_arg(self):
        """ ... when state has an invalid value """
        response = self.client.get(self.url, {'state': 123})
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertEqual(response.data, {'state': 'Invalid state'})
        response = self.client.get(self.url, {'state': 'Washington DC'})
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertEqual(response.data, {'state': 'Invalid state'})

    def test_county_limit_by_state__valid_arg(self):
        """ ... when state has a valid arg """
        response_11 = self.client.get(self.url, {'state': 11})
        self.assertEqual(response_11.status_code, status.HTTP_200_OK)
        self.assertFalse(not response_11.data['data'])

        response_DC = self.client.get(self.url, {'state': 'DC'})
        self.assertEqual(len(response_11.data['data']), 2)
        self.assertTrue(response_11.data['data'] == response_DC.data['data'])

        response_VA = self.client.get(self.url, {'state': 'VA'})
        self.assertTrue(len(response_VA.data['data']) == 1)
        self.assertFalse(response_11.data['data'] == response_VA.data['data'])

    def test_unicode(self):
        state = State.objects.first()
        county = County.objects.first()
        county_limit = CountyLimit.objects.first()

        self.assertEqual('{}'.format(state), 'Alabama')
        self.assertEqual('{}'.format(county), 'DC County 1 (333)')
        self.assertEqual(
            '{}'.format(county_limit),
            'CountyLimit {}'.format(county_limit.id))

class LoadCountyLimitsTestCase(TestCase):

    def prepare_sample_data(self, filename):
        with open(filename, 'w') as data:
            data.write("Header to be skipped\n")
            data.write('DC,01,001,01001,DC County 1,417000,271050,417000\n')
            data.write('DC,01,002,01002,DC County 2,417000,271050,417000\n')
            data.close()

    def setUp(self):
        self.c = Command()
        self.out = StringIO()
        self.c.stdout = self.c.stderr = self.out

        self.test_dir = '{}/test_folder'.format(
            os.path.dirname(os.path.realpath(__file__))
            )
        os.mkdir(self.test_dir)
        os.chdir(self.test_dir)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_handle__no_confirm(self):
        """ .. check that nothing happens when confirm is not y|Y."""
        self.c.handle(stdout=self.out)
        self.assertNotIn(self.out.getvalue(), 'Successfully loaded data from')

    def test_handle__no_arguments(self):
        """ .. check that CommandError is raised."""
        self.assertRaises(
            CommandError,
            self.c.handle,
            confirmed='y',
            stdout=self.out)

    def test_handle__bad_file(self):
        """ .. check that CommandError is raised when path to file is wrong."""
        self.assertRaises(
            CommandError,
            self.c.handle,
            'inexistent.file.csv',
            confirmed='Y',
            stdout=self.out)

    def test_handle__success(self):
        """ .. check that all countylimits are loaded."""
        fname = 'county-limits.csv'
        self.prepare_sample_data(fname)
        self.c.handle('{}/{}'.format(self.test_dir, fname), confirmed='Y')
        self.assertIn(
            'Successfully loaded data from {}/{}'.format(
                self.test_dir, fname),
            self.out.getvalue()
        )
        county_limits = CountyLimit.objects.all()
        self.assertEqual(len(county_limits), 2)
<commit_msg>try Django test class instead of DRF<commit_after>import os
import shutil

from django.core.management.base import CommandError
from rest_framework import status
from rest_framework.test import APITestCase
from django.test import TestCase
from django.utils.six import StringIO

from countylimits.models import CountyLimit, County, State
from countylimits.management.commands.load_county_limits import Command

class CountyLimitTest(TestCase):

    fixtures = ['states.json']

    def populate_db(self):
        """ Prepopulate DB with dummy data. """
        sDC = State.objects.get(state_abbr='DC')
        sVA = State.objects.get(state_abbr='VA')
        cDC1 = County(
            county_name='DC County 1',
            county_fips=333,
            state_id=sDC.id)
        cDC1.save()
        cDC2 = County(
            county_name='DC County 2',
            county_fips=444,
            state_id=sDC.id)
        cDC2.save()
        cVA1 = County(
            county_name='VA County 1',
            county_fips=555,
            state_id=sVA.id)
        cVA1.save()
        cl1 = CountyLimit(
            county=cDC1,
            fha_limit=1,
            gse_limit=1,
            va_limit=1)
        cl1.save()
        cl2 = CountyLimit(
            county=cDC2,
            fha_limit=10,
            gse_limit=10,
            va_limit=10)
        cl2.save()
        cl3 = CountyLimit(
            county=cVA1,
            fha_limit=100,
            gse_limit=100,
            va_limit=100)
        cl3.save()

    def setUp(self):
        self.url = '/oah-api/county/'
        self.populate_db()

    def test_county_limits_by_state__no_args(self):
        """ ... when state is blank """
        response = self.client.get(self.url, {})
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertEqual(
            response.data,
            {'detail': 'Required parameter state is missing'})

    def test_county_limit_by_state__invalid_arg(self):
        """ ... when state has an invalid value """
        response = self.client.get(self.url, {'state': 123})
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertEqual(response.data, {'state': 'Invalid state'})
        response = self.client.get(self.url, {'state': 'Washington DC'})
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertEqual(response.data, {'state': 'Invalid state'})

    def test_county_limit_by_state__valid_arg(self):
        """ ... when state has a valid arg """
        response_11 = self.client.get(self.url, {'state': 11})
        self.assertEqual(response_11.status_code, status.HTTP_200_OK)
        self.assertFalse(not response_11.data['data'])

        response_DC = self.client.get(self.url, {'state': 'DC'})
        self.assertEqual(len(response_11.data['data']), 2)
        self.assertTrue(response_11.data['data'] == response_DC.data['data'])

        response_VA = self.client.get(self.url, {'state': 'VA'})
        self.assertTrue(len(response_VA.data['data']) == 1)
        self.assertFalse(response_11.data['data'] == response_VA.data['data'])

    def test_unicode(self):
        state = State.objects.first()
        county = County.objects.first()
        county_limit = CountyLimit.objects.first()

        self.assertEqual('{}'.format(state), 'Alabama')
        self.assertEqual('{}'.format(county), 'DC County 1 (333)')
        self.assertEqual(
            '{}'.format(county_limit),
            'CountyLimit {}'.format(county_limit.id))

class LoadCountyLimitsTestCase(TestCase):

    def prepare_sample_data(self, filename):
        with open(filename, 'w') as data:
            data.write("Header to be skipped\n")
            data.write('DC,01,001,01001,DC County 1,417000,271050,417000\n')
            data.write('DC,01,002,01002,DC County 2,417000,271050,417000\n')
            data.close()

    def setUp(self):
        self.c = Command()
        self.out = StringIO()
        self.c.stdout = self.c.stderr = self.out

        self.test_dir = '{}/test_folder'.format(
            os.path.dirname(os.path.realpath(__file__))
            )
        os.mkdir(self.test_dir)
        os.chdir(self.test_dir)

    def tearDown(self):
        shutil.rmtree(self.test_dir)

    def test_handle__no_confirm(self):
        """ .. check that nothing happens when confirm is not y|Y."""
        self.c.handle(stdout=self.out)
        self.assertNotIn(self.out.getvalue(), 'Successfully loaded data from')

    def test_handle__no_arguments(self):
        """ .. check that CommandError is raised."""
        self.assertRaises(
            CommandError,
            self.c.handle,
            confirmed='y',
            stdout=self.out)

    def test_handle__bad_file(self):
        """ .. check that CommandError is raised when path to file is wrong."""
        self.assertRaises(
            CommandError,
            self.c.handle,
            'inexistent.file.csv',
            confirmed='Y',
            stdout=self.out)

    def test_handle__success(self):
        """ .. check that all countylimits are loaded."""
        fname = 'county-limits.csv'
        self.prepare_sample_data(fname)
        self.c.handle('{}/{}'.format(self.test_dir, fname), confirmed='Y')
        self.assertIn(
            'Successfully loaded data from {}/{}'.format(
                self.test_dir, fname),
            self.out.getvalue()
        )
        county_limits = CountyLimit.objects.all()
        self.assertEqual(len(county_limits), 2)
<|endoftext|>