AQuery/prompt.py

import os

from engine.ast import Context
if __name__ == '__main__':
    import mimetypes
    mimetypes._winreg = None

from dataclasses import dataclass
import enum
from tabnanny import check
import time
# import dbconn
import re
from typing import Callable, List, Optional
from mo_parsing import ParseException
import aquery_parser as parser
import engine
import engine.projection
import engine.ddl
import reconstruct as xengine
import subprocess
import mmap
import sys
from engine.utils import base62uuid
import atexit
import threading
import ctypes
import aquery_config
import numpy as np
from engine.utils import ws
from engine.utils import add_dll_dir

## GLOBALS BEGIN
nullstream = open(os.devnull, 'w')
help_message = '''\
Run prompt.py without supplying with any arguments to run in interactive mode.

--help, -h
    print out this help message
--version, -v
    returns current version of AQuery
--mode, -m Optional[threaded|IPC]
    execution engine run mode:
        threaded or 1 (default): run the execution engine and compiler/prompt in separate threads
        IPC, standalong or 2: run the execution engine in a new process which uses shared memory to communicate with the compiler process
--script, -s [SCRIPT_FILE]
    script mode: run the aquery script file
--parse-only, -p
    parse only: parse the file and print out the AST
'''
## GLOBALS END

## CLASSES BEGIN
class RunType(enum.Enum):
    Threaded = 0
    IPC = 1

class Backend_Type(enum.Enum):
	BACKEND_AQuery = 0
	BACKEND_MonetDB = 1
	BACKEND_MariaDB = 2

class Config:
    __all_attrs__ = ['running', 'new_query', 'server_mode', 'backend_type', 'has_dll', 'n_buffers']
    __init_attributes__ = False
    def __init_self__():
        if not Config.__init_attributes__:
            from functools import partial
            for _i, attr in enumerate(Config.__all_attrs__):
                if not hasattr(Config, attr):
                    setattr(Config, attr, property(partial(Config.getter, i = _i), partial(Config.setter, i = _i)))
            __init_attributes__ = True
            
    def __init__(self, mode, nq = 0, n_bufs = 0, bf_szs = []) -> None:
        Config.__init_self__()
        self.int_size = 4
        self.n_attrib = 6
        self.buf = bytearray((self.n_attrib + n_bufs) * self.int_size)
        self.np_buf = np.ndarray(shape=(self.n_attrib), buffer=self.buf, dtype=np.int32)
        self.new_query = nq
        self.server_mode = mode.value 
        self.running = 1
        self.backend_type = Backend_Type.BACKEND_AQuery.value
        self.has_dll = 0
        self.n_buffers = n_bufs
        
    def getter (self, *, i):
        return self.np_buf[i]
    def setter(self, v, *, i):
        self.np_buf[i] = v

    def set_bufszs(self, buf_szs):
        for i in range(min(len(buf_szs), self.n_buffers)):
            self.np_buf[i+self.n_attrib] = buf_szs[i]

    @property
    def c(self):
        return ctypes.cast(\
            (ctypes.c_char * len(self.buf)).from_buffer(self.buf), ctypes.c_char_p)

@dataclass
class PromptState():
    cleanup = True
    th = None
    send = None
    test_parser = True
    server_mode : RunType = RunType.Threaded
    server_bin = 'server.bin' if server_mode == RunType.IPC else 'server.so'
    set_ready = lambda: None
    get_ready = lambda: None
    server_status = lambda: False
    cfg : Config = None
    shm : str = ''
    server : subprocess.Popen = None
    basecmd : List[str] = None
    mm : mmap = None
    th : threading.Thread
    send : Callable = lambda *_:None
    init : Callable[['PromptState'], None] = lambda _:None
    stmts = ['']
    payloads = {}
## CLASSES END

## FUNCTIONS BEGIN
def rm(state:PromptState):
    if state.cleanup:
        state.mm.seek(0,os.SEEK_SET)
        state.mm.write(b'\x00\x00')
        state.mm.flush()
    
        try:
            time.sleep(.001)
            state.server.kill()
            time.sleep(.001)
            state.server.terminate()
        except OSError:
            pass
    
        files = os.listdir('.')
        for f in files:
            if f.endswith('.shm'):
                os.remove(f)
        state.mm.close()
        state.cleanup = False
        nullstream.close()

def init_ipc(state: PromptState):
    state.shm = base62uuid()
    if sys.platform != 'win32':
        state.shm += '.shm'
        state.basecmd = ['bash', '-c', 'rlwrap k']
        state.mm = None
        if not os.path.isfile(shm):
            # create initial file
            with open(shm, "w+b") as handle:
                handle.write(b'\x01\x00') # [running, new job]
                handle.flush()
                state.mm = mmap.mmap(handle.fileno(), 2, access=mmap.ACCESS_WRITE, offset=0)
        if state.mm is None:
            exit(1)
    else:
        state.basecmd = ['bash.exe', '-c', 'rlwrap ./k']
        state.mm = mmap.mmap(0, 2, state.shm)
        state.mm.write(b'\x01\x00')
        state.mm.flush()
    state.server = subprocess.Popen(["./server.bin", state.shm])

