Parse a larger set of CREATE FUNCTION statements

digitallyinduced · Oct 2, 2021 · 2298c29 · 2298c29
1 parent 76755e5
commit 2298c29
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 6 deletions.
diff --git a/IHP/IDE/SchemaDesigner/Compiler.hs b/IHP/IDE/SchemaDesigner/Compiler.hs
@@ -28,7 +28,7 @@ compileStatement CreateExtension { name, ifNotExists } = "CREATE EXTENSION " <>
 compileStatement AddConstraint { tableName, constraintName, constraint } = "ALTER TABLE " <> compileIdentifier tableName <> " ADD CONSTRAINT " <> compileIdentifier constraintName <> " " <> compileConstraint constraint <> ";"
 compileStatement Comment { content } = "-- " <> content
 compileStatement CreateIndex { indexName, unique, tableName, expressions, whereClause } = "CREATE" <> (if unique then " UNIQUE " else " ") <> "INDEX " <> indexName <> " ON " <> tableName <> " (" <> (intercalate ", " (map compileExpression expressions)) <> ")" <> (case whereClause of Just expression -> " WHERE " <> compileExpression expression; Nothing -> "") <> ";"
-compileStatement CreateFunction { functionName, functionBody, orReplace } = "CREATE " <> (if orReplace then "OR REPLACE " else "") <> "FUNCTION " <> functionName <> "() RETURNS TRIGGER AS $$" <> functionBody <> "$$ language plpgsql;"
+compileStatement CreateFunction { functionName, functionBody, orReplace, returns, language } = "CREATE " <> (if orReplace then "OR REPLACE " else "") <> "FUNCTION " <> functionName <> "() RETURNS " <> compilePostgresType returns <> " AS $$" <> functionBody <> "$$ language " <> language <> ";"
 compileStatement EnableRowLevelSecurity { tableName } = "ALTER TABLE " <> tableName <> " ENABLE ROW LEVEL SECURITY;"
 compileStatement UnknownStatement { raw } = raw <> ";"
 
@@ -124,6 +124,7 @@ compilePostgresType PJSONB = "JSONB"
 compilePostgresType PInet = "INET"
 compilePostgresType PTSVector = "TSVECTOR"
 compilePostgresType (PArray type_) = compilePostgresType type_ <> "[]"
+compilePostgresType PTrigger = "TRIGGER"
 compilePostgresType (PCustomType theType) = theType
 
 compileIdentifier :: Text -> Text

diff --git a/IHP/IDE/SchemaDesigner/Parser.hs b/IHP/IDE/SchemaDesigner/Parser.hs
@@ -207,6 +207,7 @@ sqlType = choice $ map optionalArray
         , jsonb
         , inet
         , tsvector
+        , trigger
         , customType
         ]
             where
