Merge branch 'native-sql-format'

2014-11-03 23:49:21 -08:00
parent 36f7ac3ff4 cc2c65a9ab
commit 8c6b34aaf2
6 changed files with 126 additions and 91 deletions
@@ -1,5 +1,5 @@
 name:                  dbapi
-version:               0.2.4.3
+version:               0.2.4.4
 synopsis:              The database is your api
 license:               MIT
 license-file:          LICENSE
@@ -26,6 +26,7 @@ executable dbapi
                     , containers, unordered-containers
                     , optparse-applicative >= 0.9.1 && < 0.10
                     , regex-base, regex-tdfa
+                     , regex-tdfa-text
                     , Ranged-sets
                     , transformers
                     , bcrypt, base64-string
@@ -45,7 +46,7 @@ Test-Suite spec
  ghc-options:         -Wall -W -Werror
  Main-Is:             Main.hs
  Other-Modules:       Dbapi, Spec, SpecHelper
-  Build-Depends:       base, hspec2
+  Build-Depends:       base, hspec2, QuickCheck
                     , hspec-wai >= 0.5.0, hspec-wai-json
                     , HDBC, HDBC-postgresql
                     , warp, wai >= 3.0.1 && < 3.0.2
@@ -60,6 +61,7 @@ Test-Suite spec
                     , regex-base
                     , string-conversions
                     , http-media, regex-tdfa
+                     , regex-tdfa-text
                     , Ranged-sets
                     , transformers
                     , bcrypt
@@ -16,7 +16,7 @@ import Data.Map (intersection, fromList, toList, Map)
 import Data.List (sort)
 import qualified Data.Set as S
 import Data.Convertible.Base (convert)
-import Data.Text (strip)
+import Data.Text (strip, Text)

 import Network.HTTP.Types.Status
 import Network.HTTP.Types.Header
