/src/gdal/ogr/ogrsf_frmts/adbc/ogradbcdataset.cpp

Source
/******************************************************************************
 *
 * Project:  GDAL
 * Purpose:  Arrow Database Connectivity driver
 * Author:   Even Rouault, <even dot rouault at spatialys.com>
 *
 ******************************************************************************
 * Copyright (c) 2024, Even Rouault <even dot rouault at spatialys.com>
 * Copyright (c) 2024, Dewey Dunnington <dewey@voltrondata.com>
 *
 * SPDX-License-Identifier: MIT
 ****************************************************************************/

#include "ogr_adbc.h"
#include "ogradbcdrivercore.h"
#include "memdataset.h"
#include "ogr_p.h"
#include "cpl_error.h"
#include "cpl_json.h"
#include "gdal_adbc.h"

#include <algorithm>

#if defined(OGR_ADBC_HAS_DRIVER_MANAGER)
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdocumentation"
#endif
#include <arrow-adbc/adbc_driver_manager.h>
#ifdef __clang__
#pragma clang diagnostic pop
#endif
#endif

// #define DEBUG_VERBOSE

#define OGR_ADBC_VERSION ADBC_VERSION_1_1_0
static_assert(sizeof(AdbcDriver) == ADBC_DRIVER_1_1_0_SIZE);

namespace
{

#if !defined(OGR_ADBC_HAS_DRIVER_MANAGER)
AdbcStatusCode OGRDuckDBLoadDriver(const char *driver_name, void *driver,
                                   struct AdbcError *error)
{
    void *load_handle = CPLGetSymbol(driver_name, "duckdb_adbc_init");
    if (!load_handle)
    {
        return ADBC_STATUS_INTERNAL;
    }

    AdbcDriverInitFunc init_func =
        reinterpret_cast<AdbcDriverInitFunc>(load_handle);
    return init_func(OGR_ADBC_VERSION, driver, error);
}
#endif

AdbcStatusCode OGRADBCLoadDriver(const char *driver_name,
                                 const char *entrypoint, void *driver,
                                 struct AdbcError *error)
{
    GDALAdbcLoadDriverFunc load_driver_override =
        GDALGetAdbcLoadDriverOverride();
    if (load_driver_override)
    {
        return load_driver_override(driver_name, entrypoint, OGR_ADBC_VERSION,
                                    driver, error);
    }
    else
    {
#if defined(OGR_ADBC_HAS_DRIVER_MANAGER)
        return AdbcLoadDriver(driver_name, entrypoint, OGR_ADBC_VERSION, driver,
                              error);
#else
        // If the driver is for DuckDB, use a minimal loading function, which
        // doesn't rely on the ADBC driver manager.
        if (strstr(driver_name, "duckdb"))
        {
            return OGRDuckDBLoadDriver(driver_name, driver, error);
        }
        return ADBC_STATUS_NOT_IMPLEMENTED;
#endif
    }
}

}  // namespace

// Helper to wrap driver callbacks
#define ADBC_CALL(func, ...) m_driver.func(__VA_ARGS__)

/************************************************************************/
/*                           ~OGRADBCDataset()                          */
/************************************************************************/

OGRADBCDataset::~OGRADBCDataset()
{
    OGRADBCDataset::FlushCache(true);

    // Layers must be closed before the connection
    m_apoLayers.clear();
    OGRADBCError error;
    if (m_connection)
        ADBC_CALL(ConnectionRelease, m_connection.get(), error);
    error.clear();
    if (m_driver.release)
    {
        ADBC_CALL(DatabaseRelease, &m_database, error);
        m_driver.release(&m_driver, error);
    }
}

/************************************************************************/
/*                           FlushCache()                               */
/************************************************************************/

CPLErr OGRADBCDataset::FlushCache(bool /* bAtClosing */)
{
    CPLErr eErr = CE_None;
    for (auto &poLayer : m_apoLayers)
    {
        auto poADBCLayer = dynamic_cast<OGRADBCLayer *>(poLayer.get());
        if (poADBCLayer)
        {
            if (!poADBCLayer->RunDeferredCreation())
                eErr = CE_Failure;
        }
    }

    return eErr;
}

/************************************************************************/
/*                           CreateLayer()                              */
/************************************************************************/

