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 : pszADBCDriverName =
390 0 : CPLGetConfigOption("OGR_ADBC_LIBDUCKDB", pszADBCDriverName);
391 0 : if (OGRADBCLoadDriver(pszADBCDriverName, "duckdb_adbc_init", &m_driver,
392 0 : error) != ADBC_STATUS_OK)
393 : {
394 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcLoadDriver() failed: %s",
395 : error.message());
396 0 : return false;
397 : }
398 : }
399 : else
400 : {
401 1 : if (OGRADBCLoadDriver(pszADBCDriverName, nullptr, &m_driver, error) !=
402 : ADBC_STATUS_OK)
403 : {
404 1 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcLoadDriver() failed: %s",
405 : error.message());
406 1 : return false;
407 : }
408 : }
409 :
410 : // Allocate the database
411 0 : if (ADBC_CALL(DatabaseNew, &m_database, error) != ADBC_STATUS_OK)
412 : {
413 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcDatabaseNew() failed: %s",
414 : error.message());
415 0 : return false;
416 : }
417 :
418 : // Set options
419 0 : if (m_bIsDuckDBDriver && pszFilename[0])
420 : {
421 : VSIStatBuf sStatBuf;
422 0 : if (!bIsParquet && VSIStat(pszFilename, &sStatBuf) != 0 &&
423 0 : strcmp(pszFilename, ":memory:") != 0)
424 : {
425 0 : CPLError(CE_Failure, CPLE_AppDefined, "%s does not exist",
426 : pszFilename);
427 0 : return false;
428 : }
429 0 : if (ADBC_CALL(DatabaseSetOption, &m_database, "path",
430 : bIsParquet ? ":memory:" : 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 0 : }
437 : }
438 0 : else if (pszFilename[0])
439 : {
440 0 : if (ADBC_CALL(DatabaseSetOption, &m_database, "uri", pszFilename,
441 0 : error) != ADBC_STATUS_OK)
442 : {
443 0 : CPLError(CE_Failure, CPLE_AppDefined,
444 : "AdbcDatabaseSetOption() failed: %s", error.message());
445 0 : return false;
446 : }
447 : }
448 :
449 0 : const auto GetAsOpenOptionOrConfigOption = [poOpenInfo](const char *pszKey)
450 : {
451 : const char *pszVal =
452 0 : CSLFetchNameValue(poOpenInfo->papszOpenOptions, pszKey);
453 0 : if (pszVal)
454 0 : return pszVal;
455 : // Below comments are for scripts/collect_config_options.py
456 : // CPLGetConfigOption("BIGQUERY_PROJECT_ID", nullptr);
457 : // CPLGetConfigOption("BIGQUERY_DATASET_ID", nullptr);
458 : // CPLGetConfigOption("BIGQUERY_JSON_CREDENTIAL_STRING", nullptr);
459 : // CPLGetConfigOption("BIGQUERY_JSON_CREDENTIAL_FILE", nullptr);
460 0 : return CPLGetConfigOption(pszKey, nullptr);
461 0 : };
462 :
463 : const char *pszBIGQUERY_PROJECT_ID =
464 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_PROJECT_ID");
465 0 : if (pszBIGQUERY_PROJECT_ID && pszBIGQUERY_PROJECT_ID[0])
466 : {
467 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
468 : "adbc.bigquery.sql.project_id", pszBIGQUERY_PROJECT_ID,
469 0 : error) != ADBC_STATUS_OK)
470 : {
471 0 : CPLError(CE_Failure, CPLE_AppDefined,
472 : "AdbcDatabaseSetOption() failed: %s", error.message());
473 0 : return false;
474 : }
475 : }
476 :
477 : const char *pszBIGQUERY_DATASET_ID =
478 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_DATASET_ID");
479 0 : if (pszBIGQUERY_DATASET_ID && pszBIGQUERY_DATASET_ID[0])
480 : {
481 0 : m_osBigQueryDatasetId = pszBIGQUERY_DATASET_ID;
482 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
483 : "adbc.bigquery.sql.dataset_id", pszBIGQUERY_DATASET_ID,
484 0 : error) != ADBC_STATUS_OK)
485 : {
486 0 : CPLError(CE_Failure, CPLE_AppDefined,
487 : "AdbcDatabaseSetOption() failed: %s", error.message());
488 0 : return false;
489 : }
490 : }
491 :
492 : const char *pszBIGQUERY_JSON_CREDENTIAL_STRING =
493 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_JSON_CREDENTIAL_STRING");
494 : const char *pszBIGQUERY_JSON_CREDENTIAL_FILE =
495 0 : GetAsOpenOptionOrConfigOption("BIGQUERY_JSON_CREDENTIAL_FILE");
496 0 : if (pszBIGQUERY_JSON_CREDENTIAL_STRING &&
497 0 : pszBIGQUERY_JSON_CREDENTIAL_STRING[0])
498 : {
499 0 : if (pszBIGQUERY_JSON_CREDENTIAL_FILE &&
500 0 : pszBIGQUERY_JSON_CREDENTIAL_FILE[0])
501 : {
502 0 : CPLError(CE_Warning, CPLE_AppDefined,
503 : "BIGQUERY_JSON_CREDENTIAL_FILE ignored when "
504 : "BIGQUERY_JSON_CREDENTIAL_STRING is set");
505 : }
506 :
507 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
508 : "adbc.bigquery.sql.auth_credentials",
509 : "adbc.bigquery.sql.auth_type.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 : if (ADBC_CALL(DatabaseSetOption, &m_database,
518 : "adbc.bigquery.sql.auth_credentials",
519 : pszBIGQUERY_JSON_CREDENTIAL_STRING,
520 0 : error) != ADBC_STATUS_OK)
521 : {
522 0 : CPLError(CE_Failure, CPLE_AppDefined,
523 : "AdbcDatabaseSetOption() failed: %s", error.message());
524 0 : return false;
525 : }
526 : }
527 0 : else if (pszBIGQUERY_JSON_CREDENTIAL_FILE &&
528 0 : pszBIGQUERY_JSON_CREDENTIAL_FILE[0])
529 : {
530 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
531 : "adbc.bigquery.sql.auth_credentials",
532 : "adbc.bigquery.sql.auth_type.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 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
541 : "adbc.bigquery.sql.auth_credentials",
542 : pszBIGQUERY_JSON_CREDENTIAL_FILE,
543 0 : error) != ADBC_STATUS_OK)
544 : {
545 0 : CPLError(CE_Failure, CPLE_AppDefined,
546 : "AdbcDatabaseSetOption() failed: %s", error.message());
547 0 : return false;
548 : }
549 : }
550 :
551 0 : if (m_bIsBigQuery)
552 : {
553 0 : if (!(pszBIGQUERY_PROJECT_ID && pszBIGQUERY_PROJECT_ID[0]))
554 : {
555 0 : CPLError(CE_Failure, CPLE_AppDefined,
556 : "Required BIGQUERY_PROJECT_ID open option not provided");
557 0 : return false;
558 : }
559 0 : if (!(pszBIGQUERY_JSON_CREDENTIAL_STRING &&
560 0 : pszBIGQUERY_JSON_CREDENTIAL_STRING[0]) &&
561 0 : !(pszBIGQUERY_JSON_CREDENTIAL_FILE &&
562 0 : pszBIGQUERY_JSON_CREDENTIAL_FILE[0]))
563 : {
564 0 : CPLError(CE_Failure, CPLE_AppDefined,
565 : "Required BIGQUERY_JSON_CREDENTIAL_STRING or "
566 : "BIGQUERY_JSON_CREDENTIAL_FILE open option not provided");
567 0 : return false;
568 : }
569 : }
570 :
571 0 : for (const auto &[pszKey, pszValue] : cpl::IterateNameValue(
572 0 : static_cast<CSLConstList>(poOpenInfo->papszOpenOptions)))
573 : {
574 0 : if (STARTS_WITH_CI(pszKey, "ADBC_OPTION_"))
575 : {
576 0 : if (ADBC_CALL(DatabaseSetOption, &m_database,
577 : pszKey + strlen("ADBC_OPTION_"), pszValue,
578 0 : error) != ADBC_STATUS_OK)
579 : {
580 0 : CPLError(CE_Failure, CPLE_AppDefined,
581 : "AdbcDatabaseSetOption() failed: %s", error.message());
582 0 : return false;
583 : }
584 : }
585 : }
586 :
587 0 : if (ADBC_CALL(DatabaseInit, &m_database, error) != ADBC_STATUS_OK)
588 : {
589 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcDatabaseInit() failed: %s",
590 : error.message());
591 0 : return false;
592 : }
593 :
594 0 : m_connection = std::make_unique<AdbcConnection>();
595 0 : if (ADBC_CALL(ConnectionNew, m_connection.get(), error) != ADBC_STATUS_OK)
596 : {
597 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcConnectionNew() failed: %s",
598 : error.message());
599 0 : return false;
600 : }
601 :
602 0 : if (ADBC_CALL(ConnectionInit, m_connection.get(), &m_database, error) !=
603 : ADBC_STATUS_OK)
604 : {
605 0 : CPLError(CE_Failure, CPLE_AppDefined, "AdbcConnectionInit() failed: %s",
606 : error.message());
607 0 : return false;
608 : }
609 :
610 0 : char **papszPreludeStatements = CSLFetchNameValueMultiple(
611 0 : poOpenInfo->papszOpenOptions, "PRELUDE_STATEMENTS");
612 0 : for (const char *pszStatement :
613 0 : cpl::Iterate(CSLConstList(papszPreludeStatements)))
614 : {
615 0 : CPL_IGNORE_RET_VAL(CreateInternalLayer(pszStatement)->GotError());
616 : }
617 0 : CSLDestroy(papszPreludeStatements);
618 0 : if (m_bIsDuckDBDriver && CPLTestBool(CPLGetConfigOption(
619 : "OGR_ADBC_AUTO_LOAD_DUCKDB_SPATIAL", "ON")))
620 : {
621 : auto poTmpLayer =
622 : CreateInternalLayer("SELECT 1 FROM duckdb_extensions() WHERE "
623 0 : "extension_name='spatial' AND loaded = false");
624 0 : if (!poTmpLayer->GotError() &&
625 0 : std::unique_ptr<OGRFeature>(poTmpLayer->GetNextFeature()) !=
626 : nullptr)
627 : {
628 0 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
629 0 : CPL_IGNORE_RET_VAL(CreateInternalLayer("LOAD spatial")->GotError());
630 : }
631 :
632 : poTmpLayer =
633 0 : CreateInternalLayer("SELECT 1 FROM duckdb_extensions() WHERE "
634 0 : "extension_name='spatial' AND loaded = true");
635 0 : m_bSpatialLoaded = !poTmpLayer->GotError() &&
636 0 : std::unique_ptr<OGRFeature>(
637 0 : poTmpLayer->GetNextFeature()) != nullptr;
638 : }
639 :
640 0 : std::string osLayerName = "RESULTSET";
641 0 : std::string osSQL;
642 0 : bool bIsParquetLayer = false;
643 0 : if (bIsParquet)
644 : {
645 0 : if (m_osParquetFilename.empty())
646 0 : m_osParquetFilename = pszFilename;
647 0 : osLayerName = CPLGetBasenameSafe(m_osParquetFilename.c_str());
648 0 : if (osLayerName == "*")
649 0 : osLayerName = CPLGetBasenameSafe(
650 0 : CPLGetDirnameSafe(m_osParquetFilename.c_str()).c_str());
651 0 : if (!pszSQL)
652 : {
653 : osSQL =
654 : CPLSPrintf("SELECT * FROM read_parquet('%s')",
655 0 : OGRDuplicateCharacter(pszFilename, '\'').c_str());
656 0 : pszSQL = osSQL.c_str();
657 0 : bIsParquetLayer = true;
658 : }
659 : }
660 :
661 0 : if (pszSQL)
662 : {
663 0 : if (pszSQL[0])
664 : {
665 0 : std::unique_ptr<OGRADBCLayer> poLayer;
666 0 : if ((bIsParquet || m_bIsDuckDBDataset) && m_bSpatialLoaded)
667 : {
668 0 : std::string osErrorMsg;
669 : {
670 0 : CPLErrorStateBackuper oBackuper(CPLQuietErrorHandler);
671 0 : poLayer = CreateLayer(pszSQL, osLayerName.c_str(), false);
672 0 : if (poLayer->GotError())
673 0 : osErrorMsg = CPLGetLastErrorMsg();
674 : }
675 0 : if (poLayer->GotError())
676 : {
677 0 : CPLDebug("ADBC",
678 : "Connecting with 'LOAD spatial' did not work "
679 : "(%s). Retrying without it",
680 : osErrorMsg.c_str());
681 0 : ADBC_CALL(ConnectionRelease, m_connection.get(), error);
682 0 : m_connection.reset();
683 :
684 0 : ADBC_CALL(DatabaseRelease, &m_database, error);
685 0 : memset(&m_database, 0, sizeof(m_database));
686 :
687 0 : if (ADBC_CALL(DatabaseNew, &m_database, error) !=
688 : ADBC_STATUS_OK)
689 : {
690 0 : CPLError(CE_Failure, CPLE_AppDefined,
691 : "AdbcDatabaseNew() failed: %s",
692 : error.message());
693 0 : return false;
694 : }
695 0 : if (ADBC_CALL(DatabaseSetOption, &m_database, "path",
696 0 : ":memory:", error) != ADBC_STATUS_OK)
697 : {
698 0 : CPLError(CE_Failure, CPLE_AppDefined,
699 : "AdbcDatabaseSetOption() failed: %s",
700 : error.message());
701 0 : return false;
702 : }
703 :
704 0 : if (ADBC_CALL(DatabaseInit, &m_database, error) !=
705 : ADBC_STATUS_OK)
706 : {
707 0 : CPLError(CE_Failure, CPLE_AppDefined,
708 : "AdbcDatabaseInit() failed: %s",
709 : error.message());
710 0 : return false;
711 : }
712 :
713 0 : m_connection = std::make_unique<AdbcConnection>();
714 0 : if (ADBC_CALL(ConnectionNew, m_connection.get(), error) !=
715 : ADBC_STATUS_OK)
716 : {
717 0 : CPLError(CE_Failure, CPLE_AppDefined,
718 : "AdbcConnectionNew() failed: %s",
719 : error.message());
720 0 : return false;
721 : }
722 :
723 0 : if (ADBC_CALL(ConnectionInit, m_connection.get(),
724 0 : &m_database, error) != ADBC_STATUS_OK)
725 : {
726 0 : CPLError(CE_Failure, CPLE_AppDefined,
727 : "AdbcConnectionInit() failed: %s",
728 : error.message());
729 0 : return false;
730 : }
731 : }
732 : }
733 0 : if (!poLayer || poLayer->GotError())
734 : {
735 0 : if (m_bIsBigQuery)
736 0 : poLayer = std::make_unique<OGRADBCBigQueryLayer>(
737 0 : this, osLayerName.c_str(), pszSQL,
738 0 : /* bInternalUse = */ false);
739 : else
740 0 : poLayer = CreateLayer(pszSQL, osLayerName.c_str(), false);
741 0 : if (poLayer->GotError())
742 0 : return false;
743 : }
744 :
745 0 : poLayer->m_bIsParquetLayer = bIsParquetLayer;
746 0 : m_apoLayers.emplace_back(std::move(poLayer));
747 : }
748 : }
749 0 : else if (m_bIsDuckDBDataset || bIsSQLite3)
750 : {
751 : auto poLayerList = CreateInternalLayer(
752 0 : "SELECT name FROM sqlite_master WHERE type IN ('table', 'view')");
753 0 : if (poLayerList->GotError() ||
754 0 : poLayerList->GetLayerDefn()->GetFieldCount() != 1)
755 : {
756 0 : return false;
757 : }
758 :
759 0 : for (const auto &poFeature : poLayerList.get())
760 : {
761 0 : const char *pszLayerName = poFeature->GetFieldAsString(0);
762 0 : if (bIsSQLite3 && EQUAL(pszLayerName, "SpatialIndex"))
763 0 : continue;
764 : const std::string osStatement =
765 : CPLSPrintf("SELECT * FROM \"%s\"",
766 0 : OGRDuplicateCharacter(pszLayerName, '"').c_str());
767 : m_apoLayers.emplace_back(
768 0 : CreateLayer(osStatement.c_str(), pszLayerName, false));
769 0 : }
770 : }
771 0 : else if (bIsPostgreSQL)
772 : {
773 : auto poLayerList = CreateInternalLayer(
774 : "SELECT n.nspname, c.relname FROM pg_class c "
775 : "JOIN pg_namespace n ON c.relnamespace = n.oid "
776 : "AND c.relkind in ('r','v','m','f') "
777 : "AND n.nspname NOT IN ('pg_catalog', 'information_schema') "
778 0 : "ORDER BY c.oid");
779 0 : if (poLayerList->GotError() ||
780 0 : poLayerList->GetLayerDefn()->GetFieldCount() != 2)
781 : {
782 0 : return false;
783 : }
784 :
785 0 : for (const auto &poFeature : poLayerList.get())
786 : {
787 0 : const char *pszNamespace = poFeature->GetFieldAsString(0);
788 0 : const char *pszTableName = poFeature->GetFieldAsString(1);
789 : const std::string osStatement =
790 : CPLSPrintf("SELECT * FROM \"%s\".\"%s\"",
791 0 : OGRDuplicateCharacter(pszNamespace, '"').c_str(),
792 0 : OGRDuplicateCharacter(pszTableName, '"').c_str());
793 :
794 0 : m_apoLayers.emplace_back(CreateLayer(
795 : osStatement.c_str(),
796 0 : CPLSPrintf("%s.%s", pszNamespace, pszTableName), false));
797 : }
798 : }
799 0 : else if (m_bIsBigQuery)
800 : {
801 0 : if (!(pszBIGQUERY_DATASET_ID && pszBIGQUERY_DATASET_ID[0]))
802 : {
803 0 : CPLError(CE_Failure, CPLE_AppDefined,
804 : "Cannot list tables when BIGQUERY_DATASET_ID open option "
805 : "is not provided");
806 0 : return false;
807 : }
808 0 : const std::string s(pszBIGQUERY_DATASET_ID);
809 0 : if (!std::all_of(s.begin(), s.end(),
810 0 : [](char c) { return std::isalnum(c) || c == '_'; }))
811 : {
812 0 : CPLError(CE_Failure, CPLE_AppDefined,
813 : "Invalid characters found in BIGQUERY_DATASET_ID value");
814 0 : return false;
815 : }
816 : auto poLayerList = CreateInternalLayer(
817 : CPLSPrintf("SELECT table_name FROM %s.INFORMATION_SCHEMA.TABLES "
818 : "ORDER BY creation_time",
819 0 : pszBIGQUERY_DATASET_ID));
820 0 : if (poLayerList->GotError() ||
821 0 : poLayerList->GetLayerDefn()->GetFieldCount() != 1)
822 : {
823 0 : return false;
824 : }
825 :
826 0 : for (const auto &poFeature : poLayerList.get())
827 : {
828 0 : const char *pszTableName = poFeature->GetFieldAsString(0);
829 : const std::string osStatement = CPLSPrintf(
830 : "SELECT * FROM `%s`.`%s`",
831 0 : OGRDuplicateCharacter(pszBIGQUERY_DATASET_ID, '`').c_str(),
832 0 : OGRDuplicateCharacter(pszTableName, '`').c_str());
833 :
834 0 : m_apoLayers.emplace_back(std::make_unique<OGRADBCBigQueryLayer>(
835 : this, pszTableName, osStatement,
836 0 : /* bInternalUse = */ false));
837 : }
838 : }
839 :
840 0 : return true;
841 : }
842 :
843 : /************************************************************************/
844 : /* ICreateLayer() */
845 : /************************************************************************/
846 :
847 0 : OGRLayer *OGRADBCDataset::ICreateLayer(const char *pszName,
848 : const OGRGeomFieldDefn *poGeomFieldDefn,
849 : CSLConstList papszOptions)
850 : {
851 0 : if (!m_bIsBigQuery)
852 : {
853 0 : CPLError(CE_Failure, CPLE_NotSupported,
854 : "CreateLayer() only supported for BigQuery");
855 0 : return nullptr;
856 : }
857 0 : if (GetAccess() != GA_Update)
858 : {
859 0 : CPLError(
860 : CE_Failure, CPLE_NotSupported,
861 : "CreateLayer() only supported on datasets opened in update mode");
862 0 : return nullptr;
863 : }
864 0 : if (m_osBigQueryDatasetId.empty())
865 : {
866 0 : CPLError(CE_Failure, CPLE_AppDefined,
867 : "Open option BIGQUERY_DATASET_ID should be set");
868 0 : return nullptr;
869 : }
870 :
871 0 : if (GetLayerByName(pszName))
872 : {
873 0 : CPLError(CE_Failure, CPLE_AppDefined, "Table %s already exists",
874 : pszName);
875 0 : return nullptr;
876 : }
877 :
878 0 : if (poGeomFieldDefn)
879 : {
880 0 : const auto poSRS = poGeomFieldDefn->GetSpatialRef();
881 0 : if (poSRS && !poSRS->IsGeographic())
882 : {
883 0 : CPLError(CE_Failure, CPLE_NotSupported,
884 : "BigQuery only supports geographic CRS. Please reproject "
885 : "your layer to one (typically EPSG:4326)");
886 0 : return nullptr;
887 : }
888 : }
889 :
890 : const std::string osStatement = CPLSPrintf(
891 : "SELECT * FROM `%s`.`%s`",
892 0 : OGRDuplicateCharacter(m_osBigQueryDatasetId.c_str(), '`').c_str(),
893 0 : OGRDuplicateCharacter(pszName, '`').c_str());
894 :
895 : const char *pszFIDColName =
896 0 : CSLFetchNameValueDef(papszOptions, "FID", "ogc_fid");
897 : auto poLayer =
898 : std::make_unique<OGRADBCBigQueryLayer>(this, pszName, osStatement,
899 0 : /* bInternalUse = */ false);
900 0 : poLayer->SetDeferredCreation(pszFIDColName, poGeomFieldDefn);
901 0 : m_apoLayers.emplace_back(std::move(poLayer));
902 0 : return m_apoLayers.back().get();
903 : }
904 :
905 : /************************************************************************/
906 : /* DeleteLayer() */
907 : /************************************************************************/
908 :
909 0 : OGRErr OGRADBCDataset::DeleteLayer(int iLayer)
910 : {
911 0 : if (!m_bIsBigQuery)
912 : {
913 0 : CPLError(CE_Failure, CPLE_NotSupported,
914 : "DeleteLayer() only supported for BigQuery");
915 0 : return OGRERR_FAILURE;
916 : }
917 0 : if (GetAccess() != GA_Update)
918 : {
919 0 : CPLError(
920 : CE_Failure, CPLE_NotSupported,
921 : "DeleteLayer() only supported on datasets opened in update mode");
922 0 : return OGRERR_FAILURE;
923 : }
924 0 : if (iLayer < 0 || static_cast<size_t>(iLayer) >= m_apoLayers.size())
925 : {
926 0 : CPLError(CE_Failure, CPLE_AppDefined, "Invalid layer index");
927 0 : return OGRERR_FAILURE;
928 : }
929 :
930 : auto poADBCLayer =
931 0 : dynamic_cast<OGRADBCBigQueryLayer *>(m_apoLayers[iLayer].get());
932 0 : if (poADBCLayer && !poADBCLayer->m_bDeferredCreation)
933 : {
934 0 : std::string osDatasetId;
935 0 : std::string osTableId;
936 0 : if (!poADBCLayer->GetBigQueryDatasetAndTableId(osDatasetId, osTableId))
937 : {
938 0 : CPLError(CE_Failure, CPLE_NotSupported,
939 : "DeleteLayer(): cannot get dataset and table ID");
940 0 : return OGRERR_FAILURE;
941 : }
942 :
943 0 : std::string osSQL = "DROP TABLE `";
944 0 : osSQL += OGRDuplicateCharacter(osDatasetId.c_str(), '`');
945 0 : osSQL += "`.`";
946 0 : osSQL += OGRDuplicateCharacter(osTableId.c_str(), '`');
947 0 : osSQL += "`";
948 : // CPLDebug("ADBC", "%s", osSQL.c_str());
949 0 : if (CreateInternalLayer(osSQL.c_str())->GotError())
950 : {
951 0 : return OGRERR_FAILURE;
952 : }
953 : }
954 :
955 0 : m_apoLayers.erase(m_apoLayers.begin() + iLayer);
956 0 : return OGRERR_NONE;
957 : }
958 :
959 : /************************************************************************/
960 : /* TestCapability() */
961 : /************************************************************************/
962 :
963 0 : int OGRADBCDataset::TestCapability(const char *pszCap) const
964 : {
965 0 : if (EQUAL(pszCap, ODsCCreateLayer) || EQUAL(pszCap, ODsCDeleteLayer))
966 0 : return m_bIsBigQuery && eAccess == GA_Update;
967 0 : return false;
968 : }
969 :
970 : /************************************************************************/
971 : /* GetLayerByName() */
972 : /************************************************************************/
973 :
974 0 : OGRLayer *OGRADBCDataset::GetLayerByName(const char *pszName)
975 : {
976 0 : OGRLayer *poLayer = GDALDataset::GetLayerByName(pszName);
977 0 : if (poLayer || !EQUAL(pszName, "table_list"))
978 0 : return poLayer;
979 :
980 0 : OGRADBCError error;
981 0 : auto objectsStream = std::make_unique<OGRArrowArrayStream>();
982 0 : ADBC_CALL(ConnectionGetObjects, m_connection.get(),
983 : ADBC_OBJECT_DEPTH_TABLES,
984 : /* catalog = */ nullptr,
985 : /* db_schema = */ nullptr,
986 : /* table_name = */ nullptr,
987 : /* table_type = */ nullptr,
988 : /* column_name = */ nullptr, objectsStream->get(), error);
989 :
990 0 : ArrowSchema schema = {};
991 0 : if (objectsStream->get_schema(&schema) != 0)
992 : {
993 0 : CPLError(CE_Failure, CPLE_AppDefined, "get_schema() failed");
994 0 : return nullptr;
995 : }
996 :
997 0 : OGRADBCLayer tmpLayer(this, "", std::move(objectsStream), &schema,
998 0 : /* bInternalUse = */ true);
999 0 : const auto tmpLayerDefn = tmpLayer.GetLayerDefn();
1000 0 : if (tmpLayerDefn->GetFieldIndex("catalog_name") < 0 ||
1001 0 : tmpLayerDefn->GetFieldIndex("catalog_db_schemas") < 0)
1002 : {
1003 0 : return nullptr;
1004 : }
1005 :
1006 : auto poTableListLayer =
1007 0 : std::make_unique<OGRMemLayer>("table_list", nullptr, wkbNone);
1008 : {
1009 0 : OGRFieldDefn oField("catalog_name", OFTString);
1010 0 : poTableListLayer->CreateField(&oField);
1011 : }
1012 : {
1013 0 : OGRFieldDefn oField("schema_name", OFTString);
1014 0 : poTableListLayer->CreateField(&oField);
1015 : }
1016 : {
1017 0 : OGRFieldDefn oField("table_name", OFTString);
1018 0 : poTableListLayer->CreateField(&oField);
1019 : }
1020 : {
1021 0 : OGRFieldDefn oField("table_type", OFTString);
1022 0 : poTableListLayer->CreateField(&oField);
1023 : }
1024 :
1025 0 : for (const auto &poFeature : tmpLayer)
1026 : {
1027 : const char *pszCatalogName =
1028 0 : poFeature->GetFieldAsString("catalog_name");
1029 : const char *pszCatalogDBSchemas =
1030 0 : poFeature->GetFieldAsString("catalog_db_schemas");
1031 0 : if (pszCatalogDBSchemas)
1032 : {
1033 0 : CPLJSONDocument oDoc;
1034 0 : if (oDoc.LoadMemory(pszCatalogDBSchemas))
1035 : {
1036 0 : auto oRoot = oDoc.GetRoot();
1037 0 : if (oRoot.GetType() == CPLJSONObject::Type::Array)
1038 : {
1039 0 : for (const auto &oSchema : oRoot.ToArray())
1040 : {
1041 0 : if (oSchema.GetType() == CPLJSONObject::Type::Object)
1042 : {
1043 : const std::string osSchemaName =
1044 0 : oSchema.GetString("schema_name");
1045 : const auto oTables =
1046 0 : oSchema.GetArray("db_schema_tables");
1047 0 : if (oTables.IsValid())
1048 : {
1049 0 : for (const auto &oTable : oTables)
1050 : {
1051 0 : if (oTable.GetType() ==
1052 : CPLJSONObject::Type::Object)
1053 : {
1054 : const std::string osTableName =
1055 0 : oTable.GetString("table_name");
1056 : const std::string osTableType =
1057 0 : oTable.GetString("table_type");
1058 0 : if (!osTableName.empty() &&
1059 0 : osTableType != "index" &&
1060 0 : osTableType != "trigger")
1061 : {
1062 : auto poFeat =
1063 : std::make_unique<OGRFeature>(
1064 : poTableListLayer
1065 0 : ->GetLayerDefn());
1066 0 : if (pszCatalogName)
1067 0 : poFeat->SetField(
1068 : "catalog_name",
1069 : pszCatalogName);
1070 0 : if (oSchema.GetObj("schema_name")
1071 0 : .IsValid())
1072 0 : poFeat->SetField(
1073 : "schema_name",
1074 : osSchemaName.c_str());
1075 0 : poFeat->SetField(
1076 : "table_name",
1077 : osTableName.c_str());
1078 0 : if (oTable.GetObj("table_type")
1079 0 : .IsValid())
1080 0 : poFeat->SetField(
1081 : "table_type",
1082 : osTableType.c_str());
1083 0 : CPL_IGNORE_RET_VAL(
1084 0 : poTableListLayer->CreateFeature(
1085 0 : std::move(poFeat)));
1086 : }
1087 : }
1088 : }
1089 : }
1090 : }
1091 : }
1092 : }
1093 : }
1094 : }
1095 : }
1096 :
1097 0 : m_apoLayers.emplace_back(std::move(poTableListLayer));
1098 0 : return m_apoLayers.back().get();
1099 : }
1100 :
1101 : #undef ADBC_CALL
|
