hive-router 0.2.0

use opentelemetry::KeyValue;
use tracing::{field::Empty, info_span, record_all, Level, Span};
use tracing_opentelemetry::OpenTelemetrySpanExt;

use crate::telemetry::{
    metrics::demand_control_metrics::DemandControlResultCode,
    traces::{
        disabled_span, is_level_enabled,
        spans::{
            attributes::{
                self, ERROR_MESSAGE, ERROR_TYPE, HIVE_ERROR_AFFECTED_PATH, HIVE_ERROR_PATH,
                HIVE_ERROR_SUBGRAPH_NAME, HIVE_KIND,
            },
            kind::{HiveEventKind, HiveSpanKind},
            observed_error::ObservedError,
            TARGET_NAME,
        },
    },
};

pub struct GraphQLParseSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLParseSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLParseSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLParseSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlParse.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.parse",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
            "cache.hit" = Empty,
            "graphql.operation.name" = Empty,
            "graphql.operation.type" = Empty,
            "graphql.document.hash" = Empty,
        );
        GraphQLParseSpan { span }
    }

    pub fn record_cache_hit(&self, hit: bool) {
        self.span.record(attributes::CACHE_HIT, hit);
    }

    pub fn record_operation_identity(&self, identity: GraphQLSpanOperationIdentity) {
        if self.span.is_disabled() {
            return;
        }

        record_all!(
            self.span,
            "graphql.operation.name" = identity.name,
            "graphql.operation.type" = identity.operation_type,
            "graphql.document.hash" = identity.client_document_hash,
        );
    }
}

pub struct GraphQLValidateSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLValidateSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLValidateSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLValidateSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlValidate.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.validate",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
            "cache.hit" = Empty,
        );
        GraphQLValidateSpan { span }
    }

    pub fn record_cache_hit(&self, hit: bool) {
        self.span.record(attributes::CACHE_HIT, hit);
    }
}

pub struct GraphQLVariableCoercionSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLVariableCoercionSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLVariableCoercionSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLVariableCoercionSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlVariableCoercion.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.variable_coercion",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
        );
        GraphQLVariableCoercionSpan { span }
    }
}

pub struct GraphQLNormalizeSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLNormalizeSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLNormalizeSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLNormalizeSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlNormalize.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.normalize",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
            "cache.hit" = Empty,
        );
        GraphQLNormalizeSpan { span }
    }

    pub fn record_cache_hit(&self, hit: bool) {
        self.span.record(attributes::CACHE_HIT, hit);
    }
}

pub struct GraphQLAuthorizeSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLAuthorizeSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLAuthorizeSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLAuthorizeSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlAuthorize.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.authorize",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
        );
        GraphQLAuthorizeSpan { span }
    }
}

pub struct GraphQLPlanSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLPlanSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLPlanSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLPlanSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlPlan.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.plan",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
            "cache.hit" = Empty,
        );
        GraphQLPlanSpan { span }
    }

    pub fn record_cache_hit(&self, hit: bool) {
        self.span.record(attributes::CACHE_HIT, hit);
    }
}

pub struct GraphQLExecuteSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLExecuteSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLExecuteSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLExecuteSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlExecute.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.execute",
            "hive.kind" = kind,
            "otel.kind" = "Internal",
        );
        GraphQLExecuteSpan { span }
    }
}

pub struct GraphQLOperationSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLOperationSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl Default for GraphQLOperationSpan {
    fn default() -> Self {
        Self::new()
    }
}

impl GraphQLOperationSpan {
    pub fn new() -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphqlOperation.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.operation",
            "hive.kind" = kind,
            "otel.status_code" = Empty,
            "otel.kind" = "Server",
            "error.type" = Empty,
            "graphql.operation.name" = Empty,
            "graphql.operation.type" = Empty,
            "graphql.operation.id" = Empty,
            "graphql.document.hash" = Empty,
            "graphql.document" = Empty,
            "cost.estimated" = Empty,
            "cost.actual" = Empty,
            "cost.delta" = Empty,
            "cost.result" = Empty,
            "hive.graphql.error.count" = Empty,
            "hive.graphql.error.codes" = Empty,
            "hive.graphql.operation.hash" = Empty,
            "hive.client.name" = Empty,
            "hive.client.version" = Empty,
            "hive.target" = Empty,
            "cost.formula_cache_hit" = Empty,
        );
        GraphQLOperationSpan { span }
    }

    pub fn record_hive_target(&self, target: Option<&str>) {
        if let Some(target) = target {
            self.span.record(attributes::HIVE_TARGET, target);
        }
    }

    pub fn record_error_count(&self, count: usize) {
        self.span
            .record(attributes::HIVE_GRAPHQL_ERROR_COUNT, count);
    }

    pub fn record_errors(&self, errors_fn: impl FnOnce() -> Vec<ObservedError>) {
        if self.is_disabled() {
            return;
        }

        let errors = errors_fn();
        record_error_codes_to_span(&self.span, &errors);
        record_error_events_to_span(&self.span, errors);
    }

    pub fn record_details(
        &self,
        document: &str,
        identity: GraphQLSpanOperationIdentity,
        client_name: Option<&str>,
        client_version: Option<&str>,
        hash: &str,
    ) {
        if self.span.is_disabled() {
            return;
        }

        record_all!(
            self.span,
            "graphql.document" = document,
            "graphql.operation.name" = identity.name,
            "graphql.operation.type" = identity.operation_type,
            "graphql.document.hash" = identity.client_document_hash,
            "hive.graphql.operation.hash" = hash,
            "hive.client.name" = client_name,
            "hive.client.version" = client_version,
        );
    }

    pub fn record_demand_control(
        &self,
        estimated: u64,
        actual: Option<u64>,
        delta: Option<i64>,
        result: &DemandControlResultCode,
    ) {
        if self.span.is_disabled() {
            return;
        }

        self.span.record(attributes::COST_ESTIMATED, estimated);

        if let Some(actual) = actual {
            self.span.record(attributes::COST_ACTUAL, actual);
        }

        if let Some(delta) = delta {
            self.span.record(attributes::COST_DELTA, delta);
        }

        self.span.record(attributes::COST_RESULT, result.as_str());
    }
}

