lablup / lablup/backend.ai

Introduce DB Source in Backend.AI repository classes

Open
#7,461 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
670
Forks
183
Avg merge
17h 7m
Merged PRs (30d)
358

Description

## Motivation

Backend.AI's Repository layer follows a **Source-based architecture** pattern that separates database access logic into dedicated `DBSource` classes. This pattern provides several critical benefits:

- **Clear Separation of Concerns**: DBSource encapsulates all database operations, isolating them from business logic in Services and Repository orchestration logic
- **Consistent Transaction Management**: DBSource public methods create their own database sessions, while private methods accept sessions as parameters, ensuring proper transaction boundaries
- **Enhanced Testability**: Repositories can be tested directly against real databases using dedicated test fixtures, improving test reliability and catching database-specific issues early
- **Better Observability**: Applying resilience patterns (MetricPolicy, RetryPolicy) to DBSource methods enables automatic collection of granular metrics for database operations
- **Type Safety and Error Handling**: DBSource methods return structured dataclasses and raise domain-specific exceptions, preventing database errors from leaking to upper layers
- **Improved Maintainability**: Database access patterns are centralized and documented in one place per domain, making the codebase easier to understand and modify

Currently, **21 out of 30 repository domains** have adopted the DB Source pattern (70% adoption rate). The remaining **9 domains** still access the database directly from Repository classes or Service layers, leading to:

- Mixed transaction management patterns across the codebase
- Database errors exposed to Service layer
- Difficulty in writing comprehensive repository tests
- Inconsistent metrics collection for database operations
- Higher coupling between business logic and data access

## Objective

Introduce the **DB Source pattern** to all remaining repository domains in Backend.AI Manager, achieving 100% adoption of the source-based architecture pattern.

**Primary Goals:**

1. **Create DBSource classes** for each domain without DB Source implementation
1. **Migrate database access logic** from Repository/Service layers into DBSource methods
1. **Refactor repository tests** to directly communicate with real databases using dedicated test fixtures
1. **Move database error handling** from Service/upper layers into DBSource methods
1. **Apply resilience policies** to DBSource methods for consistent metrics collection and retry behavior

**Success Criteria:**

- All repository domains follow the source-based architecture pattern
- Repository tests use real database fixtures instead of mocks
- Database errors are properly encapsulated within DBSource layer
- Prometheus metrics are collected for all database operations via MetricPolicy
- Transaction boundaries are clearly defined with explicit session management

## Architecture Pattern

### Current Pattern (21 domains already implemented)

```
Services Layer

Repository (orchestrates sources)

DBSource (db_source/) ← Dedicated database access layer

Database Models & PostgreSQL
```

**Example Structure:**

```
repositories/agent/
├── db_source/ # Database access layer
│ ├── __init__.py
│ └── db_source.py # AgentDBSource class
├── cache_source/ # Optional: Redis cache access
├── repository.py # AgentRepository (orchestrates sources)
└── repositories.py # Legacy compatibility wrapper
```

### Target Pattern for Remaining Domains

**DBSource Responsibilities:**

1. **Transaction Management**:
- Public methods create their own database sessions using `begin_readonly_session()` or `begin_session()`
- Private methods receive `db_sess` as a parameter to maintain the same session
- Clear read/write separation
1. **Type Safety**:
- Return structured dataclasses from public methods (not ORM Row objects)
- Private methods can return Row objects for internal reuse
- Raise domain-specific exceptions (e.g., `SessionNotFound`) instead of returning `None`
1. **Resilience Application**:
- Apply `@resilience.apply()` decorator to public methods
- Configure MetricPolicy for automatic metrics collection
- Configure RetryPolicy for transient error handling

**Repository Responsibilities:**

- Orchestrate multiple sources (db_source, cache_source, stateful_source)
- Implement cache-first patterns when cache_source exists
- Delegate database operations to db_source
- Handle cache invalidation strategies

**Example Implementation Pattern:**

```python
# repositories/domain/db_source/db_source.py
from ai.backend.common.resilience import Resilience
from ai.backend.common.resilience.policies.metrics import MetricPolicy, MetricArgs
from ai.backend.common.resilience.policies.retry import RetryPolicy, RetryArgs

domain_db_source_resilience = Resilience(
policies=[
MetricPolicy(
MetricArgs(
domain=DomainType.REPOSITORY,
layer=LayerType.DOMAIN_DB_SOURCE,
)
),
RetryPolicy(
RetryArgs(
max_retries=10,
retry_delay=0.1,
backoff_strategy=BackoffStrategy.FIXED,
non_retryable_exceptions=(BackendAIError,),
)
),
]
)

class DomainDBSource:
_db: ExtendedAsyncSAEngine

def __init__(self, db: ExtendedAsyncSAEngine) -> None:
self._db = db

@domain_db_source_resilience.apply()
async def get_item_by_id(
self,
item_id: ItemId,
) -> ItemData:
"""Public method: creates its own session"""
async with self._db.begin_readonly_session() as db_sess:
item_row = await self._fetch_item(db_sess, item_id)
if not item_row:
raise ItemNotFound(item_id)
return self._to_item_data(item_row)

async def _fetch_item(
self,
db_sess: SASession,
item_id: ItemId,
) -> Optional[ItemRow]:
"""Private method: receives session as parameter"""
stmt = sa.select(ItemRow).where(ItemRow.id == item_id)
return await db_sess.scalar(stmt)

# repositories/domain/repository.py
class DomainRepository:
_db_source: DomainDBSource

def __init__(self, db_source: DomainDBSource) -> None:
self._db_source = db_source

async def get_item(self, item_id: ItemId) -> ItemData:
"""Repository delegates to db_source"""
return await self._db_source.get_item_by_id(item_id)
```