def init_threaded(state : PromptState):
    state.cleanup = False
    if os.name == 'nt' and aquery_config.add_path_to_ldpath:
        t = os.environ['PATH'].lower().split(';')
        vars = re.compile('%.*%')
        add_dll_dir(os.path.abspath('.'))
        add_dll_dir(os.path.abspath('./lib'))
        for e in t:
            if(len(e) != 0):
                if '%' in e:
                    try:
                        m_e = vars.findall(e)
                        for m in m_e:
                            e = e.replace(m, os.environ[m[1:-1]])
                    except Exception:
                        continue
                try:
                    add_dll_dir(e)
                except Exception:
                    continue
    
    else:
        os.environ['PATH'] = os.environ['PATH'] + os.pathsep + os.path.abspath('.')
        os.environ['PATH'] = os.environ['PATH'] + os.pathsep + os.path.abspath('./lib')
    if aquery_config.run_backend:    
        server_so = ctypes.CDLL('./'+state.server_bin)
        state.send = server_so['receive_args']
        aquery_config.have_hge = server_so['have_hge']()
        if aquery_config.have_hge:
            from engine.types import get_int128_support
            get_int128_support()
        state.th = threading.Thread(target=server_so['main'], args=(-1, ctypes.POINTER(ctypes.c_char_p)(state.cfg.c)), daemon=True)
        state.th.start()

def init_prompt() -> PromptState:
    aquery_config.init_config()
    
    state = PromptState()
    if aquery_config.rebuild_backend:
        try:
            os.remove(state.server_bin) 
        except Exception as e:
            print(type(e), e)
        subprocess.call(['make', "info"])
        subprocess.call(['make', state.server_bin], stdout=nullstream)
        
    state.cfg = Config(state.server_mode)
        
    if state.server_mode == RunType.IPC:
        atexit.register(lambda: rm(state))
        state.init = init_ipc
        state.set_ready = lambda : state.mm.seek(0,os.SEEK_SET) or state.mm.write(b'\x01\x01')
        def __get_ready():
            state.mm.seek(0,os.SEEK_SET)
            return  state.mm.read(2)[1]
        state.get_ready = __get_ready
        state.server_status = lambda : state.server.poll() is not None
    else:
        state.init = init_threaded
        rm = lambda: None
        def __set_ready():
            state.cfg.new_query = 1
        state.set_ready = __set_ready
        state.get_ready = lambda: aquery_config.run_backend and state.cfg.new_query
        if aquery_config.run_backend:
            state.server_status = lambda : not state.th.is_alive()
        else:
            state.server_status = lambda : True
    state.init(state)
    return state

def save(q:str, cxt: xengine.Context):
    savecmd = re.split(r'[ \t]', q)
    if len(savecmd) > 1:
        fname = savecmd[1]
    else:
        tm = time.gmtime()
        fname = f'{tm.tm_year}{tm.tm_mon}_{tm.tm_mday}_{tm.tm_hour}:{tm.tm_min}:{tm.tm_sec}'
    if cxt:
        from typing import Optional
        def savefile(attr:str, desc:str, ext:Optional[str] = None):
            if hasattr(cxt, attr):
                attr : str = getattr(cxt, attr)
                if attr:
                    ext = ext if ext else '.' + desc
                    name = fname if fname.endswith(ext) else fname + ext
                    with open('saves/' + name, 'wb') as cfile:
                        cfile.write(attr.encode('utf-8'))
                        print(f'saved {desc} code as {name}')
        savefile('ccode', 'cpp')
        savefile('udf', 'udf', '.hpp')
        savefile('sql', 'sql')

