feat: add resumable local OpenSearch loading
This commit is contained in:
@@ -1,11 +1,15 @@
|
||||
import hashlib
|
||||
import io
|
||||
import json
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest
|
||||
import urllib.error
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
from search.minjust_opensearch import export_bulk
|
||||
from search.minjust_opensearch import bulk_batches, document_codes, export_bulk, load_bulk, request_json
|
||||
|
||||
|
||||
class MinjustOpenSearchTest(unittest.TestCase):
|
||||
@@ -14,9 +18,11 @@ class MinjustOpenSearchTest(unittest.TestCase):
|
||||
root = Path(temporary)
|
||||
document_root = root / "normalized/documents/7"
|
||||
document_root.mkdir(parents=True)
|
||||
with sqlite3.connect(root / "normalized/manifest.sqlite3") as connection:
|
||||
connection.execute("CREATE TABLE documents (code TEXT, state TEXT)")
|
||||
connection.execute("INSERT INTO documents VALUES ('7', 'success')")
|
||||
with closing(sqlite3.connect(root / "normalized/manifest.sqlite3")) as connection:
|
||||
with connection:
|
||||
connection.execute("CREATE TABLE documents (code TEXT, state TEXT)")
|
||||
connection.execute("INSERT INTO documents VALUES ('7', 'success')")
|
||||
connection.execute("INSERT INTO documents VALUES ('8', 'success')")
|
||||
(document_root / "document.json").write_text(
|
||||
json.dumps(
|
||||
{
|
||||
@@ -33,6 +39,11 @@ class MinjustOpenSearchTest(unittest.TestCase):
|
||||
),
|
||||
encoding="utf-8",
|
||||
)
|
||||
empty = root / "normalized/documents/8"
|
||||
empty.mkdir()
|
||||
(empty / "document.json").write_text(
|
||||
json.dumps({"schema_version": "1", "source_code": "8"}), encoding="utf-8"
|
||||
)
|
||||
for language, text in (("ru", "Текст \"RU\"\nстрока"), ("ky", "Кыргызча текст")):
|
||||
path = document_root / f"editions/10/{language}/fragments.json"
|
||||
path.parent.mkdir(parents=True)
|
||||
@@ -56,7 +67,7 @@ class MinjustOpenSearchTest(unittest.TestCase):
|
||||
)
|
||||
|
||||
output = root / "bulk.ndjson"
|
||||
self.assertEqual(export_bulk(root / "normalized", output, "test-index"), (1, 2))
|
||||
self.assertEqual(export_bulk(root / "normalized", output, "test-index"), (2, 2))
|
||||
content = output.read_bytes()
|
||||
self.assertTrue(content.endswith(b"\n"))
|
||||
lines = [json.loads(line) for line in content.splitlines()]
|
||||
@@ -65,6 +76,186 @@ class MinjustOpenSearchTest(unittest.TestCase):
|
||||
self.assertIn("text_ky", lines[1])
|
||||
self.assertNotIn("text_ru", lines[1])
|
||||
self.assertEqual(lines[3]["text_ru"], "Текст \"RU\"\nстрока")
|
||||
self.assertEqual(
|
||||
list(bulk_batches(iter((("7", b"a\nb\n"), ("8", b"c\nd\n"))), 4)),
|
||||
[("7", b"a\nb\n"), ("8", b"c\nd\n")],
|
||||
)
|
||||
self.assertEqual(document_codes(root / "normalized", 2, "8"), ["8"])
|
||||
with self.assertRaisesRegex(ValueError, "selected range"):
|
||||
document_codes(root / "normalized", 1, "8")
|
||||
|
||||
checkpoint = root / "checkpoint.json"
|
||||
with patch("search.minjust_opensearch.request_json") as request:
|
||||
request.side_effect = [
|
||||
{},
|
||||
{"cluster_uuid": "cluster-1"},
|
||||
{"test-index": {"settings": {"index": {"uuid": "index-1"}}}},
|
||||
{"errors": False, "items": [{"index": {}}, {"index": {}}]},
|
||||
]
|
||||
self.assertEqual(
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"test-index",
|
||||
maximum_bytes=4096,
|
||||
checkpoint=checkpoint,
|
||||
),
|
||||
(2, 2),
|
||||
)
|
||||
self.assertEqual(request.call_args_list[-1].args[3], "application/x-ndjson")
|
||||
state = json.loads(checkpoint.read_text(encoding="utf-8"))
|
||||
self.assertEqual(state["last_document_code"], "7")
|
||||
self.assertTrue(state["complete"])
|
||||
|
||||
state["last_document_code"] = "8"
|
||||
state["complete"] = False
|
||||
checkpoint.write_text(json.dumps(state), encoding="utf-8")
|
||||
with patch("search.minjust_opensearch.request_json") as request:
|
||||
request.side_effect = [
|
||||
{"cluster_uuid": "cluster-2"},
|
||||
{"test-index": {"settings": {"index": {"uuid": "index-1"}}}},
|
||||
]
|
||||
with self.assertRaisesRegex(ValueError, "does not match"):
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"test-index",
|
||||
maximum_bytes=4096,
|
||||
resume=True,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
with patch("search.minjust_opensearch.request_json") as request:
|
||||
request.side_effect = [
|
||||
{"cluster_uuid": "cluster-1"},
|
||||
{"test-index": {"settings": {"index": {"uuid": "index-1"}}}},
|
||||
]
|
||||
with self.assertRaisesRegex(ValueError, "limit"):
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"test-index",
|
||||
maximum_bytes=4096,
|
||||
limit=1,
|
||||
resume=True,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
with patch("search.minjust_opensearch.request_json") as request:
|
||||
request.side_effect = [
|
||||
{"cluster_uuid": "cluster-1"},
|
||||
{"test-index": {"settings": {"index": {"uuid": "index-1"}}}},
|
||||
]
|
||||
self.assertEqual(
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"test-index",
|
||||
maximum_bytes=4096,
|
||||
resume=True,
|
||||
checkpoint=checkpoint,
|
||||
),
|
||||
(1, 0),
|
||||
)
|
||||
self.assertTrue(all(call.args[1] == "GET" for call in request.call_args_list))
|
||||
|
||||
state["last_document_code"] = "9"
|
||||
state["complete"] = False
|
||||
checkpoint.write_text(json.dumps(state), encoding="utf-8")
|
||||
with patch("search.minjust_opensearch.request_json") as request:
|
||||
request.side_effect = [
|
||||
{"cluster_uuid": "cluster-1"},
|
||||
{"test-index": {"settings": {"index": {"uuid": "index-1"}}}},
|
||||
]
|
||||
with self.assertRaisesRegex(ValueError, "Resume document not found"):
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"test-index",
|
||||
maximum_bytes=4096,
|
||||
resume=True,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
failed_checkpoint = root / "failed-checkpoint.json"
|
||||
failure = {
|
||||
"errors": True,
|
||||
"items": [{"index": {"_id": "bad-id", "error": {"type": "mapper", "reason": "bad value"}}}],
|
||||
}
|
||||
failed_requests = [
|
||||
{},
|
||||
{"cluster_uuid": "cluster-1"},
|
||||
{"failed-index": {"settings": {"index": {"uuid": "index-2"}}}},
|
||||
failure,
|
||||
]
|
||||
with patch("search.minjust_opensearch.request_json", side_effect=failed_requests):
|
||||
with self.assertRaisesRegex(RuntimeError, "bad-id: mapper: bad value"):
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"failed-index",
|
||||
maximum_bytes=4096,
|
||||
checkpoint=failed_checkpoint,
|
||||
)
|
||||
failed_state = json.loads(failed_checkpoint.read_text(encoding="utf-8"))
|
||||
self.assertIsNone(failed_state["last_document_code"])
|
||||
self.assertFalse(failed_state["complete"])
|
||||
|
||||
state["last_document_code"] = "8"
|
||||
state["complete"] = False
|
||||
checkpoint.write_text(json.dumps(state), encoding="utf-8")
|
||||
with closing(sqlite3.connect(root / "normalized/manifest.sqlite3")) as connection:
|
||||
with connection:
|
||||
connection.execute("INSERT INTO documents VALUES ('9', 'success')")
|
||||
with patch("search.minjust_opensearch.request_json") as request:
|
||||
request.side_effect = [
|
||||
{"cluster_uuid": "cluster-1"},
|
||||
{"test-index": {"settings": {"index": {"uuid": "index-1"}}}},
|
||||
]
|
||||
with self.assertRaisesRegex(ValueError, "manifest changed"):
|
||||
load_bulk(
|
||||
root / "normalized",
|
||||
"http://127.0.0.1:9200",
|
||||
"test-index",
|
||||
maximum_bytes=4096,
|
||||
resume=True,
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
http_error = urllib.error.HTTPError(
|
||||
"http://127.0.0.1:9200/test",
|
||||
429,
|
||||
"busy",
|
||||
{},
|
||||
io.BytesIO(b"busy"),
|
||||
)
|
||||
with (
|
||||
patch("search.minjust_opensearch.urllib.request.urlopen", side_effect=[http_error, io.BytesIO(b"{}")]),
|
||||
patch("search.minjust_opensearch.time.sleep") as sleep,
|
||||
):
|
||||
self.assertEqual(
|
||||
request_json("http://127.0.0.1:9200/test", "GET", None, "application/json", attempts=2),
|
||||
{},
|
||||
)
|
||||
sleep.assert_called_once_with(1)
|
||||
|
||||
with (
|
||||
patch(
|
||||
"search.minjust_opensearch.urllib.request.urlopen",
|
||||
side_effect=[io.BytesIO(b"{"), io.BytesIO(b"{}")],
|
||||
),
|
||||
patch("search.minjust_opensearch.time.sleep") as sleep,
|
||||
):
|
||||
self.assertEqual(
|
||||
request_json(
|
||||
"http://127.0.0.1:9200/_bulk",
|
||||
"POST",
|
||||
b"{}\n{}\n",
|
||||
"application/x-ndjson",
|
||||
attempts=2,
|
||||
retry_invalid_json=True,
|
||||
),
|
||||
{},
|
||||
)
|
||||
sleep.assert_called_once_with(1)
|
||||
|
||||
bad = document_root / "editions/10/ru/fragments.json"
|
||||
fragments = json.loads(bad.read_text(encoding="utf-8"))
|
||||
@@ -83,9 +274,17 @@ class MinjustOpenSearchTest(unittest.TestCase):
|
||||
export_bulk(root / "normalized", root / "normalized/manifest.sqlite3", "test-index")
|
||||
self.assertEqual(output.read_bytes(), content)
|
||||
|
||||
mapping = json.loads(
|
||||
bad.write_text("", encoding="utf-8")
|
||||
with self.assertRaisesRegex(ValueError, "fragments.json"):
|
||||
export_bulk(root / "normalized", output, "test-index")
|
||||
self.assertEqual(output.read_bytes(), content)
|
||||
|
||||
definition = json.loads(
|
||||
(Path(__file__).parent / "search/minjust-fragments-index.json").read_text(encoding="utf-8")
|
||||
)["mappings"]
|
||||
)
|
||||
mapping = definition["mappings"]
|
||||
self.assertEqual(definition["settings"]["index"]["number_of_shards"], 1)
|
||||
self.assertEqual(definition["settings"]["index"]["number_of_replicas"], 0)
|
||||
self.assertEqual(mapping["dynamic"], "strict")
|
||||
self.assertEqual(mapping["properties"]["position"]["type"], "integer")
|
||||
self.assertEqual(mapping["properties"]["text_ru"]["analyzer"], "russian")
|
||||
|
||||
Reference in New Issue
Block a user