crank/crates/crank-registry/src/postgres/api_key.rs

use super::*;

impl PostgresRegistry {
    pub async fn list_platform_api_keys_for_agent(
        &self,
        workspace_id: &WorkspaceId,
        agent_id: &AgentId,
    ) -> Result<Vec<PlatformApiKeyRecord>, RegistryError> {
        let rows = sqlx::query(
            "select
                id,
                workspace_id,
                agent_id,
                name,
                prefix,
                scopes_json,
                status,
                created_at,
                last_used_at
             from platform_api_keys
             where workspace_id = $1
               and agent_id = $2
             order by created_at desc",
        )
        .bind(workspace_id.as_str())
        .bind(agent_id.as_str())
        .fetch_all(&self.pool)
        .await?;

        rows.into_iter()
            .map(|row| {
                Ok(PlatformApiKeyRecord {
                    api_key: PlatformApiKey {
                        id: PlatformApiKeyId::new(row.get::<String, _>("id")),
                        workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
                        agent_id: row.get::<Option<String>, _>("agent_id").map(AgentId::new),
                        name: row.get("name"),
                        prefix: row.get("prefix"),
                        scopes: deserialize_json_value(row.get::<Value, _>("scopes_json"))?,
                        status: deserialize_enum_text(&row.get::<String, _>("status"), "status")?,
                        created_at: row.get("created_at"),
                        last_used_at: row.get("last_used_at"),
                    },
                })
            })
            .collect()
    }

    pub async fn list_platform_api_keys(
        &self,
        workspace_id: &WorkspaceId,
    ) -> Result<Vec<PlatformApiKeyRecord>, RegistryError> {
        let rows = sqlx::query(
            "select
                id,
                workspace_id,
                agent_id,
                name,
                prefix,
                scopes_json,
                status,
                created_at,
                last_used_at
             from platform_api_keys
             where workspace_id = $1
             order by created_at desc",
        )
        .bind(workspace_id.as_str())
        .fetch_all(&self.pool)
        .await?;

        rows.into_iter()
            .map(|row| {
                Ok(PlatformApiKeyRecord {
                    api_key: PlatformApiKey {
                        id: PlatformApiKeyId::new(row.get::<String, _>("id")),
                        workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
                        agent_id: row.get::<Option<String>, _>("agent_id").map(AgentId::new),
                        name: row.get("name"),
                        prefix: row.get("prefix"),
                        scopes: deserialize_json_value(row.get::<Value, _>("scopes_json"))?,
                        status: deserialize_enum_text(&row.get::<String, _>("status"), "status")?,
                        created_at: row.get("created_at"),
                        last_used_at: row.get("last_used_at"),
                    },
                })
            })
            .collect()
    }

    pub async fn get_platform_api_key_by_secret_for_agent_slug(
        &self,
        workspace_slug: &str,
        agent_slug: &str,
        secret_hash: &str,
    ) -> Result<Option<PlatformApiKeyRecord>, RegistryError> {
        let row = sqlx::query(
            "select
                k.id,
                k.workspace_id,
                k.agent_id,
                k.name,
                k.prefix,
                k.scopes_json,
                k.status,
                k.created_at,
                k.last_used_at
             from platform_api_keys k
             join workspaces w on w.id = k.workspace_id
             join agents a on a.id = k.agent_id
             where w.slug = $1
               and a.slug = $2
               and k.secret_hash = $3
               and k.status = 'active'
             limit 1",
        )
        .bind(workspace_slug)
        .bind(agent_slug)
        .bind(secret_hash)
        .fetch_optional(&self.pool)
        .await?;

        row.map(|row| {
            Ok(PlatformApiKeyRecord {
                api_key: PlatformApiKey {
                    id: PlatformApiKeyId::new(row.get::<String, _>("id")),
                    workspace_id: WorkspaceId::new(row.get::<String, _>("workspace_id")),
                    agent_id: row.get::<Option<String>, _>("agent_id").map(AgentId::new),
                    name: row.get("name"),
                    prefix: row.get("prefix"),
                    scopes: deserialize_json_value(row.get::<Value, _>("scopes_json"))?,
                    status: deserialize_enum_text(&row.get::<String, _>("status"), "status")?,
                    created_at: row.get("created_at"),
                    last_used_at: row.get("last_used_at"),
                },
            })
        })
        .transpose()
    }

