lava-oushudb-dt-sql-parser/test/parser/trino/contextCollect/entityCollector.test.ts

import fs from 'fs';
import path from 'path';
import { TrinoSQL, TrinoSqlSplitListener, TrinoEntityCollector } from 'src/parser/trino';
import { ParseTreeListener } from 'antlr4ng';
import { TrinoSqlListener } from 'src/lib/trino/TrinoSqlListener';
import { EntityContextType } from 'src/parser/common/types';
import { StmtContextType } from 'src/parser/common/entityCollector';

const commonSql = fs.readFileSync(path.join(__dirname, 'fixtures', 'common.sql'), 'utf-8');

describe('Trino entity collector tests', () => {
    const trino = new TrinoSQL();
    const parseTree = trino.parse(commonSql);
    const splitListener = new TrinoSqlSplitListener();
    trino.listen(splitListener as TrinoSqlListener, parseTree);

    test('validate common sql', () => {
        expect(trino.validate(commonSql).length).toBe(0);
    });

    test('split results', () => {
        expect(splitListener.statementsContext.length).toBe(9);
    });

    test('create table like', () => {
        const testingContext = splitListener.statementsContext[0];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(2);

        const tableCreateEntity = allEntities[0];

        expect(tableCreateEntity.entityContextType).toBe(EntityContextType.TABLE_CREATE);
        expect(tableCreateEntity.text).toBe('bar');
        expect(tableCreateEntity.position).toEqual({
            startIndex: 27,
            endIndex: 29,
            line: 1,
            startColumn: 28,
            endColumn: 31,
        });

        expect(tableCreateEntity.belongStmt.stmtContextType).toBe(
            StmtContextType.CREATE_TABLE_STMT
        );
        expect(tableCreateEntity.belongStmt.position).toEqual({
            startIndex: 0,
            endIndex: 68,
            startLine: 1,
            endLine: 1,
            startColumn: 1,
            endColumn: 70,
        });

        expect(tableCreateEntity.relatedEntities.length).toBe(1);

        const beLikedEntity = allEntities[1];

        expect(tableCreateEntity.relatedEntities[0]).toBe(beLikedEntity);
        expect(beLikedEntity.text).toBe('like_table');
        expect(beLikedEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(beLikedEntity.belongStmt).toBe(tableCreateEntity.belongStmt);
    });

    test('create table as select', () => {
        const testingContext = splitListener.statementsContext[1];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(2);

        const tableCreateEntity = allEntities[0];
        const originTableEntity = allEntities[1];

        expect(tableCreateEntity.entityContextType).toBe(EntityContextType.TABLE_CREATE);
        expect(tableCreateEntity.text).toBe('foo');
        expect(tableCreateEntity.belongStmt.stmtContextType).toBe(
            StmtContextType.CREATE_TABLE_STMT
        );

        expect(tableCreateEntity.columns.length).toBe(2);
        tableCreateEntity.columns.forEach((columEntity) => {
            expect(columEntity.entityContextType).toBe(EntityContextType.COLUMN_CREATE);
            expect(columEntity.belongStmt).toBe(tableCreateEntity.belongStmt);
            expect(columEntity.text).toBe(
                commonSql.slice(columEntity.position.startIndex, columEntity.position.endIndex + 1)
            );
        });
        expect(tableCreateEntity.relatedEntities.length).toBe(1);
        expect(tableCreateEntity.relatedEntities[0]).toBe(originTableEntity);

        expect(originTableEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(originTableEntity.text).toBe('t');
        expect(originTableEntity.belongStmt.rootStmt).toBe(tableCreateEntity.belongStmt);
    });

    test('create view as select', () => {
        const testingContext = splitListener.statementsContext[2];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(2);

        const tableCreateEntity = allEntities[0];
        const originTableEntity = allEntities[1];

        expect(tableCreateEntity.entityContextType).toBe(EntityContextType.VIEW_CREATE);
        expect(tableCreateEntity.text).toBe('a');
        expect(tableCreateEntity.belongStmt.stmtContextType).toBe(StmtContextType.CREATE_VIEW_STMT);

        expect(tableCreateEntity.columns).toBeNull();
        expect(tableCreateEntity.relatedEntities.length).toBe(1);
        expect(tableCreateEntity.relatedEntities[0]).toBe(originTableEntity);

        expect(originTableEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(originTableEntity.text).toBe('t');
        expect(originTableEntity.belongStmt.rootStmt).toBe(tableCreateEntity.belongStmt);
    });

    test('create materialize view as select', () => {
        const testingContext = splitListener.statementsContext[3];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(2);

        const tableCreateEntity = allEntities[0];
        const originTableEntity = allEntities[1];

        expect(tableCreateEntity.entityContextType).toBe(EntityContextType.VIEW_CREATE);
        expect(tableCreateEntity.text).toBe('a');
        expect(tableCreateEntity.belongStmt.stmtContextType).toBe(StmtContextType.CREATE_VIEW_STMT);

        expect(tableCreateEntity.columns).toBeNull();
        expect(tableCreateEntity.relatedEntities.length).toBe(1);
        expect(tableCreateEntity.relatedEntities[0]).toBe(originTableEntity);

        expect(originTableEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(originTableEntity.text).toBe('t');
        expect(originTableEntity.belongStmt.rootStmt).toBe(tableCreateEntity.belongStmt);
    });

    test('select from table', () => {
        const testingContext = splitListener.statementsContext[4];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(1);

        const tableCreateEntity = allEntities[0];

        expect(tableCreateEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(tableCreateEntity.text).toBe('table1');
        expect(tableCreateEntity.belongStmt.stmtContextType).toBe(StmtContextType.SELECT_STMT);

        expect(tableCreateEntity.columns).toBeNull();
        expect(tableCreateEntity.relatedEntities).toBeNull();
    });

    test('insert into table as select', () => {
        const testingContext = splitListener.statementsContext[5];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(2);

        const insertTableEntity = allEntities[0];
        const fromTableEntity = allEntities[1];

        expect(insertTableEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(insertTableEntity.belongStmt.stmtContextType).toBe(StmtContextType.INSERT_STMT);
        expect(insertTableEntity.text).toBe('orders');

        expect(fromTableEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(fromTableEntity.belongStmt.stmtContextType).toBe(StmtContextType.SELECT_STMT);
        expect(fromTableEntity.text).toBe('new_orders');
        expect(fromTableEntity.belongStmt.parentStmt).toBe(insertTableEntity.belongStmt);
    });

    test('insert into table values', () => {
        const testingContext = splitListener.statementsContext[6];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(1);

        const tableEntity = allEntities[0];

        expect(tableEntity.entityContextType).toBe(EntityContextType.TABLE);
        expect(tableEntity.belongStmt.stmtContextType).toBe(StmtContextType.INSERT_STMT);
        expect(tableEntity.text).toBe('cities');
    });

    test('create schema', () => {
        const testingContext = splitListener.statementsContext[7];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(1);

        const schemaEntity = allEntities[0];

        expect(schemaEntity.entityContextType).toBe(EntityContextType.DATABASE_CREATE);
        expect(schemaEntity.belongStmt.stmtContextType).toBe(StmtContextType.CREATE_DATABASE_STMT);
        expect(schemaEntity.text).toBe('test');
    });

    test('use schema', () => {
        const testingContext = splitListener.statementsContext[8];

        const collectListener = new TrinoEntityCollector(commonSql);
        trino.listen(collectListener as ParseTreeListener, testingContext);

        const allEntities = collectListener.getEntities();

        expect(allEntities.length).toBe(1);

        const schemaEntity = allEntities[0];

        expect(schemaEntity.entityContextType).toBe(EntityContextType.DATABASE);
        expect(schemaEntity.belongStmt.stmtContextType).toBe(StmtContextType.COMMON_STMT);
        expect(schemaEntity.text).toBe('information_schema');
    });
});