Coverage for src/dataknobs_data/backends/elasticsearch.py: 15%

1"""Elasticsearch backend implementation for the data package."""

3import asyncio

4import time

5import uuid

6from typing import Any, AsyncIterator, Iterator, Optional

8from dataknobs_config import ConfigurableBase

9from dataknobs_utils.elasticsearch_utils import SimplifiedElasticsearchIndex

11from ..database import AsyncDatabase, SyncDatabase

12from ..exceptions import DatabaseError

13from ..query import Operator, Query, SortOrder

14from ..records import Record

15from ..streaming import StreamConfig, StreamResult, StreamingMixin

18class SyncElasticsearchDatabase(SyncDatabase, StreamingMixin, ConfigurableBase):

19 """Synchronous Elasticsearch database backend."""

21 def __init__(self, config: dict[str, Any] | None = None):

22 """Initialize Elasticsearch database.

24 Args:

25 config: Configuration with the following optional keys:

26 - host: Elasticsearch host (default: localhost)

27 - port: Elasticsearch port (default: 9200)

28 - index: Index name (default: "records")

29 - refresh: Whether to refresh after write operations (default: True)

30 - settings: Index settings dict

31 - mappings: Index mappings dict

32 """

33 super().__init__(config)

34 self.es_index = None # Will be initialized in connect()

35 self._connected = False

37 @classmethod

38 def from_config(cls, config: dict) -> "SyncElasticsearchDatabase":

39 """Create from config dictionary."""

40 return cls(config)

42 def connect(self) -> None:

43 """Connect to the Elasticsearch database."""

44 if self._connected:

45 return # Already connected

47 # Initialize the Elasticsearch connection and index

48 config = self.config.copy()

50 # Extract configuration

51 self.host = config.pop("host", "localhost")

52 self.port = config.pop("port", 9200)

53 self.index_name = config.pop("index", "records")

54 self.refresh = config.pop("refresh", True)

56 # Create index settings

57 settings = config.pop("settings", {

58 "number_of_shards": 1,

59 "number_of_replicas": 0,

60 })

62 # Create mappings for flexible schema

63 mappings = config.pop("mappings", {

64 "properties": {

65 "id": {"type": "keyword"},

66 "data": {"type": "object", "enabled": True},

67 "metadata": {"type": "object", "enabled": True},

68 "created_at": {"type": "date"},

69 "updated_at": {"type": "date"},

70 }

71 })

73 # Initialize the Elasticsearch index

74 self.es_index = SimplifiedElasticsearchIndex(

75 index_name=self.index_name,

76 host=self.host,

77 port=self.port,

78 settings=settings,

79 mappings=mappings,

80 )

82 # Ensure index exists

83 if not self.es_index.exists():

84 self.es_index.create()

86 self._connected = True

88 def close(self) -> None:

89 """Close the database connection."""

90 if self.es_index:

91 # ElasticsearchIndex manages its own connections

92 self._connected = False

94 def _initialize(self) -> None:

95 """Initialize method - connection setup moved to connect()."""

96 # Configuration parsing stays here if needed

97 pass

99 def _check_connection(self) -> None:

100 """Check if database is connected."""

101 if not self._connected or not self.es_index:

102 raise RuntimeError("Database not connected. Call connect() first.")

103

104 def _record_to_doc(self, record: Record, id: str | None = None) -> dict[str, Any]:

105 """Convert a Record to an Elasticsearch document."""

106 data = {}

107 for field_name, field_obj in record.fields.items():

108 data[field_name] = field_obj.value

109

110 doc = {

111 "id": id or str(uuid.uuid4()),

112 "data": data,

113 "metadata": record.metadata or {},

114 }

115

116 return doc

117

118 def _doc_to_record(self, doc: dict[str, Any]) -> Record:

119 """Convert an Elasticsearch document to a Record."""

120 data = doc.get("data", {})

121 metadata = doc.get("metadata", {})

122

123 return Record(data=data, metadata=metadata)

124

125 def create(self, record: Record) -> str:

126 """Create a new record."""

127 # Use record's ID if it has one, otherwise generate a new one

128 id = record.id if record.id else str(uuid.uuid4())

129 doc = self._record_to_doc(record, id)

130

131 # Index the document

132 response = self.es_index.index(

133 body=doc,

134 doc_id=id,

135 refresh=self.refresh,

136 )

137

138 if not response.get("_id"):

139 raise DatabaseError(f"Failed to create record: {response}")

140

141 return response["_id"]

142

143 def read(self, id: str) -> Record | None:

144 """Read a record by ID."""

145 response = self.es_index.get(doc_id=id)

146

147 if not response:

148 return None

149

150 doc = response.get("_source", {})

