diff --git a/LICENSE b/LICENSE new file mode 100644 --- /dev/null +++ b/LICENSE @@ -0,0 +1,13 @@ +Copyright (c) 2015 Dropbox, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. diff --git a/Setup.hs b/Setup.hs new file mode 100644 --- /dev/null +++ b/Setup.hs @@ -0,0 +1,2 @@ +import Distribution.Simple +main = defaultMain diff --git a/presto-hdbc.cabal b/presto-hdbc.cabal new file mode 100644 --- /dev/null +++ b/presto-hdbc.cabal @@ -0,0 +1,84 @@ + +-- Initial prestotalk.cabal generated by cabal init. For further +-- documentation, see http://haskell.org/cabal/users-guide/ + +-- The name of the package. +name: presto-hdbc + +-- The package version. See the Haskell package versioning policy (PVP) +-- for standards guiding when and how versions should be incremented. +-- http://www.haskell.org/haskellwiki/Package_versioning_policy +-- PVP summary: +-+------- breaking API changes +-- | | +----- non-breaking API additions +-- | | | +--- code changes with no API change +version: 0.1.0.0 + +-- A short (one-line) description of the package. +synopsis: An HDBC connector for Presto + +-- A longer description of the package. +description: Uses Presto's REST API to provide an HDBC interface. + +-- The license under which the package is released. +license: Apache-2.0 +license-file: LICENSE + +-- The package author(s). +author: Tom McLaughlin + +-- An email address to which users can send suggestions, bug reports, and +-- patches. +maintainer: tomm@dropbox.com + +-- A copyright notice. +copyright: Copyright (c) 2015 Dropbox, Inc. + +category: Database + +build-type: Simple + +-- Extra files to be distributed with the package, such as examples or a +-- README. +-- extra-source-files: + +-- Constraint on the version of Cabal needed to build this package. +cabal-version: >=1.10 + + +source-repository head + type: git + location: https://github.com/dropbox/presto-hdbc + + +library + -- Modules exported by the library. + exposed-modules: Database.HDBC.Presto + + -- Modules included in this library but not exported. + -- other-modules: + + -- LANGUAGE extensions used by modules in this package. + -- other-extensions: + + -- Other library packages from which modules are imported. + build-depends: HDBC + , HTTP >= 4000.2.18 + , aeson >= 0.8.0.0 + , aeson-pretty >= 0.7.1 + , base >=4.7 && <4.8 + , bytestring >= 0.10.4.0 + , either + , either-unwrap >= 1.1 + , errors >= 1.4.7 + , lens + , network-uri >= 2.6.0.1 + , scientific >= 0.3.3.0 + , text >= 1.1.1.3 + , transformers >= 0.4.1.0 + , mtl, convertible, safe, http-streams, io-streams + + -- Directories containing source files. + hs-source-dirs: src + + -- Base language which the package is written in. + default-language: Haskell2010 diff --git a/src/Database/HDBC/Presto.hs b/src/Database/HDBC/Presto.hs new file mode 100644 --- /dev/null +++ b/src/Database/HDBC/Presto.hs @@ -0,0 +1,264 @@ +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE OverloadedStrings #-} +module Database.HDBC.Presto (PrestoConnection(..), connectToPresto) where + +import Control.Applicative ((<$>)) +import Control.Concurrent (threadDelay) +import Control.Exception (throw) +import Control.Monad (when) +import Control.Monad.Trans +import Control.Monad.Trans.Either +import Data.Aeson +import qualified Data.ByteString as BN +import qualified Data.ByteString.Lazy as B +import Data.Convertible +import Data.Maybe +import Database.HDBC.PrestoData +import GHC.Word (Word16) +import Network.URI +import Safe (readMay) + +-- TODO: remove +import Data.Either.Unwrap +import qualified Data.Text.Encoding as DE +import qualified Data.Text.Lazy.Encoding as LDE +import Data.Text.Lazy as T +import Data.Text.Lazy.Encoding (encodeUtf8) +import System.IO.Error + +-- HDBC stuff +import Control.Concurrent.MVar(MVar, readMVar, modifyMVar, modifyMVar_, newMVar, withMVar, isEmptyMVar) +import Database.HDBC (IConnection(..), SqlValue, fetchAllRows, SqlError(..)) +import Database.HDBC.ColTypes as CT +import Database.HDBC.SqlValue +import Database.HDBC.Statement (Statement(..)) +import Network.Http.Client +import qualified System.IO.Streams as Streams + +-- TODO: move to utility library or find in real library +maybeToEither = flip maybe Right . Left +maybeToEitherT x y = hoistEither $ maybeToEither x y + +instance IConnection PrestoConnection where + disconnect = disconnectPresto + commit _ = fail "Not implemented" + rollback _ = fail "Not implemented" + run _ _ _ = fail "Not implemented" + prepare = preparePresto + clone _ = fail "Not implemented" + hdbcDriverName _ = "presto" + hdbcClientVer _ = fail "Not implemented" + proxiedClientName _ = fail "Not implemented" + proxiedClientVer _ = fail "Not implemented" + dbServerVer _ = fail "Not implemented" + dbTransactionSupport _ = False + getTables conn = do + showTables <- prepare conn "show tables" + _ <- execute showTables [] + rows <- fetchAllRows showTables + -- Each row should be a single string + return $ fmap ((\(SqlString s) -> s) . Prelude.head) rows + describeTable conn name = do + statement <- prepare conn $ "describe " ++ name + execute statement [] + rows <- fetchAllRows statement + + -- Presto's describe table output is of the following form: + -- Column (varchar) + -- Type (varchar) + -- Null (boolean) + -- Partition Key (boolean) + -- Comment (boolean) + + -- Pack the data into a SqlColDesc + return $ (flip fmap) rows (\((SqlString name): + (SqlString typ): + (SqlBool nullable): + (SqlBool partition): + (SqlString comment): + xs) -> + (name, SqlColDesc { + colType = prestoTypeToSqlColType typ, + colSize = Just $ if partition then 1 else 0, -- Hide the partition info here + colOctetLength = Nothing, + colDecDigits = Nothing, + colNullable = Just nullable + })) + + +data StmtState = Prepared + | Executed [PrestoColumn] [PrestoRow] + | Errored String + | Finished + +columnsFromStatement :: PrestoStatement -> [PrestoColumn] +columnsFromStatement statement = case statement of + PrestoStatement (Executed cols rows) _ _ -> cols + _ -> fail "You have to call execute before getting column names" + +type Catalog = T.Text +type Schema = T.Text + +-- Store the uri in here for now +data PrestoConnection = PrestoConnection URI Connection Catalog Schema + +data PrestoStatement = PrestoStatement StmtState PrestoConnection Text + +connectToPresto :: URI -> T.Text -> T.Text -> IO (Either Text PrestoConnection) +connectToPresto uri catalog schema = runEitherT $ do + auth <- maybeToEitherT "URI didn't contain authority" $ uriAuthority uri + port <- maybeToEitherT "Couldn't read port from URI" $ readMay $ Prelude.tail $ uriPort auth + + eitherConn <- liftIO $ tryIOError $ openConnection "localhost" port + conn <- hoistEither $ mapLeft (convert . show) eitherConn + return $ PrestoConnection uri conn catalog schema + + +fetchPrestoRow :: PrestoConnection -> MVar PrestoStatement -> IO (Maybe [SqlValue]) +fetchPrestoRow conn prestoStatementVar = modifyMVar prestoStatementVar fetch where + fetch (PrestoStatement Prepared _ _) = fail "Trying to fetch before executing statement." + + fetch statement@(PrestoStatement (Executed _ []) conn query) = do + let newStatement = PrestoStatement Finished conn query + return (newStatement, Nothing) + + fetch statement@(PrestoStatement (Executed cols rows@(row:remainder)) conn query) = do + let newStatement = PrestoStatement (Executed cols remainder) conn query + return (newStatement, Just row) + + fetch (PrestoStatement (Errored msg) _ _) = throw $ SqlError "Query failed" 0 msg + + fetch (PrestoStatement Finished _ _ ) = throw $ SqlError "Can't execute an already executed statement." 0 "" + +executePresto conn prestoStatementVar vals = modifyMVar prestoStatementVar exec where + exec stmt@(PrestoStatement (Executed _ _) _ _) = return (stmt, 0) + exec (PrestoStatement Prepared (PrestoConnection uri httpConn _ _) query) = do + result <- doPollingPrestoQuery httpConn conn uri query + let newStatement = case result of + Left msg -> throw $ SqlError "Query failed" 0 msg + Right (columns, rows) -> PrestoStatement (Executed columns rows) conn query + return (newStatement, 0) + + +preparePresto :: PrestoConnection -> String -> IO Statement +preparePresto conn query = do + prestoStatementVar <- newMVar $ PrestoStatement Prepared conn (T.pack query) + + return Statement { + executeRaw = undefined + , execute = executePresto conn prestoStatementVar + , executeMany = \_ -> fail "Not implemented" + , finish = fail "Not implemented" + , fetchRow = fetchPrestoRow conn prestoStatementVar + , originalQuery = query + , getColumnNames = do + prestoStatement <- readMVar prestoStatementVar + let columns = columnsFromStatement prestoStatement + return $ fmap prestoColumn_name columns + , describeResult = do + prestoStatement <- readMVar prestoStatementVar + let columns = columnsFromStatement prestoStatement + + return $ [(prestoColumn_name x, SqlColDesc { + colType = prestoTypeToSqlColType $ prestoColumn_type x, + colSize = Nothing, + colOctetLength = Nothing, + colDecDigits = Nothing, + colNullable = Nothing + }) | x <- columns] + } + +disconnectPresto :: PrestoConnection -> IO () +disconnectPresto conn = return () +-- TODO: the JDBC version seems to do a DELETE when you close the StatementClient + +------------------- Connection and polling stuff ------------------- + +getPrestoResponse :: (FromJSON a, ToJSON a) => Connection -> PrestoConnection -> URI -> Method -> B.ByteString -> EitherT String IO a +getPrestoResponse conn (PrestoConnection _ _ catalog schema) uri method body = do + request <- liftIO $ buildRequest $ do + http method (DE.encodeUtf8 $ convert $ uriPath uri) + setContentType "text/plain" + setAccept "*/*" + setHeader "X-Presto-User" "presto" + setHeader "X-Presto-Catalog" (DE.encodeUtf8 $ convert catalog) + setHeader "X-Presto-Schema" (DE.encodeUtf8 $ convert schema) + + bodyStream <- liftIO $ Streams.fromByteString $ convert body + liftIO $ sendRequest conn request $ inputStreamBody bodyStream + + -- maybeBytes <- liftIO $ receiveResponse conn (\_ i -> do + -- (i2, lenio) <- Streams.countInput i + -- len <- lenio + -- putStrLn $ "Bytes: " ++ show len + -- Streams.read i) + + liftIO $ receiveResponse conn jsonHandler + + +pollForResult :: Connection -> PrestoConnection -> URI -> Int -> EitherT String IO ([PrestoColumn], [[SqlValue]]) +pollForResult conn prestoConn uri attemptNo = do + liftIO $ threadDelay (attemptNo * 100000) -- Linear backoff in increments of 100ms = 100000 us + response <- getPrestoResponse conn prestoConn uri GET B.empty + + case _data response of + Nothing -> do + when (isNothing $ nextUri response) $ do + iUri <- hoistEither $ maybeToEither "No data or nextUri and couldn't parse infoUri" $ parseAbsoluteURI (infoUri response) + info :: PrestoInfo <- getPrestoResponse conn prestoConn iUri GET B.empty + left $ "No data and no nextUri: " ++ show response ++ "\n\n" ++ "Info: " ++ show info + newUri <- case parseAbsoluteURI $ fromJust $ nextUri response of + Nothing -> left "Failed to parse nextUri" + Just newUri -> right newUri + pollForResult conn prestoConn newUri (attemptNo + 1) + + Just dat -> case columns response of + Nothing -> left "Data but no columns found in response" + Just cols -> right (cols, fmap (makeSqlValues cols) dat) + + +initialUri = fromJust $ parseRelativeReference "/v1/statement" + +doPollingPrestoQuery :: Connection -> PrestoConnection -> URI -> T.Text -> IO (Either String ([PrestoColumn], [PrestoRow])) +doPollingPrestoQuery conn prestoConn _ q = runEitherT $ do + initialResponse <- getPrestoResponse conn prestoConn initialUri POST (encodeUtf8 q) + + uri <- maybeToEitherT "No nextUri found in initialResponse" $ nextUri initialResponse + + newUri <- case parseAbsoluteURI uri of + Nothing -> left "Failed to parse nextUri" + Just newUri -> right newUri + + pollForResult conn prestoConn newUri 0 + + +------------------- SQL stuff ------------------- + + +prestoTypeToSqlColType :: String -> CT.SqlTypeId +prestoTypeToSqlColType prestoType = case prestoType of + "varchar" -> SqlVarCharT + "bigint" -> SqlBigIntT + "boolean" -> SqlBitT + "double" -> SqlDoubleT + "date" -> SqlDateT + "time" -> SqlTimeT + "time_with_timezone" -> SqlTimeWithZoneT + "timestamp" -> SqlTimestampT + "timestamp_with_timezone" -> SqlTimestampWithZoneT + "interval_year_to_month" -> SqlIntervalT SqlIntervalYearToMonthT + "interval_year_to_second" -> error "Can't represent year to second as SqlTypeId (?)" + + +makeSqlValue :: PrestoColumn -> AnyValue -> SqlValue +makeSqlValue column value = case (prestoColumn_type column, value) of + ("varchar", TextValue t) -> SqlString $ T.unpack $ T.fromStrict t + ("double", NumValue s) -> SqlDouble (realToFrac s) + ("bigint", NumValue s) -> SqlInteger $ floor $ realToFrac s + ("boolean", BoolValue b) -> SqlBool b + (_, NullValue) -> SqlNull + _ -> error $ "Unrecognized Presto column type: " ++ (prestoColumn_type column) + +makeSqlValues :: [PrestoColumn] -> [AnyValue] -> [SqlValue] +makeSqlValues columns strings = [makeSqlValue column string | (column, string) <- Prelude.zip columns strings]