@@ -125,7 +125,7 @@ app conn req respond =
    ([table], "POST") ->
      jsonBodyAction req (\row -> do
        allvals <- insert ver table row conn
-        keys <- primaryKeyColumns ver (cs table) conn
+        keys <- map cs <$> primaryKeyColumns ver (cs table) conn
        let params = urlEncodeVars $ map (\t -> (fst t, "eq." <> convert (snd t) :: String)) $ toList $ filterByKeys allvals keys
        return $ responseLBS status201
          [ jsonContentType
@@ -210,15 +210,15 @@ respondWithRangedResult rr =
      | (1 + to - from) < total = status206
      | otherwise               = status200

-requestedVersion :: RequestHeaders -> Maybe String
+requestedVersion :: RequestHeaders -> Maybe Text
 requestedVersion hdrs =
  case verStr of
       Just [[_, ver]] -> Just ver
       _ -> Nothing

  where verRegex = "version[ ]*=[ ]*([0-9]+)" :: String
-        accept = cs <$> lookup hAccept hdrs :: Maybe String
-        verStr = (=~ verRegex) <$> accept :: Maybe [[String]]
+        accept = cs <$> lookup hAccept hdrs :: Maybe Text
+        verStr = (=~ verRegex) <$> accept :: Maybe [[Text]]


 addHeaders :: ResponseHeaders -> Response -> Response
@@ -11,18 +11,23 @@ module PgQuery (
 , setRole
 , resetRole
 , checkPass
+, pgFmtIdent
+, pgFmtLit
 , RangedResult(..)
 , LoginAttempt(..)
 , DbRole
 ) where

-import Data.Text (Text, splitOn, intercalate)
+import Data.Text (Text, splitOn, intercalate, replace, takeWhile)
 import Data.String.Conversions (cs)
 import Data.Functor ( (<$>) )
 import Data.Maybe (fromMaybe, mapMaybe)
 import Data.Monoid ((<>), mconcat)
 import qualified Data.Map as M

+import Text.Regex.TDFA ((=~))
+import Text.Regex.TDFA.Text ()
+
 import Control.Monad (join)

 import qualified RangeQuery as R
@@ -47,8 +52,7 @@ data RangedResult = RangedResult {
 , rrBody  :: BL.ByteString
 } deriving (Show)

-type QuotedSql = (Text, [SqlValue])
-type Schema = String
+type Schema = Text
 type DbRole = BS.ByteString

 data LoginAttempt =
@@ -58,15 +62,8 @@ data LoginAttempt =
  | LoginSuccess DbRole
  deriving (Eq, Show)

-getRows :: Schema -> String -> Net.Query -> Maybe R.NonnegRange -> Connection -> IO RangedResult
+getRows :: Schema -> Text -> Net.Query -> Maybe R.NonnegRange -> Connection -> IO RangedResult
 getRows schema table qq range conn = do
-  query <- populateSql conn
-    $ globalAndLimitedCounts schema table qq <>
-      jsonArrayRows
-      (selectStarClause schema table
-        <> whereClause qq
-        <> orderClause qq
-        <> limitClause range)
  r <- quickQuery conn (cs query) []

  return $ case r of
@@ -78,26 +75,31 @@ getRows schema table qq range conn = do

  where
    offset = fromMaybe 0 $ R.offset <$> range
+    query  = globalAndLimitedCounts schema table qq <> jsonArrayRows (
+        selectStarClause schema table
+        <> whereClause qq
+        <> orderClause qq
+        <> limitClause range)


-whereClause :: Net.Query -> QuotedSql
+whereClause :: Net.Query -> Text
 whereClause qs =
-  if null qs then ("", []) else (" where ", []) <> conjunction
+  if null qs then "" else " where " <> conjunction

  where
    cols = [ col | col <- qs, fst col `notElem` ["order"] ]
-    conjunction = mconcat $ L.intersperse (" and ", []) (map wherePred cols)
+    conjunction = mconcat $ L.intersperse " and " (map wherePred cols)


-orderClause :: Net.Query -> QuotedSql
+orderClause :: Net.Query -> Text
 orderClause qs = do
  let order = fromMaybe "" $ join $ lookup "order" qs
      terms = mapMaybe parseOrderTerm $ splitOn "," $ cs order
-      termPred = mconcat $ L.intersperse (", ", []) (map orderTermSql terms)
+      termPred = mconcat $ L.intersperse ", " (map orderTermSql terms)

  if null terms
-     then ("", [])
-     else (" order by ", []) <> termPred
+     then ""
+     else " order by " <> termPred

  where
    parseOrderTerm :: Text -> Maybe OrderTerm
@@ -109,9 +111,8 @@ orderClause qs = do
                else Nothing
           _ -> Nothing

-    orderTermSql :: OrderTerm -> QuotedSql
-    orderTermSql t =
-      ("%I " <> otDirection t, [toSql $ otColumn t])
+    orderTermSql :: OrderTerm -> Text
+    orderTermSql t = pgFmtIdent (otColumn t) <> " " <> otDirection t


 data OrderTerm = OrderTerm {
@@ -120,9 +121,9 @@ data OrderTerm = OrderTerm {
 }


-wherePred :: Net.QueryItem -> QuotedSql
+wherePred :: Net.QueryItem -> Text
 wherePred (column, predicate) =
-  ("%I " <> op <> "%L", map toSql [column, value])
+  pgFmtIdent (cs column) <> " " <> op <> " " <> pgFmtLit (cs value)

  where
    opCode:rest = BS.split '.' $ fromMaybe "." predicate
@@ -136,37 +137,38 @@ wherePred (column, predicate) =
              "neq" -> "<>"
              _     -> "="

-limitClause :: Maybe R.NonnegRange -> QuotedSql
+limitClause :: Maybe R.NonnegRange -> Text
 limitClause range =
-  (" LIMIT %s OFFSET %s ", [toSql limit, toSql offset])
+  cs $ " LIMIT " <> limit <> " OFFSET " <> show offset <> " "

  where
    limit  = fromMaybe "ALL" $ show <$> (R.limit =<< range)
    offset = fromMaybe 0     $ R.offset <$> range

-globalAndLimitedCounts :: Schema -> String -> Net.Query -> QuotedSql
+globalAndLimitedCounts :: Schema -> Text -> Net.Query -> Text
 globalAndLimitedCounts schema table qq =
-  (" select ", [])
-  <> ("(select count(1) from %I.%I ", map toSql [schema, table])
+  " select "
+  <> "(select count(1) from " <> pgFmtIdent schema <> "." <> pgFmtIdent table <> " "
  <> whereClause qq
-  <> ("), count(t), ", [])
+  <> "), count(t), "

-selectStarClause :: Schema -> String -> QuotedSql
+selectStarClause :: Schema -> Text -> Text
 selectStarClause schema table =
-  (" select * from %I.%I ", map toSql [schema, table])
+  " select * from " <> pgFmtIdent schema <> "." <> pgFmtIdent table <> " "

-jsonArrayRows :: QuotedSql -> QuotedSql
+jsonArrayRows :: Text -> Text
 jsonArrayRows q =
-  ("array_to_json(array_agg(row_to_json(t))) from (", []) <> q <> (") t", [])
+  "array_to_json(array_agg(row_to_json(t))) from (" <> q <> ") t"

 insert :: Schema -> Text -> SqlRow -> Connection -> IO (M.Map String SqlValue)
 insert schema table row conn = do
-  sql    <- populateSql conn $ insertClause schema table row
  stmt   <- prepare conn $ cs sql
  _      <- execute stmt $ sqlRowValues row
  Just m <- fetchRowMap stmt
  return m

+  where sql = insertClause schema table row
+
 addUser :: BS.ByteString -> BS.ByteString -> BS.ByteString -> Connection -> IO ()
 addUser identity pass role conn = do
  Just hashed <- hashPasswordUsingPolicy fastBcryptHashingPolicy $ cs pass
@@ -190,53 +192,63 @@ checkPass = validatePassword

 upsert :: Schema -> Text -> SqlRow -> Net.Query -> Connection -> IO (M.Map String SqlValue)
 upsert schema table row qq conn = do
-  sql    <- populateSql conn $ upsertClause schema table row qq
  stmt   <- prepare conn $ cs sql
  _      <- execute stmt $ join $ replicate 2 $ sqlRowValues row
  Just m <- fetchRowMap stmt
  return m

+  where sql = upsertClause schema table row qq
+
 placeholders :: Text -> SqlRow -> Text
 placeholders symbol = intercalate ", " . map (const symbol) . getRow

-insertClause :: Schema -> Text -> SqlRow -> QuotedSql
+insertClause :: Schema -> Text -> SqlRow -> Text
 insertClause schema table (SqlRow []) =
-  ("insert into %I.%I default values returning *", [toSql schema, toSql table])
+  "insert into " <> pgFmtIdent schema <> "." <> pgFmtIdent table <> " default values returning *"
 insertClause schema table row =
-  ("insert into %I.%I (" <> placeholders "%I" row <> ")",
-  map toSql $ cs schema : table : sqlRowColumns row)
-  <> (" values (" <> placeholders "?" row <> ") returning *", sqlRowValues row)
+  "insert into " <> pgFmtIdent schema <> "." <> pgFmtIdent table <> " (" <>
+    intercalate ", " (map pgFmtIdent (sqlRowColumns row))
+  <> ") values (" <> placeholders "?" row <> ") returning *"

-
-insertClauseViaSelect :: Schema -> Text -> SqlRow -> QuotedSql
+insertClauseViaSelect :: Schema -> Text -> SqlRow -> Text
 insertClauseViaSelect schema table row =
-    ("insert into %I.%I (" <> placeholders "%I" row <> ")",
-     map toSql $ cs schema : table : sqlRowColumns row)
-  <> (" select " <> placeholders "?" row, sqlRowValues row)
+    "insert into " <> pgFmtIdent schema <> "." <> pgFmtIdent table <> " (" <>
+      intercalate ", " (map pgFmtIdent (sqlRowColumns row))
+  <> ") select " <> placeholders "?" row

-updateClause :: Schema -> Text -> SqlRow -> QuotedSql
+updateClause :: Schema -> Text -> SqlRow -> Text
 updateClause schema table row =
-    ("update %I.%I set (" <> placeholders "%I" row <> ")",
-     map toSql $ cs schema : table : sqlRowColumns row)
-  <> (" = (" <> placeholders "?" row <> ")", [])
+    "update " <> pgFmtIdent schema <> "." <> pgFmtIdent table <> " set (" <>
+      intercalate ", " (map pgFmtIdent (sqlRowColumns row))
+  <> ") = (" <> placeholders "?" row <> ")"

-upsertClause :: Schema -> Text -> SqlRow -> Net.Query -> QuotedSql
+upsertClause :: Schema -> Text -> SqlRow -> Net.Query -> Text
 upsertClause schema table row qq =
-  ("with upsert as (", []) <> updateClause schema table row
+  "with upsert as (" <> updateClause schema table row
  <> whereClause qq
-  <> (" returning *) ", []) <> insertClauseViaSelect schema table row
-  <> (" where not exists (select * from upsert) returning *", [])
+  <> " returning *) " <> insertClauseViaSelect schema table row
+  <> " where not exists (select * from upsert) returning *"

-populateSql :: Connection -> QuotedSql -> IO Text
-populateSql conn sql = do
-  [[escaped]] <- quickQuery conn (cs q) (snd sql)
-  return $ fromSql escaped
+pgFmtIdent :: Text -> Text
+pgFmtIdent x =
+  let escaped = replace "\"" "\"\"" (trimNullChars x) in
+  if escaped =~ danger
+    then "\"" <> escaped <> "\""
+    else escaped

-  where
-    q = mconcat [ "select format('", fst sql, "', ", ph (snd sql), ")" ]
+  where danger = "^$|^[^a-z_]|[^a-z_0-9]" :: Text

-    ph :: [a] -> Text
-    ph = intercalate ", " . map (const "?::varchar")
+pgFmtLit :: Text -> Text
+pgFmtLit x =
+  let trimmed = trimNullChars x
+      escaped = "'" <> replace "'" "''" trimmed <> "'"
+      slashed = replace "\\" "\\\\" escaped in
+  if escaped =~ ("\\\\" :: Text)
+    then "E" <> slashed
+    else slashed
+
+trimNullChars :: Text -> Text
+trimNullChars = Data.Text.takeWhile (/= '\x0')

 setRole :: Connection -> DbRole -> IO ()
 setRole conn role = runRaw conn $ "set role " <> cs role
@@ -4,11 +4,13 @@ module PgStructure where

 import Data.Functor ( (<$>) )
 import Data.Maybe (mapMaybe)
+import Data.Text hiding (foldl, map, zipWith, concat)
+import Data.Monoid ((<>))
+import Data.String.Conversions (cs)

 import Control.Applicative ( (<*>) )

 import qualified Data.ByteString.Lazy as BL
-import Data.List.Split (splitOn)

 import qualified Data.Aeson as JSON
 import qualified Data.Map as Map
@@ -19,8 +21,8 @@ import Database.HDBC.PostgreSQL
 import Data.Aeson ((.=))

 data Table = Table {
-  tableSchema :: String
-, tableName :: String
+  tableSchema :: Text
+, tableName :: Text
 , tableInsertable :: Bool
 } deriving (Show)

@@ -30,17 +32,17 @@ instance JSON.ToJSON Table where
    , "name"       .= tableName v
    , "insertable" .= tableInsertable v ]

-toBool :: String -> Bool
+toBool :: Text -> Bool
 toBool = (== "YES")

 data ForeignKey = ForeignKey {
-  fkTable::String, fkCol::String
+  fkTable::Text, fkCol::Text
 } deriving (Eq, Show)

 instance JSON.ToJSON ForeignKey where
  toJSON fk = JSON.object ["table".=fkTable fk, "column".=fkCol fk]

-foreignKeys :: String -> String -> Connection -> IO (Map.Map String ForeignKey)
+foreignKeys :: Text -> Text -> Connection -> IO (Map.Map Text ForeignKey)
 foreignKeys schema table conn = do
  r <- quickQuery conn
    "select kcu.column_name, ccu.table_name AS foreign_table_name,\
@@ -59,17 +61,17 @@ foreignKeys schema table conn = do
    addKey m _ = m --should never happen

 data Column = Column {
-  colSchema :: String
-, colTable :: String
-, colName :: String
+  colSchema :: Text
+, colTable :: Text
+, colName :: Text
 , colPosition :: Int
 , colNullable :: Bool
-, colType :: String
+, colType :: Text
 , colUpdatable :: Bool
 , colMaxLen :: Maybe Int
 , colPrecision :: Maybe Int
-, colDefault :: Maybe String
-, colEnum :: Maybe [String]
+, colDefault :: Maybe Text
+, colEnum :: Maybe [Text]
 , colFK :: Maybe ForeignKey
 } deriving (Show)

@@ -89,7 +91,7 @@ instance JSON.ToJSON Column where

 data TableOptions = TableOptions {
  tblOptcolumns :: [Column]
-, tblOptpkey :: [String]
+, tblOptpkey :: [Text]
 }

 instance JSON.ToJSON TableOptions where
@@ -97,7 +99,7 @@ instance JSON.ToJSON TableOptions where
      "columns" .= tblOptcolumns t
    , "pkey"    .= tblOptpkey t ]

-tables :: String -> Connection -> IO [Table]
+tables :: Text -> Connection -> IO [Table]
 tables s conn = do
  r <- quickQuery conn
        "select table_schema, table_name,\
@@ -114,7 +116,7 @@ tables s conn = do
        (toBool (fromSql insertable))
    mkTable _ = Nothing

-columns :: String -> String -> Connection -> IO [Column]
+columns :: Text -> Text -> Connection -> IO [Column]
 columns s t conn = do
  r <- quickQuery conn
    "select info.table_schema as schema, info.table_name as table_name, \
@@ -161,23 +163,23 @@ columns s t conn = do
        (fromSql maxlen)
        (fromSql precision)
        (fromSql defVal)
-        (splitOn "," <$> fromSql enum)
-    mkColumn _ =  error $ "Incomplete column data received for table " ++
-      t ++ " in schema " ++ s ++ "."
+        (Data.Text.splitOn "," <$> fromSql enum)
+    mkColumn _ =  error $ "Incomplete column data received for table " <>
+      cs t <> " in schema " <> cs s <> "."

-printTables :: String -> Connection -> IO BL.ByteString
+printTables :: Text -> Connection -> IO BL.ByteString
 printTables schema conn = JSON.encode <$> tables schema conn

-printColumns :: String -> String -> Connection -> IO BL.ByteString
+printColumns :: Text -> Text -> Connection -> IO BL.ByteString
 printColumns schema table conn =
  JSON.encode <$> (TableOptions <$> cols <*> pkey)
  where
    cols :: IO [Column]
    cols = columns schema table conn
-    pkey :: IO [String]
+    pkey :: IO [Text]
    pkey = primaryKeyColumns schema table conn

-primaryKeyColumns :: String -> String -> Connection -> IO [String]
+primaryKeyColumns :: Text -> Text -> Connection -> IO [Text]
 primaryKeyColumns s t conn = do
  r <- quickQuery conn
        "select kc.column_name \
@@ -3,11 +3,14 @@
 module Unit.PgQuerySpec where

 import Test.Hspec
+import Test.QuickCheck
+import Test.QuickCheck.Monadic

 import Database.HDBC (IConnection, SqlValue, toSql, prepare,
  quickQuery, fromSql, execute, seState, fetchAllRowsAL)

-import PgQuery (LoginAttempt(..), insert, addUser, signInRole, checkPass)
+import PgQuery (LoginAttempt(..), insert, addUser, signInRole, checkPass
+              , pgFmtIdent, pgFmtLit)
 import Types (SqlRow(SqlRow))
 import TestTypes (fromList, incStr, incNullableStr, incInsert, incId)
 import Data.Map (toList)
@@ -83,3 +86,17 @@ spec = around dbWithSchema $ do
    it "returns nothing with bad creds" $ \conn -> do
      signInRole "not-a-user" pass conn `shouldReturn` LoginFailed
      signInRole user (pass <> "crap") conn `shouldReturn` LoginFailed
+
+  describe "pgFmtIdent" $
+    it "Does what format %I would do" $ \conn ->
+      property $ monadicIO $ do
+        fuzz <- pick arbitrary
+        [[row]] <- run $ quickALQuery conn "select format('%I', ? :: varchar)" [toSql (fuzz :: String)]
+        assert $ fromSql (snd row) == pgFmtIdent (cs fuzz)
+
+  describe "pgFmtLit" $
+    it "Does what format %L would do" $ \conn ->
+      property $ monadicIO $ do
+        fuzz <- pick arbitrary
+        [[row]] <- run $ quickALQuery conn "select format('%L', ? :: varchar)" [toSql (fuzz :: String)]
+        assert $ fromSql (snd row) == pgFmtLit (cs fuzz)
@@ -1,3 +1,5 @@
+{-# LANGUAGE OverloadedStrings #-}
+
 module Unit.PgStructureSpec where

 import Test.Hspec