Add ability to forward incoming heartbeats to another Wakatime compat…

…ible server.

fixes mujx#34

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Features
+
+- Add ability to forward incoming heartbeats to another Wakatime compatible server.
+
 ## [1.3.2] - 2021-04-11
 
 ### Bug fixes

README.md

@@ -14,6 +14,7 @@ It comes together with a dashboard which provides a graphical representation of
 
 - Import Wakatime activity using an API token and a range of dates.
 - See time spent on Github commits.
+- Forward incoming heartbeats to another Wakatime compatible server (e.g `wakatime.com`)
 - Group projects together with tags (e.g `#work`, `#personal`) and view their aggregated statistics.
 - User registration & login through the UI.
 - Leaderboards for all the users of the instance.
@@ -84,6 +85,10 @@ services:
       HAKA_ENV: "dev" # Use a json logger for production, otherwise key=value pairs.
       HAKA_HTTP_LOG: "true" # If you want to log http requests.
       GITHUB_TOKEN: "<token>" # If you want to retrieve time spent per commit. No extra scope is required.
+      # Add the following variables if you want to forward any received heartbeats to another
+      # Wakatime compatible server.
+      HAKA_REMOTE_WRITE_URL: "https://wakatime.com/api/v1/users/current/heartbeats.bulk"
+      HAKA_REMOTE_WRITE_TOKEN: "<token>"
     ports:
       - "127.0.0.1:8080:8080"
   haka_db:

app/Main.hs

@@ -12,6 +12,7 @@ import Haka.App
     AppM,
     LogState (..),
     RegistrationStatus (..),
+    RemoteWriteConfig (..),
     ServerSettings (..),
     runAppT,
   )
@@ -91,18 +92,23 @@ initApp settings unApp = do
     else do
       run (hakaPort settings) (unApp appCtx)
 
+emptyStr :: String
+emptyStr = ""
+
 getServerSettings :: IO ServerSettings
 getServerSettings = do
   p <- envAsInt "HAKA_PORT" 8080
   badgeUrl <- getEnvDefault "HAKA_BADGE_URL" "http://localhost:8080"
   dashboardPath <- envAsString "HAKA_DASHBOARD_PATH" "./dashboard/dist"
-  apiPrefix <- envAsString "HAKA_API_PREFIX" ""
+  apiPrefix <- envAsString "HAKA_API_PREFIX" emptyStr
   shieldsIOUrl <- envAsString "HAKA_SHIELDS_IO_URL" "https://img.shields.io"
   enableRegistration <- envAsBool "HAKA_ENABLE_REGISTRATION" True
   sessionExpiry <- envAsInt "HAKA_SESSION_EXPIRY" 24
   logLevel <- envAsString "HAKA_LOG_LEVEL" "info"
   rEnv <- envAsString "HAKA_ENV" "prod"
   enableHttpLog <- envAsBool "HAKA_HTTP_LOG" True
+  remoteWriteUrl <- envAsString "HAKA_REMOTE_WRITE_URL" emptyStr
+  remoteWriteToken <- envAsString "HAKA_REMOTE_WRITE_TOKEN" emptyStr
   when (sessionExpiry <= 0) (error "Session expiry should be a positive integer")
   return $
     ServerSettings
@@ -117,6 +123,15 @@ getServerSettings = do
           "prod" -> Log.Prod
           "production" -> Log.Prod
           _ -> Log.Dev,
+        hakaRemoteWriteConfig =
+          if remoteWriteUrl /= emptyStr && remoteWriteToken /= emptyStr
+            then
+              Just $
+                RemoteWriteConfig
+                  { heartbeatUrl = toText remoteWriteUrl,
+                    token = toText remoteWriteToken
+                  }
+            else Nothing,
         hakaEnableRegistration =
           if enableRegistration
             then EnabledRegistration

default.nix

@@ -2,7 +2,7 @@
 , case-insensitive, containers, contravariant-extras, cookie
 , cryptonite, fakedata, file-embed, hasql, hasql-pool, hasql-queue
 , hasql-transaction, http-client, http-client-tls, http-media
-, http-types, katip, lib, mr-env, optparse-applicative
+, http-types, katip, lib, modern-uri, mr-env, optparse-applicative
 , postgresql-binary, postgresql-simple, postgresql-simple-migration
 , random, raw-strings-qq, relude, req, safe-exceptions, servant
 , servant-client, servant-server, system-filepath, text, time, unix