std::unique_ptr<OGRADBCLayer>
OGRADBCDataset::CreateLayer(const char *pszStatement, const char *pszLayerName,
                            bool bInternalUse)
{

    CPLString osStatement(pszStatement);
    if (!m_osParquetFilename.empty())
    {
        const char *pszSrcLayerName = m_apoLayers.size() == 1
                                          ? m_apoLayers[0]->GetDescription()
                                          : pszLayerName;
        // Substitute the OGR layer name with the DuckDB expected filename,
        // single-quoted
        const std::string osFrom =
            std::string(" FROM ").append(pszSrcLayerName);
        const auto nPos = osStatement.ifind(osFrom);
        if (nPos != std::string::npos)
        {
            osStatement =
                osStatement.substr(0, nPos)
                    .append(" FROM '")
                    .append(OGRDuplicateCharacter(m_osParquetFilename, '\''))
                    .append("'")
                    .append(osStatement.substr(nPos + osFrom.size()));
        }
        else
        {
            const std::string osFrom2 =
                std::string(" FROM \"")
                    .append(OGRDuplicateCharacter(pszSrcLayerName, '"'))
                    .append("\"");
            const auto nPos2 = osStatement.ifind(osFrom2);
            if (nPos2 != std::string::npos)
            {
                osStatement =
                    osStatement.substr(0, nPos2)
                        .append(" FROM '")
                        .append(
                            OGRDuplicateCharacter(m_osParquetFilename, '\''))
                        .append("'")
                        .append(osStatement.substr(nPos2 + osFrom2.size()));
            }
        }
    }

    return std::make_unique<OGRADBCLayer>(this, pszLayerName, osStatement,
                                          bInternalUse);
}

/************************************************************************/
/*                        CreateInternalLayer()                         */
/************************************************************************/

std::unique_ptr<OGRADBCLayer>
OGRADBCDataset::CreateInternalLayer(const char *pszStatement)
{
#ifdef DEBUG_VERBOSE
    CPLDebug("ADBC", "%s", pszStatement);
#endif
    return CreateLayer(pszStatement, "temp", true);
}

/************************************************************************/
/*                             ExecuteSQL()                             */
/************************************************************************/

OGRLayer *OGRADBCDataset::ExecuteSQL(const char *pszStatement,
                                     OGRGeometry *poSpatialFilter,
                                     const char *pszDialect)
{
    for (auto &poLayer : m_apoLayers)
    {
        auto poADBCLayer = dynamic_cast<OGRADBCLayer *>(poLayer.get());
        if (poADBCLayer)
            poADBCLayer->RunDeferredCreation();
    }

    if (pszDialect && pszDialect[0] != 0 && !EQUAL(pszDialect, "NATIVE"))
    {
        return GDALDataset::ExecuteSQL(pszStatement, poSpatialFilter,
                                       pszDialect);
    }

    std::string osStatement(pszStatement);
    for (const char *pszPrefix : {"SELECT * FROM ", "SELECT COUNT(*) FROM "})
    {
        if (m_bIsBigQuery && STARTS_WITH_CI(pszStatement, pszPrefix))
        {
            const auto nPos = osStatement.find(' ', strlen(pszPrefix));
            const std::string osTableName = osStatement.substr(
                strlen(pszPrefix),
                nPos == std::string::npos ? nPos : nPos - strlen(pszPrefix));
            auto poADBCLayer = dynamic_cast<OGRADBCBigQueryLayer *>(
                GetLayerByName(osTableName.c_str()));
            if (poADBCLayer)
            {
                std::string osDatasetId;
                std::string osTableId;
                if (poADBCLayer->GetBigQueryDatasetAndTableId(osDatasetId,
                                                              osTableId))
                {
                    std::string osNewStatement = pszPrefix;
                    osNewStatement += '`';
                    osNewStatement += OGRDuplicateCharacter(osDatasetId, '`');
                    osNewStatement += "`.`";
                    osNewStatement += OGRDuplicateCharacter(osTableId, '`');
                    osNewStatement += '`';
                    if (nPos != std::string::npos)
                        osNewStatement += osStatement.substr(nPos);
                    osStatement = std::move(osNewStatement);
                }
            }
            break;
        }
    }

    const char *pszLayerName = "RESULTSET";
    std::unique_ptr<OGRADBCLayer> poLayer;
    if (m_bIsBigQuery)
        poLayer = std::make_unique<OGRADBCBigQueryLayer>(
            this, pszLayerName, osStatement,
            /* bInternalUse = */ false);
    else
        poLayer = CreateLayer(osStatement.c_str(), pszLayerName, false);
    if (poLayer->GotError())
        return nullptr;
    if (poSpatialFilter)
    {
        if (poLayer->GetGeomType() == wkbNone)
            return nullptr;
        poLayer->SetSpatialFilter(poSpatialFilter);
    }
    return poLayer.release();
}

/************************************************************************/
/*                       IsParquetExtension()                           */
/************************************************************************/

static bool IsParquetExtension(const char *pszStr)
{
    const std::string osExt = CPLGetExtensionSafe(pszStr);
    return EQUAL(osExt.c_str(), "parquet") || EQUAL(osExt.c_str(), "parq");
}

/************************************************************************/
/*                                Open()                                */
/************************************************************************/

