obelisk 0.35.3

Deterministic workflow engine
use crate::args::Generate;
use crate::args::shadow::PKG_VERSION;
use crate::command::server::{VerifyParams, verify_config_compile_link};
use crate::command::termination_notifier::termination_notifier;
use crate::config::config_holder::ConfigHolder;
use crate::config::config_holder::ConfigSource;
use crate::init::{self};
use crate::project_dirs;
use crate::wit_printer::{OutputToFile, process_pkg_with_deps};
use anyhow::Context;
use concepts::{ComponentType, ExecutionId, prefixed_ulid::DeploymentId};
use directories::{BaseDirs, ProjectDirs};
use hashbrown::HashSet;
use std::sync::Arc;
use std::{borrow::Cow, path::PathBuf};
use tokio::sync::watch;
use utils::{wasm_tools::WasmComponent, wit};

impl Generate {
    pub(crate) async fn run(self) -> Result<(), anyhow::Error> {
        match self {
            #[cfg(debug_assertions)]
            Generate::ConfigSchema { output } => generate_toml_schema(output),
            Generate::Config { config, overwrite } => {
                let config_file = ConfigHolder::generate_default_config(config, overwrite).await?;
                println!("Generated {config_file:?}");
                Ok(())
            }

            Generate::WitExtensions {
                component_type,
                input_wit_directory,
                output_directory,
                force,
            } => {
                generate_exported_extension_wits(
                    input_wit_directory,
                    output_directory,
                    component_type,
                    force,
                )
                .await
            }
            Generate::WitSupport {
                component_type,
                output_directory,
            } => generate_support_wits(component_type, output_directory).await,
            Generate::WitDeps {
                config,
                output_directory,
                overwrite,
            } => {
                generate_wit_deps(
                    project_dirs(),
                    BaseDirs::new(),
                    config,
                    output_directory,
                    overwrite,
                )
                .await
            }
            Generate::ExecutionId => {
                println!("{}", ExecutionId::generate());
                Ok(())
            }
        }
    }
}

#[cfg(debug_assertions)]
pub(crate) fn generate_toml_schema(output: Option<PathBuf>) -> Result<(), anyhow::Error> {
    use std::{
        fs::File,
        io::{BufWriter, Write as _, stdout},
    };
    let schema = schemars::schema_for!(crate::config::toml::ConfigToml);
    if let Some(output) = output {
        // Save to a file
        let mut writer = BufWriter::new(File::create(&output)?);
        serde_json::to_writer_pretty(&mut writer, &schema)?;
        writer.write_all(b"\n")?;
        writer.flush()?; // Do not swallow errors
    } else {
        serde_json::to_writer_pretty(stdout().lock(), &schema)?;
    }
    Ok(())
}

const HEADER: &str = "// Generated by Obelisk";

pub(crate) async fn generate_exported_extension_wits(
    input_wit_directory: PathBuf,
    output_directory: PathBuf,
    component_type: ComponentType,
    force: bool,
) -> Result<(), anyhow::Error> {
    let wasm_component = WasmComponent::new_from_wit_folder(&input_wit_directory, component_type)?;
    let pkgs_to_wits = wasm_component.exported_extension_wits()?;
    for (pkg_fqn, new_content) in pkgs_to_wits {
        let pkg_folder = output_directory.join(pkg_fqn.to_string().replace(':', "_"));
        let wit_file = pkg_folder.join(format!("{}.wit", pkg_fqn.package_name));

        let old_content = tokio::fs::read_to_string(&wit_file)
            .await
            .unwrap_or_default();

        let old_content = if force {
            None
        } else {
            Some(strip_header(&old_content))
        };
        if old_content.as_ref() != Some(&new_content) {
            let new_content = format!("{HEADER} {PKG_VERSION}\n{new_content}");
            tokio::fs::create_dir_all(&pkg_folder)
                .await
                .with_context(|| format!("cannot write {pkg_folder:?}"))?;
            tokio::fs::write(&wit_file, new_content.as_bytes())
                .await
                .with_context(|| format!("cannot write {wit_file:?}"))?;
            println!("{wit_file:?} created or updated");
        } else {
            println!("{wit_file:?} is up to date");
        }
    }
    Ok(())
}