@@ -19,7 +19,7 @@ mkDerivation {
     aeson base base64-bytestring bytestring case-insensitive containers
     contravariant-extras cookie cryptonite file-embed hasql hasql-pool
     hasql-queue hasql-transaction http-client http-client-tls
-    http-media http-types katip mr-env optparse-applicative
+    http-media http-types katip modern-uri mr-env optparse-applicative
     postgresql-binary postgresql-simple postgresql-simple-migration
     raw-strings-qq relude req safe-exceptions servant servant-server
     system-filepath text time unix uuid uuid-types vector wai wai-extra

hakatime.cabal

@@ -115,6 +115,7 @@ library
                      , postgresql-simple
                      , postgresql-simple-migration
                      , vector
+                     , modern-uri
 
   ghc-options:         -Wall
                        -Wcompat

src/Haka/App.hs

@@ -6,6 +6,7 @@
 module Haka.App
   ( AppCtx (..),
     LogState (..),
+    RemoteWriteConfig (..),
     AppM,
     runAppT,
     mkAppT,
@@ -124,6 +125,13 @@ data RegistrationStatus
   = EnabledRegistration
   | DisabledRegistration
 
+data RemoteWriteConfig = RemoteWriteConfig
+  { -- Wakatime compatible endpoint that can receive heartbeats.
+    heartbeatUrl :: Text,
+    -- Authentication token to use with the request.
+    token :: Text
+  }
+
 -- | Server configuration settings.
 data ServerSettings = ServerSettings
   { -- | Where the service will listen to.
@@ -146,5 +154,7 @@ data ServerSettings = ServerSettings
     -- | Verbosity level.
     hakaLogLevel :: String,
     -- | Whether to log the HTTP requests.
-    hakaHasHttpLogger :: Bool
+    hakaHasHttpLogger :: Bool,
+    -- | Configuration regarding remote write to a Wakatime compatible server.
+    hakaRemoteWriteConfig :: Maybe RemoteWriteConfig
   }

src/Haka/Handlers/Heartbeats.hs

@@ -8,21 +8,24 @@ module Haka.Handlers.Heartbeats
   )
 where
 
-import Control.Exception.Safe (throw, try)
+import Control.Exception.Safe (MonadThrow, catchAny, throw, try)
 import Data.Aeson (ToJSON)
+import Data.ByteString.Base64 (encode)
 import Data.Text (toUpper)
 import Data.Time.Calendar (Day)
 import Filesystem.Path (splitExtension)
 import Filesystem.Path.CurrentOS (fromText)
-import Haka.App (AppCtx (..), AppM)
+import Haka.App (AppCtx (..), AppM, RemoteWriteConfig (..), ServerSettings (..))
 import qualified Haka.Database as Db
 import Haka.Errors (HeartbeatApiResponse (..))
 import qualified Haka.Errors as Err
 import Haka.Types
 import Hasql.Pool (Pool)
 import Katip
+import qualified Network.HTTP.Req as R
 import qualified Relude.Unsafe as Unsafe
 import Servant
+import Text.URI (mkURI)
 
 data User = User
   { name :: Text,
@@ -113,10 +116,16 @@ heartbeatHandler ::
   HeartbeatPayload ->
   AppM HeartbeatApiResponse
 heartbeatHandler _ Nothing _ = throw Err.missingAuthError
-heartbeatHandler machineId (Just token) heartbeat = do
+heartbeatHandler machineId (Just tkn) heartbeat = do
   logFM InfoS "received a heartbeat"
   p <- asks pool
-  res <- storeHeartbeats p token machineId [heartbeat]
+
+  settings <- asks srvSettings
+
+  _ <- catchAny (remoteWriteHeartbeats (hakaRemoteWriteConfig settings) machineId [heartbeat]) print
+
+  res <- storeHeartbeats p tkn machineId [heartbeat]
+
   mkResponse res
 
 multiHeartbeatHandler ::
@@ -127,10 +136,15 @@ multiHeartbeatHandler ::
   [HeartbeatPayload] ->
   AppM HeartbeatApiResponse
 multiHeartbeatHandler _ Nothing _ = throw Err.missingAuthError
-multiHeartbeatHandler machineId (Just token) heartbeats = do
+multiHeartbeatHandler machineId (Just tkn) heartbeats = do
   logFM InfoS ("received " <> showLS (length heartbeats) <> " heartbeats")
   p <- asks pool
-  res <- storeHeartbeats p token machineId heartbeats
+  settings <- asks srvSettings
+
+  _ <- catchAny (remoteWriteHeartbeats (hakaRemoteWriteConfig settings) machineId heartbeats) print
+
+  res <- storeHeartbeats p tkn machineId heartbeats
+
   mkResponse res
 
 -- | Construct an API Heartbeat response depending on the size of the response.
@@ -174,9 +188,34 @@ storeHeartbeats ::
   Maybe Text ->
   [HeartbeatPayload] ->
   AppM (Either Db.DatabaseException [Int64])
-storeHeartbeats p token machineId heartbeats = do
+storeHeartbeats p tkn machineId heartbeats = do
   let updatedHeartbeats = map ((\beat -> beat {machine = machineId}) . addMissingLang) heartbeats
 
-  try $ liftIO $ Db.processHeartbeatRequest p token updatedHeartbeats
+  try $ liftIO $ Db.processHeartbeatRequest p tkn updatedHeartbeats
 
 -- TODO: Discard timestamps from the future
+
+remoteWriteHeartbeats ::
+  (R.MonadHttp m, MonadThrow m, KatipContext m) =>
+  Maybe RemoteWriteConfig ->
+  Maybe Text ->
+  [HeartbeatPayload] ->
+  m ()
+remoteWriteHeartbeats Nothing _ _ = pure ()
+remoteWriteHeartbeats (Just conf) machineHeader heartbeats = do
+  let machHeader = maybe mempty (R.header "X-Machine-Name" . encodeUtf8) machineHeader
+
+  remoteUrl <- mkURI $ heartbeatUrl conf
+
+  let header =
+        R.header "Authorization" ("Basic " <> encode (encodeUtf8 (token conf))) <> machHeader
+
+  logFM DebugS ("Sending data to " <> ls (heartbeatUrl conf))
+
+  case R.useHttpsURI remoteUrl of
+    Nothing -> logFM ErrorS "Malformed remote write URL was given"
+    Just (url, _) -> do
+      _ <- R.req R.POST url (R.ReqBodyJson heartbeats) R.ignoreResponse header
+      pure ()
+
+  pure ()