backend/feat: clickhouse subqueries

nammayatri · Dec 4, 2024 · baa7659 · baa7659
1 parent 0d44f4c
commit baa7659
Show file tree

Hide file tree

Showing 6 changed files with 241 additions and 97 deletions.
diff --git a/lib/mobility-core/src/Kernel/Storage/ClickhouseV2.hs b/lib/mobility-core/src/Kernel/Storage/ClickhouseV2.hs
@@ -17,5 +17,6 @@ module Kernel.Storage.ClickhouseV2 (module Reexport) where
 import Kernel.Storage.ClickhouseV2.ClickhouseDb as Reexport
 import Kernel.Storage.ClickhouseV2.ClickhouseTable as Reexport
 import Kernel.Storage.ClickhouseV2.ClickhouseValue as Reexport
+import Kernel.Storage.ClickhouseV2.Internal.ClickhouseColumns ()
 import Kernel.Storage.ClickhouseV2.Operators as Reexport
 import Kernel.Storage.ClickhouseV2.Queries as Reexport
diff --git a/lib/mobility-core/src/Kernel/Storage/ClickhouseV2/Internal/ClickhouseColumns.hs b/lib/mobility-core/src/Kernel/Storage/ClickhouseV2/Internal/ClickhouseColumns.hs
@@ -11,6 +11,7 @@
 
   General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
 -}
+{-# OPTIONS_GHC -Wno-orphans #-}
 
 module Kernel.Storage.ClickhouseV2.Internal.ClickhouseColumns where
 
@@ -23,18 +24,10 @@ import Kernel.Storage.ClickhouseV2.ClickhouseTable
 import Kernel.Storage.ClickhouseV2.ClickhouseValue
 import Kernel.Storage.ClickhouseV2.Internal.Types
 
-data Select a db table cols gr ord where
-  Select :: (ClickhouseTable table, ClickhouseColumns a cols) => cols -> GroupBy a gr -> Q db table cols ord -> Select a db table cols gr ord
-
-class ClickhouseColumns (a :: IsAggregated) cols where
-  type ColumnsType a cols
-  showClickhouseColumns :: Proxy a -> cols -> String
-  parseColumns :: Proxy a -> cols -> A.Value -> Either String (ColumnsType a cols)
-
 instance (FromJSON (ColumnsType 'NOT_AGG (Columns 'NOT_AGG t)), ClickhouseTable t) => ClickhouseColumns 'NOT_AGG (Columns 'NOT_AGG t) where
   type ColumnsType 'NOT_AGG (Columns 'NOT_AGG t) = t Identity
-  showClickhouseColumns _ _ = "*"
-  parseColumns _ _ = eitherResult . A.fromJSON
+  showClickhouseColumns _ _ _ = "*"
+  parseColumns _ _ val _ = eitherResult . A.fromJSON $ val
 
 -- should be all AGG columns or all NOT_AGG columns
 instance (ClickhouseValue v) => ClickhouseColumns a (Column a t v) where
@@ -75,117 +68,121 @@ parseColumns1 ::
   ClickhouseValue v1 =>
   Column a t v1 ->
   A.Value ->
+  SubQueryLevel ->
   Either String v1
-parseColumns1 c1 json = do
+parseColumns1 c1 json l = do
   mapResult <- eitherResult . A.fromJSON @(A.KeyMap (Value NotSpecified)) $ json
-  parseValueFromMap @a @t @v1 1 c1 mapResult
+  parseValueFromMap @a @t @v1 1 c1 mapResult l
 
 parseColumns2 ::
   forall a t v1 v2.
   (C2 ClickhouseValue v1 v2) =>
   T2 (Column a t) v1 v2 ->
   A.Value ->
+  SubQueryLevel ->
   Either String (v1, v2)
-parseColumns2 (c1, c2) json = do
+parseColumns2 (c1, c2) json l = do
   mapResult <- eitherResult . A.fromJSON @(A.KeyMap (Value NotSpecified)) $ json
-  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult
-  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult
+  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult l
+  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult l
   pure (v1, v2)
 
 parseColumns3 ::
   forall a t v1 v2 v3.
   (C3 ClickhouseValue v1 v2 v3) =>
   T3 (Column a t) v1 v2 v3 ->
   A.Value ->
