sqlx.py

# sqlx 是一种扩展 sql 的语言
# 目标是打造 "易读易写 方便维护" 的 sql 脚本
# 语法参考 test.sqlx

# Python3.3

import os
import sys
import re
import traceback
import pprint
import random
import copy


VERSION = '0.2.0'


# 构建后添加的头部文字
HEADER = '-- ======== Generated By Sqlx ========\n-- https://github.com/taojy123/sqlx'

# sqlx 语法注释标记
COMMENT_PREFIX = '-- !'

# 目前支持的关系运算符
OPERATORS = ['>', '<', '>=', '<=', '==', '!=']

# 定义转义符
ESCAPE = '\\'


class SqlxException(Exception):

    def __init__(self, message):
        self.message = message

    def __str__(self):
        return self.message


def make_sure(flag, message='something went wrong!'):
    if not flag:
        raise SqlxException(message)


def escape(content, escape_map=None):
    # 转义处理，先只传原文，语法处理后再次调用该函数，将第一次访问的 escape_map 传入
    if escape_map is None:
        escape_map = {}
        items = re.findall(r'\\\S', content)
        for item in items:
            value = item[1]
            # 如果出现重复随机数会出错，概率太低，不特殊处理了
            n = random.randint(1, 99999999)
            # 为兼容 python3.3 只能以 .format 方式代替 f''
            key = '[escape{n}]'.format(**locals())
            escape_map[key] = value
            content = content.replace(item, key)
        return content, escape_map
    else:
        make_sure(isinstance(escape_map, dict))
        for key, value in escape_map.items():
            content = content.replace(key, value)
        return content
    

def remove_space_line(content):
    # 移除空行
    new_lines = []
    for line in content.splitlines():
        if line.strip():
            new_lines.append(line)
    return '\n'.join(new_lines)


def remove_gap(content, n):
    # 移除过多的空行 比如传入参数为 5 ,则会将文本中 5 个空行替换为 1 个空行
    target = '\n' * n
    while target in content:
        content = content.replace(target, '\n')
    return content

def get_indent(s):
    # 获取字符串前有多少个前导空格
    return len(s) - len(s.lstrip())


