Refactor

packages/server/src/complete/AstUtils.ts

@@ -0,0 +1,182 @@
+import {
+  ColumnRefNode,
+  FromTableNode,
+  SelectStatement,
+  NodeRange,
+} from '@joe-re/sql-parser'
+import log4js from 'log4js'
+import { Table } from '../database_libs/AbstractClient'
+import { Pos } from './complete'
+
+const logger = log4js.getLogger()
+
+function isNotEmpty<T>(value: T | null | undefined): value is T {
+  return value === null || value === undefined ? false : true
+}
+
+export function getColumnRefByPos(
+  columns: ColumnRefNode[],
+  pos: Pos
+): ColumnRefNode | null {
+  return (
+    columns.find(
+      (v) =>
+        // guard against ColumnRefNode that don't have a location,
+        // for example sql functions that are not known to the parser
+        v.location &&
+        v.location.start.line === pos.line + 1 &&
+        v.location.start.column <= pos.column &&
+        v.location.end.line === pos.line + 1 &&
+        v.location.end.column >= pos.column
+    ) ?? null
+  )
+}
+
+export function isPosInLocation(location: NodeRange, pos: Pos) {
+  return (
+    location.start.line === pos.line + 1 &&
+    location.start.column <= pos.column &&
+    location.end.line === pos.line + 1 &&
+    location.end.column >= pos.column
+  )
+}
+
+export function createTablesFromFromNodes(fromNodes: FromTableNode[]): Table[] {
+  return fromNodes.reduce((p, c) => {
+    if (c.type !== 'subquery') {
+      return p
+    }
+    if (!Array.isArray(c.subquery.columns)) {
+      return p
+    }
+    const columns = c.subquery.columns
+      .map((v) => {
+        if (typeof v === 'string') {
+          return null
+        }
+        return {
+          columnName:
+            v.as || (v.expr.type === 'column_ref' && v.expr.column) || '',
+          description: 'alias',
+        }
+      })
+      .filter(isNotEmpty)
+    return p.concat({
+      database: null,
+      catalog: null,
+      columns: columns ?? [],
+      tableName: c.as ?? '',
+    })
+  }, [] as Table[])
+}
+
+export function findColumnAtPosition(
+  ast: SelectStatement,
+  pos: Pos
+): ColumnRefNode | null {
+  const columns = ast.columns
+  if (Array.isArray(columns)) {
+    // columns in select clause
+    const columnRefs = columns
+      .map((col) => col.expr)
+      .filter((expr): expr is ColumnRefNode => expr.type === 'column_ref')
+    if (ast.type === 'select' && ast.where?.expression) {
+      if (ast.where.expression.type === 'column_ref') {
+        // columns in where clause
+        columnRefs.push(ast.where.expression)
+      }
+    }
+    // column at position
+    const columnRef = getColumnRefByPos(columnRefs, pos)
+    if (logger.isDebugEnabled()) logger.debug(JSON.stringify(columnRef))
+    return columnRef ?? null
+  } else if (columns.type == 'star') {
+    if (ast.type === 'select' && ast.where?.expression) {
+      // columns in where clause
+      const columnRefs =
+        ast.where.expression.type === 'column_ref' ? [ast.where.expression] : []
+      // column at position
+      const columnRef = getColumnRefByPos(columnRefs, pos)
+      if (logger.isDebugEnabled()) logger.debug(JSON.stringify(columnRef))
+      return columnRef ?? null
+    }
+  }
+  return null
+}
+
+/**
+ * Recursively pull out the FROM nodes (including sub-queries)
+ * @param tableNodes
+ * @returns
+ */
+export function getAllNestedFromNodes(
+  tableNodes: FromTableNode[]
+): FromTableNode[] {
+  return tableNodes.flatMap((tableNode) => {
+    let result = [tableNode]
+    if (tableNode.type == 'subquery') {
+      const subTableNodes = tableNode.subquery.from?.tables || []
+      result = result.concat(getAllNestedFromNodes(subTableNodes))
+    }
+    return result
+  })
+}
+
+/**
+ * Finds the most deeply nested FROM node that have a range encompasing the position.
+ * In cases such as SELECT * FROM T1 JOIN (SELECT * FROM (SELECT * FROM T2 <pos>))
+ * We will get a list of nodes like this
+ * SELECT * FROM T1
+ * (SELECT * FROM
+ *    (SELECT * FROM T2))
+ * The idea is to reverse the list so that the most nested queries come first. Then
+ * apply a filter to keep only the FROM nodes which encompass the position and take
+ * the first one from that resulting list.
+ * @param fromNodes
+ * @param pos
+ * @returns
+ */
+export function getNearestFromTableFromPos(
+  fromNodes: FromTableNode[],
+  pos: Pos
+): FromTableNode | null {
+  return (
+    fromNodes
+      .reverse()
+      .filter((tableNode) => isPosInLocation(tableNode.location, pos))
+      .shift() ?? null
+  )
+}
+
+/**
+ * Test if the given table matches the fromNode.
+ * @param fromNode
+ * @param table
+ * @returns
+ */
+export function isTableMatch(fromNode: FromTableNode, table: Table): boolean {
+  switch (fromNode.type) {
+    case 'subquery': {
+      if (fromNode.as && fromNode.as !== table.tableName) {
+        return false
+      }
+      break
+    }
+    case 'table': {
+      if (fromNode.table && fromNode.table !== table.tableName) {
+        return false
+      }
+      if (fromNode.db && fromNode.db !== table.database) {
+        return false
+      }
+      if (fromNode.catalog && fromNode.catalog !== table.catalog) {
+        return false
+      }
+      break
+    }
+    default: {
+      return false
+    }
+  }
+  return true
+}