151 return self._doc_to_record(doc)

152

153 def update(self, id: str, record: Record) -> bool:

154 """Update an existing record."""

155 doc = self._record_to_doc(record, id)

156

157 # Update the document

158 success = self.es_index.update(

159 doc_id=id,

160 body={"doc": doc},

161 refresh=self.refresh,

162 )

163

164 return success

165

166 def delete(self, id: str) -> bool:

167 """Delete a record by ID."""

168 success = self.es_index.delete(doc_id=id)

169

170 # Refresh if needed

171 if success and self.refresh:

172 self.es_index.refresh()

173

174 return success

175

176 def exists(self, id: str) -> bool:

177 """Check if a record exists."""

178 return self.es_index.exists(doc_id=id)

179

180 def upsert(self, id: str, record: Record) -> str:

181 """Update or insert a record with a specific ID."""

182 doc = self._record_to_doc(record, id)

183 response = self.es_index.index(body=doc, doc_id=id, refresh=self.refresh)

184

185 if response.get("_id"):

186 return id

187 else:

188 raise DatabaseError(f"Failed to upsert record {id}: {response}")

189

190 def create_batch(self, records: list[Record]) -> list[str]:

191 """Create multiple records in batch with a single refresh."""

192 ids = []

193 for record in records:

194 # Generate ID

195 id = str(uuid.uuid4())

196 doc = self._record_to_doc(record, id)

197

198 # Index without refresh (we'll refresh once at the end)

199 response = self.es_index.index(body=doc, doc_id=id, refresh=False)

200

201 if response.get("_id"):

202 ids.append(id)

203 else:

204 ids.append(None)

205

206 # Single refresh after all documents are indexed

207 if self.refresh and any(ids):

208 self.es_index.refresh()

209

210 return ids

211

212 def read_batch(self, ids: list[str]) -> list[Record | None]:

213 """Read multiple records in batch."""

214 records = []

215 for id in ids:

216 record = self.read(id)

217 records.append(record)

218 return records

219

220 def delete_batch(self, ids: list[str]) -> list[bool]:

221 """Delete multiple records in batch with a single refresh."""

222 results = []

223 for id in ids:

224 # Delete without refresh (we'll refresh once at the end)

225 success = self.es_index.delete(doc_id=id)

226 results.append(success)

227

228 # Single refresh after all documents are deleted

229 if self.refresh and any(results):

230 self.es_index.refresh()

231

232 return results

233

234 def search(self, query: Query) -> list[Record]:

235 """Search for records matching a query."""

236 # Build Elasticsearch query from Query object

237 es_query = {"bool": {"must": []}}

238

239 # Apply filters

240 for filter_obj in query.filters:

241 field_path = f"data.{filter_obj.field}"

242

243 # For string fields in exact match queries, use .keyword suffix

244 # LIKE and REGEX need to use the text field, not keyword

245 if filter_obj.operator in [Operator.EQ, Operator.NEQ, Operator.IN, Operator.NOT_IN]:

246 if isinstance(filter_obj.value, str) or (

247 isinstance(filter_obj.value, list) and

248 filter_obj.value and

249 isinstance(filter_obj.value[0], str)

250 ):

251 field_path = f"{field_path}.keyword"

252 elif filter_obj.operator == Operator.LIKE:

253 # Wildcard needs .keyword for proper matching

254 if isinstance(filter_obj.value, str):

255 field_path = f"{field_path}.keyword"

256

257 if filter_obj.operator == Operator.EQ:

258 # Handle boolean values correctly

259 value = str(filter_obj.value).lower() if isinstance(filter_obj.value, bool) else filter_obj.value

260 es_query["bool"]["must"].append({"term": {field_path: value}})

261 elif filter_obj.operator == Operator.NEQ:

262 value = str(filter_obj.value).lower() if isinstance(filter_obj.value, bool) else filter_obj.value

263 es_query["bool"]["must"].append({"bool": {"must_not": {"term": {field_path: value}}}})

264 elif filter_obj.operator == Operator.GT:

265 es_query["bool"]["must"].append({"range": {field_path: {"gt": filter_obj.value}}})

266 elif filter_obj.operator == Operator.GTE:

267 es_query["bool"]["must"].append({"range": {field_path: {"gte": filter_obj.value}}})

268 elif filter_obj.operator == Operator.LT:

269 es_query["bool"]["must"].append({"range": {field_path: {"lt": filter_obj.value}}})

270 elif filter_obj.operator == Operator.LTE:

271 es_query["bool"]["must"].append({"range": {field_path: {"lte": filter_obj.value}}})

272 elif filter_obj.operator == Operator.LIKE:

273 # Convert SQL LIKE pattern to Elasticsearch wildcard

274 # Wildcard queries should use the keyword field for exact matching

275 pattern = filter_obj.value.replace("%", "*").replace("_", "?")

276 # Use the base field path for LIKE (already has .keyword added above if string)

277 es_query["bool"]["must"].append({"wildcard": {field_path: pattern}})

278 elif filter_obj.operator == Operator.IN:

279 es_query["bool"]["must"].append({"terms": {field_path: filter_obj.value}})

280 elif filter_obj.operator == Operator.NOT_IN:

281 es_query["bool"]["must"].append({"bool": {"must_not": {"terms": {field_path: filter_obj.value}}}})

282 elif filter_obj.operator == Operator.EXISTS:

283 es_query["bool"]["must"].append({"exists": {"field": field_path}})

284 elif filter_obj.operator == Operator.NOT_EXISTS:

285 es_query["bool"]["must"].append({"bool": {"must_not": {"exists": {"field": field_path}}}})

286 elif filter_obj.operator == Operator.REGEX:

287 es_query["bool"]["must"].append({"regexp": {field_path: filter_obj.value}})

288 elif filter_obj.operator == Operator.BETWEEN:

289 # Use Elasticsearch's native range query for efficient BETWEEN

290 if isinstance(filter_obj.value, (list, tuple)) and len(filter_obj.value) == 2:

291 lower, upper = filter_obj.value

292 es_query["bool"]["must"].append({

293 "range": {

294 field_path: {

295 "gte": lower,

296 "lte": upper

297 }

298 }

299 })

300 elif filter_obj.operator == Operator.NOT_BETWEEN:

301 # NOT BETWEEN using bool must_not with range

302 if isinstance(filter_obj.value, (list, tuple)) and len(filter_obj.value) == 2:

303 lower, upper = filter_obj.value

304 es_query["bool"]["must"].append({

305 "bool": {

306 "must_not": {

307 "range": {

308 field_path: {

309 "gte": lower,

310 "lte": upper

311 }

312 }

313 }

314 }

315 })

316

317 # If no filters, match all

318 if not es_query["bool"]["must"]:

319 es_query = {"match_all": {}}

320

321 # Build sort

322 sort = []

323 if query.sort_specs:

324 for sort_spec in query.sort_specs:

325 field_path = f"data.{sort_spec.field}"

326 # Don't add .keyword if user already specified it or for common numeric fields

327 # This is a heuristic - ideally we'd check the mapping

328 numeric_fields = ['age', 'salary', 'balance', 'count', 'score', 'amount', 'price', 'index', 'id', 'number', 'total', 'quantity']

329 if (not sort_spec.field.endswith('.keyword') and

330 not sort_spec.field.endswith('.raw') and

331 sort_spec.field.lower() not in numeric_fields):

332 # Likely a text field, add .keyword for sorting

333 field_path = f"data.{sort_spec.field}.keyword"

334 order = "desc" if sort_spec.order == SortOrder.DESC else "asc"

335 sort.append({field_path: {"order": order}})

336

337 # Build search body

338 search_body = {"query": es_query}

339 if sort:

340 search_body["sort"] = sort

341 if query.limit_value:

342 search_body["size"] = query.limit_value

343 if query.offset_value:

344 search_body["from"] = query.offset_value

345

346 # Execute search

347 response = self.es_index.search(body=search_body)

348

349 # Check if the response is valid (has the expected structure)

350 # An empty result set is still a valid response

351 if not hasattr(response, 'json') or response.json is None:

352 raise DatabaseError(f"Invalid search response: {response}")

353

354 # Check for actual errors in the response

355 if 'error' in response.json:

356 raise DatabaseError(f"Failed to search records: {response.json['error']}")

357

358 # Parse results

359 records = []

360 hits = response.json.get("hits", {}).get("hits", [])

361 for hit in hits:

362 doc = hit.get("_source", {})

363 records.append(self._doc_to_record(doc))

364

365 # Apply field projection if specified

366 if query.fields:

367 for record in records:

368 # Keep only specified fields

369 field_names = list(record.fields.keys())

370 for field_name in field_names:

371 if field_name not in query.fields:

372 del record.fields[field_name]

373

374 return records

375

376 def _count_all(self) -> int:

377 """Count all records in the database."""

378 self._check_connection()

379 return self.es_index.count()

380

381 def count(self, query: Query | None = None) -> int:

382 """Count records matching a query using efficient Elasticsearch count.

383

384 Args:

385 query: Optional search query (counts all if None)

386

387 Returns:

388 Number of matching records

389 """

390 if not query or not query.filters:

391 return self._count_all()

392

393 # Build Elasticsearch query from Query object (same as search)

