Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance.
Project price only 1 $
You can buy this project and download/modify it how often you want.
io.kaizensolutions.virgil.trace4cats.zio.extras.TracedCQLExecutor.scala Maven / Gradle / Ivy
package io.kaizensolutions.virgil.trace4cats.zio.extras
import com.datastax.oss.driver.api.core.metrics.Metrics
import com.datastax.oss.driver.api.core.{CqlSession, CqlSessionBuilder}
import trace4cats.model.{AttributeValue, SampleDecision, SpanKind, SpanStatus}
import io.kaizensolutions.trace4cats.zio.extras.{ZSpan, ZTracer}
import io.kaizensolutions.virgil.codecs.DecoderException
import io.kaizensolutions.virgil.configuration.PageState
import io.kaizensolutions.virgil.internal.CqlStatementRenderer
import io.kaizensolutions.virgil.internal.Proofs.*
import io.kaizensolutions.virgil.*
import zio.*
import zio.stream.*
import scala.collection.mutable
/**
* A Traced version of CQLExecutor that will trace and enrich Spans with
* metadata about the query
* @param tracer
* is the ZTracer to use for tracing
* @param underlying
* is the underlying CQLExecutor
* @param dropMarkerFromSpan
* if true, the marker will be dropped from the span (use this if you don't
* want sensitive data to be in the span)
*/
class TracedCQLExecutor(underlying: CQLExecutor, tracer: ZTracer, dropMarkerFromSpan: String => Boolean)
extends CQLExecutor {
override def execute[A](in: CQL[A])(implicit trace: Trace): Stream[Throwable, A] =
ZStream.unwrap(tracer.withSpan(extractQueryString(in), SpanKind.Internal) { span =>
val enrichSpanWithBindMarkers =
if (span.isSampled) enrichSpan(in, span, dropMarkerFromSpan)
else ZIO.unit
val spannedStream =
underlying
.execute(in)
.tapError(enrichSpanWithErrorInformation(span))
.tapErrorCause(cause =>
if (cause.isDie) span.setStatus(SpanStatus.Internal(cause.prettyPrint))
else ZIO.unit
)
enrichSpanWithBindMarkers *> ZIO.succeed(spannedStream)
})
override def executeMutation(in: CQL[MutationResult])(implicit trace: Trace): Task[MutationResult] =
tracer.withSpan(extractQueryString(in), SpanKind.Internal) { span =>
val isSampled = span.context.traceFlags.sampled == SampleDecision.Include
val enrichSpanWithBindMarkers =
if (isSampled) enrichSpan(in, span, dropMarkerFromSpan)
else ZIO.unit
enrichSpanWithBindMarkers *>
underlying
.executeMutation(in)
.tapError(enrichSpanWithErrorInformation(span))
.tapDefect(enrichSpanWithDefectInformation(span))
}
override def executePage[A](in: CQL[A], pageState: Option[PageState])(implicit
ev: A =:!= MutationResult,
trace: Trace
): Task[Paged[A]] = {
val query = extractQueryString(in)
val pageNr = pageState.map(_.toString()).getOrElse("begin")
tracer.withSpan(s"page-$pageNr: $query", SpanKind.Internal) { span =>
val isSampled = span.context.traceFlags.sampled == SampleDecision.Include
val enrichSpanWithBindMarkers =
if (isSampled) enrichSpan(in, span, dropMarkerFromSpan)
else ZIO.unit
enrichSpanWithBindMarkers *>
span.put("virgil.query-paged", AttributeValue.BooleanValue(true)) *>
underlying
.executePage(in, pageState)
.tapError(enrichSpanWithErrorInformation(span))
.tapDefect(enrichSpanWithDefectInformation(span))
}
}
private def enrichSpanWithDefectInformation[E](currentSpan: ZSpan)(cause: Cause[E]): UIO[Unit] = {
val addInfo = currentSpan.context.traceFlags.sampled.toBoolean
if (addInfo) currentSpan.setStatus(SpanStatus.Internal(cause.prettyPrint))
else ZIO.unit
}
private def enrichSpanWithErrorInformation(currentSpan: ZSpan)(error: Throwable): UIO[Unit] = {
val addInfo = currentSpan.context.traceFlags.sampled.toBoolean
val enrich = error match {
case e: DecoderException.PrimitiveReadFailure =>
currentSpan.putAll(
Map(
"virgil.error.message" -> AttributeValue.StringValue(e.message),
"virgil.error.cause" -> AttributeValue.StringValue(e.cause.getMessage)
)
)
case e: DecoderException.StructureReadFailure =>
currentSpan.putAll(
Map(
"virgil.error.message" -> AttributeValue.StringValue(e.message),
"virgil.error.cause" -> AttributeValue.StringValue(e.cause.getMessage),
"virgil.error.actual-db-structure" -> AttributeValue.StringValue(e.debugStructure),
"virgil.error.field" -> AttributeValue
.StringValue(e.field.map(_.toString).getOrElse("field-not-available"))
)
)
case e =>
val errorMessage =
if (e != null && e.getMessage != null)
Option("virgil.error.message" -> AttributeValue.StringValue(e.getMessage))
else None
val cause =
if (e != null && e.getCause != null && e.getCause.getMessage != null)
Option("virgil.error.cause" -> AttributeValue.StringValue(e.getCause.getMessage))
else None
val attrMap: Map[String, AttributeValue] = (errorMessage ++ cause).toMap
if (attrMap.nonEmpty) currentSpan.putAll(attrMap)
else ZIO.unit
}
if (addInfo) enrich
else ZIO.unit
}
private def extractQueryString[A](in: CQL[A]) = {
in.cqlType match {
case mutation: CQLType.Mutation =>
val (queryStr, _) = CqlStatementRenderer.render(mutation)
queryStr
case CQLType.Batch(mutations, _) =>
val sb = new mutable.StringBuilder()
sb.append("BATCH(")
mutations.foreach { mut =>
val (queryStr, _) = CqlStatementRenderer.render(mut)
sb.append(queryStr)
sb.append(";")
}
sb.append(")")
sb.toString()
case q @ CQLType.Query(_, _, _) =>
val (queryStr, _) = CqlStatementRenderer.render(q)
queryStr
}
}
private def enrichSpan[A](in: CQL[A], span: ZSpan, dropMarkerFromSpan: String => Boolean) =
in.cqlType match {
case q @ CQLType.Query(_, _, pullMode) =>
val attrMap = mutable.Map.empty[String, AttributeValue]
val (_, markers) = CqlStatementRenderer.render(q)
attrMap += ("virgil.query-type" -> AttributeValue.StringValue("query"))
attrMap += ("virgil.elements-to-pull" -> AttributeValue.StringValue(pullMode.toString))
markers.underlying.foreach { m =>
val (name, marker) = m
if (dropMarkerFromSpan(name.name)) ()
else attrMap += (s"virgil.bind-markers.$name" -> AttributeValue.StringValue(marker.value.toString))
}
span.putAll(attrMap.toMap)
case mutation: CQLType.Mutation =>
val attrMap = mutable.Map.empty[String, AttributeValue]
val (_, markers) = CqlStatementRenderer.render(mutation)
attrMap += ("virgil.query-type" -> AttributeValue.StringValue("mutation"))
markers.underlying.foreach { m =>
val (name, marker) = m
if (dropMarkerFromSpan(name.name)) ()
else attrMap += (s"virgil.bind-markers.$name" -> AttributeValue.StringValue(marker.value.toString))
}
span.putAll(attrMap.toMap)
case CQLType.Batch(mutations, batchType: BatchType) =>
var queryCounter = 0
val attrMap = mutable.Map.empty[String, AttributeValue]
attrMap += ("virgil.batch-type" -> AttributeValue.StringValue(batchType.toString))
attrMap += ("virgil.query-type" -> AttributeValue.StringValue("batch-mutation"))
mutations.foreach { mut =>
val (qs, markers) = CqlStatementRenderer.render(mut)
attrMap += (s"virgil.query.$queryCounter" -> AttributeValue.StringValue(qs))
markers.underlying.foreach { m =>
val (name, marker) = m
if (dropMarkerFromSpan(name.name)) ()
else
attrMap += (s"virgil.bind-markers.$queryCounter.$name" -> AttributeValue.StringValue(
marker.value.toString
))
}
queryCounter += 1
}
span.putAll(attrMap.toMap)
}
override def metrics: UIO[Option[Metrics]] =
tracer.span("metrics")(underlying.metrics)
}
object TracedCQLExecutor {
def apply(
underlying: CQLExecutor,
tracer: ZTracer,
dropMarkerFromSpan: String => Boolean
): TracedCQLExecutor = new TracedCQLExecutor(underlying, tracer, dropMarkerFromSpan)
def apply(
builder: CqlSessionBuilder,
tracer: ZTracer,
dropMarkerFromSpan: String => Boolean = _ => false
): RIO[Scope, CQLExecutor] =
CQLExecutor(builder)
.map(new TracedCQLExecutor(_, tracer, dropMarkerFromSpan))
def fromCqlSession(
session: CqlSession,
tracer: ZTracer,
dropMarkerFromSpan: String => Boolean = _ => false
): CQLExecutor =
new TracedCQLExecutor(CQLExecutor.fromCqlSession(session), tracer, dropMarkerFromSpan)
val layer: URLayer[CQLExecutor & ZTracer, CQLExecutor] =
ZLayer.fromZIO(
for {
tracer <- ZIO.service[ZTracer]
cql <- ZIO.service[CQLExecutor]
} yield new TracedCQLExecutor(cql, tracer, _ => false)
)
}
