Line data    Source code

       1             : /******************************************************************************
       2             :  *
       3             :  * Project:  GDAL
       4             :  * Purpose:  Arrow Database Connectivity driver
       5             :  * Author:   Even Rouault, <even dot rouault at spatialys.com>
       6             :  *
       7             :  ******************************************************************************
       8             :  * Copyright (c) 2024, Even Rouault <even dot rouault at spatialys.com>
       9             :  *
      10             :  * SPDX-License-Identifier: MIT
      11             :  ****************************************************************************/
      12             : 
      13             : #ifndef OGR_ADBC_INCLUDED
      14             : #define OGR_ADBC_INCLUDED
      15             : 
      16             : #include "gdal_priv.h"
      17             : #include "ogrsf_frmts.h"
      18             : #include "ogrlayerarrow.h"
      19             : 
      20             : #include "ogr_adbc_internal.h"
      21             : 
      22             : /************************************************************************/
      23             : /*                OGRArrowArrayToOGRFeatureAdapterLayer                 */
      24             : /************************************************************************/
      25             : 
      26             : class OGRArrowArrayToOGRFeatureAdapterLayer final : public OGRLayer
      27             : {
      28             :     friend class OGRADBCLayer;
      29             :     OGRFeatureDefn *m_poLayerDefn = nullptr;
      30             :     std::vector<std::unique_ptr<OGRFeature>> m_apoFeatures{};
      31             : 
      32             :     CPL_DISALLOW_COPY_ASSIGN(OGRArrowArrayToOGRFeatureAdapterLayer)
      33             : 
      34             :   public:
      35           0 :     explicit OGRArrowArrayToOGRFeatureAdapterLayer(const char *pszName)
      36           0 :     {
      37           0 :         m_poLayerDefn = new OGRFeatureDefn(pszName);
      38           0 :         m_poLayerDefn->SetGeomType(wkbNone);
      39           0 :         m_poLayerDefn->Reference();
      40           0 :     }
      41             : 
      42             :     ~OGRArrowArrayToOGRFeatureAdapterLayer() override;
      43             : 
      44             :     using OGRLayer::GetLayerDefn;
      45             : 
      46           0 :     const OGRFeatureDefn *GetLayerDefn() const override
      47             :     {
      48           0 :         return m_poLayerDefn;
      49             :     }
      50             : 
      51           0 :     void ResetReading() override
      52             :     {
      53           0 :     }
      54             : 
      55           0 :     OGRFeature *GetNextFeature() override
      56             :     {
      57           0 :         return nullptr;
      58             :     }
      59             : 
      60           0 :     int TestCapability(const char *pszCap) const override
      61             :     {
      62           0 :         return EQUAL(pszCap, OLCCreateField) ||
      63           0 :                EQUAL(pszCap, OLCSequentialWrite);
      64             :     }
      65             : 
      66           0 :     OGRErr CreateField(const OGRFieldDefn *poFieldDefn, int) override
      67             :     {
      68           0 :         m_poLayerDefn->AddFieldDefn(poFieldDefn);
      69           0 :         return OGRERR_NONE;
      70             :     }
      71             : 
      72           0 :     OGRErr CreateGeomField(const OGRGeomFieldDefn *poGeomFieldDefn,
      73             :                            int) override
      74             :     {
      75           0 :         m_poLayerDefn->AddGeomFieldDefn(poGeomFieldDefn);
      76           0 :         return OGRERR_NONE;
      77             :     }
      78             : 
      79           0 :     OGRErr ICreateFeature(OGRFeature *poFeature) override
      80             :     {
      81             :         m_apoFeatures.emplace_back(
      82           0 :             std::unique_ptr<OGRFeature>(poFeature->Clone()));
      83           0 :         return OGRERR_NONE;
      84             :     }
      85             : };
      86             : 
      87             : /************************************************************************/
      88             : /*                             OGRADBCLayer                             */
      89             : /************************************************************************/
      90             : 
      91             : class OGRADBCDataset;
      92             : 
      93             : class OGRADBCLayer /* non final */ : public OGRLayer,
      94             :                                      public OGRGetNextFeatureThroughRaw<
      95             :                                          OGRADBCLayer>
      96             : {
      97             :   public:
      98             :     //! Describe the bbox column of a geometry column
      99             :     struct GeomColBBOX
     100             :     {
     101             :         std::string osXMin{};  // empty if no bbox column
     102             :         std::string osYMin{};
     103             :         std::string osXMax{};
     104             :         std::string osYMax{};
     105             :     };
     106             : 
     107             :   protected:
     108             :     friend class OGRADBCDataset;
     109             : 
     110             :     OGRADBCDataset *m_poDS = nullptr;
     111             :     const std::string m_osBaseStatement{};    // as provided by user
     112             :     std::string m_osModifiedBaseStatement{};  // above tuned to use ST_AsWKB()
     113             :     std::string m_osModifiedSelect{};         // SELECT part of above
     114             :     std::string m_osAttributeFilter{};
     115             :     std::unique_ptr<AdbcStatement> m_statement{};
     116             :     std::unique_ptr<OGRArrowArrayToOGRFeatureAdapterLayer> m_poAdapterLayer{};
     117             :     std::unique_ptr<OGRArrowArrayStream> m_stream{};
     118             :     bool m_bInternalUse = false;
     119             :     bool m_bLayerDefinitionError = false;
     120             : 
     121             :     struct ArrowSchema m_schema{};
     122             : 
     123             :     bool m_bEOF = false;
     124             :     size_t m_nIdx = 0;
     125             :     GIntBig m_nFeatureID = 0;
     126             :     GIntBig m_nMaxFeatureID = -1;
     127             :     bool m_bIsParquetLayer = false;
     128             : 
     129             :     std::vector<GeomColBBOX>
     130             :         m_geomColBBOX{};                     // same size as GetGeomFieldCount()
     131             :     std::vector<OGREnvelope3D> m_extents{};  // same size as GetGeomFieldCount()
     132             :     std::string m_osFIDColName{};
     133             : 
     134             :     OGRFeature *GetNextRawFeature();
     135             :     bool GetArrowStreamInternal(struct ArrowArrayStream *out_stream);
     136             :     GIntBig GetFeatureCountSelectCountStar();
     137             :     GIntBig GetFeatureCountArrow();
     138             :     GIntBig GetFeatureCountParquet();
     139             : 
     140             :     bool BuildLayerDefnInit();
     141             :     virtual void BuildLayerDefn();
     142             :     bool ReplaceStatement(const char *pszNewStatement);
     143             :     bool UpdateStatement();
     144             :     virtual std::string GetCurrentStatement() const;
     145             : 
     146           0 :     virtual bool RunDeferredCreation()
     147             :     {
     148           0 :         return true;
     149             :     }
     150             : 
     151             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCLayer)
     152             : 
     153             :   public:
     154             :     OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
     155             :                  const std::string &osStatement, bool bInternalUse);
     156             :     OGRADBCLayer(OGRADBCDataset *poDS, const char *pszName,
     157             :                  std::unique_ptr<OGRArrowArrayStream> poStream,
     158             :                  ArrowSchema *schema, bool bInternalUse);
     159             :     ~OGRADBCLayer() override;
     160             : 
     161             :     bool GotError();
     162             : 
     163             :     using OGRLayer::GetLayerDefn;
     164             :     const OGRFeatureDefn *GetLayerDefn() const override;
     165             : 
     166           0 :     const char *GetName() const override
     167             :     {
     168           0 :         return GetDescription();
     169             :     }
     170             : 
     171             :     void ResetReading() override;
     172           0 :     DEFINE_GET_NEXT_FEATURE_THROUGH_RAW(OGRADBCLayer)
     173             :     int TestCapability(const char *) const override;
     174             :     GDALDataset *GetDataset() override;
     175             :     bool GetArrowStream(struct ArrowArrayStream *out_stream,
     176             :                         CSLConstList papszOptions = nullptr) override;
     177             :     GIntBig GetFeatureCount(int bForce) override;
     178             : 
     179             :     OGRErr SetAttributeFilter(const char *pszFilter) override;
     180             :     OGRErr ISetSpatialFilter(int iGeomField,
     181             :                              const OGRGeometry *poGeom) override;
     182             : 
     183             :     OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
     184             :                       bool bForce) override;
     185             :     OGRErr IGetExtent3D(int iGeomField, OGREnvelope3D *psExtent,
     186             :                         bool bForce) override;
     187             : 
     188             :     const char *GetFIDColumn() const override;
     189             : };
     190             : 
     191             : /************************************************************************/
     192             : /*                         OGRADBCBigQueryLayer                         */
     193             : /************************************************************************/
     194             : 
     195             : class OGRADBCBigQueryLayer final : public OGRADBCLayer
     196             : {
     197             :   private:
     198             :     friend class OGRADBCDataset;
     199             : 
     200             :     bool m_bDeferredCreation = false;
     201             : 
     202             :     void BuildLayerDefn() override;
     203             :     bool RunDeferredCreation() override;
     204             :     std::string GetCurrentStatement() const override;
     205             :     bool GetBigQueryDatasetAndTableId(std::string &osDatasetId,
     206             :                                       std::string &osTableId) const;
     207             : 
     208             :   public:
     209             :     OGRADBCBigQueryLayer(OGRADBCDataset *poDS, const char *pszName,
     210             :                          const std::string &osStatement, bool bInternalUse);
     211             : 
     212             :     int TestCapability(const char *) const override;
     213             :     GIntBig GetFeatureCount(int bForce) override;
     214             :     OGRErr IGetExtent(int iGeomField, OGREnvelope *psExtent,
     215             :                       bool bForce) override;
     216             :     OGRErr SetAttributeFilter(const char *pszFilter) override;
     217             : 
     218             :     OGRErr CreateField(const OGRFieldDefn *poField, int bApproxOK) override;
     219             :     OGRErr ICreateFeature(OGRFeature *poFeature) override;
     220             :     OGRErr ISetFeature(OGRFeature *poFeature) override;
     221             :     OGRErr DeleteFeature(GIntBig nFID) override;
     222             : 
     223             :     void SetDeferredCreation(const char *pszFIDColName,
     224             :                              const OGRGeomFieldDefn *poGeomFieldDefn);
     225             : };
     226             : 
     227             : /************************************************************************/
     228             : /*                            OGRADBCDataset                            */
     229             : /************************************************************************/
     230             : 
     231             : class OGRADBCDataset final : public GDALDataset
     232             : {
     233             :     friend class OGRADBCLayer;
     234             : 
     235             :     AdbcDriver m_driver{};
     236             :     AdbcDatabase m_database{};
     237             :     std::unique_ptr<AdbcConnection> m_connection{};
     238             :     std::vector<std::unique_ptr<OGRLayer>> m_apoLayers{};
     239             :     std::string m_osParquetFilename{};
     240             :     bool m_bIsDuckDBDataset = false;
     241             :     bool m_bIsDuckDBDriver = false;
     242             :     bool m_bSpatialLoaded = false;
     243             :     bool m_bIsBigQuery = false;
     244             :     std::string m_osBigQueryDatasetId{};
     245             : 
     246             :   public:
     247           1 :     OGRADBCDataset() = default;
     248             :     ~OGRADBCDataset() override;
     249             : 
     250             :     CPLErr FlushCache(bool bAtClosing) override;
     251             : 
     252             :     bool Open(const GDALOpenInfo *poOpenInfo);
     253             : 
     254           0 :     int GetLayerCount() const override
     255             :     {
     256           0 :         return static_cast<int>(m_apoLayers.size());
     257             :     }
     258             : 
     259           0 :     const OGRLayer *GetLayer(int idx) const override
     260             :     {
     261           0 :         return (idx >= 0 && idx < GetLayerCount()) ? m_apoLayers[idx].get()
     262           0 :                                                    : nullptr;
     263             :     }
     264             : 
     265             :     OGRLayer *GetLayerByName(const char *pszName) override;
     266             : 
     267             :     std::unique_ptr<OGRADBCLayer> CreateLayer(const char *pszStatement,
     268             :                                               const char *pszLayerName,
     269             :                                               bool bInternalUse);
     270             : 
     271             :     std::unique_ptr<OGRADBCLayer>
     272             :     CreateInternalLayer(const char *pszStatement) CPL_WARN_UNUSED_RESULT;
     273             : 
     274             :     OGRLayer *ICreateLayer(const char *pszName,
     275             :                            const OGRGeomFieldDefn *poGeomFieldDefn,
     276             :                            CSLConstList papszOptions) override;
     277             :     OGRErr DeleteLayer(int iLayer) override;
     278             : 
     279             :     OGRLayer *ExecuteSQL(const char *pszStatement, OGRGeometry *poSpatialFilter,
     280             :                          const char *pszDialect) override;
     281             : 
     282             :     int TestCapability(const char *pszCap) const override;
     283             : };
     284             : 
     285             : /************************************************************************/
     286             : /*                             OGRADBCError                             */
     287             : /************************************************************************/
     288             : 
     289             : struct OGRADBCError
     290             : {
     291             :     AdbcError error{ADBC_ERROR_INIT};
     292             : 
     293             :     inline OGRADBCError() = default;
     294             : 
     295           2 :     inline ~OGRADBCError()
     296           2 :     {
     297           2 :         clear();
     298           2 :     }
     299             : 
     300           3 :     inline void clear()
     301             :     {
     302           3 :         if (error.release)
     303           0 :             error.release(&error);
     304           3 :         memset(&error, 0, sizeof(error));
     305           3 :     }
     306             : 
     307           1 :     inline const char *message() const
     308             :     {
     309           1 :         return error.message ? error.message : "";
     310             :     }
     311             : 
     312           1 :     inline operator AdbcError *()
     313             :     {
     314           1 :         return &error;
     315             :     }
     316             : 
     317             :     CPL_DISALLOW_COPY_ASSIGN(OGRADBCError)
     318             : };
     319             : 
     320             : #endif  // OGR_ADBC_INCLUDED