fn strip_header(old_content: &str) -> String {
    let old_content = match old_content.strip_prefix(HEADER) {
        Some(wit) => {
            if let Some((_, wit)) = wit.split_once('\n') {
                Cow::Borrowed(wit)
            } else {
                Cow::Borrowed(wit)
            }
        }
        None => Cow::Borrowed(old_content),
    };
    let old_content = match old_content.strip_prefix(&format!("/{HEADER}")) {
        // Bug in wasm_tools is turning // into ///
        Some(wit) => {
            if let Some((_, wit)) = wit.split_once('\n') {
                Cow::Borrowed(wit)
            } else {
                Cow::Borrowed(wit)
            }
        }
        None => old_content,
    };
    old_content.into_owned()
}

pub(crate) async fn generate_support_wits(
    component_type: ComponentType,
    output_directory: PathBuf,
) -> Result<(), anyhow::Error> {
    let files = match component_type {
        ComponentType::ActivityWasm => {
            vec![
                wit::WIT_OBELISK_ACTIVITY_PACKAGE_PROCESS,
                wit::WIT_OBELISK_LOG_PACKAGE,
            ]
        }
        ComponentType::ActivityStub | ComponentType::ActivityExternal => vec![],
        ComponentType::Workflow => vec![
            wit::WIT_OBELISK_TYPES_PACKAGE,
            wit::WIT_OBELISK_WORKFLOW_PACKAGE,
            wit::WIT_OBELISK_LOG_PACKAGE,
        ],
        ComponentType::WebhookEndpoint => {
            vec![
                wit::WIT_OBELISK_TYPES_PACKAGE, // Needed for -schedule ext functions.
                wit::WIT_OBELISK_LOG_PACKAGE,
            ]
        }
    };
    for [folder, filename, content] in files {
        let output_directory = output_directory.join(folder);
        let path = output_directory.join(filename);
        if let Ok(actual) = tokio::fs::read_to_string(&path).await
            && actual == content
        {
            println!("{path:?} is up to date");
        } else {
            tokio::fs::create_dir_all(&output_directory)
                .await
                .with_context(|| format!("cannot write {output_directory:?}"))?;
            tokio::fs::write(&path, content)
                .await
                .with_context(|| format!("cannot write {path:?}"))?;
            println!("{path:?} created or updated");
        }
    }
    Ok(())
}

pub(crate) async fn generate_wit_deps(
    project_dirs: Option<ProjectDirs>,
    base_dirs: Option<BaseDirs>,
    config: Option<ConfigSource>,
    output_directory: PathBuf,
    overwrite: bool,
) -> Result<(), anyhow::Error> {
    let config_holder = ConfigHolder::new(project_dirs, base_dirs, config, true)?;
    let mut config = config_holder.load_config().await?;
    let _guard = init::init(&mut config)?;
    let (termination_sender, mut termination_watcher) = watch::channel(());
    tokio::spawn(async move { termination_notifier(termination_sender).await });
    let (compiled_and_linked, _component_source_map) = Box::pin(verify_config_compile_link(
        config,
        Arc::new(config_holder.path_prefixes),
        DeploymentId::generate(),
        VerifyParams {
            ignore_missing_env_vars: true,
            clean_cache: false,
            clean_codegen_cache: false,
            suppress_type_checking_errors: true, // Just extracting WITs, not running components
        },
        &mut termination_watcher,
    ))
    .await?;
    tokio::fs::create_dir_all(&output_directory)
        .await
        .with_context(|| format!("cannot create the output directory {output_directory:?}"))?;
    let include_exports = true;
    let mut output = OutputToFile::default();
    for component in &compiled_and_linked
        .component_registry_ro
        .list(include_exports)
    {
        if let Some(wit) = &component.wit
            && let Some(importable) = &component.workflow_or_activity_config
        {
            assert!(
                component.component_id.component_type.is_activity()
                    || component.component_id.component_type == ComponentType::Workflow
            );
            let mut already_processed_packages = HashSet::new();

            let packages: HashSet<_> = importable
                .exports_hierarchy_ext
                .iter()
                .map(|ifc_fqns| ifc_fqns.ifc_fqn.pkg_fqn_name())
                .collect();

            for pkg in packages {
                output = process_pkg_with_deps(wit, &pkg, &mut already_processed_packages, output)
                    .await?;
            }
        }
    }
    output.write(&output_directory, overwrite).await?;
    Ok(())
}