@@ -341,6 +342,10 @@ sqlType = choice $ map optionalArray
                     arrayType <- typeParser;
                     (try do symbol' "[]"; pure $ PArray arrayType) <|> pure arrayType
 
+                trigger = do
+                    try (symbol' "TRIGGER")
+                    pure PTrigger
+
                 customType = do
                     theType <- try (takeWhile1P (Just "Custom type") (\c -> isAlphaNum c || c == '_'))
                     pure (PCustomType theType)
@@ -439,15 +444,15 @@ createFunction = do
     functionName <- identifier
     lexeme "()"
     lexeme "RETURNS"
-    lexeme "TRIGGER"
+    returns <- sqlType
     lexeme "AS"
     space
     functionBody <- cs <$> between (char '$' >> char '$') (char '$' >> char '$') (many (anySingleBut '$'))
     space
-    lexeme "language"
-    lexeme "plpgsql"
+    lexeme "language" <|> lexeme "LANGUAGE"
+    language <- symbol "plpgsql" <|> symbol "SQL"
     char ';'
-    pure CreateFunction { functionName, functionBody, orReplace }
+    pure CreateFunction { functionName, functionBody, orReplace, returns, language }
 
 -- | Triggers are not currently used by IHP, therefore they're implemented using UnknownStatement
 -- This avoid errors when having custom triggers in Schema.sql

diff --git a/IHP/IDE/SchemaDesigner/Types.hs b/IHP/IDE/SchemaDesigner/Types.hs
@@ -23,7 +23,7 @@ data Statement
     -- | CREATE UNIQUE INDEX name ON table (column [, ...]);
     | CreateIndex { indexName :: Text, unique :: Bool, tableName :: Text, expressions :: [Expression], whereClause :: Maybe Expression }
     -- | CREATE OR REPLACE FUNCTION functionName() RETURNS TRIGGER AS $$functionBody$$ language plpgsql;
-    | CreateFunction { functionName :: Text, functionBody :: Text, orReplace :: Bool }
+    | CreateFunction { functionName :: Text, functionBody :: Text, orReplace :: Bool, returns :: PostgresType, language :: Text }
     -- | ALTER TABLE tableName ENABLE ROW LEVEL SECURITY;
     | EnableRowLevelSecurity { tableName :: Text }
     deriving (Eq, Show)
@@ -130,5 +130,6 @@ data PostgresType
     | PInet
     | PTSVector
     | PArray PostgresType
+    | PTrigger
     | PCustomType Text
     deriving (Eq, Show)
diff --git a/Test/IDE/SchemaDesigner/CompilerSpec.hs b/Test/IDE/SchemaDesigner/CompilerSpec.hs
@@ -438,6 +438,8 @@ tests = do
                     { functionName = "notify_did_insert_webrtc_connection"
                     , functionBody = " BEGIN PERFORM pg_notify('did_insert_webrtc_connection', json_build_object('id', NEW.id, 'floor_id', NEW.floor_id, 'source_user_id', NEW.source_user_id, 'target_user_id', NEW.target_user_id)::text); RETURN NEW; END; "
                     , orReplace = True
+                    , returns = PTrigger
+                    , language = "plpgsql"
                     }
 
             compileSql [statement] `shouldBe` sql
@@ -449,6 +451,8 @@ tests = do
                     { functionName = "notify_did_insert_webrtc_connection"
                     , functionBody = " BEGIN PERFORM pg_notify('did_insert_webrtc_connection', json_build_object('id', NEW.id, 'floor_id', NEW.floor_id, 'source_user_id', NEW.source_user_id, 'target_user_id', NEW.target_user_id)::text); RETURN NEW; END; "
                     , orReplace = False
+                    , returns = PTrigger
+                    , language = "plpgsql"
                     }
 
             compileSql [statement] `shouldBe` sql

diff --git a/Test/IDE/SchemaDesigner/ParserSpec.hs b/Test/IDE/SchemaDesigner/ParserSpec.hs
@@ -444,13 +444,17 @@ tests = do
                     { functionName = "notify_did_insert_webrtc_connection"
                     , functionBody = " BEGIN PERFORM pg_notify('did_insert_webrtc_connection', json_build_object('id', NEW.id, 'floor_id', NEW.floor_id, 'source_user_id', NEW.source_user_id, 'target_user_id', NEW.target_user_id)::text); RETURN NEW; END; "
                     , orReplace = True
+                    , returns = PTrigger
+                    , language = "plpgsql"
                     }
 
         it "should parse a CREATE FUNCTION ..() RETURNS TRIGGER .." do
             parseSql "CREATE FUNCTION notify_did_insert_webrtc_connection() RETURNS TRIGGER AS $$ BEGIN PERFORM pg_notify('did_insert_webrtc_connection', json_build_object('id', NEW.id, 'floor_id', NEW.floor_id, 'source_user_id', NEW.source_user_id, 'target_user_id', NEW.target_user_id)::text); RETURN NEW; END; $$ language plpgsql;" `shouldBe` CreateFunction
                     { functionName = "notify_did_insert_webrtc_connection"
                     , functionBody = " BEGIN PERFORM pg_notify('did_insert_webrtc_connection', json_build_object('id', NEW.id, 'floor_id', NEW.floor_id, 'source_user_id', NEW.source_user_id, 'target_user_id', NEW.target_user_id)::text); RETURN NEW; END; "
                     , orReplace = False
+                    , returns = PTrigger
+                    , language = "plpgsql"
                     }
 
         it "should parse unsupported SQL as a unknown statement" do