127d9860df
<!-- .github/pull_request_template.md -->
## Description
Add info on dataset database handler used for dataset database
## Type of Change
<!-- Please check the relevant option -->
- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Code refactoring
- [ ] Performance improvement
- [ ] Other (please specify):
## Screenshots/Videos (if applicable)
<!-- Add screenshots or videos to help explain your changes -->
## Pre-submission Checklist
<!-- Please check all boxes that apply before submitting your PR -->
- [ ] **I have tested my changes thoroughly before submitting this PR**
- [ ] **This PR contains minimal changes necessary to address the
issue/feature**
- [ ] My code follows the project's coding standards and style
guidelines
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have added necessary documentation (if applicable)
- [ ] All new and existing tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been
submitted already
- [ ] I have linked any relevant issues in the description
- [ ] My commits have clear and descriptive messages
## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit
* **New Features**
* Datasets now record their assigned vector and graph database handlers,
allowing per-dataset backend selection.
* **Chores**
* Database schema expanded to store handler identifiers per dataset.
* Deletion/cleanup processes now use dataset-level handler info for
accurate removal across backends.
* **Tests**
* Tests updated to include and validate the new handler fields in
dataset creation outputs.
<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
36 lines
1.6 KiB
Python
36 lines
1.6 KiB
Python
from cognee.modules.users.models.DatasetDatabase import DatasetDatabase
|
|
|
|
|
|
async def _get_vector_db_connection_info(dataset_database: DatasetDatabase) -> DatasetDatabase:
|
|
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
|
supported_dataset_database_handlers,
|
|
)
|
|
|
|
handler = supported_dataset_database_handlers[dataset_database.vector_dataset_database_handler]
|
|
return await handler["handler_instance"].resolve_dataset_connection_info(dataset_database)
|
|
|
|
|
|
async def _get_graph_db_connection_info(dataset_database: DatasetDatabase) -> DatasetDatabase:
|
|
from cognee.infrastructure.databases.dataset_database_handler.supported_dataset_database_handlers import (
|
|
supported_dataset_database_handlers,
|
|
)
|
|
|
|
handler = supported_dataset_database_handlers[dataset_database.graph_dataset_database_handler]
|
|
return await handler["handler_instance"].resolve_dataset_connection_info(dataset_database)
|
|
|
|
|
|
async def resolve_dataset_database_connection_info(
|
|
dataset_database: DatasetDatabase,
|
|
) -> DatasetDatabase:
|
|
"""
|
|
Resolve the connection info for the given DatasetDatabase instance.
|
|
Resolve both vector and graph database connection info and return the updated DatasetDatabase instance.
|
|
|
|
Args:
|
|
dataset_database: DatasetDatabase instance
|
|
Returns:
|
|
DatasetDatabase instance with resolved connection info
|
|
"""
|
|
dataset_database = await _get_vector_db_connection_info(dataset_database)
|
|
dataset_database = await _get_graph_db_connection_info(dataset_database)
|
|
return dataset_database
|