Add Lua Scripting support to nri-redis

nix/mk-shell.nix

@@ -6,11 +6,12 @@ let
     with pkgs.haskell.lib;
     overrideCabal hpkg (drv: { enableSeparateBinOutput = false; });
   # It is still necessary to run `hpack --force` into packages home dirs
-  haskell-language-server = pkgs.haskellPackages.haskell-language-server.override {
-    hls-ormolu-plugin = pkgs.haskellPackages.hls-ormolu-plugin.override {
-      ormolu = (workaround140774 pkgs.haskellPackages.ormolu);
+  haskell-language-server =
+    pkgs.haskellPackages.haskell-language-server.override {
+      hls-ormolu-plugin = pkgs.haskellPackages.hls-ormolu-plugin.override {
+        ormolu = (workaround140774 pkgs.haskellPackages.ormolu);
+      };
     };
-  };
 
 in pkgs.mkShell {
   buildInputs = [
@@ -60,6 +61,7 @@ in pkgs.mkShell {
         text
         text-zipper
         time
+        th-test-utils
         unordered-containers
         uuid
         vector

nri-postgresql/setup-postgres.sh

@@ -23,3 +23,7 @@ psql -c "CREATE TABLE test_table2 (enum_array_col test_enum[] NOT NULL)" || true
 
 ## Setup for test/Test.hs
 psql -c "CREATE TABLE constraints_table (user_id int PRIMARY KEY)" || true
+
+## Setup for test/ObservabilitySpec.hs
+createuser -s postgres
+psql -c "GRANT ALL PRIVILEGES ON DATABASE testdb TO postgres;" || true

nri-postgresql/test/ObservabilitySpec.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE CPP #-}
 {-# LANGUAGE QuasiQuotes #-}
 
 module ObservabilitySpec
@@ -42,7 +43,13 @@ tests postgres =
             )
             |> spanForTask
         Debug.toString span
-          |> Expect.equalToContentsOf "test/golden-results/observability-spec-postgres-reporting"
+          |> Expect.equalToContentsOf
+#if __GLASGOW_HASKELL__ >= 902
+            "test/golden-results/observability-spec-postgres-reporting-ghc-9"
+#else
+            "test/golden-results/observability-spec-postgres-reporting-ghc-8"
+#endif
+
     ]
 
 spanForTask :: Show e => Task e () -> Expect.Expectation' Platform.TracingSpan

nri-postgresql/test/golden-results/observability-spec-postgres-reporting → nri-postgresql/test/golden-results/observability-spec-postgres-reporting-ghc-8

@@ -9,9 +9,9 @@ TracingSpan
             { srcLocPackage = "main"
             , srcLocModule = "ObservabilitySpec"
             , srcLocFile = "test/ObservabilitySpec.hs"
-            , srcLocStartLine = 53
+            , srcLocStartLine = 60
             , srcLocStartCol = 7
-            , srcLocEndLine = 57
+            , srcLocEndLine = 64
             , srcLocEndCol = 40
             }
         )
@@ -32,9 +32,9 @@ TracingSpan
                     { srcLocPackage = "main"
                     , srcLocModule = "ObservabilitySpec"
                     , srcLocFile = "test/ObservabilitySpec.hs"
-                    , srcLocStartLine = 35
+                    , srcLocStartLine = 36
                     , srcLocStartCol = 11
-                    , srcLocEndLine = 42
+                    , srcLocEndLine = 43
                     , srcLocEndCol = 14
                     }
                 )

nri-postgresql/test/golden-results/observability-spec-postgres-reporting-ghc-9

@@ -0,0 +1,76 @@
+TracingSpan
+  { name = "test-root"
+  , started = MonotonicTime { inMicroseconds = 0 }
+  , finished = MonotonicTime { inMicroseconds = 0 }
+  , frame =
+      Just
+        ( "rootTracingSpanIO"
+        , SrcLoc
+            { srcLocPackage = "main"
+            , srcLocModule = "ObservabilitySpec"
+            , srcLocFile = "test/ObservabilitySpec.hs"
+            , srcLocStartLine = 60
+            , srcLocStartCol = 7
+            , srcLocEndLine = 60
+            , srcLocEndCol = 33
+            }
+        )
+  , details = Nothing
+  , summary = Nothing
+  , succeeded = Succeeded
+  , containsFailures = False
+  , allocated = 0
+  , children =
+      [ TracingSpan
+          { name = "Postgresql Query"
+          , started = MonotonicTime { inMicroseconds = 0 }
+          , finished = MonotonicTime { inMicroseconds = 0 }
+          , frame =
+              Just
+                ( "doQuery"
+                , SrcLoc
+                    { srcLocPackage = "main"
+                    , srcLocModule = "ObservabilitySpec"
+                    , srcLocFile = "test/ObservabilitySpec.hs"
+                    , srcLocStartLine = 36
+                    , srcLocStartCol = 11
+                    , srcLocEndLine = 36
+                    , srcLocEndCol = 27
+                    }
+                )
+          , details =
+              Just
+                "{\"query\":\"Secret *****\",\"query template\":\"!SELECT 1::bigint\",\"sql operation\":\"UNKNOWN\",\"queried relation\":\"!SELECT 1::bigint\",\"database type\":\"PostgreSQL\",\"host\":\"/mock/db/path.sock\",\"database\":\"mock-db-name\",\"rows returned\":1}"
+          , summary = Just "UNKNOWN !SELECT 1::bigint"
+          , succeeded = Succeeded
+          , containsFailures = False
+          , allocated = 0
+          , children =
+              [ TracingSpan
+                  { name = "acquiring Postgres connection from pool"
+                  , started = MonotonicTime { inMicroseconds = 0 }
+                  , finished = MonotonicTime { inMicroseconds = 0 }
+                  , frame =
+                      Just
+                        ( "withContext"
+                        , SrcLoc
+                            { srcLocPackage = "main"
+                            , srcLocModule = "Postgres"
+                            , srcLocFile = "src/Postgres.hs"
+                            , srcLocStartLine = 225
+                            , srcLocStartCol = 9
+                            , srcLocEndLine = 225
+                            , srcLocEndCol = 24
+                            }
+                        )
+                  , details = Just "{}"
+                  , summary = Just "acquiring Postgres connection from pool"
+                  , succeeded = Succeeded
+                  , containsFailures = False
+                  , allocated = 0
+                  , children = []
+                  }
+              ]
+          }
+      ]
+  }

nri-redis/nri-redis.cabal

@@ -40,6 +40,7 @@ library
       Redis.Codec
       Redis.Handler
       Redis.Internal
+      Redis.Script
       Redis.Settings
       Paths_nri_redis
   hs-source-dirs:
@@ -68,14 +69,18 @@ library
     , bytestring >=0.10.8.2 && <0.12
     , conduit >=1.3.0 && <1.4
     , containers >=0.6.0.1 && <0.7
+    , cryptohash-sha1 >=0.11.101.0 && <0.12
+    , haskell-src-meta >=0.8.12 && <0.9
     , hedis >=0.14.0 && <0.16
+    , megaparsec >=9.2.2 && <9.4
     , modern-uri >=0.3.1.0 && <0.4
     , nri-env-parser >=0.1.0.0 && <0.2
     , nri-observability >=0.1.0 && <0.2
     , nri-prelude >=0.1.0.0 && <0.7
     , pcre-light >=0.4.1.0 && <0.4.2
     , resourcet >=1.2.0 && <1.3
     , safe-exceptions >=0.1.7.0 && <1.3
+    , template-haskell >=2.16 && <3.0
     , text >=1.2.3.1 && <2.1
     , unordered-containers >=0.2.0.0 && <0.3
     , uuid >=1.3.0 && <1.4
@@ -87,6 +92,7 @@ test-suite tests
   other-modules:
       Helpers
       Spec.Redis
+      Spec.Redis.Script
       Spec.Settings
       NonEmptyDict
       Redis
@@ -96,6 +102,7 @@ test-suite tests
       Redis.Hash
       Redis.Internal
       Redis.List
+      Redis.Script
       Redis.Set
       Redis.Settings
       Redis.SortedSet
@@ -127,14 +134,18 @@ test-suite tests
     , bytestring >=0.10.8.2 && <0.12
     , conduit >=1.3.0 && <1.4
     , containers >=0.6.0.1 && <0.7
+    , cryptohash-sha1 >=0.11.101.0 && <0.12
+    , haskell-src-meta >=0.8.12 && <0.9
     , hedis >=0.14.0 && <0.16
+    , megaparsec >=9.2.2 && <9.4
     , modern-uri >=0.3.1.0 && <0.4
     , nri-env-parser >=0.1.0.0 && <0.2
     , nri-observability >=0.1.0 && <0.2
     , nri-prelude >=0.1.0.0 && <0.7
     , pcre-light >=0.4.1.0 && <0.4.2
     , resourcet >=1.2.0 && <1.3
     , safe-exceptions >=0.1.7.0 && <1.3
+    , template-haskell >=2.16 && <3.0
     , text >=1.2.3.1 && <2.1
     , unordered-containers >=0.2.0.0 && <0.3
     , uuid >=1.3.0 && <1.4

nri-redis/package.yaml

@@ -20,8 +20,11 @@ dependencies:
   - bytestring >= 0.10.8.2 && < 0.12
   - conduit >= 1.3.0 && < 1.4
   - containers >= 0.6.0.1 && < 0.7
+  - cryptohash-sha1 >= 0.11.101.0 && < 0.12
+  - haskell-src-meta >= 0.8.12 && < 0.9
   # hedis 14 introduces redis-cluster support
   - hedis >= 0.14.0 && < 0.16
+  - megaparsec >= 9.2.2 && < 9.4
   - modern-uri >= 0.3.1.0 && < 0.4
   - nri-env-parser >= 0.1.0.0 && < 0.2
   - nri-observability >= 0.1.0 && < 0.2
@@ -30,6 +33,7 @@ dependencies:
   - resourcet >= 1.2.0 && < 1.3
   - safe-exceptions >= 0.1.7.0 && < 1.3
   - text >= 1.2.3.1 && < 2.1
+  - template-haskell >= 2.16 && < 3.0
   - unordered-containers >=0.2.0.0 && <0.3
   - uuid >=1.3.0 && < 1.4
 library:

nri-redis/src/Redis.hs

@@ -48,6 +48,11 @@ module Redis
     Internal.map3,
     Internal.sequence,
     Internal.foldWithScan,
+
+    -- * Lua Scripting
+    script,
+    ScriptParam (..),
+    Internal.eval,
   )
 where
 
