{-# LANGUAGE QuasiQuotes #-}

module Hasura.Backends.BigQuery.Schema.Introspection
  ( listAllTables,
  )
where

import Data.Aeson (toJSON)
import Data.HashMap.Strict.InsOrd qualified as InsOrdHashMap
import Data.String.Interpolate (i)
import Data.Text.Lazy qualified as LT
import Hasura.Backends.BigQuery.Execute (BigQuery (..), OutputValue (..), RecordSet (..), ShowDetails (..))
import Hasura.Backends.BigQuery.Execute qualified as Execute
import Hasura.Backends.BigQuery.Source (BigQueryDataset (..), BigQuerySourceConfig (..))
import Hasura.Backends.BigQuery.Types (TableName (..))
import Hasura.Base.Error (QErr, throw500, throw500WithDetail)
import Hasura.Prelude
import Hasura.RQL.Types.BackendType qualified as Backend
import Hasura.RQL.Types.Common (SourceName)
import Hasura.RQL.Types.Metadata (MetadataM)
import Hasura.RQL.Types.SchemaCache (CacheRM, askSourceConfig)

-- | List all tables, tracked or untracked, on a given BigQuery source. All
-- given datasets' tables will be included.
listAllTables :: (CacheRM m, MetadataM m, MonadError QErr m, MonadIO m) => SourceName -> m [TableName]
listAllTables :: forall (m :: * -> *).
(CacheRM m, MetadataM m, MonadError QErr m, MonadIO m) =>
SourceName -> m [TableName]
listAllTables SourceName
sourceName = do
  BigQuerySourceConfig
sourceConfig <- forall (b :: BackendType) (m :: * -> *).
(CacheRM m, MonadError QErr m, Backend b, MetadataM m) =>
SourceName -> m (SourceConfig b)
askSourceConfig @'Backend.BigQuery SourceName
sourceName

  let queryPerDataset :: BigQueryDataset -> LT.Text
      queryPerDataset :: BigQueryDataset -> Text
queryPerDataset (BigQueryDataset Text
dataset) =
        [i|
          select table_name, table_schema
          FROM #{dataset}.INFORMATION_SCHEMA.TABLES
        |]

      query :: LT.Text
      query :: Text
query =
        Text -> [Text] -> Text
LT.intercalate Text
"union all"
          ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ (BigQueryDataset -> Text) -> [BigQueryDataset] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map BigQueryDataset -> Text
queryPerDataset
          ([BigQueryDataset] -> [Text]) -> [BigQueryDataset] -> [Text]
forall a b. (a -> b) -> a -> b
$ BigQuerySourceConfig -> [BigQueryDataset]
_scDatasets BigQuerySourceConfig
sourceConfig

  (Job
_, RecordSet
recordSet) <-
    BigQueryConnection
-> BigQuery -> m (Either ExecuteProblem (Job, RecordSet))
forall (m :: * -> *).
MonadIO m =>
BigQueryConnection
-> BigQuery -> m (Either ExecuteProblem (Job, RecordSet))
Execute.streamBigQuery (BigQuerySourceConfig -> BigQueryConnection
_scConnection BigQuerySourceConfig
sourceConfig) (Text -> InsOrdHashMap ParameterName Parameter -> BigQuery
BigQuery Text
query InsOrdHashMap ParameterName Parameter
forall a. Monoid a => a
mempty)
      m (Either ExecuteProblem (Job, RecordSet))
-> (ExecuteProblem -> m (Job, RecordSet)) -> m (Job, RecordSet)
forall (m :: * -> *) e a.
Monad m =>
m (Either e a) -> (e -> m a) -> m a
`onLeftM` \ExecuteProblem
err -> Text -> Value -> m (Job, RecordSet)
forall (m :: * -> *) a. QErrM m => Text -> Value -> m a
throw500WithDetail (ShowDetails -> ExecuteProblem -> Text
Execute.executeProblemMessage ShowDetails
HideDetails ExecuteProblem
err) (ExecuteProblem -> Value
forall a. ToJSON a => a -> Value
toJSON ExecuteProblem
err)

  [(Text, Text)]
results <- [InsOrdHashMap FieldNameText OutputValue]
-> (InsOrdHashMap FieldNameText OutputValue -> m (Text, Text))
-> m [(Text, Text)]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
t a -> (a -> f b) -> f (t b)
for (Vector (InsOrdHashMap FieldNameText OutputValue)
-> [InsOrdHashMap FieldNameText OutputValue]
forall a. Vector a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (RecordSet -> Vector (InsOrdHashMap FieldNameText OutputValue)
rows RecordSet
recordSet)) \InsOrdHashMap FieldNameText OutputValue
row -> do
    Text
tableName <- case FieldNameText
-> InsOrdHashMap FieldNameText OutputValue -> Maybe OutputValue
forall k v. (Eq k, Hashable k) => k -> InsOrdHashMap k v -> Maybe v
InsOrdHashMap.lookup FieldNameText
"table_name" InsOrdHashMap FieldNameText OutputValue
row of
      Just (TextOutputValue Text
tableName) -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
tableName
      Maybe OutputValue
_ -> Text -> m Text
forall (m :: * -> *) a. QErrM m => Text -> m a
throw500 Text
"Unexpected BigQuery introspection result (table_name)"

    Text
tableNameSchema <- case FieldNameText
-> InsOrdHashMap FieldNameText OutputValue -> Maybe OutputValue
forall k v. (Eq k, Hashable k) => k -> InsOrdHashMap k v -> Maybe v
InsOrdHashMap.lookup FieldNameText
"table_schema" InsOrdHashMap FieldNameText OutputValue
row of
      Just (TextOutputValue Text
tableNameSchema) -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
tableNameSchema
      Maybe OutputValue
_ -> Text -> m Text
forall (m :: * -> *) a. QErrM m => Text -> m a
throw500 Text
"Unexpected BigQuery introspection result (table_schema)"

    (Text, Text) -> m (Text, Text)
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text
tableName, Text
tableNameSchema)

  [TableName] -> m [TableName]
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [TableName {Text
tableName :: Text
tableNameSchema :: Text
$sel:tableName:TableName :: Text
$sel:tableNameSchema:TableName :: Text
..} | (Text
tableName, Text
tableNameSchema) <- [(Text, Text)]
results]