sql_database.go

package main

import (
	"database/sql"
	"os"

	_ "github.com/mattn/go-sqlite3"
)

type SQLDataBase struct {
	db_path string
	db      *sql.DB
}

func NewSQLDataBase(db_path string) (*SQLDataBase, error) {
	// Create database
	db, err := sql.Open("sqlite3", db_path)
	if err != nil {
		return nil, err
	}

	// Create tables
	_, err = db.Exec(`
		PRAGMA foreign_keys=ON;
		BEGIN TRANSACTION;
		CREATE TABLE IF NOT EXISTS words(
			id INTEGER PRIMARY KEY, 
			word STRING UNIQUE
		);
		CREATE TABLE IF NOT EXISTS documents(
			id INTEGER PRIMARY KEY, 
			title STRING,
			url STRING UNIQUE, 
			wordcount INTEGER
		);
		CREATE TABLE IF NOT EXISTS frequencies(
			id INTEGER PRIMARY KEY, 
			word_id INTEGER, 
			doc_id INTEGER, 
			freq INTEGER, 
			FOREIGN KEY(word_id) REFERENCES words(id), 
			FOREIGN KEY(doc_id) REFERENCES documents(id),
			UNIQUE(word_id, doc_id)
		);
		COMMIT;
	`)
	if err != nil {
		return nil, err
	}

	return &SQLDataBase{db_path, db}, nil
}

