horizon-sdk 6.7.2

//! Horizon SDK: canonical Rust data access layer for the Horizon platform.
//!
//! # Quick Start
//!
//! ```rust,no_run
//! use horizon_sdk::HorizonSdk;
//!
//! # async fn example() -> Result<(), horizon_sdk::types::error::HorizonError> {
//! let sdk = HorizonSdk::builder()
//!     .with_postgres_url("postgresql://user:pass@localhost/horizon")
//!     .build()
//!     .await?;
//! # Ok(())
//! # }
//! ```
//!
//! # Offline builds
//!
//! The SDK uses `sqlx` compile-time-checked queries. The published crate bundles
//! a `.sqlx` query cache so it compiles without a database: when `DATABASE_URL`
//! is unset, `sqlx` uses the bundled cache. Set `SQLX_OFFLINE=true` to force
//! offline compilation even when `DATABASE_URL` is present.

#![allow(
    clippy::multiple_inherent_impl,
    reason = "builder methods are grouped by subsystem for readability"
)]

pub mod bench_support;
pub mod coverage;
pub mod iceberg;
pub mod kafka;
pub mod postgres;
pub mod types;

use std::fmt;

use tracing::info;

use crate::coverage::client::CoverageClient;
use crate::iceberg::discovery::discover_table_metadata;
use crate::iceberg::partition::PartitionValue;
use crate::iceberg::repository::IcebergRepository;
use crate::kafka::producer::KafkaProducer;
use crate::postgres::pool::create_pool;
use crate::postgres::repository::PostgresRepository;
use crate::types::config::{CoverageConfig, IcebergConfig, KafkaConfig, PostgresConfig, SdkConfig};
use crate::types::error::{ConfigError, HorizonError, Result};
use crate::types::model::{DataRow, MetadataRow};

/// Top-level SDK facade composing all Horizon subsystems.
///
/// Constructed via [`HorizonSdk::builder()`]. Postgres is required;
/// Kafka/Iceberg and Coverage are optional subsystems.
///
/// When Kafka and Iceberg are both configured, data row and metadata row
/// writes are dual-written: first to Postgres, then to Kafka (as Arrow IPC)
/// for eventual Iceberg ingestion via `horizon-iceberg-sink`.
pub struct HorizonSdk {
    /// Coverage HTTP client, when the service is configured.
    possible_coverage_client: Option<CoverageClient>,
    /// Iceberg repository for data row dual-write via Kafka.
    possible_data_row_iceberg: Option<IcebergRepository>,
    /// Iceberg repository for metadata row dual-write via Kafka.
    possible_metadata_row_iceberg: Option<IcebergRepository>,
    /// Required `PostgreSQL` repository for all CRUD operations.
    postgres_repo: PostgresRepository,
}

impl HorizonSdk {
    /// Start building a new SDK instance.
    #[must_use]
    pub fn builder() -> HorizonSdkBuilder {
        HorizonSdkBuilder::default()
    }

    /// Access the coverage client, if configured.
    #[must_use]
    pub const fn coverage(&self) -> Option<&CoverageClient> {
        self.possible_coverage_client.as_ref()
    }

    /// Insert a data row into Postgres. If Iceberg is configured, also
    /// produces the row as Arrow IPC to the data row topic.
    ///
    /// # Consistency
    ///
    /// The Postgres insert commits before the Iceberg publish is
    /// attempted. If the Iceberg leg fails, the Postgres row is NOT
    /// rolled back — the error returns as
    /// [`HorizonError::PartialDualWrite`] so callers can distinguish
    /// "nothing happened" (Postgres leg failed) from "half applied"
    /// (Postgres committed, Iceberg did not).
    ///
    /// # Errors
    ///
    /// Returns `PostgresError::Query` if the Postgres insert fails, or
    /// `HorizonError::PartialDualWrite` wrapping the Iceberg/Kafka
    /// failure if the dual-write second leg fails.
    pub async fn create_data_row(&self, row: &DataRow) -> Result<()> {
        self.postgres_repo.insert_data_row(row).await?;

        if let Some(iceberg_repo) = &self.possible_data_row_iceberg {
            iceberg_repo
                .insert(row, &data_row_partition_values(row))
                .await
                .map_err(|source| HorizonError::PartialDualWrite {
                    source: Box::new(source),
                })?;
        }

        Ok(())
    }