def render(content, var_map, func_map, local_map=None):
    # render sqlt content to sql

    key_map = {}
    key_map.update(var_map)
    if local_map:
        key_map.update(local_map) 

    # 处理 for 循环，暂时不支持嵌套
    for_funcs = re.findall(r'(\{\s*%\s*for\s+(.+?)\s+in\s+(.+?)\s*%\s*\}(.*?)\{\s*%\s*endfor\s*%\s*\})', content, re.S)
    for full_func, for_names, for_values, for_content in for_funcs:
        for_names = for_names.split('|')
        for_values = for_values.split(',')
        for_values = [t.split('|') for t in for_values]
        # {% for n|m in 1|a,2|b,3|c %} ... {% endfor %}
        # => 
        # for_names = ['n', 'm']
        # for_values = [['1', 'a'], ['2', 'b'], ['3', 'c']]
        rendered_funcs = []
        for values in for_values:
            local_map = {}
            local_map.update(key_map)
            for for_name, for_value in zip(for_names, values):
                local_map[for_name] = for_value
            # local_map => {n: 1, m: a}
            rendered_func = render(for_content, var_map, func_map, local_map)
            rendered_func = remove_space_line(rendered_func)
            rendered_funcs.append(rendered_func)

        rendered_funcs = '\n'.join(rendered_funcs)
        content = content.replace(full_func, rendered_funcs)

    # 处理 if 判断，暂时不支持嵌套
    if_funcs = re.findall(r'(\{\s*%\s*if(.+?)%\s*\}(.*?)\{\s*%\s*endif\s*%\s*\})', content, re.S)
    for full_func, condition, if_content in if_funcs:
        # {% if a > b %} ... {% else %} ... {% endif %}
        if_content = re.sub(r'\{\s*%\s*else\s*%\s*\}', r'{% else %}', if_content)
        ts = if_content.split(r'{% else %}')
        make_sure(len(ts) in [1, 2], '{full_func} 内容编写错误!'.format(**locals()))
        if_content = ts[0]
        if len(ts) == 2:
            else_content = ts[1]
        else:
            else_content = ''

        a1 = a2 = None
        for op in OPERATORS:
            if op in condition:
                make_sure(condition.count(op) == 1, '{condition} 判定条件编写错误!'.format(**locals()))
                a1, a2 = condition.split(op)
                a1 = a1.strip()
                a2 = a2.strip()
                break
        make_sure(a1 and a2, '{condition} 未找到合法的关系运算符!'.format(**locals()))

        # 判断项默认以字符串类型比较
        # 如果以 $ 开头后接变量名，则转为对应的变量值
        if a1.startswith('$') and a1[1:] in key_map:
            a1 = key_map[a1[1:]]

        if a2.startswith('$') and a2[1:] in key_map:
            a2 = key_map[a2[1:]]

        # 先尝试将两个变量转为数字类型再比较
        try:
            a1 = float(a1)
        except ValueError as e:
            pass
        try:
            a2 = float(a2)
        except ValueError as e:
            pass

        a1 = repr(a1)
        a2 = repr(a2)
        s = '{a1} {op} {a2}'.format(**locals())

        try:
            result = eval(s)
        except Exception as e:
            print(condition)
            print(s)
            raise e

        make_sure(result in (True, False))

        if result:
            rendering_content = if_content
        else:
            rendering_content = else_content

        rendered_func = render(rendering_content, var_map, func_map, key_map)
        rendered_func = remove_space_line(rendered_func)
        content = content.replace(full_func, rendered_func)

    # 处理 var 替换和 func (func) 替换
    tags = re.findall(r'\{.+?\}', content)
    tags = set(tags)
    rendered_map = {}
    for tag in tags:
        key = tag.strip('{}').strip()
        if '(' not in key:
            make_sure(key in key_map, '`{tag}` var 引用未找到!'.format(**locals()))
            value = key_map[key]
            # 对于简单的变量替换，直接 replace 就行了
            content = content.replace(tag, value)
        else:
            rs = re.findall(r'(.+?)\((.*?)\)', key)
            make_sure(len(rs) == 1, '`{tag}` func 引用语法不正确!'.format(**locals())) 
            func_name, params = rs[0]
            make_sure(func_name in func_map, '`{tag}` func 引用未找到!'.format(**locals()))
            func_content = func_map[func_name]['content']
            param_names = func_map[func_name]['params']
            params = params.split(',')
            params = [param.strip() for param in params if param.strip()]
            make_sure(len(param_names) == len(params), '{tag} func 参数数量不正确!'.format(**locals()))
            local_map = copy.copy(key_map)
            for name, value in zip(param_names, params):
                if value.startswith('$') and value[1:] in key_map:
                    value = key_map[value[1:]]
                local_map[name] = value
            rendered_func = render(func_content, var_map, func_map, local_map)
            # 对于块替换，为了更好的视觉体验，先将渲染后的块内容保存下来，接下来用到
            rendered_map[tag] = rendered_func

    lines = content.splitlines()
    new_lines = []
    for line in lines:
        for tag in rendered_map.keys():
            if tag in line:
                # 遍历每一行，替换行中的块内容，并加上合适的缩进
                # 例如 `select * from {myfunc()} where 1=1` 渲染后得到:
                # select * from 
                #     (
                #         SELECT
                #             id, name
                #         FROM
                #             mytable
                #     ) AS myfunc
                # where 1=1
                n = get_indent(line)
                n = n - 4 if n >= 4 else 0  # 假设用户会在 func 中缩进 4 格
                rendered_func = rendered_map[tag]
                rendered_func = rendered_func.replace('\n', '\n' + ' ' * n)
                rendered_func = '\n' + ' ' * n + rendered_func + '\n' + ' ' * n
                # 先尝试替换 tag 两边有空格的情况
                tag2 = ' {tag} '.format(**locals())
                line = line.replace(tag2 , rendered_func)
                line = line.replace(tag, rendered_func)
        new_lines.append(line)
    content = '\n'.join(new_lines)

    return content


def handle_import(content, path, var_map, func_map):
    # 处理 import 和 sqlx 注释
    # 通过 import 可以引入现有的 sqlx 脚本文件作，但只能导入其中的 var 和 func
    # 如果在当前脚本有重复同名变量或 func，会被覆盖以当前脚本为准
    # import xxx

    make_sure(isinstance(var_map, dict))
    make_sure(isinstance(func_map, dict))
    if not path:
        path = os.getcwd()
    make_sure(os.path.isdir(path), '{path} 脚本所在目录不正确!'.format(**locals()))

    # 插入第一行空行
    new_lines = ['']
    lines = content.splitlines()
    for line in lines:

        if line.startswith(COMMENT_PREFIX):
            continue

        if COMMENT_PREFIX in line:
            i = line.find(COMMENT_PREFIX)
            line = line[:i]

        if line.lower().startswith('import '):
            items = line.split()
            make_sure(len(items) == 2, '`{line}` import 语法不正确!'.format(**locals()))
            var, script_name = items
            script_name += '.sqlx'
            script_path = os.path.join(path, script_name)
            make_sure(os.path.isfile(script_path), '{script_path} 导入模块路径不正确!'.format(**locals()))
            
            script_content = open(script_path, encoding='utf8').read()
            script_content = handle_import(script_content, path, var_map, func_map)
            script_content = handle_var(script_content, var_map)
            script_content = handle_func(script_content, func_map)

            continue

        new_lines.append(line)

    sqlx_content = '\n'.join(new_lines)

    # pprint.pprint(var_map)
    # pprint.pprint(func_map)
    return sqlx_content