def main(running = lambda:True, next = input, state = None):
    if state is None:
        state = init_prompt()
    q = ''
    payload = None
    keep = True
    cxt = engine.initialize()
    while running():
        try:
            if state.server_status():
                state.init()
            while state.get_ready():
                time.sleep(.00001)
            print("> ", end="")
            q = next().lower()
            if q == 'exec': # generate build and run (AQuery Engine)
                state.cfg.backend_type = Backend_Type.BACKEND_AQuery.value
                cxt = engine.exec(state.stmts, cxt, keep)
                if subprocess.call(['make', 'snippet'], stdout = nullstream) == 0:
                    state.set_ready()
                continue
            
            elif q == 'xexec': # generate build and run (MonetDB Engine)
                state.cfg.backend_type = Backend_Type.BACKEND_MonetDB.value
                cxt = xengine.exec(state.stmts, cxt, keep)
                if state.server_mode == RunType.Threaded:
                    # assignment to avoid auto gc
                    # sqls =  [s.strip() for s in cxt.sql.split(';')]
                    qs = [ctypes.c_char_p(bytes(q, 'utf-8')) for q in cxt.queries if len(q)]
                    sz = len(qs)
                    payload = (ctypes.c_char_p*sz)(*qs)
                    state.payload = payload
                    try:
                        state.send(sz, payload)
                    except TypeError as e:
                        print(e)
                if cxt.udf is not None:
                    with open('udf.hpp', 'wb') as outfile:
                        outfile.write(cxt.udf.encode('utf-8'))
                    
                if cxt.has_dll:
                    with open('out.cpp', 'wb') as outfile:
                        outfile.write((cxt.finalize()).encode('utf-8'))
                    subprocess.call(['make', 'snippet'], stdout = nullstream)
                    state.cfg.has_dll = 1
                else:
                    state.cfg.has_dll = 0
                state.set_ready()
                
                continue
            
            elif q == 'dbg':
                import code
                from copy import deepcopy
                var = {**globals(), **locals()}
                sh = code.InteractiveConsole(var)
                try:
                    sh.interact(banner = 'debugging session began.', exitmsg = 'debugging session ended.')
                except BaseException as e: 
                # don't care about anything happened in interactive console
                    print(e)
                continue
            elif q.startswith('log'):
                qs = re.split(r'[ \t]', q)
                if len(qs) > 1:
                    cxt.log_level = qs[1]
                else:
                    cxt.print(cxt.log_level)
                continue
            elif q == 'k':
                subprocess.call(state.basecmd)
                continue
            elif q == 'print':
                cxt.print(state.stmts)
                continue
            elif q.startswith('save'):
                save(q, cxt)
                continue
            elif q == 'keep':
                keep = not keep
                continue
            elif q == 'format' or q == 'fmt':
                subprocess.call(['clang-format', 'out.cpp'])
            elif q == 'exit':
                rm(state)
                exit()
            elif q == 'r': # build and run
                if subprocess.call(['make', 'snippet']) == 0:
                    state.set_ready()
                continue
            elif q == 'rr': # run
                state.set_ready()
                continue
            elif q == 'script':
                # TODO: script mode
                pass
            elif q.startswith('save2'):
                filename = re.split(r'[ \t]', q)
                if (len(filename) > 1):
                    filename = filename[1]
                else:
                    filename = f'out_{base62uuid(4)}.cpp'
                with open(filename, 'wb') as outfile:
                    outfile.write((cxt.finalize()).encode('utf-8'))
                continue
            trimed = ws.sub(' ', q.lower()).split(' ') 
            if trimed[0].startswith('f'):
                fn = 'stock.a' if len(trimed) <= 1 or len(trimed[1]) == 0 \
                                else trimed[1]
                    
                with open(fn, 'r') as file:
                    contents = file.read()#.lower()
                state.stmts = parser.parse(contents)
                continue
            state.stmts = parser.parse(q)
            cxt.Info(state.stmts)
        except ParseException as e:
            print(e)
            continue
        except (ValueError, FileNotFoundError) as e:
            print(e)
        except (KeyboardInterrupt):
            break
        except:
            import code, traceback
            sh = code.InteractiveConsole({**globals(), **locals()})
            sh.interact(banner = traceback.format_exc(), exitmsg = 'debugging session ended.')
            save('', cxt)
            rm(state)
            raise    
    rm(state)
## FUNCTIONS END

## MAIN
if __name__ == '__main__':
    state = None
    nextcmd = ''
    def check_param(param:List[str], args = False) -> bool:
        global nextcmd
        for p in param:
            if p in sys.argv:
                if args:
                    return True
                pos = sys.argv.index(p)
                if len(sys.argv) > pos + 1:
                    nextcmd = sys.argv[pos + 1]
                    return True
        return False
    if check_param(['-h', '--help'], True):
        print(help_message)
        exit()
    
    if len(sys.argv) == 2:
        nextcmd = sys.argv[1]
        if nextcmd.startswith('-'):
            nextcmd = ''
    
    nextcmd = 'test.aquery'
    if nextcmd or check_param(['-s', '--script']):
        with open(nextcmd) as file:
            nextcmd = file.readline()
            from engine.utils import _Counter
            if file.name.endswith('aquery') or nextcmd.strip() == '#!aquery':
                state = init_prompt()
                while(nextcmd):
                    while(not ws.sub('', nextcmd) or nextcmd.strip().startswith('#')):
                        nextcmd = file.readline()
                    cnt = _Counter(1)
                    main(lambda : cnt.inc(-1) > 0, lambda:nextcmd.strip(), state)
                    nextcmd = file.readline()
        
    if check_param(['-p', '--parse']):
        with open(nextcmd, 'r') as file:
            contents = file.read()
            print(parser.parse(contents))
    
    if check_param(['-m', '--mode', '--run-type']):
        nextcmd = nextcmd.lower()
        ipc_string = ['ipc', 'proc', '2', 'standalong']
        thread_string = ['thread', '1']
        if any([s in nextcmd for s in ipc_string]):
            server_mode = RunType.IPC
        elif any([s in nextcmd for s in thread_string]):
            server_mode = RunType.Threaded       
                 
    main(state=state)