Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  Zarr driver
       5             :  * Author:   Even Rouault <even dot rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2021, Even Rouault <even dot rouault at spatialys.com>
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #include "zarr.h"
      14             : #include "vsikerchunk.h"
      15             : 
      16             : #include "cpl_json.h"
      17             : 
      18             : /************************************************************************/
      19             : /*              ZarrSharedResource::ZarrSharedResource()                */
      20             : /************************************************************************/
      21             : 
      22        1055 : ZarrSharedResource::ZarrSharedResource(const std::string &osRootDirectoryName,
      23        1055 :                                        bool bUpdatable)
      24        1055 :     : m_bUpdatable(bUpdatable)
      25             : {
      26        1055 :     m_oObj.Add("zarr_consolidated_format", 1);
      27        1055 :     m_oObj.Add("metadata", CPLJSONObject());
      28             : 
      29        1055 :     m_osRootDirectoryName = osRootDirectoryName;
      30        1055 :     if (!m_osRootDirectoryName.empty() && m_osRootDirectoryName.back() == '/')
      31             :     {
      32           0 :         m_osRootDirectoryName.pop_back();
      33             :     }
      34        2110 :     m_poPAM = std::make_shared<GDALPamMultiDim>(
      35        3165 :         CPLFormFilenameSafe(m_osRootDirectoryName.c_str(), "pam", nullptr));
      36        1055 : }
      37             : 
      38             : /************************************************************************/
      39             : /*              ZarrSharedResource::Create()                            */
      40             : /************************************************************************/
      41             : 
      42             : std::shared_ptr<ZarrSharedResource>
      43        1055 : ZarrSharedResource::Create(const std::string &osRootDirectoryName,
      44             :                            bool bUpdatable)
      45             : {
      46             :     return std::shared_ptr<ZarrSharedResource>(
      47        1055 :         new ZarrSharedResource(osRootDirectoryName, bUpdatable));
      48             : }
      49             : 
      50             : /************************************************************************/
      51             : /*              ZarrSharedResource::~ZarrSharedResource()               */
      52             : /************************************************************************/
      53             : 
      54        1055 : ZarrSharedResource::~ZarrSharedResource()
      55             : {
      56             :     // We try to clean caches at dataset closing, especially for Parquet
      57             :     // references, since closing Parquet datasets when the virtual file
      58             :     // systems are destroyed can be too late and cause crashes.
      59        1055 :     VSIKerchunkFileSystemsCleanCache();
      60             : 
      61        1055 :     if (m_bZMetadataModified)
      62             :     {
      63         182 :         CPLJSONDocument oDoc;
      64         182 :         oDoc.SetRoot(m_oObj);
      65         182 :         oDoc.Save(CPLFormFilenameSafe(m_osRootDirectoryName.c_str(),
      66             :                                       ".zmetadata", nullptr));
      67             :     }
      68        1055 : }
      69             : 
      70             : /************************************************************************/
      71             : /*             ZarrSharedResource::OpenRootGroup()                      */
      72             : /************************************************************************/
      73             : 
      74         740 : std::shared_ptr<ZarrGroupBase> ZarrSharedResource::OpenRootGroup()
      75             : {
      76             :     {
      77        1480 :         auto poRG = ZarrV2Group::Create(shared_from_this(), std::string(), "/");
      78         740 :         poRG->SetUpdatable(m_bUpdatable);
      79         740 :         poRG->SetDirectoryName(m_osRootDirectoryName);
      80             : 
      81             :         const std::string osZarrayFilename(CPLFormFilenameSafe(
      82         740 :             m_osRootDirectoryName.c_str(), ".zarray", nullptr));
      83             :         VSIStatBufL sStat;
      84         740 :         const auto nErrorCount = CPLGetErrorCounter();
      85         740 :         if (VSIStatL(osZarrayFilename.c_str(), &sStat) == 0)
      86             :         {
      87         518 :             CPLJSONDocument oDoc;
      88         259 :             if (!oDoc.Load(osZarrayFilename))
      89           0 :                 return nullptr;
      90         518 :             const auto oRoot = oDoc.GetRoot();
      91         259 :             if (oRoot["_NCZARR_ARRAY"].IsValid())
      92             :             {
      93             :                 // If opening a NCZarr array, initialize its group from NCZarr
      94             :                 // metadata.
      95             :                 const std::string osGroupFilename(CPLFormFilenameSafe(
      96           3 :                     CPLGetDirnameSafe(m_osRootDirectoryName.c_str()).c_str(),
      97           3 :                     ".zgroup", nullptr));
      98           3 :                 if (VSIStatL(osGroupFilename.c_str(), &sStat) == 0)
      99             :                 {
     100           2 :                     CPLJSONDocument oDocGroup;
     101           2 :                     if (oDocGroup.Load(osGroupFilename))
     102             :                     {
     103           2 :                         if (!poRG->InitFromZGroup(oDocGroup.GetRoot()))
     104           1 :                             return nullptr;
     105             :                     }
     106             :                 }
     107             :             }
     108             :             const std::string osArrayName(
     109         516 :                 CPLGetBasenameSafe(m_osRootDirectoryName.c_str()));
     110         774 :             if (!poRG->LoadArray(osArrayName, osZarrayFilename, oRoot, false,
     111         774 :                                  CPLJSONObject()))
     112          39 :                 return nullptr;
     113             : 
     114         219 :             return poRG;
     115             :         }
     116         489 :         else if (CPLGetErrorCounter() > nErrorCount &&
     117           8 :                  strstr(CPLGetLastErrorMsg(),
     118             :                         "Generation of Kerchunk Parquet cache"))
     119             :         {
     120           6 :             return nullptr;
     121             :         }
     122             : 
     123             :         const std::string osZmetadataFilename(CPLFormFilenameSafe(
     124         475 :             m_osRootDirectoryName.c_str(), ".zmetadata", nullptr));
     125         475 :         if (CPLTestBool(CSLFetchNameValueDef(GetOpenOptions(), "USE_ZMETADATA",
     126         944 :                                              "YES")) &&
     127         469 :             VSIStatL(osZmetadataFilename.c_str(), &sStat) == 0)
     128             :         {
     129         171 :             if (!m_bZMetadataEnabled)
     130             :             {
     131         171 :                 CPLJSONDocument oDoc;
     132         171 :                 if (!oDoc.Load(osZmetadataFilename))
     133           0 :                     return nullptr;
     134             : 
     135         171 :                 m_bZMetadataEnabled = true;
     136         171 :                 m_oObj = oDoc.GetRoot();
     137             :             }
     138         171 :             poRG->InitFromZMetadata(m_oObj);
     139             : 
     140         171 :             return poRG;
     141             :         }
     142             : 
     143             :         const std::string osGroupFilename(CPLFormFilenameSafe(
     144         304 :             m_osRootDirectoryName.c_str(), ".zgroup", nullptr));
     145         304 :         if (VSIStatL(osGroupFilename.c_str(), &sStat) == 0)
     146             :         {
     147         174 :             CPLJSONDocument oDoc;
     148          87 :             if (!oDoc.Load(osGroupFilename))
     149           0 :                 return nullptr;
     150             : 
     151          87 :             if (!poRG->InitFromZGroup(oDoc.GetRoot()))
     152           3 :                 return nullptr;
     153          84 :             return poRG;
     154             :         }
     155             :     }
     156             : 
     157             :     // Zarr v3
     158         434 :     auto poRG_V3 = ZarrV3Group::Create(shared_from_this(), std::string(), "/",
     159         868 :                                        m_osRootDirectoryName);
     160         217 :     poRG_V3->SetUpdatable(m_bUpdatable);
     161             : 
     162             :     const std::string osZarrJsonFilename(CPLFormFilenameSafe(
     163         434 :         m_osRootDirectoryName.c_str(), "zarr.json", nullptr));
     164             :     VSIStatBufL sStat;
     165         217 :     if (VSIStatL(osZarrJsonFilename.c_str(), &sStat) == 0)
     166             :     {
     167         428 :         CPLJSONDocument oDoc;
     168         214 :         if (!oDoc.Load(osZarrJsonFilename))
     169           0 :             return nullptr;
     170         428 :         const auto oRoot = oDoc.GetRoot();
     171         214 :         if (oRoot.GetInteger("zarr_format") != 3)
     172             :         {
     173           0 :             CPLError(CE_Failure, CPLE_AppDefined,
     174             :                      "Unhandled zarr_format value");
     175           0 :             return nullptr;
     176             :         }
     177         642 :         const std::string osNodeType = oRoot.GetString("node_type");
     178         214 :         if (osNodeType == "array")
     179             :         {
     180             :             const std::string osArrayName(
     181         106 :                 CPLGetBasenameSafe(m_osRootDirectoryName.c_str()));
     182          53 :             poRG_V3->SetExplored();
     183          53 :             if (!poRG_V3->LoadArray(osArrayName, osZarrJsonFilename, oRoot))
     184          32 :                 return nullptr;
     185             : 
     186          21 :             return poRG_V3;
     187             :         }
     188         161 :         else if (osNodeType == "group")
     189             :         {
     190         161 :             return poRG_V3;
     191             :         }
     192             :         else
     193             :         {
     194           0 :             CPLError(CE_Failure, CPLE_AppDefined, "Unhandled node_type value");
     195           0 :             return nullptr;
     196             :         }
     197             :     }
     198             : 
     199             :     // No explicit zarr.json in root directory ? Then recurse until we find
     200             :     // one.
     201           3 :     auto psDir = VSIOpenDir(m_osRootDirectoryName.c_str(), -1, nullptr);
     202           3 :     if (!psDir)
     203           3 :         return nullptr;
     204           0 :     bool bZarrJsonFound = false;
     205           0 :     while (const VSIDIREntry *psEntry = VSIGetNextDirEntry(psDir))
     206             :     {
     207           0 :         if (!VSI_ISDIR(psEntry->nMode) &&
     208           0 :             strcmp(CPLGetFilename(psEntry->pszName), "zarr.json") == 0)
     209             :         {
     210           0 :             bZarrJsonFound = true;
     211           0 :             break;
     212             :         }
     213           0 :     }
     214           0 :     VSICloseDir(psDir);
     215           0 :     if (bZarrJsonFound)
     216           0 :         return poRG_V3;
     217             : 
     218           0 :     return nullptr;
     219             : }
     220             : 
     221             : /************************************************************************/
     222             : /*             ZarrSharedResource::SetZMetadataItem()                   */
     223             : /************************************************************************/
     224             : 
     225         962 : void ZarrSharedResource::SetZMetadataItem(const std::string &osFilename,
     226             :                                           const CPLJSONObject &obj)
     227             : {
     228         962 :     if (m_bZMetadataEnabled)
     229             :     {
     230        1686 :         CPLString osNormalizedFilename(osFilename);
     231         843 :         osNormalizedFilename.replaceAll('\\', '/');
     232         843 :         CPLAssert(STARTS_WITH(osNormalizedFilename.c_str(),
     233             :                               (m_osRootDirectoryName + '/').c_str()));
     234         843 :         m_bZMetadataModified = true;
     235             :         const char *pszKey =
     236         843 :             osNormalizedFilename.c_str() + m_osRootDirectoryName.size() + 1;
     237        1686 :         auto oMetadata = m_oObj["metadata"];
     238         843 :         oMetadata.DeleteNoSplitName(pszKey);
     239         843 :         oMetadata.AddNoSplitName(pszKey, obj);
     240             :     }
     241         962 : }
     242             : 
     243             : /************************************************************************/
     244             : /*         ZarrSharedResource::DeleteZMetadataItemRecursive()           */
     245             : /************************************************************************/
     246             : 
     247          12 : void ZarrSharedResource::DeleteZMetadataItemRecursive(
     248             :     const std::string &osFilename)
     249             : {
     250          12 :     if (m_bZMetadataEnabled)
     251             :     {
     252           8 :         CPLString osNormalizedFilename(osFilename);
     253           4 :         osNormalizedFilename.replaceAll('\\', '/');
     254           4 :         CPLAssert(STARTS_WITH(osNormalizedFilename.c_str(),
     255             :                               (m_osRootDirectoryName + '/').c_str()));
     256           4 :         m_bZMetadataModified = true;
     257             :         const char *pszKey =
     258           4 :             osNormalizedFilename.c_str() + m_osRootDirectoryName.size() + 1;
     259             : 
     260          12 :         auto oMetadata = m_oObj["metadata"];
     261          28 :         for (auto &item : oMetadata.GetChildren())
     262             :         {
     263          24 :             if (STARTS_WITH(item.GetName().c_str(), pszKey))
     264             :             {
     265          14 :                 oMetadata.DeleteNoSplitName(item.GetName());
     266             :             }
     267             :         }
     268             :     }
     269          12 : }
     270             : 
     271             : /************************************************************************/
     272             : /*             ZarrSharedResource::RenameZMetadataRecursive()           */
     273             : /************************************************************************/
     274             : 
     275          12 : void ZarrSharedResource::RenameZMetadataRecursive(
     276             :     const std::string &osOldFilename, const std::string &osNewFilename)
     277             : {
     278          12 :     if (m_bZMetadataEnabled)
     279             :     {
     280           8 :         CPLString osNormalizedOldFilename(osOldFilename);
     281           4 :         osNormalizedOldFilename.replaceAll('\\', '/');
     282           4 :         CPLAssert(STARTS_WITH(osNormalizedOldFilename.c_str(),
     283             :                               (m_osRootDirectoryName + '/').c_str()));
     284             : 
     285           8 :         CPLString osNormalizedNewFilename(osNewFilename);
     286           4 :         osNormalizedNewFilename.replaceAll('\\', '/');
     287           4 :         CPLAssert(STARTS_WITH(osNormalizedNewFilename.c_str(),
     288             :                               (m_osRootDirectoryName + '/').c_str()));
     289             : 
     290           4 :         m_bZMetadataModified = true;
     291             : 
     292             :         const char *pszOldKeyRadix =
     293           4 :             osNormalizedOldFilename.c_str() + m_osRootDirectoryName.size() + 1;
     294             :         const char *pszNewKeyRadix =
     295           4 :             osNormalizedNewFilename.c_str() + m_osRootDirectoryName.size() + 1;
     296             : 
     297          12 :         auto oMetadata = m_oObj["metadata"];
     298          32 :         for (auto &item : oMetadata.GetChildren())
     299             :         {
     300          28 :             if (STARTS_WITH(item.GetName().c_str(), pszOldKeyRadix))
     301             :             {
     302          15 :                 oMetadata.DeleteNoSplitName(item.GetName());
     303          30 :                 std::string osNewKey(pszNewKeyRadix);
     304          15 :                 osNewKey += (item.GetName().c_str() + strlen(pszOldKeyRadix));
     305          15 :                 oMetadata.AddNoSplitName(osNewKey, item);
     306             :             }
     307             :         }
     308             :     }
     309          12 : }
     310             : 
     311             : /************************************************************************/
     312             : /*             ZarrSharedResource::UpdateDimensionSize()                */
     313             : /************************************************************************/
     314             : 
     315           7 : void ZarrSharedResource::UpdateDimensionSize(
     316             :     const std::shared_ptr<GDALDimension> &poDim)
     317             : {
     318          14 :     auto poRG = m_poWeakRootGroup.lock();
     319           7 :     if (!poRG)
     320           0 :         poRG = OpenRootGroup();
     321           7 :     if (poRG)
     322             :     {
     323           7 :         poRG->UpdateDimensionSize(poDim);
     324             :     }
     325             :     else
     326             :     {
     327           0 :         CPLError(CE_Failure, CPLE_AppDefined, "UpdateDimensionSize() failed");
     328             :     }
     329           7 :     poRG.reset();
     330           7 : }
     331             : 
     332             : /************************************************************************/
     333             : /*             ZarrSharedResource::AddArrayInLoading()                  */
     334             : /************************************************************************/
     335             : 
     336         883 : bool ZarrSharedResource::AddArrayInLoading(const std::string &osZarrayFilename)
     337             : {
     338             :     // Prevent too deep or recursive array loading
     339         883 :     if (m_oSetArrayInLoading.find(osZarrayFilename) !=
     340        1766 :         m_oSetArrayInLoading.end())
     341             :     {
     342           1 :         CPLError(CE_Failure, CPLE_AppDefined,
     343             :                  "Attempt at recursively loading %s", osZarrayFilename.c_str());
     344           1 :         return false;
     345             :     }
     346         882 :     if (m_oSetArrayInLoading.size() == 32)
     347             :     {
     348           1 :         CPLError(CE_Failure, CPLE_AppDefined,
     349             :                  "Too deep call stack in LoadArray()");
     350           1 :         return false;
     351             :     }
     352         881 :     m_oSetArrayInLoading.insert(osZarrayFilename);
     353         881 :     return true;
     354             : }
     355             : 
     356             : /************************************************************************/
     357             : /*             ZarrSharedResource::RemoveArrayInLoading()               */
     358             : /************************************************************************/
     359             : 
     360         881 : void ZarrSharedResource::RemoveArrayInLoading(
     361             :     const std::string &osZarrayFilename)
     362             : {
     363         881 :     m_oSetArrayInLoading.erase(osZarrayFilename);
     364         881 : }