{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE TemplateHaskell #-}

module Hasura.Backends.BigQuery.Connection
  ( BigQueryProblem,
    resolveConfigurationInput,
    resolveConfigurationInputs,
    resolveConfigurationJson,
    initConnection,
    runBigQuery,
  )
where

import Control.Concurrent.MVar
import Control.Exception
import Control.Retry qualified as Retry
import Crypto.Hash.Algorithms (SHA256 (..))
import Crypto.PubKey.RSA.PKCS15 (signSafer)
import Crypto.PubKey.RSA.Types as Cry (Error)
import Data.Aeson qualified as J
import Data.Aeson.Casing qualified as J
import Data.Aeson.TH qualified as J
import Data.ByteArray.Encoding qualified as BAE
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as B8
import Data.ByteString.Lazy qualified as BL
import Data.Environment qualified as Env
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Text.Encoding.Error qualified as TE
import Data.Time.Clock
import Data.Time.Clock.POSIX (getPOSIXTime)
import Hasura.Backends.BigQuery.Source
import Hasura.Backends.MSSQL.Connection qualified as MSSQLConn (getEnv)
import Hasura.Base.Error
import Hasura.Prelude
import Network.HTTP.Client
import Network.HTTP.Simple
import Network.HTTP.Types

newtype Scope = Scope {Scope -> Text
unScope :: T.Text}
  deriving (Int -> Scope -> ShowS
[Scope] -> ShowS
Scope -> String
(Int -> Scope -> ShowS)
-> (Scope -> String) -> ([Scope] -> ShowS) -> Show Scope
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [Scope] -> ShowS
$cshowList :: [Scope] -> ShowS
show :: Scope -> String
$cshow :: Scope -> String
showsPrec :: Int -> Scope -> ShowS
$cshowsPrec :: Int -> Scope -> ShowS
Show, Scope -> Scope -> Bool
(Scope -> Scope -> Bool) -> (Scope -> Scope -> Bool) -> Eq Scope
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: Scope -> Scope -> Bool
$c/= :: Scope -> Scope -> Bool
== :: Scope -> Scope -> Bool
$c== :: Scope -> Scope -> Bool
Eq, String -> Scope
(String -> Scope) -> IsString Scope
forall a. (String -> a) -> IsString a
fromString :: String -> Scope
$cfromString :: String -> Scope
IsString)

data GoogleAccessTokenRequest = GoogleAccessTokenRequest
  { GoogleAccessTokenRequest -> Text
_gatrGrantType :: Text,
    GoogleAccessTokenRequest -> Text
_gatrAssertion :: Text
  }
  deriving (Int -> GoogleAccessTokenRequest -> ShowS
[GoogleAccessTokenRequest] -> ShowS
GoogleAccessTokenRequest -> String
(Int -> GoogleAccessTokenRequest -> ShowS)
-> (GoogleAccessTokenRequest -> String)
-> ([GoogleAccessTokenRequest] -> ShowS)
-> Show GoogleAccessTokenRequest
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [GoogleAccessTokenRequest] -> ShowS
$cshowList :: [GoogleAccessTokenRequest] -> ShowS
show :: GoogleAccessTokenRequest -> String
$cshow :: GoogleAccessTokenRequest -> String
showsPrec :: Int -> GoogleAccessTokenRequest -> ShowS
$cshowsPrec :: Int -> GoogleAccessTokenRequest -> ShowS
Show, GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool
(GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool)
-> (GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool)
-> Eq GoogleAccessTokenRequest
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool
$c/= :: GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool
== :: GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool
$c== :: GoogleAccessTokenRequest -> GoogleAccessTokenRequest -> Bool
Eq)

$(J.deriveJSON (J.aesonDrop 5 J.snakeCase) {J.omitNothingFields = False} ''GoogleAccessTokenRequest)

mkTokenRequest :: Text -> GoogleAccessTokenRequest
mkTokenRequest :: Text -> GoogleAccessTokenRequest
mkTokenRequest = Text -> Text -> GoogleAccessTokenRequest
GoogleAccessTokenRequest Text
"urn:ietf:params:oauth:grant-type:jwt-bearer"

data TokenProblem
  = BearerTokenDecodeProblem TE.UnicodeException
  | BearerTokenSignsaferProblem Cry.Error
  | TokenFetchProblem JSONException
  | TokenRequestNonOK Status
  deriving (Int -> TokenProblem -> ShowS
[TokenProblem] -> ShowS
TokenProblem -> String
(Int -> TokenProblem -> ShowS)
-> (TokenProblem -> String)
-> ([TokenProblem] -> ShowS)
-> Show TokenProblem
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [TokenProblem] -> ShowS
$cshowList :: [TokenProblem] -> ShowS
show :: TokenProblem -> String
$cshow :: TokenProblem -> String
showsPrec :: Int -> TokenProblem -> ShowS
$cshowsPrec :: Int -> TokenProblem -> ShowS
Show, (forall x. TokenProblem -> Rep TokenProblem x)
-> (forall x. Rep TokenProblem x -> TokenProblem)
-> Generic TokenProblem
forall x. Rep TokenProblem x -> TokenProblem
forall x. TokenProblem -> Rep TokenProblem x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep TokenProblem x -> TokenProblem
$cfrom :: forall x. TokenProblem -> Rep TokenProblem x
Generic)

tokenProblemMessage :: TokenProblem -> Text
tokenProblemMessage :: TokenProblem -> Text
tokenProblemMessage = \case
  BearerTokenDecodeProblem UnicodeException
_ -> Text
"Cannot decode bearer token"
  BearerTokenSignsaferProblem Error
_ -> Text
"Cannot sign bearer token"
  TokenFetchProblem JSONException
_ -> Text
"JSON exception occurred while fetching token"
  TokenRequestNonOK Status
status -> Text
"HTTP request to fetch token failed with status " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Status -> Text
forall a. Show a => a -> Text
tshow Status
status

data ServiceAccountProblem
  = ServiceAccountFileDecodeProblem String
  deriving (Int -> ServiceAccountProblem -> ShowS
[ServiceAccountProblem] -> ShowS
ServiceAccountProblem -> String
(Int -> ServiceAccountProblem -> ShowS)
-> (ServiceAccountProblem -> String)
-> ([ServiceAccountProblem] -> ShowS)
-> Show ServiceAccountProblem
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [ServiceAccountProblem] -> ShowS
$cshowList :: [ServiceAccountProblem] -> ShowS
show :: ServiceAccountProblem -> String
$cshow :: ServiceAccountProblem -> String
showsPrec :: Int -> ServiceAccountProblem -> ShowS
$cshowsPrec :: Int -> ServiceAccountProblem -> ShowS
Show)

instance Exception ServiceAccountProblem

resolveConfigurationJson ::
  (QErrM m, J.FromJSON a) =>
  Env.Environment ->
  ConfigurationJSON a -> -- REVIEW: Can this be made polymorphic?
  m (Either String a)
resolveConfigurationJson :: Environment -> ConfigurationJSON a -> m (Either String a)
resolveConfigurationJson Environment
env = \case
  FromYamlJSON a
s -> Either String a -> m (Either String a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String a -> m (Either String a))
-> (a -> Either String a) -> a -> m (Either String a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either String a
forall a b. b -> Either a b
Right (a -> m (Either String a)) -> a -> m (Either String a)
forall a b. (a -> b) -> a -> b
$ a
s
  FromEnvJSON Text
v -> do
    Text
fileContents <- Environment -> Text -> m Text
forall (m :: * -> *). QErrM m => Environment -> Text -> m Text
MSSQLConn.getEnv Environment
env Text
v
    case ByteString -> Either String a
forall a. FromJSON a => ByteString -> Either String a
J.eitherDecode (ByteString -> Either String a)
-> (Text -> ByteString) -> Text -> Either String a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BL.fromStrict (ByteString -> ByteString)
-> (Text -> ByteString) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
TE.encodeUtf8 (Text -> Either String a) -> Text -> Either String a
forall a b. (a -> b) -> a -> b
$ Text
fileContents of
      Left String
e -> Either String a -> m (Either String a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String a -> m (Either String a))
-> (String -> Either String a) -> String -> m (Either String a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Either String a
forall a b. a -> Either a b
Left (String -> m (Either String a)) -> String -> m (Either String a)
forall a b. (a -> b) -> a -> b
$ String
e
      Right a
sa -> Either String a -> m (Either String a)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String a -> m (Either String a))
-> (a -> Either String a) -> a -> m (Either String a)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Either String a
forall a b. b -> Either a b
Right (a -> m (Either String a)) -> a -> m (Either String a)
forall a b. (a -> b) -> a -> b
$ a
sa

resolveConfigurationInput ::
  QErrM m =>
  Env.Environment ->
  ConfigurationInput ->
  m Text
resolveConfigurationInput :: Environment -> ConfigurationInput -> m Text
resolveConfigurationInput Environment
env = \case
  FromYaml Text
s -> Text -> m Text
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
s
  FromEnv Text
v -> Environment -> Text -> m Text
forall (m :: * -> *). QErrM m => Environment -> Text -> m Text
MSSQLConn.getEnv Environment
env Text
v

resolveConfigurationInputs ::
  QErrM m =>
  Env.Environment ->
  ConfigurationInputs ->
  m [Text]
resolveConfigurationInputs :: Environment -> ConfigurationInputs -> m [Text]
resolveConfigurationInputs Environment
env = \case
  FromYamls [Text]
a -> [Text] -> m [Text]
forall (f :: * -> *) a. Applicative f => a -> f a
pure [Text]
a
  FromEnvs Text
v -> (Text -> Bool) -> [Text] -> [Text]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool) -> (Text -> Bool) -> Text -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Bool
T.null) ([Text] -> [Text]) -> (Text -> [Text]) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text -> [Text]
T.splitOn Text
"," (Text -> [Text]) -> m Text -> m [Text]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Environment -> Text -> m Text
forall (m :: * -> *). QErrM m => Environment -> Text -> m Text
MSSQLConn.getEnv Environment
env Text
v

