DuckDB and MotherDuck support

autoload/db_ui/dbout.vim

@@ -60,15 +60,15 @@ function! db_ui#dbout#get_cell_value() abort
   endif
 
   let cell_line_number = s:get_cell_line_number(scheme)
-  let cell_range = s:get_cell_range(cell_line_number, getcurpos(), scheme)
-  let field_value = getline('.')[(cell_range.from):(cell_range.to)]
-  let start_spaces = len(matchstr(field_value, '^[[:blank:]]*'))
-  let end_spaces = len(matchstr(field_value, '[[:blank:]]*$'))
+  let cell_range = s:get_cell_range(cell_line_number, getcursorcharpos(), scheme)
+  let field_value = strcharpart(getline('.'), cell_range.from, cell_range.to - cell_range.from + 1)
+  let start_spaces = strcharlen(matchstr(field_value, '^[[:blank:]]*'))
+  let end_spaces = strcharlen(matchstr(field_value, '[[:blank:]]*$'))
   let old_selection = &selection
   set selection=inclusive
   let from = cell_range.from + start_spaces + 1
   let to = cell_range.to - end_spaces + 1
-  call cursor(line('.'), from)
+  call setcursorcharpos(line('.'), from)
   let motion = max([(to - from), 0])
   let cmd = 'normal!v'
   if motion > 0
@@ -111,24 +111,30 @@ function! db_ui#dbout#yank_header() abort
   endif
 
   let cell_line_number = s:get_cell_line_number(scheme)
-  let table_line = '-'
-  let column_line = getline(cell_line_number-1)
+  let table_line = get(scheme, 'table_line', '-')
+  let table_edge_offset = get(scheme, 'table_edge_offset', 0)
+  let header_rows = get(scheme, 'header_rows', 1)
   let underline = getline(cell_line_number)
-  let from = 0
-  let to = 0
-  let i = 0
-  let columns=[]
-  let lastcol = strlen(underline)
-  while i <= lastcol
-    if underline[i] !=? table_line || i == lastcol
-      let to = i-1
-      call add(columns, trim(column_line[from:to]))
-      let from = i+1
-    endif
-    let i += 1
-  endwhile
-  let csv_columns = join(columns, ', ')
-  call setreg(v:register, csv_columns)
+  let lastcol = strcharlen(underline)
+  let csv_rows = []
+  for j in range(header_rows, 1, -1)
+    let column_line = getline(cell_line_number-j)
+    let from = table_edge_offset
+    let to = table_edge_offset
+    let i = table_edge_offset
+    let columns = []
+    while i <= lastcol - table_edge_offset
+      if strcharpart(underline, i, 1) !=? table_line || i == lastcol
+        let to = i-1
+        call add(columns, trim(strcharpart(column_line, from, to - from + 1)))
+        let from = i+1
+      endif
+      let i += 1
+    endwhile
+    let csv_columns = join(columns, ', ')
+    call add(csv_rows, csv_columns)
+  endfor
+  call setreg(v:register, join(csv_rows, "\n"))
 endfunction
 
 function! s:get_cell_range(cell_line_number, curpos, scheme) abort
@@ -137,20 +143,20 @@ function! s:get_cell_range(cell_line_number, curpos, scheme) abort
   endif
 
   let line = getline(a:cell_line_number)
-  let table_line = '-'
+  let table_line = get(a:scheme, 'table_line', '-')
 
   let col = a:curpos[2] - 1
   let from = 0
 
-  while col >= 0 && line[col] ==? table_line
+  while col >= 0 && strcharpart(line, col, 1) ==? table_line
     let from = col
     let col -= 1
   endwhile
 
   let col = a:curpos[2] - 1
   let to = 0
 
-  while col <= len(line) && line[col] ==? table_line
+  while col <= strcharlen(line) && strcharpart(line, col, 1) ==? table_line
     let to = col
     let col += 1
   endwhile

autoload/db_ui/schemas.vim

@@ -58,6 +58,29 @@ let s:postgresql = {
       \ 'quote': 1,
       \ }
 
