|
12 | 12 | import logging |
13 | 13 | import os |
14 | 14 | from pathlib import Path |
| 15 | +from typing import Any |
15 | 16 |
|
16 | 17 | import aiosqlite |
17 | 18 | import sqlite_vec |
@@ -391,3 +392,188 @@ async def store_embeddings( |
391 | 392 | # Convert to streaming format |
392 | 393 | chunk_embedding_pairs = zip(chunks, embeddings, strict=False) |
393 | 394 | await store_embeddings_streaming(db_path, chunk_embedding_pairs) |
| 395 | + |
| 396 | + |
| 397 | +async def store_trait_implementations( |
| 398 | + db_path: Path, trait_implementations: list[dict[str, Any]], crate_id: int |
| 399 | +) -> None: |
| 400 | + """Store trait implementations in the database. |
| 401 | + |
| 402 | + Args: |
| 403 | + db_path: Path to the database file |
| 404 | + trait_implementations: List of trait implementation dictionaries |
| 405 | + crate_id: Database crate ID for foreign key constraint |
| 406 | + """ |
| 407 | + if not trait_implementations: |
| 408 | + logger.debug("No trait implementations to store") |
| 409 | + return |
| 410 | + |
| 411 | + async with aiosqlite.connect(db_path) as db: |
| 412 | + try: |
| 413 | + # Prepare batch insert for trait implementations |
| 414 | + trait_impl_data = [] |
| 415 | + for impl in trait_implementations: |
| 416 | + trait_impl_data.append(( |
| 417 | + crate_id, |
| 418 | + impl.get("trait_path", ""), |
| 419 | + impl.get("impl_type_path", ""), |
| 420 | + impl.get("generic_params"), |
| 421 | + impl.get("where_clauses"), |
| 422 | + 1 if impl.get("is_blanket", False) else 0, |
| 423 | + 1 if impl.get("is_negative", False) else 0, |
| 424 | + impl.get("impl_signature"), |
| 425 | + None, # source_location - not available from rustdoc JSON |
| 426 | + impl.get("stability_level", "stable"), |
| 427 | + impl.get("item_id", "") |
| 428 | + )) |
| 429 | + |
| 430 | + # Batch insert with proper SQLite syntax |
| 431 | + await db.executemany(""" |
| 432 | + INSERT OR IGNORE INTO trait_implementations ( |
| 433 | + crate_id, trait_path, impl_type_path, generic_params, |
| 434 | + where_clauses, is_blanket, is_negative, impl_signature, |
| 435 | + source_location, stability_level, item_id |
| 436 | + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| 437 | + """, trait_impl_data) |
| 438 | + |
| 439 | + await db.commit() |
| 440 | + logger.info(f"Stored {len(trait_implementations)} trait implementations for crate {crate_id}") |
| 441 | + |
| 442 | + except Exception as e: |
| 443 | + logger.error(f"Error storing trait implementations: {e}") |
| 444 | + await db.rollback() |
| 445 | + raise |
| 446 | + |
| 447 | + |
| 448 | +async def store_trait_definitions( |
| 449 | + db_path: Path, trait_definitions: list[dict[str, Any]], crate_id: int |
| 450 | +) -> None: |
| 451 | + """Store trait definitions in the database. |
| 452 | + |
| 453 | + Note: This stores in trait_implementations table with special markers |
| 454 | + for trait definitions themselves. |
| 455 | + |
| 456 | + Args: |
| 457 | + db_path: Path to the database file |
| 458 | + trait_definitions: List of trait definition dictionaries |
| 459 | + crate_id: Database crate ID for foreign key constraint |
| 460 | + """ |
| 461 | + if not trait_definitions: |
| 462 | + logger.debug("No trait definitions to store") |
| 463 | + return |
| 464 | + |
| 465 | + async with aiosqlite.connect(db_path) as db: |
| 466 | + try: |
| 467 | + # Store trait definitions as special entries |
| 468 | + # We could create a separate traits table, but for MVP we'll use |
| 469 | + # the existing structure with a marker |
| 470 | + trait_def_data = [] |
| 471 | + for trait_def in trait_definitions: |
| 472 | + # Store trait definition as impl of itself |
| 473 | + trait_def_data.append(( |
| 474 | + crate_id, |
| 475 | + trait_def.get("trait_path", ""), |
| 476 | + f"_TRAIT_DEF_{trait_def.get('trait_path', '')}", # Special marker |
| 477 | + trait_def.get("generic_params"), |
| 478 | + json.dumps(trait_def.get("supertraits", [])) if trait_def.get("supertraits") else None, |
| 479 | + 0, # not blanket |
| 480 | + 0, # not negative |
| 481 | + f"trait {trait_def.get('trait_path', '').split('::')[-1]}" if trait_def.get("trait_path") else "", |
| 482 | + None, # source_location |
| 483 | + trait_def.get("stability_level", "stable"), |
| 484 | + trait_def.get("item_id", "") |
| 485 | + )) |
| 486 | + |
| 487 | + await db.executemany(""" |
| 488 | + INSERT OR IGNORE INTO trait_implementations ( |
| 489 | + crate_id, trait_path, impl_type_path, generic_params, |
| 490 | + where_clauses, is_blanket, is_negative, impl_signature, |
| 491 | + source_location, stability_level, item_id |
| 492 | + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) |
| 493 | + """, trait_def_data) |
| 494 | + |
| 495 | + await db.commit() |
| 496 | + logger.info(f"Stored {len(trait_definitions)} trait definitions for crate {crate_id}") |
| 497 | + |
| 498 | + except Exception as e: |
| 499 | + logger.error(f"Error storing trait definitions: {e}") |
| 500 | + await db.rollback() |
| 501 | + raise |
| 502 | + |
| 503 | + |
| 504 | +async def store_enhanced_items_streaming( |
| 505 | + db_path: Path, enhanced_items_stream, crate_id: int |
| 506 | +) -> None: |
| 507 | + """Store enhanced rustdoc items with trait extraction support. |
| 508 | + |
| 509 | + This function handles both regular items and trait-specific data |
| 510 | + from the enhanced rustdoc parser. |
| 511 | + |
| 512 | + Args: |
| 513 | + db_path: Path to the database file |
| 514 | + enhanced_items_stream: Stream of enhanced items with trait data |
| 515 | + crate_id: Database crate ID for trait storage |
| 516 | + """ |
| 517 | + from .signature_extractor import ( |
| 518 | + extract_signature, |
| 519 | + extract_deprecated, |
| 520 | + extract_visibility, |
| 521 | + ) |
| 522 | + from .code_examples import extract_code_examples |
| 523 | + |
| 524 | + regular_items = [] |
| 525 | + trait_implementations = [] |
| 526 | + trait_definitions = [] |
| 527 | + modules_data = None |
| 528 | + |
| 529 | + # Collect items from stream |
| 530 | + async for item in enhanced_items_stream: |
| 531 | + if "_trait_impl" in item: |
| 532 | + # This is trait implementation data |
| 533 | + trait_implementations.append(item["_trait_impl"]) |
| 534 | + elif "_trait_def" in item: |
| 535 | + # This is trait definition data |
| 536 | + trait_definitions.append(item["_trait_def"]) |
| 537 | + elif "_modules" in item: |
| 538 | + # Module hierarchy data |
| 539 | + modules_data = item["_modules"] |
| 540 | + else: |
| 541 | + # Regular item for embedding storage - enhance with metadata |
| 542 | + try: |
| 543 | + item["signature"] = extract_signature(item) |
| 544 | + item["deprecated"] = extract_deprecated(item) |
| 545 | + item["visibility"] = extract_visibility(item) |
| 546 | + item["examples"] = extract_code_examples(item.get("doc", "")) |
| 547 | + regular_items.append(item) |
| 548 | + except Exception as e: |
| 549 | + logger.warning(f"Error enhancing item metadata: {e}") |
| 550 | + regular_items.append(item) # Store anyway |
| 551 | + |
| 552 | + # Store trait implementations and definitions first |
| 553 | + if trait_implementations: |
| 554 | + await store_trait_implementations(db_path, trait_implementations, crate_id) |
| 555 | + |
| 556 | + if trait_definitions: |
| 557 | + await store_trait_definitions(db_path, trait_definitions, crate_id) |
| 558 | + |
| 559 | + # Store module hierarchy if present |
| 560 | + if modules_data: |
| 561 | + try: |
| 562 | + # Import and store modules |
| 563 | + from docsrs_mcp.database.storage import store_modules |
| 564 | + await store_modules(db_path, crate_id, modules_data) |
| 565 | + logger.info(f"Stored module hierarchy with {len(modules_data)} modules") |
| 566 | + except Exception as e: |
| 567 | + logger.warning(f"Error storing modules: {e}") |
| 568 | + |
| 569 | + # Generate embeddings and store regular items if any |
| 570 | + if regular_items: |
| 571 | + logger.info(f"Processing {len(regular_items)} regular items for embedding") |
| 572 | + chunk_embedding_pairs = generate_embeddings_streaming(regular_items) |
| 573 | + await store_embeddings_streaming(db_path, chunk_embedding_pairs) |
| 574 | + |
| 575 | + logger.info( |
| 576 | + f"Enhanced storage complete: {len(regular_items)} items, " |
| 577 | + f"{len(trait_implementations)} trait impls, " |
| 578 | + f"{len(trait_definitions)} trait defs" |
| 579 | + ) |
0 commit comments