initConnection :: MonadIO m => ServiceAccount -> Text -> Maybe RetryOptions -> m BigQueryConnection
initConnection :: ServiceAccount
-> Text -> Maybe RetryOptions -> m BigQueryConnection
initConnection ServiceAccount
_bqServiceAccount Text
_bqProjectId Maybe RetryOptions
_bqRetryOptions = do
  MVar (Maybe TokenResp)
_bqAccessTokenMVar <- IO (MVar (Maybe TokenResp)) -> m (MVar (Maybe TokenResp))
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (MVar (Maybe TokenResp)) -> m (MVar (Maybe TokenResp)))
-> IO (MVar (Maybe TokenResp)) -> m (MVar (Maybe TokenResp))
forall a b. (a -> b) -> a -> b
$ Maybe TokenResp -> IO (MVar (Maybe TokenResp))
forall a. a -> IO (MVar a)
newMVar Maybe TokenResp
forall a. Maybe a
Nothing -- `runBigQuery` initializes the token
  BigQueryConnection -> m BigQueryConnection
forall (f :: * -> *) a. Applicative f => a -> f a
pure BigQueryConnection :: ServiceAccount
-> Text
-> Maybe RetryOptions
-> MVar (Maybe TokenResp)
-> BigQueryConnection
BigQueryConnection {Maybe RetryOptions
MVar (Maybe TokenResp)
Text
ServiceAccount
_bqAccessTokenMVar :: MVar (Maybe TokenResp)
_bqRetryOptions :: Maybe RetryOptions
_bqProjectId :: Text
_bqServiceAccount :: ServiceAccount
_bqAccessTokenMVar :: MVar (Maybe TokenResp)
_bqRetryOptions :: Maybe RetryOptions
_bqProjectId :: Text
_bqServiceAccount :: ServiceAccount
..}

