diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog + +## 0.1.0.0 + +Initial release. diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,28 @@ +BSD 3-Clause License + +Copyright (c) 2026, Kronor + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/README.md b/README.md new file mode 100644 --- /dev/null +++ b/README.md @@ -0,0 +1,71 @@ +# hasql-opentelemetry + +OpenTelemetry tracing for [hasql](https://hackage.haskell.org/package/hasql) 2. + +hasql 2 talks to Postgres through a [pqi](https://hackage.haskell.org/package/pqi) +adapter, which is a record of `IO` closures. This package wraps that adapter, so +each command sent to Postgres becomes a client span. Your statements and +sessions stay as they are. + +```haskell +import Hasql.Connection qualified as Connection +import Hasql.OpenTelemetry qualified as HasqlOtel +import Pqi.Ffi qualified + +Connection.acquire (HasqlOtel.traced Pqi.Ffi.adapter) settings +``` + +Pass the same wrapped adapter to `hasql-pool`. It works with both `pqi-ffi` and `pqi-native`. + +## What a span carries + +| Attribute | Value | +|-----------------------------------|--------------------------------------------------------------| +| `db.system.name` | `postgresql` | +| `db.namespace` | database name | +| `server.address`, `server.port` | from the connection | +| `db.query.text` | the statement, whitespace collapsed, cut at 1024 characters | +| `db.operation.name` | leading keyword (`SELECT`, `INSERT`, …) | +| `db.postgresql.prepared_statement`| name it was prepared under, for executions of prepared statements | +| `db.response.returned_rows` | rows returned | +| `db.postgresql.affected_rows` | rows affected by a command that returns none | +| `db.response.status_code`, `error.type` | SQLSTATE, on failure | + +Spans are timed from the send to the last result, and this holds in pipeline +mode too, which is how hasql 2 runs every roundtrip. Prepares appear as separate +`PREPARE …` spans. + +By default spans are only started when a span is already active on the calling +thread, so pool warm-up and health checks don't start their own traces. Set +`requireParentSpan = False` to trace everything. + +## What is never recorded + +**Parameter values** and the **text of server error reports** are never +recorded. Postgres often puts the offending value in the error message (for +example, `invalid input syntax for type uuid: "…"`), so a failure records only +its SQLSTATE. + +SQL built at runtime with literal values spliced into it is recorded exactly as +sent. Pass such values as parameters instead. + +## Datadog + +`tracedWith defaultConfig { datadogCompat = True }` changes three things so +Datadog APM renders these spans as database queries: + +- span names become `postgresql.query`, `postgresql.prepare` and + `postgresql.describe`; +- `span.type=sql` is added; +- the statement is also set as `resource.name` and `db.statement`. + +## Testing + +The test suite runs real sessions against Postgres, configured through the +usual `PGHOST` / `PGPORT` / `PGUSER` / `PGPASSWORD` / `PGDATABASE` variables. +Without `PGHOST` it is skipped. + +```sh +docker run --rm -d -p 5432:5432 -e POSTGRES_PASSWORD=postgres postgres:18 +PGHOST=localhost PGPASSWORD=postgres cabal test +``` diff --git a/hasql-opentelemetry.cabal b/hasql-opentelemetry.cabal new file mode 100644 --- /dev/null +++ b/hasql-opentelemetry.cabal @@ -0,0 +1,72 @@ +cabal-version: 3.0 +name: hasql-opentelemetry +version: 0.1.0.0 +synopsis: OpenTelemetry tracing for hasql +description: + Wraps the pqi adapter that hasql 2 connects through, so that every command + sent to Postgres becomes an OpenTelemetry client span carrying the SQL, + the row count and, on failure, the SQLSTATE. Spans are timed from the send + to the last result, including in pipeline mode. Parameter values and error + messages are never recorded. + . + See "Hasql.OpenTelemetry" for usage. + +category: Database, Telemetry +homepage: https://github.com/kronor-io/hasql-opentelemetry +bug-reports: https://github.com/kronor-io/hasql-opentelemetry/issues +license: BSD-3-Clause +license-file: LICENSE +author: Kronor +maintainer: lorenzo@kronor.io +copyright: 2026 Kronor +build-type: Simple +tested-with: GHC ==9.12.4 +extra-doc-files: + CHANGELOG.md + README.md + +source-repository head + type: git + location: https://github.com/kronor-io/hasql-opentelemetry + +common common + default-language: GHC2021 + default-extensions: + BlockArguments + LambdaCase + OverloadedStrings + + ghc-options: + -Wall -Wcompat -Widentities -Wincomplete-record-updates + -Wincomplete-uni-patterns -Wpartial-fields -Wredundant-constraints + +library + import: common + hs-source-dirs: src + exposed-modules: Hasql.OpenTelemetry + build-depends: + , base >=4.18 && <5 + , bytestring >=0.11 && <0.13 + , containers >=0.6 && <0.9 + , hs-opentelemetry-api ^>=1.0 + , pqi ^>=1.1 + , text >=2.0 && <2.2 + , unordered-containers >=0.2.14 && <0.3 + +test-suite hasql-opentelemetry-test + import: common + type: exitcode-stdio-1.0 + hs-source-dirs: test + main-is: Main.hs + ghc-options: -threaded + build-depends: + , base + , hasql ^>=2.0 + , hasql-opentelemetry + , hs-opentelemetry-api + , hs-opentelemetry-exporter-in-memory ^>=1.0 + , hspec >=2.10 && <2.12 + , pqi + , pqi-ffi ^>=1.0 + , pqi-native ^>=1.0 + , text diff --git a/src/Hasql/OpenTelemetry.hs b/src/Hasql/OpenTelemetry.hs new file mode 100644 --- /dev/null +++ b/src/Hasql/OpenTelemetry.hs @@ -0,0 +1,525 @@ +-- | OpenTelemetry tracing for hasql, by decorating its pqi adapter. +-- +-- 'Pqi.Adapter' and 'Pqi.Connection' are records of @IO@ closures rather than +-- opaque handles, so an adapter can be wrapped: 'traced' takes whichever +-- backend you picked (@pqi-ffi@ or @pqi-native@) and returns one whose +-- connections open a client span per command. Nothing above changes — hasql +-- keeps calling the same fields — and the span carries what an application +-- level span cannot: the SQL that actually went to Postgres. +-- +-- > import Hasql.Connection qualified as Connection +-- > import Hasql.OpenTelemetry qualified as HasqlOtel +-- > import Pqi.Ffi qualified +-- > +-- > Connection.acquire (HasqlOtel.traced Pqi.Ffi.adapter) settings +-- +-- The same adapter goes to @hasql-pool@. Spans are reported to the global +-- tracer provider, and by default only under a span that is already active on +-- the calling thread (see 'requireParentSpan'). +-- +-- == Where the SQL comes from +-- +-- hasql 2 sends a statement twice over: once as @sendPrepare name sql@ and +-- from then on as @sendQueryPrepared name params@, which never names the SQL +-- again. The prepare is therefore the only place a connection ever sees the +-- text, and the attributes derived from it are cached per statement name for +-- the life of the connection. +-- +-- == How a span is timed +-- +-- Every hasql roundtrip is pipelined: @enterPipelineMode@, the sends, +-- @pipelineSync@, then a @getResult@ loop. A send therefore returns long +-- before the server has answered, so timing the send alone would measure +-- nothing. Instead a span starts at the send and is queued; @getResult@ +-- returning 'Nothing' — libpq's end-of-this-command marker, in pipeline mode +-- and out of it — ends the one at the head of the queue. Commands are +-- answered in the order they were sent, so the queue stays aligned as long as +-- every result-producing send is wrapped, which is why all of them are, not +-- just the ones hasql happens to call today. +-- +-- Interrupted sessions break that pairing: hasql drains and cancels, and a +-- drain reads one 'Nothing' for a pipeline that may have had several commands +-- in it. Leaving pipeline mode therefore closes whatever is still queued, +-- marked @db.postgresql.results_unread@, which bounds a leak at one roundtrip. +-- +-- == What is never recorded +-- +-- Parameter values, and no text at all from a server error report. hasql +-- passes parameters out of band of the SQL, and they routinely hold personal +-- data and secrets; only the statement text, which is constant for a +-- statement written in the source, reaches a span. A failure records the +-- SQLSTATE and stops there, because every human-readable field of the report +-- can carry the offending value — not just the detail. Postgres puts it in +-- the primary message whenever a cast fails (@invalid input syntax for type +-- uuid: \"...\"@), and a @raise exception \'... %\'@ in PL\/pgSQL +-- interpolates whatever it is given into the same field. +-- +-- Statement text built at runtime with literals spliced into it is recorded +-- as sent. Keep such values in parameters. +module Hasql.OpenTelemetry ( + traced, + tracedWith, + Config (..), + defaultConfig, +) where + +import Control.Exception (bracket, mask, onException) +import Control.Monad (when, (>=>)) +import Data.ByteString (ByteString) +import Data.ByteString.Char8 qualified as BS8 +import Data.Char qualified as Char +import Data.Foldable (for_, traverse_) +import Data.HashMap.Strict (HashMap) +import Data.HashMap.Strict qualified as HashMap +import Data.IORef (IORef, atomicModifyIORef', modifyIORef', newIORef, readIORef, writeIORef) +import Data.Maybe (catMaybes, fromMaybe, isNothing) +import Data.Sequence (Seq) +import Data.Sequence qualified as Seq +import Data.Text (Text) +import Data.Text qualified as T +import Data.Text.Encoding (decodeUtf8Lenient) +import OpenTelemetry.Attributes (Attribute, ToAttribute (toAttribute)) +import OpenTelemetry.Context qualified as Context +import OpenTelemetry.Context.ThreadLocal qualified as ThreadLocal +import OpenTelemetry.Trace.Core (Span, Tracer) +import OpenTelemetry.Trace.Core qualified as Otel +import Pqi qualified +import Text.Read (readMaybe) + + +-- | How spans are reported. Start from 'defaultConfig' and override fields. +data Config = Config + { tracerName :: Text + -- ^ The instrumentation scope spans are reported under. + , maxStatementLength :: Int + -- ^ Longer statements are cut to this many characters. The text is sent + -- on every execution of a statement, so this is what a span costs in + -- ingestion. + , requireParentSpan :: Bool + -- ^ Only trace commands issued under a span that is already active on the + -- calling thread. Keeps the volume proportional to traced work: pool + -- warm-up, @LISTEN@ connections and health-check pings would otherwise + -- each start a trace of their own. + , datadogCompat :: Bool + -- ^ Shape spans the way Datadog APM renders database calls: span names + -- @postgresql.query@ \/ @postgresql.prepare@ \/ @postgresql.describe@, + -- @span.type=sql@, and the statement also as @resource.name@ and + -- @db.statement@. Datadog reads @span.type@ off the span before deriving + -- one, and @sql@ is what makes it show the statement as a query; it + -- derives a resource from @db.statement@ only after @resource.name@ has + -- failed to supply one, so both are set. + } + + +-- | Semantic-convention span names and attributes, spans only under an +-- existing parent, statements cut at 1024 characters. +defaultConfig :: Config +defaultConfig = + Config + { tracerName = "hasql-opentelemetry" + , maxStatementLength = 1024 + , requireParentSpan = True + , datadogCompat = False + } + + +-- | Decorate an adapter so that every connection it opens reports its +-- commands to OpenTelemetry, with 'defaultConfig'. +traced :: Pqi.Adapter -> Pqi.Adapter +traced = tracedWith defaultConfig + + +-- | 'traced' with an explicit t'Config'. +tracedWith :: Config -> Pqi.Adapter -> Pqi.Adapter +tracedWith config adapter = + adapter + { Pqi.connectdb = Pqi.connectdb adapter >=> instrument config + , Pqi.connectStart = Pqi.connectStart adapter >=> instrument config + } + + +-- | Per-connection tracing state. One of these is closed over by the wrapped +-- closures of a single 'Pqi.Connection'. +data ConnectionTracer = ConnectionTracer + { config :: Config + , tracer :: Tracer + , connection :: Pqi.Connection + -- ^ The undecorated connection, which is what every wrapper calls through + -- to. + , connectionAttributes :: IORef (Maybe [(Text, Attribute)]) + -- ^ Read from the connection on first use. Deferred rather than read at + -- connect time because @connectStart@ returns before the connection + -- exists. + , statements :: IORef (HashMap ByteString Statement) + -- ^ Each prepared statement, by the name it was prepared under. + -- Deliberately unbounded: hasql's own statement cache is per-connection + -- and unbounded too and goes on executing under the names it holds, so + -- evicting here would silently strip the SQL off statements that are + -- still running. Both are bounded by the number of distinct statements in + -- the program and the connection's lifetime, and 'forget' clears this one + -- where hasql resets its own. + , pending :: IORef (Seq (Maybe Span)) + -- ^ Commands sent whose results have not been read yet, oldest first. + -- 'Nothing' for a command that is not being traced, which still has to + -- take its turn or the queue stops matching the results. + } + + +-- | What the span of a command learns from its SQL. +data Statement = Statement + { operation :: Maybe Text + , attributes :: [(Text, Attribute)] + } + + +data Command = Query | Prepare | Describe + + +-- | How many unanswered commands to keep spans for before dropping the +-- oldest. A hasql pipeline holds three; anything beyond that means the queue +-- has lost its place and the span would never be ended otherwise. +maxPendingCommands :: Int +maxPendingCommands = 64 + + +instrument :: Config -> Pqi.Connection -> IO Pqi.Connection +instrument cfg pq + | Pqi.isNullConnection pq = pure pq + | otherwise = do + provider <- Otel.getGlobalTracerProvider + attributeCache <- newIORef Nothing + statementCache <- newIORef mempty + pendingCommands <- newIORef mempty + let ct = + ConnectionTracer + { config = cfg + , tracer = Otel.makeTracer provider (Otel.instrumentationLibrary (tracerName cfg) "") Otel.tracerOptions + , connection = pq + , connectionAttributes = attributeCache + , statements = statementCache + , pending = pendingCommands + } + ofSql = pure . parseStatement cfg + ofName = preparedStatement ct + pure + pq + { Pqi.exec = \sql -> + tracedSync ct Query (ofSql sql) (Pqi.exec pq sql) + , Pqi.execParams = \sql params resultFormat -> + tracedSync ct Query (ofSql sql) (Pqi.execParams pq sql params resultFormat) + , Pqi.prepare = \name sql oids -> do + rememberStatement ct name sql + tracedSync ct Prepare (ofSql sql) (Pqi.prepare pq name sql oids) + , Pqi.execPrepared = \name params resultFormat -> + tracedSync ct Query (ofName name) (Pqi.execPrepared pq name params resultFormat) + , Pqi.describePrepared = \name -> + tracedSync ct Describe (ofName name) (Pqi.describePrepared pq name) + , Pqi.describePortal = \name -> + tracedSync ct Describe (ofName name) (Pqi.describePortal pq name) + , Pqi.sendQuery = \sql -> + tracedSend ct Query (ofSql sql) (Pqi.sendQuery pq sql) + , Pqi.sendQueryParams = \sql params resultFormat -> + tracedSend ct Query (ofSql sql) (Pqi.sendQueryParams pq sql params resultFormat) + , Pqi.sendPrepare = \name sql oids -> do + rememberStatement ct name sql + tracedSend ct Prepare (ofSql sql) (Pqi.sendPrepare pq name sql oids) + , Pqi.sendQueryPrepared = \name params resultFormat -> + tracedSend ct Query (ofName name) (Pqi.sendQueryPrepared pq name params resultFormat) + , Pqi.sendDescribePrepared = \name -> + tracedSend ct Describe (ofName name) (Pqi.sendDescribePrepared pq name) + , Pqi.sendDescribePortal = \name -> + tracedSend ct Describe (ofName name) (Pqi.sendDescribePortal pq name) + , -- The sync produces a result of its own, so it takes a turn in the + -- queue. It is never worth a span: it measures the pipeline, which + -- the commands in it already do. + Pqi.pipelineSync = do + sent <- Pqi.pipelineSync pq + when sent (enqueue ct Nothing) + pure sent + , -- Masked because the read advances the connection's protocol state: + -- an exception landing between it and the bookkeeping would leave + -- the queue one command behind, and every later result would be + -- attributed to the wrong statement. + Pqi.getResult = mask \restore -> do + result <- restore (Pqi.getResult pq) + maybe (commandAnswered ct) (resultReceived ct) result + pure result + , Pqi.exitPipelineMode = do + left <- Pqi.exitPipelineMode pq + when left (abandonPending ct) + pure left + , Pqi.reset = do + Pqi.reset pq + forget ct + , Pqi.resetStart = do + started <- Pqi.resetStart pq + when started (forget ct) + pure started + , Pqi.finish = do + abandonPending ct + Pqi.finish pq + } + + +-- * Spans + + +-- | The span name for a command. Without 'datadogCompat' it follows the +-- database semantic conventions: the operation where one is known, the +-- database system otherwise. +spanName :: Config -> Command -> Statement -> Text +spanName cfg command statement + | datadogCompat cfg = case command of + Query -> "postgresql.query" + Prepare -> "postgresql.prepare" + Describe -> "postgresql.describe" + | otherwise = case command of + Query -> fromMaybe "postgresql" (operation statement) + Prepare -> maybe "PREPARE" ("PREPARE " <>) (operation statement) + Describe -> "DESCRIBE" + + +-- | Start a span for a command, unless this connection is not being traced +-- right now. The span is deliberately left out of the thread-local context: +-- it is a leaf, and several of them can be open at once on one thread. +startSpan :: ConnectionTracer -> Command -> IO Statement -> IO (Maybe Span) +startSpan ct command mkStatement = do + context <- ThreadLocal.getContext + if requireParentSpan (config ct) && isNothing (Context.lookupSpan context) + then pure Nothing + else do + shared <- sharedAttributes ct + statement <- mkStatement + fmap Just $ + Otel.createSpanWithoutCallStack (tracer ct) context (spanName (config ct) command statement) $ + Otel.addAttributesToSpanArguments + (HashMap.fromList (shared <> attributes statement)) + Otel.defaultSpanArguments{Otel.kind = Otel.Client} + + +-- | Wrap a command that sends and reads its own results. +tracedSync :: + ConnectionTracer -> + Command -> + IO Statement -> + IO (Maybe Pqi.Result) -> + IO (Maybe Pqi.Result) +tracedSync ct command mkStatement act = + bracket + (startSpan ct command mkStatement) + (traverse_ (`Otel.endSpan` Nothing)) + \mSpan -> do + result <- act + for_ mSpan (for_ result . annotateResult) + pure result + + +-- | Wrap a command that only queues bytes for the server. The span stays open +-- until its results have been read off the connection. +-- +-- Masked for the same reason as @getResult@: a command the backend accepted +-- and did not queue here is one the queue will never account for. +tracedSend :: ConnectionTracer -> Command -> IO Statement -> IO Bool -> IO Bool +tracedSend ct command mkStatement send = mask \restore -> do + mSpan <- startSpan ct command mkStatement + sent <- restore send `onException` failSpan mSpan "interrupted before the command was sent" + if sent + then True <$ enqueue ct mSpan + else False <$ failSpan mSpan "the command could not be sent" + + +-- | Add a command to the queue of the unanswered. +enqueue :: ConnectionTracer -> Maybe Span -> IO () +enqueue ct entry = do + queued <- readIORef (pending ct) + case queued of + oldest Seq.:<| rest + | Seq.length queued >= maxPendingCommands -> do + abandonSpan oldest "no result was matched to this command" + writeIORef (pending ct) (rest Seq.|> entry) + _ -> writeIORef (pending ct) (queued Seq.|> entry) + + +-- | End the span of the oldest unanswered command: @getResult@ has returned +-- the 'Nothing' that closes it. +commandAnswered :: ConnectionTracer -> IO () +commandAnswered ct = + readIORef (pending ct) >>= \case + Seq.Empty -> pure () + oldest Seq.:<| rest -> do + writeIORef (pending ct) rest + for_ oldest (`Otel.endSpan` Nothing) + + +-- | Record what a result says about the command it belongs to, which is the +-- oldest unanswered one. +resultReceived :: ConnectionTracer -> Pqi.Result -> IO () +resultReceived ct result = + readIORef (pending ct) >>= \case + Just span_ Seq.:<| _ -> annotateResult span_ result + _ -> pure () + + +annotateResult :: Span -> Pqi.Result -> IO () +annotateResult span_ result = + Pqi.resultStatus result >>= \case + Pqi.TuplesOk -> do + rows <- Pqi.ntuples result + Otel.addAttribute span_ "db.response.returned_rows" (fromIntegral rows :: Int) + Pqi.CommandOk -> do + affected <- Pqi.cmdTuples result + for_ (affected >>= rowCount) (Otel.addAttribute span_ "db.postgresql.affected_rows") + Pqi.FatalError -> recordFailure + Pqi.BadResponse -> recordFailure + Pqi.PipelineAbort -> recordFailure + -- 'Pqi.NonfatalError' is libpq's notice-or-warning status, not a + -- failed command. Marking it would report errors for statements + -- Postgres was content with. + _ -> pure () + where + -- The SQLSTATE is the whole of the error report that is safe to send: it + -- is a five-character code from the protocol, and every other field can + -- carry the value that caused the failure. + recordFailure = do + sqlstate <- fmap decodeUtf8Lenient <$> Pqi.resultErrorField result Pqi.DiagSqlstate + for_ sqlstate \code -> do + Otel.addAttribute span_ "db.response.status_code" code + Otel.addAttribute span_ "error.type" code + Otel.setStatus span_ . Otel.Error $ + maybe + "postgres reported an error" + ("postgres reported an error, SQLSTATE " <>) + sqlstate + + +-- | The rows a command affected, which @cmdTuples@ reports as a bare decimal. +-- Only a clean parse of the whole field counts: a command that touches no rows +-- reports an empty one. +rowCount :: ByteString -> Maybe Int +rowCount raw = case BS8.readInt raw of + Just (rows, rest) | BS8.null rest -> Just rows + _ -> Nothing + + +-- | End a span for a command that never reached the backend. Its outcome is +-- known, and it is a failure. +failSpan :: Maybe Span -> Text -> IO () +failSpan mSpan reason = + for_ mSpan \span_ -> do + Otel.setStatus span_ (Otel.Error reason) + Otel.endSpan span_ Nothing + + +-- | End a span whose command reached the backend but whose result we never +-- saw. Deliberately not an error: the statement most likely succeeded and we +-- simply stopped reading, which happens on every interrupted session. The +-- duration is a lower bound, hence the marker. +abandonSpan :: Maybe Span -> Text -> IO () +abandonSpan mSpan reason = + for_ mSpan \span_ -> do + Otel.addAttribute span_ "db.postgresql.results_unread" reason + Otel.endSpan span_ Nothing + + +-- | Close every span still waiting for a result. Called where the connection +-- can no longer answer them in order. +abandonPending :: ConnectionTracer -> IO () +abandonPending ct = do + queued <- atomicModifyIORef' (pending ct) (Seq.empty,) + for_ queued (`abandonSpan` "the pipeline ended before the result arrived") + + +-- | Drop everything remembered about a connection that has been reset: its +-- prepared statements are gone from the server and its queued commands will +-- never be answered. +forget :: ConnectionTracer -> IO () +forget ct = do + abandonPending ct + writeIORef (statements ct) mempty + + +-- * Attributes + + +-- | A statement referred to by the name it was prepared under. The SQL is +-- whatever the prepare on this connection recorded; a miss leaves the span +-- with the name alone, which still says which statement it was. +preparedStatement :: ConnectionTracer -> ByteString -> IO Statement +preparedStatement ct name = do + remembered <- HashMap.lookup name <$> readIORef (statements ct) + let statement = fromMaybe (Statement Nothing []) remembered + pure + statement + { attributes = + ("db.postgresql.prepared_statement", toAttribute (decodeUtf8Lenient name)) + : attributes statement + } + + +rememberStatement :: ConnectionTracer -> ByteString -> ByteString -> IO () +rememberStatement ct name sql = + modifyIORef' (statements ct) (HashMap.insert name (parseStatement (config ct) sql)) + + +parseStatement :: Config -> ByteString -> Statement +parseStatement cfg sql = + Statement + { operation = op + , attributes = + ("db.query.text", text) + : maybe [] (\o -> [("db.operation.name", toAttribute o)]) op + <> if datadogCompat cfg + then [("db.statement", text), ("resource.name", text)] + else [] + } + where + text = toAttribute statement + op = operationOf statement + -- One line, so a backend that groups by the statement shows the query + -- rather than the first line of its indentation. + statement = + cutTo (maxStatementLength cfg) + . T.unwords + . T.words + $ decodeUtf8Lenient sql + + +-- | The leading keyword. Anything that does not start with a bare word is not +-- worth guessing at. +operationOf :: Text -> Maybe Text +operationOf statement = case T.words statement of + word : _ | T.all Char.isAlpha word -> Just (T.toUpper word) + _ -> Nothing + + +cutTo :: Int -> Text -> Text +cutTo limit text + | T.length text > limit = T.take limit text <> "…" + | otherwise = text + + +-- | Attributes of the connection, read once. Connections are used by one +-- thread at a time, so the read-then-write needs no more than an 'IORef'. +sharedAttributes :: ConnectionTracer -> IO [(Text, Attribute)] +sharedAttributes ct = + readIORef (connectionAttributes ct) >>= \case + Just cached -> pure cached + Nothing -> do + fresh <- readConnectionAttributes (config ct) (connection ct) + writeIORef (connectionAttributes ct) (Just fresh) + pure fresh + + +readConnectionAttributes :: Config -> Pqi.Connection -> IO [(Text, Attribute)] +readConnectionAttributes cfg pq = do + name <- Pqi.db pq + host <- Pqi.host pq + port <- Pqi.port pq + pure $ + ("db.system.name", toAttribute ("postgresql" :: Text)) + : [("span.type", toAttribute ("sql" :: Text)) | datadogCompat cfg] + <> catMaybes + [ ("db.namespace",) . toAttribute . decodeUtf8Lenient <$> name + , ("server.address",) . toAttribute . decodeUtf8Lenient <$> host + , ("server.port",) . toAttribute <$> (readMaybe . BS8.unpack =<< port :: Maybe Int) + ] diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,203 @@ +-- | Runs real hasql sessions through a traced adapter against a live Postgres +-- and inspects the spans that come out. The connection is configured by the +-- usual libpq environment (@PGHOST@, @PGPORT@, @PGUSER@, @PGPASSWORD@, +-- @PGDATABASE@); without @PGHOST@ the suite is skipped. +module Main (main) where + +import Control.Exception (bracket) +import Control.Monad (forM_, void) +import Data.Foldable (for_) +import Data.IORef (IORef, readIORef, writeIORef) +import Data.Int (Int64) +import Data.List (sortOn) +import Data.Text (Text) +import Data.Text qualified as T +import Hasql.Connection qualified as Connection +import Hasql.Connection.Settings qualified as Settings +import Hasql.Decoders qualified as Decoders +import Hasql.Encoders qualified as Encoders +import Hasql.OpenTelemetry (Config (..), defaultConfig, tracedWith) +import Hasql.Session qualified as Session +import Hasql.Statement qualified as Statement +import OpenTelemetry.Attributes (Attribute (..), PrimitiveAttribute (..), lookupAttribute) +import OpenTelemetry.Context qualified as Context +import OpenTelemetry.Context.ThreadLocal qualified as ThreadLocal +import OpenTelemetry.Exporter.InMemory.Span (inMemoryListExporter) +import OpenTelemetry.Trace.Core qualified as Otel +import Pqi qualified +import Pqi.Ffi qualified +import Pqi.Native qualified +import System.Environment (lookupEnv) +import Test.Hspec + + +main :: IO () +main = + lookupEnv "PGHOST" >>= \case + Nothing -> putStrLn "PGHOST is not set, skipping the integration tests" + Just _ -> do + (processor, exported) <- inMemoryListExporter + provider <- Otel.createTracerProvider [processor] Otel.emptyTracerProviderOptions + Otel.setGlobalTracerProvider provider + hspec do + for_ [Pqi.Ffi.adapter, Pqi.Native.adapter] \adapter -> + describe (T.unpack (Pqi.name adapter)) (spec adapter exported) + + +-- | What a test needs to know about an exported span. +data Exported = Exported + { name :: Text + , attribute :: Text -> Maybe Attribute + , status :: Otel.SpanStatus + , start :: Otel.Timestamp + } + + +spec :: Pqi.Adapter -> IORef [Otel.ImmutableSpan] -> Spec +spec adapter exported = do + it "reports nothing outside a parent span" do + spans <- capture exported do + withConnection defaultConfig adapter \c -> run c (Session.statement 1 addOne) + map name spans `shouldBe` [] + + it "reports one span per command, with the statement on each" do + spans <- capture exported do + withConnection defaultConfig adapter \c -> underParent do + forM_ [1, 2, 3] \n -> run c (Session.statement n addOne) + -- The first execution prepares the statement; the rest reuse it. + map name spans `shouldBe` ["PREPARE SELECT", "SELECT", "SELECT", "SELECT"] + for_ spans \s -> do + text (attribute s "db.query.text") `shouldBe` Just "select $1 :: int8 + 1" + text (attribute s "db.system.name") `shouldBe` Just "postgresql" + attribute s "db.postgresql.results_unread" `shouldBe` Nothing + status s `shouldBe` Otel.Unset + int (attribute (spans !! 1) "db.response.returned_rows") `shouldBe` Just 1 + + it "records the SQLSTATE of a failure and none of its message" do + spans <- capture exported do + withConnection defaultConfig adapter \c -> underParent do + void (Connection.use c (Session.statement "not a uuid" castToUuid)) + run c (Session.statement 41 addOne) + let failed = filter (\s -> status s /= Otel.Unset) spans + map (text . (`attribute` "db.response.status_code")) failed `shouldBe` [Just "22P02"] + map status failed `shouldBe` [Otel.Error "postgres reported an error, SQLSTATE 22P02"] + -- The value that failed the cast must not appear anywhere. + for_ spans \s -> for_ ["db.query.text", "db.response.status_code", "error.type"] \key -> + fmap (T.isInfixOf "not a uuid") (text (attribute s key)) `shouldNotBe` Just True + -- And the queue is still aligned for the statement after the failure. + map name (filter (\s -> text (attribute s "db.query.text") == Just "select $1 :: int8 + 1") spans) + `shouldSatisfy` (\names -> "SELECT" `elem` names) + for_ spans \s -> attribute s "db.postgresql.results_unread" `shouldBe` Nothing + + it "shapes spans for Datadog when asked" do + spans <- capture exported do + withConnection defaultConfig{datadogCompat = True} adapter \c -> underParent do + run c (Session.statement 1 addOne) + map name spans `shouldBe` ["postgresql.prepare", "postgresql.query"] + for_ spans \s -> do + text (attribute s "span.type") `shouldBe` Just "sql" + text (attribute s "resource.name") `shouldBe` Just "select $1 :: int8 + 1" + text (attribute s "db.statement") `shouldBe` Just "select $1 :: int8 + 1" + + it "flattens and cuts long statements" do + spans <- capture exported do + withConnection defaultConfig{maxStatementLength = 10} adapter \c -> underParent do + run c (Session.script "select\n 1") + map (text . (`attribute` "db.query.text")) spans `shouldBe` [Just "select 1"] + spans' <- capture exported do + withConnection defaultConfig{maxStatementLength = 6} adapter \c -> underParent do + run c (Session.script "select\n 1") + map (text . (`attribute` "db.query.text")) spans' `shouldBe` [Just "select…"] + + +addOne :: Statement.Statement Int64 Int64 +addOne = + Statement.preparable + "select $1 :: int8 + 1" + (Encoders.param (Encoders.nonNullable Encoders.int8)) + (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int8))) + + +castToUuid :: Statement.Statement Text () +castToUuid = + Statement.preparable + "select $1 :: text :: uuid" + (Encoders.param (Encoders.nonNullable Encoders.text)) + Decoders.noResult + + +withConnection :: Config -> Pqi.Adapter -> (Connection.Connection -> IO a) -> IO a +withConnection config adapter = + bracket + do + settings <- settingsFromEnv + Connection.acquire (tracedWith config adapter) settings >>= either (fail . show) pure + Connection.release + + +-- | The libpq environment, spelled out because not every adapter reads it. +settingsFromEnv :: IO Settings.Settings +settingsFromEnv = do + host <- env "PGHOST" "localhost" + port <- env "PGPORT" "5432" + user <- env "PGUSER" "postgres" + password <- env "PGPASSWORD" "" + database <- env "PGDATABASE" "postgres" + pure $ + mconcat + [ Settings.hostAndPort host (read (T.unpack port)) + , Settings.user user + , Settings.password password + , Settings.dbname database + ] + where + env key fallback = maybe fallback T.pack <$> lookupEnv key + + +run :: Connection.Connection -> Session.Session a -> IO () +run c session = Connection.use c session >>= either (fail . show) (const (pure ())) + + +-- | Run under a fresh root span, as an application's request handler would. +underParent :: IO a -> IO a +underParent act = do + provider <- Otel.getGlobalTracerProvider + let tracer = Otel.makeTracer provider "test" Otel.tracerOptions + bracket + (Otel.createSpan tracer Context.empty "parent" Otel.defaultSpanArguments) + (\parent -> Otel.endSpan parent Nothing >> ThreadLocal.adjustContext (const Context.empty)) + \parent -> do + ThreadLocal.adjustContext (Context.insertSpan parent) + act + + +-- | The database spans an action reported, in the order they started. +capture :: IORef [Otel.ImmutableSpan] -> IO () -> IO [Exported] +capture exported act = do + writeIORef exported [] + act + spans <- readIORef exported + described <- traverse describeSpan spans + pure . sortOn start $ filter (\s -> name s /= "parent") described + where + describeSpan s = do + hot <- readIORef (Otel.spanHot s) + pure + Exported + { name = Otel.hotName hot + , attribute = lookupAttribute (Otel.hotAttributes hot) + , status = Otel.hotStatus hot + , start = Otel.spanStart s + } + + +text :: Maybe Attribute -> Maybe Text +text = \case + Just (AttributeValue (TextAttribute t)) -> Just t + _ -> Nothing + + +int :: Maybe Attribute -> Maybe Int64 +int = \case + Just (AttributeValue (IntAttribute i)) -> Just i + _ -> Nothing