use super::{
BaseMethods, ConnectionManager, ConnectionOptions, Data, Defaults, Group, Groups, Host, Hosts,
ResolvedConnectionParams, TransformFunction, TransformFunctionOptions,
};
use crate::{CustomTreeMap, NatString, State};
use dashmap::DashMap;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
pub struct Inventory {
pub(crate) hosts: Hosts,
pub(crate) groups: Option<Groups>,
pub(crate) defaults: Option<Defaults>,
#[serde(skip)]
transform_function: Option<TransformFunction>,
transform_function_options: Option<TransformFunctionOptions>,
#[serde(skip)]
#[schemars(skip)]
connections: Arc<ConnectionManager>,
#[serde(skip)]
#[schemars(skip)]
host_cache: DashMap<NatString, Host>,
#[serde(skip)]
#[schemars(skip)]
group_cache: DashMap<NatString, Group>,
#[serde(skip)]
#[schemars(skip)]
resolved_host_cache: DashMap<NatString, Host>,
#[serde(skip)]
#[schemars(skip)]
resolved_params_cache: DashMap<(NatString, String), ResolvedConnectionParams>,
#[serde(skip)]
#[schemars(skip)]
state: Arc<State>,
}
impl BaseMethods for Inventory {}
impl Inventory {
pub fn builder() -> InventoryBuilder {
InventoryBuilder::new()
}
pub fn hosts(&self) -> HostsView<'_> {
HostsView { inventory: self }
}
pub fn state(&self) -> &State {
self.state.as_ref()
}
pub fn hosts_raw(&self) -> &Hosts {
&self.hosts
}
pub fn groups(&self) -> Option<GroupsView<'_>> {
self.groups.as_ref().map(|groups| GroupsView {
inventory: self,
groups,
})
}
pub fn groups_raw(&self) -> Option<&Groups> {
self.groups.as_ref()
}
pub fn defaults(&self) -> Option<Defaults> {
self.defaults
.as_ref()
.map(|defaults| self.transform_defaults_value(defaults))
}
pub fn defaults_raw(&self) -> Option<&Defaults> {
self.defaults.as_ref()
}
pub fn transform_function_options(&self) -> Option<&TransformFunctionOptions> {
self.transform_function_options.as_ref()
}
pub fn connections(&self) -> &ConnectionManager {
&self.connections
}
#[cfg(test)]
pub(crate) fn resolved_host_cache_len(&self) -> usize {
self.resolved_host_cache.len()
}
#[cfg(test)]
pub(crate) fn resolved_params_cache_len(&self) -> usize {
self.resolved_params_cache.len()
}
pub fn resolve_host(&self, name: &str) -> Option<Host> {
let key = NatString::new(name.to_string());
if let Some(entry) = self.resolved_host_cache.get(&key) {
return Some(entry.value().clone());
}
let host = self.hosts.get(name)?;
let mut resolved = Host::new();
if let Some(defaults) = self.defaults.as_ref() {
merge_defaults_into_host(&mut resolved, defaults);
}
let mut group_stack = std::collections::HashSet::new();
let mut group_cache = std::collections::HashMap::new();
if let Some(groups) = host.groups.as_ref() {
for group_name in groups.iter() {
if let Some(group) =
self.resolve_group_internal(group_name, &mut group_stack, &mut group_cache)
{
merge_group_into_host(&mut resolved, &group);
}
}
}
merge_host_into_host(&mut resolved, host);
let resolved = self.transform_host_value(&resolved);
self.resolved_host_cache.insert(key, resolved.clone());
Some(resolved)
}
pub fn resolve_connection_params(
&self,
name: &str,
connection_type: &str,
) -> Option<ResolvedConnectionParams> {
let key = (
NatString::new(name.to_string()),
connection_type.to_string(),
);
if let Some(entry) = self.resolved_params_cache.get(&key) {
return Some(entry.value().clone());
}
let host = self.resolve_host(name)?;
let resolved = host.resolve_connection_params(connection_type);
self.resolved_params_cache.insert(key, resolved.clone());
Some(resolved)
}
pub(crate) fn resolve_group_internal(
&self,
name: &str,
stack: &mut std::collections::HashSet<String>,
memo: &mut std::collections::HashMap<String, Group>,
) -> Option<Group> {
if let Some(cached) = memo.get(name) {
return Some(cached.clone());
}
if !stack.insert(name.to_string()) {
return None;
}
let group = self.groups.as_ref()?.get(name)?;
let mut resolved = empty_group();
if let Some(parent_groups) = group.groups.as_ref() {
for parent in parent_groups.iter() {
if let Some(parent_group) = self.resolve_group_internal(parent, stack, memo) {
merge_group_into_group(&mut resolved, &parent_group);
}
}
}
merge_group_into_group(&mut resolved, group);
stack.remove(name);
memo.insert(name.to_string(), resolved.clone());
Some(resolved)
}
fn transform_value<T: Clone>(
&self,
value: &T,
apply: impl FnOnce(&TransformFunction, &T, Option<&TransformFunctionOptions>) -> T,
) -> T {
match &self.transform_function {
Some(transform) => apply(transform, value, self.transform_function_options.as_ref()),
None => value.clone(),
}
}
fn transform_host_value(&self, host: &Host) -> Host {
self.transform_value(host, |transform, host, options| {
transform.transform_host(host, options)
})
}
fn transform_group_value(&self, group: &Group) -> Group {
self.transform_value(group, |transform, group, options| {
transform.transform_group(group, options)
})
}
fn transform_defaults_value(&self, defaults: &Defaults) -> Defaults {
self.transform_value(defaults, |transform, defaults, options| {
transform.transform_defaults(defaults, options)
})
}
fn cached_value<T: Clone>(
&self,
key: &NatString,
cache: &DashMap<NatString, T>,
raw: &T,
transform: impl FnOnce(&T) -> T,
) -> T {
if let Some(entry) = cache.get(key) {
return entry.value().clone();
}
let transformed = transform(raw);
cache.insert(key.clone(), transformed.clone());
transformed
}
fn cached_host_value(&self, key: &NatString, host: &Host) -> Host {
self.cached_value(key, &self.host_cache, host, |host| {
self.transform_host_value(host)
})
}
fn cached_group_value(&self, key: &NatString, group: &Group) -> Group {
self.cached_value(key, &self.group_cache, group, |group| {
self.transform_group_value(group)
})
}
fn host_in_scope(&self, name: &str) -> bool {
self.state.is_in_scope(name)
}
fn host_key_in_scope(&self, key: &NatString) -> bool {
self.state.is_in_scope_key(key)
}
}
fn empty_group() -> Group {
Group {
hostname: None,
port: None,
username: None,
password: None,
platform: None,
groups: None,
data: None,
connection_options: None,
}
}
trait OverlayFields {
fn hostname(&self) -> &Option<String>;
fn port(&self) -> &Option<u16>;
fn username(&self) -> &Option<String>;
fn password(&self) -> &Option<String>;
fn platform(&self) -> &Option<String>;
fn data(&self) -> &Option<Data>;
fn connection_options(&self) -> &Option<CustomTreeMap<ConnectionOptions>>;
}
impl OverlayFields for Defaults {
fn hostname(&self) -> &Option<String> {
&self.hostname
}
fn port(&self) -> &Option<u16> {
&self.port
}
fn username(&self) -> &Option<String> {
&self.username
}
fn password(&self) -> &Option<String> {
&self.password
}
fn platform(&self) -> &Option<String> {
&self.platform
}
fn data(&self) -> &Option<Data> {
&self.data
}
fn connection_options(&self) -> &Option<CustomTreeMap<ConnectionOptions>> {
&self.connection_options
}
}
impl OverlayFields for Group {
fn hostname(&self) -> &Option<String> {
&self.hostname
}
fn port(&self) -> &Option<u16> {
&self.port
}
fn username(&self) -> &Option<String> {
&self.username
}
fn password(&self) -> &Option<String> {
&self.password
}
fn platform(&self) -> &Option<String> {
&self.platform
}
fn data(&self) -> &Option<Data> {
&self.data
}
fn connection_options(&self) -> &Option<CustomTreeMap<ConnectionOptions>> {
&self.connection_options
}
}
impl OverlayFields for Host {
fn hostname(&self) -> &Option<String> {
&self.hostname
}
fn port(&self) -> &Option<u16> {
&self.port
}
fn username(&self) -> &Option<String> {
&self.username
}
fn password(&self) -> &Option<String> {
&self.password
}
fn platform(&self) -> &Option<String> {
&self.platform
}
fn data(&self) -> &Option<Data> {
&self.data
}
fn connection_options(&self) -> &Option<CustomTreeMap<ConnectionOptions>> {
&self.connection_options
}
}
fn merge_overlay_into_host<T: OverlayFields>(target: &mut Host, source: &T) {
merge_option(&mut target.hostname, source.hostname());
merge_option(&mut target.port, source.port());
merge_option(&mut target.username, source.username());
merge_option(&mut target.password, source.password());
merge_option(&mut target.platform, source.platform());
merge_data(&mut target.data, source.data());
merge_connection_options(&mut target.connection_options, source.connection_options());
}
fn merge_overlay_into_group<T: OverlayFields>(target: &mut Group, source: &T) {
merge_option(&mut target.hostname, source.hostname());
merge_option(&mut target.port, source.port());
merge_option(&mut target.username, source.username());
merge_option(&mut target.password, source.password());
merge_option(&mut target.platform, source.platform());
merge_data(&mut target.data, source.data());
merge_connection_options(&mut target.connection_options, source.connection_options());
}
fn merge_defaults_into_host(target: &mut Host, defaults: &Defaults) {
merge_overlay_into_host(target, defaults);
}
fn merge_group_into_host(target: &mut Host, group: &Group) {
merge_overlay_into_host(target, group);
}
fn merge_host_into_host(target: &mut Host, host: &Host) {
merge_overlay_into_host(target, host);
if host.groups.is_some() {
target.groups = host.groups.clone();
}
}
fn merge_group_into_group(target: &mut Group, group: &Group) {
merge_overlay_into_group(target, group);
if group.groups.is_some() {
target.groups = group.groups.clone();
}
}
fn merge_option<T: Clone>(target: &mut Option<T>, source: &Option<T>) {
if let Some(value) = source.as_ref() {
*target = Some(value.clone());
}
}
pub(crate) fn merge_data(target: &mut Option<Data>, source: &Option<Data>) {
match (target.as_mut(), source.as_ref()) {
(Some(target_data), Some(source_data)) => {
if let (Some(target_obj), Some(source_obj)) =
(target_data.as_object_mut(), source_data.as_object())
{
for (key, value) in source_obj {
target_obj.insert(key.clone(), value.clone());
}
} else {
*target = Some(source_data.clone());
}
}
(None, Some(source_data)) => {
*target = Some(source_data.clone());
}
_ => {}
}
}
pub(crate) fn merge_connection_options(
target: &mut Option<CustomTreeMap<ConnectionOptions>>,
source: &Option<CustomTreeMap<ConnectionOptions>>,
) {
let Some(source_map) = source.as_ref() else {
return;
};
if target.is_none() {
*target = Some(CustomTreeMap::new());
}
let target_map = target.as_mut().expect("target map initialized");
for (name, options) in source_map.iter() {
if let Some(existing) = target_map.get_mut(name.as_str()) {
merge_connection_options_fields(existing, options);
} else {
target_map.insert(name.as_str(), options.clone());
}
}
}
pub(crate) fn merge_connection_options_fields(
target: &mut ConnectionOptions,
source: &ConnectionOptions,
) {
if source.hostname.is_some() {
target.hostname = source.hostname.clone();
}
if source.port.is_some() {
target.port = source.port;
}
if source.username.is_some() {
target.username = source.username.clone();
}
if source.password.is_some() {
target.password = source.password.clone();
}
if source.platform.is_some() {
target.platform = source.platform.clone();
}
if source.extras.is_some() {
target.extras = source.extras.clone();
}
}
pub struct HostsView<'a> {
inventory: &'a Inventory,
}
impl<'a> HostsView<'a> {
pub fn len(&self) -> usize {
self.inventory
.hosts
.keys()
.filter(|key| self.inventory.host_key_in_scope(key))
.count()
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn keys(&self) -> impl Iterator<Item = &'a NatString> {
self.inventory
.hosts
.keys()
.filter(|key| self.inventory.host_key_in_scope(key))
}
pub fn get(&self, name: &str) -> Option<Host> {
if !self.inventory.host_in_scope(name) {
return None;
}
let key = NatString::new(name.to_string());
self.inventory
.hosts
.get(name)
.map(|host| self.inventory.cached_host_value(&key, host))
}
pub fn iter(&self) -> impl Iterator<Item = (&'a NatString, Host)> {
self.inventory.hosts.iter().filter_map(|(id, host)| {
if self.inventory.host_key_in_scope(id) {
Some((id, self.inventory.cached_host_value(id, host)))
} else {
None
}
})
}
}
pub struct GroupsView<'a> {
inventory: &'a Inventory,
groups: &'a Groups,
}
impl<'a> GroupsView<'a> {
pub fn len(&self) -> usize {
self.groups.len()
}
pub fn is_empty(&self) -> bool {
self.groups.is_empty()
}
pub fn keys(&self) -> impl Iterator<Item = &'a NatString> {
self.groups.keys()
}
pub fn get(&self, name: &str) -> Option<Group> {
let key = NatString::new(name.to_string());
self.groups
.get(name)
.map(|group| self.inventory.cached_group_value(&key, group))
}
pub fn iter(&self) -> impl Iterator<Item = (&'a NatString, Group)> {
self.groups
.iter()
.map(|(id, group)| (id, self.inventory.cached_group_value(id, group)))
}
}
impl Default for Inventory {
fn default() -> Self {
Inventory {
hosts: Hosts::new(),
groups: None,
defaults: None,
transform_function: None,
transform_function_options: None,
connections: Arc::new(ConnectionManager::default()),
host_cache: DashMap::new(),
group_cache: DashMap::new(),
resolved_host_cache: DashMap::new(),
resolved_params_cache: DashMap::new(),
state: Arc::new(State::new()),
}
}
}
pub struct InventoryBuilder {
pub hosts: Option<Hosts>,
pub groups: Option<Groups>,
pub defaults: Option<Defaults>,
pub transform_function: Option<TransformFunction>,
pub transform_function_options: Option<TransformFunctionOptions>,
pub connections: Option<Arc<ConnectionManager>>,
}
impl InventoryBuilder {
pub fn new() -> InventoryBuilder {
InventoryBuilder {
hosts: None,
groups: None,
defaults: None,
transform_function: None,
transform_function_options: None,
connections: None,
}
}
pub fn hosts(mut self, hosts: Hosts) -> Self {
self.hosts = Some(hosts);
self
}
pub fn groups(mut self, groups: Groups) -> Self {
self.groups = Some(groups);
self
}
pub fn defaults(mut self, defaults: Defaults) -> Self {
self.defaults = Some(defaults);
self
}
pub fn transform_function(mut self, transform: TransformFunction) -> Self {
self.transform_function = Some(transform);
self
}
pub fn transform_function_options(mut self, options: TransformFunctionOptions) -> Self {
self.transform_function_options = Some(options);
self
}
pub fn connections(mut self, connections: ConnectionManager) -> Self {
self.connections = Some(Arc::new(connections));
self
}
pub fn build(self) -> Inventory {
let hosts = self.hosts.unwrap_or_default();
let state = State::new();
for key in hosts.keys() {
state.mark_in_scope_key(key);
}
Inventory {
hosts,
groups: self.groups,
defaults: self.defaults,
transform_function: self.transform_function,
transform_function_options: self.transform_function_options,
connections: self
.connections
.unwrap_or_else(|| Arc::new(ConnectionManager::default())),
host_cache: DashMap::new(),
group_cache: DashMap::new(),
resolved_host_cache: DashMap::new(),
resolved_params_cache: DashMap::new(),
state: Arc::new(state),
}
}
}
impl Default for InventoryBuilder {
fn default() -> Self {
Self::new()
}
}