+  SubQueryLevel ->
   Either String (v1, v2, v3)
-parseColumns3 (c1, c2, c3) json = do
+parseColumns3 (c1, c2, c3) json l = do
   mapResult <- eitherResult . A.fromJSON @(A.KeyMap (Value NotSpecified)) $ json
-  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult
-  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult
-  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult
+  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult l
+  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult l
+  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult l
   pure (v1, v2, v3)
 
 parseColumns4 ::
   forall a t v1 v2 v3 v4.
   (C4 ClickhouseValue v1 v2 v3 v4) =>
   T4 (Column a t) v1 v2 v3 v4 ->
   A.Value ->
+  SubQueryLevel ->
   Either String (v1, v2, v3, v4)
-parseColumns4 (c1, c2, c3, c4) json = do
+parseColumns4 (c1, c2, c3, c4) json l = do
   mapResult <- eitherResult . A.fromJSON @(A.KeyMap (Value NotSpecified)) $ json
-  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult
-  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult
-  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult
-  v4 <- parseValueFromMap @a @t @v4 4 c4 mapResult
+  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult l
+  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult l
+  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult l
+  v4 <- parseValueFromMap @a @t @v4 4 c4 mapResult l
   pure (v1, v2, v3, v4)
 
 parseColumns5 ::
   forall a t v1 v2 v3 v4 v5.
   (C5 ClickhouseValue v1 v2 v3 v4 v5) =>
   T5 (Column a t) v1 v2 v3 v4 v5 ->
   A.Value ->
+  SubQueryLevel ->
   Either String (v1, v2, v3, v4, v5)
-parseColumns5 (c1, c2, c3, c4, c5) json = do
+parseColumns5 (c1, c2, c3, c4, c5) json l = do
   mapResult <- eitherResult . A.fromJSON @(A.KeyMap (Value NotSpecified)) $ json
-  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult
-  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult
-  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult
-  v4 <- parseValueFromMap @a @t @v4 4 c4 mapResult
-  v5 <- parseValueFromMap @a @t @v5 5 c5 mapResult
+  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult l
+  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult l
+  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult l
+  v4 <- parseValueFromMap @a @t @v4 4 c4 mapResult l
+  v5 <- parseValueFromMap @a @t @v5 5 c5 mapResult l
   pure (v1, v2, v3, v4, v5)
 
 parseColumns6 ::
   forall a t v1 v2 v3 v4 v5 v6.
   (C6 ClickhouseValue v1 v2 v3 v4 v5 v6) =>
   T6 (Column a t) v1 v2 v3 v4 v5 v6 ->
   A.Value ->
+  SubQueryLevel ->
   Either String (v1, v2, v3, v4, v5, v6)
-parseColumns6 (c1, c2, c3, c4, c5, c6) json = do
+parseColumns6 (c1, c2, c3, c4, c5, c6) json l = do
   mapResult <- eitherResult . A.fromJSON @(A.KeyMap (Value NotSpecified)) $ json
-  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult
-  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult
-  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult
-  v4 <- parseValueFromMap @a @t @v4 4 c4 mapResult
-  v5 <- parseValueFromMap @a @t @v5 5 c5 mapResult
-  v6 <- parseValueFromMap @a @t @v6 6 c6 mapResult
+  v1 <- parseValueFromMap @a @t @v1 1 c1 mapResult l
+  v2 <- parseValueFromMap @a @t @v2 2 c2 mapResult l
+  v3 <- parseValueFromMap @a @t @v3 3 c3 mapResult l
+  v4 <- parseValueFromMap @a @t @v4 4 c4 mapResult l
+  v5 <- parseValueFromMap @a @t @v5 5 c5 mapResult l
+  v6 <- parseValueFromMap @a @t @v6 6 c6 mapResult l
   pure (v1, v2, v3, v4, v5, v6)
 
 -- FIXME should parse Numbers also
 parseValueFromMap ::
   forall a t v.
   (ClickhouseValue v) =>
-  Int ->
+  ColumnNumber ->
   Column a t v ->
   A.KeyMap (Value NotSpecified) ->
+  SubQueryLevel ->
   Either String v