getAccessToken :: MonadIO m => ServiceAccount -> m (Either TokenProblem TokenResp)
getAccessToken :: ServiceAccount -> m (Either TokenProblem TokenResp)
getAccessToken ServiceAccount
sa = do
  Either TokenProblem ByteString
eJwt <- ServiceAccount -> [Scope] -> m (Either TokenProblem ByteString)
forall (m :: * -> *).
MonadIO m =>
ServiceAccount -> [Scope] -> m (Either TokenProblem ByteString)
encodeBearerJWT ServiceAccount
sa [Scope
"https://www.googleapis.com/auth/cloud-platform"]
  case Either TokenProblem ByteString
eJwt of
    Left TokenProblem
tokenProblem -> Either TokenProblem TokenResp -> m (Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem TokenResp
 -> m (Either TokenProblem TokenResp))
-> (TokenProblem -> Either TokenProblem TokenResp)
-> TokenProblem
-> m (Either TokenProblem TokenResp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenProblem -> Either TokenProblem TokenResp
forall a b. a -> Either a b
Left (TokenProblem -> m (Either TokenProblem TokenResp))
-> TokenProblem -> m (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$ TokenProblem
tokenProblem
    Right ByteString
jwt ->
      case ByteString -> Either UnicodeException Text
TE.decodeUtf8' ByteString
jwt of
        Left UnicodeException
unicodeEx -> Either TokenProblem TokenResp -> m (Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem TokenResp
 -> m (Either TokenProblem TokenResp))
-> (UnicodeException -> Either TokenProblem TokenResp)
-> UnicodeException
-> m (Either TokenProblem TokenResp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenProblem -> Either TokenProblem TokenResp
forall a b. a -> Either a b
Left (TokenProblem -> Either TokenProblem TokenResp)
-> (UnicodeException -> TokenProblem)
-> UnicodeException
-> Either TokenProblem TokenResp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UnicodeException -> TokenProblem
BearerTokenDecodeProblem (UnicodeException -> m (Either TokenProblem TokenResp))
-> UnicodeException -> m (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$ UnicodeException
unicodeEx
        Right Text
assertion -> do
          Response (Either JSONException TokenResp)
tokenFetchResponse :: Response (Either JSONException TokenResp) <-
            Request -> m (Response (Either JSONException TokenResp))
forall (m :: * -> *) a.
(MonadIO m, FromJSON a) =>
Request -> m (Response (Either JSONException a))
httpJSONEither (Request -> m (Response (Either JSONException TokenResp)))
-> Request -> m (Response (Either JSONException TokenResp))
forall a b. (a -> b) -> a -> b
$
              GoogleAccessTokenRequest -> Request -> Request
forall a. ToJSON a => a -> Request -> Request
setRequestBodyJSON (Text -> GoogleAccessTokenRequest
mkTokenRequest Text
assertion) (Request -> Request) -> Request -> Request
forall a b. (a -> b) -> a -> b
$
                String -> Request
parseRequest_ (String
"POST " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
tokenURL)
          if Response (Either JSONException TokenResp) -> Int
forall a. Response a -> Int
getResponseStatusCode Response (Either JSONException TokenResp)
tokenFetchResponse Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/= Int
200
            then Either TokenProblem TokenResp -> m (Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem TokenResp
 -> m (Either TokenProblem TokenResp))
-> (Response (Either JSONException TokenResp)
    -> Either TokenProblem TokenResp)
-> Response (Either JSONException TokenResp)
-> m (Either TokenProblem TokenResp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenProblem -> Either TokenProblem TokenResp
forall a b. a -> Either a b
Left (TokenProblem -> Either TokenProblem TokenResp)
-> (Response (Either JSONException TokenResp) -> TokenProblem)
-> Response (Either JSONException TokenResp)
-> Either TokenProblem TokenResp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Status -> TokenProblem
TokenRequestNonOK (Status -> TokenProblem)
-> (Response (Either JSONException TokenResp) -> Status)
-> Response (Either JSONException TokenResp)
-> TokenProblem
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Response (Either JSONException TokenResp) -> Status
forall a. Response a -> Status
getResponseStatus (Response (Either JSONException TokenResp)
 -> m (Either TokenProblem TokenResp))
-> Response (Either JSONException TokenResp)
-> m (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$ Response (Either JSONException TokenResp)
tokenFetchResponse
            else case Response (Either JSONException TokenResp)
-> Either JSONException TokenResp
forall a. Response a -> a
getResponseBody Response (Either JSONException TokenResp)
tokenFetchResponse of
              Left JSONException
jsonEx -> Either TokenProblem TokenResp -> m (Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem TokenResp
 -> m (Either TokenProblem TokenResp))
-> (JSONException -> Either TokenProblem TokenResp)
-> JSONException
-> m (Either TokenProblem TokenResp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenProblem -> Either TokenProblem TokenResp
forall a b. a -> Either a b
Left (TokenProblem -> Either TokenProblem TokenResp)
-> (JSONException -> TokenProblem)
-> JSONException
-> Either TokenProblem TokenResp
forall b c a. (b -> c) -> (a -> b) -> a -> c
. JSONException -> TokenProblem
TokenFetchProblem (JSONException -> m (Either TokenProblem TokenResp))
-> JSONException -> m (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$ JSONException
jsonEx
              Right tr :: TokenResp
tr@TokenResp {Integer
_trExpiresAt :: TokenResp -> Integer
_trExpiresAt :: Integer
_trExpiresAt} -> do
                -- We add the current POSIXTime and store the POSIX "moment" at
                -- which this token will expire, so that at the site where
                -- we need to check if a token is nearing expiry, we only
                -- need to compare it with the _then_ "current" POSIXTime.
                POSIXTime
expiresAt <- (Integer -> POSIXTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
_trExpiresAt POSIXTime -> POSIXTime -> POSIXTime
forall a. Num a => a -> a -> a
+) (POSIXTime -> POSIXTime) -> m POSIXTime -> m POSIXTime
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO POSIXTime -> m POSIXTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO POSIXTime
getPOSIXTime
                Either TokenProblem TokenResp -> m (Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem TokenResp
 -> m (Either TokenProblem TokenResp))
-> (TokenResp -> Either TokenProblem TokenResp)
-> TokenResp
-> m (Either TokenProblem TokenResp)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenResp -> Either TokenProblem TokenResp
forall a b. b -> Either a b
Right (TokenResp -> m (Either TokenProblem TokenResp))
-> TokenResp -> m (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$ TokenResp
tr {_trExpiresAt :: Integer
_trExpiresAt = POSIXTime -> Integer
forall a b. (RealFrac a, Integral b) => a -> b
truncate POSIXTime
expiresAt}
  where
    -- TODO: use jose for jwt encoding
    b64EncodeJ :: (J.ToJSON a) => a -> BS.ByteString
    b64EncodeJ :: a -> ByteString
b64EncodeJ = ByteString -> ByteString
base64 (ByteString -> ByteString) -> (a -> ByteString) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BL.toStrict (ByteString -> ByteString) -> (a -> ByteString) -> a -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> ByteString
forall a. ToJSON a => a -> ByteString
J.encode
    base64 :: BS.ByteString -> BS.ByteString
    base64 :: ByteString -> ByteString
base64 = Base -> ByteString -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
Base -> bin -> bout
BAE.convertToBase Base
BAE.Base64URLUnpadded
    tokenURL :: String
    tokenURL :: String
tokenURL = String
"https://www.googleapis.com/oauth2/v4/token"
    maxTokenLifetime :: Int
    maxTokenLifetime :: Int
maxTokenLifetime = Int
3600
    truncateEquals :: B8.ByteString -> B8.ByteString
    truncateEquals :: ByteString -> ByteString
truncateEquals ByteString
bs =
      case ByteString -> Maybe (ByteString, Char)
B8.unsnoc ByteString
bs of
        Maybe (ByteString, Char)
Nothing -> ByteString
forall a. Monoid a => a
mempty
        Just (ByteString
bs', Char
x)
          | Char
x Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'=' -> ByteString
bs'
          | Bool
otherwise -> ByteString
bs
    encodeBearerJWT :: (MonadIO m) => ServiceAccount -> [Scope] -> m (Either TokenProblem BS.ByteString)
    encodeBearerJWT :: ServiceAccount -> [Scope] -> m (Either TokenProblem ByteString)
encodeBearerJWT ServiceAccount {Text
PKey
_saProjectId :: ServiceAccount -> Text
_saPrivateKey :: ServiceAccount -> PKey
_saClientEmail :: ServiceAccount -> Text
_saProjectId :: Text
_saPrivateKey :: PKey
_saClientEmail :: Text
..} [Scope]
scopes = do
      ByteString
inp <- Int -> ByteString
mkSigInput (Int -> ByteString)
-> (POSIXTime -> Int) -> POSIXTime -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. POSIXTime -> Int
forall a b. (RealFrac a, Integral b) => a -> b
truncate (POSIXTime -> ByteString) -> m POSIXTime -> m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO POSIXTime -> m POSIXTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO IO POSIXTime
getPOSIXTime
      Either Error ByteString
signRes <- IO (Either Error ByteString) -> m (Either Error ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either Error ByteString) -> m (Either Error ByteString))
-> IO (Either Error ByteString) -> m (Either Error ByteString)
forall a b. (a -> b) -> a -> b
$ Maybe SHA256
-> PrivateKey -> ByteString -> IO (Either Error ByteString)
forall hashAlg (m :: * -> *).
(HashAlgorithmASN1 hashAlg, MonadRandom m) =>
Maybe hashAlg
-> PrivateKey -> ByteString -> m (Either Error ByteString)
signSafer (SHA256 -> Maybe SHA256
forall a. a -> Maybe a
Just SHA256
SHA256) (PKey -> PrivateKey
unPKey PKey
_saPrivateKey) ByteString
inp
      case Either Error ByteString
signRes of
        Left Error
e -> Either TokenProblem ByteString
-> m (Either TokenProblem ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem ByteString
 -> m (Either TokenProblem ByteString))
-> (Error -> Either TokenProblem ByteString)
-> Error
-> m (Either TokenProblem ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenProblem -> Either TokenProblem ByteString
forall a b. a -> Either a b
Left (TokenProblem -> Either TokenProblem ByteString)
-> (Error -> TokenProblem)
-> Error
-> Either TokenProblem ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Error -> TokenProblem
BearerTokenSignsaferProblem (Error -> m (Either TokenProblem ByteString))
-> Error -> m (Either TokenProblem ByteString)
forall a b. (a -> b) -> a -> b
$ Error
e
        Right ByteString
sig -> Either TokenProblem ByteString
-> m (Either TokenProblem ByteString)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either TokenProblem ByteString
 -> m (Either TokenProblem ByteString))
-> (ByteString -> Either TokenProblem ByteString)
-> ByteString
-> m (Either TokenProblem ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Either TokenProblem ByteString
forall a b. b -> Either a b
Right (ByteString -> m (Either TokenProblem ByteString))
-> ByteString -> m (Either TokenProblem ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString
inp ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"." ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
truncateEquals (ByteString -> ByteString
base64 ByteString
sig)
      where
        mkSigInput :: Int -> BS.ByteString
        mkSigInput :: Int -> ByteString
mkSigInput Int
n = ByteString
header ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"." ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
payload
          where
            header :: ByteString
header =
              Value -> ByteString
forall a. ToJSON a => a -> ByteString
b64EncodeJ (Value -> ByteString) -> Value -> ByteString
forall a b. (a -> b) -> a -> b
$
                [Pair] -> Value
J.object
                  [ Key
"alg" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= (Text
"RS256" :: T.Text),
                    Key
"typ" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= (Text
"JWT" :: T.Text)
                  ]
            payload :: ByteString
payload =
              Value -> ByteString
forall a. ToJSON a => a -> ByteString
b64EncodeJ (Value -> ByteString) -> Value -> ByteString
forall a b. (a -> b) -> a -> b
$
                [Pair] -> Value
J.object
                  [ Key
"aud" Key -> String -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= String
tokenURL,
                    Key
"scope" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= Text -> [Text] -> Text
T.intercalate Text
" " ((Scope -> Text) -> [Scope] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map Scope -> Text
unScope [Scope]
scopes),
                    Key
"iat" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= Int
n,
                    Key
"exp" Key -> Int -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
maxTokenLifetime),
                    Key
"iss" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= Text
_saClientEmail
                  ]

-- | Get a usable token. If the token has expired refresh it.
getUsableToken :: MonadIO m => BigQueryConnection -> m (Either TokenProblem TokenResp)
getUsableToken :: BigQueryConnection -> m (Either TokenProblem TokenResp)
getUsableToken BigQueryConnection {ServiceAccount
_bqServiceAccount :: ServiceAccount
_bqServiceAccount :: BigQueryConnection -> ServiceAccount
_bqServiceAccount, MVar (Maybe TokenResp)
_bqAccessTokenMVar :: MVar (Maybe TokenResp)
_bqAccessTokenMVar :: BigQueryConnection -> MVar (Maybe TokenResp)
_bqAccessTokenMVar} =
  IO (Either TokenProblem TokenResp)
-> m (Either TokenProblem TokenResp)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Either TokenProblem TokenResp)
 -> m (Either TokenProblem TokenResp))
-> IO (Either TokenProblem TokenResp)
-> m (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$
    MVar (Maybe TokenResp)
-> (Maybe TokenResp
    -> IO (Maybe TokenResp, Either TokenProblem TokenResp))
-> IO (Either TokenProblem TokenResp)
forall a b. MVar a -> (a -> IO (a, b)) -> IO b
modifyMVar MVar (Maybe TokenResp)
_bqAccessTokenMVar ((Maybe TokenResp
  -> IO (Maybe TokenResp, Either TokenProblem TokenResp))
 -> IO (Either TokenProblem TokenResp))
-> (Maybe TokenResp
    -> IO (Maybe TokenResp, Either TokenProblem TokenResp))
-> IO (Either TokenProblem TokenResp)
forall a b. (a -> b) -> a -> b
$ \Maybe TokenResp
mTokenResp -> do
      case Maybe TokenResp
mTokenResp of
        Maybe TokenResp
Nothing -> do
          Either TokenProblem TokenResp
refreshedToken <- ServiceAccount -> IO (Either TokenProblem TokenResp)
forall (m :: * -> *).
MonadIO m =>
ServiceAccount -> m (Either TokenProblem TokenResp)
getAccessToken ServiceAccount
_bqServiceAccount
          case Either TokenProblem TokenResp
refreshedToken of
            Left TokenProblem
e -> (Maybe TokenResp, Either TokenProblem TokenResp)
-> IO (Maybe TokenResp, Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe TokenResp
forall a. Maybe a
Nothing, TokenProblem -> Either TokenProblem TokenResp
forall a b. a -> Either a b
Left TokenProblem
e)
            Right TokenResp
t -> (Maybe TokenResp, Either TokenProblem TokenResp)
-> IO (Maybe TokenResp, Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TokenResp -> Maybe TokenResp
forall a. a -> Maybe a
Just TokenResp
t, TokenResp -> Either TokenProblem TokenResp
forall a b. b -> Either a b
Right TokenResp
t)
        Just t :: TokenResp
t@TokenResp {GoogleAccessToken
_trAccessToken :: TokenResp -> GoogleAccessToken
_trAccessToken :: GoogleAccessToken
_trAccessToken, Integer
_trExpiresAt :: Integer
_trExpiresAt :: TokenResp -> Integer
_trExpiresAt} -> do
          POSIXTime
pt <- IO POSIXTime -> IO POSIXTime
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO POSIXTime -> IO POSIXTime) -> IO POSIXTime -> IO POSIXTime
forall a b. (a -> b) -> a -> b
$ IO POSIXTime
getPOSIXTime
          if (POSIXTime
pt POSIXTime -> POSIXTime -> Bool
forall a. Ord a => a -> a -> Bool
>= Integer -> POSIXTime
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
_trExpiresAt POSIXTime -> POSIXTime -> POSIXTime
forall a. Num a => a -> a -> a
- (POSIXTime
10 :: NominalDiffTime)) -- when posix-time is greater than expires-at-minus-threshold
            then do
              Either TokenProblem TokenResp
refreshedToken' <- ServiceAccount -> IO (Either TokenProblem TokenResp)
forall (m :: * -> *).
MonadIO m =>
ServiceAccount -> m (Either TokenProblem TokenResp)
getAccessToken ServiceAccount
_bqServiceAccount
              case Either TokenProblem TokenResp
refreshedToken' of
                Left TokenProblem
e -> (Maybe TokenResp, Either TokenProblem TokenResp)
-> IO (Maybe TokenResp, Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TokenResp -> Maybe TokenResp
forall a. a -> Maybe a
Just TokenResp
t, TokenProblem -> Either TokenProblem TokenResp
forall a b. a -> Either a b
Left TokenProblem
e)
                Right TokenResp
t' -> (Maybe TokenResp, Either TokenProblem TokenResp)
-> IO (Maybe TokenResp, Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TokenResp -> Maybe TokenResp
forall a. a -> Maybe a
Just TokenResp
t', TokenResp -> Either TokenProblem TokenResp
forall a b. b -> Either a b
Right TokenResp
t')
            else (Maybe TokenResp, Either TokenProblem TokenResp)
-> IO (Maybe TokenResp, Either TokenProblem TokenResp)
forall (f :: * -> *) a. Applicative f => a -> f a
pure (TokenResp -> Maybe TokenResp
forall a. a -> Maybe a
Just TokenResp
t, TokenResp -> Either TokenProblem TokenResp
forall a b. b -> Either a b
Right TokenResp
t)

data BigQueryProblem
  = TokenProblem TokenProblem
  deriving (Int -> BigQueryProblem -> ShowS
[BigQueryProblem] -> ShowS
BigQueryProblem -> String
(Int -> BigQueryProblem -> ShowS)
-> (BigQueryProblem -> String)
-> ([BigQueryProblem] -> ShowS)
-> Show BigQueryProblem
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [BigQueryProblem] -> ShowS
$cshowList :: [BigQueryProblem] -> ShowS
show :: BigQueryProblem -> String
$cshow :: BigQueryProblem -> String
showsPrec :: Int -> BigQueryProblem -> ShowS
$cshowsPrec :: Int -> BigQueryProblem -> ShowS
Show, (forall x. BigQueryProblem -> Rep BigQueryProblem x)
-> (forall x. Rep BigQueryProblem x -> BigQueryProblem)
-> Generic BigQueryProblem
forall x. Rep BigQueryProblem x -> BigQueryProblem
forall x. BigQueryProblem -> Rep BigQueryProblem x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cto :: forall x. Rep BigQueryProblem x -> BigQueryProblem
$cfrom :: forall x. BigQueryProblem -> Rep BigQueryProblem x
Generic)

instance J.ToJSON BigQueryProblem where
  toJSON :: BigQueryProblem -> Value
toJSON (TokenProblem TokenProblem
tokenProblem) =
    [Pair] -> Value
J.object [Key
"token_problem" Key -> Text -> Pair
forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
J..= TokenProblem -> Text
tokenProblemMessage TokenProblem
tokenProblem]

runBigQuery ::
  (MonadIO m) =>
  BigQueryConnection ->
  Request ->
  m (Either BigQueryProblem (Response BL.ByteString))
runBigQuery :: BigQueryConnection
-> Request -> m (Either BigQueryProblem (Response ByteString))
runBigQuery BigQueryConnection
conn Request
req = do
  Either TokenProblem TokenResp
eToken <- BigQueryConnection -> m (Either TokenProblem TokenResp)
forall (m :: * -> *).
MonadIO m =>
BigQueryConnection -> m (Either TokenProblem TokenResp)
getUsableToken BigQueryConnection
conn
  case Either TokenProblem TokenResp
eToken of
    Left TokenProblem
e -> Either BigQueryProblem (Response ByteString)
-> m (Either BigQueryProblem (Response ByteString))
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either BigQueryProblem (Response ByteString)
 -> m (Either BigQueryProblem (Response ByteString)))
-> (TokenProblem -> Either BigQueryProblem (Response ByteString))
-> TokenProblem
-> m (Either BigQueryProblem (Response ByteString))
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BigQueryProblem -> Either BigQueryProblem (Response ByteString)
forall a b. a -> Either a b
Left (BigQueryProblem -> Either BigQueryProblem (Response ByteString))
-> (TokenProblem -> BigQueryProblem)
-> TokenProblem
-> Either BigQueryProblem (Response ByteString)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TokenProblem -> BigQueryProblem
TokenProblem (TokenProblem -> m (Either BigQueryProblem (Response ByteString)))
-> TokenProblem -> m (Either BigQueryProblem (Response ByteString))
forall a b. (a -> b) -> a -> b
$ TokenProblem
e
    Right TokenResp {GoogleAccessToken
_trAccessToken :: GoogleAccessToken
_trAccessToken :: TokenResp -> GoogleAccessToken
_trAccessToken, Integer
_trExpiresAt :: Integer
_trExpiresAt :: TokenResp -> Integer
_trExpiresAt} -> do
      let req' :: Request
req' = HeaderName -> [ByteString] -> Request -> Request
setRequestHeader HeaderName
"Authorization" [ByteString
"Bearer " ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> (Text -> ByteString
TE.encodeUtf8 (Text -> ByteString)
-> (GoogleAccessToken -> Text) -> GoogleAccessToken -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. GoogleAccessToken -> Text
coerce) GoogleAccessToken
_trAccessToken] Request
req
      -- TODO: Make this catch the HTTP exceptions
      Response ByteString -> Either BigQueryProblem (Response ByteString)
forall a b. b -> Either a b
Right (Response ByteString
 -> Either BigQueryProblem (Response ByteString))
-> m (Response ByteString)
-> m (Either BigQueryProblem (Response ByteString))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case BigQueryConnection -> Maybe RetryOptions
_bqRetryOptions BigQueryConnection
conn of
        Just RetryOptions
opts -> RetryOptions -> m (Response ByteString) -> m (Response ByteString)
forall (m :: * -> *) body.
MonadIO m =>
RetryOptions -> m (Response body) -> m (Response body)
withGoogleApiRetries RetryOptions
opts (Request -> m (Response ByteString)
forall (m :: * -> *).
MonadIO m =>
Request -> m (Response ByteString)
httpLBS Request
req')
        Maybe RetryOptions
Nothing -> Request -> m (Response ByteString)
forall (m :: * -> *).
MonadIO m =>
Request -> m (Response ByteString)
httpLBS Request
req'

-- | Uses up to specified number retries for Google API requests with the specified base delay, uses full jitter backoff,
-- see https://aws.amazon.com/ru/blogs/architecture/exponential-backoff-and-jitter/
-- HTTP statuses for transient errors were taken from
-- https://github.com/googleapis/python-api-core/blob/34ebdcc251d4f3d7d496e8e0b78847645a06650b/google/api_core/retry.py#L112-L115
withGoogleApiRetries :: (MonadIO m) => RetryOptions -> m (Response body) -> m (Response body)
withGoogleApiRetries :: RetryOptions -> m (Response body) -> m (Response body)
withGoogleApiRetries RetryOptions {Int
Microseconds
_retryNumRetries :: RetryOptions -> Int
_retryBaseDelay :: RetryOptions -> Microseconds
_retryNumRetries :: Int
_retryBaseDelay :: Microseconds
..} m (Response body)
action =
  RetryPolicyM m
-> (RetryStatus -> Response body -> m Bool)
-> (RetryStatus -> m (Response body))
-> m (Response body)
forall (m :: * -> *) b.
MonadIO m =>
RetryPolicyM m
-> (RetryStatus -> b -> m Bool) -> (RetryStatus -> m b) -> m b
Retry.retrying RetryPolicyM m
retryPolicy RetryStatus -> Response body -> m Bool
forall (f :: * -> *) p body.
Applicative f =>
p -> Response body -> f Bool
checkStatus (m (Response body) -> RetryStatus -> m (Response body)
forall a b. a -> b -> a
const m (Response body)
action)
  where
    baseDelay :: Int
baseDelay = Integer -> Int
forall a. Num a => Integer -> a
fromInteger (Integer -> Int) -> (DiffTime -> Integer) -> DiffTime -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DiffTime -> Integer
diffTimeToMicroSeconds (DiffTime -> Int) -> DiffTime -> Int
forall a b. (a -> b) -> a -> b
$ Microseconds -> DiffTime
microseconds Microseconds
_retryBaseDelay
    retryPolicy :: RetryPolicyM m
retryPolicy = Int -> RetryPolicyM m
forall (m :: * -> *). MonadIO m => Int -> RetryPolicyM m
Retry.fullJitterBackoff Int
baseDelay RetryPolicyM m -> RetryPolicyM m -> RetryPolicyM m
forall a. Semigroup a => a -> a -> a
<> Int -> RetryPolicy
Retry.limitRetries Int
_retryNumRetries
    checkStatus :: p -> Response body -> f Bool
checkStatus p
_ Response body
resp =
      Bool -> f Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> f Bool) -> Bool -> f Bool
forall a b. (a -> b) -> a -> b
$ Response body -> Status
forall a. Response a -> Status
responseStatus Response body
resp Status -> [Status] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Status
tooManyRequests429, Status
internalServerError500, Status
serviceUnavailable503]