Initial word table SQLite importer
This commit is contained in:
赵义仑
606
main.go
Normal file
606
main.go
Normal file
@@ -0,0 +1,606 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"archive/zip"
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"database/sql"
|
||||
"encoding/hex"
|
||||
"encoding/xml"
|
||||
"errors"
|
||||
"flag"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
_ "modernc.org/sqlite"
|
||||
)
|
||||
|
||||
type Table struct {
|
||||
Index int
|
||||
Rows []Row
|
||||
}
|
||||
|
||||
type Row struct {
|
||||
Index int
|
||||
Cells []Cell
|
||||
}
|
||||
|
||||
type Cell struct {
|
||||
RowIndex int
|
||||
CellIndex int
|
||||
ColIndex int
|
||||
ColSpan int
|
||||
VMerge string
|
||||
Text string
|
||||
}
|
||||
|
||||
type Field struct {
|
||||
RecordIndex int
|
||||
RowIndex int
|
||||
ColIndex int
|
||||
Key string
|
||||
Value string
|
||||
SourceKind string
|
||||
}
|
||||
|
||||
func main() {
|
||||
if err := run(os.Args[1:]); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "error:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
|
||||
func run(args []string) error {
|
||||
fs := flag.NewFlagSet("wordtable2sqlite", flag.ContinueOnError)
|
||||
dbPath := fs.String("db", "word_tables.sqlite", "SQLite database path")
|
||||
reset := fs.Bool("reset", false, "remove the target database before importing")
|
||||
quiet := fs.Bool("quiet", false, "suppress per-file summary output")
|
||||
if err := fs.Parse(args); err != nil {
|
||||
return err
|
||||
}
|
||||
docxPaths := fs.Args()
|
||||
if len(docxPaths) == 0 {
|
||||
return errors.New("usage: wordtable2sqlite -db output.sqlite [-reset] file1.docx [file2.docx ...]")
|
||||
}
|
||||
|
||||
if *reset {
|
||||
if err := os.Remove(*dbPath); err != nil && !errors.Is(err, os.ErrNotExist) {
|
||||
return fmt.Errorf("reset database: %w", err)
|
||||
}
|
||||
}
|
||||
if dir := filepath.Dir(*dbPath); dir != "." {
|
||||
if err := os.MkdirAll(dir, 0o755); err != nil {
|
||||
return fmt.Errorf("create database directory: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
db, err := sql.Open("sqlite", *dbPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("open sqlite database: %w", err)
|
||||
}
|
||||
defer db.Close()
|
||||
|
||||
ctx := context.Background()
|
||||
for _, docxPath := range docxPaths {
|
||||
before, err := dbStats(ctx, db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := importDocument(ctx, db, docxPath); err != nil {
|
||||
return err
|
||||
}
|
||||
after, err := dbStats(ctx, db)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !*quiet {
|
||||
fmt.Printf("imported %s: %d tables, %d cells, %d fields\n",
|
||||
docxPath,
|
||||
after.Tables-before.Tables,
|
||||
after.Cells-before.Cells,
|
||||
after.Fields-before.Fields,
|
||||
)
|
||||
}
|
||||
}
|
||||
if !*quiet {
|
||||
fmt.Printf("database: %s\n", *dbPath)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type stats struct {
|
||||
Tables int
|
||||
Cells int
|
||||
Fields int
|
||||
}
|
||||
|
||||
func dbStats(ctx context.Context, db *sql.DB) (stats, error) {
|
||||
if err := ensureSchema(ctx, db); err != nil {
|
||||
return stats{}, err
|
||||
}
|
||||
var s stats
|
||||
for _, item := range []struct {
|
||||
query string
|
||||
dst *int
|
||||
}{
|
||||
{`SELECT COUNT(*) FROM tables`, &s.Tables},
|
||||
{`SELECT COUNT(*) FROM cells`, &s.Cells},
|
||||
{`SELECT COUNT(*) FROM fields`, &s.Fields},
|
||||
} {
|
||||
if err := db.QueryRowContext(ctx, item.query).Scan(item.dst); err != nil {
|
||||
return stats{}, err
|
||||
}
|
||||
}
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func importDocument(ctx context.Context, db *sql.DB, docxPath string) error {
|
||||
if err := ensureSchema(ctx, db); err != nil {
|
||||
return err
|
||||
}
|
||||
tables, err := readDocxTables(docxPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
sum, err := fileSHA256(docxPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
tx, err := db.BeginTx(ctx, nil)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer tx.Rollback()
|
||||
|
||||
res, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO documents(path, filename, sha256, imported_at) VALUES (?, ?, ?, ?)`,
|
||||
docxPath,
|
||||
filepath.Base(docxPath),
|
||||
sum,
|
||||
time.Now().Format(time.RFC3339),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("insert document: %w", err)
|
||||
}
|
||||
documentID, err := res.LastInsertId()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for tableIndex, table := range tables {
|
||||
cellCount := 0
|
||||
for _, row := range table.Rows {
|
||||
cellCount += len(row.Cells)
|
||||
}
|
||||
res, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO tables(document_id, table_index, row_count, cell_count) VALUES (?, ?, ?, ?)`,
|
||||
documentID,
|
||||
tableIndex,
|
||||
len(table.Rows),
|
||||
cellCount,
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("insert table %d: %w", tableIndex, err)
|
||||
}
|
||||
tableID, err := res.LastInsertId()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
for rowIndex, row := range table.Rows {
|
||||
for cellIndex, cell := range row.Cells {
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO cells(document_id, table_id, table_index, row_index, cell_index, col_index, col_span, v_merge, text)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
documentID,
|
||||
tableID,
|
||||
tableIndex,
|
||||
rowIndex,
|
||||
cellIndex,
|
||||
cell.ColIndex,
|
||||
cell.ColSpan,
|
||||
cell.VMerge,
|
||||
cell.Text,
|
||||
); err != nil {
|
||||
return fmt.Errorf("insert cell table=%d row=%d cell=%d: %w", tableIndex, rowIndex, cellIndex, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for _, field := range inferFields(table) {
|
||||
if _, err := tx.ExecContext(ctx,
|
||||
`INSERT INTO fields(document_id, table_id, table_index, record_index, row_index, col_index, key, value, source_kind)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)`,
|
||||
documentID,
|
||||
tableID,
|
||||
tableIndex,
|
||||
field.RecordIndex,
|
||||
field.RowIndex,
|
||||
field.ColIndex,
|
||||
field.Key,
|
||||
field.Value,
|
||||
field.SourceKind,
|
||||
); err != nil {
|
||||
return fmt.Errorf("insert field %q: %w", field.Key, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return tx.Commit()
|
||||
}
|
||||
|
||||
func ensureSchema(ctx context.Context, db execer) error {
|
||||
statements := []string{
|
||||
`PRAGMA foreign_keys = ON`,
|
||||
`CREATE TABLE IF NOT EXISTS documents (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
path TEXT NOT NULL,
|
||||
filename TEXT NOT NULL,
|
||||
sha256 TEXT NOT NULL,
|
||||
imported_at TEXT NOT NULL
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS tables (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
|
||||
table_index INTEGER NOT NULL,
|
||||
row_count INTEGER NOT NULL,
|
||||
cell_count INTEGER NOT NULL
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS cells (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
|
||||
table_id INTEGER NOT NULL REFERENCES tables(id) ON DELETE CASCADE,
|
||||
table_index INTEGER NOT NULL,
|
||||
row_index INTEGER NOT NULL,
|
||||
cell_index INTEGER NOT NULL,
|
||||
col_index INTEGER NOT NULL,
|
||||
col_span INTEGER NOT NULL,
|
||||
v_merge TEXT NOT NULL DEFAULT '',
|
||||
text TEXT NOT NULL
|
||||
)`,
|
||||
`CREATE TABLE IF NOT EXISTS fields (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
document_id INTEGER NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
|
||||
table_id INTEGER NOT NULL REFERENCES tables(id) ON DELETE CASCADE,
|
||||
table_index INTEGER NOT NULL,
|
||||
record_index INTEGER NOT NULL,
|
||||
row_index INTEGER NOT NULL,
|
||||
col_index INTEGER NOT NULL,
|
||||
key TEXT NOT NULL,
|
||||
value TEXT NOT NULL,
|
||||
source_kind TEXT NOT NULL
|
||||
)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_cells_doc_table ON cells(document_id, table_index, row_index, cell_index)`,
|
||||
`CREATE INDEX IF NOT EXISTS idx_fields_doc_key ON fields(document_id, key)`,
|
||||
}
|
||||
for _, statement := range statements {
|
||||
if _, err := db.ExecContext(ctx, statement); err != nil {
|
||||
return fmt.Errorf("schema statement failed: %w", err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type execer interface {
|
||||
ExecContext(context.Context, string, ...any) (sql.Result, error)
|
||||
}
|
||||
|
||||
func readDocxTables(path string) ([]Table, error) {
|
||||
zr, err := zip.OpenReader(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("open docx zip: %w", err)
|
||||
}
|
||||
defer zr.Close()
|
||||
|
||||
for _, file := range zr.File {
|
||||
if file.Name != "word/document.xml" {
|
||||
continue
|
||||
}
|
||||
rc, err := file.Open()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer rc.Close()
|
||||
data, err := io.ReadAll(rc)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return parseDocumentXML(data)
|
||||
}
|
||||
return nil, errors.New("word/document.xml not found in docx")
|
||||
}
|
||||
|
||||
func parseDocumentXML(data []byte) ([]Table, error) {
|
||||
decoder := xml.NewDecoder(bytes.NewReader(data))
|
||||
var tables []Table
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return tables, nil
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
start, ok := token.(xml.StartElement)
|
||||
if !ok || start.Name.Local != "tbl" {
|
||||
continue
|
||||
}
|
||||
table, err := parseTable(decoder)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
table.Index = len(tables)
|
||||
tables = append(tables, table)
|
||||
}
|
||||
}
|
||||
|
||||
func parseTable(decoder *xml.Decoder) (Table, error) {
|
||||
var table Table
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if err != nil {
|
||||
return table, err
|
||||
}
|
||||
switch t := token.(type) {
|
||||
case xml.StartElement:
|
||||
if t.Name.Local == "tr" {
|
||||
row, err := parseRow(decoder, len(table.Rows))
|
||||
if err != nil {
|
||||
return table, err
|
||||
}
|
||||
table.Rows = append(table.Rows, row)
|
||||
}
|
||||
case xml.EndElement:
|
||||
if t.Name.Local == "tbl" {
|
||||
return table, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func parseRow(decoder *xml.Decoder, rowIndex int) (Row, error) {
|
||||
row := Row{Index: rowIndex}
|
||||
logicalCol := 0
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if err != nil {
|
||||
return row, err
|
||||
}
|
||||
switch t := token.(type) {
|
||||
case xml.StartElement:
|
||||
if t.Name.Local == "tc" {
|
||||
cell, err := parseCell(decoder)
|
||||
if err != nil {
|
||||
return row, err
|
||||
}
|
||||
cell.RowIndex = rowIndex
|
||||
cell.CellIndex = len(row.Cells)
|
||||
cell.ColIndex = logicalCol
|
||||
if cell.ColSpan < 1 {
|
||||
cell.ColSpan = 1
|
||||
}
|
||||
logicalCol += cell.ColSpan
|
||||
row.Cells = append(row.Cells, cell)
|
||||
}
|
||||
case xml.EndElement:
|
||||
if t.Name.Local == "tr" {
|
||||
return row, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func parseCell(decoder *xml.Decoder) (Cell, error) {
|
||||
cell := Cell{ColSpan: 1}
|
||||
var b strings.Builder
|
||||
inText := false
|
||||
tcDepth := 1
|
||||
for {
|
||||
token, err := decoder.Token()
|
||||
if err != nil {
|
||||
return cell, err
|
||||
}
|
||||
switch t := token.(type) {
|
||||
case xml.StartElement:
|
||||
switch t.Name.Local {
|
||||
case "tc":
|
||||
tcDepth++
|
||||
case "gridSpan":
|
||||
if tcDepth == 1 {
|
||||
if span, err := strconv.Atoi(attr(t, "val")); err == nil && span > 0 {
|
||||
cell.ColSpan = span
|
||||
}
|
||||
}
|
||||
case "vMerge":
|
||||
if tcDepth == 1 {
|
||||
cell.VMerge = attr(t, "val")
|
||||
if cell.VMerge == "" {
|
||||
cell.VMerge = "continue"
|
||||
}
|
||||
}
|
||||
case "t":
|
||||
inText = true
|
||||
case "tab":
|
||||
b.WriteByte('\t')
|
||||
case "br", "cr":
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
case xml.CharData:
|
||||
if inText {
|
||||
b.Write([]byte(t))
|
||||
}
|
||||
case xml.EndElement:
|
||||
switch t.Name.Local {
|
||||
case "t":
|
||||
inText = false
|
||||
case "p":
|
||||
if b.Len() > 0 && !strings.HasSuffix(b.String(), "\n") {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
case "tc":
|
||||
tcDepth--
|
||||
if tcDepth == 0 {
|
||||
cell.Text = normalizeCellText(b.String())
|
||||
return cell, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func attr(start xml.StartElement, local string) string {
|
||||
for _, a := range start.Attr {
|
||||
if a.Name.Local == local {
|
||||
return a.Value
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func normalizeCellText(s string) string {
|
||||
s = strings.ReplaceAll(s, "\r\n", "\n")
|
||||
s = strings.ReplaceAll(s, "\r", "\n")
|
||||
lines := strings.Split(s, "\n")
|
||||
out := make([]string, 0, len(lines))
|
||||
for _, line := range lines {
|
||||
line = strings.TrimSpace(strings.Join(strings.Fields(line), " "))
|
||||
if line != "" {
|
||||
out = append(out, line)
|
||||
}
|
||||
}
|
||||
return strings.Join(out, "\n")
|
||||
}
|
||||
|
||||
func inferFields(table Table) []Field {
|
||||
headerIndex := detectHeaderRow(table)
|
||||
if headerIndex >= 0 {
|
||||
return inferHeaderFields(table, headerIndex)
|
||||
}
|
||||
return inferKeyValueFields(table)
|
||||
}
|
||||
|
||||
func detectHeaderRow(table Table) int {
|
||||
for i := 0; i < len(table.Rows)-1; i++ {
|
||||
row := table.Rows[i]
|
||||
if isTitleRow(row) {
|
||||
continue
|
||||
}
|
||||
if nonEmptyCellCount(row) >= 2 && nonEmptyCellCount(table.Rows[i+1]) >= 2 {
|
||||
return i
|
||||
}
|
||||
}
|
||||
return -1
|
||||
}
|
||||
|
||||
func isTitleRow(row Row) bool {
|
||||
nonEmpty := nonEmptyCellCount(row)
|
||||
if nonEmpty != 1 || len(row.Cells) == 0 {
|
||||
return false
|
||||
}
|
||||
return len(row.Cells) == 1 || row.Cells[0].ColSpan > 1
|
||||
}
|
||||
|
||||
func nonEmptyCellCount(row Row) int {
|
||||
count := 0
|
||||
for _, cell := range row.Cells {
|
||||
if strings.TrimSpace(cell.Text) != "" {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func inferHeaderFields(table Table, headerIndex int) []Field {
|
||||
headers := table.Rows[headerIndex].Cells
|
||||
var fields []Field
|
||||
recordIndex := 0
|
||||
for rowIndex := headerIndex + 1; rowIndex < len(table.Rows); rowIndex++ {
|
||||
row := table.Rows[rowIndex]
|
||||
if nonEmptyCellCount(row) == 0 {
|
||||
continue
|
||||
}
|
||||
for i, header := range headers {
|
||||
if i >= len(row.Cells) {
|
||||
continue
|
||||
}
|
||||
key := cleanFieldKey(header.Text)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
value := strings.TrimSpace(row.Cells[i].Text)
|
||||
fields = append(fields, Field{
|
||||
RecordIndex: recordIndex,
|
||||
RowIndex: rowIndex,
|
||||
ColIndex: row.Cells[i].ColIndex,
|
||||
Key: key,
|
||||
Value: value,
|
||||
SourceKind: "header_row",
|
||||
})
|
||||
}
|
||||
recordIndex++
|
||||
}
|
||||
return fields
|
||||
}
|
||||
|
||||
func inferKeyValueFields(table Table) []Field {
|
||||
var fields []Field
|
||||
recordIndex := 0
|
||||
for rowIndex, row := range table.Rows {
|
||||
cells := nonEmptyCells(row)
|
||||
if len(cells) < 2 || len(cells)%2 != 0 {
|
||||
continue
|
||||
}
|
||||
for i := 0; i < len(cells); i += 2 {
|
||||
key := cleanFieldKey(cells[i].Text)
|
||||
if key == "" {
|
||||
continue
|
||||
}
|
||||
fields = append(fields, Field{
|
||||
RecordIndex: recordIndex,
|
||||
RowIndex: rowIndex,
|
||||
ColIndex: cells[i+1].ColIndex,
|
||||
Key: key,
|
||||
Value: cells[i+1].Text,
|
||||
SourceKind: "key_value",
|
||||
})
|
||||
}
|
||||
recordIndex++
|
||||
}
|
||||
return fields
|
||||
}
|
||||
|
||||
func nonEmptyCells(row Row) []Cell {
|
||||
out := make([]Cell, 0, len(row.Cells))
|
||||
for _, cell := range row.Cells {
|
||||
if strings.TrimSpace(cell.Text) != "" {
|
||||
out = append(out, cell)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cleanFieldKey(s string) string {
|
||||
s = strings.TrimSpace(s)
|
||||
s = strings.TrimRight(s, "：: \t\r\n")
|
||||
return strings.TrimSpace(s)
|
||||
}
|
||||
|
||||
func fileSHA256(path string) (string, error) {
|
||||
f, err := os.Open(path)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer f.Close()
|
||||
h := sha256.New()
|
||||
if _, err := io.Copy(h, f); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hex.EncodeToString(h.Sum(nil)), nil
|
||||
}
|
||||
Reference in New Issue
Block a user