-parseValueFromMap n column mapResult = do
+parseValueFromMap n column mapResult l = do
   let columnName = showColumn column
-  val <- case A.lookup (fromString $ getColumnSynonym n) mapResult of
-    Nothing -> Left $ "Key \"" <> getColumnSynonym n <> "\" for column \"" <> columnName <> "\" did not found"
+  val <- case A.lookup (fromString $ getColumnSynonym n l) mapResult of
+    Nothing -> Left $ "Key \"" <> getColumnSynonym n l <> "\" for column \"" <> columnName <> "\" did not found"
     Just val -> pure $ coerce @(Value NotSpecified) @(Value v) val
-  either (\err -> Left $ "Failed to parse key \"" <> getColumnSynonym n <> "\" for column \"" <> columnName <> "\": " <> err) pure $
+  either (\err -> Left $ "Failed to parse key \"" <> getColumnSynonym n l <> "\" for column \"" <> columnName <> "\": " <> err) pure $
     getExcept $ fromClickhouseValue @v val
 
-zipColumnsWithSynonyms1 :: Column a t v1 -> String
+zipColumnsWithSynonyms1 :: Column a t v1 -> SubQueryLevel -> String
 zipColumnsWithSynonyms1 c1 = zipColumns [showColumn c1]
 
-zipColumnsWithSynonyms2 :: T2 (Column a t) v1 v2 -> String
+zipColumnsWithSynonyms2 :: T2 (Column a t) v1 v2 -> SubQueryLevel -> String
 zipColumnsWithSynonyms2 (c1, c2) = zipColumns [showColumn c1, showColumn c2]
 
-zipColumnsWithSynonyms3 :: T3 (Column a t) v1 v2 v3 -> String
+zipColumnsWithSynonyms3 :: T3 (Column a t) v1 v2 v3 -> SubQueryLevel -> String
 zipColumnsWithSynonyms3 (c1, c2, c3) = zipColumns [showColumn c1, showColumn c2, showColumn c3]
 
-zipColumnsWithSynonyms4 :: T4 (Column a t) v1 v2 v3 v4 -> String
+zipColumnsWithSynonyms4 :: T4 (Column a t) v1 v2 v3 v4 -> SubQueryLevel -> String
 zipColumnsWithSynonyms4 (c1, c2, c3, c4) = zipColumns [showColumn c1, showColumn c2, showColumn c3, showColumn c4]
 
-zipColumnsWithSynonyms5 :: T5 (Column a t) v1 v2 v3 v4 v5 -> String
+zipColumnsWithSynonyms5 :: T5 (Column a t) v1 v2 v3 v4 v5 -> SubQueryLevel -> String
 zipColumnsWithSynonyms5 (c1, c2, c3, c4, c5) = zipColumns [showColumn c1, showColumn c2, showColumn c3, showColumn c4, showColumn c5]
 
-zipColumnsWithSynonyms6 :: T6 (Column a t) v1 v2 v3 v4 v5 v6 -> String
+zipColumnsWithSynonyms6 :: T6 (Column a t) v1 v2 v3 v4 v5 v6 -> SubQueryLevel -> String
 zipColumnsWithSynonyms6 (c1, c2, c3, c4, c5, c6) = zipColumns [showColumn c1, showColumn c2, showColumn c3, showColumn c4, showColumn c5, showColumn c6]
 
-zipColumns :: [String] -> String
-zipColumns columns = List.intercalate ", " $ zipWith (\n column -> column <> " as " <> getColumnSynonym n) [1 ..] columns
-
-getColumnSynonym :: Int -> String
-getColumnSynonym n = "res" <> show n
+zipColumns :: [String] -> SubQueryLevel -> String
+zipColumns columns l = List.intercalate ", " $ zipWith (\n column -> column <> " as " <> getColumnSynonym n l) [1 ..] columns
diff --git a/lib/mobility-core/src/Kernel/Storage/ClickhouseV2/Internal/ClickhouseQuery.hs b/lib/mobility-core/src/Kernel/Storage/ClickhouseV2/Internal/ClickhouseQuery.hs
@@ -11,8 +11,7 @@
 
   General Public License along with this program. If not, see <https://www.gnu.org/licenses/>.
 -}