def handle_var(content, var_map):
    # 处理 var
    # var a = xxx
    make_sure(isinstance(var_map, dict))
    new_lines = []
    lines = content.splitlines()
    for line in lines:
        # 兼容老版本 define 写法 
        if line.lower().startswith('var ') or line.lower().startswith('define '):
            # `var a = xxx` 与 `var a xxx` 两种写法都可以
            line = line.replace('=', ' ', 1)  
            items = line.split()
            make_sure(len(items) == 3, '`{line}` var 语法不正确!'.format(**locals()))
            var, key, value = items
            var_map[key] = value
            continue
        new_lines.append(line)
        
    sqlx_content = '\n'.join(new_lines)

    # pprint.pprint(var_map)
    return sqlx_content


def handle_func(content, func_map):
    make_sure(isinstance(func_map, dict))
    # 处理 func
    # func foo()
    #   ...
    # end

    # foo() 后面加不加:都可以
    func_pattern = r'\nfunc\s+(.+?)\((.*?)\)[:\s]*\n(.*?)\nend'
    funcs = re.findall(func_pattern, content, re.S)
    content = re.sub(func_pattern, '', content, flags=re.S)

    # 兼容老版本 block foo() ... endblock 写法
    func_pattern_old = r'\nblock\s+(.+?)\((.*?)\)[:\s]*\n(.*?)\nendblock'  
    funcs += re.findall(func_pattern_old, content, re.S)
    content = re.sub(func_pattern_old, '', content, flags=re.S)

    for func in funcs:
        func_name, params, func_content = func
        params = params.split(',')
        params = [param.strip() for param in params if param.strip()]
        func_map[func_name] = {
            'params': params,
            'content': func_content,
        }
    # pprint.pprint(func_map)
    return content


def build(content, pretty=False, path=''):
    # build sqlx content to sql
    
    content, escape_map = escape(content)
    
    var_map = {}
    func_map = {}

    content = handle_import(content, path, var_map, func_map)
    content = handle_var(content, var_map)
    content = handle_func(content, func_map)

    sql = render(content, var_map, func_map)
    sql = sql.strip()
    sql = escape(sql, escape_map)
    sql = remove_gap(sql, 5)
    header = HEADER
    sql = '{header}\n\n{sql}\n'.format(**locals())

    # print(sql)

    if pretty:
        import sqlformat
        sql = sqlformat.sqlformat(sql)

    return sql


def auto(path='.', pretty=False):

    import pyperclip
    
    # pip intall sqlx
    # sqlx [path/to/sqlxfiles]
    version = VERSION
    print('==== sqlx v{version} ===='.format(**locals()))

    args = sys.argv
    if len(args) > 1:
        path = args[1]

    if 'pretty' in args:
        pretty = True

    if os.path.isdir(path):
        files = os.listdir(path)
        files = [file for file in files if file.endswith('.sqlx')]
    elif os.path.isfile(path) and path.endswith('.sqlx'):
        files = [path]
    else:
        print('Usage: sqlx path/to/sqlxfiles')
        return 1

    for file in files:
        # build xx.sqlx to dist/xx.sql

        print('building', file)

        # sqlx 脚本所在目录
        dirname, filename = os.path.split(file)
        
        # 要生成的 sql 所在目录
        distname = os.path.join(dirname, 'dist')
        
        # 要生成的 sql 文件路径
        filename = os.path.join(distname, filename[:-1])

        if not os.path.isdir(distname):
            os.makedirs(distname)

        sqlx_content = ''
        for encoding in ['utf8', 'gbk']:
            try:
                sqlx_content = open(file, encoding=encoding).read()
                break
            except Exception as e:
                encoding = None
                
        if not encoding:
            print(file, 'read failed!')
            continue

        sql_content = build(sqlx_content, pretty, dirname)
        copied = ''
        if os.path.isfile(filename):
            old_content = open(filename, encoding=encoding).read()
            if sql_content != old_content:
                pyperclip.copy(sql_content)
                copied = 'and Copied'
        else:
            pyperclip.copy(sql_content)
            copied = 'and Copied'
        open(filename, 'w', encoding=encoding).write(sql_content)
        print('{filename} Built {copied}'.format(**locals()))


if __name__ == '__main__':
    try:
        auto()
    except:
        traceback.print_exc()
        print('See https://github.com/taojy123/sqlx/blob/master/README.md for help')
    input('Press Enter to Exit..')