Skip to content

Commit

Permalink
feat: completion for non default schema
Browse files Browse the repository at this point in the history
non default schema tables where not considered as candidates for
completion
  • Loading branch information
Nicolás De los Santos committed Apr 6, 2024
1 parent e89c70f commit 63c11a2
Show file tree
Hide file tree
Showing 5 changed files with 353 additions and 84 deletions.
72 changes: 55 additions & 17 deletions internal/completer/candidates.go
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,17 @@ func (c *Completer) columnCandidates(targetTables []*parseutil.TableInfo, parent
if table.Name != parent.Name && table.Alias != parent.Name {
continue
}
columns, ok := c.DBCache.ColumnDescs(table.Name)

var columns []*database.ColumnDesc
var ok bool

if table.DatabaseSchema != "" {
columns, ok = c.DBCache.ColumnDatabase(table.DatabaseSchema, table.Name)

} else {
columns, ok = c.DBCache.ColumnDescs(table.Name)
}

if !ok {
continue
}
Expand Down Expand Up @@ -113,7 +123,13 @@ func (c *Completer) ReferencedTableCandidates(targetTables []*parseutil.TableInf

for _, targetTable := range targetTables {
includeTables := []*parseutil.TableInfo{}
for _, table := range c.DBCache.SortedTables() {
var schemaTables []string
if (targetTable.DatabaseSchema != "") {
schemaTables, _ = c.DBCache.SortedTablesByDBName(targetTable.DatabaseSchema)
} else {
schemaTables = c.DBCache.SortedTables()
}
for _, table := range schemaTables {
if table == targetTable.Name {
includeTables = append(includeTables, targetTable)
}
Expand All @@ -129,24 +145,29 @@ func (c *Completer) TableCandidates(parent *completionParent, targetTables []*pa

switch parent.Type {
case ParentTypeNone:
excludeTables := []string{}
for _, table := range c.DBCache.SortedTables() {
isExclude := false
for _, targetTable := range targetTables {
if table == targetTable.Name {
isExclude = true
targetTablesMap := make(map[string]*parseutil.TableInfo)
for _, targetTable := range targetTables {
targetTablesMap[targetTable.Name] = targetTable
}
for schemaKey, schema := range c.DBCache.Schemas {
excludeTables := []string{}

tables := c.DBCache.SchemaTables[schemaKey]
for _, table := range tables {
_, isExclude := targetTablesMap[table]
if isExclude {
continue
}
excludeTables = append(excludeTables, table)
}
if isExclude {
continue
}
excludeTables = append(excludeTables, table)

schemaCandidates := generateTableCandidatesBySchema(schema, excludeTables, c.DBCache)
candidates = append(candidates, schemaCandidates...)
}
candidates = append(candidates, generateTableCandidates(excludeTables, c.DBCache)...)
case ParentTypeSchema:
tables, ok := c.DBCache.SortedTablesByDBName(parent.Name)
if ok {
candidates = append(candidates, generateTableCandidates(tables, c.DBCache)...)
candidates = append(candidates, generateTableCandidatesBySchema(parent.Name, tables, c.DBCache)...)
}
case ParentTypeTable:
// pass
Expand Down Expand Up @@ -320,14 +341,24 @@ func generateForeignKeyCandidate(target string,
}

func generateTableCandidates(tables []string, dbCache *database.DBCache) []lsp.CompletionItem {
return generateTableCandidatesBySchema(dbCache.DefaultSchema, tables, dbCache)
}

func generateTableCandidatesBySchema(schemaName string, tables []string, dbCache *database.DBCache) []lsp.CompletionItem {
candidates := []lsp.CompletionItem{}
for _, tableName := range tables {
var label string
if schemaName != dbCache.DefaultSchema {
label = fmt.Sprintf("%s.%s", schemaName, tableName)
} else {
label = tableName
}
candidate := lsp.CompletionItem{
Label: tableName,
Label: label,
Kind: lsp.ClassCompletion,
Detail: "table",
}
cols, ok := dbCache.ColumnDescs(tableName)
cols, ok := dbCache.ColumnDatabase(schemaName, tableName)
if ok {
candidate.Documentation = lsp.MarkupContent{
Kind: lsp.Markdown,
Expand All @@ -353,7 +384,14 @@ func generateTableCandidatesByInfos(tables []*parseutil.TableInfo, dbCache *data
Kind: lsp.ClassCompletion,
Detail: detail,
}
cols, ok := dbCache.ColumnDescs(table.Name)
var cols []*database.ColumnDesc
var ok bool

if table.DatabaseSchema != "" {
cols, ok = dbCache.ColumnDatabase(table.DatabaseSchema, table.Name)
} else {
cols, ok = dbCache.ColumnDescs(table.Name)
}
if ok {
candidate.Documentation = lsp.MarkupContent{
Kind: lsp.Markdown,
Expand Down
13 changes: 12 additions & 1 deletion internal/completer/completer.go
Original file line number Diff line number Diff line change
Expand Up @@ -411,8 +411,19 @@ func getCompletionTypes(nw *parseutil.NodeWalker) *CompletionContext {

func filterCandidates(candidates []lsp.CompletionItem, lastWord string) []lsp.CompletionItem {
filtered := []lsp.CompletionItem{}
withBackQuote := strings.HasPrefix(lastWord, "`")

for _, candidate := range candidates {
if strings.HasPrefix(strings.ToUpper(candidate.Label), strings.ToUpper(lastWord)) {
label := strings.ToUpper(candidate.Label)

if !withBackQuote && candidate.Kind != lsp.SnippetCompletion {
index := strings.LastIndex(label, ".")
if index != -1 {
label = label[index+1:]
}
}

if strings.HasPrefix(label, strings.ToUpper(lastWord)) {
filtered = append(filtered, candidate)
}
}
Expand Down
18 changes: 9 additions & 9 deletions internal/database/cache.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ func NewDBCacheUpdater(repo DBRepository) *DBCacheGenerator {
func (u *DBCacheGenerator) GenerateDBCachePrimary(ctx context.Context) (*DBCache, error) {
var err error
dbCache := &DBCache{}
dbCache.defaultSchema, err = u.repo.CurrentSchema(ctx)
dbCache.DefaultSchema, err = u.repo.CurrentSchema(ctx)
if err != nil {
return nil, err
}
Expand All @@ -32,13 +32,13 @@ func (u *DBCacheGenerator) GenerateDBCachePrimary(ctx context.Context) (*DBCache
dbCache.Schemas[strings.ToUpper(index)] = element
}

if dbCache.defaultSchema == "" {
if dbCache.DefaultSchema == "" {
var topKey string
for k := range dbCache.Schemas {
topKey = k
continue
}
dbCache.defaultSchema = dbCache.Schemas[topKey]
dbCache.DefaultSchema = dbCache.Schemas[topKey]
}
schemaTables, err := u.repo.SchemaTables(ctx)
if err != nil {
Expand All @@ -49,11 +49,11 @@ func (u *DBCacheGenerator) GenerateDBCachePrimary(ctx context.Context) (*DBCache
dbCache.SchemaTables[strings.ToUpper(index)] = element
}

dbCache.ColumnsWithParent, err = u.genColumnCacheCurrent(ctx, dbCache.defaultSchema)
dbCache.ColumnsWithParent, err = u.genColumnCacheCurrent(ctx, dbCache.DefaultSchema)
if err != nil {
return nil, err
}
dbCache.ForeignKeys, err = u.genForeignKeysCache(ctx, dbCache.defaultSchema)
dbCache.ForeignKeys, err = u.genForeignKeysCache(ctx, dbCache.DefaultSchema)
if err != nil {
return nil, err
}
Expand Down Expand Up @@ -128,7 +128,7 @@ func genColumnMap(columnDescs []*ColumnDesc) map[string][]*ColumnDesc {
}

type DBCache struct {
defaultSchema string
DefaultSchema string
Schemas map[string]string
SchemaTables map[string][]string
ColumnsWithParent map[string][]*ColumnDesc
Expand Down Expand Up @@ -156,12 +156,12 @@ func (dc *DBCache) SortedTablesByDBName(dbName string) (tbls []string, ok bool)
}

func (dc *DBCache) SortedTables() []string {
tbls, _ := dc.SortedTablesByDBName(dc.defaultSchema)
tbls, _ := dc.SortedTablesByDBName(dc.DefaultSchema)
return tbls
}

func (dc *DBCache) ColumnDescs(tableName string) (cols []*ColumnDesc, ok bool) {
cols, ok = dc.ColumnsWithParent[columnDatabaseKey(dc.defaultSchema, tableName)]
cols, ok = dc.ColumnsWithParent[columnDatabaseKey(dc.DefaultSchema, tableName)]
return
}

Expand All @@ -171,7 +171,7 @@ func (dc *DBCache) ColumnDatabase(dbName, tableName string) (cols []*ColumnDesc,
}

func (dc *DBCache) Column(tableName, colName string) (*ColumnDesc, bool) {
cols, ok := dc.ColumnsWithParent[columnDatabaseKey(dc.defaultSchema, tableName)]
cols, ok := dc.ColumnsWithParent[columnDatabaseKey(dc.DefaultSchema, tableName)]
if !ok {
return nil, false
}
Expand Down
Loading

0 comments on commit 63c11a2

Please sign in to comment.