bool OGRADBCDataset::Open(const GDALOpenInfo *poOpenInfo)
{
    OGRADBCError error;

    const char *pszFilename = poOpenInfo->pszFilename;
    std::unique_ptr<GDALOpenInfo> poTmpOpenInfo;
    if (STARTS_WITH(pszFilename, "ADBC:"))
    {
        pszFilename += strlen("ADBC:");
        if (pszFilename[0])
        {
            poTmpOpenInfo = std::make_unique<GDALOpenInfo>(pszFilename,
                                                           poOpenInfo->eAccess);
            poTmpOpenInfo->papszOpenOptions = poOpenInfo->papszOpenOptions;
            poOpenInfo = poTmpOpenInfo.get();
        }
    }
    const char *pszADBCDriverName =
        CSLFetchNameValue(poOpenInfo->papszOpenOptions, "ADBC_DRIVER");
    m_bIsDuckDBDataset = OGRADBCDriverIsDuckDB(poOpenInfo);
    const bool bIsSQLite3 =
        (pszADBCDriverName && EQUAL(pszADBCDriverName, "adbc_driver_sqlite")) ||
        OGRADBCDriverIsSQLite3(poOpenInfo);
    bool bIsParquet =
        OGRADBCDriverIsParquet(poOpenInfo) || IsParquetExtension(pszFilename);
    m_bIsBigQuery =
        pszADBCDriverName && strstr(pszADBCDriverName, "adbc_driver_bigquery");
    const char *pszSQL = CSLFetchNameValue(poOpenInfo->papszOpenOptions, "SQL");
    if (!bIsParquet && pszSQL)
    {
        CPLString osSQL(pszSQL);
        auto iPos = osSQL.find("FROM '");
        if (iPos != std::string::npos)
        {
            iPos += strlen("FROM '");
            const auto iPos2 = osSQL.find("'", iPos);
            if (iPos2 != std::string::npos)
            {
                std::string osFilename = osSQL.substr(iPos, iPos2 - iPos);
                if (IsParquetExtension(osFilename.c_str()))
                {
                    m_osParquetFilename = std::move(osFilename);
                    bIsParquet = true;
                }
            }
        }
    }
    const bool bIsPostgreSQL = STARTS_WITH(pszFilename, "postgresql://");

    if (!pszADBCDriverName)
    {
        if (m_bIsDuckDBDataset || bIsParquet)
        {
            pszADBCDriverName =
#ifdef _WIN32
                "duckdb.dll"
#elif defined(__MACH__) && defined(__APPLE__)
                "libduckdb.dylib"
#else
                "libduckdb.so"
#endif
                ;
        }
        else if (bIsPostgreSQL)
            pszADBCDriverName = "adbc_driver_postgresql";
        else if (bIsSQLite3)
        {
            pszADBCDriverName = "adbc_driver_sqlite";
        }
        else if (CSLFetchNameValue(poOpenInfo->papszOpenOptions,
                                   "BIGQUERY_PROJECT_ID") ||
                 CSLFetchNameValue(poOpenInfo->papszOpenOptions,
                                   "BIGQUERY_DATASET_ID") ||
                 CSLFetchNameValue(poOpenInfo->papszOpenOptions,
                                   "BIGQUERY_JSON_CREDENTIAL_STRING") ||
                 CSLFetchNameValue(poOpenInfo->papszOpenOptions,
                                   "BIGQUERY_JSON_CREDENTIAL_FILE"))
        {
            m_bIsBigQuery = true;
            pszADBCDriverName = "adbc_driver_bigquery";
        }
        else
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "Open option ADBC_DRIVER must be specified");
            return false;
        }
    }

    if (poOpenInfo->eAccess == GA_Update && !m_bIsBigQuery)
    {
        return false;
    }

    eAccess = poOpenInfo->eAccess;

    m_bIsDuckDBDriver =
        m_bIsDuckDBDataset || bIsParquet ||
        (pszADBCDriverName && strstr(pszADBCDriverName, "duckdb"));

    // Load the driver
    if (m_bIsDuckDBDriver)
    {
        if (OGRADBCLoadDriver(pszADBCDriverName, "duckdb_adbc_init", &m_driver,
                              error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined, "AdbcLoadDriver() failed: %s",
                     error.message());
            return false;
        }
    }
    else
    {
        if (OGRADBCLoadDriver(pszADBCDriverName, nullptr, &m_driver, error) !=
            ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined, "AdbcLoadDriver() failed: %s",
                     error.message());
            return false;
        }
    }

    // Allocate the database
    if (ADBC_CALL(DatabaseNew, &m_database, error) != ADBC_STATUS_OK)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "AdbcDatabaseNew() failed: %s",
                 error.message());
        return false;
    }

    // Set options
    if (m_bIsDuckDBDriver && pszFilename[0])
    {
        VSIStatBuf sStatBuf;
        if (!bIsParquet && VSIStat(pszFilename, &sStatBuf) != 0 &&
            strcmp(pszFilename, ":memory:") != 0)
        {
            CPLError(CE_Failure, CPLE_AppDefined, "%s does not exist",
                     pszFilename);
            return false;
        }
        if (ADBC_CALL(DatabaseSetOption, &m_database, "path",
                      bIsParquet ? ":memory:" : pszFilename,
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }
    }
    else if (pszFilename[0])
    {
        if (ADBC_CALL(DatabaseSetOption, &m_database, "uri", pszFilename,
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }
    }

    const auto GetAsOpenOptionOrConfigOption = [poOpenInfo](const char *pszKey)
    {
        const char *pszVal =
            CSLFetchNameValue(poOpenInfo->papszOpenOptions, pszKey);
        if (pszVal)
            return pszVal;
        // Below comments are for scripts/collect_config_options.py
        // CPLGetConfigOption("BIGQUERY_PROJECT_ID", nullptr);
        // CPLGetConfigOption("BIGQUERY_DATASET_ID", nullptr);
        // CPLGetConfigOption("BIGQUERY_JSON_CREDENTIAL_STRING", nullptr);
        // CPLGetConfigOption("BIGQUERY_JSON_CREDENTIAL_FILE", nullptr);
        return CPLGetConfigOption(pszKey, nullptr);
    };

    const char *pszBIGQUERY_PROJECT_ID =
        GetAsOpenOptionOrConfigOption("BIGQUERY_PROJECT_ID");
    if (pszBIGQUERY_PROJECT_ID && pszBIGQUERY_PROJECT_ID[0])
    {
        if (ADBC_CALL(DatabaseSetOption, &m_database,
                      "adbc.bigquery.sql.project_id", pszBIGQUERY_PROJECT_ID,
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }
    }

    const char *pszBIGQUERY_DATASET_ID =
        GetAsOpenOptionOrConfigOption("BIGQUERY_DATASET_ID");
    if (pszBIGQUERY_DATASET_ID && pszBIGQUERY_DATASET_ID[0])
    {
        m_osBigQueryDatasetId = pszBIGQUERY_DATASET_ID;
        if (ADBC_CALL(DatabaseSetOption, &m_database,
                      "adbc.bigquery.sql.dataset_id", pszBIGQUERY_DATASET_ID,
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }
    }

    const char *pszBIGQUERY_JSON_CREDENTIAL_STRING =
        GetAsOpenOptionOrConfigOption("BIGQUERY_JSON_CREDENTIAL_STRING");
    const char *pszBIGQUERY_JSON_CREDENTIAL_FILE =
        GetAsOpenOptionOrConfigOption("BIGQUERY_JSON_CREDENTIAL_FILE");
    if (pszBIGQUERY_JSON_CREDENTIAL_STRING &&
        pszBIGQUERY_JSON_CREDENTIAL_STRING[0])
    {
        if (pszBIGQUERY_JSON_CREDENTIAL_FILE &&
            pszBIGQUERY_JSON_CREDENTIAL_FILE[0])
        {
            CPLError(CE_Warning, CPLE_AppDefined,
                     "BIGQUERY_JSON_CREDENTIAL_FILE ignored when "
                     "BIGQUERY_JSON_CREDENTIAL_STRING is set");
        }

        if (ADBC_CALL(DatabaseSetOption, &m_database,
                      "adbc.bigquery.sql.auth_credentials",
                      "adbc.bigquery.sql.auth_type.json_credential_string",
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }

        if (ADBC_CALL(DatabaseSetOption, &m_database,
                      "adbc.bigquery.sql.auth_credentials",
                      pszBIGQUERY_JSON_CREDENTIAL_STRING,
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }
    }
    else if (pszBIGQUERY_JSON_CREDENTIAL_FILE &&
             pszBIGQUERY_JSON_CREDENTIAL_FILE[0])
    {
        if (ADBC_CALL(DatabaseSetOption, &m_database,
                      "adbc.bigquery.sql.auth_credentials",
                      "adbc.bigquery.sql.auth_type.json_credential_file",
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }

        if (ADBC_CALL(DatabaseSetOption, &m_database,
                      "adbc.bigquery.sql.auth_credentials",
                      pszBIGQUERY_JSON_CREDENTIAL_FILE,
                      error) != ADBC_STATUS_OK)
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "AdbcDatabaseSetOption() failed: %s", error.message());
            return false;
        }
    }

    if (m_bIsBigQuery)
    {
        if (!(pszBIGQUERY_PROJECT_ID && pszBIGQUERY_PROJECT_ID[0]))
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "Required BIGQUERY_PROJECT_ID open option not provided");
            return false;
        }
        if (!(pszBIGQUERY_JSON_CREDENTIAL_STRING &&
              pszBIGQUERY_JSON_CREDENTIAL_STRING[0]) &&
            !(pszBIGQUERY_JSON_CREDENTIAL_FILE &&
              pszBIGQUERY_JSON_CREDENTIAL_FILE[0]))
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "Required BIGQUERY_JSON_CREDENTIAL_STRING or "
                     "BIGQUERY_JSON_CREDENTIAL_FILE open option not provided");
            return false;
        }
    }

    for (const auto &[pszKey, pszValue] : cpl::IterateNameValue(
             static_cast<CSLConstList>(poOpenInfo->papszOpenOptions)))
    {
        if (STARTS_WITH_CI(pszKey, "ADBC_OPTION_"))
        {
            if (ADBC_CALL(DatabaseSetOption, &m_database,
                          pszKey + strlen("ADBC_OPTION_"), pszValue,
                          error) != ADBC_STATUS_OK)
            {
                CPLError(CE_Failure, CPLE_AppDefined,
                         "AdbcDatabaseSetOption() failed: %s", error.message());
                return false;
            }
        }
    }

    if (ADBC_CALL(DatabaseInit, &m_database, error) != ADBC_STATUS_OK)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "AdbcDatabaseInit() failed: %s",
                 error.message());
        return false;
    }

    m_connection = std::make_unique<AdbcConnection>();
    if (ADBC_CALL(ConnectionNew, m_connection.get(), error) != ADBC_STATUS_OK)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "AdbcConnectionNew() failed: %s",
                 error.message());
        return false;
    }

    if (ADBC_CALL(ConnectionInit, m_connection.get(), &m_database, error) !=
        ADBC_STATUS_OK)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "AdbcConnectionInit() failed: %s",
                 error.message());
        return false;
    }

    char **papszPreludeStatements = CSLFetchNameValueMultiple(
        poOpenInfo->papszOpenOptions, "PRELUDE_STATEMENTS");
    for (const char *pszStatement :
         cpl::Iterate(CSLConstList(papszPreludeStatements)))
    {
        CPL_IGNORE_RET_VAL(CreateInternalLayer(pszStatement)->GotError());
    }
    CSLDestroy(papszPreludeStatements);
    if (m_bIsDuckDBDriver && CPLTestBool(CPLGetConfigOption(
                                 "OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", "ON")))
    {
        auto poTmpLayer =
            CreateInternalLayer("SELECT 1 FROM duckdb_extensions() WHERE "
                                "extension_name='spatial' AND loaded = false");
        if (!poTmpLayer->GotError() &&
            std::unique_ptr<OGRFeature>(poTmpLayer->GetNextFeature()) !=
                nullptr)
        {
            CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
            CPL_IGNORE_RET_VAL(CreateInternalLayer("LOAD spatial")->GotError());
        }

        poTmpLayer =
            CreateInternalLayer("SELECT 1 FROM duckdb_extensions() WHERE "
                                "extension_name='spatial' AND loaded = true");
        m_bSpatialLoaded = !poTmpLayer->GotError() &&
                           std::unique_ptr<OGRFeature>(
                               poTmpLayer->GetNextFeature()) != nullptr;
    }

    std::string osLayerName = "RESULTSET";
    std::string osSQL;
    bool bIsParquetLayer = false;
    if (bIsParquet)
    {
        if (m_osParquetFilename.empty())
            m_osParquetFilename = pszFilename;
        osLayerName = CPLGetBasenameSafe(m_osParquetFilename.c_str());
        if (osLayerName == "*")
            osLayerName = CPLGetBasenameSafe(
                CPLGetDirnameSafe(m_osParquetFilename.c_str()).c_str());
        if (!pszSQL)
        {
            osSQL =
                CPLSPrintf("SELECT * FROM read_parquet('%s')",
                           OGRDuplicateCharacter(pszFilename, '\'').c_str());
            pszSQL = osSQL.c_str();
            bIsParquetLayer = true;
        }
    }

    if (pszSQL)
    {
        if (pszSQL[0])
        {
            std::unique_ptr<OGRADBCLayer> poLayer;
            if ((bIsParquet || m_bIsDuckDBDataset) && m_bSpatialLoaded)
            {
                std::string osErrorMsg;
                {
                    CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
                    poLayer = CreateLayer(pszSQL, osLayerName.c_str(), false);
                    if (poLayer->GotError())
                        osErrorMsg = CPLGetLastErrorMsg();
                }
                if (poLayer->GotError())
                {
                    CPLDebug("ADBC",
                             "Connecting with 'LOAD spatial' did not work "
                             "(%s). Retrying without it",
                             osErrorMsg.c_str());
                    ADBC_CALL(ConnectionRelease, m_connection.get(), error);
                    m_connection.reset();

                    ADBC_CALL(DatabaseRelease, &m_database, error);
                    memset(&m_database, 0, sizeof(m_database));

                    if (ADBC_CALL(DatabaseNew, &m_database, error) !=
                        ADBC_STATUS_OK)
                    {
                        CPLError(CE_Failure, CPLE_AppDefined,
                                 "AdbcDatabaseNew() failed: %s",
                                 error.message());
                        return false;
                    }
                    if (ADBC_CALL(DatabaseSetOption, &m_database, "path",
                                  ":memory:", error) != ADBC_STATUS_OK)
                    {
                        CPLError(CE_Failure, CPLE_AppDefined,
                                 "AdbcDatabaseSetOption() failed: %s",
                                 error.message());
                        return false;
                    }

                    if (ADBC_CALL(DatabaseInit, &m_database, error) !=
                        ADBC_STATUS_OK)
                    {
                        CPLError(CE_Failure, CPLE_AppDefined,
                                 "AdbcDatabaseInit() failed: %s",
                                 error.message());
                        return false;
                    }

                    m_connection = std::make_unique<AdbcConnection>();
                    if (ADBC_CALL(ConnectionNew, m_connection.get(), error) !=
                        ADBC_STATUS_OK)
                    {
                        CPLError(CE_Failure, CPLE_AppDefined,
                                 "AdbcConnectionNew() failed: %s",
                                 error.message());
                        return false;
                    }

                    if (ADBC_CALL(ConnectionInit, m_connection.get(),
                                  &m_database, error) != ADBC_STATUS_OK)
                    {
                        CPLError(CE_Failure, CPLE_AppDefined,
                                 "AdbcConnectionInit() failed: %s",
                                 error.message());
                        return false;
                    }
                }
            }
            if (!poLayer || poLayer->GotError())
            {
                if (m_bIsBigQuery)
                    poLayer = std::make_unique<OGRADBCBigQueryLayer>(
                        this, osLayerName.c_str(), pszSQL,
                        /* bInternalUse = */ false);
                else
                    poLayer = CreateLayer(pszSQL, osLayerName.c_str(), false);
                if (poLayer->GotError())
                    return false;
            }

            poLayer->m_bIsParquetLayer = bIsParquetLayer;
            m_apoLayers.emplace_back(std::move(poLayer));
        }
    }
    else if (m_bIsDuckDBDataset || bIsSQLite3)
    {
        auto poLayerList = CreateInternalLayer(
            "SELECT name FROM sqlite_master WHERE type IN ('table', 'view')");
        if (poLayerList->GotError() ||
            poLayerList->GetLayerDefn()->GetFieldCount() != 1)
        {
            return false;
        }

        for (const auto &poFeature : poLayerList.get())
        {
            const char *pszLayerName = poFeature->GetFieldAsString(0);
            if (bIsSQLite3 && EQUAL(pszLayerName, "SpatialIndex"))
                continue;
            const std::string osStatement =
                CPLSPrintf("SELECT * FROM \"%s\"",
                           OGRDuplicateCharacter(pszLayerName, '"').c_str());
            m_apoLayers.emplace_back(
                CreateLayer(osStatement.c_str(), pszLayerName, false));
        }
    }
    else if (bIsPostgreSQL)
    {
        auto poLayerList = CreateInternalLayer(
            "SELECT n.nspname, c.relname FROM pg_class c "
            "JOIN pg_namespace n ON c.relnamespace = n.oid "
            "AND c.relkind in ('r','v','m','f') "
            "AND n.nspname NOT IN ('pg_catalog', 'information_schema') "
            "ORDER BY c.oid");
        if (poLayerList->GotError() ||
            poLayerList->GetLayerDefn()->GetFieldCount() != 2)
        {
            return false;
        }

        for (const auto &poFeature : poLayerList.get())
        {
            const char *pszNamespace = poFeature->GetFieldAsString(0);
            const char *pszTableName = poFeature->GetFieldAsString(1);
            const std::string osStatement =
                CPLSPrintf("SELECT * FROM \"%s\".\"%s\"",
                           OGRDuplicateCharacter(pszNamespace, '"').c_str(),
                           OGRDuplicateCharacter(pszTableName, '"').c_str());

            m_apoLayers.emplace_back(CreateLayer(
                osStatement.c_str(),
                CPLSPrintf("%s.%s", pszNamespace, pszTableName), false));
        }
    }
    else if (m_bIsBigQuery)
    {
        if (!(pszBIGQUERY_DATASET_ID && pszBIGQUERY_DATASET_ID[0]))
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "Cannot list tables when BIGQUERY_DATASET_ID open option "
                     "is not provided");
            return false;
        }
        const std::string s(pszBIGQUERY_DATASET_ID);
        if (!std::all_of(s.begin(), s.end(),
                         [](char c) { return std::isalnum(c) || c == '_'; }))
        {
            CPLError(CE_Failure, CPLE_AppDefined,
                     "Invalid characters found in BIGQUERY_DATASET_ID value");
            return false;
        }
        auto poLayerList = CreateInternalLayer(
            CPLSPrintf("SELECT table_name FROM %s.INFORMATION_SCHEMA.TABLES "
                       "ORDER BY creation_time",
                       pszBIGQUERY_DATASET_ID));
        if (poLayerList->GotError() ||
            poLayerList->GetLayerDefn()->GetFieldCount() != 1)
        {
            return false;
        }

        for (const auto &poFeature : poLayerList.get())
        {
            const char *pszTableName = poFeature->GetFieldAsString(0);
            const std::string osStatement = CPLSPrintf(
                "SELECT * FROM `%s`.`%s`",
                OGRDuplicateCharacter(pszBIGQUERY_DATASET_ID, '`').c_str(),
                OGRDuplicateCharacter(pszTableName, '`').c_str());

            m_apoLayers.emplace_back(std::make_unique<OGRADBCBigQueryLayer>(
                this, pszTableName, osStatement,
                /* bInternalUse = */ false));
        }
    }

    return true;
}

