Create TemplateHaskell API for writing Lua Scripts

nri-redis/nri-redis.cabal

@@ -69,6 +69,7 @@ library
     , bytestring >=0.10.8.2 && <0.12
     , conduit >=1.3.0 && <1.4
     , containers >=0.6.0.1 && <0.7
+    , haskell-src-meta >=0.8.12 && <0.9
     , hedis >=0.14.0 && <0.16
     , megaparsec >=9.2.2 && <9.3
     , modern-uri >=0.3.1.0 && <0.4
@@ -132,6 +133,7 @@ test-suite tests
     , bytestring >=0.10.8.2 && <0.12
     , conduit >=1.3.0 && <1.4
     , containers >=0.6.0.1 && <0.7
+    , haskell-src-meta >=0.8.12 && <0.9
     , hedis >=0.14.0 && <0.16
     , megaparsec >=9.2.2 && <9.3
     , modern-uri >=0.3.1.0 && <0.4

nri-redis/package.yaml

@@ -20,6 +20,7 @@ dependencies:
   - bytestring >= 0.10.8.2 && < 0.12
   - conduit >= 1.3.0 && < 1.4
   - containers >= 0.6.0.1 && < 0.7
+  - 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.3

nri-redis/src/Redis/Script.hs

@@ -1,6 +1,14 @@
-module Redis.Script (Script (..), script, evalString, mapKeys, keysTouchedByScript, paramNames, paramValues, parser, Tokens (..)) where
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE IncoherentInstances #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE UndecidableInstances #-}
 
+module Redis.Script (Script (..), script, evalString, mapKeys, keysTouchedByScript, paramNames, paramValues, parser, Tokens (..), ScriptParam (..), printScript) where
+
+import Data.Either (Either (..))
 import Data.Void (Void)
+import qualified GHC.TypeLits
+import Language.Haskell.Meta.Parse (parseExp)
 import qualified Language.Haskell.TH as TH
 import qualified Language.Haskell.TH.Quote as QQ
 import qualified Set
@@ -11,16 +19,44 @@ import Prelude (notElem, pure, (<*))
 import qualified Prelude
 
 data Script result = Script