394 es_query = {"bool": {"must": []}}

395

396 for filter_obj in query.filters:

397 field_path = f"data.{filter_obj.field}"

398

399 # For string fields in exact match queries, use .keyword suffix

400 # LIKE and REGEX need different handling

401 if filter_obj.operator in [Operator.EQ, Operator.NEQ, Operator.IN, Operator.NOT_IN]:

402 if isinstance(filter_obj.value, str) or (

403 isinstance(filter_obj.value, list) and

404 filter_obj.value and

405 isinstance(filter_obj.value[0], str)

406 ):

407 field_path = f"{field_path}.keyword"

408 elif filter_obj.operator == Operator.LIKE:

409 # Wildcard needs .keyword for proper matching

410 if isinstance(filter_obj.value, str):

411 field_path = f"{field_path}.keyword"

412

413 if filter_obj.operator == Operator.EQ:

414 # Handle boolean values correctly

415 value = str(filter_obj.value).lower() if isinstance(filter_obj.value, bool) else filter_obj.value

416 es_query["bool"]["must"].append({"term": {field_path: value}})

417 elif filter_obj.operator == Operator.NEQ:

418 value = str(filter_obj.value).lower() if isinstance(filter_obj.value, bool) else filter_obj.value

419 es_query["bool"]["must"].append({"bool": {"must_not": {"term": {field_path: value}}}})

420 elif filter_obj.operator == Operator.GT:

421 es_query["bool"]["must"].append({"range": {field_path: {"gt": filter_obj.value}}})

422 elif filter_obj.operator == Operator.GTE:

423 es_query["bool"]["must"].append({"range": {field_path: {"gte": filter_obj.value}}})

424 elif filter_obj.operator == Operator.LT:

425 es_query["bool"]["must"].append({"range": {field_path: {"lt": filter_obj.value}}})

426 elif filter_obj.operator == Operator.LTE:

427 es_query["bool"]["must"].append({"range": {field_path: {"lte": filter_obj.value}}})

428 elif filter_obj.operator == Operator.LIKE:

429 pattern = filter_obj.value.replace("%", "*").replace("_", "?")

430 es_query["bool"]["must"].append({"wildcard": {field_path: pattern}})

431 elif filter_obj.operator == Operator.IN:

432 es_query["bool"]["must"].append({"terms": {field_path: filter_obj.value}})

433 elif filter_obj.operator == Operator.NOT_IN:

434 es_query["bool"]["must"].append({"bool": {"must_not": {"terms": {field_path: filter_obj.value}}}})

435 elif filter_obj.operator == Operator.EXISTS:

436 es_query["bool"]["must"].append({"exists": {"field": field_path}})

437 elif filter_obj.operator == Operator.NOT_EXISTS:

438 es_query["bool"]["must"].append({"bool": {"must_not": {"exists": {"field": field_path}}}})

439 elif filter_obj.operator == Operator.REGEX:

440 es_query["bool"]["must"].append({"regexp": {field_path: filter_obj.value}})

441

442 # If no filters were added, use match_all

443 if not es_query["bool"]["must"]:

444 es_query = {"match_all": {}}

445

446 # Count with the query

447 return self.es_index.count(body={"query": es_query})

448

449 def clear(self) -> int:

450 """Clear all records from the database."""

451 self._check_connection()

452 # Get count before deletion

453 count = self._count_all()

454

455 # Delete by query - delete all documents

456 response = self.es_index.delete_by_query(

457 body={"query": {"match_all": {}}}

458 )

459

460 # Refresh if needed

461 if self.refresh:

462 self.es_index.refresh()

463

464 return response.get("deleted", count)

465

466 def close(self) -> None:

467 """Close the database connection."""

468 # ElasticsearchIndex manages its own connections

469 pass

470

471 def stream_read(

472 self,

473 query: Optional[Query] = None,

474 config: Optional[StreamConfig] = None

475 ) -> Iterator[Record]:

476 """Stream records from Elasticsearch."""

477 config = config or StreamConfig()

478

479 # Use search to get all matching records

480 if query:

481 records = self.search(query)

482 else:

483 records = self.search(Query())

484

485 # Yield records in batches for consistency

486 for i in range(0, len(records), config.batch_size):

487 batch = records[i:i + config.batch_size]

488 for record in batch:

489 yield record

490

491 def stream_write(

492 self,

493 records: Iterator[Record],

494 config: Optional[StreamConfig] = None

495 ) -> StreamResult:

496 """Stream records into Elasticsearch."""

497 # Use the default implementation from mixin

498 return self._default_stream_write(records, config)

499

500

501# Import the native async implementation

502from .elasticsearch_async import AsyncElasticsearchDatabase