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 : * Copyright (c) 2024, Dewey Dunnington <dewey@voltrondata.com>
10 : *
11 : * SPDX-License-Identifier: MIT
12 : ****************************************************************************/
13 :
14 : #include "ogr_adbc.h"
15 : #include "ogradbcdrivercore.h"
16 : #include "memdataset.h"
17 : #include "ogr_p.h"
18 : #include "cpl_error.h"
19 : #include "cpl_json.h"
20 : #include "gdal_adbc.h"
21 :
22 : #include <algorithm>
23 :
24 : #if defined(OGR_ADBC_HAS_DRIVER_MANAGER)
25 : #ifdef __clang__
26 : #pragma clang diagnostic push
27 : #pragma clang diagnostic ignored "-Wdocumentation"
28 : #endif
29 : #include <arrow-adbc/adbc_driver_manager.h>
30 : #ifdef __clang__
31 : #pragma clang diagnostic pop
32 : #endif
33 : #endif
34 :
35 : // #define DEBUG_VERBOSE
36 :
37 : #define OGR_ADBC_VERSION ADBC_VERSION_1_1_0
38 : static_assert(sizeof(AdbcDriver) == ADBC_DRIVER_1_1_0_SIZE);
39 :
40 : namespace
41 : {
42 :
43 : #if !defined(OGR_ADBC_HAS_DRIVER_MANAGER)
44 0 : AdbcStatusCode OGRDuckDBLoadDriver(const char *driver_name, void *driver,
45 : struct AdbcError *error)
46 : {
47 0 : void *load_handle = CPLGetSymbol(driver_name, "duckdb_adbc_init");
48 0 : if (!load_handle)
49 : {
50 0 : return ADBC_STATUS_INTERNAL;
51 : }
52 :
53 0 : AdbcDriverInitFunc init_func =
54 : reinterpret_cast<AdbcDriverInitFunc>(load_handle);
55 0 : return init_func(OGR_ADBC_VERSION, driver, error);
56 : }
57 : #endif
58 :
59 1 : AdbcStatusCode OGRADBCLoadDriver(const char *driver_name,
60 : const char *entrypoint, void *driver,
61 : struct AdbcError *error)
62 : {
63 : GDALAdbcLoadDriverFunc load_driver_override =
64 1 : GDALGetAdbcLoadDriverOverride();
65 1 : if (load_driver_override)
66 : {
67 0 : return load_driver_override(driver_name, entrypoint, OGR_ADBC_VERSION,
68 0 : driver, error);
69 : }
70 : else
71 : {
72 : #if defined(OGR_ADBC_HAS_DRIVER_MANAGER)
73 : return AdbcLoadDriver(driver_name, entrypoint, OGR_ADBC_VERSION, driver,
74 : error);
75 : #else
76 : // If the driver is for DuckDB, use a minimal loading function, which
77 : // doesn't rely on the ADBC driver manager.
78 1 : if (strstr(driver_name, "duckdb"))
79 : {
80 0 : return OGRDuckDBLoadDriver(driver_name, driver, error);
81 : }
82 1 : return ADBC_STATUS_NOT_IMPLEMENTED;
83 : #endif
84 : }
85 : }
86 :
87 : } // namespace
88 :
89 : // Helper to wrap driver callbacks
90 : #define ADBC_CALL(func, ...) m_driver.func(__VA_ARGS__)
91 :
92 : /************************************************************************/
93 : /* ~OGRADBCDataset() */
94 : /************************************************************************/
95 :
96 2 : OGRADBCDataset::~OGRADBCDataset()
97 : {
98 1 : OGRADBCDataset::FlushCache(true);
99 :
100 : // Layers must be closed before the connection
101 1 : m_apoLayers.clear();
102 1 : OGRADBCError error;
103 1 : if (m_connection)
104 0 : ADBC_CALL(ConnectionRelease, m_connection.get(), error);
105 1 : error.clear();
106 1 : if (m_driver.release)
107 : {
108 0 : ADBC_CALL(DatabaseRelease, &m_database, error);
109 0 : m_driver.release(&m_driver, error);
110 : }
111 2 : }
112 :
113 : /************************************************************************/
114 : /* FlushCache() */
115 : /************************************************************************/
116 :
117 1 : CPLErr OGRADBCDataset::FlushCache(bool /* bAtClosing */)
118 : {
119 1 : CPLErr eErr = CE_None;
120 1 : for (auto &poLayer : m_apoLayers)
121 : {
122 0 : auto poADBCLayer = dynamic_cast<OGRADBCLayer *>(poLayer.get());
123 0 : if (poADBCLayer)
124 : {
125 0 : if (!poADBCLayer->RunDeferredCreation())
126 0 : eErr = CE_Failure;
127 : }
128 : }
129 :
130 1 : return eErr;
131 : }
132 :
133 : /************************************************************************/
134 : /* CreateLayer() */
135 : /************************************************************************/
136 :
137 : std::unique_ptr<OGRADBCLayer>
138 0 : OGRADBCDataset::CreateLayer(const char *pszStatement, const char *pszLayerName,
139 : bool bInternalUse)
140 : {
141 :
142 0 : CPLString osStatement(pszStatement);
143 0 : if (!m_osParquetFilename.empty())
144 : {
145 0 : const char *pszSrcLayerName = m_apoLayers.size() == 1
146 0 : ? m_apoLayers[0]->GetDescription()
147 0 : : pszLayerName;
148 : // Substitute the OGR layer name with the DuckDB expected filename,
149 : // single-quoted
150 : const std::string osFrom =
151 0 : std::string(" FROM ").append(pszSrcLayerName);
152 0 : const auto nPos = osStatement.ifind(osFrom);
153 0 : if (nPos != std::string::npos)
154 : {
155 : osStatement =
156 0 : osStatement.substr(0, nPos)
157 0 : .append(" FROM '")
158 0 : .append(OGRDuplicateCharacter(m_osParquetFilename, '\''))
159 0 : .append("'")
160 0 : .append(osStatement.substr(nPos + osFrom.size()));
161 : }
162 : else
163 : {
164 : const std::string osFrom2 =
165 0 : std::string(" FROM \"")
166 0 : .append(OGRDuplicateCharacter(pszSrcLayerName, '"'))
167 0 : .append("\"");
168 0 : const auto nPos2 = osStatement.ifind(osFrom2);
169 0 : if (nPos2 != std::string::npos)
170 : {
171 : osStatement =
172 0 : osStatement.substr(0, nPos2)
173 0 : .append(" FROM '")
174 : .append(
175 0 : OGRDuplicateCharacter(m_osParquetFilename, '\''))
176 0 : .append("'")
177 0 : .append(osStatement.substr(nPos2 + osFrom2.size()));
178 : }
179 : }
180 : }
181 :
182 : return std::make_unique<OGRADBCLayer>(this, pszLayerName, osStatement,
183 0 : bInternalUse);
184 : }
185 :
186 : /************************************************************************/
187 : /* CreateInternalLayer() */
188 : /************************************************************************/
189 :
190 : std::unique_ptr<OGRADBCLayer>
191 0 : OGRADBCDataset::CreateInternalLayer(const char *pszStatement)
192 : {
193 : #ifdef DEBUG_VERBOSE
194 : CPLDebug("ADBC", "%s", pszStatement);
195 : #endif
196 0 : return CreateLayer(pszStatement, "temp", true);
197 : }
198 :
199 : /************************************************************************/
200 : /* ExecuteSQL() */
201 : /************************************************************************/
202 :
203 0 : OGRLayer *OGRADBCDataset::ExecuteSQL(const char *pszStatement,
204 : OGRGeometry *poSpatialFilter,
205 : const char *pszDialect)
206 : {
207 0 : for (auto &poLayer : m_apoLayers)
208 : {
209 0 : auto poADBCLayer = dynamic_cast<OGRADBCLayer *>(poLayer.get());
210 0 : if (poADBCLayer)
211 0 : poADBCLayer->RunDeferredCreation();
212 : }
213 :
214 0 : if (pszDialect && pszDialect[0] != 0 && !EQUAL(pszDialect, "NATIVE"))
215 : {
216 0 : return GDALDataset::ExecuteSQL(pszStatement, poSpatialFilter,
217 0 : pszDialect);
218 : }
219 :
220 0 : std::string osStatement(pszStatement);
221 0 : for (const char *pszPrefix : {"SELECT * FROM ", "SELECT COUNT(*) FROM "})
222 : {
223 0 : if (m_bIsBigQuery && STARTS_WITH_CI(pszStatement, pszPrefix))
224 : {
225 0 : const auto nPos = osStatement.find(' ', strlen(pszPrefix));
226 : const std::string osTableName = osStatement.substr(
227 : strlen(pszPrefix),
228 0 : nPos == std::string::npos ? nPos : nPos - strlen(pszPrefix));
229 0 : auto poADBCLayer = dynamic_cast<OGRADBCBigQueryLayer *>(
230 0 : GetLayerByName(osTableName.c_str()));
231 0 : if (poADBCLayer)
232 : {
233 0 : std::string osDatasetId;
234 0 : std::string osTableId;
235 0 : if (poADBCLayer->GetBigQueryDatasetAndTableId(osDatasetId,
236 : osTableId))
237 : {
238 0 : std::string osNewStatement = pszPrefix;
239 0 : osNewStatement += '`';
240 0 : osNewStatement += OGRDuplicateCharacter(osDatasetId, '`');
241 0 : osNewStatement += "`.`";
242 0 : osNewStatement += OGRDuplicateCharacter(osTableId, '`');
243 0 : osNewStatement += '`';
244 0 : if (nPos != std::string::npos)
245 0 : osNewStatement += osStatement.substr(nPos);
246 0 : osStatement = std::move(osNewStatement);
247 : }
248 : }
249 0 : break;
250 : }
251 : }
252 :
253 0 : const char *pszLayerName = "RESULTSET";
254 0 : std::unique_ptr<OGRADBCLayer> poLayer;
255 0 : if (m_bIsBigQuery)
256 0 : poLayer = std::make_unique<OGRADBCBigQueryLayer>(
257 : this, pszLayerName, osStatement,
258 0 : /* bInternalUse = */ false);
259 : else
260 0 : poLayer = CreateLayer(osStatement.c_str(), pszLayerName, false);
261 0 : if (poLayer->GotError())
262 0 : return nullptr;
263 0 : if (poSpatialFilter)
264 : {
265 0 : if (poLayer->GetGeomType() == wkbNone)
266 0 : return nullptr;
267 0 : poLayer->SetSpatialFilter(poSpatialFilter);
268 : }
269 0 : return poLayer.release();
270 : }
271 :
272 : /************************************************************************/
273 : /* IsParquetExtension() */
274 : /************************************************************************/
275 :
276 1 : static bool IsParquetExtension(const char *pszStr)
277 : {
278 1 : const std::string osExt = CPLGetExtensionSafe(pszStr);
279 2 : return EQUAL(osExt.c_str(), "parquet") || EQUAL(osExt.c_str(), "parq");
280 : }
281 :
282 : /************************************************************************/
283 : /* Open() */
284 : /************************************************************************/
285 :
286 1 : bool OGRADBCDataset::Open(const GDALOpenInfo *poOpenInfo)
287 : {
288 1 : OGRADBCError error;
289 :
290 1 : const char *pszFilename = poOpenInfo->pszFilename;
291 1 : std::unique_ptr<GDALOpenInfo> poTmpOpenInfo;
292 1 : if (STARTS_WITH(pszFilename, "ADBC:"))
293 : {
294 1 : pszFilename += strlen("ADBC:");
295 1 : if (pszFilename[0])
296 : {
297 0 : poTmpOpenInfo = std::make_unique<GDALOpenInfo>(pszFilename,
298 0 : poOpenInfo->eAccess);
299 0 : poTmpOpenInfo->papszOpenOptions = poOpenInfo->papszOpenOptions;
300 0 : poOpenInfo = poTmpOpenInfo.get();
301 : }
302 : }
303 : const char *pszADBCDriverName =
304 1 : CSLFetchNameValue(poOpenInfo->papszOpenOptions, "ADBC_DRIVER");
305 1 : m_bIsDuckDBDataset = OGRADBCDriverIsDuckDB(poOpenInfo);
306 : const bool bIsSQLite3 =
307 2 : (pszADBCDriverName && EQUAL(pszADBCDriverName, "adbc_driver_sqlite")) ||
308 1 : OGRADBCDriverIsSQLite3(poOpenInfo);
309 : bool bIsParquet =
310 1 : OGRADBCDriverIsParquet(poOpenInfo) || IsParquetExtension(pszFilename);
311 1 : m_bIsBigQuery =
312 1 : pszADBCDriverName && strstr(pszADBCDriverName, "adbc_driver_bigquery");
313 1 : const char *pszSQL = CSLFetchNameValue(poOpenInfo->papszOpenOptions, "SQL");
314 1 : if (!bIsParquet && pszSQL)
315 : {
316 0 : CPLString osSQL(pszSQL);
317 0 : auto iPos = osSQL.find("FROM '");
318 0 : if (iPos != std::string::npos)
319 : {
320 0 : iPos += strlen("FROM '");
321 0 : const auto iPos2 = osSQL.find("'", iPos);
322 0 : if (iPos2 != std::string::npos)
323 : {
324 0 : std::string osFilename = osSQL.substr(iPos, iPos2 - iPos);
325 0 : if (IsParquetExtension(osFilename.c_str()))
326 : {
327 0 : m_osParquetFilename = std::move(osFilename);
328 0 : bIsParquet = true;
329 : }
330 : }
331 : }
332 : }
333 1 : const bool bIsPostgreSQL = STARTS_WITH(pszFilename, "postgresql://");
334 :
335 1 : if (!pszADBCDriverName)
336 : {
337 0 : if (m_bIsDuckDBDataset || bIsParquet)
338 : {
339 0 : pszADBCDriverName =
340 : #ifdef _WIN32
341 : "duckdb.dll"
342 : #elif defined(__MACH__) && defined(__APPLE__)
343 : "libduckdb.dylib"
344 : #else
345 : "libduckdb.so"
346 : #endif
347 : ;
348 : }
349 0 : else if (bIsPostgreSQL)
350 0 : pszADBCDriverName = "adbc_driver_postgresql";
351 0 : else if (bIsSQLite3)
352 : {
353 0 : pszADBCDriverName = "adbc_driver_sqlite";
354 : }
355 0 : else if (CSLFetchNameValue(poOpenInfo->papszOpenOptions,
356 0 : "BIGQUERY_PROJECT_ID") ||
357 0 : CSLFetchNameValue(poOpenInfo->papszOpenOptions,
358 0 : "BIGQUERY_DATASET_ID") ||
359 0 : CSLFetchNameValue(poOpenInfo->papszOpenOptions,
360 0 : "BIGQUERY_JSON_CREDENTIAL_STRING") ||
361 0 : CSLFetchNameValue(poOpenInfo->papszOpenOptions,
362 : "BIGQUERY_JSON_CREDENTIAL_FILE"))
363 : {
364 0 : m_bIsBigQuery = true;
365 0 : pszADBCDriverName = "adbc_driver_bigquery";
366 : }
367 : else
368 : {
369 0 : CPLError(CE_Failure, CPLE_AppDefined,
370 : "Open option ADBC_DRIVER must be specified");
371 0 : return false;
372 : }
373 : }
374 :
375 1 : if (poOpenInfo->eAccess == GA_Update && !m_bIsBigQuery)
376 : {
377 0 : return false;
378 : }
379 :
380 1 : eAccess = poOpenInfo->eAccess;
381 :
382 1 : m_bIsDuckDBDriver =
383 2 : m_bIsDuckDBDataset || bIsParquet ||
384 1 : (pszADBCDriverName && strstr(pszADBCDriverName, "duckdb"));
385 :
386 : // Load the driver
387 1 : if (m_bIsDuckDBDriver)
388 : {
389 0 : if (OGRADBCLoadDriver(pszADBCDriverName, "duckdb_adbc_init", &m_driver,
390 0 : error) != ADBC_STATUS_OK)
391 : {
392 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcLoadDriver() failed: %s",
393 : error.message());
394 0 : return false;
395 : }
396 : }
397 : else
398 : {
399 1 : if (OGRADBCLoadDriver(pszADBCDriverName, nullptr, &m_driver, error) !=
400 : ADBC_STATUS_OK)
401 : {
402 1 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcLoadDriver() failed: %s",
403 : error.message());
404 1 : return false;
405 : }
406 : }
407 :
408 : // Allocate the database
409 0 : if (ADBC_CALL(DatabaseNew, &m_database, error) != ADBC_STATUS_OK)
410 : {
411 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcDatabaseNew() failed: %s",
412 : error.message());
413 0 : return false;
414 : }
415 :
416 : // Set options
417 0 : if (m_bIsDuckDBDriver)
418 : {
419 0 : if (ADBC_CALL(DatabaseSetOption, &m_database, "path",
420 : bIsParquet ? ":memory:" : pszFilename,
421 0 : error) != ADBC_STATUS_OK)
422 : {
423 0 : CPLError(CE_Failure, CPLE_AppDefined,
424 : "AdbcDatabaseSetOption() failed: %s", error.message());
425 0 : return false;
426 : }
427 : }
428 0 : else if (pszFilename[0])
429 : {
430 0 : if (ADBC_CALL(DatabaseSetOption, &m_database, "uri", pszFilename,
431 0 : error) != ADBC_STATUS_OK)
432 : {
433 0 : CPLError(CE_Failure, CPLE_AppDefined,
434 : "AdbcDatabaseSetOption() failed: %s", error.message());
435 0 : return false;
436 : }
437 : }
438 :
439 0 : const auto GetAsOpenOptionOrConfigOption = [poOpenInfo](const char *pszKey)
440 : {
441 : const char *pszVal =
442 0 : CSLFetchNameValue(poOpenInfo->papszOpenOptions, pszKey);
443 0 : if (pszVal)
444 0 : return pszVal;
445 : // Below comments are for scripts/collect_config_options.py
446 : // CPLGetConfigOption("BIGQUERY_PROJECT_ID", nullptr);
447 : // CPLGetConfigOption("BIGQUERY_DATASET_ID", nullptr);
448 : // CPLGetConfigOption("BIGQUERY_JSON_CREDENTIAL_STRING", nullptr);
449 : // CPLGetConfigOption("BIGQUERY_JSON_CREDENTIAL_FILE", nullptr);
450 0 : return CPLGetConfigOption(pszKey, nullptr);
451 0 : };
452 :
453 : const char *pszBIGQUERY_PROJECT_ID =
454 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_PROJECT_ID");
455 0 : if (pszBIGQUERY_PROJECT_ID && pszBIGQUERY_PROJECT_ID[0])
456 : {
457 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
458 : "adbc.bigquery.sql.project_id", pszBIGQUERY_PROJECT_ID,
459 0 : error) != ADBC_STATUS_OK)
460 : {
461 0 : CPLError(CE_Failure, CPLE_AppDefined,
462 : "AdbcDatabaseSetOption() failed: %s", error.message());
463 0 : return false;
464 : }
465 : }
466 :
467 : const char *pszBIGQUERY_DATASET_ID =
468 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_DATASET_ID");
469 0 : if (pszBIGQUERY_DATASET_ID && pszBIGQUERY_DATASET_ID[0])
470 : {
471 0 : m_osBigQueryDatasetId = pszBIGQUERY_DATASET_ID;
472 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
473 : "adbc.bigquery.sql.dataset_id", pszBIGQUERY_DATASET_ID,
474 0 : error) != ADBC_STATUS_OK)
475 : {
476 0 : CPLError(CE_Failure, CPLE_AppDefined,
477 : "AdbcDatabaseSetOption() failed: %s", error.message());
478 0 : return false;
479 : }
480 : }
481 :
482 : const char *pszBIGQUERY_JSON_CREDENTIAL_STRING =
483 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_JSON_CREDENTIAL_STRING");
484 : const char *pszBIGQUERY_JSON_CREDENTIAL_FILE =
485 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_JSON_CREDENTIAL_FILE");
486 0 : if (pszBIGQUERY_JSON_CREDENTIAL_STRING &&
487 0 : pszBIGQUERY_JSON_CREDENTIAL_STRING[0])
488 : {
489 0 : if (pszBIGQUERY_JSON_CREDENTIAL_FILE &&
490 0 : pszBIGQUERY_JSON_CREDENTIAL_FILE[0])
491 : {
492 0 : CPLError(CE_Warning, CPLE_AppDefined,
493 : "BIGQUERY_JSON_CREDENTIAL_FILE ignored when "
494 : "BIGQUERY_JSON_CREDENTIAL_STRING is set");
495 : }
496 :
497 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
498 : "adbc.bigquery.sql.auth_credentials",
499 : "adbc.bigquery.sql.auth_type.json_credential_string",
500 0 : error) != ADBC_STATUS_OK)
501 : {
502 0 : CPLError(CE_Failure, CPLE_AppDefined,
503 : "AdbcDatabaseSetOption() failed: %s", error.message());
504 0 : return false;
505 : }
506 :
507 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
508 : "adbc.bigquery.sql.auth_credentials",
509 : pszBIGQUERY_JSON_CREDENTIAL_STRING,
510 0 : error) != ADBC_STATUS_OK)
511 : {
512 0 : CPLError(CE_Failure, CPLE_AppDefined,
513 : "AdbcDatabaseSetOption() failed: %s", error.message());
514 0 : return false;
515 : }
516 : }
517 0 : else if (pszBIGQUERY_JSON_CREDENTIAL_FILE &&
518 0 : pszBIGQUERY_JSON_CREDENTIAL_FILE[0])
519 : {
520 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
521 : "adbc.bigquery.sql.auth_credentials",
522 : "adbc.bigquery.sql.auth_type.json_credential_file",
523 0 : error) != ADBC_STATUS_OK)
524 : {
525 0 : CPLError(CE_Failure, CPLE_AppDefined,
526 : "AdbcDatabaseSetOption() failed: %s", error.message());
527 0 : return false;
528 : }
529 :
530 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
531 : "adbc.bigquery.sql.auth_credentials",
532 : pszBIGQUERY_JSON_CREDENTIAL_FILE,
533 0 : error) != ADBC_STATUS_OK)
534 : {
535 0 : CPLError(CE_Failure, CPLE_AppDefined,
536 : "AdbcDatabaseSetOption() failed: %s", error.message());
537 0 : return false;
538 : }
539 : }
540 :
541 0 : if (m_bIsBigQuery)
542 : {
543 0 : if (!(pszBIGQUERY_PROJECT_ID && pszBIGQUERY_PROJECT_ID[0]))
544 : {
545 0 : CPLError(CE_Failure, CPLE_AppDefined,
546 : "Required BIGQUERY_PROJECT_ID open option not provided");
547 0 : return false;
548 : }
549 0 : if (!(pszBIGQUERY_JSON_CREDENTIAL_STRING &&
550 0 : pszBIGQUERY_JSON_CREDENTIAL_STRING[0]) &&
551 0 : !(pszBIGQUERY_JSON_CREDENTIAL_FILE &&
552 0 : pszBIGQUERY_JSON_CREDENTIAL_FILE[0]))
553 : {
554 0 : CPLError(CE_Failure, CPLE_AppDefined,
555 : "Required BIGQUERY_JSON_CREDENTIAL_STRING or "
556 : "BIGQUERY_JSON_CREDENTIAL_FILE open option not provided");
557 0 : return false;
558 : }
559 : }
560 :
561 0 : for (const auto &[pszKey, pszValue] : cpl::IterateNameValue(
562 0 : static_cast<CSLConstList>(poOpenInfo->papszOpenOptions)))
563 : {
564 0 : if (STARTS_WITH_CI(pszKey, "ADBC_OPTION_"))
565 : {
566 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
567 : pszKey + strlen("ADBC_OPTION_"), pszValue,
568 0 : error) != ADBC_STATUS_OK)
569 : {
570 0 : CPLError(CE_Failure, CPLE_AppDefined,
571 : "AdbcDatabaseSetOption() failed: %s", error.message());
572 0 : return false;
573 : }
574 : }
575 : }
576 :
577 0 : if (ADBC_CALL(DatabaseInit, &m_database, error) != ADBC_STATUS_OK)
578 : {
579 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcDatabaseInit() failed: %s",
580 : error.message());
581 0 : return false;
582 : }
583 :
584 0 : m_connection = std::make_unique<AdbcConnection>();
585 0 : if (ADBC_CALL(ConnectionNew, m_connection.get(), error) != ADBC_STATUS_OK)
586 : {
587 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcConnectionNew() failed: %s",
588 : error.message());
589 0 : return false;
590 : }
591 :
592 0 : if (ADBC_CALL(ConnectionInit, m_connection.get(), &m_database, error) !=
593 : ADBC_STATUS_OK)
594 : {
595 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcConnectionInit() failed: %s",
596 : error.message());
597 0 : return false;
598 : }
599 :
600 0 : char **papszPreludeStatements = CSLFetchNameValueMultiple(
601 0 : poOpenInfo->papszOpenOptions, "PRELUDE_STATEMENTS");
602 0 : for (const char *pszStatement :
603 0 : cpl::Iterate(CSLConstList(papszPreludeStatements)))
604 : {
605 0 : CPL_IGNORE_RET_VAL(CreateInternalLayer(pszStatement)->GotError());
606 : }
607 0 : CSLDestroy(papszPreludeStatements);
608 0 : if (m_bIsDuckDBDriver && CPLTestBool(CPLGetConfigOption(
609 : "OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", "ON")))
610 : {
611 : auto poTmpLayer =
612 : CreateInternalLayer("SELECT 1 FROM duckdb_extensions() WHERE "
613 0 : "extension_name='spatial' AND loaded = false");
614 0 : if (!poTmpLayer->GotError() &&
615 0 : std::unique_ptr<OGRFeature>(poTmpLayer->GetNextFeature()) !=
616 : nullptr)
617 : {
618 0 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
619 0 : CPL_IGNORE_RET_VAL(CreateInternalLayer("LOAD spatial")->GotError());
620 : }
621 :
622 : poTmpLayer =
623 0 : CreateInternalLayer("SELECT 1 FROM duckdb_extensions() WHERE "
624 0 : "extension_name='spatial' AND loaded = true");
625 0 : m_bSpatialLoaded = !poTmpLayer->GotError() &&
626 0 : std::unique_ptr<OGRFeature>(
627 0 : poTmpLayer->GetNextFeature()) != nullptr;
628 : }
629 :
630 0 : std::string osLayerName = "RESULTSET";
631 0 : std::string osSQL;
632 0 : bool bIsParquetLayer = false;
633 0 : if (bIsParquet)
634 : {
635 0 : if (m_osParquetFilename.empty())
636 0 : m_osParquetFilename = pszFilename;
637 0 : osLayerName = CPLGetBasenameSafe(m_osParquetFilename.c_str());
638 0 : if (osLayerName == "*")
639 0 : osLayerName = CPLGetBasenameSafe(
640 0 : CPLGetDirnameSafe(m_osParquetFilename.c_str()).c_str());
641 0 : if (!pszSQL)
642 : {
643 : osSQL =
644 : CPLSPrintf("SELECT * FROM read_parquet('%s')",
645 0 : OGRDuplicateCharacter(pszFilename, '\'').c_str());
646 0 : pszSQL = osSQL.c_str();
647 0 : bIsParquetLayer = true;
648 : }
649 : }
650 :
651 0 : if (pszSQL)
652 : {
653 0 : if (pszSQL[0])
654 : {
655 0 : std::unique_ptr<OGRADBCLayer> poLayer;
656 0 : if ((bIsParquet || m_bIsDuckDBDataset) && m_bSpatialLoaded)
657 : {
658 0 : std::string osErrorMsg;
659 : {
660 0 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
661 0 : poLayer = CreateLayer(pszSQL, osLayerName.c_str(), false);
662 0 : if (poLayer->GotError())
663 0 : osErrorMsg = CPLGetLastErrorMsg();
664 : }
665 0 : if (poLayer->GotError())
666 : {
667 0 : CPLDebug("ADBC",
668 : "Connecting with 'LOAD spatial' did not work "
669 : "(%s). Retrying without it",
670 : osErrorMsg.c_str());
671 0 : ADBC_CALL(ConnectionRelease, m_connection.get(), error);
672 0 : m_connection.reset();
673 :
674 0 : ADBC_CALL(DatabaseRelease, &m_database, error);
675 0 : memset(&m_database, 0, sizeof(m_database));
676 :
677 0 : if (ADBC_CALL(DatabaseNew, &m_database, error) !=
678 : ADBC_STATUS_OK)
679 : {
680 0 : CPLError(CE_Failure, CPLE_AppDefined,
681 : "AdbcDatabaseNew() failed: %s",
682 : error.message());
683 0 : return false;
684 : }
685 0 : if (ADBC_CALL(DatabaseSetOption, &m_database, "path",
686 0 : ":memory:", error) != ADBC_STATUS_OK)
687 : {
688 0 : CPLError(CE_Failure, CPLE_AppDefined,
689 : "AdbcDatabaseSetOption() failed: %s",
690 : error.message());
691 0 : return false;
692 : }
693 :
694 0 : if (ADBC_CALL(DatabaseInit, &m_database, error) !=
695 : ADBC_STATUS_OK)
696 : {
697 0 : CPLError(CE_Failure, CPLE_AppDefined,
698 : "AdbcDatabaseInit() failed: %s",
699 : error.message());
700 0 : return false;
701 : }
702 :
703 0 : m_connection = std::make_unique<AdbcConnection>();
704 0 : if (ADBC_CALL(ConnectionNew, m_connection.get(), error) !=
705 : ADBC_STATUS_OK)
706 : {
707 0 : CPLError(CE_Failure, CPLE_AppDefined,
708 : "AdbcConnectionNew() failed: %s",
709 : error.message());
710 0 : return false;
711 : }
712 :
713 0 : if (ADBC_CALL(ConnectionInit, m_connection.get(),
714 0 : &m_database, error) != ADBC_STATUS_OK)
715 : {
716 0 : CPLError(CE_Failure, CPLE_AppDefined,
717 : "AdbcConnectionInit() failed: %s",
718 : error.message());
719 0 : return false;
720 : }
721 : }
722 : }
723 0 : if (!poLayer || poLayer->GotError())
724 : {
725 0 : if (m_bIsBigQuery)
726 0 : poLayer = std::make_unique<OGRADBCBigQueryLayer>(
727 0 : this, osLayerName.c_str(), pszSQL,
728 0 : /* bInternalUse = */ false);
729 : else
730 0 : poLayer = CreateLayer(pszSQL, osLayerName.c_str(), false);
731 0 : if (poLayer->GotError())
732 0 : return false;
733 : }
734 :
735 0 : poLayer->m_bIsParquetLayer = bIsParquetLayer;
736 0 : m_apoLayers.emplace_back(std::move(poLayer));
737 : }
738 : }
739 0 : else if (m_bIsDuckDBDataset || bIsSQLite3)
740 : {
741 : auto poLayerList = CreateInternalLayer(
742 0 : "SELECT name FROM sqlite_master WHERE type IN ('table', 'view')");
743 0 : if (poLayerList->GotError() ||
744 0 : poLayerList->GetLayerDefn()->GetFieldCount() != 1)
745 : {
746 0 : return false;
747 : }
748 :
749 0 : for (const auto &poFeature : poLayerList.get())
750 : {
751 0 : const char *pszLayerName = poFeature->GetFieldAsString(0);
752 0 : if (bIsSQLite3 && EQUAL(pszLayerName, "SpatialIndex"))
753 0 : continue;
754 : const std::string osStatement =
755 : CPLSPrintf("SELECT * FROM \"%s\"",
756 0 : OGRDuplicateCharacter(pszLayerName, '"').c_str());
757 : m_apoLayers.emplace_back(
758 0 : CreateLayer(osStatement.c_str(), pszLayerName, false));
759 0 : }
760 : }
761 0 : else if (bIsPostgreSQL)
762 : {
763 : auto poLayerList = CreateInternalLayer(
764 : "SELECT n.nspname, c.relname FROM pg_class c "
765 : "JOIN pg_namespace n ON c.relnamespace = n.oid "
766 : "AND c.relkind in ('r','v','m','f') "
767 : "AND n.nspname NOT IN ('pg_catalog', 'information_schema') "
768 0 : "ORDER BY c.oid");
769 0 : if (poLayerList->GotError() ||
770 0 : poLayerList->GetLayerDefn()->GetFieldCount() != 2)
771 : {
772 0 : return false;
773 : }
774 :
775 0 : for (const auto &poFeature : poLayerList.get())
776 : {
777 0 : const char *pszNamespace = poFeature->GetFieldAsString(0);
778 0 : const char *pszTableName = poFeature->GetFieldAsString(1);
779 : const std::string osStatement =
780 : CPLSPrintf("SELECT * FROM \"%s\".\"%s\"",
781 0 : OGRDuplicateCharacter(pszNamespace, '"').c_str(),
782 0 : OGRDuplicateCharacter(pszTableName, '"').c_str());
783 :
784 0 : m_apoLayers.emplace_back(CreateLayer(
785 : osStatement.c_str(),
786 0 : CPLSPrintf("%s.%s", pszNamespace, pszTableName), false));
787 : }
788 : }
789 0 : else if (m_bIsBigQuery)
790 : {
791 0 : if (!(pszBIGQUERY_DATASET_ID && pszBIGQUERY_DATASET_ID[0]))
792 : {
793 0 : CPLError(CE_Failure, CPLE_AppDefined,
794 : "Cannot list tables when BIGQUERY_DATASET_ID open option "
795 : "is not provided");
796 0 : return false;
797 : }
798 0 : const std::string s(pszBIGQUERY_DATASET_ID);
799 0 : if (!std::all_of(s.begin(), s.end(),
800 0 : [](char c) { return std::isalnum(c) || c == '_'; }))
801 : {
802 0 : CPLError(CE_Failure, CPLE_AppDefined,
803 : "Invalid characters found in BIGQUERY_DATASET_ID value");
804 0 : return false;
805 : }
806 : auto poLayerList = CreateInternalLayer(
807 : CPLSPrintf("SELECT table_name FROM %s.INFORMATION_SCHEMA.TABLES "
808 : "ORDER BY creation_time",
809 0 : pszBIGQUERY_DATASET_ID));
810 0 : if (poLayerList->GotError() ||
811 0 : poLayerList->GetLayerDefn()->GetFieldCount() != 1)
812 : {
813 0 : return false;
814 : }
815 :
816 0 : for (const auto &poFeature : poLayerList.get())
817 : {
818 0 : const char *pszTableName = poFeature->GetFieldAsString(0);
819 : const std::string osStatement = CPLSPrintf(
820 : "SELECT * FROM `%s`.`%s`",
821 0 : OGRDuplicateCharacter(pszBIGQUERY_DATASET_ID, '`').c_str(),
822 0 : OGRDuplicateCharacter(pszTableName, '`').c_str());
823 :
824 0 : m_apoLayers.emplace_back(std::make_unique<OGRADBCBigQueryLayer>(
825 : this, pszTableName, osStatement,
826 0 : /* bInternalUse = */ false));
827 : }
828 : }
829 :
830 0 : return true;
831 : }
832 :
833 : /************************************************************************/
834 : /* ICreateLayer() */
835 : /************************************************************************/
836 :
837 0 : OGRLayer *OGRADBCDataset::ICreateLayer(const char *pszName,
838 : const OGRGeomFieldDefn *poGeomFieldDefn,
839 : CSLConstList papszOptions)
840 : {
841 0 : if (!m_bIsBigQuery)
842 : {
843 0 : CPLError(CE_Failure, CPLE_NotSupported,
844 : "CreateLayer() only supported for BigQuery");
845 0 : return nullptr;
846 : }
847 0 : if (GetAccess() != GA_Update)
848 : {
849 0 : CPLError(
850 : CE_Failure, CPLE_NotSupported,
851 : "CreateLayer() only supported on datasets opened in update mode");
852 0 : return nullptr;
853 : }
854 0 : if (m_osBigQueryDatasetId.empty())
855 : {
856 0 : CPLError(CE_Failure, CPLE_AppDefined,
857 : "Open option BIGQUERY_DATASET_ID should be set");
858 0 : return nullptr;
859 : }
860 :
861 0 : if (GetLayerByName(pszName))
862 : {
863 0 : CPLError(CE_Failure, CPLE_AppDefined, "Table %s already exists",
864 : pszName);
865 0 : return nullptr;
866 : }
867 :
868 0 : if (poGeomFieldDefn)
869 : {
870 0 : const auto poSRS = poGeomFieldDefn->GetSpatialRef();
871 0 : if (poSRS && !poSRS->IsGeographic())
872 : {
873 0 : CPLError(CE_Failure, CPLE_NotSupported,
874 : "BigQuery only supports geographic CRS. Please reproject "
875 : "your layer to one (typically EPSG:4326)");
876 0 : return nullptr;
877 : }
878 : }
879 :
880 : const std::string osStatement = CPLSPrintf(
881 : "SELECT * FROM `%s`.`%s`",
882 0 : OGRDuplicateCharacter(m_osBigQueryDatasetId.c_str(), '`').c_str(),
883 0 : OGRDuplicateCharacter(pszName, '`').c_str());
884 :
885 : const char *pszFIDColName =
886 0 : CSLFetchNameValueDef(papszOptions, "FID", "ogc_fid");
887 : auto poLayer =
888 : std::make_unique<OGRADBCBigQueryLayer>(this, pszName, osStatement,
889 0 : /* bInternalUse = */ false);
890 0 : poLayer->SetDeferredCreation(pszFIDColName, poGeomFieldDefn);
891 0 : m_apoLayers.emplace_back(std::move(poLayer));
892 0 : return m_apoLayers.back().get();
893 : }
894 :
895 : /************************************************************************/
896 : /* DeleteLayer() */
897 : /************************************************************************/
898 :
899 0 : OGRErr OGRADBCDataset::DeleteLayer(int iLayer)
900 : {
901 0 : if (!m_bIsBigQuery)
902 : {
903 0 : CPLError(CE_Failure, CPLE_NotSupported,
904 : "DeleteLayer() only supported for BigQuery");
905 0 : return OGRERR_FAILURE;
906 : }
907 0 : if (GetAccess() != GA_Update)
908 : {
909 0 : CPLError(
910 : CE_Failure, CPLE_NotSupported,
911 : "DeleteLayer() only supported on datasets opened in update mode");
912 0 : return OGRERR_FAILURE;
913 : }
914 0 : if (iLayer < 0 || static_cast<size_t>(iLayer) >= m_apoLayers.size())
915 : {
916 0 : CPLError(CE_Failure, CPLE_AppDefined, "Invalid layer index");
917 0 : return OGRERR_FAILURE;
918 : }
919 :
920 : auto poADBCLayer =
921 0 : dynamic_cast<OGRADBCBigQueryLayer *>(m_apoLayers[iLayer].get());
922 0 : if (poADBCLayer && !poADBCLayer->m_bDeferredCreation)
923 : {
924 0 : std::string osDatasetId;
925 0 : std::string osTableId;
926 0 : if (!poADBCLayer->GetBigQueryDatasetAndTableId(osDatasetId, osTableId))
927 : {
928 0 : CPLError(CE_Failure, CPLE_NotSupported,
929 : "DeleteLayer(): cannot get dataset and table ID");
930 0 : return OGRERR_FAILURE;
931 : }
932 :
933 0 : std::string osSQL = "DROP TABLE `";
934 0 : osSQL += OGRDuplicateCharacter(osDatasetId.c_str(), '`');
935 0 : osSQL += "`.`";
936 0 : osSQL += OGRDuplicateCharacter(osTableId.c_str(), '`');
937 0 : osSQL += "`";
938 : // CPLDebug("ADBC", "%s", osSQL.c_str());
939 0 : if (CreateInternalLayer(osSQL.c_str())->GotError())
940 : {
941 0 : return OGRERR_FAILURE;
942 : }
943 : }
944 :
945 0 : m_apoLayers.erase(m_apoLayers.begin() + iLayer);
946 0 : return OGRERR_NONE;
947 : }
948 :
949 : /************************************************************************/
950 : /* TestCapability() */
951 : /************************************************************************/
952 :
953 0 : int OGRADBCDataset::TestCapability(const char *pszCap) const
954 : {
955 0 : if (EQUAL(pszCap, ODsCCreateLayer) || EQUAL(pszCap, ODsCDeleteLayer))
956 0 : return m_bIsBigQuery && eAccess == GA_Update;
957 0 : return false;
958 : }
959 :
960 : /************************************************************************/
961 : /* GetLayerByName() */
962 : /************************************************************************/
963 :
964 0 : OGRLayer *OGRADBCDataset::GetLayerByName(const char *pszName)
965 : {
966 0 : OGRLayer *poLayer = GDALDataset::GetLayerByName(pszName);
967 0 : if (poLayer || !EQUAL(pszName, "table_list"))
968 0 : return poLayer;
969 :
970 0 : OGRADBCError error;
971 0 : auto objectsStream = std::make_unique<OGRArrowArrayStream>();
972 0 : ADBC_CALL(ConnectionGetObjects, m_connection.get(),
973 : ADBC_OBJECT_DEPTH_TABLES,
974 : /* catalog = */ nullptr,
975 : /* db_schema = */ nullptr,
976 : /* table_name = */ nullptr,
977 : /* table_type = */ nullptr,
978 : /* column_name = */ nullptr, objectsStream->get(), error);
979 :
980 0 : ArrowSchema schema = {};
981 0 : if (objectsStream->get_schema(&schema) != 0)
982 : {
983 0 : CPLError(CE_Failure, CPLE_AppDefined, "get_schema() failed");
984 0 : return nullptr;
985 : }
986 :
987 0 : OGRADBCLayer tmpLayer(this, "", std::move(objectsStream), &schema,
988 0 : /* bInternalUse = */ true);
989 0 : const auto tmpLayerDefn = tmpLayer.GetLayerDefn();
990 0 : if (tmpLayerDefn->GetFieldIndex("catalog_name") < 0 ||
991 0 : tmpLayerDefn->GetFieldIndex("catalog_db_schemas") < 0)
992 : {
993 0 : return nullptr;
994 : }
995 :
996 : auto poTableListLayer =
997 0 : std::make_unique<OGRMemLayer>("table_list", nullptr, wkbNone);
998 : {
999 0 : OGRFieldDefn oField("catalog_name", OFTString);
1000 0 : poTableListLayer->CreateField(&oField);
1001 : }
1002 : {
1003 0 : OGRFieldDefn oField("schema_name", OFTString);
1004 0 : poTableListLayer->CreateField(&oField);
1005 : }
1006 : {
1007 0 : OGRFieldDefn oField("table_name", OFTString);
1008 0 : poTableListLayer->CreateField(&oField);
1009 : }
1010 : {
1011 0 : OGRFieldDefn oField("table_type", OFTString);
1012 0 : poTableListLayer->CreateField(&oField);
1013 : }
1014 :
1015 0 : for (const auto &poFeature : tmpLayer)
1016 : {
1017 : const char *pszCatalogName =
1018 0 : poFeature->GetFieldAsString("catalog_name");
1019 : const char *pszCatalogDBSchemas =
1020 0 : poFeature->GetFieldAsString("catalog_db_schemas");
1021 0 : if (pszCatalogDBSchemas)
1022 : {
1023 0 : CPLJSONDocument oDoc;
1024 0 : if (oDoc.LoadMemory(pszCatalogDBSchemas))
1025 : {
1026 0 : auto oRoot = oDoc.GetRoot();
1027 0 : if (oRoot.GetType() == CPLJSONObject::Type::Array)
1028 : {
1029 0 : for (const auto &oSchema : oRoot.ToArray())
1030 : {
1031 0 : if (oSchema.GetType() == CPLJSONObject::Type::Object)
1032 : {
1033 : const std::string osSchemaName =
1034 0 : oSchema.GetString("schema_name");
1035 : const auto oTables =
1036 0 : oSchema.GetArray("db_schema_tables");
1037 0 : if (oTables.IsValid())
1038 : {
1039 0 : for (const auto &oTable : oTables)
1040 : {
1041 0 : if (oTable.GetType() ==
1042 : CPLJSONObject::Type::Object)
1043 : {
1044 : const std::string osTableName =
1045 0 : oTable.GetString("table_name");
1046 : const std::string osTableType =
1047 0 : oTable.GetString("table_type");
1048 0 : if (!osTableName.empty() &&
1049 0 : osTableType != "index" &&
1050 0 : osTableType != "trigger")
1051 : {
1052 : OGRFeature oFeat(
1053 : poTableListLayer
1054 0 : ->GetLayerDefn());
1055 0 : if (pszCatalogName)
1056 0 : oFeat.SetField("catalog_name",
1057 : pszCatalogName);
1058 0 : if (oSchema.GetObj("schema_name")
1059 0 : .IsValid())
1060 0 : oFeat.SetField(
1061 : "schema_name",
1062 : osSchemaName.c_str());
1063 0 : oFeat.SetField("table_name",
1064 : osTableName.c_str());
1065 0 : if (oTable.GetObj("table_type")
1066 0 : .IsValid())
1067 0 : oFeat.SetField(
1068 : "table_type",
1069 : osTableType.c_str());
1070 0 : CPL_IGNORE_RET_VAL(
1071 0 : poTableListLayer->CreateFeature(
1072 : &oFeat));
1073 : }
1074 : }
1075 : }
1076 : }
1077 : }
1078 : }
1079 : }
1080 : }
1081 : }
1082 : }
1083 :
1084 0 : m_apoLayers.emplace_back(std::move(poTableListLayer));
1085 0 : return m_apoLayers.back().get();
1086 : }
1087 :
1088 : #undef ADBC_CALL
|
