Egutierrez
93 lines
2.4 KiB
Go
93 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"database/sql"
|
|
"fmt"
|
|
"os/exec"
|
|
"strconv"
|
|
"strings"
|
|
)
|
|
|
|
type tableInfo struct {
|
|
Name string
|
|
RowCount int64
|
|
}
|
|
|
|
// listTablesForDatabase returns the list of tables for a given DB file.
|
|
// Supports kind="sqlite" (uses database/sql + sqlite3 driver) and
|
|
// kind="duckdb" (uses python venv subprocess since go-duckdb is not linked
|
|
// in sqlite_api — it lives in the root module, not here).
|
|
func listTablesForDatabase(kind, path string) ([]tableInfo, error) {
|
|
switch kind {
|
|
case "sqlite":
|
|
return listSQLiteTables(path)
|
|
case "duckdb":
|
|
return listDuckDBTablesViaCLI(path)
|
|
default:
|
|
return nil, fmt.Errorf("unsupported kind: %s", kind)
|
|
}
|
|
}
|
|
|
|
func listSQLiteTables(path string) ([]tableInfo, error) {
|
|
db, err := sql.Open("sqlite3", "file:"+path+"?mode=ro")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer db.Close()
|
|
|
|
rows, err := db.Query(`
|
|
SELECT name FROM sqlite_master
|
|
WHERE type='table'
|
|
AND name NOT LIKE 'sqlite_%'
|
|
AND name != '_migrations'
|
|
ORDER BY name`)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
|
|
var out []tableInfo
|
|
for rows.Next() {
|
|
var name string
|
|
if err := rows.Scan(&name); err != nil {
|
|
continue
|
|
}
|
|
var cnt int64
|
|
row := db.QueryRow(`SELECT count(*) FROM "` + name + `"`)
|
|
_ = row.Scan(&cnt)
|
|
out = append(out, tableInfo{Name: name, RowCount: cnt})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func listDuckDBTablesViaCLI(path string) ([]tableInfo, error) {
|
|
// Use the python venv's duckdb module to avoid linking go-duckdb.
|
|
pyBin := "/home/lucas/fn_registry/python/.venv/bin/python3"
|
|
script := fmt.Sprintf(
|
|
`import duckdb,sys
|
|
c=duckdb.connect(%q, read_only=True)
|
|
for (name,) in c.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='main' ORDER BY table_name").fetchall():
|
|
cnt=c.execute(f'SELECT count(*) FROM "{name}"').fetchone()[0]
|
|
print(f"{name}\t{cnt}")
|
|
`, path)
|
|
cmd := exec.Command(pyBin, "-c", script)
|
|
out, err := cmd.CombinedOutput()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("duckdb via python: %w (output: %s)", err, string(out))
|
|
}
|
|
|
|
var result []tableInfo
|
|
for _, line := range strings.Split(strings.TrimSpace(string(out)), "\n") {
|
|
if line == "" {
|
|
continue
|
|
}
|
|
parts := strings.SplitN(line, "\t", 2)
|
|
if len(parts) != 2 {
|
|
continue
|
|
}
|
|
cnt, _ := strconv.ParseInt(strings.TrimSpace(parts[1]), 10, 64)
|
|
result = append(result, tableInfo{Name: parts[0], RowCount: cnt})
|
|
}
|
|
return result, nil
|
|
}
|