szczanieckiej/traffic/index.go

// SPDX-FileCopyrightText: Adam Evyčędo
//
// SPDX-License-Identifier: AGPL-3.0-or-later

package traffic

import (
	"database/sql"
	"errors"
	"fmt"
	"path/filepath"
	"strings"
)

var indices map[string]map[string]*sql.DB

func OpenIndices(context Context) error {
	indexDbName := filepath.Join(context.DataHome, context.FeedID, string(context.Version), FILE_INDEX_DB)
	indexDb, err := sql.Open("sqlite3", indexDbName)
	if err != nil {
		return fmt.Errorf("while opening index: %w", err)
	}
	if indices == nil {
		indices = map[string]map[string]*sql.DB{}
	}
	if indices[context.FeedID] == nil {
		indices[context.FeedID] = map[string]*sql.DB{}
	}
	indices[context.FeedID][string(context.Version)] = indexDb

	return nil
}

func ExistsStopForCode(code string, context Context) (bool, error) {
	return existsOffsetForID(code, context, INDEX_STOP_CODE)
}

func existsOffsetForID(id string, context Context, table string) (bool, error) {
	_, err := getOffsetByID(id, context, table)
	if err == nil {
		return true, nil
	} else {
		if errors.Is(err, sql.ErrNoRows) {
			return false, nil
		} else {
			return false, err
		}
	}
}

func getLineOffsetByID(id string, context Context) (uint, error) {
	return getOffsetByID(id, context, INDEX_LINE_CODE)
}

func getStopOffsetByCode(id string, context Context) (uint, error) {
	return getOffsetByID(id, context, INDEX_STOP_CODE)
}

func getTripOffset(id string, context Context) (uint, error) {
	return getOffsetByID(id, context, INDEX_TRIPS)
}

func getOffsetByID(id string, context Context, table string) (uint, error) {
	indexDb := indices[context.FeedID][string(context.Version)]
	row := indexDb.QueryRow("select offset from "+table+" where code == ?", id)
	var offset uint
	err := row.Scan(&offset)
	return offset, err
}

func getTripsOffsets(ids []string, context Context) ([]uint, error) {
	return getOffsetsByIDs(ids, context, INDEX_TRIPS)
}

func getOffsetsByIDs(ids []string, context Context, table string) ([]uint, error) {
	indexDb := indices[context.FeedID][string(context.Version)]
	offsets := []uint{}
	placeholders := make([]string, len(ids))
	idValues := make([]any, len(ids))
	for i := range ids {
		placeholders[i] = "?"
		idValues[i] = ids[i]
	}
	rows, err := indexDb.Query("select offset from "+table+" where code in ("+strings.Join(placeholders, ", ")+")", idValues...)
	if err != nil {
		return offsets, fmt.Errorf("while selecting: %w", err)
	}
	for rows.Next() {
		var offset uint
		err := rows.Scan(&offset)
		if err != nil {
			return offsets, fmt.Errorf("while scanning: %w", err)
		}
		offsets = append(offsets, offset)
	}
	return offsets, err

}

func getLineOffsetsByName(name string, context Context) ([]uint, error) {
	indexDb := indices[context.FeedID][string(context.Version)]
	offsets := []uint{}
	rows, err := indexDb.Query("select offset from "+INDEX_LINES+" where name == ?", name)
	if err != nil {
		return offsets, fmt.Errorf("while selecting: %w", err)
	}
	for rows.Next() {
		var offset uint
		err := rows.Scan(&offset)
		if err != nil {
			return offsets, fmt.Errorf("while scanning: %w", err)
		}
		offsets = append(offsets, offset)
	}
	return offsets, err
}

func getNameIndex(feedID string, validity Validity) (NameIndexDB, error) {
	indexDb := indices[feedID][string(validity)]
	tx, err := indexDb.Begin()
	if err != nil {
		return NameIndexDB{}, fmt.Errorf("while beginnging transaction: %w", err)
	}

	return NameIndexDB{
		tx:    tx,
		table: INDEX_LINES,
	}, nil
}

func getLineOffsetByRowID(i int, context Context) (uint, error) {
	return getOffsetByRowID(i, context, INDEX_LINES)
}

func getStopOffsetByRowID(i int, context Context) (uint, error) {
	return getOffsetByRowID(i, context, INDEX_STOPS)
}

func getOffsetByRowID(i int, context Context, table string) (uint, error) {
	indexDb := indices[context.FeedID][string(context.Version)]
	row := indexDb.QueryRow("select offset from "+table+" where rowid == ?", i+1)
	var offset uint
	err := row.Scan(&offset)
	return offset, err
}