@@ -60,6 +65,7 @@ import qualified NonEmptyDict
 import qualified Redis.Codec as Codec
 import qualified Redis.Handler as Handler
 import qualified Redis.Internal as Internal
+import Redis.Script (ScriptParam (..), script)
 import qualified Redis.Settings as Settings
 import qualified Prelude
 

nri-redis/src/Redis/Handler.hs

@@ -18,6 +18,7 @@ import qualified Dict
 import qualified GHC.Stack as Stack
 import qualified Platform
 import qualified Redis.Internal as Internal
+import qualified Redis.Script as Script
 import qualified Redis.Settings as Settings
 import qualified Set
 import qualified Text
@@ -76,6 +77,9 @@ timeoutAfterMilliseconds milliseconds handler' =
           >> Task.timeout milliseconds Internal.TimeoutError,
       Internal.doTransaction =
         Stack.withFrozenCallStack (Internal.doTransaction handler')
+          >> Task.timeout milliseconds Internal.TimeoutError,
+      Internal.doEval =
+        Stack.withFrozenCallStack (Internal.doEval handler')
           >> Task.timeout milliseconds Internal.TimeoutError
     }
 
@@ -94,7 +98,10 @@ defaultExpiryKeysAfterSeconds secs handler' =
               |> Stack.withFrozenCallStack (Internal.doQuery handler'),
           Internal.doTransaction = \query' ->
             wrapWithExpire query'
-              |> Stack.withFrozenCallStack (Internal.doTransaction handler')
+              |> Stack.withFrozenCallStack (Internal.doTransaction handler'),
+          Internal.doEval = \script' ->
+            -- We can't guarantee auto-expire for EVAL, so we just run it as-is
+            Stack.withFrozenCallStack (Internal.doEval handler' script')
         }
 
 acquireHandler :: Text -> Settings.Settings -> IO (Internal.Handler' x, Connection)
@@ -131,6 +138,8 @@ acquireHandler namespace settings = do
                           Database.Redis.TxError err -> Right (Err (Internal.RedisError (Text.fromList err)))
                     )
                   |> Stack.withFrozenCallStack (platformRedis (Internal.cmds query) connection anything),
+          Internal.doEval = \script' ->
+            Stack.withFrozenCallStack (platformRedisScript script' connection anything),
           Internal.namespace = namespace,
           Internal.maxKeySize = Settings.maxKeySize settings
         },
@@ -364,15 +373,7 @@ platformRedis cmds connection anything action =
             Ok a -> a
             Err err -> Err err
       )
-    |> Exception.handle (\(_ :: Database.Redis.ConnectionLostException) -> pure <| Err Internal.ConnectionLost)
-    |> Exception.handleAny
-      ( \err ->
-          Exception.displayException err
-            |> Text.fromList
-            |> Internal.LibraryError
-            |> Err
-            |> pure
-      )
+    |> handleExceptions
     |> Platform.doAnything anything
     |> Stack.withFrozenCallStack Internal.traceQuery cmds (connectionHost connection) (connectionPort connection)
 
@@ -383,5 +384,70 @@ toResult reply =
     Left err -> Err (Internal.RedisError ("Redis library got back a value with a type it didn't expect: " ++ Text.fromList (Prelude.show err)))
     Right r -> Ok r
 
+handleExceptions :: IO (Result Internal.Error value) -> IO (Result Internal.Error value)
+handleExceptions =
+  Exception.handle (\(_ :: Database.Redis.ConnectionLostException) -> pure <| Err Internal.ConnectionLost)
+    >> Exception.handleAny
+      ( \err ->
+          Exception.displayException err
+            |> Text.fromList
+            |> Internal.LibraryError
+            |> Err
+            |> pure
+      )
+
+-- | Run a script in Redis trying to leverage the script cache
+platformRedisScript ::
+  (Stack.HasCallStack, Database.Redis.RedisResult a) =>
+  Script.Script a ->
+  Connection ->
+  Platform.DoAnythingHandler ->
+  Task Internal.Error a
+platformRedisScript script connection anything = do
+  -- Try EVALSHA
+  evalsha script connection anything
+    |> Task.onError
+      ( \err ->
+          case err of
+            Internal.RedisError "NOSCRIPT No matching script. Please use EVAL." -> do
+              -- If it fails with NOSCRIPT, load the script and try again
+              loadScript script connection anything
+              evalsha script connection anything
+            _ -> Task.fail err
+      )
+
+evalsha ::
+  (Stack.HasCallStack, Database.Redis.RedisResult a) =>
+  Script.Script a ->
+  Connection ->
+  Platform.DoAnythingHandler ->
+  Task Internal.Error a
+evalsha script connection anything =
+  Database.Redis.evalsha
+    (toB (Script.luaScriptHash script))
+    (map toB (Script.keys script))
+    (map toB (Log.unSecret (Script.arguments script)))
+    |> Database.Redis.runRedis (connectionHedis connection)
+    |> map toResult
+    |> handleExceptions
+    |> Platform.doAnything anything
+    |> Stack.withFrozenCallStack Internal.traceQuery [Script.evalShaString script] (connectionHost connection) (connectionPort connection)
+
+loadScript ::
+  Stack.HasCallStack =>
+  Script.Script a ->
+  Connection ->
+  Platform.DoAnythingHandler ->
+  Task Internal.Error ()
+loadScript script connection anything = do
+  Database.Redis.scriptLoad (toB (Script.luaScript script))
+    |> Database.Redis.runRedis (connectionHedis connection)
+    |> map toResult
+    |> handleExceptions
+    -- The result is the hash, which we already have. No sense in decoding it.
+    |> map (map (\_ -> ()))
+    |> Platform.doAnything anything
+    |> Stack.withFrozenCallStack Internal.traceQuery [Script.scriptLoadString script] (connectionHost connection) (connectionPort connection)
+
 toB :: Text -> Data.ByteString.ByteString
 toB = Data.Text.Encoding.encodeUtf8