-{-# LANGUAGE DerivingStrategies #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# OPTIONS_GHC -Wno-orphans #-}
 
 module Kernel.Storage.ClickhouseV2.Internal.ClickhouseQuery
   ( ClickhouseQuery (toClickhouseQuery),
@@ -27,35 +26,33 @@ import Kernel.Prelude
 import Kernel.Storage.ClickhouseV2.ClickhouseDb
 import Kernel.Storage.ClickhouseV2.ClickhouseTable
 import Kernel.Storage.ClickhouseV2.ClickhouseValue
-import Kernel.Storage.ClickhouseV2.Internal.ClickhouseColumns
 import Kernel.Storage.ClickhouseV2.Internal.Types
 import Kernel.Utils.JSON (camelToSnakeCase)
 
-newtype RawQuery = RawQuery {getRawQuery :: String}
-  deriving newtype (IsString, Semigroup, Monoid)
-
-class ClickhouseQuery expr where
-  toClickhouseQuery :: expr -> RawQuery
-
-instance (ClickhouseDb db, ClickhouseTable t, ClickhouseColumns a cols, ClickhouseQuery gr, ClickhouseQuery ord) => ClickhouseQuery (Select a db t cols gr ord) where
+instance
+  ( ClickhouseDb db,
+    ClickhouseTable t,
+    ClickhouseColumns a cols,
+    ClickhouseQuery gr,
+    ClickhouseQuery ord,
+    ClickhouseQuery (AvailableColumns db t acols)
+  ) =>
+  ClickhouseQuery (Select a db t cols gr ord acols)
+  where
   toClickhouseQuery (Select cols groupBy q) = do
-    -- should we add table name modifier?
-    let tableName = dropBeforeDot $ camelToSnakeCase . dropTSuffix . show $ typeRep (Proxy @t)
     let selectModifier = case getSelectModifier (Proxy @t) of
           NO_SELECT_MODIFIER -> ""
           SELECT_FINAL_MODIFIER -> " FINAL "
     "SELECT "
-      <> RawQuery (showClickhouseColumns @a @cols (Proxy @a) cols)
+      <> RawQuery (showClickhouseColumns @a @cols (Proxy @a) cols q.subQueryLevelQ)
       <> " FROM "
-      <> fromString tableName
+      <> toClickhouseQuery @(AvailableColumns db t acols) q.tableQ
       <> selectModifier
-      <> toClickhouseQuery @(Where t) (q.whereQ cols)
+      <> mkMaybeClause @(Where t) (q.whereQ <&> ($ cols))
       <> toClickhouseQuery @(GroupBy a gr) groupBy
       <> mkMaybeClause @(OrderBy ord) (q.orderByQ <&> ($ cols))
       <> mkMaybeClause @Limit q.limitQ
       <> mkMaybeClause @Offset q.offsetQ
-    where
-      dropTSuffix str = take (length str - 1) str
 
 mkMaybeClause :: forall expr. ClickhouseQuery expr => Maybe expr -> RawQuery
 mkMaybeClause = maybe mempty (toClickhouseQuery @expr)
@@ -141,3 +138,13 @@ instance ClickhouseTable t => ClickhouseQuery (T5 (Column a t) v1 v2 v3 v4 v5) w
 
 instance ClickhouseTable t => ClickhouseQuery (T6 (Column a t) v1 v2 v3 v4 v5 v6) where
   toClickhouseQuery (c1, c2, c3, c4, c5, c6) = intercalate ", " [toClickhouseQuery c1, toClickhouseQuery c2, toClickhouseQuery c3, toClickhouseQuery c4, toClickhouseQuery c5, toClickhouseQuery c6]
+
+instance ClickhouseTable t => ClickhouseQuery (AvailableAllColumns db t) where
+  toClickhouseQuery _ = do
+    let tableName = dropBeforeDot $ camelToSnakeCase . dropTSuffix . show $ typeRep (Proxy @t)
+    fromString tableName
+    where
+      dropTSuffix str = take (length str - 1) str
+
+instance ClickhouseQuery (AvailableSubSelectColumns db t subcols) where
+  toClickhouseQuery (AvailableColumns (SubSelectColumns subSelect)) = addBrackets . toClickhouseQuery $ subSelect