    /// Insert a batch of data rows into Postgres. If Iceberg is
    /// configured, also produces each row as Arrow IPC to the data row
    /// topic.
    ///
    /// # Consistency
    ///
    /// Same as [`create_data_row`](Self::create_data_row): Postgres
    /// commits first; if the Iceberg batch publish fails, a
    /// [`HorizonError::PartialDualWrite`] is returned and the Postgres
    /// rows are NOT rolled back.
    ///
    /// # Errors
    ///
    /// Returns `PostgresError::Query` if the Postgres insert fails, or
    /// `HorizonError::PartialDualWrite` wrapping the Iceberg/Kafka
    /// failure if the dual-write second leg fails.
    pub async fn create_data_row_batch(&self, rows: &[DataRow]) -> Result<()> {
        self.postgres_repo.insert_data_row_batch(rows).await?;

        if let Some(iceberg_repo) = &self.possible_data_row_iceberg {
            iceberg_repo
                .insert_batch(rows, data_row_partition_values)
                .await
                .map_err(|source| HorizonError::PartialDualWrite {
                    source: Box::new(source),
                })?;
        }

        Ok(())
    }

    /// Insert a batch of data rows into Iceberg ONLY, bypassing Postgres.
    ///
    /// Batch counterpart to
    /// [`create_data_row_iceberg_only`](Self::create_data_row_iceberg_only).
    /// Returns after every message is acknowledged.
    ///
    /// # Errors
    ///
    /// Returns [`ConfigError::MissingSubsystem`] if the Iceberg write path is
    /// not configured, or the underlying Iceberg/Kafka error if serialization
    /// or delivery of any partition's message fails.
    pub async fn create_data_row_batch_iceberg_only(&self, rows: &[DataRow]) -> Result<()> {
        let iceberg_repo =
            self.possible_data_row_iceberg
                .as_ref()
                .ok_or(ConfigError::MissingSubsystem {
                    subsystem: "Iceberg",
                    hint: "call with_iceberg() and with_kafka() on the builder",
                })?;

        iceberg_repo
            .insert_batch(rows, data_row_partition_values)
            .await
    }

    /// Insert a data row into Iceberg ONLY, bypassing Postgres entirely.
    ///
    /// Unlike [`create_data_row`](Self::create_data_row), this does not write
    /// to Postgres. The row is serialized to Arrow IPC and produced to Kafka
    /// for eventual ingestion by `horizon-iceberg-sink`. The call returns
    /// after Kafka acknowledgement (acks=all).
    ///
    /// Because Postgres — the system of record — is skipped, a missing Iceberg
    /// write path is treated as an error here rather than silently ignored
    /// (the dual-write methods skip Iceberg silently when it is unconfigured).
    ///
    /// # Errors
    ///
    /// Returns [`ConfigError::MissingSubsystem`] if the Iceberg write path is
    /// not configured (call [`with_iceberg`](HorizonSdkBuilder::with_iceberg)
    /// and [`with_kafka`](HorizonSdkBuilder::with_kafka) on the builder), or
    /// the underlying Iceberg/Kafka error if serialization or delivery fails.
    pub async fn create_data_row_iceberg_only(&self, row: &DataRow) -> Result<()> {
        let iceberg_repo =
            self.possible_data_row_iceberg
                .as_ref()
                .ok_or(ConfigError::MissingSubsystem {
                    subsystem: "Iceberg",
                    hint: "call with_iceberg() and with_kafka() on the builder",
                })?;

        iceberg_repo
            .insert(row, &data_row_partition_values(row))
            .await
    }

    /// Insert a metadata row into Postgres. If Iceberg is configured,
    /// also produces the row as Arrow IPC to the metadata row topic.
    ///
    /// # Consistency
    ///
    /// Same as [`create_data_row`](Self::create_data_row): Postgres
    /// commits first; Iceberg failure surfaces as
    /// [`HorizonError::PartialDualWrite`] without rolling back.
    ///
    /// # Errors
    ///
    /// Returns `PostgresError::Query` if the Postgres insert fails, or
    /// `HorizonError::PartialDualWrite` wrapping the Iceberg/Kafka
    /// failure if the dual-write second leg fails.
    pub async fn create_metadata_row(&self, row: &MetadataRow) -> Result<()> {
        self.postgres_repo.insert_metadata_row(row).await?;

        if let Some(iceberg_repo) = &self.possible_metadata_row_iceberg {
            iceberg_repo
                .insert(row, &metadata_row_partition_values(row))
                .await
                .map_err(|source| HorizonError::PartialDualWrite {
                    source: Box::new(source),
                })?;
        }

        Ok(())
    }

