diff --git a/alembic/versions/b6061da886ee_add_encrypted_columns.py b/alembic/versions/b6061da886ee_add_encrypted_columns.py new file mode 100644 index 00000000..a32c1d94 --- /dev/null +++ b/alembic/versions/b6061da886ee_add_encrypted_columns.py @@ -0,0 +1,39 @@ +"""add encrypted columns + +Revision ID: b6061da886ee +Revises: 89b595051e48 +Create Date: 2025-10-06 14:55:32.554544 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "b6061da886ee" +down_revision: Union[str, None] = "89b595051e48" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("agent_environment_variables", sa.Column("value_enc", sa.Text(), nullable=True)) + op.add_column("mcp_oauth", sa.Column("authorization_code_enc", sa.Text(), nullable=True)) + op.add_column("providers", sa.Column("api_key_enc", sa.Text(), nullable=True)) + op.add_column("providers", sa.Column("access_key_enc", sa.Text(), nullable=True)) + op.add_column("sandbox_environment_variables", sa.Column("value_enc", sa.Text(), nullable=True)) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("sandbox_environment_variables", "value_enc") + op.drop_column("providers", "access_key_enc") + op.drop_column("providers", "api_key_enc") + op.drop_column("mcp_oauth", "authorization_code_enc") + op.drop_column("agent_environment_variables", "value_enc") + # ### end Alembic commands ### diff --git a/letta/orm/mcp_oauth.py b/letta/orm/mcp_oauth.py index 163514b1..8cc4f645 100644 --- a/letta/orm/mcp_oauth.py +++ b/letta/orm/mcp_oauth.py @@ -35,6 +35,7 @@ class MCPOAuth(SqlalchemyBase, OrganizationMixin, UserMixin): # OAuth flow data authorization_url: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="OAuth authorization URL") authorization_code: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="OAuth authorization code") + authorization_code_enc: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="Encrypted OAuth authorization code") # Token data access_token: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="OAuth access token") diff --git a/letta/orm/provider.py b/letta/orm/provider.py index b46a95b8..c6a3cadc 100644 --- a/letta/orm/provider.py +++ b/letta/orm/provider.py @@ -1,6 +1,6 @@ -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional -from sqlalchemy import UniqueConstraint +from sqlalchemy import Text, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship from letta.orm.mixins import OrganizationMixin @@ -33,5 +33,9 @@ class Provider(SqlalchemyBase, OrganizationMixin): region: Mapped[str] = mapped_column(nullable=True, doc="Region used for requests to the provider.") api_version: Mapped[str] = mapped_column(nullable=True, doc="API version used for requests to the provider.") + # encrypted columns + api_key_enc: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="Encrypted API key or secret key for the provider.") + access_key_enc: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="Encrypted access key for the provider.") + # relationships organization: Mapped["Organization"] = relationship("Organization", back_populates="providers") diff --git a/letta/orm/sandbox_config.py b/letta/orm/sandbox_config.py index a3d22b18..a2ff2799 100644 --- a/letta/orm/sandbox_config.py +++ b/letta/orm/sandbox_config.py @@ -1,7 +1,7 @@ import uuid from typing import TYPE_CHECKING, Dict, List, Optional -from sqlalchemy import JSON, Enum as SqlEnum, Index, String, UniqueConstraint +from sqlalchemy import JSON, Enum as SqlEnum, Index, String, Text, UniqueConstraint from sqlalchemy.orm import Mapped, mapped_column, relationship from letta.orm.mixins import AgentMixin, OrganizationMixin, SandboxConfigMixin @@ -49,6 +49,9 @@ class SandboxEnvironmentVariable(SqlalchemyBase, OrganizationMixin, SandboxConfi value: Mapped[str] = mapped_column(String, nullable=False, doc="The value of the environment variable.") description: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="An optional description of the environment variable.") + # encrypted columns + value_enc: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="Encrypted value of the environment variable.") + # relationships organization: Mapped["Organization"] = relationship("Organization", back_populates="sandbox_environment_variables") sandbox_config: Mapped["SandboxConfig"] = relationship("SandboxConfig", back_populates="sandbox_environment_variables") @@ -71,5 +74,8 @@ class AgentEnvironmentVariable(SqlalchemyBase, OrganizationMixin, AgentMixin): value: Mapped[str] = mapped_column(String, nullable=False, doc="The value of the environment variable.") description: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="An optional description of the environment variable.") + # encrypted columns + value_enc: Mapped[Optional[str]] = mapped_column(Text, nullable=True, doc="Encrypted value of the environment variable.") + organization: Mapped["Organization"] = relationship("Organization", back_populates="agent_environment_variables") agent: Mapped[List["Agent"]] = relationship("Agent", back_populates="tool_exec_environment_variables")