pub struct GraphQLSubgraphOperationSpan {
    pub span: Span,
}

impl std::ops::Deref for GraphQLSubgraphOperationSpan {
    type Target = Span;
    fn deref(&self) -> &Self::Target {
        &self.span
    }
}

impl GraphQLSubgraphOperationSpan {
    pub fn new(subgraph_name: &str, document: &str) -> Self {
        if !is_level_enabled(Level::INFO) {
            return Self {
                span: disabled_span(),
            };
        }

        let kind: &'static str = HiveSpanKind::GraphQLSubgraphOperation.into();
        let span = info_span!(
            target: TARGET_NAME,
            "graphql.subgraph.operation",
            "hive.kind" = kind,
            "otel.status_code" = Empty,
            "otel.kind" = "Client",
            "error.type" = Empty,
            "graphql.operation.name" = Empty,
            "graphql.operation.type" = Empty,
            "graphql.document.hash" = Empty,
            "graphql.document" = document,
            "hive.graphql.error.count" = Empty,
            "hive.graphql.error.codes" = Empty,
            // Hive Console Attributes
            "hive.graphql.subgraph.name" = subgraph_name,
        );
        GraphQLSubgraphOperationSpan { span }
    }

    pub fn record_error_count(&self, count: usize) {
        self.span
            .record(attributes::HIVE_GRAPHQL_ERROR_COUNT, count);
    }

    pub fn record_errors(&self, errors_fn: impl FnOnce() -> Vec<ObservedError>) {
        if self.is_disabled() {
            return;
        }

        let errors = errors_fn();
        record_error_codes_to_span(&self.span, &errors);
        record_error_events_to_span(&self.span, errors);
    }

    pub fn record_operation_identity(&self, identity: GraphQLSpanOperationIdentity) {
        record_all!(
            self.span,
            "graphql.operation.name" = identity.name,
            "graphql.operation.type" = identity.operation_type,
            "graphql.document.hash" = identity.client_document_hash,
        );
    }
}

fn record_error_codes_to_span(span: &Span, errors: &[ObservedError]) {
    let mut codes: Vec<&str> = errors.iter().filter_map(|e| e.code.as_deref()).collect();

    if codes.is_empty() {
        return;
    }

    codes.sort_unstable();
    codes.dedup();

    span.record(attributes::HIVE_GRAPHQL_ERROR_CODES, codes.join(","));
}

fn record_error_events_to_span(span: &Span, errors: Vec<ObservedError>) {
    if errors.is_empty() {
        return;
    }

    for error in errors {
        let message = &error.message;
        let mut attributes: Vec<KeyValue> = Vec::with_capacity(6);
        let kind: &'static str = HiveEventKind::GraphQLError.into();

        attributes.push(KeyValue::new(HIVE_KIND, kind));
        attributes.push(KeyValue::new(ERROR_MESSAGE, message.to_string()));
        attributes.push(KeyValue::new(
            ERROR_TYPE,
            error.code.unwrap_or(String::from("unknown")).to_string(),
        ));

        if let Some(service_name) = error.service_name {
            attributes.push(KeyValue::new(
                HIVE_ERROR_SUBGRAPH_NAME,
                service_name.to_string(),
            ));
        }

        if let Some(affected_path) = error.affected_path {
            attributes.push(KeyValue::new(
                HIVE_ERROR_AFFECTED_PATH,
                affected_path.to_string(),
            ));
        }

        if let Some(path) = error.path {
            attributes.push(KeyValue::new(HIVE_ERROR_PATH, path));
        }

        span.add_event(message.to_string(), attributes);
    }
}

pub struct GraphQLSpanOperationIdentity<'a> {
    pub name: Option<&'a str>,
    pub operation_type: &'a str,
    /// Hash of the original document sent to the router, by the client.
    pub client_document_hash: &'a str,
}