    /// Insert a batch of metadata rows into Postgres. If Iceberg is
    /// configured, also produces each row as Arrow IPC to the metadata
    /// row topic.
    ///
    /// # Consistency
    ///
    /// Same as [`create_data_row`](Self::create_data_row): Postgres
    /// commits first; Iceberg batch failure surfaces as
    /// [`HorizonError::PartialDualWrite`] without rolling back.
    ///
    /// # Errors
    ///
    /// Returns `PostgresError::Query` if the Postgres insert fails, or
    /// `HorizonError::PartialDualWrite` wrapping the Iceberg/Kafka
    /// failure if the dual-write second leg fails.
    pub async fn create_metadata_row_batch(&self, rows: &[MetadataRow]) -> Result<()> {
        self.postgres_repo.insert_metadata_row_batch(rows).await?;

        if let Some(iceberg_repo) = &self.possible_metadata_row_iceberg {
            iceberg_repo
                .insert_batch(rows, metadata_row_partition_values)
                .await
                .map_err(|source| HorizonError::PartialDualWrite {
                    source: Box::new(source),
                })?;
        }

        Ok(())
    }

    /// Whether the Iceberg write path is enabled.
    ///
    /// The Iceberg repository owns the Kafka producer internally, so this
    /// single predicate gates the Postgres-plus-Iceberg dual-write path —
    /// there is no Kafka-only mode to branch on separately.
    #[must_use]
    pub const fn has_iceberg(&self) -> bool {
        self.possible_data_row_iceberg.is_some()
    }

    /// Access the `PostgreSQL` repository for direct CRUD operations.
    #[must_use]
    pub const fn postgres(&self) -> &PostgresRepository {
        &self.postgres_repo
    }
}

impl fmt::Debug for HorizonSdk {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("HorizonSdk")
            .field("coverage_enabled", &self.possible_coverage_client.is_some())
            .field("iceberg_enabled", &self.has_iceberg())
            .field("postgres", &self.postgres_repo)
            .finish_non_exhaustive()
    }
}

/// Builder for constructing an [`HorizonSdk`] instance.
///
/// Postgres is required — call either [`with_postgres_config`](Self::with_postgres_config)
/// or [`with_postgres_url`](Self::with_postgres_url) before [`build`](Self::build).
/// All other subsystems are optional.
#[derive(Default)]
pub struct HorizonSdkBuilder {
    /// Coverage `BentoML` service configuration.
    coverage_config: Option<CoverageConfig>,
    /// Iceberg catalog configuration for schema discovery.
    iceberg_config: Option<IcebergConfig>,
    /// Kafka producer configuration for dual-write.
    kafka_config: Option<KafkaConfig>,
    /// Organization ID for multi-tenant row-level security.
    organization_id: Option<uuid::Uuid>,
    /// Structured `PostgreSQL` connection configuration.
    postgres_config: Option<PostgresConfig>,
    /// `PostgreSQL` connection URL (alternative to structured config).
    postgres_url: Option<String>,
}

impl HorizonSdkBuilder {
    /// Build the SDK, connecting to all configured subsystems.
    ///
    /// # Errors
    ///
    /// Returns `ConfigError::MissingSubsystem` if Postgres is not configured,
    /// or subsystem-specific errors if connections fail.
    pub async fn build(self) -> Result<HorizonSdk> {
        if self.iceberg_config.is_some() && self.kafka_config.is_none() {
            return Err(ConfigError::MissingSubsystem {
                subsystem: "Kafka",
                hint: "required when Iceberg is configured for the write path",
            }
            .into());
        }

        let pg_config = self.resolve_postgres_config()?;
        let pool = create_pool(&pg_config).await?;
        let org_id = self.organization_id.or(pg_config.possible_organization_id);
        let postgres_repo = PostgresRepository::new(pool, org_id);

        info!("initialized Postgres repository");

        let (possible_data_row_iceberg, possible_metadata_row_iceberg) =
            if let Some(iceberg_config) = &self.iceberg_config {
                let kafka_config = self.kafka_config.as_ref().ok_or({
                    ConfigError::MissingSubsystem {
                        subsystem: "Kafka",
                        hint: "required when Iceberg is configured for the write path",
                    }
                })?;

                let (dr_field_ids, dr_partition) =
                    discover_table_metadata(iceberg_config, &iceberg_config.data_row_table).await?;
                let dr_producer = KafkaProducer::new(kafka_config, &kafka_config.data_row_topic)?;
                let dr_repo = IcebergRepository::new(dr_producer, dr_field_ids, dr_partition);

                let (mr_field_ids, mr_partition) =
                    discover_table_metadata(iceberg_config, &iceberg_config.metadata_row_table)
                        .await?;
                let mr_producer =
                    KafkaProducer::new(kafka_config, &kafka_config.metadata_row_topic)?;
                let mr_repo = IcebergRepository::new(mr_producer, mr_field_ids, mr_partition);

                (Some(dr_repo), Some(mr_repo))
            } else {
                (None, None)
            };

        let possible_coverage_client = self
            .coverage_config
            .as_ref()
            .map(CoverageClient::new)
            .transpose()?;

        info!(
            dual_write = possible_data_row_iceberg.is_some(),
            coverage = possible_coverage_client.is_some(),
            "HorizonSdk initialized"
        );

        Ok(HorizonSdk {
            possible_coverage_client,
            possible_data_row_iceberg,
            possible_metadata_row_iceberg,
            postgres_repo,
        })
    }