/************************************************************************/
/*                          ICreateLayer()                              */
/************************************************************************/

OGRLayer *OGRADBCDataset::ICreateLayer(const char *pszName,
                                       const OGRGeomFieldDefn *poGeomFieldDefn,
                                       CSLConstList papszOptions)
{
    if (!m_bIsBigQuery)
    {
        CPLError(CE_Failure, CPLE_NotSupported,
                 "CreateLayer() only supported for BigQuery");
        return nullptr;
    }
    if (GetAccess() != GA_Update)
    {
        CPLError(
            CE_Failure, CPLE_NotSupported,
            "CreateLayer() only supported on datasets opened in update mode");
        return nullptr;
    }
    if (m_osBigQueryDatasetId.empty())
    {
        CPLError(CE_Failure, CPLE_AppDefined,
                 "Open option BIGQUERY_DATASET_ID should be set");
        return nullptr;
    }

    if (GetLayerByName(pszName))
    {
        CPLError(CE_Failure, CPLE_AppDefined, "Table %s already exists",
                 pszName);
        return nullptr;
    }

    if (poGeomFieldDefn)
    {
        const auto poSRS = poGeomFieldDefn->GetSpatialRef();
        if (poSRS && !poSRS->IsGeographic())
        {
            CPLError(CE_Failure, CPLE_NotSupported,
                     "BigQuery only supports geographic CRS. Please reproject "
                     "your layer to one (typically EPSG:4326)");
            return nullptr;
        }
    }

    const std::string osStatement = CPLSPrintf(
        "SELECT * FROM `%s`.`%s`",
        OGRDuplicateCharacter(m_osBigQueryDatasetId.c_str(), '`').c_str(),
        OGRDuplicateCharacter(pszName, '`').c_str());

    const char *pszFIDColName =
        CSLFetchNameValueDef(papszOptions, "FID", "ogc_fid");
    auto poLayer =
        std::make_unique<OGRADBCBigQueryLayer>(this, pszName, osStatement,
                                               /* bInternalUse = */ false);
    poLayer->SetDeferredCreation(pszFIDColName, poGeomFieldDefn);
    m_apoLayers.emplace_back(std::move(poLayer));
    return m_apoLayers.back().get();
}

