Tables.hs

-- Copyright (c) 2017 Uber Technologies, Inc.
--
-- Permission is hereby granted, free of charge, to any person obtaining a copy
-- of this software and associated documentation files (the "Software"), to deal
-- in the Software without restriction, including without limitation the rights
-- to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-- copies of the Software, and to permit persons to whom the Software is
-- furnished to do so, subject to the following conditions:
--
-- The above copyright notice and this permission notice shall be included in
-- all copies or substantial portions of the Software.
--
-- THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-- IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-- FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-- AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-- LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-- OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-- THE SOFTWARE.

{-# LANGUAGE FlexibleContexts #-}
module Database.Sql.Util.Tables where

import           Data.Foldable

import           Data.Set (Set)
import qualified Data.Set as S

import           Data.Map (Map)
import qualified Data.Map as M

import           Data.Ord
import           Data.Monoid
import           Data.Maybe (catMaybes)

import           Data.List.NonEmpty (NonEmpty(..))
import           Data.Text.Lazy (Text)
import qualified Data.Text.Lazy as TL

import           Control.Monad.Identity
import           Control.Monad.Writer
import           Control.Monad.Reader
import           Control.Monad.State

import           Database.Sql.Type
import           Database.Sql.Position


-- This describes the usage of a particular table, whether it was read from or
-- written to. Note that if a table is used in both modes, it will show up
-- twice as each mode.
data UsageMode = ReadData | ReadMeta | WriteData | WriteMeta | Unknown
  deriving (Show, Eq, Ord)

data TableUse = TableUse UsageMode FullyQualifiedTableName
  deriving (Show, Eq, Ord)


getTables :: HasTables q => q -> Set FullyQualifiedTableName
getTables = S.map tableFromUsage . getUsages
  where
   tableFromUsage (TableUse _ t) = t

getUsages :: HasTables q => q -> Set TableUse
getUsages = execWriter . flip runReaderT Unknown . goTables

class HasTables q where
    goTables :: q -> ReaderT UsageMode (Writer (Set TableUse)) ()

-- Note that vertica and hive statements have their own table usage instances
-- in their type file. Changes made here should also reflect there.
instance HasTables (Statement d ResolvedNames a) where
  goTables (QueryStmt q) = goTables q
  goTables (InsertStmt i) = goTables i
  goTables (UpdateStmt u) = goTables u
  goTables (DeleteStmt d) = goTables d
  goTables (TruncateStmt t) = goTables t
  goTables (CreateTableStmt c) = goTables c
  goTables (AlterTableStmt a) = goTables a
  goTables (DropTableStmt d) = goTables d
  goTables (CreateViewStmt c) = goTables c
  goTables (DropViewStmt d) = goTables d
  goTables (CreateSchemaStmt _) = return ()
  goTables (GrantStmt _) = return ()
  goTables (RevokeStmt _) = return ()
  goTables (BeginStmt _) = return ()
  goTables (CommitStmt _) = return ()
  goTables (RollbackStmt _) = return ()
  goTables (ExplainStmt _ s) = goTables s
  goTables (EmptyStmt _) = return ()

instance HasTables (Query ResolvedNames a) where
    goTables (QuerySelect _ select) = goTables select
    goTables (QueryExcept _ _ lhs rhs) = mapM_ goTables [lhs, rhs]
    goTables (QueryUnion _ _ _ lhs rhs) = mapM_ goTables [lhs, rhs]
    goTables (QueryIntersect _ _ lhs rhs) = mapM_ goTables [lhs, rhs]
    goTables (QueryWith _ ctes query) = mapM_ goTables $ query : map cteQuery ctes
    goTables (QueryOrder _ orders query) = mapM_ goTables orders >> goTables query
    goTables (QueryLimit _ _ query) = goTables query
    goTables (QueryOffset _ _ query) = goTables query

instance HasTables (Select ResolvedNames a) where
    goTables (Select {..}) = sequence_
        [ goTables selectCols
        , maybe (return ()) goTables selectFrom
        , maybe (return ()) goTables selectWhere
        , maybe (return ()) goTables selectTimeseries
        , maybe (return ()) goTables selectGroup
        , maybe (return ()) goTables selectHaving
        , maybe (return ()) goTables selectNamedWindow
        ]


emitTable :: (MonadWriter (Set TableUse) m,
              MonadReader UsageMode m)
          => FQTableName a -> m ()
emitTable t = do
  r <- ask
  tell . S.singleton $ TableUse r $ fqtnToFQTN t

instance (HasTables a, HasTables b) => HasTables (a, b) where
    goTables (a, b) = goTables a >> goTables b

instance HasTables (RColumnRef a) where
    goTables (RColumnRef fqcn) = emitTable . runIdentity $ columnNameTable fqcn
    goTables (RColumnAlias _) = return () -- the table that introduced it will get iterated over anyway. (shrug)

instance HasTables (FQTableName a) where
    goTables fqtn = emitTable fqtn

instance HasTables (RTableName a) where
    goTables (RTableName table _) = emitTable table

instance HasTables (SelectColumns ResolvedNames a) where
    goTables (SelectColumns _ columns) = mapM_ goTables columns

instance HasTables (SelectFrom ResolvedNames a) where
    goTables (SelectFrom _ tablishes) = local (\_ -> ReadData) $ mapM_ goTables tablishes

instance HasTables (SelectWhere ResolvedNames a) where
    goTables (SelectWhere _ condition) = goTables condition

instance HasTables (SelectTimeseries ResolvedNames a) where
    goTables (SelectTimeseries _ _ _ partition expr) = do
        goTables partition
        goTables expr

instance HasTables (PositionOrExpr ResolvedNames a) where
    goTables (PositionOrExprPosition _ _ _) = return ()
    goTables (PositionOrExprExpr expr) = goTables expr

instance HasTables (GroupingElement ResolvedNames a) where
    goTables (GroupingElementExpr _ posOrExpr) = goTables posOrExpr
    goTables (GroupingElementSet _ exprs) = mapM_ goTables exprs

instance HasTables (SelectGroup ResolvedNames a) where
    goTables (SelectGroup _ groupingElements) = mapM_ goTables groupingElements

instance HasTables (SelectHaving ResolvedNames a) where
    goTables (SelectHaving _ havings) = mapM_ goTables havings

instance HasTables (SelectNamedWindow ResolvedNames a) where
    goTables (SelectNamedWindow _ windows) = mapM_ goTables windows

instance HasTables (NamedWindowExpr ResolvedNames a) where
    goTables (NamedWindowExpr _ _ windowExpr) = goTables windowExpr
    goTables (NamedPartialWindowExpr _ _ partial) = goTables partial

instance HasTables (WindowExpr ResolvedNames a) where
    goTables (WindowExpr _ mPartition orders _) = do
        goTables mPartition
        mapM_ goTables orders

instance HasTables (PartialWindowExpr ResolvedNames a) where
    goTables (PartialWindowExpr _ _ mPartition orders _) = do
        goTables mPartition
        mapM_ goTables orders

instance HasTables (Selection ResolvedNames a) where
    goTables (SelectStar _ _ _) = return ()
    goTables (SelectExpr _ _ expr) = goTables expr

instance HasTables (Insert ResolvedNames a) where
    goTables Insert{..} = do
        local (\_ -> WriteData) $ goTables insertTable
        goTables insertValues

instance HasTables (InsertValues ResolvedNames a) where
    goTables (InsertExprValues _ e) = goTables e
    goTables (InsertSelectValues q) = goTables q
    goTables (InsertDefaultValues _) = return ()
    goTables (InsertDataFromFile _ _) = return ()

instance HasTables (DefaultExpr ResolvedNames a) where
    goTables (DefaultValue _) = return ()
    goTables (ExprValue e) = goTables e

instance HasTables a => HasTables (NonEmpty a) where
    goTables ne = mapM_ goTables ne

instance HasTables a => HasTables (Maybe a) where
    goTables Nothing = return ()
    goTables (Just a) = goTables a

instance HasTables (Update ResolvedNames a) where
    goTables Update{..} = do
        local (\_ -> WriteData) $ goTables updateTable
        local (\_ -> ReadData) $ goTables updateSetExprs
        local (\_ -> ReadData) $ goTables updateFrom
        goTables updateWhere

instance HasTables (Delete ResolvedNames a) where
    goTables (Delete _ table expr) = do
        local (\_ -> WriteData) $ goTables table
        goTables expr

instance HasTables (CreateTable d ResolvedNames a) where
    goTables CreateTable{createTableName = RCreateTableName table _, ..} = do
        -- TODO handle createTableExtra, and the dialect instances
        local (\_ -> WriteData) $ emitTable table
        goTables createTableDefinition

instance HasTables (TableDefinition d ResolvedNames a) where
    goTables (TableColumns _ s) = goTables s
    goTables (TableLike _ table) = goTables table
    goTables (TableAs _ _ query) = goTables query
    goTables (TableNoColumnInfo _) = return ()

instance HasTables (ColumnOrConstraint d ResolvedNames a) where
    goTables (ColumnOrConstraintColumn c) = goTables c
    goTables (ColumnOrConstraintConstraint _) = return ()

instance HasTables (ColumnDefinition d ResolvedNames a) where
    goTables ColumnDefinition{..} = goTables columnDefinitionDefault

instance HasTables (Truncate ResolvedNames a) where
    goTables (Truncate _ tn) = local (\_ -> WriteData) $ goTables tn

instance HasTables (AlterTable ResolvedNames a) where
    goTables (AlterTableRenameTable _ tl tr) = do
      -- Since renaming table both reads and writes (drops) the source,
      -- we report twice.
      local (\_ -> ReadData) $ goTables tl
      local (\_ -> WriteData) $ goTables tl
      local (\_ -> WriteData) $ goTables tr
    goTables (AlterTableRenameColumn _ t _ _) =
      local (\_ -> WriteMeta) $ goTables t
    goTables (AlterTableAddColumns _ t _) =
      local (\_ -> WriteData) $ goTables t

instance HasTables (DropTable ResolvedNames a) where
    goTables DropTable{dropTableNames = tables} =
      mapM_
      (\case
          RDropExistingTableName table _ -> local (\_ -> WriteData) $ emitTable table
          RDropMissingTableName _ -> pure ()
      )
      tables

instance HasTables (CreateView ResolvedNames a) where
    goTables CreateView{createViewName = RCreateTableName view _, ..} = do
        local (\_ -> WriteMeta) $ emitTable view
        goTables createViewQuery

instance HasTables (DropView ResolvedNames a) where
    goTables DropView{dropViewName = RDropExistingTableName view _} = local (\_ -> WriteMeta) $ emitTable view
    goTables DropView{dropViewName = RDropMissingTableName _} = pure ()

instance HasTables (Tablish ResolvedNames a) where
    goTables (TablishTable _ _ (RTableRef fqtn _)) = emitTable fqtn
    goTables (TablishTable _ _ (RTableAlias _)) = return ()
    goTables (TablishSubQuery _ _ query) = goTables query
    goTables (TablishLateralView _ LateralView{..} lhs) = goTables lhs >> mapM_ goTables lateralViewExprs
    goTables (TablishJoin _ _ cond outer inner) = do
        case cond of
            JoinOn expr -> goTables expr
            JoinNatural _ _ -> return ()
            JoinUsing _ _ -> return ()

        goTables outer
        goTables inner

instance HasTables (Expr ResolvedNames a) where
    goTables (BinOpExpr _ _ lhs rhs) = mapM_ goTables [lhs, rhs]
    goTables (CaseExpr _ whens else_) =
        let whens' = foldl' (\ xs (x, y) -> x:y:xs) [] whens
         in mapM_ goTables (maybe id (:) else_ whens')

    goTables (UnOpExpr _ _ operand) = goTables operand
    goTables (LikeExpr _ _ escape pattern expr) = mapM_ goTables $ catMaybes
        [ fmap escapeExpr escape
        , Just (patternExpr pattern)
        , Just expr
        ]

    goTables (ConstantExpr _ _) = return ()
    goTables (ColumnExpr _ _) = return ()
    goTables (InListExpr _ exprs expr) = mapM_ goTables $ expr : exprs
    goTables (InSubqueryExpr _ query expr) = do
        goTables expr
        goTables query

    goTables (BetweenExpr _ expr start end) =
        mapM_ goTables [expr, start, end]

    goTables (OverlapsExpr _ (s1, e1) (s2, e2)) =
        mapM_ goTables [s1, e1, s2, e2]

    goTables (FunctionExpr _ _ _ args params mFilter mOver) = do
        mapM_ goTables $ args ++ map snd params
        maybe (return ()) goTables mFilter
        maybe (return ()) goTables mOver

    goTables (AtTimeZoneExpr _ expr tz) = mapM_ goTables [expr, tz]
    goTables (SubqueryExpr _ query) = goTables query
    goTables (ArrayExpr _ values) = mapM_ goTables values
    goTables (ExistsExpr _ query) = goTables query
    goTables (FieldAccessExpr _ expr _) = goTables expr
    goTables (ArrayAccessExpr _ expr subscript) = do
        goTables expr
        goTables subscript
    goTables (TypeCastExpr _ _ expr _) = goTables expr
    goTables (VariableSubstitutionExpr _) = return ()

instance HasTables (Filter ResolvedNames a) where
    goTables (Filter _ expr) = goTables expr

instance HasTables (OverSubExpr ResolvedNames a) where
    goTables (OverWindowExpr _ windowExpr) = goTables windowExpr
    goTables (OverWindowName _ _) = return ()
    goTables (OverPartialWindowExpr _ partial) = goTables partial

instance HasTables (Partition ResolvedNames a) where
    goTables (PartitionBy _ exprs) = mapM_ goTables exprs
    goTables (PartitionBest _) = return ()
    goTables (PartitionNodes _) = return ()

instance HasTables (Order ResolvedNames a) where
    goTables (Order _ posOrExpr _ _) = goTables posOrExpr


data Open = Open { openRange :: Range
                 , openNumber :: RangeNumber
                 } deriving (Eq, Show)

data Close = Close { closeRange :: Range
                   , closeNumber :: RangeNumber
                   } deriving (Eq, Show)

class Positioned a where
    position :: a -> Position

instance Ord Open where
    compare = comparing (start . openRange)
                <> flip (comparing $ end . openRange)

instance Positioned Open where
    position = start . openRange

instance Ord Close where
    compare = comparing (end . closeRange)
                <> flip (comparing $ start . closeRange)

instance Positioned Close where
    position = end . closeRange


newtype RangeNumber = RangeNumber Integer deriving (Eq, Ord, Show)
newtype NodeNumber = NodeNumber Integer deriving (Eq, Ord, Show)

getRanges :: Query RawNames Range
        -> ( Query RawNames NodeNumber
           , Set Open, Set Close
           , Map NodeNumber RangeNumber )
getRanges query =
    let (query', (_, m)) = runState (mapM numberNodes query) (0, M.empty)
        (ranges, (_, mapping)) =
            runState (mapM numberRanges m) (0, M.empty)

        (opens, closes) =
            let makeOpenAndClose r rnum =
                    ( S.singleton (Open r rnum)
                    , S.singleton (Close r rnum)
                    )
             in M.foldMapWithKey makeOpenAndClose ranges

     in (query', opens, closes, mapping)

  where
    numberNodes r = state $ \ (i, m) ->
        let m' = M.insertWith S.union r (S.singleton $ NodeNumber i) m
         in (NodeNumber i, (i+1, m'))

    numberRanges s = state $ \ (i, m) ->
        let m' = M.union (M.fromSet (const $ RangeNumber i) s) m
         in (RangeNumber i, (i+1, m'))


spliceMarkers :: Monoid a => (Open -> a)
                          -> (Close -> a)
                          -> (Text -> a)
                          -> Set Open
                          -> Set Close
                          -> Text
                          -> a

spliceMarkers renderOpen renderClose renderText = go 0
  where
    go offset opens closes text =
        case (S.minView opens, S.minView closes) of
            (Just (o, opens'), Just (c, closes'))
                | offset == positionOffset (position c) ->
                    renderClose c <> go offset opens closes' text

                | offset == positionOffset (position o) ->
                    renderOpen o <> go offset opens' closes text

                | otherwise ->
                    let offset' = min (positionOffset $ position c)
                                      (positionOffset $ position o)
                        (chunk, rest) = TL.splitAt (offset' - offset) text
                     in renderText chunk <> go offset' opens closes rest

            (Just _, Nothing) -> error $ unwords
                [ "remaining opens when all closes are exhausted"
                , "- this should not be possible"
                ]

            (Nothing, Just (c, closes'))
                | offset == positionOffset (position c) ->
                    renderClose c <> go offset opens closes' text

                | otherwise ->
                    let offset' = positionOffset $ position c
                        (chunk, rest) = TL.splitAt (offset' - offset) text
                     in renderText chunk <> go offset' opens closes rest

            (Nothing, Nothing) -> renderText text