packages/server/src/complete/CompletionItemUtils.ts

@@ -0,0 +1,39 @@
+import { CompletionItem, CompletionItemKind } from 'vscode-languageserver-types'
+import { DbFunction } from '../database_libs/AbstractClient'
+
+export const ICONS = {
+  KEYWORD: CompletionItemKind.Text,
+  COLUMN: CompletionItemKind.Interface,
+  TABLE: CompletionItemKind.Field,
+  FUNCTION: CompletionItemKind.Property,
+  ALIAS: CompletionItemKind.Variable,
+  UTILITY: CompletionItemKind.Event,
+}
+
+export function toCompletionItemForFunction(f: DbFunction): CompletionItem {
+  const item: CompletionItem = {
+    label: f.name,
+    detail: 'function',
+    kind: ICONS.FUNCTION,
+    documentation: f.description,
+  }
+  return item
+}
+
+export function toCompletionItemForAlias(alias: string): CompletionItem {
+  const item: CompletionItem = {
+    label: alias,
+    detail: 'alias',
+    kind: ICONS.ALIAS,
+  }
+  return item
+}
+
+export function toCompletionItemForKeyword(name: string): CompletionItem {
+  const item: CompletionItem = {
+    label: name,
+    kind: ICONS.KEYWORD,
+    detail: 'keyword',
+  }
+  return item
+}

packages/server/src/complete/Identifier.ts

@@ -1,5 +1,5 @@
 import { CompletionItem, CompletionItemKind } from 'vscode-languageserver-types'
