sql-parser/run_test.py

from pathlib import Path
import sys
import time
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
from watchdog.events import LoggingEventHandler
import asyncio.subprocess as subprocess
import asyncio
from watchfiles import awatch, watch
from termcolor import colored
from datetime import datetime
import orjson
import os


async def assert_sql(sql, target):
    p = await subprocess.create_subprocess_exec(
        "xmake",
        "run",
        "sql-parser",
        sql,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    stdout, stderr = await p.communicate()
    if stderr:
        print(colored(stderr, "yellow"))
    if b"error" in stdout:
        print(stdout.decode("utf-8"))
        print(datetime.now(), "-" * 40)
        assert False, "sql-parser error"
    try:
        output = orjson.loads(stdout)
    except Exception as e:
        output = {"error": e, "output": stdout.decode("utf-8")}

    assert (
        output == target
    ), f"""{colored("sql-parser error", "red")}
input: {colored(sql, "yellow")}
expect: {colored(target, "green")}
actual: {colored(output, "red")}
"""


async def assert_sqls():
    await assert_sql(
        "create table asd;", [{"type": "create_table", "table_name": "asd", "cols": []}]
    )
    await assert_sql(
        "create table tb (col1 INT, col2 string, col3 FLOAT);",
        [
            {
                "type": "create_table",
                "table_name": "tb",
                "cols": [
                    {
                        "type": "create_column",
                        "column_name": "col1",
                        "data_type": "INT",
                        "primary_key": False,
                    },
                    {
                        "type": "create_column",
                        "column_name": "col2",
                        "data_type": "string",
                        "primary_key": False,
                    },
                    {
                        "type": "create_column",
                        "column_name": "col3",
                        "data_type": "FLOAT",
                        "primary_key": False,
                    },
                ],
            }
        ],
    )
    await assert_sql(
        """
    create table tb1 (
        col1 int primary key, 
        col2 FLOAT
    );
    """,
        [
            {
                "type": "create_table",
                "table_name": "tb1",
                "cols": [
                    {
                        "type": "create_column",
                        "column_name": "col1",
                        "data_type": "int",
                        "primary_key": True,
                    },
                    {
                        "type": "create_column",
                        "column_name": "col2",
                        "data_type": "FLOAT",
                        "primary_key": False,
                    },
                ],
            }
        ],
    )
    await assert_sql(
        """
        create table tb2 (
            x float,
            y int,
            z int
        );
        """,
        [
            {
                "type": "create_table",
                "table_name": "tb2",
                "cols": [
                    {
                        "type": "create_column",
                        "column_name": "x",
                        "data_type": "float",
                        "primary_key": False,
                    },
                    {
                        "type": "create_column",
                        "column_name": "y",
                        "data_type": "int",
                        "primary_key": False,
                    },
                    {
                        "type": "create_column",
                        "column_name": "z",
                        "data_type": "int",
                        "primary_key": False,
                    },
                ],
            }
        ],
    )
    await assert_sql(
        """insert into tb1 values (1, 'foo');""",
        [
            {
                "type": "insert",
                "table_name": "tb1",
                "values": [
                    {"type": "int", "value": 1},
                    {"type": "string", "value": "'foo'"},
                ],
            }
        ],
    )
    await assert_sql(
        """insert into tb1 values (2, 'foo', 'zxc', 1234.234);""",
        [
            {
                "type": "insert",
                "table_name": "tb1",
                "values": [
                    {"type": "int", "value": 2},
                    {"type": "string", "value": "'foo'"},
                    {"type": "string", "value": "'zxc'"},
                    {"type": "float", "value": 1234.234},
                ],
            }
        ],
    )
    await assert_sql(
        "update tb1 set col1=3 where col1=2 and col2=col3;",
        [
            {
                "type": "update",
                "table_name": "tb1",
                "set": [
                    {
                        "type": "assign_const",
                        "left": {"type": "identifier", "value": "col1"},
                        "right": {"type": "int", "value": 3},
                    }
                ],
                "where": [
                    {
                        "type": "where_condition",
                        "left": {"type": "identifier", "value": "col1"},
                        "right": {"type": "int", "value": 2},
                    },
                    {
                        "type": "where_condition",
                        "left": {"type": "identifier", "value": "col2"},
                        "right": {"type": "identifier", "value": "col3"},
                    },
                ],
            }
        ],
    )
    await assert_sql(
        "delete from tb1 where c1 = 1 and c2= 3;",
        [
            {
                "type": "delete",
                "table_name": "tb1",
                "where": [
                    {
                        "type": "where_condition",
                        "left": {"type": "identifier", "value": "c1"},
                        "right": {"type": "int", "value": 1},
                    },
                    {
                        "type": "where_condition",
                        "left": {"type": "identifier", "value": "c2"},
                        "right": {"type": "int", "value": 3},
                    },
                ],
            }
        ],
    )


async def on_modified(event):
    p = await subprocess.create_subprocess_shell(
        "xmake", stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    stdout, _ = await p.communicate()
    if b"error" in stdout:
        print(stdout.decode("utf-8"))
        print(datetime.now(), "-" * 40)
        return

    try:
        await assert_sqls()
    except Exception as e:
        print(e)
    else:
        print(datetime.now(), colored("all tests right!", "green"))


async def restart():
    async for _ in awatch(__file__):
        print("restart")
        os.execl("/bin/python", Path(__file__).as_posix(), Path(__file__).as_posix())


async def watch_src():
    async for changes in awatch("src"):
        print(datetime.now(), "re run...")
        await asyncio.wait_for(on_modified(changes), 20)


async def main():
    try:
        await assert_sqls()
    except Exception as e:
        print(e)
    await asyncio.gather(restart(), watch_src())


if __name__ == "__main__":
    asyncio.run(main())