/************************************************************************/
/*                           DeleteLayer()                              */
/************************************************************************/

OGRErr OGRADBCDataset::DeleteLayer(int iLayer)
{
    if (!m_bIsBigQuery)
    {
        CPLError(CE_Failure, CPLE_NotSupported,
                 "DeleteLayer() only supported for BigQuery");
        return OGRERR_FAILURE;
    }
    if (GetAccess() != GA_Update)
    {
        CPLError(
            CE_Failure, CPLE_NotSupported,
            "DeleteLayer() only supported on datasets opened in update mode");
        return OGRERR_FAILURE;
    }
    if (iLayer < 0 || static_cast<size_t>(iLayer) >= m_apoLayers.size())
    {
        CPLError(CE_Failure, CPLE_AppDefined, "Invalid layer index");
        return OGRERR_FAILURE;
    }

    auto poADBCLayer =
        dynamic_cast<OGRADBCBigQueryLayer *>(m_apoLayers[iLayer].get());
    if (poADBCLayer && !poADBCLayer->m_bDeferredCreation)
    {
        std::string osDatasetId;
        std::string osTableId;
        if (!poADBCLayer->GetBigQueryDatasetAndTableId(osDatasetId, osTableId))
        {
            CPLError(CE_Failure, CPLE_NotSupported,
                     "DeleteLayer(): cannot get dataset and table ID");
            return OGRERR_FAILURE;
        }

        std::string osSQL = "DROP TABLE `";
        osSQL += OGRDuplicateCharacter(osDatasetId.c_str(), '`');
        osSQL += "`.`";
        osSQL += OGRDuplicateCharacter(osTableId.c_str(), '`');
        osSQL += "`";
        // CPLDebug("ADBC", "%s", osSQL.c_str());
        if (CreateInternalLayer(osSQL.c_str())->GotError())
        {
            return OGRERR_FAILURE;
        }
    }

    m_apoLayers.erase(m_apoLayers.begin() + iLayer);
    return OGRERR_NONE;
}