## Expected Sub Issues

### Domains Requiring DB Source Implementation (9 domains)

**Priority Order** (based on usage frequency and architectural importance):

1. **session** - Core domain for compute session management
- High usage frequency
- Complex queries and transaction patterns
- Critical for scheduling and resource allocation
1. **user** - User account and credential management
- High usage frequency
- Authentication and authorization flows
- Shared across many services
1. **vfolder** - Virtual folder and persistent storage
- High usage frequency
- File operations and permission management
- Integration with storage backends
1. **group** - Group and domain organization
- Medium usage frequency
- RBAC and multi-tenancy support
- Relationships with users and resources
1. **domain** - Domain-level configuration and isolation
- Medium usage frequency
- Multi-tenancy boundary management
- Resource quota enforcement
1. **container_registry** - Container image registry management
- Medium usage frequency
- Image scanning and synchronization
- Version control and metadata
1. **project_resource_policy** - Project-level resource policies
- Medium usage frequency
- Resource quota management
- Policy inheritance patterns
1. **model_serving** - Model serving infrastructure management
- Lower usage frequency (newer feature)
- ML model deployment and lifecycle
- Endpoint management
1. **metric** - Container metrics data access
- Special case: Prometheus metric querying
- Time-series data retrieval
- Already has dedicated [documentation\|src/ai/backend/manager/repositories/metric/README.md]

### Implementation Tasks Per Domain

For each domain above, the following tasks need to be completed:

1. **Create DB Source Structure**
- Create `db_source/` subdirectory under the repository domain
- Implement `DBSource` class with resilience policies
- Define domain-specific exceptions
1. **Migrate Database Logic**
- Move all database queries from Repository to DBSource
- Move database queries from Service layer to DBSource (if any)
- Implement proper transaction boundaries (public vs private methods)
- Convert methods to use structured return types (dataclasses)
1. **Update Repository Class**
- Inject DBSource as a dependency
- Delegate database operations to DBSource
- Maintain cache orchestration logic (if cache_source exists)
1. **Refactor Tests**
- Create test fixtures that use real database connections
- Replace mock-based tests with database integration tests
- Add test cases for error scenarios (not found, duplicates, etc.)
- Verify transaction behavior and rollback scenarios
1. **Error Handling Migration**
- Move database exception handling into DBSource
- Raise domain-specific exceptions instead of returning None
- Update Service layer to handle domain exceptions
1. **Documentation**
- Document DBSource public methods with clear docstrings
- Add examples of common query patterns
- Update domain-specific README if it exists

## Impact

### Benefits

**Architecture:**

- Unified data access pattern across all repository domains
- Clear separation of concerns between orchestration and data access
- Explicit transaction boundaries and session management
- Easier to understand and maintain codebase structure

**Observability:**

- Comprehensive Prometheus metrics for all database operations via MetricPolicy
- Consistent metric labels: `domain`, `layer`, `operation`, `success`
- Automatic tracking of operation duration, retry counts, and error rates
- Better visibility into database performance and bottlenecks

**Testing:**

- Repository tests directly verify database behavior with real PostgreSQL
- Higher confidence in data access correctness
- Earlier detection of database-specific issues (constraints, transactions, etc.)
- Reduced reliance on mocks and test doubles

**Resilience:**

- Automatic retry handling for transient database errors
- Configurable retry policies per domain
- Proper handling of non-retryable business logic errors
- Improved system stability under database load

**Type Safety:**

- Stronger type guarantees with structured return types
- Prevention of ORM object leakage to upper layers
- Better IDE support and refactoring capabilities
- Compile-time detection of data access issues

### Migration Strategy

**Recommended Approach:**

1. **Incremental Migration**: Implement DB Source for one domain at a time
1. **Test-Driven**: Write comprehensive repository tests before refactoring
1. **Backward Compatibility**: Maintain existing Repository public APIs during migration
1. **Validation**: Run full integration tests after each domain migration
1. **Documentation**: Update architecture docs as pattern becomes standardized

**Risk Mitigation:**

- Create sub-tasks for each domain to track progress independently
- Extensive testing at each step to prevent regressions
- Code reviews focusing on transaction boundaries and error handling
- Gradual rollout allows early detection of issues

## References

- [Repository Layer Architecture\|src/ai/backend/manager/repositories/README.md]
- [Manager Architecture Overview\|src/ai/backend/manager/README.md]
- [Agent Repository Example\|src/ai/backend/manager/repositories/agent/] - Reference implementation with db_source
- [Scheduler Repository Example\|src/ai/backend/manager/repositories/scheduler/] - Reference implementation with db_source
- [Metric Repository Documentation\|src/ai/backend/manager/repositories/metric/README.md] - Special case for Prometheus queries

JIRA Issue: BA-3467

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.