From 044e3865acf1adc8e3ee916db04b56ef5a124521 Mon Sep 17 00:00:00 2001
From: Diogo Biazus <diogob@gmail.com>
Date: Tue, 20 Oct 2015 01:03:11 -0400
Subject: [PATCH] Removes 'Prefer: return=jwt' header and chooses jwt return
 based on function type

---
 src/PostgREST/App.hs         |  3 ++-
 src/PostgREST/PgStructure.hs | 22 ++++++++++++++++++----
 2 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
index be2f366ea..50106fac2 100644
--- a/src/PostgREST/App.hs
+++ b/src/PostgREST/App.hs
@@ -175,9 +175,10 @@ app dbstructure conf reqBody req =
           let call = B.Stmt "select " V.empty True <>
                 asJson (callProc qi $ fromMaybe M.empty (decode reqBody))
           bodyJson :: Maybe (Identity Value) <- H.maybeEx call
+          returnJWT <- doesProcReturnJWT schema proc
           return $ responseLBS status200 [jsonH]
                  (let body = fromMaybe emptyArray $ runIdentity <$> bodyJson in
-                    if hasPrefer "return=jwt"
+                    if returnJWT
                     then "{\"token\":\"" <> cs (tokenJWT jwtSecret body) <> "\"}"
                     else cs $ encode body)
         else return $ responseLBS status404 [] ""
diff --git a/src/PostgREST/PgStructure.hs b/src/PostgREST/PgStructure.hs
index 8ed08873b..a33d7829d 100644
--- a/src/PostgREST/PgStructure.hs
+++ b/src/PostgREST/PgStructure.hs
@@ -13,25 +13,39 @@ import           Data.Monoid
 import           Data.Text             (Text, split)
 import qualified Hasql                 as H
 import qualified Hasql.Postgres        as P
+import qualified Hasql.Backend         as B
 import           PostgREST.PgQuery     ()
 import           PostgREST.Types
 
 import           GHC.Exts              (groupWith)
 import           Prelude
 
+doesProc :: forall c s. B.CxValue c Int =>
+            (Text -> Text -> B.Stmt c) -> Text -> Text -> H.Tx c s Bool
+doesProc stmt schema proc = do
+  row :: Maybe (Identity Int) <- H.maybeEx $ stmt schema proc
+  return $ isJust row
 
 doesProcExist :: Text -> Text -> H.Tx P.Postgres s Bool
-doesProcExist schema proc = do
-  row :: Maybe (Identity Int) <- H.maybeEx $ [H.stmt|
+doesProcExist = doesProc [H.stmt|
       SELECT 1
       FROM   pg_catalog.pg_namespace n
       JOIN   pg_catalog.pg_proc p
       ON     pronamespace = n.oid
       WHERE  nspname = ?
       AND    proname = ?
-    |] schema proc
-  return $ isJust row
+    |]
 
+doesProcReturnJWT :: Text -> Text -> H.Tx P.Postgres s Bool
+doesProcReturnJWT = doesProc [H.stmt|
+      SELECT 1
+      FROM   pg_catalog.pg_namespace n
+      JOIN   pg_catalog.pg_proc p
+      ON     pronamespace = n.oid
+      WHERE  nspname = ?
+      AND    proname = ?
+      AND    pg_catalog.pg_get_function_result(p.oid) = 'jwt'
+    |]
 
 tableFromRow :: (Text, Text, Bool, Maybe Text) -> Table
 tableFromRow (s, n, i, a) = Table s n i (parseAcl a)