/************************************************************************/
/*                         TestCapability()                             */
/************************************************************************/

int OGRADBCDataset::TestCapability(const char *pszCap) const
{
    if (EQUAL(pszCap, ODsCCreateLayer) || EQUAL(pszCap, ODsCDeleteLayer))
        return m_bIsBigQuery && eAccess == GA_Update;
    return false;
}

/************************************************************************/
/*                         GetLayerByName()                             */
/************************************************************************/

OGRLayer *OGRADBCDataset::GetLayerByName(const char *pszName)
{
    OGRLayer *poLayer = GDALDataset::GetLayerByName(pszName);
    if (poLayer || !EQUAL(pszName, "table_list"))
        return poLayer;

    OGRADBCError error;
    auto objectsStream = std::make_unique<OGRArrowArrayStream>();
    ADBC_CALL(ConnectionGetObjects, m_connection.get(),
              ADBC_OBJECT_DEPTH_TABLES,
              /* catalog = */ nullptr,
              /* db_schema = */ nullptr,
              /* table_name = */ nullptr,
              /* table_type = */ nullptr,
              /* column_name = */ nullptr, objectsStream->get(), error);

    ArrowSchema schema = {};
    if (objectsStream->get_schema(&schema) != 0)
    {
        CPLError(CE_Failure, CPLE_AppDefined, "get_schema() failed");
        return nullptr;
    }

    OGRADBCLayer tmpLayer(this, "", std::move(objectsStream), &schema,
                          /* bInternalUse = */ true);
    const auto tmpLayerDefn = tmpLayer.GetLayerDefn();
    if (tmpLayerDefn->GetFieldIndex("catalog_name") < 0 ||
        tmpLayerDefn->GetFieldIndex("catalog_db_schemas") < 0)
    {
        return nullptr;
    }

    auto poTableListLayer =
        std::make_unique<OGRMemLayer>("table_list", nullptr, wkbNone);
    {
        OGRFieldDefn oField("catalog_name", OFTString);
        poTableListLayer->CreateField(&oField);
    }
    {
        OGRFieldDefn oField("schema_name", OFTString);
        poTableListLayer->CreateField(&oField);
    }
    {
        OGRFieldDefn oField("table_name", OFTString);
        poTableListLayer->CreateField(&oField);
    }
    {
        OGRFieldDefn oField("table_type", OFTString);
        poTableListLayer->CreateField(&oField);
    }

    for (const auto &poFeature : tmpLayer)
    {
        const char *pszCatalogName =
            poFeature->GetFieldAsString("catalog_name");
        const char *pszCatalogDBSchemas =
            poFeature->GetFieldAsString("catalog_db_schemas");
        if (pszCatalogDBSchemas)
        {
            CPLJSONDocument oDoc;
            if (oDoc.LoadMemory(pszCatalogDBSchemas))
            {
                auto oRoot = oDoc.GetRoot();
                if (oRoot.GetType() == CPLJSONObject::Type::Array)
                {
                    for (const auto &oSchema : oRoot.ToArray())
                    {
                        if (oSchema.GetType() == CPLJSONObject::Type::Object)
                        {
                            const std::string osSchemaName =
                                oSchema.GetString("schema_name");
                            const auto oTables =
                                oSchema.GetArray("db_schema_tables");
                            if (oTables.IsValid())
                            {
                                for (const auto &oTable : oTables)
                                {
                                    if (oTable.GetType() ==
                                        CPLJSONObject::Type::Object)
                                    {
                                        const std::string osTableName =
                                            oTable.GetString("table_name");
                                        const std::string osTableType =
                                            oTable.GetString("table_type");
                                        if (!osTableName.empty() &&
                                            osTableType != "index" &&
                                            osTableType != "trigger")
                                        {
                                            auto poFeat =
                                                std::make_unique<OGRFeature>(
                                                    poTableListLayer
                                                        ->GetLayerDefn());
                                            if (pszCatalogName)
                                                poFeat->SetField(
                                                    "catalog_name",
                                                    pszCatalogName);
                                            if (oSchema.GetObj("schema_name")
                                                    .IsValid())
                                                poFeat->SetField(
                                                    "schema_name",
                                                    osSchemaName.c_str());
                                            poFeat->SetField(
                                                "table_name",
                                                osTableName.c_str());
                                            if (oTable.GetObj("table_type")
                                                    .IsValid())
                                                poFeat->SetField(
                                                    "table_type",
                                                    osTableType.c_str());
                                            CPL_IGNORE_RET_VAL(
                                                poTableListLayer->CreateFeature(
                                                    std::move(poFeat)));
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    m_apoLayers.emplace_back(std::move(poTableListLayer));
    return m_apoLayers.back().get();
}

#undef ADBC_CALL

Coverage Report

Created: 2025-12-31 08:30