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