-import { makeTableAlias } from './utils'
+import { makeTableAlias } from './StringUtils'
 
 export const ICONS = {
   KEYWORD: CompletionItemKind.Text,
@@ -40,7 +40,7 @@ export class Identifier {
 
   toCompletionItem(): CompletionItem {
     const idx = this.lastToken.lastIndexOf('.')
-    const label = this.identifier.substr(idx + 1)
+    const label = this.identifier.substring(idx + 1)
     let kindName: string
     let tableAlias = ''
     if (this.kind === ICONS.TABLE) {

packages/server/src/complete/StringUtils.ts

@@ -0,0 +1,56 @@
+import { FromTableNode } from '@joe-re/sql-parser'
+import { Table } from '../database_libs/AbstractClient'
+import { Pos } from './complete'
+
+export function makeTableAlias(tableName: string): string {
+  if (tableName.length > 3) {
+    return tableName.substring(0, 3)
+  }
+  return tableName
+}
+
+export function getRidOfAfterPosString(sql: string, pos: Pos): string {
+  return sql
+    .split('\n')
+    .filter((_v, idx) => pos.line >= idx)
+    .map((v, idx) => (idx === pos.line ? v.slice(0, pos.column) : v))
+    .join('\n')
+}
+
+// Gets the last token from the given string considering that tokens can contain dots.
+export function getLastToken(sql: string): string {
+  const match = sql.match(/^(?:.|\s)*[^A-z0-9\\.:'](.*?)$/)
+  if (match) {
+    let prevToken = ''
+    let currentToken = match[1]
+    while (currentToken != prevToken) {
+      prevToken = currentToken
+      currentToken = prevToken.replace(/\[.*?\]/, '')
+    }
+    return currentToken
+  }
+  return sql
+}
+
+export function makeTableName(table: Table): string {
+  if (table.catalog) {
+    return table.catalog + '.' + table.database + '.' + table.tableName
+  } else if (table.database) {
+    return table.database + '.' + table.tableName
+  }
+  return table.tableName
+}
+
+export function getAliasFromFromTableNode(node: FromTableNode): string {
+  if (node.as) {
+    return node.as
+  }
+  if (node.type === 'table') {
+    return node.table
+  }
+  return ''
+}
+
+export function makeColumnName(alias: string, columnName: string) {
+  return alias ? alias + '.' + columnName : columnName
+}

packages/server/src/complete/candidates/createAliasCandidates.ts

@@ -0,0 +1,13 @@
+import { CompletionItem } from 'vscode-languageserver-types'
+import { FromTableNode } from '@joe-re/sql-parser'
+import { toCompletionItemForAlias } from '../CompletionItemUtils'
+
+export function createAliasCandidates(
+  fromNodes: FromTableNode[],
+  token: string
+): CompletionItem[] {
+  return fromNodes
+    .map((fromNode) => fromNode.as)
+    .filter((aliasName) => aliasName && aliasName.startsWith(token))
+    .map((aliasName) => toCompletionItemForAlias(aliasName || ''))
+}

packages/server/src/complete/candidates/createBasicKeywordCandidates.ts

@@ -0,0 +1,16 @@
+import { toCompletionItemForKeyword } from '../CompletionItemUtils'
+
+const CLAUSES: string[] = [
+  'SELECT',
+  'WHERE',
+  'ORDER BY',
+  'GROUP BY',
+  'LIMIT',
+  '--',
+  '/*',
+  '(',
+]
+
+export function createBasicKeywordCandidates() {
+  return CLAUSES.map((v) => toCompletionItemForKeyword(v))
+}

packages/server/src/complete/candidates/createColumnCandidates.ts

@@ -0,0 +1,51 @@
+import { CompletionItem } from 'vscode-languageserver-types'
+import { FromTableNode } from '@joe-re/sql-parser'
+import { Table } from '../../database_libs/AbstractClient'
+import { getAliasFromFromTableNode, makeColumnName } from '../StringUtils'
+import { isTableMatch } from '../AstUtils'
+import { ICONS } from '../CompletionItemUtils'
+import { Identifier } from '../Identifier'
+
+export function createCandidatesForColumnsOfAnyTable(
+  tables: Table[],
+  lastToken: string
+): CompletionItem[] {
+  return tables
+    .flatMap((table) => table.columns)
+    .map((column) => {
+      return new Identifier(
+        lastToken,
+        column.columnName,
+        column.description,
+        ICONS.TABLE
+      )
+    })
+    .filter((item) => item.matchesLastToken())
+    .map((item) => item.toCompletionItem())
+}
+
+export function createCandidatesForScopedColumns(
+  fromNodes: FromTableNode[],
+  tables: Table[],
+  lastToken: string
+): CompletionItem[] {
+  return tables
+    .flatMap((table) => {
+      return fromNodes
+        .filter((fromNode) => isTableMatch(fromNode, table))
+        .map(getAliasFromFromTableNode)
+        .filter((alias) => lastToken.startsWith(alias + '.'))
+        .flatMap((alias) =>
+          table.columns.map((col) => {
+            return new Identifier(
+              lastToken,
+              makeColumnName(alias, col.columnName),
+              col.description,
+              ICONS.COLUMN
+            )
+          })
+        )
+    })
+    .filter((item) => item.matchesLastToken())
+    .map((item) => item.toCompletionItem())
+}