use crate::call::AsyncCall;
use crate::canister::{Argument, CanisterBuilder};
use crate::Canister;
use async_trait::async_trait;
use candid::{CandidType, Deserialize};
use delay::Waiter;
use ic_agent::export::Principal;
use ic_agent::{Agent, AgentError, RequestId};
use std::fmt::Debug;
use std::str::FromStr;
pub mod attributes;
pub use attributes::ComputeAllocation;
pub use attributes::MemoryAllocation;
pub struct ManagementCanister;
impl ManagementCanister {
pub fn create(agent: &Agent) -> Canister<ManagementCanister> {
Canister::builder()
.with_agent(agent)
.with_canister_id(Principal::management_canister())
.with_interface(ManagementCanister)
.build()
.unwrap()
}
pub fn with_agent(agent: &Agent) -> CanisterBuilder<ManagementCanister> {
Canister::builder()
.with_agent(agent)
.with_canister_id(Principal::management_canister())
.with_interface(ManagementCanister)
}
}
#[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
#[serde(rename_all = "snake_case")]
pub enum CanisterStatus {
Running,
Stopping,
Stopped,
}
impl std::fmt::Display for CanisterStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
std::fmt::Debug::fmt(self, f)
}
}
#[derive(Copy, Clone, CandidType, Deserialize, Eq, PartialEq)]
pub enum InstallMode {
#[serde(rename = "install")]
Install,
#[serde(rename = "reinstall")]
Reinstall,
#[serde(rename = "upgrade")]
Upgrade,
}
impl FromStr for InstallMode {
type Err = String;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"install" => Ok(InstallMode::Install),
"reinstall" => Ok(InstallMode::Reinstall),
"upgrade" => Ok(InstallMode::Upgrade),
&_ => Err(format!("Invalid install mode: {}", s)),
}
}
}
pub struct InstallCodeBuilder<'agent, 'canister: 'agent, T> {
canister: &'canister Canister<'agent, T>,
canister_id: Principal,
wasm: &'canister [u8],
arg: Argument,
mode: Option<InstallMode>,
compute_allocation: Option<ComputeAllocation>,
memory_allocation: Option<MemoryAllocation>,
}
impl<'agent, 'canister: 'agent, T> InstallCodeBuilder<'agent, 'canister, T> {
pub fn builder(
canister: &'canister Canister<'agent, T>,
canister_id: &Principal,
wasm: &'canister [u8],
) -> Self {
Self {
canister,
canister_id: canister_id.clone(),
wasm,
arg: Default::default(),
mode: None,
compute_allocation: None,
memory_allocation: None,
}
}
pub fn with_arg<Argument: CandidType + Sync + Send>(
mut self,
arg: Argument,
) -> InstallCodeBuilder<'agent, 'canister, T> {
self.arg.push_idl_arg(arg);
self
}
pub fn with_raw_arg(mut self, arg: Vec<u8>) -> InstallCodeBuilder<'agent, 'canister, T> {
self.arg.set_raw_arg(arg);
self
}
pub fn with_mode(self, mode: InstallMode) -> Self {
Self {
mode: Some(mode),
..self
}
}
pub fn with_optional_compute_allocation<C: Into<ComputeAllocation>>(
self,
compute_allocation: Option<C>,
) -> Self {
Self {
compute_allocation: compute_allocation.map(|ca| ca.into()),
..self
}
}
pub fn with_compute_allocation<C: Into<ComputeAllocation>>(
self,
compute_allocation: C,
) -> Self {
self.with_optional_compute_allocation(Some(compute_allocation))
}
pub fn with_optional_memory_allocation<C: Into<MemoryAllocation>>(
self,
memory_allocation: Option<C>,
) -> Self {
Self {
memory_allocation: memory_allocation.map(|ma| ma.into()),
..self
}
}
pub fn with_memory_allocation<C: Into<MemoryAllocation>>(self, memory_allocation: C) -> Self {
Self {
memory_allocation: Some(memory_allocation.into()),
..self
}
}
pub fn build(self) -> Result<impl 'agent + AsyncCall<()>, AgentError> {
#[derive(candid::CandidType)]
struct CanisterInstall {
mode: InstallMode,
canister_id: Principal,
wasm_module: Vec<u8>,
arg: Vec<u8>,
compute_allocation: Option<u8>,
memory_allocation: Option<u64>,
}
Ok(self
.canister
.update_("install_code")
.with_arg(CanisterInstall {
mode: self.mode.unwrap_or(InstallMode::Install),
canister_id: self.canister_id.clone(),
wasm_module: self.wasm.to_owned(),
arg: self.arg.serialize()?,
compute_allocation: self.compute_allocation.map(|ca| ca.into()),
memory_allocation: self.memory_allocation.map(|ma| ma.into()),
})
.build())
}
pub async fn call(self) -> Result<RequestId, AgentError> {
self.build()?.call().await
}
pub async fn call_and_wait<W>(self, waiter: W) -> Result<(), AgentError>
where
W: Waiter,
{
self.build()?.call_and_wait(waiter).await
}
}
#[async_trait]
impl<'agent, 'canister: 'agent, T: Sync> AsyncCall<()>
for InstallCodeBuilder<'agent, 'canister, T>
{
async fn call(self) -> Result<RequestId, AgentError> {
self.build()?.call().await
}
async fn call_and_wait<W>(self, waiter: W) -> Result<(), AgentError>
where
W: Waiter,
{
self.build()?.call_and_wait(waiter).await
}
}
impl<'agent> Canister<'agent, ManagementCanister> {
pub fn canister_status<'canister: 'agent>(
&'canister self,
canister_id: &Principal,
) -> impl 'agent + AsyncCall<(CanisterStatus,)> {
#[derive(CandidType)]
struct In {
canister_id: Principal,
}
#[derive(Deserialize)]
struct Out {
status: CanisterStatus,
}
self.update_("canister_status")
.with_arg(In {
canister_id: canister_id.clone(),
})
.build()
.map(|result: (Out,)| (result.0.status,))
}
pub fn create_canister<'canister: 'agent>(
&'canister self,
) -> impl 'agent + AsyncCall<(Principal,)> {
#[derive(Deserialize)]
struct Out {
canister_id: Principal,
}
self.update_("create_canister")
.build()
.map(|result: (Out,)| (result.0.canister_id,))
}
pub fn delete_canister<'canister: 'agent>(
&'canister self,
canister_id: &Principal,
) -> impl 'agent + AsyncCall<()> {
#[derive(CandidType)]
struct Argument {
canister_id: Principal,
}
self.update_("delete_canister")
.with_arg(Argument {
canister_id: canister_id.clone(),
})
.build()
}
pub fn start_canister<'canister: 'agent>(
&'canister self,
canister_id: &Principal,
) -> impl 'agent + AsyncCall<()> {
#[derive(CandidType)]
struct Argument {
canister_id: Principal,
}
self.update_("start_canister")
.with_arg(Argument {
canister_id: canister_id.clone(),
})
.build()
}
pub fn stop_canister<'canister: 'agent>(
&'canister self,
canister_id: &Principal,
) -> impl 'agent + AsyncCall<()> {
#[derive(CandidType)]
struct Argument {
canister_id: Principal,
}
self.update_("stop_canister")
.with_arg(Argument {
canister_id: canister_id.clone(),
})
.build()
}
pub fn install_code<'canister: 'agent>(
&'canister self,
canister_id: &Principal,
wasm: &'canister [u8],
) -> InstallCodeBuilder<'agent, 'canister, ManagementCanister> {
InstallCodeBuilder::builder(self, canister_id, wasm)
}
}