-  { -- | The Lua script to be executed
+  { -- | The Lua script to be executed with @args placeholders for Redis
     luaScript :: Text,
-    -- | The parameters that fill the placeholders in this query
-    params :: Log.Secret [Param],
     -- | The script string as extracted from a `script` quasi quote.
-    quasiQuotedString :: Text
+    quasiQuotedString :: Text,
+    -- | The parameters that fill the placeholders in this query
+    params :: Log.Secret [EvaluatedParam]
   }
   deriving (Eq, Show)
 
-data Param = Param
+-- | A type for enforcing parameters used in [script|${ ... }|] are either tagged as Key or Literal.
+--
+-- We need keys to be tagged, otherwise we can't implement `mapKeys` and enforce namespacing
+-- in Redis APIs.
+--
+-- We make this extra generic to allow us to provide nice error messages using TypeError in a
+-- type class below.
+data ScriptParam
+  = forall a. (Show a) => Key a
+  | forall a. (Show a) => Literal a
+
+class HasScriptParam a where
+  getScriptParam :: a -> ScriptParam
+
+instance HasScriptParam ScriptParam where
+  getScriptParam = Prelude.id
+
+-- | This instance is used to provide a helpful error message when a user tries to use a type
+-- other than a ScriptParam in a [script|${ ... }|] quasi quote.
+--
+-- It is what forces us to have IncoherentInstances and UndecidedInstances enabled.
+instance
+  {-# OVERLAPPABLE #-}
+  GHC.TypeLits.TypeError (GHC.TypeLits.Text "[script| ${..} ] interpolation only supports Key or Literal inputs.") =>
+  HasScriptParam x
+  where
+  getScriptParam = Prelude.error "This won't ever hit bc this generates a compile-time error."
+
+data EvaluatedParam = EvaluatedParam
   { kind :: ParamKind,
     name :: Text,
     value :: Text
@@ -39,15 +75,54 @@ script =
       QQ.quoteDec = Prelude.error "script not supported in declarations"
     }
 
-qqScript :: Prelude.String -> TH.ExpQ
+qqScript :: Prelude.String -> TH.Q TH.Exp
 qqScript scriptWithVars = do
   let str = Text.fromList scriptWithVars
-  let _expr = P.parse parser "" str
-  -- let parsedScript = case expr str of
-  --       Left err -> Prelude.error <| "Failed to parse script: " ++ err
-  --       Right parsed ->
-  -- [|parsedScript|]
-  Debug.todo "qqScript"
+  let parseResult = P.parse parser "" str
+  case parseResult of
+    Left err -> Prelude.error <| "Failed to parse script: " ++ P.errorBundlePretty err
+    Right tokens -> do
+      let luaScript' = tokensToScript tokens
+      let paramsExp =
+            tokens
+              |> List.filterMap
+                ( \t ->
+                    case t of
+                      ScriptVariable name -> Just name
+                      _ -> Nothing
+                )
+              -- Parse haskell syntax between ${} and convert to a TH Exp
+              |> List.indexedMap (\idx exp -> varToExp exp |> scriptParamExpression (Prelude.fromIntegral idx))
+              |> TH.ListE
+              |> TH.AppE (TH.VarE 'Log.mkSecret)
+      scriptConstructor <- [|Script luaScript' str|]
+      pure <| scriptConstructor `TH.AppE` paramsExp
+
+scriptParamExpression :: Prelude.Integer -> TH.Exp -> TH.Exp
+scriptParamExpression idx exp =
+  (TH.VarE 'evaluateScriptParam) `TH.AppE` TH.LitE (TH.IntegerL idx) `TH.AppE` exp
+
+evaluateScriptParam :: HasScriptParam a => Int -> a -> EvaluatedParam
+evaluateScriptParam idx scriptParam =
+  case getScriptParam scriptParam of
+    Key a ->
+      EvaluatedParam
+        { kind = RedisKey,
+          name = "arg" ++ Text.fromInt idx,
+          value = Debug.toString a
+        }
+    Literal a ->
+      EvaluatedParam
+        { kind = ArbitraryValue,
+          name = "arg" ++ Text.fromInt idx,
+          value = Debug.toString a
+        }
+
+varToExp :: Text -> TH.Exp
+varToExp var =
+  case parseExp (Text.toList var) of
+    Left err -> Prelude.error <| "Failed to parse variable: " ++ err
+    Right exp -> exp
 
 data Tokens
   = ScriptText Text
@@ -74,6 +149,17 @@ parseVariable = do
   _ <- PC.char '}'
   pure <| ScriptVariable <| Text.trim name
 
+tokensToScript :: List Tokens -> Text
+tokensToScript tokens =
+  tokens
+    |> List.indexedMap
+      ( \idx t ->
+          case t of
+            ScriptText text -> text
+            ScriptVariable _ -> "@arg" ++ Text.fromInt idx
+      )
+    |> Text.join ""
+
 -- | EVAL script numkeys [key [key ...]] [arg [arg ...]]
 evalString :: Script a -> Text
 evalString = Debug.todo "evalString"
@@ -101,3 +187,8 @@ paramValues script' =
     |> params
     |> Log.unSecret
     |> List.map value
+
+printScript :: Script a -> Text
+printScript Script {luaScript, quasiQuotedString, params} =
+  let printableParams = Log.unSecret params
+  in "Script { luaScript = \"" ++ luaScript ++ "\", quasiQuotedString = \"" ++ quasiQuotedString ++ "\", params = " ++ Debug.toString printableParams ++ " }"

nri-redis/test/Spec/Redis/Script.hs

@@ -1,3 +1,5 @@
+{-# LANGUAGE QuasiQuotes #-}
+
 module Spec.Redis.Script (tests) where
 
 import Data.Either (Either (..))
@@ -10,7 +12,8 @@ tests :: Test.Test
 tests =
   Test.describe
     "Redis.Script"
-    [ Test.describe "parser" parserTests
+    [ Test.describe "parser" parserTests,
+      Test.describe "th tests" thTests
     ]
 
 parserTests :: List Test.Test
@@ -104,6 +107,23 @@ parserTests =
           )
   ]
 
+thTests :: List Test.Test
+thTests =
+  [ Test.test "just text" <| \_ ->
+      [script|some text|]
+        |> printScript
+        |> Expect.equal "Script { luaScript = \"some text\", quasiQuotedString = \"some text\", params = [] }",
+    Test.test "one key argument" <| \_ ->
+      [script|${Key "hi"}|]
+        |> printScript
+        |> Expect.equal "Script { luaScript = \"@arg0\", quasiQuotedString = \"${Key \"hi\"}\", params = [ EvaluatedParam\n    { kind = RedisKey , name = \"arg0\" , value = \"\\\"hi\\\"\" }\n] }"
+    -- We can't test for compile-time errors, but manually test our helpful error message, uncomment
+    -- the lines below:
+    -- Test.test "compilation error" <| \_ ->
+    --   [script|${123}|]
+    --     |> Expect.equal "Doesn't matter, this won't compile"
+  ]
+
 mapLeft :: (a -> c) -> Either a b -> Either c b
 mapLeft f (Left a) = Left (f a)
 mapLeft _ (Right b) = Right b