    /// Resolve `PostgreSQL` configuration from structured config or URL.
    fn resolve_postgres_config(&self) -> Result<PostgresConfig> {
        if let Some(config) = &self.postgres_config {
            return Ok(config.clone());
        }

        if let Some(url) = &self.postgres_url {
            let parsed = url::Url::parse(url).map_err(|source| ConfigError::InvalidUrl {
                field: "Postgres",
                url: url.clone(),
                source,
            })?;

            let host = parsed.host_str().unwrap_or("localhost").to_owned();
            let port = parsed.port().unwrap_or(5432);
            let user = if parsed.username().is_empty() {
                "postgres".to_owned()
            } else {
                parsed.username().to_owned()
            };
            let password = parsed.password().unwrap_or_default().to_owned();
            let database = parsed.path().trim_start_matches('/').to_owned();
            let sslmode = parsed
                .query_pairs()
                .find(|(k, _)| k == "sslmode")
                .map_or_else(|| "require".to_owned(), |(_, v)| v.into_owned());

            return Ok(PostgresConfig {
                database,
                host,
                password,
                port,
                possible_organization_id: None,
                role: None,
                sslmode,
                user,
            });
        }

        Err(ConfigError::MissingSubsystem {
            subsystem: "Postgres",
            hint: "call with_postgres_config() or with_postgres_url()",
        }
        .into())
    }

    /// Load configuration from an [`SdkConfig`] struct.
    ///
    /// This sets all subsystems at once. Individual `with_*` calls
    /// after this will override the corresponding subsystem.
    #[must_use]
    pub fn with_config(mut self, config: SdkConfig) -> Self {
        self.coverage_config = config.coverage;
        self.iceberg_config = config.iceberg;
        self.kafka_config = config.kafka;
        self.postgres_config = Some(config.postgres);
        self
    }

    /// Set coverage service configuration.
    #[must_use]
    pub fn with_coverage(mut self, config: CoverageConfig) -> Self {
        self.coverage_config = Some(config);
        self
    }

    /// Set Iceberg configuration for schema discovery.
    #[must_use]
    pub fn with_iceberg(mut self, config: IcebergConfig) -> Self {
        self.iceberg_config = Some(config);
        self
    }

    /// Set Kafka configuration for dual-write.
    #[must_use]
    pub fn with_kafka(mut self, config: KafkaConfig) -> Self {
        self.kafka_config = Some(config);
        self
    }

    /// Set the organization ID for multi-tenant row-level security.
    #[must_use]
    pub const fn with_organization_id(mut self, id: uuid::Uuid) -> Self {
        self.organization_id = Some(id);
        self
    }

    /// Set Postgres configuration from a structured config.
    #[must_use]
    pub fn with_postgres_config(mut self, config: PostgresConfig) -> Self {
        self.postgres_config = Some(config);
        self
    }

    /// Set Postgres connection from a URL string.
    ///
    /// Format: `postgresql://user:password@host:port/database?sslmode=require`.
    #[must_use]
    pub fn with_postgres_url(mut self, url: &str) -> Self {
        self.postgres_url = Some(url.to_owned());
        self
    }
}

/// Load SDK configuration from layered sources.
///
/// Priority (highest to lowest):
/// 1. Environment variables (prefixed with `HORIZON_`)
/// 2. TOML configuration file (if `config_path` is provided)
/// 3. Struct defaults
///
/// # Errors
///
/// Returns `ConfigError::Load` if the configuration cannot be loaded.
pub fn load_config(config_path: Option<&str>) -> Result<SdkConfig> {
    let mut builder = config::Config::builder();

    if let Some(path) = config_path {
        builder = builder.add_source(config::File::with_name(path));
    }

    builder = builder.add_source(
        config::Environment::with_prefix("HORIZON")
            .separator("__")
            .try_parsing(true),
    );

    builder
        .build()
        .and_then(config::Config::try_deserialize)
        .map_err(|source| ConfigError::from(source).into())
}

