DB-GPT/dbgpt/serve/flow/models/variables_adapter.py

from typing import Type

from sqlalchemy.orm import Session

from dbgpt.core.interface.storage import StorageItemAdapter
from dbgpt.core.interface.variables import StorageVariables, VariablesIdentifier

from .models import VariablesEntity


class VariablesAdapter(StorageItemAdapter[StorageVariables, VariablesEntity]):
    """Variables adapter.

    Convert between storage format and database model.
    """

    def to_storage_format(self, item: StorageVariables) -> VariablesEntity:
        """Convert to storage format."""
        return VariablesEntity(
            key=item.key,
            name=item.name,
            label=item.label,
            value=item.value,
            value_type=item.value_type,
            category=item.category,
            encryption_method=item.encryption_method,
            salt=item.salt,
            scope=item.scope,
            scope_key=item.scope_key,
            sys_code=item.sys_code,
            user_name=item.user_name,
            description=item.description,
        )

    def from_storage_format(self, model: VariablesEntity) -> StorageVariables:
        """Convert from storage format."""
        return StorageVariables(
            key=model.key,
            name=model.name,
            label=model.label,
            value=model.value,
            value_type=model.value_type,
            category=model.category,
            encryption_method=model.encryption_method,
            salt=model.salt,
            scope=model.scope,
            scope_key=model.scope_key,
            sys_code=model.sys_code,
            user_name=model.user_name,
            description=model.description,
        )

    def get_query_for_identifier(
        self,
        storage_format: Type[VariablesEntity],
        resource_id: VariablesIdentifier,
        **kwargs,
    ):
        """Get query for identifier."""
        session: Session = kwargs.get("session")
        if session is None:
            raise Exception("session is None")
        query_obj = session.query(VariablesEntity)
        for key, value in resource_id.to_dict().items():
            if value is None:
                continue
            query_obj = query_obj.filter(getattr(VariablesEntity, key) == value)

        # enabled must be True
        query_obj = query_obj.filter(VariablesEntity.enabled == 1)
        return query_obj