func (sdb *SQLDataBase) GetAllDocs() (map[string]struct{}, error) {
	// Create freq map and get rows
	m := make(map[string]struct{})
	rows, err := sdb.db.Query(`SELECT url FROM documents`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Loop through rows to make map
	for rows.Next() {
		var url string
		err = rows.Scan(&url)
		if err != nil {
			return nil, err
		}
		m[url] = struct{}{}
	}
	return m, nil
}

func (sdb *SQLDataBase) GetAllFreq(doc string) (map[string]uint, error) {
	// Start transaction
	tx, err := sdb.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Commit()

	// Get doc id
	doc_id, err := sdb.docToId(tx, doc)
	if err == sql.ErrNoRows {
		return make(map[string]uint), nil
	}
	if err != nil {
		return nil, err
	}

	// Create freq map and get rows
	m := make(map[string]uint)
	rows, err := tx.Query(`SELECT word_id, freq FROM frequencies WHERE doc_id = ?`, doc_id)
	if err == sql.ErrNoRows {
		return make(map[string]uint), nil
	}
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Loop through rows to make map
	for rows.Next() {
		var word_id, freq int
		err = rows.Scan(&word_id, &freq)
		if err != nil {
			return nil, err
		}
		word, err := sdb.idToWord(tx, word_id)
		if err != nil {
			return nil, err
		}
		m[word] = uint(freq)
	}
	return m, nil
}

func (sdb *SQLDataBase) GetDocsWithWord(word string) (map[string]struct {
	title      string
	word_count uint
	freq       uint
}, error) {
	// Start transaction
	tx, err := sdb.db.Begin()
	if err != nil {
		return nil, err
	}
	defer tx.Commit()

	// Get word id
	word_id, err := sdb.wordToId(tx, word)
	if err == sql.ErrNoRows {
		return make(map[string]struct {
			title      string
			word_count uint
			freq       uint
		}), nil
	}
	if err != nil {
		return nil, err
	}

	// Create freq map and get rows
	m := make(map[string]struct {
		title      string
		word_count uint
		freq       uint
	})
	rows, err := tx.Query(`SELECT doc_id, freq FROM frequencies WHERE word_id = ?`, word_id)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	// Loop through rows to make map
	for rows.Next() {
		var doc_id, freq int
		err = rows.Scan(&doc_id, &freq)
		if err != nil {
			return nil, err
		}
		doc, err := sdb.idToDoc(tx, doc_id)
		if err != nil {
			return nil, err
		}
		m[doc.url] = struct {
			title      string
			word_count uint
			freq       uint
		}{
			doc.title,
			doc.word_count,
			uint(freq),
		}
	}
	return m, nil
}

func (sdb *SQLDataBase) AddDocWithWords(title string, url string, wordcount uint, words map[string]uint) error {
	// Start transaction
	tx, err := sdb.db.Begin()
	if err != nil {
		return err
	}
	defer tx.Commit()

	// Add doc
	_, err = tx.Exec(`INSERT INTO documents(title, url, wordcount) VALUES(?, ?, ?)`, title, url, wordcount)
	if err != nil {
		return err
	}

	// Get doc id
	doc_id, err := sdb.docToId(tx, url)
	if err != nil {
		return err
	}

	// Add all words
	for word, freq := range words {
		// Get word id and add word to table if new
		word_id, err := sdb.wordToId(tx, word)
		if err == sql.ErrNoRows {
			_, err = tx.Exec(`INSERT INTO words(word) VALUES(?)`, word)
			if err != nil {
				return err
			}
			word_id, err = sdb.wordToId(tx, word)
		}
		if err != nil {
			return err
		}

		// Add freq
		_, err = tx.Exec(`INSERT OR REPLACE INTO frequencies(doc_id, word_id, freq) VALUES(?, ?, ?)`, doc_id, word_id, freq)
		if err != nil {
			return err
		}
	}

	return nil
}

func (sdb *SQLDataBase) ContainsDoc(doc string) (bool, error) {
	ids, err := sdb.db.Query(`SELECT id FROM documents WHERE url = ?`, doc)
	if err != nil {
		return false, err
	}
	defer ids.Close()
	if ids.Next() {
		return true, nil
	} else {
		return false, nil
	}
}

func (sdb *SQLDataBase) Size() (uint, error) {
	var count int
	row := sdb.db.QueryRow(`SELECT COUNT(*) FROM documents`)
	err := row.Scan(&count)
	if err != nil {
		return 0, err
	}
	return uint(count), nil
}

// Delete the database
func (sdb *SQLDataBase) DeleteDataBase() {
	sdb.Close()
	os.Remove(sdb.db_path)
}

// Close closes the database connection
func (sdb *SQLDataBase) Close() error {
	if sdb.db != nil {
		return sdb.db.Close()
	}
	return nil
}

// -------- Internal Methods -------- //

func (sdb *SQLDataBase) docToId(tx *sql.Tx, doc string) (int, error) {
	ids, err := tx.Query(`SELECT id FROM documents WHERE url = ?`, doc)
	if err != nil {
		return 0, err
	}
	defer ids.Close()
	var id int
	if ids.Next() {
		err = ids.Scan(&id)
		if err != nil {
			return 0, err
		}
	} else {
		return 0, sql.ErrNoRows
	}
	return id, nil
}

func (sdb *SQLDataBase) idToDoc(tx *sql.Tx, id int) (struct {
	title      string
	url        string
	word_count uint
}, error) {
	docs, err := tx.Query(`SELECT title, url, wordcount FROM documents WHERE id = ?`, id)
	if err != nil {
		return struct {
			title      string
			url        string
			word_count uint
		}{}, err
	}
	defer docs.Close()
	var url, title string
	var word_count int
	if docs.Next() {
		err = docs.Scan(&title, &url, &word_count)
		if err != nil {
			return struct {
				title      string
				url        string
				word_count uint
			}{}, err
		}
	} else {
		return struct {
			title      string
			url        string
			word_count uint
		}{}, sql.ErrNoRows
	}
	return struct {
		title      string
		url        string
		word_count uint
	}{title, url, uint(word_count)}, nil
}

func (sdb *SQLDataBase) wordToId(tx *sql.Tx, word string) (int, error) {
	ids, err := tx.Query(`SELECT id FROM words WHERE word = ?`, word)
	if err != nil {
		return 0, err
	}
	defer ids.Close()
	var id int
	if ids.Next() {
		err = ids.Scan(&id)
		if err != nil {
			return 0, err
		}
	} else {
		return 0, sql.ErrNoRows
	}
	return id, nil
}

func (sdb *SQLDataBase) idToWord(tx *sql.Tx, id int) (string, error) {
	words, err := tx.Query(`SELECT word FROM words WHERE id = ?`, id)
	if err != nil {
		return "", err
	}
	defer words.Close()
	var word string
	if words.Next() {
		err = words.Scan(&word)
		if err != nil {
			return "", err
		}
	} else {
		return "", sql.ErrNoRows
	}
	return word, nil
}