setup_grammar.py

import datetime
import logging
import subprocess
import sys
import re
from enum import Enum
from pathlib import Path

import requests


def setup_logging():
    logging.basicConfig(stream=sys.stderr, level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s")


logger = logging.getLogger(__name__)


class Antlr4Target(Enum):
    js = 'JavaScript'
    python = 'Python3'


build_options = {
    'antlr4_compiled_target_output': {
        Antlr4Target.js: 'cratedb_sqlparse_js',
        Antlr4Target.python: 'cratedb_sqlparse_py'
    },

    'antlr4_compiled_target_subdir': 'cratedb_sqlparse/generated_parser',

    # List of '.g4' files that will be built
    'files': [
        {
            'url': 'https://github.com/crate/crate/raw/{version}/libs/sql-parser/src/main/antlr/io/crate/sql/parser/antlr/SqlBaseLexer.g4',
            'filename': 'SqlBaseLexer.g4'
        },
        {
            'url': 'https://github.com/crate/crate/raw/{version}/libs/sql-parser/src/main/antlr/io/crate/sql/parser/antlr/SqlBaseParser.g4',
            'filename': 'SqlBaseParser.g4'
        }
    ]
}

PARSER_COMPILE_PATH = Path(__file__).parent


def download_cratedb_grammar(version='master'):
    """
    Downloads CrateDB's `version` grammar files.

    Version should match a tag; for a list of tags run:
        $ curl https://api.github.com/repos/crate/crate/tags | jq -r '.[] | .name'
    """
    for file in build_options['files']:
        url = file['url'].format(version=version)
        logger.info(f"Downloading grammar: {url}")
        response = requests.get(url)

        # We annotate the CrateDB branch and date of download to the Grammar files for reference.
        text = f'/* crate_branch={version}, at={datetime.datetime.now()}, annotatedby=cratedb_sqlparse */\n' + response.text

        outfile = PARSER_COMPILE_PATH / file['filename']
        logger.info(f"Writing downloaded grammar: {outfile}")
        outfile.write_text(text)


def compile_grammar(target: Antlr4Target):
    """
    Compiles antlr4 files into `target` code.
    """

    base_dir = build_options['antlr4_compiled_target_output'][target]
    sub_dir = build_options['antlr4_compiled_target_subdir']

    for file in build_options['files']:
        outfile = PARSER_COMPILE_PATH / base_dir / sub_dir / file['filename']
        logger.info(f"Compiling grammar: {outfile}")
        subprocess.check_call(
            [
                'antlr4',
                f'-Dlanguage={target.value}',
                '-visitor',
                '-o',
                str(PARSER_COMPILE_PATH / base_dir / sub_dir),
                file['filename']
            ]
        )


def patch_lexer(target: Antlr4Target):
    """
    Patches the lexer file, removing bad syntax generated by Antlr4.
    """

    logger.info(f"Patching lexer type: {target}")

    REMOVE_LINES = [
        'import io.crate.sql.AbstractSqlBaseLexer;',
    ]

    # If more targets are added, this needs to be improved.
    extension = 'py' if target == Antlr4Target.python else 'js'

    base_dir = build_options['antlr4_compiled_target_output'][target]
    sub_dir = build_options['antlr4_compiled_target_subdir']
    file_name = build_options['files'][0]['filename'].replace('g4', extension)

    lexer_file = Path(PARSER_COMPILE_PATH / base_dir / sub_dir / file_name)
    logger.info(f"Patching lexer file: {lexer_file}")

    text = lexer_file.read_text()

    for text_to_remove in REMOVE_LINES:
        text = text.replace(text_to_remove, '')

    lexer_file.write_text(text)


def set_version(target: Antlr4Target, version: str):
    """
    Specifies the compiled version to the target package,
    depending on the package the strategy differs.
    """
    base_dir = build_options['antlr4_compiled_target_output'][target]
    sub_dir = build_options['antlr4_compiled_target_subdir']

    target_path = (PARSER_COMPILE_PATH / base_dir / sub_dir).parent

    version = f'"{version}"'  # Escape quotes on echo command.

    index_file = ''
    variable = ''

    if target == Antlr4Target.python:
        index_file = '__init__.py'
        variable = '__cratedb_version__'

    if target == Antlr4Target.js:
        index_file = 'index.js'
        variable = 'export const __cratedb_version__'

    with open(target_path / index_file, "r+") as f:
        content = f.read()

        # Removes the current content on disk.
        f.seek(0)
        f.truncate()

        updated_content = re.sub(f'({variable} = )"(.*)"', r'\1' + version, content)

        f.write(updated_content)

    logger.info(f'Updated  {variable} to {version} in {index_file}')


if __name__ == '__main__':
    """
    Invoke the grammar compiler / generator.
    """
    setup_logging()

    input_target = sys.argv[1]
    version = '6.3.2'

    if input_target.startswith("py"):
        target = Antlr4Target.python
    elif input_target.startswith("js") or input_target.startswith("javascript"):
        target = Antlr4Target.js
    else:
        raise NotImplementedError(f"Parser generator for target {input_target} not implemented")

    download_cratedb_grammar(version)
    compile_grammar(target)
    patch_lexer(target)
    set_version(target, version)