refactor: table list to table map

PostgREST · Apr 18, 2022 · cdcc175 · cdcc175
1 parent f995799
commit cdcc175
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 34 deletions.
diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
@@ -417,7 +417,7 @@ handleDelete identifier context@(RequestContext _ ctxDbStructure ApiRequest{..}
 
 handleInfo :: Monad m => QualifiedIdentifier -> RequestContext -> Handler m Wai.Response
 handleInfo identifier RequestContext{..} =
-  case findTable (qiSchema identifier) (qiName identifier) $ dbTables ctxDbStructure of
+  case findTable identifier $ dbTables ctxDbStructure of
     Just table ->
       return $ Wai.responseLBS HTTP.status200 [allOrigins, allowH table] mempty
     Nothing ->
@@ -491,7 +491,7 @@ handleOpenApi headersOnly tSchema (RequestContext conf@AppConfig{..} dbStructure
            <*> SQL.statement tSchema (DbStructure.schemaDescription configDbPreparedStatements)
       OAIgnorePriv ->
         OpenAPI.encode conf dbStructure
-              (filter (\x -> tableSchema x == tSchema) $ DbStructure.dbTables dbStructure)
+              (M.filterWithKey (\(QualifiedIdentifier sch _) _ ->  sch == tSchema) $ DbStructure.dbTables dbStructure)
               (M.filterWithKey (\(QualifiedIdentifier sch _) _ ->  sch == tSchema) $ DbStructure.dbProcs dbStructure)
           <$> SQL.statement tSchema (DbStructure.schemaDescription configDbPreparedStatements)
       OADisabled ->

diff --git a/src/PostgREST/DbStructure.hs b/src/PostgREST/DbStructure.hs
@@ -57,14 +57,14 @@ import PostgREST.DbStructure.Relationship (Cardinality (..),
                                            Junction (..),
                                            PrimaryKey (..),
                                            Relationship (..))
-import PostgREST.DbStructure.Table        (Column (..), Table (..))
+import PostgREST.DbStructure.Table        (Column (..), Table (..), TablesMap)
 
 import Protolude
 import Protolude.Unsafe (unsafeHead)
 
 
 data DbStructure = DbStructure
-  { dbTables        :: [Table]
+  { dbTables        :: TablesMap
   , dbColumns       :: [Column]
   , dbRelationships :: [Relationship]
   , dbPrimaryKeys   :: [PrimaryKey]
@@ -80,11 +80,11 @@ tableCols dbs tSchema tName = filter (\Column{colTable=Table{tableSchema=s, tabl
 tablePKCols :: DbStructure -> Schema -> TableName -> [Text]
 tablePKCols dbs tSchema tName = pkName <$> filter (\pk -> tSchema == (tableSchema . pkTable) pk && tName == (tableName . pkTable) pk) (dbPrimaryKeys dbs)
 
-findTable :: Schema -> TableName -> [Table] -> Maybe Table
-findTable tSchema tName = find (\tbl -> tableSchema tbl == tSchema && tableName tbl == tName)
+findTable :: QualifiedIdentifier -> TablesMap -> Maybe Table
+findTable identifier tbls = M.lookup identifier tbls
 
-findIfView :: QualifiedIdentifier -> [Table] -> Bool
-findIfView identifier tbls = maybe False tableIsView (findTable (qiSchema identifier) (qiName identifier) tbls)
+findIfView :: QualifiedIdentifier -> TablesMap -> Bool
+findIfView identifier tbls = maybe False tableIsView $ findTable identifier tbls
 
 -- | The source table column a view column refers to
 type SourceColumn = (Column, ViewColumn)
@@ -119,7 +119,7 @@ queryDbStructure schemas extraSearchPath prepared = do
 removeInternal :: [Schema] -> DbStructure -> DbStructure
 removeInternal schemas dbStruct =
   DbStructure {
-      dbTables        = filter (\x -> tableSchema x `elem` schemas) $ dbTables dbStruct
+      dbTables        = M.filterWithKey (\(QualifiedIdentifier sch _) _ -> sch `elem` schemas) $ dbTables dbStruct
     , dbColumns       = filter (\x -> tableSchema (colTable x) `elem` schemas) (dbColumns dbStruct)
     , dbRelationships = filter (\x -> tableSchema (relTable x) `elem` schemas &&
                                       tableSchema (relForeignTable x) `elem` schemas &&
@@ -132,9 +132,9 @@ removeInternal schemas dbStruct =
       M2M Junction{junTable} -> tableSchema junTable `notElem` schemas
       _                      -> False
 
-decodeTables :: HD.Result [Table]
+decodeTables :: HD.Result TablesMap
 decodeTables =
-  HD.rowList tblRow
+ M.fromList . map (\tbl@Table{tableSchema, tableName} -> (QualifiedIdentifier tableSchema tableName, tbl)) <$> HD.rowList tblRow
  where
   tblRow = Table <$> column HD.text
                  <*> column HD.text
@@ -144,7 +144,7 @@ decodeTables =
                  <*> column HD.bool
                  <*> column HD.bool
 
-decodeColumns :: [Table] -> HD.Result [Column]
+decodeColumns :: TablesMap -> HD.Result [Column]
 decodeColumns tables =
   mapMaybe (columnFromRow tables) <$> HD.rowList colRow
  where
@@ -160,7 +160,7 @@ decodeColumns tables =
       <*> nullableColumn HD.text
       <*> nullableColumn HD.text
 
-decodeRels :: [Table] -> [Column] -> HD.Result [Relationship]
+decodeRels :: TablesMap -> [Column] -> HD.Result [Relationship]
 decodeRels tables cols =
   mapMaybe (relFromRow tables cols) <$> HD.rowList relRow
  where
@@ -173,7 +173,7 @@ decodeRels tables cols =
     <*> column HD.text
     <*> arrayColumn HD.text
 
-decodePks :: [Table] -> HD.Result [PrimaryKey]
+decodePks :: TablesMap -> HD.Result [PrimaryKey]
 decodePks tables =
   mapMaybe (pkFromRow tables) <$> HD.rowList pkRow
  where
@@ -333,7 +333,7 @@ schemaDescription =
       where
         n.nspname = $1 |]
 
-accessibleTables :: PgVersion -> Bool -> SQL.Statement Schema [Table]
+accessibleTables :: PgVersion -> Bool -> SQL.Statement Schema TablesMap
 accessibleTables pgVer =
   SQL.Statement sql (param HE.text) decodeTables
  where
@@ -426,7 +426,7 @@ addViewPrimaryKeys srcCols = concatMap (\pk ->
                 filter (\(col, _) -> colTable col == pkTable pk && colName col == pkName pk) srcCols in
   pk : viewPks)
 
-allTables :: PgVersion -> Bool -> SQL.Statement () [Table]
+allTables :: PgVersion -> Bool -> SQL.Statement () TablesMap
 allTables pgVer =
   SQL.Statement sql HE.noParams decodeTables
   where
@@ -483,7 +483,7 @@ tablesSqlQuery getAll pgVer = [q|
       )|]
     relIsPartition = if pgVer >= pgVersion100 then " AND not c.relispartition " else mempty
 
-allColumns :: [Table] -> Bool -> SQL.Statement [Schema] [Column]
+allColumns :: TablesMap -> Bool -> SQL.Statement [Schema] [Column]
 allColumns tabs =
  SQL.Statement sql (arrayParam HE.text) (decodeColumns tabs)
  where
@@ -604,19 +604,19 @@ allColumns tabs =
     ) AS enum_info ON (info.udt_name = enum_info.n)
     ORDER BY schema, position |]
 
-columnFromRow :: [Table] ->
+columnFromRow :: TablesMap ->
                  (Text,        Text,        Text,
                   Maybe Text,  Bool,        Text,
                   Maybe Int32, Maybe Text,  Maybe Text)
                  -> Maybe Column
 columnFromRow tabs (s, t, n, desc, nul, typ, l, d, e) = buildColumn <$> table
   where
     buildColumn tbl = Column tbl n desc nul typ l d (parseEnum e)
-    table = find (\tbl -> tableSchema tbl == s && tableName tbl == t) tabs
+    table = findTable (QualifiedIdentifier s t) tabs
     parseEnum :: Maybe Text -> [Text]
     parseEnum = maybe [] (split (==','))
 
-allM2ORels :: [Table] -> [Column] -> Bool -> SQL.Statement () [Relationship]
+allM2ORels :: TablesMap -> [Column] -> Bool -> SQL.Statement () [Relationship]
 allM2ORels tabs cols =
   SQL.Statement sql HE.noParams (decodeRels tabs cols)
  where
@@ -643,17 +643,17 @@ allM2ORels tabs cols =
     WHERE confrelid != 0
     ORDER BY (conrelid, column_info.nums) |]
 
-relFromRow :: [Table] -> [Column] -> (Text, Text, Text, [Text], Text, Text, [Text]) -> Maybe Relationship
+relFromRow :: TablesMap -> [Column] -> (Text, Text, Text, [Text], Text, Text, [Text]) -> Maybe Relationship
 relFromRow allTabs allCols (rs, rt, cn, rcs, frs, frt, frcs) =
   Relationship <$> table <*> cols <*> tableF <*> colsF <*> pure (M2O cn)
   where
     findCol s t c = find (\col -> tableSchema (colTable col) == s && tableName (colTable col) == t && colName col == c) allCols
-    table  = findTable rs rt allTabs
-    tableF = findTable frs frt allTabs
+    table  = findTable (QualifiedIdentifier rs rt) allTabs
+    tableF = findTable (QualifiedIdentifier frs frt) allTabs
     cols  = mapM (findCol rs rt) rcs
     colsF = mapM (findCol frs frt) frcs
 
-allPrimaryKeys :: [Table] -> Bool -> SQL.Statement () [PrimaryKey]
+allPrimaryKeys :: TablesMap -> Bool -> SQL.Statement () [PrimaryKey]
 allPrimaryKeys tabs =
   SQL.Statement sql HE.noParams (decodePks tabs)
  where
@@ -729,9 +729,9 @@ allPrimaryKeys tabs =
     WHERE
         key_col_usage.table_schema NOT IN ('pg_catalog', 'information_schema') |]
 
-pkFromRow :: [Table] -> (Schema, Text, Text) -> Maybe PrimaryKey
+pkFromRow :: TablesMap -> (Schema, Text, Text) -> Maybe PrimaryKey
 pkFromRow tabs (s, t, n) = PrimaryKey <$> table <*> pure n
-  where table = find (\tbl -> tableSchema tbl == s && tableName tbl == t) tabs
+  where table = findTable (QualifiedIdentifier s t) tabs
 
 -- returns all the primary and foreign key columns which are referenced in views
 pfkSourceColumns :: [Column] -> Bool -> SQL.Statement ([Schema], [Schema]) [SourceColumn]

diff --git a/src/PostgREST/DbStructure/Table.hs b/src/PostgREST/DbStructure/Table.hs
@@ -5,9 +5,11 @@ module PostgREST.DbStructure.Table
   ( Column(..)
   , Table(..)
   , tableQi
+  , TablesMap
   ) where
 
 import qualified Data.Aeson as JSON
+import qualified Data.HashMap.Strict as M
 
 import PostgREST.DbStructure.Identifiers (FieldName,
                                           QualifiedIdentifier (..),
@@ -50,8 +52,4 @@ data Column = Column
 instance Eq Column where
   Column{colTable=t1,colName=n1} == Column{colTable=t2,colName=n2} = t1 == t2 && n1 == n2
 
-data PrimaryKey = PrimaryKey
-  { pkTable :: Table
-  , pkName  :: Text
-  }
-  deriving (Generic, JSON.ToJSON)
+type TablesMap = M.HashMap QualifiedIdentifier Table
diff --git a/src/PostgREST/OpenAPI.hs b/src/PostgREST/OpenAPI.hs
@@ -34,20 +34,20 @@ import PostgREST.DbStructure.Proc         (ProcDescription (..),
 import PostgREST.DbStructure.Relationship (Cardinality (..),
                                            PrimaryKey (..),
                                            Relationship (..))
-import PostgREST.DbStructure.Table        (Column (..), Table (..))
+import PostgREST.DbStructure.Table        (Column (..), Table (..), TablesMap)
 import PostgREST.Version                  (docsVersion, prettyVersion)
 
 import PostgREST.ContentType
 
 import Protolude hiding (Proxy, get)
 
-encode :: AppConfig -> DbStructure -> [Table] -> M.HashMap k [ProcDescription] -> Maybe Text -> LBS.ByteString
+encode :: AppConfig -> DbStructure -> TablesMap -> M.HashMap k [ProcDescription] -> Maybe Text -> LBS.ByteString
 encode conf dbStructure tables procs schemaDescription =
   JSON.encode $
     postgrestSpec
       (dbRelationships dbStructure)
       (concat $ M.elems procs)
-      (openApiTableInfo dbStructure <$> tables)
+      (openApiTableInfo dbStructure <$> (snd <$> M.toList tables))
       (proxyUri conf)
       schemaDescription
       (dbPrimaryKeys dbStructure)