+let s:duckdb_list_schema_query = "
+    \ SELECT schema_name AS schema_name
+    \ FROM duckdb_schemas() AS s
+    \ INNER JOIN (SELECT current_database() AS database_name) USING (database_name);"
+
+let s:duckdb_tables = "
+    \ SELECT t.schema_name AS table_schema, t.table_name
+    \ FROM duckdb_tables() AS t
+    \ INNER JOIN (SELECT current_database() AS database_name) USING (database_name);"
+
+let s:duckdb = {
+      \ 'args': ['-list', '-c'],
+      \ 'schemes_query': s:duckdb_list_schema_query,
+      \ 'schemes_tables_query': s:duckdb_tables,
+      \ 'parse_results': {results,min_len -> s:results_parser(filter(results, '!empty(v:val)')[1:-1], '|', min_len)},
+      \ 'cell_line_number': 4,
+      \ 'cell_line_pattern': '^├─\+┼─\+',
+      \ 'table_line': '─',
+      \ 'table_edge_offset': 1,
+      \ 'header_rows': 2,
+      \ 'quote': 1,
+      \ }
+
 let s:sqlserver_foreign_keys_query = "
       \ SELECT TOP 1 c2.table_name as foreign_table_name, kcu2.column_name as foreign_column_name, kcu2.table_schema as foreign_table_schema
       \ from   information_schema.table_constraints c
@@ -240,6 +263,8 @@ let s:schemas = {
       \ 'mariadb': s:mysql,
       \ 'oracle': s:oracle,
       \ 'bigquery': s:bigquery,
+      \ 'duckdb': s:duckdb,
+      \ 'md': s:duckdb,
       \ 'clickhouse': s:clickhouse,
       \ }
 

autoload/db_ui/table_helpers.vim

@@ -40,6 +40,28 @@ let s:mysql = {
       \ 'Primary Keys': "SELECT * FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS WHERE TABLE_SCHEMA = '{schema}' AND TABLE_NAME = '{table}' AND CONSTRAINT_TYPE = 'PRIMARY KEY'",
       \ }
 
+let s:duckdb_table_comments_query = "
+    \SELECT\n
+    \    t.comment AS table_comment,\n
+    \    s.comment AS schema_comment,\n
+    \    d.comment AS database_comment\n
+    \FROM duckdb_tables() AS t\n
+    \INNER JOIN (SELECT current_database() AS database_name) USING (database_name)\n
+    \INNER JOIN duckdb_schemas() AS s ON t.schema_oid = s.oid\n
+    \INNER JOIN duckdb_databases() AS d ON t.database_oid = d.database_oid\n
+    \WHERE t.schema_name = '{schema}' AND t.table_name = '{table}'"
+
+let s:duckdb = {
+      \ 'List': 'SELECT * FROM {optional_schema}"{table}" LIMIT 200',
+      \ 'Columns': "SELECT * FROM duckdb_columns() INNER JOIN (SELECT current_database() AS database_name) USING (database_name) where table_name = '{table}' and schema_name = '{schema}'",
+      \ 'Describe': 'DESCRIBE {optional_schema}"{table}"',
+      \ 'Indexes': "SELECT * FROM duckdb_indexes() INNER JOIN (SELECT current_database() AS database_name) USING (database_name) where table_name = '{table}' and schema_name = '{schema}'",
+      \ 'Constraints': "SELECT * FROM duckdb_constraints() INNER JOIN (SELECT current_database() AS database_name) USING (database_name) where table_name = '{table}' and schema_name = '{schema}'",
+      \ 'Table Comments': s:duckdb_table_comments_query,
+      \ 'Column Comments': "SELECT column_name, comment FROM duckdb_columns() INNER JOIN (SELECT current_database() AS database_name) USING (database_name) where table_name = '{table}' and schema_name = '{schema}'",
+      \ 'Summarize': 'SUMMARIZE {optional_schema}"{table}"',
+      \ }
+
 let s:oracle_from = "
       \FROM all_constraints N\n
       \JOIN all_cons_columns L\n\t
@@ -194,6 +216,8 @@ let s:helpers = {
       \ 'mariadb': s:mysql,
       \ 'oracle': s:oracle,
       \ 'sqlite': s:sqlite,
+      \ 'duckdb': s:duckdb,
+      \ 'md': s:duckdb,
       \ 'sqlserver': s:sqlserver,
       \ 'clickhouse': s:clickhouse,
       \ 'mongodb': { 'List': '{table}.find()'},