    pub async fn create_platform_api_key(
        &self,
        request: CreatePlatformApiKeyRequest<'_>,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "insert into platform_api_keys (
                id,
                workspace_id,
                agent_id,
                name,
                prefix,
                secret_hash,
                scopes_json,
                status,
                created_at,
                last_used_at,
                revoked_at
             ) values (
                $1, $2, $3, $4, $5, $6, $7, $8, $9::timestamptz, $10::timestamptz, $11::timestamptz
             )",
        )
        .bind(request.api_key.id.as_str())
        .bind(request.api_key.workspace_id.as_str())
        .bind(request.api_key.agent_id.as_ref().map(AgentId::as_str))
        .bind(&request.api_key.name)
        .bind(&request.api_key.prefix)
        .bind(request.secret_hash)
        .bind(Json(serialize_json_value(&request.api_key.scopes)?))
        .bind(serialize_enum_text(&request.api_key.status, "status")?)
        .bind(request.api_key.created_at)
        .bind(request.api_key.last_used_at)
        .bind(Option::<&str>::None)
        .execute(&self.pool)
        .await;

        match result {
            Ok(_) => Ok(()),
            Err(sqlx::Error::Database(error))
                if error.constraint() == Some("platform_api_keys_workspace_name_idx") =>
            {
                Err(RegistryError::Storage(sqlx::Error::Database(error)))
            }
            Err(error) => Err(RegistryError::Storage(error)),
        }
    }

    pub async fn revoke_platform_api_key(
        &self,
        workspace_id: &WorkspaceId,
        key_id: &PlatformApiKeyId,
        revoked_at: &time::OffsetDateTime,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "update platform_api_keys
                set status = $1,
                    revoked_at = $2::timestamptz
              where workspace_id = $3 and id = $4",
        )
        .bind(serialize_enum_text(
            &PlatformApiKeyStatus::Revoked,
            "status",
        )?)
        .bind(revoked_at)
        .bind(workspace_id.as_str())
        .bind(key_id.as_str())
        .execute(&self.pool)
        .await?;

        if result.rows_affected() == 0 {
            return Err(RegistryError::PlatformApiKeyNotFound {
                key_id: key_id.as_str().to_owned(),
            });
        }

        Ok(())
    }

    pub async fn revoke_platform_api_key_for_agent(
        &self,
        workspace_id: &WorkspaceId,
        agent_id: &AgentId,
        key_id: &PlatformApiKeyId,
        revoked_at: &time::OffsetDateTime,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "update platform_api_keys
                set status = $1,
                    revoked_at = $2::timestamptz
              where workspace_id = $3 and agent_id = $4 and id = $5",
        )
        .bind(serialize_enum_text(
            &PlatformApiKeyStatus::Revoked,
            "status",
        )?)
        .bind(revoked_at)
        .bind(workspace_id.as_str())
        .bind(agent_id.as_str())
        .bind(key_id.as_str())
        .execute(&self.pool)
        .await?;

        if result.rows_affected() == 0 {
            return Err(RegistryError::PlatformApiKeyNotFound {
                key_id: key_id.as_str().to_owned(),
            });
        }

        Ok(())
    }

    pub async fn delete_platform_api_key(
        &self,
        workspace_id: &WorkspaceId,
        key_id: &PlatformApiKeyId,
    ) -> Result<(), RegistryError> {
        let result =
            sqlx::query("delete from platform_api_keys where workspace_id = $1 and id = $2")
                .bind(workspace_id.as_str())
                .bind(key_id.as_str())
                .execute(&self.pool)
                .await?;

        if result.rows_affected() == 0 {
            return Err(RegistryError::PlatformApiKeyNotFound {
                key_id: key_id.as_str().to_owned(),
            });
        }

        Ok(())
    }

    pub async fn delete_platform_api_key_for_agent(
        &self,
        workspace_id: &WorkspaceId,
        agent_id: &AgentId,
        key_id: &PlatformApiKeyId,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "delete from platform_api_keys where workspace_id = $1 and agent_id = $2 and id = $3",
        )
        .bind(workspace_id.as_str())
        .bind(agent_id.as_str())
        .bind(key_id.as_str())
        .execute(&self.pool)
        .await?;

        if result.rows_affected() == 0 {
            return Err(RegistryError::PlatformApiKeyNotFound {
                key_id: key_id.as_str().to_owned(),
            });
        }

        Ok(())
    }

    pub async fn touch_platform_api_key(
        &self,
        workspace_id: &WorkspaceId,
        key_id: &PlatformApiKeyId,
        used_at: &time::OffsetDateTime,
    ) -> Result<(), RegistryError> {
        let result = sqlx::query(
            "update platform_api_keys
                set last_used_at = $3::timestamptz
              where workspace_id = $1 and id = $2",
        )
        .bind(workspace_id.as_str())
        .bind(key_id.as_str())
        .bind(used_at)
        .execute(&self.pool)
        .await?;

        if result.rows_affected() == 0 {
            return Err(RegistryError::PlatformApiKeyNotFound {
                key_id: key_id.as_str().to_owned(),
            });
        }

        Ok(())
    }
}