feat: Add location search
parent
f520473ed3
commit
26af8af8a9
@ -0,0 +1,99 @@
|
|||||||
|
package citydb
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"database/sql"
|
||||||
|
"fmt"
|
||||||
|
"strings"
|
||||||
|
|
||||||
|
"github.com/doug-martin/goqu/v9"
|
||||||
|
|
||||||
|
_ "modernc.org/sqlite"
|
||||||
|
|
||||||
|
"prayertimes/pkg/prayer"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Provider struct {
|
||||||
|
db *goqu.Database
|
||||||
|
}
|
||||||
|
|
||||||
|
func New(db *goqu.Database) Provider {
|
||||||
|
return Provider{db: db}
|
||||||
|
}
|
||||||
|
|
||||||
|
func Open(path string) (Provider, error) {
|
||||||
|
conn, err := sql.Open("sqlite", path)
|
||||||
|
if err != nil {
|
||||||
|
return Provider{}, fmt.Errorf("failed to open cities database: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
if err := conn.Ping(); err != nil {
|
||||||
|
return Provider{}, fmt.Errorf("failed to connect to cities database: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
return New(goqu.New("sqlite3", conn)), nil
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p Provider) Close() error {
|
||||||
|
return p.db.Db.(*sql.DB).Close()
|
||||||
|
}
|
||||||
|
|
||||||
|
type locationRow struct {
|
||||||
|
ID int `db:"geoname_id"`
|
||||||
|
Name string `db:"name"`
|
||||||
|
ASCIIName string `db:"ascii_name"`
|
||||||
|
AlternateNames string `db:"alternate_names"`
|
||||||
|
CountryCode string `db:"country_code"`
|
||||||
|
Latitude float64 `db:"latitude"`
|
||||||
|
Longitude float64 `db:"longitude"`
|
||||||
|
}
|
||||||
|
|
||||||
|
func (p Provider) SearchLocations(ctx context.Context, query string) ([]prayer.Location, error) {
|
||||||
|
query = strings.TrimSpace(query)
|
||||||
|
if query == "" {
|
||||||
|
return []prayer.Location{}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
pattern := "%" + query + "%"
|
||||||
|
|
||||||
|
q := p.db.
|
||||||
|
From("cities").
|
||||||
|
Select(
|
||||||
|
goqu.I("geoname_id"),
|
||||||
|
goqu.I("name"),
|
||||||
|
goqu.I("ascii_name"),
|
||||||
|
goqu.COALESCE(goqu.I("alternate_names"), "").As("alternate_names"),
|
||||||
|
goqu.I("country_code"),
|
||||||
|
goqu.I("latitude"),
|
||||||
|
goqu.I("longitude"),
|
||||||
|
).
|
||||||
|
Where(
|
||||||
|
goqu.Or(
|
||||||
|
goqu.L("name LIKE ? COLLATE NOCASE", pattern),
|
||||||
|
goqu.L("ascii_name LIKE ? COLLATE NOCASE", pattern),
|
||||||
|
goqu.L("alternate_names LIKE ? COLLATE NOCASE", pattern),
|
||||||
|
),
|
||||||
|
).
|
||||||
|
Order(goqu.I("population").Desc(), goqu.I("name").Asc()).
|
||||||
|
Limit(50)
|
||||||
|
|
||||||
|
var rows []locationRow
|
||||||
|
if err := q.ScanStructsContext(ctx, &rows); err != nil {
|
||||||
|
return nil, fmt.Errorf("failed to query locations from cities database: %w", err)
|
||||||
|
}
|
||||||
|
|
||||||
|
locations := make([]prayer.Location, 0, len(rows))
|
||||||
|
for _, row := range rows {
|
||||||
|
locations = append(locations, prayer.Location{
|
||||||
|
ID: row.ID,
|
||||||
|
Name: row.Name,
|
||||||
|
ASCIIName: row.ASCIIName,
|
||||||
|
AlternateNames: row.AlternateNames,
|
||||||
|
CountryCode: row.CountryCode,
|
||||||
|
Latitude: row.Latitude,
|
||||||
|
Longitude: row.Longitude,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
return locations, nil
|
||||||
|
}
|
||||||
@ -0,0 +1,182 @@
|
|||||||
|
#!/usr/bin/env python3
|
||||||
|
"""Download GeoNames cities5000 from https://download.geonames.org/export/dump/ and seed a SQLite database."""
|
||||||
|
|
||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import argparse
|
||||||
|
import io
|
||||||
|
import sqlite3
|
||||||
|
import sys
|
||||||
|
import urllib.request
|
||||||
|
import zipfile
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
SOURCE_URL = "https://download.geonames.org/export/dump/cities5000.zip"
|
||||||
|
ARCHIVE_ENTRY = "cities5000.txt"
|
||||||
|
|
||||||
|
CREATE_TABLE_SQL = """
|
||||||
|
CREATE TABLE IF NOT EXISTS cities (
|
||||||
|
geoname_id INTEGER PRIMARY KEY,
|
||||||
|
name TEXT NOT NULL,
|
||||||
|
ascii_name TEXT NOT NULL,
|
||||||
|
alternate_names TEXT,
|
||||||
|
latitude REAL NOT NULL,
|
||||||
|
longitude REAL NOT NULL,
|
||||||
|
feature_class TEXT,
|
||||||
|
feature_code TEXT,
|
||||||
|
country_code TEXT,
|
||||||
|
cc2 TEXT,
|
||||||
|
admin1_code TEXT,
|
||||||
|
admin2_code TEXT,
|
||||||
|
admin3_code TEXT,
|
||||||
|
admin4_code TEXT,
|
||||||
|
population INTEGER,
|
||||||
|
elevation INTEGER,
|
||||||
|
dem INTEGER,
|
||||||
|
timezone TEXT,
|
||||||
|
modification_date TEXT
|
||||||
|
)
|
||||||
|
"""
|
||||||
|
|
||||||
|
CREATE_INDEXES_SQL = (
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_cities_name ON cities(name)",
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_cities_ascii_name ON cities(ascii_name)",
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_cities_alternate_names ON cities(alternate_names)",
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_cities_country_code ON cities(country_code)",
|
||||||
|
"CREATE INDEX IF NOT EXISTS idx_cities_population ON cities(population DESC)",
|
||||||
|
)
|
||||||
|
|
||||||
|
INSERT_SQL = """
|
||||||
|
INSERT OR REPLACE INTO cities (
|
||||||
|
geoname_id,
|
||||||
|
name,
|
||||||
|
ascii_name,
|
||||||
|
alternate_names,
|
||||||
|
latitude,
|
||||||
|
longitude,
|
||||||
|
feature_class,
|
||||||
|
feature_code,
|
||||||
|
country_code,
|
||||||
|
cc2,
|
||||||
|
admin1_code,
|
||||||
|
admin2_code,
|
||||||
|
admin3_code,
|
||||||
|
admin4_code,
|
||||||
|
population,
|
||||||
|
elevation,
|
||||||
|
dem,
|
||||||
|
timezone,
|
||||||
|
modification_date
|
||||||
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||||
|
"""
|
||||||
|
|
||||||
|
|
||||||
|
def parse_args() -> argparse.Namespace:
|
||||||
|
parser = argparse.ArgumentParser(description="Seed GeoNames cities5000 into SQLite")
|
||||||
|
parser.add_argument("--db", required=True, help="Target SQLite database path")
|
||||||
|
parser.add_argument(
|
||||||
|
"--batch-size", type=int, default=1000, help="Insert batch size"
|
||||||
|
)
|
||||||
|
return parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
def to_int(value: str) -> int | None:
|
||||||
|
value = value.strip()
|
||||||
|
if value == "":
|
||||||
|
return None
|
||||||
|
return int(value)
|
||||||
|
|
||||||
|
|
||||||
|
def to_float(value: str) -> float:
|
||||||
|
return float(value.strip())
|
||||||
|
|
||||||
|
|
||||||
|
def row_from_columns(columns: list[str]) -> tuple:
|
||||||
|
return (
|
||||||
|
int(columns[0]),
|
||||||
|
columns[1],
|
||||||
|
columns[2],
|
||||||
|
columns[3],
|
||||||
|
to_float(columns[4]),
|
||||||
|
to_float(columns[5]),
|
||||||
|
columns[6],
|
||||||
|
columns[7],
|
||||||
|
columns[8],
|
||||||
|
columns[9],
|
||||||
|
columns[10],
|
||||||
|
columns[11],
|
||||||
|
columns[12],
|
||||||
|
columns[13],
|
||||||
|
to_int(columns[14]),
|
||||||
|
to_int(columns[15]),
|
||||||
|
to_int(columns[16]),
|
||||||
|
columns[17],
|
||||||
|
columns[18],
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def download_archive() -> bytes:
|
||||||
|
with urllib.request.urlopen(SOURCE_URL) as response:
|
||||||
|
return response.read()
|
||||||
|
|
||||||
|
|
||||||
|
def stream_city_rows(archive_bytes: bytes):
|
||||||
|
with zipfile.ZipFile(io.BytesIO(archive_bytes)) as zf:
|
||||||
|
with zf.open(ARCHIVE_ENTRY) as raw_file:
|
||||||
|
for raw_line in raw_file:
|
||||||
|
line = raw_line.decode("utf-8").rstrip("\n")
|
||||||
|
if not line:
|
||||||
|
continue
|
||||||
|
|
||||||
|
columns = line.split("\t")
|
||||||
|
if len(columns) != 19:
|
||||||
|
continue
|
||||||
|
|
||||||
|
yield row_from_columns(columns)
|
||||||
|
|
||||||
|
|
||||||
|
def seed_database(db_path: Path, batch_size: int) -> int:
|
||||||
|
archive = download_archive()
|
||||||
|
|
||||||
|
inserted = 0
|
||||||
|
batch = []
|
||||||
|
|
||||||
|
db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
with sqlite3.connect(str(db_path)) as conn:
|
||||||
|
cursor = conn.cursor()
|
||||||
|
cursor.execute("PRAGMA journal_mode=WAL")
|
||||||
|
cursor.execute(CREATE_TABLE_SQL)
|
||||||
|
for sql in CREATE_INDEXES_SQL:
|
||||||
|
cursor.execute(sql)
|
||||||
|
|
||||||
|
for row in stream_city_rows(archive):
|
||||||
|
batch.append(row)
|
||||||
|
if len(batch) >= batch_size:
|
||||||
|
cursor.executemany(INSERT_SQL, batch)
|
||||||
|
inserted += len(batch)
|
||||||
|
batch.clear()
|
||||||
|
|
||||||
|
if batch:
|
||||||
|
cursor.executemany(INSERT_SQL, batch)
|
||||||
|
inserted += len(batch)
|
||||||
|
|
||||||
|
conn.commit()
|
||||||
|
|
||||||
|
return inserted
|
||||||
|
|
||||||
|
|
||||||
|
def main() -> int:
|
||||||
|
args = parse_args()
|
||||||
|
try:
|
||||||
|
inserted = seed_database(Path(args.db), args.batch_size)
|
||||||
|
except Exception as exc: # noqa: BLE001
|
||||||
|
print(f"failed to seed cities database: {exc}", file=sys.stderr)
|
||||||
|
return 1
|
||||||
|
|
||||||
|
print(f"seeded {inserted} city rows into {args.db}")
|
||||||
|
return 0
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
raise SystemExit(main())
|
||||||
Loading…
Reference in New Issue