/// Partition values for the Iceberg `data_row` table.
///
/// Must match the table's `partitioning` clause exactly:
/// `[data_stream_id, day(datetime), data_type, specification_id]`.
/// The `day` transform is hardcoded because the Iceberg sink expects the
/// Kafka key to be the formatted day string, not the raw timestamp.
fn data_row_partition_values(row: &DataRow) -> Vec<Option<PartitionValue>> {
    vec![
        Some(PartitionValue::from(row.data_stream_id)),
        Some(PartitionValue::from(
            row.datetime.format("%Y-%m-%d").to_string(),
        )),
        Some(PartitionValue::from(row.data_type.as_str())),
        Some(PartitionValue::from(row.specification_id)),
    ]
}

/// Partition values for the Iceberg `metadata_row` table.
///
/// Must match the table's `partitioning` clause exactly:
/// `[data_stream_id, day(datetime)]`.
fn metadata_row_partition_values(row: &MetadataRow) -> Vec<Option<PartitionValue>> {
    vec![
        Some(PartitionValue::from(row.data_stream_id)),
        Some(PartitionValue::from(
            row.datetime.format("%Y-%m-%d").to_string(),
        )),
    ]
}

#[cfg(test)]
#[allow(
    clippy::absolute_paths,
    clippy::assertions_on_result_states,
    clippy::panic_in_result_fn,
    clippy::single_call_fn,
    clippy::unwrap_used,
    reason = "unit tests use assertions, unwrap, and test helpers freely"
)]
mod tests {
    use super::*;

    fn parse_postgres_url(url: &str) -> Result<PostgresConfig> {
        HorizonSdkBuilder::default()
            .with_postgres_url(url)
            .resolve_postgres_config()
    }

    #[test]
    fn builder_organization_id_propagates() {
        let id = uuid::Uuid::new_v4();
        let builder = HorizonSdk::builder().with_organization_id(id);
        assert_eq!(builder.organization_id, Some(id));
    }

    #[test]
    fn builder_requires_postgres() {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        let result = rt.block_on(HorizonSdk::builder().build());
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("Postgres configuration is required"));
    }

    #[test]
    fn builder_with_config_sets_all_subsystems() {
        let config = SdkConfig {
            coverage: Some(CoverageConfig {
                base_url: "http://localhost:3000".to_owned(),
                timeout_secs: 30,
            }),
            iceberg: None,
            kafka: Some(KafkaConfig {
                bootstrap_servers: "localhost:9092".to_owned(),
                compression_type: "lz4".to_owned(),
                data_row_topic: "horizon.data_row".to_owned(),
                metadata_row_topic: "horizon.metadata_row".to_owned(),
                sasl_mechanism: "PLAIN".to_owned(),
                sasl_password: None,
                sasl_username: None,
                security_protocol: "PLAINTEXT".to_owned(),
            }),
            postgres: PostgresConfig {
                database: "horizon".to_owned(),
                host: "localhost".to_owned(),
                password: "test".to_owned(),
                port: 5432,
                possible_organization_id: None,
                role: None,
                sslmode: "disable".to_owned(),
                user: "test".to_owned(),
            },
        };
        let builder = HorizonSdk::builder().with_config(config);
        assert!(builder.postgres_config.is_some());
        assert!(builder.kafka_config.is_some());
        assert!(builder.iceberg_config.is_none());
        assert!(builder.coverage_config.is_some());
    }

    #[test]
    fn parse_postgres_url_defaults() {
        let config = parse_postgres_url("postgresql://localhost/horizon").unwrap();
        assert_eq!(config.host, "localhost");
        assert_eq!(config.port, 5432);
        assert_eq!(config.user, "postgres");
        assert_eq!(config.password, "");
        assert_eq!(config.database, "horizon");
        assert_eq!(config.sslmode, "require");
    }

    #[test]
    fn parse_postgres_url_full() {
        let config =
            parse_postgres_url("postgresql://user:pass@db.example.com:5433/mydb?sslmode=disable")
                .unwrap();
        assert_eq!(config.host, "db.example.com");
        assert_eq!(config.port, 5433);
        assert_eq!(config.user, "user");
        assert_eq!(config.password, "pass");
        assert_eq!(config.database, "mydb");
        assert_eq!(config.sslmode, "disable");
    }

    #[test]
    fn parse_postgres_url_invalid() {
        let result = parse_postgres_url("not a url");
        assert!(result.is_err());
        let err = result.unwrap_err().to_string();
        assert!(err.contains("invalid Postgres URL"));
    }
}