import { IPCMessageReader , IPCMessageWriter , createConnection , Connection , TextDocuments , InitializeResult , Diagnostic , DiagnosticSeverity , TextDocumentPositionParams , CompletionItem , CompletionItemKind , SignatureHelp , SignatureInformation , ParameterInformation , TextDocumentSyncKind } from 'vscode-languageserver/node' ; import { PgClient } from '../common/connection' ; import * as fs from 'fs' ; import { Validator } from './validator' ; import { IConnection as IDBConnection } from '../common/IConnection' ; import { BackwardIterator } from '../common/backwordIterator' ; import { SqlQueryManager } from '../queries' ; import { TextDocument } from 'vscode-languageserver-textdocument' ; export interface ISetConnection { connection : IDBConnection ; documentUri ? : string ; } export interface ExplainResults { rowCount : number ; command : string ; rows ? : any [ ] ; fields ? : any [ ] ; } export interface DBSchema { name : string ; } export interface DBField { attisdropped : boolean ; attname : string ; attnum : number ; attrelid : string ; data_type : string ; } export interface DBTable { schemaname : string ; tablename : string ; is_table : boolean ; columns : DBField [ ] ; } export interface DBFunctionsRaw { schema : string ; name : string ; result_type : string ; argument_types : string ; type : string ; description : string ; } export interface DBFunctionArgList { args : string [ ] ; description : string ; } export interface DBFunction { schema : string ; name : string ; result_type : string ; overloads : DBFunctionArgList [ ] ; type : string ; } export interface Ident { isQuoted : boolean ; name : string ; } let schemaCache = [ ] ; let tableCache = [ ] ; let functionCache = [ ] ; let keywordCache = [ ] ; let databaseCache = [ ] ; /**
 * To Debug the language server
 *
 * 1. Start the extension via F5
 * 2. Under vscode Debug pane, switch to "Attach to Language Server"
 * 3. F5
  */ let connection = createConnection ( new IPCMessageReader ( process ) , new IPCMessageWriter ( process ) ) ; let dbConnection = null , dbConnOptions = null ; console . log = connection . console . log . bind ( connection . console ) ; console . error = connection . console . error . bind ( connection . console ) ; let documents = new TextDocuments < TextDocument > ( TextDocument ) ; documents . listen ( connection ) ; let shouldSendDiagnosticRelatedInformation = false ; connection . onInitialize ( ( _params ) : InitializeResult => { shouldSendDiagnosticRelatedInformation = _params . capabilities && _params . capabilities . textDocument && _params . capabilities . textDocument . publishDiagnostics && _params . capabilities . textDocument . publishDiagnostics . relatedInformation ; return { capabilities : { textDocumentSync : TextDocumentSyncKind . Full , completionProvider : { triggerCharacters : [ ' ' , '.' , '"' ] } , signatureHelpProvider : { triggerCharacters : [ '(' , ',' ] } } } ; } ) ; function dbConnectionEnded ( ) { dbConnection = null ; tableCache = [ ] ; functionCache = [ ] ; keywordCache = [ ] ; databaseCache = [ ] ; } async function setupDBConnection ( connectionOptions , uri ) { if ( connectionOptions ) { if ( connectionOptions . certPath && fs . existsSync ( connectionOptions . certPath ) ) { connectionOptions . ssl = { ca : fs . readFileSync ( connectionOptions . certPath ) . toString ( ) , rejectUnauthorized : false , } ; } if ( ! connectionOptions . database ) { connectionOptions = { label : connectionOptions . label , host : connectionOptions . host , user : connectionOptions . user , password : connectionOptions . password , port : connectionOptions . port , database : 'postgres' , multipleStatements : connectionOptions . multipleStatements , certPath : connectionOptions . certPath , ssl : connectionOptions . ssl } ; } if ( connectionOptions . ssl === true ) { connectionOptions . ssl = { rejectUnauthorized : false } ; } dbConnection = new PgClient ( connectionOptions ) ; await dbConnection . connect ( ) ; const versionRes = await dbConnection . query ( `SELECT current_setting('server_version_num') as ver_num;` ) ; let versionNumber = parseInt ( versionRes . rows [ 0 ] . ver_num ) ; dbConnection . pg_version = versionNumber ; dbConnection . on ( 'end' , dbConnectionEnded ) ; loadCompletionCache ( connectionOptions ) ; if ( uri ) { let document = documents . get ( uri ) ; if ( document && document . languageId === 'postgres' ) { validateTextDocument ( document ) ; } } } dbConnOptions = connectionOptions ; } async function loadCompletionCache ( connectionOptions ) { if ( ! connectionOptions || ! dbConnection ) return ; let vQueries = SqlQueryManager . getVersionQueries ( dbConnection . pg_version ) ; try { if ( connectionOptions . database ) { let schemas = await dbConnection . query ( `
        SELECT nspname as name
        FROM pg_namespace
        WHERE
          nspname not in ('information_schema', 'pg_catalog', 'pg_toast')
          AND nspname not like 'pg_temp_%'
          AND nspname not like 'pg_toast_temp_%'
          AND has_schema_privilege(oid, 'CREATE, USAGE')
        ORDER BY nspname;
        ` ) ; schemaCache = schemas . rows ; } } catch ( err ) { console . log ( err . message ) ; } try { if ( connectionOptions . database ) { let tablesAndColumns = await dbConnection . query ( `
        SELECT
          tbl.schemaname,
          tbl.tablename,
          tbl.quoted_name,
          tbl.is_table,
          json_agg(a) as columns
        FROM
          (
            SELECT
              n.nspname as schemaname,
              c.relname as tablename,
              (quote_ident(n.nspname) || '.' || quote_ident(c.relname)) as quoted_name,
              true as is_table
            FROM
              pg_catalog.pg_class c
              JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE
              c.relkind = 'r'
              AND n.nspname not in ('information_schema', 'pg_catalog', 'pg_toast')
              AND n.nspname not like 'pg_temp_%'
              AND n.nspname not like 'pg_toast_temp_%'
              AND has_schema_privilege(n.oid, 'USAGE') = true
              AND has_table_privilege(quote_ident(n.nspname) || '.' || quote_ident(c.relname), 'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER') = true
            union all
            SELECT
              n.nspname as schemaname,
              c.relname as tablename,
              (quote_ident(n.nspname) || '.' || quote_ident(c.relname)) as quoted_name,
              false as is_table
            FROM
              pg_catalog.pg_class c
              JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
            WHERE
              c.relkind in ('v', 'm')
              AND n.nspname not in ('information_schema', 'pg_catalog', 'pg_toast')
              AND n.nspname not like 'pg_temp_%'
              AND n.nspname not like 'pg_toast_temp_%'
              AND has_schema_privilege(n.oid, 'USAGE') = true
              AND has_table_privilege(quote_ident(n.nspname) || '.' || quote_ident(c.relname), 'SELECT, INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER') = true
          ) as tbl
          LEFT JOIN (
            SELECT
              attrelid,
              attname,
              format_type(atttypid, atttypmod) as data_type,
              attnum,
              attisdropped
            FROM
              pg_attribute
          ) as a ON (
            a.attrelid = tbl.quoted_name::regclass
            AND a.attnum > 0
            AND NOT a.attisdropped
            AND has_column_privilege(tbl.quoted_name, a.attname, 'SELECT, INSERT, UPDATE, REFERENCES')
          )
        GROUP BY schemaname, tablename, quoted_name, is_table;
        ` ) ; tableCache = tablesAndColumns . rows ; } } catch ( err ) { console . log ( err . message ) ; } try { let functions = await dbConnection . query ( vQueries . GetAllFunctions ) ; functions . rows . forEach ( ( fn ) => { let existing = functionCache . find ( f => f . name === fn . name ) ; if ( ! existing ) { existing = { name : fn . name , schema : fn . schema , result_type : fn . result_type , type : fn . type , overloads : [ ] } ; functionCache . push ( existing ) ; } let args = fn . argument_types . split ( ',' ) . filter ( a => a ) . map < string > ( a => a . trim ( ) ) ; existing . overloads . push ( { args , description : fn . description } ) ; } ) ; } catch ( err ) { console . log ( err . message ) ; } try { let keywords = await dbConnection . query ( `select * from pg_get_keywords();` ) ; keywordCache = keywords . rows . map < string > ( rw => rw . word . toLocaleUpperCase ( ) ) ; } catch ( err ) { console . log ( err . message ) ; } try { let databases = await dbConnection . query ( `
    SELECT datname
    FROM pg_database
    WHERE
      datistemplate = false
      AND has_database_privilege(quote_ident(datname), 'TEMP, CONNECT') = true
    ;` ) ; databaseCache = databases . rows . map < string > ( rw => rw . datname ) ; } catch ( err ) { console . log ( err . message ) ; } } connection . onRequest ( 'set_connection' , async function ( ) { let newConnection = arguments [ 0 ] ; if ( ! dbConnOptions && ! newConnection . connection ) { return ; } if ( dbConnOptions && newConnection . connection && newConnection . connection . host === dbConnOptions . host && newConnection . connection . database === dbConnOptions . database && newConnection . connection . port === dbConnOptions . port && newConnection . connection . user === dbConnOptions . user ) { return ; } if ( dbConnection && ! dbConnection . is_ended ) { await dbConnection . end ( ) ; } setupDBConnection ( newConnection . connection , newConnection . documentUri ) . catch ( err => { console . log ( err . message ) ; } ) ; } ) ; documents . onDidOpen ( ( change ) => { validateTextDocument ( change . document ) ; } ) ; documents . onDidChangeContent ( ( change ) => { validateTextDocument ( change . document ) ; } ) ; var _NL = '\n' . charCodeAt ( 0 ) ; async function validateTextDocument ( textDocument ) { let diagnostics = [ ] ; if ( dbConnection ) { let sqlText = textDocument . getText ( ) ; if ( ! sqlText ) { connection . sendDiagnostics ( { uri : textDocument . uri , diagnostics } ) ; return ; } for ( let sql of Validator . prepare_sql ( sqlText ) ) { if ( ! sql . statement ) continue ; let errColumnMod = 0 ; if ( sql . statement . trim ( ) . toUpperCase ( ) . startsWith ( 'EXPLAIN ' ) ) { let match = sql . statement . match ( /\s*?EXPLAIN\s/gi ) ; if ( match ) { for ( let i = 0 ; i < match [ 0 ] . length ; i ++ ) { let ch = match [ 0 ] . charCodeAt ( i ) ; errColumnMod ++ ; if ( ch === _NL ) { errColumnMod = 1 ; sql . line ++ ; } } sql . statement = sql . statement . replace ( /\s*?EXPLAIN\s/gi , '' ) ; } } try { const results = await dbConnection . query ( `EXPLAIN ${ sql . statement }` ) ; } catch ( err ) { let errPosition = parseInt ( err . position ) - 9 + errColumnMod ; let errLine = 0 ; while ( errPosition > sql . lines [ errLine ] . length ) { errPosition -= ( sql . lines [ errLine ] . length + 1 ) ; errLine ++ ; } let spacePos = errPosition ; if ( errPosition < sql . lines [ errLine ] . length ) { spacePos = sql . lines [ errLine ] . indexOf ( ' ' , errPosition ) ; if ( spacePos < 0 ) { spacePos = sql . lines [ errLine ] . length ; } } if ( errLine === 0 ) { errPosition += sql . column ; } diagnostics . push ( { severity : DiagnosticSeverity . Error , range : { start : { line : sql . line + errLine , character : errPosition } , end : { line : sql . line + errLine , character : spacePos } } , message : err . message , source : dbConnOptions . host } ) ; } } } connection . sendDiagnostics ( { uri : textDocument . uri , diagnostics } ) ; } export interface FieldCompletionItem extends CompletionItem { tables ? : string [ ] ; } connection . onCompletion ( ( e ) : CompletionItem [ ] => { let items = [ ] ; let scenarioFound = false ; let document = documents . get ( e . textDocument . uri ) ; if ( ! document ) return items ; let iterator = new BackwardIterator ( document , e . position . character - 1 , e . position . line ) ; if ( e . context . triggerCharacter === '"' ) { let startingQuotedIdent = iterator . isFowardDQuote ( ) ; if ( ! startingQuotedIdent ) return items ; iterator . next ( ) ; if ( iterator . isNextPeriod ( ) ) { let ident = iterator . readIdent ( ) ; let isQuotedIdent = false ; if ( ident . match ( /^\".*?\"$/ ) ) { isQuotedIdent = true ; ident = fixQuotedIdent ( ident ) ; } let table = tableCache . find ( tbl => { return ( isQuotedIdent && tbl . tablename === ident ) || ( ! isQuotedIdent && tbl . tablename . toLocaleLowerCase ( ) == ident . toLocaleLowerCase ( ) ) ; } ) ; if ( ! table ) return items ; table . columns . forEach ( field => { items . push ( { label : field . attname , kind : CompletionItemKind . Property , detail : field . data_type } ) ; } ) ; } else { tableCache . forEach ( table => { items . push ( { label : table . tablename , kind : CompletionItemKind . Class } ) ; } ) ; } return items ; } if ( e . context . triggerCharacter === '.' ) { let idents = readIdents ( iterator , 3 ) ; let pos = 0 ; let schema = schemaCache . find ( sch => { return ( idents [ pos ] . isQuoted && sch . name === idents [ pos ] . name ) || ( ! idents [ pos ] . isQuoted && sch . name . toLocaleLowerCase ( ) == idents [ pos ] . name . toLocaleLowerCase ( ) ) ; } ) ; if ( ! schema ) { schema = schemaCache . find ( sch => { return sch . name == "public" ; } ) ; } else { pos ++ ; } if ( idents . length == pos ) { tableCache . forEach ( tbl => { if ( tbl . schemaname != schema . name ) { return ; } items . push ( { label : tbl . tablename , kind : CompletionItemKind . Class , detail : tbl . schemaname !== "public" ? tbl . schemaname : null } ) ; } ) ; return items ; } let table = tableCache . find ( tbl => { return tbl . schemaname == schema . name && ( idents [ pos ] . isQuoted && tbl . tablename === idents [ pos ] . name ) || ( ! idents [ pos ] . isQuoted && tbl . tablename . toLocaleLowerCase ( ) == idents [ pos ] . name . toLocaleLowerCase ( ) ) ; } ) ; if ( table ) { table . columns . forEach ( field => { items . push ( { label : field . attname , kind : CompletionItemKind . Property , detail : field . data_type } ) ; } ) ; } return items ; } if ( ! scenarioFound ) { schemaCache . forEach ( schema => { items . push ( { label : schema . name , kind : CompletionItemKind . Module } ) ; } ) ; tableCache . forEach ( table => { items . push ( { label : table . tablename , detail : table . schemaname !== "public" ? table . schemaname : null , kind : table . is_table ? CompletionItemKind . Class : CompletionItemKind . Interface , insertText : table . schemaname == "public" ? table . tablename : table . schemaname + "." + table . tablename } ) ; table . columns . forEach ( field => { let foundItem = items . find ( i => i . label === field . attname && i . kind === CompletionItemKind . Field && i . detail === field . data_type ) ; if ( foundItem ) { foundItem . tables . push ( table . tablename ) ; foundItem . tables . sort ( ) ; foundItem . documentation = foundItem . tables . join ( ', ' ) ; } else { items . push ( { label : field . attname , kind : CompletionItemKind . Field , detail : field . data_type , documentation : table . tablename , tables : [ table . tablename ] } ) ; } } ) ; } ) ; functionCache . forEach ( fn => { items . push ( { label : fn . name , kind : CompletionItemKind . Function , detail : fn . result_type , documentation : fn . overloads [ 0 ] . description } ) ; } ) ; keywordCache . forEach ( keyword => { items . push ( { label : keyword , kind : CompletionItemKind . Keyword } ) ; } ) ; databaseCache . forEach ( database => { items . push ( { label : database , kind : CompletionItemKind . Module } ) ; } ) ; } return items ; } ) ; connection . onSignatureHelp ( ( positionParams ) : SignatureHelp => { let document = documents . get ( positionParams . textDocument . uri ) ; let activeSignature = null , activeParameter = null , signatures = [ ] ; if ( document ) { let iterator = new BackwardIterator ( document , positionParams . position . character - 1 , positionParams . position . line ) ; let paramCount = iterator . readArguments ( ) ; if ( paramCount < 0 ) return null ; let ident = iterator . readIdent ( ) ; if ( ! ident || ident . match ( /^\".*?\"$/ ) ) return null ; let fn = functionCache . find ( f => f . name . toLocaleLowerCase ( ) === ident . toLocaleLowerCase ( ) ) ; if ( ! fn ) return null ; let overloads = fn . overloads . filter ( o => o . args . length >= paramCount ) ; if ( ! overloads || ! overloads . length ) return null ; overloads . forEach ( overload => { signatures . push ( { label : `${ fn . name }( ${ overload . args . join ( ' , ' ) } )` , documentation : overload . description , parameters : overload . args . map < ParameterInformation > ( v => { return { label : v } ; } ) } ) ; } ) ; activeSignature = 0 ; activeParameter = Math . min ( paramCount , overloads [ 0 ] . args . length - 1 ) ; } return { signatures , activeSignature , activeParameter } ; } ) ; function fixQuotedIdent ( str ) { return str . replace ( /^\"/ , '' ) . replace ( /\"$/ , '' ) . replace ( /\"\"/ , '"' ) ; } function readIdents ( iterator , maxlvl ) { return iterator . readIdents ( maxlvl ) . map < Ident > ( name => { let isQuoted = false ; if ( name . match ( /^\".*?\"$/ ) ) { isQuoted = true ; name = fixQuotedIdent ( name ) ; } return { isQuoted : isQuoted , name : name } ; } ) ; } connection . listen ( ) ;