├── .editorconfig ├── .github └── workflows │ ├── MainDistributionPipeline.yml │ └── scheduled-build.yml ├── .gitignore ├── .gitmodules ├── .vscode └── settings.json ├── CMakeLists.txt ├── LICENSE ├── Makefile ├── docs ├── README.md └── duckdb-airport-1.jpg ├── extension_config.cmake ├── scripts ├── bootstrap-template.py └── extension-upload.sh ├── src ├── airport_action.cpp ├── airport_constraints.cpp ├── airport_extension.cpp ├── airport_flight_exception.cpp ├── airport_flight_statistics.cpp ├── airport_flight_stream.cpp ├── airport_json_common.cpp ├── airport_json_serializer.cpp ├── airport_list_flights.cpp ├── airport_optimizer.cpp ├── airport_request_headers.cpp ├── airport_scalar_function.cpp ├── airport_schema_utils.cpp ├── airport_secrets.cpp ├── airport_take_flight.cpp ├── include │ ├── airport_constraints.hpp │ ├── airport_extension.hpp │ ├── airport_flight_exception.hpp │ ├── airport_flight_statistics.hpp │ ├── airport_flight_stream.hpp │ ├── airport_json_common.hpp │ ├── airport_json_serializer.hpp │ ├── airport_location_descriptor.hpp │ ├── airport_macros.hpp │ ├── airport_optimizer.hpp │ ├── airport_request_headers.hpp │ ├── airport_scalar_function.hpp │ ├── airport_schema_utils.hpp │ ├── airport_secrets.hpp │ ├── airport_take_flight.hpp │ └── storage │ │ ├── airport_alter_parameters.hpp │ │ ├── airport_catalog.hpp │ │ ├── airport_catalog_api.hpp │ │ ├── airport_catalog_set.hpp │ │ ├── airport_catalog_set_base.hpp │ │ ├── airport_curl_pool.hpp │ │ ├── airport_delete.hpp │ │ ├── airport_delete_parameterized.hpp │ │ ├── airport_exchange.hpp │ │ ├── airport_insert.hpp │ │ ├── airport_scalar_function_set.hpp │ │ ├── airport_schema_entry.hpp │ │ ├── airport_schema_set.hpp │ │ ├── airport_table_entry.hpp │ │ ├── airport_table_function_set.hpp │ │ ├── airport_table_set.hpp │ │ ├── airport_transaction.hpp │ │ ├── airport_transaction_manager.hpp │ │ ├── airport_update.hpp │ │ └── airport_update_parameterized.hpp └── storage │ ├── airport_catalog.cpp │ ├── airport_catalog_api.cpp │ ├── airport_catalog_set.cpp │ ├── airport_clear_cache.cpp │ ├── airport_curl_pool.cpp │ ├── airport_delete.cpp │ ├── airport_delete_parameterized.cpp │ ├── airport_exchange.cpp │ ├── airport_insert.cpp │ ├── airport_scalar_function_set.cpp │ ├── airport_schema_entry.cpp │ ├── airport_schema_set.cpp │ ├── airport_table_entry.cpp │ ├── airport_table_function_set.cpp │ ├── airport_table_set.cpp │ ├── airport_transaction.cpp │ ├── airport_transaction_manager.cpp │ ├── airport_update.cpp │ └── airport_update_parameterized.cpp ├── test ├── README.md └── sql │ └── airport.test └── vcpkg.json /.editorconfig: -------------------------------------------------------------------------------- 1 | duckdb/.editorconfig -------------------------------------------------------------------------------- /.github/workflows/MainDistributionPipeline.yml: -------------------------------------------------------------------------------- 1 | # 2 | # This workflow calls the main distribution pipeline from DuckDB to build, test and (optionally) release the extension 3 | # 4 | name: Main Extension Distribution Pipeline 5 | on: 6 | push: 7 | pull_request: 8 | workflow_dispatch: 9 | schedule: 10 | - cron: '0 2 * * *' # Runs every night at 02:00 UTC 11 | 12 | concurrency: 13 | group: ${{ github.workflow }}-${{ github.ref }}-${{ github.head_ref || '' }}-${{ github.base_ref || '' }}-${{ github.ref != 'refs/heads/main' || github.sha }} 14 | cancel-in-progress: false 15 | 16 | jobs: 17 | duckdb-stable-build: 18 | name: Build extension binaries 19 | uses: duckdb/extension-ci-tools/.github/workflows/_extension_distribution.yml@main 20 | with: 21 | duckdb_version: main 22 | extension_name: airport 23 | ci_tools_version: main 24 | extra_toolchains: parser_tools 25 | exclude_archs: "wasm_mvp;wasm_eh;wasm_threads;" 26 | secrets: 27 | VCPKG_CACHING_AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} 28 | VCPKG_CACHING_AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} 29 | 30 | # duckdb-stable-deploy: 31 | # name: Deploy extension binaries 32 | # needs: duckdb-stable-build 33 | # uses: duckdb/extension-ci-tools/.github/workflows/_extension_distribution.yml@main 34 | # secrets: inherit 35 | # with: 36 | # duckdb_version: main 37 | # extension_name: airport 38 | # ci_tools_version: main 39 | # deploy_latest: ${{ startsWith(github.ref, 'refs/tags/v') || github.ref == 'refs/heads/main' }} 40 | # exclude_archs: "wasm_mvp;wasm_eh;wasm_threads;windows_amd64_rtools" 41 | -------------------------------------------------------------------------------- /.github/workflows/scheduled-build.yml: -------------------------------------------------------------------------------- 1 | name: Scheduled Trigger for v1.2 2 | 3 | # on: 4 | # schedule: 5 | # - cron: '0 12 * * *' # Runs at 12:00 UTC every day 6 | # workflow_dispatch: # Allows manual trigger 7 | 8 | jobs: 9 | trigger: 10 | runs-on: ubuntu-latest 11 | permissions: 12 | actions: write # Allow triggering workflows 13 | steps: 14 | - name: Checkout repository # Required for gh to work 15 | uses: actions/checkout@v4 16 | 17 | - name: Install GitHub CLI 18 | run: | 19 | sudo apt update && sudo apt install gh -y 20 | 21 | - name: Authenticate GH CLI 22 | run: | 23 | echo "${{ secrets.GITHUB_TOKEN }}" | gh auth login --with-token 24 | 25 | - name: Trigger Workflow on my-branch 26 | run: | 27 | gh workflow run MainDistributionPipeline.yml --ref v1.2 28 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | vcpkg_installed/ 2 | build 3 | .idea 4 | cmake-build-debug 5 | duckdb_unittest_tempdir/ 6 | .DS_Store 7 | testext 8 | test/python/__pycache__/ 9 | .Rhistory 10 | vcpkg 11 | -------------------------------------------------------------------------------- /.gitmodules: -------------------------------------------------------------------------------- 1 | [submodule "duckdb"] 2 | path = duckdb 3 | url = https://github.com/duckdb/duckdb 4 | branch = main 5 | [submodule "extension-ci-tools"] 6 | path = extension-ci-tools 7 | url = https://github.com/duckdb/extension-ci-tools.git 8 | branch = main 9 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | { 2 | "files.exclude": { 3 | "vcpkg/downloads": true, 4 | "vcpkg/packages": true, 5 | "vcpkg/ports": true, 6 | "vcpkg/scripts": true, 7 | "vcpkg/toolsrc": true, 8 | "vcpkg/versions": true, 9 | "vcpkg/triplets": true, 10 | "vcpkg_installed/**": true, 11 | "duckdb_unittest_tempdir": true, 12 | "build/**": true, 13 | ".tmp/**": true, 14 | "**/.git": true, 15 | "**/.svn": true, 16 | "**/.hg": true, 17 | "**/.trunk": true, 18 | "**/CVS": true, 19 | "**/.DS_Store": true, 20 | "**/Thumbs.db": true 21 | }, 22 | "cmake.environment": { 23 | "GEN": "ninja", 24 | "VCPKG_TOOLCHAIN_PATH": "${workspaceFolder}/vcpkg/scripts/buildsystems/vcpkg.cmake", 25 | }, 26 | "workbench.colorCustomizations": { 27 | "activityBar.activeBackground": "#81cb55", 28 | "activityBar.background": "#81cb55", 29 | "activityBar.foreground": "#15202b", 30 | "activityBar.inactiveForeground": "#15202b99", 31 | "activityBarBadge.background": "#3c6ec3", 32 | "activityBarBadge.foreground": "#e7e7e7", 33 | "commandCenter.border": "#15202b99", 34 | "sash.hoverBorder": "#81cb55", 35 | "statusBar.background": "#66b538", 36 | "statusBar.foreground": "#15202b", 37 | "statusBarItem.hoverBackground": "#508e2c", 38 | "statusBarItem.remoteBackground": "#66b538", 39 | "statusBarItem.remoteForeground": "#15202b", 40 | "titleBar.activeBackground": "#66b538", 41 | "titleBar.activeForeground": "#15202b", 42 | "titleBar.inactiveBackground": "#66b53899", 43 | "titleBar.inactiveForeground": "#15202b99" 44 | }, 45 | "peacock.color": "#66b538", 46 | "files.associations": { 47 | "optional": "cpp", 48 | "*.inc": "cpp", 49 | "__hash_table": "cpp", 50 | "__split_buffer": "cpp", 51 | "__tree": "cpp", 52 | "array": "cpp", 53 | "bitset": "cpp", 54 | "deque": "cpp", 55 | "hash_map": "cpp", 56 | "initializer_list": "cpp", 57 | "list": "cpp", 58 | "map": "cpp", 59 | "queue": "cpp", 60 | "regex": "cpp", 61 | "set": "cpp", 62 | "span": "cpp", 63 | "stack": "cpp", 64 | "string": "cpp", 65 | "string_view": "cpp", 66 | "unordered_map": "cpp", 67 | "unordered_set": "cpp", 68 | "valarray": "cpp", 69 | "vector": "cpp", 70 | "*.ipp": "cpp", 71 | "stdexcept": "cpp", 72 | "iomanip": "cpp", 73 | "__bit_reference": "cpp", 74 | "__locale": "cpp", 75 | "__node_handle": "cpp", 76 | "__verbose_abort": "cpp", 77 | "cctype": "cpp", 78 | "charconv": "cpp", 79 | "cmath": "cpp", 80 | "condition_variable": "cpp", 81 | "cstddef": "cpp", 82 | "cstdint": "cpp", 83 | "cstdio": "cpp", 84 | "cstdlib": "cpp", 85 | "cstring": "cpp", 86 | "ctime": "cpp", 87 | "cwchar": "cpp", 88 | "forward_list": "cpp", 89 | "fstream": "cpp", 90 | "future": "cpp", 91 | "ios": "cpp", 92 | "iosfwd": "cpp", 93 | "iostream": "cpp", 94 | "istream": "cpp", 95 | "limits": "cpp", 96 | "locale": "cpp", 97 | "mutex": "cpp", 98 | "new": "cpp", 99 | "ostream": "cpp", 100 | "print": "cpp", 101 | "ratio": "cpp", 102 | "sstream": "cpp", 103 | "streambuf": "cpp", 104 | "tuple": "cpp", 105 | "typeinfo": "cpp", 106 | "variant": "cpp", 107 | "algorithm": "cpp", 108 | "__threading_support": "cpp", 109 | "any": "cpp", 110 | "cfenv": "cpp", 111 | "cinttypes": "cpp", 112 | "clocale": "cpp", 113 | "codecvt": "cpp", 114 | "complex": "cpp", 115 | "csignal": "cpp", 116 | "cstdarg": "cpp", 117 | "cwctype": "cpp", 118 | "execution": "cpp", 119 | "memory": "cpp", 120 | "shared_mutex": "cpp", 121 | "source_location": "cpp", 122 | "strstream": "cpp", 123 | "typeindex": "cpp" 124 | } 125 | } -------------------------------------------------------------------------------- /CMakeLists.txt: -------------------------------------------------------------------------------- 1 | cmake_minimum_required(VERSION 3.5) 2 | 3 | # Set extension name here 4 | set(TARGET_NAME airport) 5 | 6 | set(EXTENSION_NAME ${TARGET_NAME}_extension) 7 | set(LOADABLE_EXTENSION_NAME ${TARGET_NAME}_loadable_extension) 8 | 9 | find_package(Arrow REQUIRED) 10 | find_package(ArrowFlight REQUIRED) 11 | find_package(CURL REQUIRED) 12 | find_package(msgpack-cxx CONFIG REQUIRED) 13 | 14 | project(${TARGET_NAME}) 15 | include_directories(src/include) 16 | 17 | set(EXTENSION_SOURCES 18 | src/airport_flight_exception.cpp 19 | src/airport_extension.cpp 20 | src/airport_flight_stream.cpp 21 | src/airport_json_common.cpp 22 | src/airport_json_serializer.cpp 23 | src/airport_list_flights.cpp 24 | src/airport_secrets.cpp 25 | src/airport_take_flight.cpp 26 | src/storage/airport_catalog_api.cpp 27 | src/storage/airport_catalog_set.cpp 28 | src/storage/airport_catalog.cpp 29 | src/storage/airport_clear_cache.cpp 30 | src/storage/airport_schema_entry.cpp 31 | src/storage/airport_schema_set.cpp 32 | src/storage/airport_table_entry.cpp 33 | src/storage/airport_table_set.cpp 34 | src/storage/airport_transaction_manager.cpp 35 | src/storage/airport_transaction.cpp 36 | src/storage/airport_curl_pool.cpp 37 | src/storage/airport_exchange.cpp 38 | src/storage/airport_delete.cpp 39 | src/storage/airport_insert.cpp 40 | src/storage/airport_update.cpp 41 | src/storage/airport_delete_parameterized.cpp 42 | src/storage/airport_update_parameterized.cpp 43 | src/airport_request_headers.cpp 44 | src/airport_optimizer.cpp 45 | src/airport_constraints.cpp 46 | src/airport_scalar_function.cpp 47 | src/airport_flight_statistics.cpp 48 | src/airport_schema_utils.cpp 49 | src/airport_action.cpp 50 | src/storage/airport_scalar_function_set.cpp 51 | src/storage/airport_table_function_set.cpp 52 | ) 53 | 54 | 55 | build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES}) 56 | 57 | #target_link_libraries(${EXTENSION_NAME} OpenSSL::SSL OpenSSL::Crypto) 58 | #target_link_libraries(${LOADABLE_EXTENSION_NAME} OpenSSL::SSL OpenSSL::Crypto) 59 | 60 | target_link_libraries(${EXTENSION_NAME} 61 | Arrow::arrow_static 62 | ArrowFlight::arrow_flight_static 63 | gRPC::grpc++ 64 | CURL::libcurl 65 | OpenSSL::SSL 66 | OpenSSL::Crypto 67 | msgpack-cxx 68 | ) 69 | 70 | #target_link_libraries(${LOADABLE_TARGET_NAME} PRIVATE Arrow::arrow_shared ArrowFlight::arrow_flight_shared CURL::libcurl OpenSSL::SSL OpenSSL::Crypto) 71 | 72 | build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES}) 73 | 74 | target_link_libraries(${LOADABLE_EXTENSION_NAME} 75 | Arrow::arrow_static 76 | ArrowFlight::arrow_flight_static 77 | gRPC::grpc++ 78 | CURL::libcurl 79 | OpenSSL::SSL 80 | OpenSSL::Crypto 81 | msgpack-cxx 82 | ) 83 | 84 | 85 | install( 86 | TARGETS ${EXTENSION_NAME} 87 | EXPORT "${DUCKDB_EXPORT_SET}" 88 | LIBRARY DESTINATION "${INSTALL_LIB_DIR}" 89 | ARCHIVE DESTINATION "${INSTALL_LIB_DIR}") 90 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | Copyright 2024-2025 Rusty Conover - https://rusty.today 2 | 3 | Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: 4 | 5 | The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. 6 | 7 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 8 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | PROJ_DIR := $(dir $(abspath $(lastword $(MAKEFILE_LIST)))) 2 | 3 | # Configuration of extension 4 | EXT_NAME=airport 5 | EXT_CONFIG=${PROJ_DIR}extension_config.cmake 6 | 7 | # Include the Makefile from extension-ci-tools 8 | include extension-ci-tools/makefiles/duckdb_extension.Makefile -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- 1 | # Airport Extension for DuckDB 2 | 3 | The **Airport** extension brings [Arrow Flight](https://arrow.apache.org/docs/format/Flight.html) support to [DuckDB](https://duckdb.org), enabling DuckDB to query, modify, and store data via Arrow Flight servers. A DuckDB extension is a plugin that expands DuckDB's core functionality by adding new capabilities. 4 | 5 | To understand the rationale behind the development of this extension, check out the [motivation for creating the extension](https://airport.query.farm/motivation.html). 6 | 7 | # Documentation 8 | 9 | Visit the [documentation for this extension](https://airport.query.farm). 10 | 11 | # Building the extension 12 | 13 | ```sh 14 | # Clone this repo with submodules. 15 | # duckdb and extension-ci-tools are submodules. 16 | git clone --recursive git@github.com:Query-farm/duckdb-airport-extension 17 | 18 | # Clone the vcpkg repo 19 | git clone https://github.com/Microsoft/vcpkg.git 20 | 21 | # Bootstrap vcpkg 22 | ./vcpkg/bootstrap-vcpkg.sh 23 | export VCPKG_TOOLCHAIN_PATH=`pwd`/vcpkg/scripts/buildsystems/vcpkg.cmake 24 | 25 | # Build the extension 26 | make 27 | 28 | # If you have ninja installed, you can use it to speed up the build 29 | # GEN=ninja make 30 | ``` 31 | 32 | The main binaries that will be built are: 33 | ```sh 34 | ./build/release/duckdb 35 | ./build/release/test/unittest 36 | ./build/release/extension/airport/airport.duckdb_extension 37 | ``` 38 | 39 | - `duckdb` is the binary for the duckdb shell with the extension code automatically loaded. 40 | - `unittest` is the test runner of duckdb. Again, the extension is already linked into the binary. 41 | - `airport.duckdb_extension` is the loadable binary as it would be distributed. 42 | 43 | ## Running the extension 44 | 45 | To run the extension code, simply start the shell with `./build/release/duckdb`. This duckdb shell will have the extension pre-loaded. 46 | 47 | Now we can use the features from the extension directly in DuckDB. 48 | 49 | ## Running the tests 50 | Different tests can be created for DuckDB extensions. The primary way of testing DuckDB extensions should be the SQL tests in `./test/sql`. These SQL tests can be run using: 51 | 52 | ```sh 53 | make test 54 | ``` 55 | 56 | -------------------------------------------------------------------------------- /docs/duckdb-airport-1.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Query-farm/airport/5d6247c519d2a19894e7479d385b8d3557724c08/docs/duckdb-airport-1.jpg -------------------------------------------------------------------------------- /extension_config.cmake: -------------------------------------------------------------------------------- 1 | # This file is included by DuckDB's build system. It specifies which extension to load 2 | 3 | # Extension from this repo 4 | duckdb_extension_load(airport 5 | SOURCE_DIR ${CMAKE_CURRENT_LIST_DIR} 6 | LOAD_TESTS 7 | ) 8 | 9 | # Any extra extensions that should be built 10 | # e.g.: duckdb_extension_load(json) -------------------------------------------------------------------------------- /scripts/bootstrap-template.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/python3 2 | 3 | import sys, os, shutil, re 4 | from pathlib import Path 5 | 6 | shutil.copyfile(f'docs/NEXT_README.md', f'README.md') 7 | os.remove(f'docs/NEXT_README.md') 8 | os.remove(f'docs/README.md') 9 | 10 | if (len(sys.argv) != 2): 11 | raise Exception('usage: python3 bootstrap-template.py ') 12 | 13 | name_extension = sys.argv[1] 14 | 15 | def is_snake_case(s): 16 | # Define the regex pattern for snake case with numbers 17 | pattern = r'^[a-z0-9]+(_[a-z0-9]+)*$' 18 | 19 | # Use re.match to check if the string matches the pattern 20 | if re.match(pattern, s): 21 | return True 22 | else: 23 | return False 24 | 25 | if name_extension[0].isdigit(): 26 | raise Exception('Please dont start your extension name with a number.') 27 | 28 | if not is_snake_case(name_extension): 29 | raise Exception('Please enter the name of your extension in valid snake_case containing only lower case letters and numbers') 30 | 31 | def to_camel_case(snake_str): 32 | return "".join(x.capitalize() for x in snake_str.lower().split("_")) 33 | 34 | def replace(file_name, to_find, to_replace): 35 | with open(file_name, 'r', encoding="utf8") as file : 36 | filedata = file.read() 37 | filedata = filedata.replace(to_find, to_replace) 38 | with open(file_name, 'w', encoding="utf8") as file: 39 | file.write(filedata) 40 | 41 | files_to_search = [] 42 | files_to_search.extend(Path('./.github').rglob('./**/*.yml')) 43 | files_to_search.extend(Path('./test').rglob('./**/*.test')) 44 | files_to_search.extend(Path('./src').rglob('./**/*.hpp')) 45 | files_to_search.extend(Path('./src').rglob('./**/*.cpp')) 46 | files_to_search.extend(Path('./src').rglob('./**/*.txt')) 47 | files_to_search.extend(Path('./src').rglob('./*.md')) 48 | 49 | def replace_everywhere(to_find, to_replace): 50 | for path in files_to_search: 51 | replace(path, to_find, to_replace) 52 | replace(path, to_find.capitalize(), to_camel_case(to_replace)) 53 | replace(path, to_find.upper(), to_replace.upper()) 54 | 55 | replace("./CMakeLists.txt", to_find, to_replace) 56 | replace("./Makefile", to_find, to_replace) 57 | replace("./Makefile", to_find.capitalize(), to_camel_case(to_replace)) 58 | replace("./Makefile", to_find.upper(), to_replace.upper()) 59 | replace("./README.md", to_find, to_replace) 60 | replace("./extension_config.cmake", to_find, to_replace) 61 | 62 | replace_everywhere("airport", name_extension) 63 | replace_everywhere("airport", name_extension.capitalize()) 64 | replace_everywhere("", name_extension) 65 | 66 | string_to_replace = name_extension 67 | string_to_find = "airport" 68 | 69 | # rename files 70 | os.rename(f'test/sql/{string_to_find}.test', f'test/sql/{string_to_replace}.test') 71 | os.rename(f'src/{string_to_find}_extension.cpp', f'src/{string_to_replace}_extension.cpp') 72 | os.rename(f'src/include/{string_to_find}_extension.hpp', f'src/include/{string_to_replace}_extension.hpp') 73 | 74 | # remove template-specific files 75 | os.remove('.github/workflows/ExtensionTemplate.yml') 76 | 77 | # finally, remove this bootstrap file 78 | os.remove(__file__) -------------------------------------------------------------------------------- /scripts/extension-upload.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # Extension upload script 4 | 5 | # Usage: ./extension-upload.sh 6 | # : Name of the extension 7 | # : Version (commit / version tag) of the extension 8 | # : Version (commit / version tag) of DuckDB 9 | # : Architecture target of the extension binary 10 | # : S3 bucket to upload to 11 | # : Set this as the latest version ("true" / "false", default: "false") 12 | # : Set this as a versioned version that will prevent its deletion 13 | 14 | set -e 15 | 16 | if [[ $4 == wasm* ]]; then 17 | ext="/tmp/extension/$1.duckdb_extension.wasm" 18 | else 19 | ext="/tmp/extension/$1.duckdb_extension" 20 | fi 21 | 22 | echo $ext 23 | 24 | script_dir="$(dirname "$(readlink -f "$0")")" 25 | 26 | # calculate SHA256 hash of extension binary 27 | cat $ext > $ext.append 28 | 29 | if [[ $4 == wasm* ]]; then 30 | # 0 for custom section 31 | # 113 in hex = 275 in decimal, total lenght of what follows (1 + 16 + 2 + 256) 32 | # [1(continuation) + 0010011(payload) = \x93, 0(continuation) + 10(payload) = \x02] 33 | echo -n -e '\x00' >> $ext.append 34 | echo -n -e '\x93\x02' >> $ext.append 35 | # 10 in hex = 16 in decimal, lenght of name, 1 byte 36 | echo -n -e '\x10' >> $ext.append 37 | echo -n -e 'duckdb_signature' >> $ext.append 38 | # the name of the WebAssembly custom section, 16 bytes 39 | # 100 in hex, 256 in decimal 40 | # [1(continuation) + 0000000(payload) = ff, 0(continuation) + 10(payload)], 41 | # for a grand total of 2 bytes 42 | echo -n -e '\x80\x02' >> $ext.append 43 | fi 44 | 45 | # (Optionally) Sign binary 46 | if [ "$DUCKDB_EXTENSION_SIGNING_PK" != "" ]; then 47 | echo "$DUCKDB_EXTENSION_SIGNING_PK" > private.pem 48 | $script_dir/../duckdb/scripts/compute-extension-hash.sh $ext.append > $ext.hash 49 | openssl pkeyutl -sign -in $ext.hash -inkey private.pem -pkeyopt digest:sha256 -out $ext.sign 50 | rm -f private.pem 51 | fi 52 | 53 | # Signature is always there, potentially defaulting to 256 zeros 54 | truncate -s 256 $ext.sign 55 | 56 | # append signature to extension binary 57 | cat $ext.sign >> $ext.append 58 | 59 | # compress extension binary 60 | if [[ $4 == wasm_* ]]; then 61 | brotli < $ext.append > "$ext.compressed" 62 | else 63 | gzip < $ext.append > "$ext.compressed" 64 | fi 65 | 66 | set -e 67 | 68 | # Abort if AWS key is not set 69 | if [ -z "$AWS_ACCESS_KEY_ID" ]; then 70 | echo "No AWS key found, skipping.." 71 | exit 0 72 | fi 73 | 74 | # upload versioned version 75 | if [[ $7 = 'true' ]]; then 76 | if [[ $4 == wasm* ]]; then 77 | aws s3 cp $ext.compressed s3://$5/$1/$2/$3/$4/$1.duckdb_extension.wasm --acl public-read --content-encoding br --content-type="application/wasm" 78 | else 79 | aws s3 cp $ext.compressed s3://$5/$1/$2/$3/$4/$1.duckdb_extension.gz --acl public-read 80 | fi 81 | fi 82 | 83 | # upload to latest version 84 | if [[ $6 = 'true' ]]; then 85 | if [[ $4 == wasm* ]]; then 86 | aws s3 cp $ext.compressed s3://$5/$3/$4/$1.duckdb_extension.wasm --acl public-read --content-encoding br --content-type="application/wasm" 87 | else 88 | aws s3 cp $ext.compressed s3://$5/$3/$4/$1.duckdb_extension.gz --acl public-read 89 | fi 90 | fi 91 | -------------------------------------------------------------------------------- /src/airport_action.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "duckdb.hpp" 3 | #include "duckdb/common/exception.hpp" 4 | 5 | // Arrow includes. 6 | #include 7 | #include "duckdb/main/extension_util.hpp" 8 | 9 | #include "duckdb/main/secret/secret_manager.hpp" 10 | 11 | #include "airport_json_common.hpp" 12 | #include "airport_json_serializer.hpp" 13 | #include "airport_macros.hpp" 14 | #include "airport_secrets.hpp" 15 | #include "airport_request_headers.hpp" 16 | #include "storage/airport_catalog.hpp" 17 | 18 | namespace flight = arrow::flight; 19 | 20 | namespace duckdb 21 | { 22 | 23 | struct ActionBindData : public TableFunctionData 24 | { 25 | // This is is the location of the server 26 | std::string server_location; 27 | 28 | // This is the auth token. 29 | std::string auth_token; 30 | 31 | std::string action_name; 32 | 33 | // The parameters that will be passed to the action. 34 | std::optional parameter; 35 | 36 | std::unordered_map> user_supplied_headers; 37 | 38 | explicit ActionBindData(std::string server_location, 39 | std::string auth_token, 40 | std::string action_name, 41 | std::optional parameter, 42 | std::unordered_map> user_supplied_headers) 43 | : server_location(std::move(server_location)), 44 | auth_token(std::move(auth_token)), 45 | action_name(std::move(action_name)), 46 | parameter(std::move(parameter)), 47 | user_supplied_headers(std::move(user_supplied_headers)) 48 | { 49 | } 50 | }; 51 | 52 | struct ActionGlobalState : public GlobalTableFunctionState 53 | { 54 | public: 55 | const std::shared_ptr flight_client_; 56 | 57 | // If the action returns a lot of items, we need to keep the result stream here. 58 | std::unique_ptr result_stream; 59 | 60 | explicit ActionGlobalState(std::shared_ptr flight_client) : flight_client_(flight_client) 61 | { 62 | } 63 | 64 | idx_t MaxThreads() const override 65 | { 66 | return 1; 67 | } 68 | 69 | static unique_ptr Init(ClientContext &context, TableFunctionInitInput &input) 70 | { 71 | const auto &bind_data = input.bind_data->Cast(); 72 | 73 | auto flight_client = AirportAPI::FlightClientForLocation(bind_data.server_location); 74 | 75 | return make_uniq(flight_client); 76 | } 77 | }; 78 | 79 | static unique_ptr do_action_bind( 80 | ClientContext &context, 81 | TableFunctionBindInput &input, 82 | vector &return_types, vector &names) 83 | { 84 | if (input.inputs.size() < 2) 85 | { 86 | throw BinderException("airport_action requires at least 2 arguments"); 87 | } 88 | 89 | auto server_location = input.inputs[0].ToString(); 90 | auto action_name = input.inputs[1].ToString(); 91 | 92 | std::optional parameter = std::nullopt; 93 | if (input.inputs.size() > 2) 94 | { 95 | parameter = input.inputs[2].ToString(); 96 | } 97 | 98 | string auth_token = ""; 99 | string secret_name = ""; 100 | std::unordered_map> user_supplied_headers; 101 | 102 | for (auto &kv : input.named_parameters) 103 | { 104 | auto loption = StringUtil::Lower(kv.first); 105 | if (loption == "auth_token") 106 | { 107 | auth_token = StringValue::Get(kv.second); 108 | } 109 | else if (loption == "secret") 110 | { 111 | secret_name = StringValue::Get(kv.second); 112 | } 113 | else if (loption == "headers") 114 | { 115 | // Now we need to parse out the map contents. 116 | auto &children = duckdb::MapValue::GetChildren(kv.second); 117 | 118 | for (auto &value_pair : children) 119 | { 120 | auto &child_struct = duckdb::StructValue::GetChildren(value_pair); 121 | auto key = StringValue::Get(child_struct[0]); 122 | auto value = StringValue::Get(child_struct[1]); 123 | 124 | user_supplied_headers[key].push_back(value); 125 | } 126 | } 127 | } 128 | 129 | auth_token = AirportAuthTokenForLocation(context, server_location, secret_name, auth_token); 130 | 131 | auto ret = make_uniq(server_location, 132 | auth_token, 133 | action_name, 134 | parameter, 135 | user_supplied_headers); 136 | 137 | return_types.emplace_back(LogicalType(LogicalTypeId::BLOB)); 138 | names.emplace_back("result"); 139 | 140 | return ret; 141 | } 142 | 143 | static void do_action(ClientContext &context, TableFunctionInput &data, DataChunk &output) 144 | { 145 | auto &bind_data = data.bind_data->Cast(); 146 | auto &global_state = data.global_state->Cast(); 147 | 148 | auto &server_location = bind_data.server_location; 149 | 150 | if (global_state.result_stream == nullptr) 151 | { 152 | // Now send a list flights request. 153 | arrow::flight::FlightCallOptions call_options; 154 | airport_add_standard_headers(call_options, bind_data.server_location); 155 | 156 | // FIXME: this will fail with large filter sizes, so its best not to pass it here. 157 | call_options.headers.emplace_back("airport-action-name", bind_data.action_name); 158 | airport_add_authorization_header(call_options, bind_data.auth_token); 159 | // printf("Calling with filters: %s\n", bind_data.json_filters.c_str()); 160 | 161 | arrow::flight::Action action{ 162 | bind_data.action_name, 163 | bind_data.parameter.has_value() ? std::make_shared( 164 | reinterpret_cast(bind_data.parameter.value().data()), 165 | bind_data.parameter.value().size()) 166 | : std::make_shared(nullptr, 0)}; 167 | 168 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(global_state.result_stream, 169 | global_state.flight_client_->DoAction(call_options, action), 170 | server_location, 171 | "airport_action"); 172 | } 173 | 174 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto action_result, global_state.result_stream->Next(), server_location, "airport_action next item"); 175 | 176 | if (action_result == nullptr) 177 | { 178 | // There are no results on the stream. 179 | AIRPORT_ARROW_ASSERT_OK_LOCATION(global_state.result_stream->Drain(), server_location, "airport_action drain"); 180 | output.SetCardinality(0); 181 | return; 182 | } 183 | 184 | FlatVector::GetData(output.data[0])[0] = StringVector::AddStringOrBlob(output.data[0], 185 | action_result->body->ToString()); 186 | 187 | output.SetCardinality(1); 188 | } 189 | 190 | void AirportAddActionFlightFunction(DatabaseInstance &instance) 191 | { 192 | auto do_action_functions = TableFunctionSet("airport_action"); 193 | 194 | auto with_parameter = TableFunction( 195 | "airport_action", 196 | {LogicalType::VARCHAR, LogicalType::VARCHAR, LogicalType::VARCHAR}, 197 | do_action, 198 | do_action_bind, 199 | ActionGlobalState::Init); 200 | 201 | with_parameter.named_parameters["auth_token"] = LogicalType::VARCHAR; 202 | with_parameter.named_parameters["secret"] = LogicalType::VARCHAR; 203 | with_parameter.named_parameters["headers"] = LogicalType::MAP(LogicalType::VARCHAR, LogicalType::VARCHAR); 204 | 205 | do_action_functions.AddFunction(with_parameter); 206 | 207 | auto without_parameter = TableFunction( 208 | "airport_action", 209 | {LogicalType::VARCHAR, LogicalType::VARCHAR}, 210 | do_action, 211 | do_action_bind, 212 | ActionGlobalState::Init); 213 | 214 | without_parameter.named_parameters["auth_token"] = LogicalType::VARCHAR; 215 | without_parameter.named_parameters["secret"] = LogicalType::VARCHAR; 216 | without_parameter.named_parameters["headers"] = LogicalType::MAP(LogicalType::VARCHAR, LogicalType::VARCHAR); 217 | 218 | do_action_functions.AddFunction(without_parameter); 219 | 220 | ExtensionUtil::RegisterFunction(instance, do_action_functions); 221 | } 222 | 223 | } 224 | -------------------------------------------------------------------------------- /src/airport_constraints.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb/common/arrow/schema_metadata.hpp" 2 | #include "duckdb/planner/expression/bound_reference_expression.hpp" 3 | #include "duckdb/planner/constraints/bound_not_null_constraint.hpp" 4 | #include "duckdb/planner/constraints/bound_check_constraint.hpp" 5 | #include "duckdb/planner/constraints/bound_unique_constraint.hpp" 6 | #include "duckdb/parser/constraints/not_null_constraint.hpp" 7 | #include "duckdb/parser/constraints/check_constraint.hpp" 8 | 9 | #include "duckdb/storage/table/append_state.hpp" 10 | 11 | #include "airport_extension.hpp" 12 | #include "storage/airport_insert.hpp" 13 | #include "storage/airport_catalog.hpp" 14 | #include "storage/airport_transaction.hpp" 15 | #include "storage/airport_table_entry.hpp" 16 | 17 | namespace duckdb 18 | { 19 | // Need to Adapt this for Airport. 20 | 21 | static void VerifyNotNullConstraint(TableCatalogEntry &table, Vector &vector, idx_t count, const string &col_name) 22 | { 23 | if (!VectorOperations::HasNull(vector, count)) 24 | { 25 | return; 26 | } 27 | 28 | throw ConstraintException("NOT NULL constraint failed: %s.%s", table.name, col_name); 29 | } 30 | 31 | // To avoid throwing an error at SELECT, instead this moves the error detection to INSERT 32 | // static void VerifyGeneratedExpressionSuccess(ClientContext &context, TableCatalogEntry &table, DataChunk &chunk, 33 | // Expression &expr, column_t index) 34 | // { 35 | // auto &col = table.GetColumn(LogicalIndex(index)); 36 | // D_ASSERT(col.Generated()); 37 | // ExpressionExecutor executor(context, expr); 38 | // Vector result(col.Type()); 39 | // try 40 | // { 41 | // executor.ExecuteExpression(chunk, result); 42 | // } 43 | // catch (InternalException &ex) 44 | // { 45 | // throw; 46 | // } 47 | // catch (std::exception &ex) 48 | // { 49 | // ErrorData error(ex); 50 | // throw ConstraintException("Incorrect value for generated column \"%s %s AS (%s)\" : %s", col.Name(), 51 | // col.Type().ToString(), col.GeneratedExpression().ToString(), error.RawMessage()); 52 | // } 53 | // } 54 | 55 | static void VerifyCheckConstraint(ClientContext &context, TableCatalogEntry &table, Expression &expr, 56 | DataChunk &chunk, CheckConstraint &check) 57 | { 58 | ExpressionExecutor executor(context, expr); 59 | Vector result(LogicalType::INTEGER); 60 | try 61 | { 62 | executor.ExecuteExpression(chunk, result); 63 | } 64 | catch (std::exception &ex) 65 | { 66 | ErrorData error(ex); 67 | throw ConstraintException("CHECK constraint failed on table %s with expression %s (Error: %s)", table.name, check.ToString(), error.RawMessage()); 68 | } 69 | catch (...) 70 | { 71 | // LCOV_EXCL_START 72 | throw ConstraintException("CHECK constraint failed on table %s with expression %s (Unknown Error)", table.name, check.ToString()); 73 | } // LCOV_EXCL_STOP 74 | UnifiedVectorFormat vdata; 75 | result.ToUnifiedFormat(chunk.size(), vdata); 76 | 77 | auto dataptr = UnifiedVectorFormat::GetData(vdata); 78 | for (idx_t i = 0; i < chunk.size(); i++) 79 | { 80 | auto idx = vdata.sel->get_index(i); 81 | if (vdata.validity.RowIsValid(idx) && dataptr[idx] == 0) 82 | { 83 | throw ConstraintException("CHECK constraint failed on table %s with expression: %s", table.name, check.ToString()); 84 | } 85 | } 86 | } 87 | 88 | std::optional findStringIndex(const std::vector &vec, const std::string &target) 89 | { 90 | auto it = std::find(vec.begin(), vec.end(), target); 91 | if (it != vec.end()) 92 | { 93 | return std::distance(vec.begin(), it); // Return the index 94 | } 95 | return std::nullopt; // Return empty if not found 96 | } 97 | 98 | void AirportVerifyAppendConstraints(ConstraintState &state, ClientContext &context, DataChunk &chunk, 99 | optional_ptr conflict_manager, 100 | vector column_names) 101 | { 102 | auto &table = state.table; 103 | 104 | // Airport currently doesn't have generated columns 105 | 106 | // if (table.HasGeneratedColumns()) { 107 | // // Verify that the generated columns expression work with the inserted values 108 | // auto binder = Binder::CreateBinder(context); 109 | // physical_index_set_t bound_columns; 110 | // CheckBinder generated_check_binder(*binder, context, table.name, table.GetColumns(), bound_columns); 111 | // for (auto &col : table.GetColumns().Logical()) { 112 | // if (!col.Generated()) { 113 | // continue; 114 | // } 115 | // D_ASSERT(col.Type().id() != LogicalTypeId::ANY); 116 | // generated_check_binder.target_type = col.Type(); 117 | // auto to_be_bound_expression = col.GeneratedExpression().Copy(); 118 | // auto bound_expression = generated_check_binder.Bind(to_be_bound_expression); 119 | // VerifyGeneratedExpressionSuccess(context, table, chunk, *bound_expression, col.Oid()); 120 | // } 121 | // } 122 | 123 | // if (HasUniqueIndexes(info->indexes)) { 124 | // VerifyUniqueIndexes(info->indexes, context, chunk, conflict_manager); 125 | // } 126 | 127 | auto &constraints = table.GetConstraints(); 128 | for (idx_t i = 0; i < state.bound_constraints.size(); i++) 129 | { 130 | auto &base_constraint = constraints[i]; 131 | switch (base_constraint->type) 132 | { 133 | case ConstraintType::NOT_NULL: 134 | { 135 | // auto &bound_not_null = constraint->Cast(); 136 | auto ¬_null = base_constraint->Cast(); 137 | auto &col = table.GetColumns().GetColumn(LogicalIndex(not_null.index)); 138 | 139 | // Since the column names that are send can be in a different order 140 | // look them up by name. 141 | auto col_index = findStringIndex(column_names, col.Name()); 142 | if (col_index.has_value()) 143 | { 144 | VerifyNotNullConstraint(table, chunk.data[col_index.value()], chunk.size(), col.Name()); 145 | } 146 | break; 147 | } 148 | case ConstraintType::CHECK: 149 | { 150 | auto &constraint = state.bound_constraints[i]; 151 | 152 | auto &check = constraint->Cast(); 153 | auto &unbound_constraint = base_constraint->Cast(); 154 | VerifyCheckConstraint(context, table, *check.expression, chunk, unbound_constraint); 155 | break; 156 | } 157 | case ConstraintType::UNIQUE: 158 | { 159 | // These were handled earlier on 160 | break; 161 | } 162 | // case ConstraintType::FOREIGN_KEY: { 163 | // auto &bfk = constraint->Cast(); 164 | // if (bfk.info.type == ForeignKeyType::FK_TYPE_FOREIGN_KEY_TABLE || 165 | // bfk.info.type == ForeignKeyType::FK_TYPE_SELF_REFERENCE_TABLE) { 166 | // VerifyAppendForeignKeyConstraint(bfk, context, chunk); 167 | // } 168 | // break; 169 | //} 170 | default: 171 | throw NotImplementedException("Constraint type not implemented!"); 172 | } 173 | } 174 | } 175 | 176 | } -------------------------------------------------------------------------------- /src/airport_extension.cpp: -------------------------------------------------------------------------------- 1 | #define DUCKDB_EXTENSION_MAIN 2 | 3 | #include "airport_extension.hpp" 4 | #include "duckdb.hpp" 5 | #include "duckdb/main/extension_util.hpp" 6 | 7 | #include "duckdb/main/secret/secret_manager.hpp" 8 | #include "duckdb/parser/parsed_data/attach_info.hpp" 9 | #include "duckdb/storage/storage_extension.hpp" 10 | #include "duckdb/parser/parsed_data/create_macro_info.hpp" 11 | #include "duckdb/function/table_macro_function.hpp" 12 | #include "duckdb/catalog/default/default_functions.hpp" 13 | #include "storage/airport_catalog.hpp" 14 | #include "storage/airport_transaction_manager.hpp" 15 | #include "airport_secrets.hpp" 16 | #include "airport_optimizer.hpp" 17 | #include "airport_scalar_function.hpp" 18 | #include 19 | 20 | namespace duckdb 21 | { 22 | 23 | static unique_ptr CreateAirportSecretFunction(ClientContext &, CreateSecretInput &input) 24 | { 25 | // apply any overridden settings 26 | vector prefix_paths; 27 | 28 | auto scope = input.scope; 29 | if (scope.empty()) 30 | { 31 | throw InternalException("No scope set Airport create secret (should start with grpc://): '%s'", input.type); 32 | } 33 | 34 | auto result = make_uniq(scope, "airport", "config", input.name); 35 | for (const auto &named_param : input.options) 36 | { 37 | auto lower_name = StringUtil::Lower(named_param.first); 38 | 39 | if (lower_name == "auth_token") 40 | { 41 | result->secret_map["auth_token"] = named_param.second.ToString(); 42 | } 43 | else 44 | { 45 | throw InternalException("Unknown named parameter passed to CreateAirportSecretFunction: " + lower_name); 46 | } 47 | } 48 | 49 | //! Set redact keys 50 | result->redact_keys = {"auth_token"}; 51 | 52 | return result; 53 | } 54 | 55 | static void AirportSetSecretParameters(CreateSecretFunction &function) 56 | { 57 | function.named_parameters["auth_token"] = LogicalType::VARCHAR; 58 | } 59 | 60 | static unique_ptr AirportCatalogAttach(StorageExtensionInfo *storage_info, ClientContext &context, 61 | AttachedDatabase &db, const string &name, AttachInfo &info, 62 | AccessMode access_mode) 63 | { 64 | string secret_name; 65 | string auth_token; 66 | string location; 67 | 68 | // check if we have a secret provided 69 | for (auto &entry : info.options) 70 | { 71 | auto lower_name = StringUtil::Lower(entry.first); 72 | if (lower_name == "type") 73 | { 74 | continue; 75 | } 76 | else if (lower_name == "secret") 77 | { 78 | secret_name = entry.second.ToString(); 79 | } 80 | else if (lower_name == "auth_token") 81 | { 82 | auth_token = entry.second.ToString(); 83 | } 84 | else if (lower_name == "location") 85 | { 86 | location = entry.second.ToString(); 87 | } 88 | else 89 | { 90 | throw BinderException("Unrecognized option for Airport ATTACH: %s", entry.first); 91 | } 92 | } 93 | 94 | auth_token = AirportAuthTokenForLocation(context, location, secret_name, auth_token); 95 | 96 | if (location.empty()) 97 | { 98 | throw BinderException("No location provided for Airport ATTACH."); 99 | } 100 | 101 | return make_uniq(db, info.path, access_mode, AirportAttachParameters(location, auth_token, secret_name, "")); 102 | } 103 | 104 | static unique_ptr CreateTransactionManager(StorageExtensionInfo *storage_info, AttachedDatabase &db, 105 | Catalog &catalog) 106 | { 107 | auto &airport_catalog = catalog.Cast(); 108 | return make_uniq(db, airport_catalog); 109 | } 110 | 111 | class AirportCatalogStorageExtension : public StorageExtension 112 | { 113 | public: 114 | AirportCatalogStorageExtension() 115 | { 116 | attach = AirportCatalogAttach; 117 | create_transaction_manager = CreateTransactionManager; 118 | } 119 | }; 120 | 121 | inline void get_user_agent(DataChunk &args, ExpressionState &state, Vector &result) 122 | { 123 | D_ASSERT(args.ColumnCount() == 0); 124 | Value val(airport_user_agent()); 125 | result.Reference(val); 126 | } 127 | 128 | void AirportAddUserAgentFunction(DatabaseInstance &instance) 129 | { 130 | ExtensionUtil::RegisterFunction( 131 | instance, 132 | ScalarFunction( 133 | "airport_user_agent", 134 | {}, 135 | LogicalType::VARCHAR, 136 | get_user_agent)); 137 | } 138 | 139 | static void RegisterTableMacro(DatabaseInstance &db, const string &name, const string &query, 140 | const vector ¶ms, const child_list_t &named_params) 141 | { 142 | Parser parser; 143 | parser.ParseQuery(query); 144 | const auto &stmt = parser.statements.back(); 145 | auto &node = stmt->Cast().node; 146 | 147 | auto func = make_uniq(std::move(node)); 148 | for (auto ¶m : params) 149 | { 150 | func->parameters.push_back(make_uniq(param)); 151 | } 152 | 153 | for (auto ¶m : named_params) 154 | { 155 | func->default_parameters[param.first] = make_uniq(param.second); 156 | } 157 | 158 | CreateMacroInfo info(CatalogType::TABLE_MACRO_ENTRY); 159 | info.schema = DEFAULT_SCHEMA; 160 | info.name = name; 161 | info.temporary = true; 162 | info.internal = true; 163 | info.macros.push_back(std::move(func)); 164 | 165 | ExtensionUtil::RegisterFunction(db, info); 166 | } 167 | 168 | static void AirportAddListDatabasesMacro(DatabaseInstance &instance) 169 | { 170 | child_list_t named_params = { 171 | // {"auth_token", Value()}, 172 | // {"secret", Value()}, 173 | // {"headers", Value()}, 174 | }; 175 | 176 | RegisterTableMacro( 177 | instance, 178 | "airport_databases", 179 | "select * from airport_take_flight(server_location, ['__databases'])", 180 | // "select * from airport_take_flight(server_location, ['__databases'], auth_token=auth_token, secret=secret, headers=headers)", 181 | {"server_location"}, 182 | named_params); 183 | } 184 | 185 | static void LoadInternal(DatabaseInstance &instance) 186 | { 187 | curl_global_init(CURL_GLOBAL_DEFAULT); 188 | 189 | AirportAddListFlightsFunction(instance); 190 | AirportAddTakeFlightFunction(instance); 191 | AirportAddUserAgentFunction(instance); 192 | AirportAddActionFlightFunction(instance); 193 | 194 | // So to create a new macro for airport_list_databases 195 | // that calls airport_take_flight with a fixed flight descriptor 196 | // of PATH /__databases 197 | 198 | AirportAddListDatabasesMacro(instance); 199 | 200 | SecretType secret_type; 201 | secret_type.name = "airport"; 202 | secret_type.deserializer = KeyValueSecret::Deserialize; 203 | secret_type.default_provider = "config"; 204 | 205 | ExtensionUtil::RegisterSecretType(instance, secret_type); 206 | 207 | CreateSecretFunction airport_secret_function = {"airport", "config", CreateAirportSecretFunction}; 208 | AirportSetSecretParameters(airport_secret_function); 209 | ExtensionUtil::RegisterFunction(instance, airport_secret_function); 210 | 211 | auto &config = DBConfig::GetConfig(instance); 212 | config.storage_extensions["airport"] = make_uniq(); 213 | 214 | OptimizerExtension airport_optimizer; 215 | airport_optimizer.optimize_function = AirportOptimizer::Optimize; 216 | config.optimizer_extensions.push_back(std::move(airport_optimizer)); 217 | } 218 | 219 | void AirportExtension::Load(DuckDB &db) 220 | { 221 | LoadInternal(*db.instance); 222 | } 223 | std::string AirportExtension::Name() 224 | { 225 | return "airport"; 226 | } 227 | 228 | std::string AirportExtension::Version() const 229 | { 230 | return "user-agent=" + airport_user_agent() + ",client=2025050401"; 231 | } 232 | 233 | } // namespace duckdb 234 | 235 | extern "C" 236 | { 237 | DUCKDB_EXTENSION_API void airport_init(duckdb::DatabaseInstance &db) 238 | { 239 | duckdb::DuckDB db_wrapper(db); 240 | db_wrapper.LoadExtension(); 241 | } 242 | 243 | DUCKDB_EXTENSION_API const char *airport_version() 244 | { 245 | return duckdb::DuckDB::LibraryVersion(); 246 | } 247 | } 248 | 249 | #ifndef DUCKDB_EXTENSION_MAIN 250 | #error DUCKDB_EXTENSION_MAIN not defined 251 | #endif 252 | -------------------------------------------------------------------------------- /src/airport_flight_exception.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | #include "airport_flight_exception.hpp" 3 | 4 | #include "airport_extension.hpp" 5 | #include "duckdb.hpp" 6 | 7 | // Arrow includes. 8 | #include 9 | #include 10 | #include 11 | #include "yyjson.hpp" 12 | #include "airport_secrets.hpp" 13 | 14 | using namespace duckdb_yyjson; // NOLINT 15 | 16 | namespace duckdb 17 | { 18 | 19 | static std::string joinWithSlash(const std::vector &vec) 20 | { 21 | if (vec.empty()) 22 | return ""; // Handle empty vector case 23 | 24 | std::ostringstream result; 25 | for (size_t i = 0; i < vec.size(); ++i) 26 | { 27 | result << vec[i]; 28 | if (i != vec.size() - 1) 29 | { 30 | result << '/'; 31 | } 32 | } 33 | return result.str(); 34 | } 35 | 36 | // Function to split string based on a delimiter 37 | static std::vector split(const std::string &input, const std::string &delimiter) 38 | { 39 | std::vector parts; 40 | size_t pos = 0; 41 | size_t start = 0; 42 | 43 | while ((pos = input.find(delimiter, start)) != std::string::npos) 44 | { 45 | parts.push_back(input.substr(start, pos - start)); 46 | start = pos + delimiter.length(); 47 | } 48 | 49 | // Add the last part after the final delimiter 50 | parts.push_back(input.substr(start)); 51 | 52 | return parts; 53 | } 54 | 55 | static void extract_grpc_debug_context(const string &value, std::unordered_map &result) 56 | { 57 | if (value.find("gRPC client debug context: ") != std::string::npos) 58 | { 59 | // Extract the gRPC context from the error message 60 | std::vector parts = split(value, "gRPC client debug context: "); 61 | if (parts.size() > 1) 62 | { 63 | result["grpc_context"] = parts[1]; // Store the gRPC context in the result map 64 | } 65 | } 66 | } 67 | 68 | unordered_map AirportFlightException::extract_extra_info(const string &status, const unordered_map &extra_info) 69 | { 70 | auto arrow_status_string = status; 71 | 72 | // Create an unordered_map to store the key-value pairs 73 | std::unordered_map result = extra_info; 74 | 75 | extract_grpc_debug_context(arrow_status_string, result); 76 | 77 | return result; 78 | } 79 | 80 | unordered_map AirportFlightException::extract_extra_info(const arrow::Status &status, const unordered_map &extra_info) 81 | { 82 | auto arrow_status_string = status.message(); 83 | 84 | // Create an unordered_map to store the key-value pairs 85 | std::unordered_map result = extra_info; 86 | 87 | extract_grpc_debug_context(arrow_status_string, result); 88 | 89 | std::shared_ptr flight_status = flight::FlightStatusDetail::UnwrapStatus(status); 90 | // The extra info is just some bytes, but in my servers its going to be a json document 91 | // that can only have string values 92 | if (flight_status != nullptr) 93 | { 94 | auto extra_info_string = flight_status->extra_info(); 95 | if (!extra_info_string.empty()) 96 | { 97 | // Parse the JSON string 98 | yyjson_doc *doc = yyjson_read(extra_info_string.c_str(), extra_info_string.size(), 0); 99 | if (doc) 100 | { 101 | // Get the root of the JSON document 102 | yyjson_val *root = yyjson_doc_get_root(doc); 103 | 104 | // Ensure the root is an object (dictionary) 105 | if (root && yyjson_is_obj(root)) 106 | { 107 | size_t idx, max; 108 | yyjson_val *key, *value; 109 | yyjson_obj_foreach(root, idx, max, key, value) 110 | { 111 | if (yyjson_is_str(key) && yyjson_is_str(value)) 112 | { 113 | // Add the key-value pair to the unordered_map 114 | result[yyjson_get_str(key)] = yyjson_get_str(value); 115 | } 116 | } 117 | } 118 | } 119 | // Free the JSON document after parsing 120 | yyjson_doc_free(doc); 121 | } 122 | } 123 | return result; 124 | } 125 | 126 | static string extract_grpc_context_prefix(const string &value) 127 | { 128 | if (value.find("gRPC client debug context: ") != std::string::npos) 129 | { 130 | // Extract the gRPC context from the error message 131 | std::vector parts = split(value, "gRPC client debug context: "); 132 | if (parts.size() > 1) 133 | { 134 | return parts[0]; // Return the part before the gRPC context 135 | } 136 | } 137 | return value; // Return the original value if no gRPC context is found 138 | } 139 | 140 | static string build_error_message(const string &location, const string &descriptor, const string &msg, const string &status) 141 | { 142 | // Handle other descriptor types if needed 143 | string descriptor_joiner = "/"; 144 | if (location.back() == '/') 145 | { 146 | descriptor_joiner = ""; 147 | } 148 | 149 | string extra_msg = ""; 150 | if (!msg.empty()) 151 | { 152 | extra_msg = msg + " "; 153 | } 154 | 155 | return "Airport @ " + location + descriptor_joiner + descriptor + ": " + extra_msg + status; 156 | } 157 | 158 | static string descriptor_to_string(const flight::FlightDescriptor &descriptor) 159 | { 160 | if (descriptor.type == flight::FlightDescriptor::DescriptorType::PATH) 161 | { 162 | return joinWithSlash(descriptor.path); 163 | } 164 | else if (descriptor.type == flight::FlightDescriptor::DescriptorType::CMD) 165 | { 166 | return "'" + descriptor.cmd + "'"; 167 | } 168 | else 169 | { 170 | throw NotImplementedException("Unsupported Arrow Flight descriptor type"); 171 | } 172 | } 173 | 174 | string AirportFlightException::produce_flight_error_message(const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg) 175 | { 176 | // So its kind of silly to have the grpc context 177 | auto arrow_status_string = extract_grpc_context_prefix(status.message()); 178 | 179 | string descriptor_string = descriptor_to_string(descriptor); 180 | 181 | return build_error_message(location, descriptor_string, msg, arrow_status_string); 182 | } 183 | 184 | string AirportFlightException::produce_flight_error_message(const string &location, const flight::FlightDescriptor &descriptor, const string &status, const string &msg) 185 | { 186 | // So its kind of silly to have the grpc context 187 | auto arrow_status_string = extract_grpc_context_prefix(status); 188 | 189 | string descriptor_string = descriptor_to_string(descriptor); 190 | 191 | return build_error_message(location, descriptor_string, msg, arrow_status_string); 192 | } 193 | 194 | string AirportFlightException::produce_flight_error_message(const string &location, const arrow::Status &status, const string &msg) 195 | { 196 | // So its kind of silly to have the grpc context 197 | auto arrow_status_string = extract_grpc_context_prefix(status.message()); 198 | 199 | return build_error_message(location, "", msg, arrow_status_string); 200 | } 201 | 202 | string AirportFlightException::produce_flight_error_message(const string &location, const string &msg) 203 | { 204 | // So its kind of silly to have the grpc context 205 | return build_error_message(location, "", msg, ""); 206 | } 207 | 208 | AirportFlightException::AirportFlightException(const string &location, const arrow::Status &status, const string &msg) : Exception(ExceptionType::IO, produce_flight_error_message(location, status, msg), extract_extra_info(status, {})) 209 | { 210 | } 211 | 212 | AirportFlightException::AirportFlightException(const string &location, const string &msg) : Exception(ExceptionType::IO, produce_flight_error_message(location, msg)) 213 | { 214 | } 215 | 216 | AirportFlightException::AirportFlightException(const string &location, const arrow::Status &status, const string &msg, const unordered_map &extra_info) 217 | : Exception(ExceptionType::IO, produce_flight_error_message(location, status, msg), extract_extra_info(status, extra_info)) 218 | { 219 | } 220 | 221 | AirportFlightException::AirportFlightException(const string &location, const flight::FlightDescriptor &descriptor, const string &status, const string &msg) : Exception(ExceptionType::IO, produce_flight_error_message(location, descriptor, status, msg), extract_extra_info(status, {})) 222 | { 223 | } 224 | 225 | AirportFlightException::AirportFlightException(const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg) : Exception(ExceptionType::IO, produce_flight_error_message(location, descriptor, status, msg), extract_extra_info(status, {})) 226 | { 227 | } 228 | 229 | AirportFlightException::AirportFlightException(const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg, const unordered_map &extra_info) 230 | : Exception(ExceptionType::IO, produce_flight_error_message(location, descriptor, status, msg), extract_extra_info(status, extra_info)) 231 | { 232 | } 233 | } -------------------------------------------------------------------------------- /src/airport_json_common.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_json_common.hpp" 2 | 3 | #include "duckdb/common/exception/binder_exception.hpp" 4 | 5 | namespace duckdb { 6 | 7 | using AirportJSONPathType = AirportJSONCommon::JSONPathType; 8 | 9 | string AirportJSONCommon::ValToString(yyjson_val *val, idx_t max_len) { 10 | AirportJSONAllocator json_allocator(Allocator::DefaultAllocator()); 11 | idx_t len; 12 | auto data = AirportJSONCommon::WriteVal(val, json_allocator.GetYYAlc(), len); 13 | if (max_len < len) { 14 | return string(data, max_len) + "..."; 15 | } else { 16 | return string(data, len); 17 | } 18 | } 19 | 20 | void AirportJSONCommon::ThrowValFormatError(string error_string, yyjson_val *val) { 21 | error_string = StringUtil::Format(error_string, AirportJSONCommon::ValToString(val)); 22 | throw InvalidInputException(error_string); 23 | } 24 | 25 | string ThrowPathError(const char *ptr, const char *end, const bool binder) { 26 | ptr--; 27 | auto msg = StringUtil::Format("JSON path error near '%s'", string(ptr, end - ptr)); 28 | if (binder) { 29 | throw BinderException(msg); 30 | } else { 31 | throw InvalidInputException(msg); 32 | } 33 | } 34 | 35 | struct JSONKeyReadResult { 36 | public: 37 | static inline JSONKeyReadResult Empty() { 38 | return {idx_t(0), string()}; 39 | } 40 | 41 | static inline JSONKeyReadResult WildCard() { 42 | return {1, "*"}; 43 | } 44 | 45 | inline bool IsValid() { 46 | return chars_read != 0; 47 | } 48 | 49 | inline bool IsWildCard() { 50 | return key == "*"; 51 | } 52 | 53 | public: 54 | idx_t chars_read; 55 | string key; 56 | }; 57 | 58 | static inline JSONKeyReadResult ReadString(const char *ptr, const char *const end, const bool escaped) { 59 | const char *const before = ptr; 60 | if (escaped) { 61 | auto key = make_unsafe_uniq_array(end - ptr); 62 | idx_t key_len = 0; 63 | 64 | bool backslash = false; 65 | while (ptr != end) { 66 | if (backslash) { 67 | if (*ptr != '"' && *ptr != '\\') { 68 | key[key_len++] = '\\'; 69 | } 70 | backslash = false; 71 | } else { 72 | if (*ptr == '"') { 73 | break; 74 | } else if (*ptr == '\\') { 75 | backslash = true; 76 | ptr++; 77 | continue; 78 | } 79 | } 80 | key[key_len++] = *ptr++; 81 | } 82 | if (ptr == end || backslash) { 83 | return JSONKeyReadResult::Empty(); 84 | } else { 85 | return {idx_t(ptr - before), string(key.get(), key_len)}; 86 | } 87 | } else { 88 | while (ptr != end) { 89 | if (*ptr == '.' || *ptr == '[') { 90 | break; 91 | } 92 | ptr++; 93 | } 94 | return {idx_t(ptr - before), string(before, ptr - before)}; 95 | } 96 | } 97 | 98 | static inline idx_t ReadInteger(const char *ptr, const char *const end, idx_t &idx) { 99 | static constexpr auto IDX_T_SAFE_DIG = 19; 100 | static constexpr auto IDX_T_MAX = ((idx_t)(~(idx_t)0)); 101 | 102 | const char *const before = ptr; 103 | idx = 0; 104 | for (idx_t i = 0; i < IDX_T_SAFE_DIG; i++) { 105 | if (ptr == end) { 106 | // No closing ']' 107 | return 0; 108 | } 109 | if (*ptr == ']') { 110 | break; 111 | } 112 | uint8_t add = (uint8_t)(*ptr - '0'); 113 | if (add <= 9) { 114 | idx = add + idx * 10; 115 | } else { 116 | // Not a digit 117 | return 0; 118 | } 119 | ptr++; 120 | } 121 | // Invalid if overflow 122 | return idx >= (idx_t)IDX_T_MAX ? 0 : ptr - before; 123 | } 124 | 125 | static inline JSONKeyReadResult ReadKey(const char *ptr, const char *const end) { 126 | D_ASSERT(ptr != end); 127 | if (*ptr == '*') { // Wildcard 128 | return JSONKeyReadResult::WildCard(); 129 | } 130 | bool escaped = false; 131 | if (*ptr == '"') { 132 | ptr++; // Skip past opening '"' 133 | escaped = true; 134 | } 135 | auto result = ReadString(ptr, end, escaped); 136 | if (!result.IsValid()) { 137 | return result; 138 | } 139 | if (escaped) { 140 | result.chars_read += 2; // Account for surrounding quotes 141 | } 142 | return result; 143 | } 144 | 145 | static inline bool ReadArrayIndex(const char *&ptr, const char *const end, idx_t &array_index, bool &from_back) { 146 | D_ASSERT(ptr != end); 147 | from_back = false; 148 | if (*ptr == '*') { // Wildcard 149 | ptr++; 150 | if (ptr == end || *ptr != ']') { 151 | return false; 152 | } 153 | array_index = DConstants::INVALID_INDEX; 154 | } else { 155 | if (*ptr == '#') { // SQLite syntax to index from back of array 156 | ptr++; // Skip over '#' 157 | if (ptr == end) { 158 | return false; 159 | } 160 | if (*ptr == ']') { 161 | // [#] always returns NULL in SQLite, so we return an array index that will do the same 162 | array_index = NumericLimits::Maximum(); 163 | ptr++; 164 | return true; 165 | } 166 | if (*ptr != '-') { 167 | return false; 168 | } 169 | from_back = true; 170 | } 171 | if (*ptr == '-') { 172 | ptr++; // Skip over '-' 173 | from_back = true; 174 | } 175 | auto idx_len = ReadInteger(ptr, end, array_index); 176 | if (idx_len == 0) { 177 | return false; 178 | } 179 | ptr += idx_len; 180 | } 181 | ptr++; // Skip past closing ']' 182 | return true; 183 | } 184 | 185 | AirportJSONPathType AirportJSONCommon::ValidatePath(const char *ptr, const idx_t &len, const bool binder) { 186 | D_ASSERT(len >= 1 && *ptr == '$'); 187 | AirportJSONPathType path_type = JSONPathType::REGULAR; 188 | const char *const end = ptr + len; 189 | ptr++; // Skip past '$' 190 | while (ptr != end) { 191 | const auto &c = *ptr++; 192 | if (ptr == end) { 193 | ThrowPathError(ptr, end, binder); 194 | } 195 | switch (c) { 196 | case '.': { // Object field 197 | auto key = ReadKey(ptr, end); 198 | if (!key.IsValid()) { 199 | ThrowPathError(ptr, end, binder); 200 | } else if (key.IsWildCard()) { 201 | path_type = JSONPathType::WILDCARD; 202 | } 203 | ptr += key.chars_read; 204 | break; 205 | } 206 | case '[': { // Array index 207 | idx_t array_index; 208 | bool from_back; 209 | if (!ReadArrayIndex(ptr, end, array_index, from_back)) { 210 | ThrowPathError(ptr, end, binder); 211 | } 212 | if (array_index == DConstants::INVALID_INDEX) { 213 | path_type = JSONPathType::WILDCARD; 214 | } 215 | break; 216 | } 217 | default: 218 | ThrowPathError(ptr, end, binder); 219 | } 220 | } 221 | return path_type; 222 | } 223 | 224 | yyjson_val *AirportJSONCommon::GetPath(yyjson_val *val, const char *ptr, const idx_t &len) { 225 | // Path has been validated at this point 226 | const char *const end = ptr + len; 227 | ptr++; // Skip past '$' 228 | while (val != nullptr && ptr != end) { 229 | const auto &c = *ptr++; 230 | D_ASSERT(ptr != end); 231 | switch (c) { 232 | case '.': { // Object field 233 | if (!unsafe_yyjson_is_obj(val)) { 234 | return nullptr; 235 | } 236 | auto key_result = ReadKey(ptr, end); 237 | D_ASSERT(key_result.IsValid()); 238 | ptr += key_result.chars_read; 239 | val = yyjson_obj_getn(val, key_result.key.c_str(), key_result.key.size()); 240 | break; 241 | } 242 | case '[': { // Array index 243 | if (!unsafe_yyjson_is_arr(val)) { 244 | return nullptr; 245 | } 246 | idx_t array_index; 247 | bool from_back; 248 | #ifdef DEBUG 249 | bool success = 250 | #endif 251 | ReadArrayIndex(ptr, end, array_index, from_back); 252 | #ifdef DEBUG 253 | D_ASSERT(success); 254 | #endif 255 | if (from_back && array_index != 0) { 256 | array_index = unsafe_yyjson_get_len(val) - array_index; 257 | } 258 | val = yyjson_arr_get(val, array_index); 259 | break; 260 | } 261 | default: // LCOV_EXCL_START 262 | throw InternalException( 263 | "Invalid JSON Path encountered in JSONCommon::GetPath, call JSONCommon::ValidatePath first!"); 264 | } // LCOV_EXCL_STOP 265 | } 266 | return val; 267 | } 268 | 269 | void GetWildcardPathInternal(yyjson_val *val, const char *ptr, const char *const end, vector &vals) { 270 | while (val != nullptr && ptr != end) { 271 | const auto &c = *ptr++; 272 | D_ASSERT(ptr != end); 273 | switch (c) { 274 | case '.': { // Object field 275 | if (!unsafe_yyjson_is_obj(val)) { 276 | return; 277 | } 278 | auto key_result = ReadKey(ptr, end); 279 | D_ASSERT(key_result.IsValid()); 280 | ptr += key_result.chars_read; 281 | if (key_result.IsWildCard()) { // Wildcard 282 | size_t idx, max; 283 | yyjson_val *key, *obj_val; 284 | yyjson_obj_foreach(val, idx, max, key, obj_val) { 285 | GetWildcardPathInternal(obj_val, ptr, end, vals); 286 | } 287 | return; 288 | } 289 | val = yyjson_obj_getn(val, key_result.key.c_str(), key_result.key.size()); 290 | break; 291 | } 292 | case '[': { // Array index 293 | if (!unsafe_yyjson_is_arr(val)) { 294 | return; 295 | } 296 | idx_t array_index; 297 | bool from_back; 298 | #ifdef DEBUG 299 | bool success = 300 | #endif 301 | ReadArrayIndex(ptr, end, array_index, from_back); 302 | #ifdef DEBUG 303 | D_ASSERT(success); 304 | #endif 305 | 306 | if (array_index == DConstants::INVALID_INDEX) { // Wildcard 307 | size_t idx, max; 308 | yyjson_val *arr_val; 309 | yyjson_arr_foreach(val, idx, max, arr_val) { 310 | GetWildcardPathInternal(arr_val, ptr, end, vals); 311 | } 312 | return; 313 | } 314 | if (from_back && array_index != 0) { 315 | array_index = unsafe_yyjson_get_len(val) - array_index; 316 | } 317 | val = yyjson_arr_get(val, array_index); 318 | break; 319 | } 320 | default: // LCOV_EXCL_START 321 | throw InternalException( 322 | "Invalid JSON Path encountered in GetWildcardPathInternal, call JSONCommon::ValidatePath first!"); 323 | } // LCOV_EXCL_STOP 324 | } 325 | if (val != nullptr) { 326 | vals.emplace_back(val); 327 | } 328 | } 329 | 330 | void AirportJSONCommon::GetWildcardPath(yyjson_val *val, const char *ptr, const idx_t &len, vector &vals) { 331 | // Path has been validated at this point 332 | const char *const end = ptr + len; 333 | ptr++; // Skip past '$' 334 | GetWildcardPathInternal(val, ptr, end, vals); 335 | } 336 | 337 | } // namespace duckdb 338 | -------------------------------------------------------------------------------- /src/airport_json_serializer.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_json_serializer.hpp" 2 | #include "duckdb/common/types/blob.hpp" 3 | 4 | namespace duckdb { 5 | 6 | void AirportJsonSerializer::PushValue(yyjson_mut_val *val) { 7 | auto current = Current(); 8 | // Array case, just append the value 9 | if (yyjson_mut_is_arr(current)) { 10 | yyjson_mut_arr_append(current, val); 11 | } 12 | // Object case, use the currently set tag. 13 | else if (yyjson_mut_is_obj(current)) { 14 | yyjson_mut_obj_add(current, current_tag, val); 15 | } 16 | // Else throw 17 | else { 18 | throw InternalException("Cannot add value to non-array/object json value"); 19 | } 20 | } 21 | 22 | void AirportJsonSerializer::OnPropertyBegin(const field_id_t, const char *tag) { 23 | current_tag = yyjson_mut_strcpy(doc, tag); 24 | } 25 | 26 | void AirportJsonSerializer::OnPropertyEnd() { 27 | } 28 | 29 | void AirportJsonSerializer::OnOptionalPropertyBegin(const field_id_t, const char *tag, bool) { 30 | current_tag = yyjson_mut_strcpy(doc, tag); 31 | } 32 | 33 | void AirportJsonSerializer::OnOptionalPropertyEnd(bool) { 34 | } 35 | 36 | //------------------------------------------------------------------------- 37 | // Nested Types 38 | //------------------------------------------------------------------------- 39 | void AirportJsonSerializer::OnNullableBegin(bool present) { 40 | if (!present && !skip_if_null) { 41 | WriteNull(); 42 | } 43 | } 44 | 45 | void AirportJsonSerializer::OnNullableEnd() { 46 | } 47 | 48 | void AirportJsonSerializer::OnListBegin(idx_t count) { 49 | auto new_value = yyjson_mut_arr(doc); 50 | // We always push a value to the stack, we just don't add it as a child to the current value 51 | // if skipping empty. Even though it is "unnecessary" to create an empty value just to discard it, 52 | // this allows the rest of the code to keep on like normal. 53 | if (!(count == 0 && skip_if_empty)) { 54 | PushValue(new_value); 55 | } 56 | stack.push_back(new_value); 57 | } 58 | 59 | void AirportJsonSerializer::OnListEnd() { 60 | stack.pop_back(); 61 | } 62 | 63 | void AirportJsonSerializer::OnObjectBegin() { 64 | auto new_value = yyjson_mut_obj(doc); 65 | PushValue(new_value); 66 | stack.push_back(new_value); 67 | } 68 | 69 | void AirportJsonSerializer::OnObjectEnd() { 70 | auto obj = Current(); 71 | auto count = yyjson_mut_obj_size(obj); 72 | 73 | stack.pop_back(); 74 | 75 | if (count == 0 && skip_if_empty && !stack.empty()) { 76 | // remove obj from parent since it was empty 77 | auto parent = Current(); 78 | if (yyjson_mut_is_arr(parent)) { 79 | size_t idx; 80 | size_t max; 81 | yyjson_mut_val *item; 82 | size_t found; 83 | yyjson_mut_arr_foreach(parent, idx, max, item) { 84 | if (item == obj) { 85 | found = idx; 86 | } 87 | } 88 | yyjson_mut_arr_remove(parent, found); 89 | } else if (yyjson_mut_is_obj(parent)) { 90 | size_t idx; 91 | size_t max; 92 | yyjson_mut_val *item; 93 | yyjson_mut_val *key; 94 | const char *found; 95 | yyjson_mut_obj_foreach(parent, idx, max, key, item) { 96 | if (item == obj) { 97 | found = yyjson_mut_get_str(key); 98 | } 99 | } 100 | yyjson_mut_obj_remove_key(parent, found); 101 | } 102 | } 103 | } 104 | 105 | //------------------------------------------------------------------------- 106 | // Primitive Types 107 | //------------------------------------------------------------------------- 108 | void AirportJsonSerializer::WriteNull() { 109 | if (skip_if_null) { 110 | return; 111 | } 112 | auto val = yyjson_mut_null(doc); 113 | PushValue(val); 114 | } 115 | 116 | void AirportJsonSerializer::WriteValue(uint8_t value) { 117 | auto val = yyjson_mut_uint(doc, value); 118 | PushValue(val); 119 | } 120 | 121 | void AirportJsonSerializer::WriteValue(int8_t value) { 122 | auto val = yyjson_mut_sint(doc, value); 123 | PushValue(val); 124 | } 125 | 126 | void AirportJsonSerializer::WriteValue(uint16_t value) { 127 | auto val = yyjson_mut_uint(doc, value); 128 | PushValue(val); 129 | } 130 | 131 | void AirportJsonSerializer::WriteValue(int16_t value) { 132 | auto val = yyjson_mut_sint(doc, value); 133 | PushValue(val); 134 | } 135 | 136 | void AirportJsonSerializer::WriteValue(uint32_t value) { 137 | auto val = yyjson_mut_uint(doc, value); 138 | PushValue(val); 139 | } 140 | 141 | void AirportJsonSerializer::WriteValue(int32_t value) { 142 | auto val = yyjson_mut_sint(doc, value); 143 | PushValue(val); 144 | } 145 | 146 | void AirportJsonSerializer::WriteValue(uint64_t value) { 147 | auto val = yyjson_mut_uint(doc, value); 148 | PushValue(val); 149 | } 150 | 151 | void AirportJsonSerializer::WriteValue(int64_t value) { 152 | auto val = yyjson_mut_sint(doc, value); 153 | PushValue(val); 154 | } 155 | 156 | void AirportJsonSerializer::WriteValue(hugeint_t value) { 157 | auto val = yyjson_mut_obj(doc); 158 | PushValue(val); 159 | stack.push_back(val); 160 | WriteProperty(100, "upper", value.upper); 161 | WriteProperty(101, "lower", value.lower); 162 | stack.pop_back(); 163 | } 164 | 165 | void AirportJsonSerializer::WriteValue(uhugeint_t value) { 166 | auto val = yyjson_mut_obj(doc); 167 | PushValue(val); 168 | stack.push_back(val); 169 | WriteProperty(100, "upper", value.upper); 170 | WriteProperty(101, "lower", value.lower); 171 | stack.pop_back(); 172 | } 173 | 174 | void AirportJsonSerializer::WriteValue(float value) { 175 | auto val = yyjson_mut_real(doc, value); 176 | PushValue(val); 177 | } 178 | 179 | void AirportJsonSerializer::WriteValue(double value) { 180 | auto val = yyjson_mut_real(doc, value); 181 | PushValue(val); 182 | } 183 | 184 | void AirportJsonSerializer::WriteValue(const string &value) { 185 | if (skip_if_empty && value.empty()) { 186 | return; 187 | } 188 | auto val = yyjson_mut_strncpy(doc, value.c_str(), value.size()); 189 | PushValue(val); 190 | } 191 | 192 | void AirportJsonSerializer::WriteValue(const string_t value) { 193 | if (skip_if_empty && value.GetSize() == 0) { 194 | return; 195 | } 196 | auto val = yyjson_mut_strncpy(doc, value.GetData(), value.GetSize()); 197 | PushValue(val); 198 | } 199 | 200 | void AirportJsonSerializer::WriteValue(const char *value) { 201 | if (skip_if_empty && strlen(value) == 0) { 202 | return; 203 | } 204 | auto val = yyjson_mut_strcpy(doc, value); 205 | PushValue(val); 206 | } 207 | 208 | void AirportJsonSerializer::WriteValue(bool value) { 209 | auto val = yyjson_mut_bool(doc, value); 210 | PushValue(val); 211 | } 212 | 213 | void AirportJsonSerializer::WriteDataPtr(const_data_ptr_t ptr, idx_t count) { 214 | auto blob = Blob::ToBlob(string_t(const_char_ptr_cast(ptr), count)); 215 | auto val = yyjson_mut_strcpy(doc, blob.c_str()); 216 | PushValue(val); 217 | } 218 | 219 | } // namespace duckdb 220 | -------------------------------------------------------------------------------- /src/airport_optimizer.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "storage/airport_catalog.hpp" 3 | #include "airport_optimizer.hpp" 4 | #include "duckdb/planner/operator/logical_delete.hpp" 5 | #include "duckdb/planner/operator/logical_update.hpp" 6 | #include "duckdb/planner/operator/logical_get.hpp" 7 | #include "duckdb/planner/operator/logical_filter.hpp" 8 | #include "duckdb/planner/operator/logical_projection.hpp" 9 | #include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp" 10 | #include "airport_flight_stream.hpp" 11 | #include "storage/airport_table_entry.hpp" 12 | 13 | namespace duckdb 14 | { 15 | static void MarkAirportTakeFlightAsSkipProducing(unique_ptr &op) 16 | { 17 | auto markAirportTakeFlightSkipResults = [](LogicalGet &get) 18 | { 19 | if (get.function.name == "airport_take_flight") 20 | { 21 | auto &bind_data = get.bind_data->Cast(); 22 | bind_data.skip_producing_result_for_update_or_delete = true; 23 | } 24 | }; 25 | 26 | reference child = *op->children[0]; 27 | if (child.get().type == LogicalOperatorType::LOGICAL_GET) 28 | { 29 | markAirportTakeFlightSkipResults(child.get().Cast()); 30 | } 31 | else if (child.get().type == LogicalOperatorType::LOGICAL_FILTER || 32 | child.get().type == LogicalOperatorType::LOGICAL_PROJECTION) 33 | { 34 | for (auto &child_node : op->children) 35 | { 36 | MarkAirportTakeFlightAsSkipProducing(child_node); 37 | } 38 | } 39 | else 40 | { 41 | throw NotImplementedException("Unsupported child type for LogicalUpdate type is" + LogicalOperatorToString(child.get().type)); 42 | } 43 | } 44 | 45 | void OptimizeAirportUpdate(unique_ptr &op) 46 | { 47 | if (op->type != LogicalOperatorType::LOGICAL_UPDATE) 48 | return; 49 | 50 | auto &update = op->Cast(); 51 | auto &airport_table = update.table.Cast(); 52 | 53 | // If the table produced rowids we cannot optimize it. 54 | if (airport_table.GetRowIdType() != LogicalType::SQLNULL) 55 | return; 56 | 57 | MarkAirportTakeFlightAsSkipProducing(op); 58 | } 59 | 60 | void OptimizeAirportDelete(unique_ptr &op) 61 | { 62 | if (op->type != LogicalOperatorType::LOGICAL_DELETE) 63 | return; 64 | 65 | auto &del = op->Cast(); 66 | auto &airport_table = del.table.Cast(); 67 | 68 | // If the table produced rowids we cannot optimize it. 69 | if (airport_table.GetRowIdType() != LogicalType::SQLNULL) 70 | return; 71 | 72 | MarkAirportTakeFlightAsSkipProducing(op); 73 | } 74 | 75 | void AirportOptimizer::Optimize(OptimizerExtensionInput &input, unique_ptr &plan) 76 | { 77 | OptimizeAirportUpdate(plan); 78 | OptimizeAirportDelete(plan); 79 | } 80 | } -------------------------------------------------------------------------------- /src/airport_request_headers.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_request_headers.hpp" 2 | #include 3 | #include "duckdb/common/types/uuid.hpp" 4 | #include 5 | 6 | // Indicate the version of the caller. 7 | #define AIRPORT_USER_AGENT "airport/20240820" 8 | 9 | namespace duckdb 10 | { 11 | 12 | string airport_trace_id() 13 | { 14 | return UUID::ToString(UUID::GenerateRandomUUID()); 15 | } 16 | 17 | std::string airport_user_agent() 18 | { 19 | return AIRPORT_USER_AGENT; 20 | } 21 | 22 | // Generate a random session id for each time that DuckDB starts, 23 | // this can be useful on the server side for tracking sessions. 24 | static const std::string airport_session_id = UUID::ToString(UUID::GenerateRandomUUID()); 25 | 26 | static void 27 | airport_add_headers(std::vector> &headers, const std::string &server_location) 28 | { 29 | headers.emplace_back("airport-user-agent", AIRPORT_USER_AGENT); 30 | headers.emplace_back("authority", server_location); 31 | headers.emplace_back("airport-client-session-id", airport_session_id); 32 | } 33 | 34 | void airport_add_standard_headers(arrow::flight::FlightCallOptions &options, const std::string &server_location) 35 | { 36 | airport_add_headers(options.headers, server_location); 37 | } 38 | 39 | void airport_add_authorization_header(arrow::flight::FlightCallOptions &options, const std::string &auth_token) 40 | { 41 | if (auth_token.empty()) 42 | { 43 | return; 44 | } 45 | 46 | options.headers.emplace_back("authorization", "Bearer " + auth_token); 47 | } 48 | 49 | static std::string join_vector_of_strings(const std::vector &vec, const char joiner) 50 | { 51 | if (vec.empty()) 52 | return ""; 53 | 54 | return std::accumulate( 55 | std::next(vec.begin()), vec.end(), vec.front(), 56 | [joiner](const std::string &a, const std::string &b) 57 | { 58 | return a + joiner + b; 59 | }); 60 | } 61 | 62 | void airport_add_flight_path_header(arrow::flight::FlightCallOptions &options, 63 | const arrow::flight::FlightDescriptor &descriptor) 64 | { 65 | if (descriptor.type == arrow::flight::FlightDescriptor::PATH) 66 | { 67 | auto path_parts = descriptor.path; 68 | std::string joined_path_parts = join_vector_of_strings(path_parts, '/'); 69 | options.headers.emplace_back("airport-flight-path", joined_path_parts); 70 | } 71 | } 72 | 73 | void airport_add_trace_id_header(arrow::flight::FlightCallOptions &options, 74 | const string &trace_id) 75 | { 76 | options.headers.emplace_back("airport-trace-id", trace_id); 77 | } 78 | 79 | void airport_add_normal_headers(arrow::flight::FlightCallOptions &options, 80 | const AirportTakeFlightParameters ¶ms, 81 | const string &trace_id, 82 | std::optional descriptor) 83 | { 84 | airport_add_standard_headers(options, params.server_location()); 85 | airport_add_authorization_header(options, params.auth_token()); 86 | airport_add_trace_id_header(options, trace_id); 87 | 88 | for (const auto &header_pair : params.user_supplied_headers()) 89 | { 90 | for (const auto &header_value : header_pair.second) 91 | { 92 | options.headers.emplace_back(header_pair.first, header_value); 93 | } 94 | } 95 | 96 | if (descriptor.has_value()) 97 | { 98 | const auto &flight_descriptor = descriptor.value(); 99 | airport_add_flight_path_header(options, flight_descriptor); 100 | } 101 | } 102 | } -------------------------------------------------------------------------------- /src/airport_schema_utils.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "duckdb.hpp" 3 | 4 | #include "duckdb/common/arrow/schema_metadata.hpp" 5 | #include "airport_take_flight.hpp" 6 | 7 | namespace duckdb 8 | { 9 | 10 | bool AirportFieldMetadataIsRowId(const char *metadata) 11 | { 12 | if (metadata == nullptr) 13 | { 14 | return false; 15 | } 16 | ArrowSchemaMetadata column_metadata(metadata); 17 | auto comment = column_metadata.GetOption("is_rowid"); 18 | if (!comment.empty()) 19 | { 20 | return true; 21 | } 22 | return false; 23 | } 24 | 25 | void AirportExamineSchema( 26 | ClientContext &context, 27 | const ArrowSchemaWrapper &schema_root, 28 | ArrowTableType *arrow_table, 29 | vector *return_types, 30 | vector *names, 31 | vector *duckdb_type_names, 32 | idx_t *rowid_column_index, 33 | bool skip_rowid_column) 34 | { 35 | if (rowid_column_index) 36 | { 37 | *rowid_column_index = COLUMN_IDENTIFIER_ROW_ID; 38 | } 39 | 40 | auto &config = DBConfig::GetConfig(context); 41 | 42 | const idx_t num_columns = static_cast(schema_root.arrow_schema.n_children); 43 | 44 | if (num_columns > 0) 45 | { 46 | if (return_types) 47 | { 48 | return_types->reserve(num_columns); 49 | } 50 | if (names) 51 | { 52 | names->reserve(num_columns); 53 | } 54 | if (duckdb_type_names) 55 | { 56 | duckdb_type_names->reserve(num_columns); 57 | } 58 | } 59 | 60 | for (idx_t col_idx = 0; col_idx < num_columns; col_idx++) 61 | { 62 | auto &schema = *schema_root.arrow_schema.children[col_idx]; 63 | if (!schema.release) 64 | { 65 | throw InvalidInputException("AirportExamineSchema: released schema passed"); 66 | } 67 | auto arrow_type = ArrowType::GetArrowLogicalType(config, schema); 68 | 69 | // Determine if the column is the rowid column by looking at the metadata 70 | // on the column. 71 | bool is_rowid_column = false; 72 | if (AirportFieldMetadataIsRowId(schema.metadata)) 73 | { 74 | is_rowid_column = true; 75 | if (rowid_column_index) 76 | { 77 | *rowid_column_index = col_idx; 78 | } 79 | } 80 | 81 | if (schema.dictionary) 82 | { 83 | auto dictionary_type = ArrowType::GetArrowLogicalType(config, *schema.dictionary); 84 | arrow_type->SetDictionary(std::move(dictionary_type)); 85 | } 86 | 87 | const idx_t column_id = is_rowid_column ? COLUMN_IDENTIFIER_ROW_ID : col_idx; 88 | 89 | const string column_name = AirportNameForField(schema.name, col_idx); 90 | 91 | if (!skip_rowid_column || !is_rowid_column) 92 | { 93 | auto duck_type = arrow_type->GetDuckType(); 94 | if (return_types) 95 | { 96 | return_types->push_back(duck_type); 97 | } 98 | if (names) 99 | { 100 | names->push_back(std::move(column_name)); 101 | } 102 | if (duckdb_type_names) 103 | { 104 | duckdb_type_names->push_back(duck_type.ToString()); 105 | } 106 | } 107 | 108 | if (arrow_table) 109 | { 110 | arrow_table->AddColumn(column_id, std::move(arrow_type)); 111 | } 112 | } 113 | QueryResult::DeduplicateColumns(*names); 114 | } 115 | 116 | } -------------------------------------------------------------------------------- /src/airport_secrets.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | #include "duckdb/common/exception.hpp" 3 | 4 | #include "duckdb/main/secret/secret_manager.hpp" 5 | #include "airport_secrets.hpp" 6 | 7 | namespace duckdb 8 | { 9 | 10 | unique_ptr AirportGetSecretByName(ClientContext &context, const string &secret_name) 11 | { 12 | auto &secret_manager = SecretManager::Get(context); 13 | auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context); 14 | // FIXME: this should be adjusted once the `GetSecretByName` API supports this 15 | // use case 16 | auto secret_entry = secret_manager.GetSecretByName(transaction, secret_name, "memory"); 17 | if (secret_entry) 18 | { 19 | return secret_entry; 20 | } 21 | secret_entry = secret_manager.GetSecretByName(transaction, secret_name, "local_file"); 22 | if (secret_entry) 23 | { 24 | return secret_entry; 25 | } 26 | return nullptr; 27 | } 28 | 29 | SecretMatch AirportGetSecretByPath(ClientContext &context, const string &path) 30 | { 31 | auto &secret_manager = SecretManager::Get(context); 32 | auto transaction = CatalogTransaction::GetSystemCatalogTransaction(context); 33 | return secret_manager.LookupSecret(transaction, path, "airport"); 34 | } 35 | 36 | string AirportAuthTokenForLocation(ClientContext &context, const string &server_location, const string &secret_name, const string &auth_token) 37 | { 38 | if (!auth_token.empty()) 39 | { 40 | return auth_token; 41 | } 42 | 43 | if (!secret_name.empty()) 44 | { 45 | auto secret_entry = AirportGetSecretByName(context, secret_name); 46 | if (!secret_entry) 47 | { 48 | throw BinderException("Secret with name \"%s\" not found", secret_name); 49 | } 50 | 51 | const auto &kv_secret = dynamic_cast(*secret_entry->secret); 52 | 53 | if (auth_token.empty()) 54 | { 55 | Value input_val = kv_secret.TryGetValue("auth_token"); 56 | if (!input_val.IsNull()) 57 | { 58 | return input_val.ToString(); 59 | } 60 | return ""; 61 | } 62 | } 63 | 64 | auto secret_match = AirportGetSecretByPath(context, server_location); 65 | if (secret_match.HasMatch()) 66 | { 67 | const auto &kv_secret = dynamic_cast(*secret_match.secret_entry->secret); 68 | 69 | Value input_val = kv_secret.TryGetValue("auth_token"); 70 | if (!input_val.IsNull()) 71 | { 72 | return input_val.ToString(); 73 | } 74 | return ""; 75 | } 76 | return ""; 77 | } 78 | } -------------------------------------------------------------------------------- /src/include/airport_constraints.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "duckdb/common/arrow/schema_metadata.hpp" 3 | #include "duckdb/storage/table/append_state.hpp" 4 | 5 | #include "airport_extension.hpp" 6 | #include "storage/airport_insert.hpp" 7 | #include "storage/airport_catalog.hpp" 8 | #include "storage/airport_transaction.hpp" 9 | #include "storage/airport_table_entry.hpp" 10 | 11 | namespace duckdb 12 | { 13 | void AirportVerifyAppendConstraints(ConstraintState &state, ClientContext &context, DataChunk &chunk, 14 | optional_ptr conflict_manager, 15 | vector column_names); 16 | } 17 | -------------------------------------------------------------------------------- /src/include/airport_extension.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb.hpp" 4 | #include 5 | #include 6 | #include 7 | 8 | #ifdef _WIN32 9 | 10 | // Windows defines min and max macros that mess up std::min/max 11 | #ifndef NOMINMAX 12 | #define NOMINMAX 13 | #endif 14 | 15 | // #define WIN32_LEAN_AND_MEAN 16 | 17 | // Set Windows 7 as a conservative minimum for Apache Arrow 18 | // #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x601 19 | // #undef _WIN32_WINNT 20 | // #endif 21 | // #ifndef _WIN32_WINNT 22 | // #define _WIN32_WINNT 0x601 23 | // #endif 24 | 25 | #ifdef max 26 | #undef max 27 | #endif 28 | #ifdef min 29 | #undef min 30 | #endif 31 | 32 | // // The Windows API defines macros from *File resolving to either 33 | // // *FileA or *FileW. Need to undo them. 34 | // #ifdef CopyFile 35 | // #undef CopyFile 36 | // #endif 37 | // #ifdef CreateFile 38 | // #undef CreateFile 39 | // #endif 40 | // #ifdef DeleteFile 41 | // #undef DeleteFile 42 | // #endif 43 | 44 | // // Other annoying Windows macro definitions... 45 | // #ifdef IN 46 | // #undef IN 47 | // #endif 48 | // #ifdef OUT 49 | // #undef OUT 50 | // #endif 51 | 52 | // Note that we can't undefine OPTIONAL, because it can be used in other 53 | // Windows headers... 54 | 55 | #endif // _WIN32 56 | 57 | namespace duckdb 58 | { 59 | 60 | class AirportExtension : public Extension 61 | { 62 | public: 63 | void Load(DuckDB &db) override; 64 | std::string Name() override; 65 | std::string Version() const override; 66 | }; 67 | 68 | void AirportAddListFlightsFunction(DatabaseInstance &instance); 69 | void AirportAddTakeFlightFunction(DatabaseInstance &instance); 70 | void AirportAddActionFlightFunction(DatabaseInstance &instance); 71 | 72 | } // namespace duckdb 73 | -------------------------------------------------------------------------------- /src/include/airport_flight_exception.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/common/exception.hpp" 4 | #include 5 | #include 6 | 7 | #include 8 | namespace flight = arrow::flight; 9 | 10 | namespace duckdb 11 | { 12 | 13 | // A custom exception class for Airport that can include the flight descriptions 14 | // and server location, it is important to surface this information on errors because 15 | // if multiple servers are being used, its helpful to help the user to diagnose 16 | // where the error is being caused. 17 | class AirportFlightException : public Exception 18 | { 19 | 20 | private: 21 | static string produce_flight_error_message(const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg); 22 | static string produce_flight_error_message(const string &location, const flight::FlightDescriptor &descriptor, const string &status, const string &msg); 23 | static string produce_flight_error_message(const string &location, const arrow::Status &status, const string &msg); 24 | static string produce_flight_error_message(const string &location, const string &msg); 25 | 26 | static unordered_map extract_extra_info(const arrow::Status &status, const unordered_map &extra_info); 27 | static unordered_map extract_extra_info(const string &status, const unordered_map &extra_info); 28 | 29 | public: 30 | DUCKDB_API explicit AirportFlightException(const string &location, const arrow::Status &status, const string &msg); 31 | DUCKDB_API explicit AirportFlightException(const string &location, const string &msg); 32 | DUCKDB_API explicit AirportFlightException(const string &location, const arrow::Status &status, const string &msg, const unordered_map &extra_info); 33 | 34 | DUCKDB_API explicit AirportFlightException(const string &location, const flight::FlightDescriptor &descriptor, const string &status, const string &msg); 35 | DUCKDB_API explicit AirportFlightException(const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg); 36 | DUCKDB_API explicit AirportFlightException(const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg, const unordered_map &extra_info); 37 | explicit AirportFlightException(ExceptionType exception_type, const string &location, const flight::FlightDescriptor &descriptor, const arrow::Status &status, const string &msg) : Exception(exception_type, produce_flight_error_message(location, descriptor, status, msg), extract_extra_info(status, {})) 38 | { 39 | } 40 | 41 | explicit AirportFlightException(ExceptionType exception_type, const string &location, const arrow::Status &status, const string &msg) : Exception(exception_type, produce_flight_error_message(location, status, msg), extract_extra_info(status, {})) 42 | { 43 | } 44 | }; 45 | } -------------------------------------------------------------------------------- /src/include/airport_flight_statistics.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "airport_extension.hpp" 3 | #include "duckdb.hpp" 4 | 5 | namespace duckdb 6 | { 7 | unique_ptr AirportTakeFlightStatistics(ClientContext &context, const FunctionData *bind_data, column_t column_index); 8 | 9 | } -------------------------------------------------------------------------------- /src/include/airport_json_serializer.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "airport_json_common.hpp" 4 | #include "duckdb/common/serializer/serializer.hpp" 5 | 6 | namespace duckdb { 7 | 8 | struct AirportJsonSerializer : Serializer { 9 | private: 10 | yyjson_mut_doc *doc; 11 | yyjson_mut_val *current_tag; 12 | vector stack; 13 | 14 | // Skip writing property if null 15 | bool skip_if_null = false; 16 | // Skip writing property if empty string, empty list or empty map. 17 | bool skip_if_empty = false; 18 | 19 | // Get the current json value 20 | inline yyjson_mut_val *Current() { 21 | return stack.back(); 22 | }; 23 | 24 | // Either adds a value to the current object with the current tag, or appends it to the current array 25 | void PushValue(yyjson_mut_val *val); 26 | 27 | public: 28 | explicit AirportJsonSerializer(yyjson_mut_doc *doc, bool skip_if_null, bool skip_if_empty, bool skip_if_default) 29 | : doc(doc), stack({yyjson_mut_obj(doc)}), skip_if_null(skip_if_null), skip_if_empty(skip_if_empty) { 30 | options.serialize_enum_as_string = true; 31 | options.serialize_default_values = !skip_if_default; 32 | } 33 | 34 | template 35 | static yyjson_mut_val *Serialize(T &value, yyjson_mut_doc *doc, bool skip_if_null, bool skip_if_empty, 36 | bool skip_if_default) { 37 | AirportJsonSerializer serializer(doc, skip_if_null, skip_if_empty, skip_if_default); 38 | value.Serialize(serializer); 39 | return serializer.GetRootObject(); 40 | } 41 | 42 | yyjson_mut_val *GetRootObject() { 43 | D_ASSERT(stack.size() == 1); // or we forgot to pop somewhere 44 | return stack.front(); 45 | }; 46 | 47 | //===--------------------------------------------------------------------===// 48 | // Nested Types Hooks 49 | //===--------------------------------------------------------------------===// 50 | void OnPropertyBegin(const field_id_t field_id, const char *tag) final; 51 | void OnPropertyEnd() final; 52 | void OnOptionalPropertyBegin(const field_id_t field_id, const char *tag, bool present) final; 53 | void OnOptionalPropertyEnd(bool present) final; 54 | 55 | void OnListBegin(idx_t count) final; 56 | void OnListEnd() final; 57 | void OnObjectBegin() final; 58 | void OnObjectEnd() final; 59 | void OnNullableBegin(bool present) final; 60 | void OnNullableEnd() final; 61 | 62 | //===--------------------------------------------------------------------===// 63 | // Primitive Types 64 | //===--------------------------------------------------------------------===// 65 | void WriteNull() final; 66 | void WriteValue(uint8_t value) final; 67 | void WriteValue(int8_t value) final; 68 | void WriteValue(uint16_t value) final; 69 | void WriteValue(int16_t value) final; 70 | void WriteValue(uint32_t value) final; 71 | void WriteValue(int32_t value) final; 72 | void WriteValue(uint64_t value) final; 73 | void WriteValue(int64_t value) final; 74 | void WriteValue(hugeint_t value) final; 75 | void WriteValue(uhugeint_t value) final; 76 | void WriteValue(float value) final; 77 | void WriteValue(double value) final; 78 | void WriteValue(const string_t value) final; 79 | void WriteValue(const string &value) final; 80 | void WriteValue(const char *value) final; 81 | void WriteValue(bool value) final; 82 | void WriteDataPtr(const_data_ptr_t ptr, idx_t count) final; 83 | }; 84 | 85 | } // namespace duckdb 86 | -------------------------------------------------------------------------------- /src/include/airport_location_descriptor.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb.hpp" 4 | #include "arrow/flight/client.h" 5 | 6 | namespace flight = arrow::flight; 7 | 8 | namespace duckdb 9 | { 10 | 11 | // Capture a server_location and a Flight descriptor, make this a struct 12 | // since many other classes/structs use this combination of members. 13 | struct AirportLocationDescriptor 14 | { 15 | public: 16 | AirportLocationDescriptor(const string &server_location, 17 | const flight::FlightDescriptor &descriptor) 18 | : server_location_(server_location), descriptor_(descriptor) 19 | { 20 | } 21 | 22 | const string &server_location() const 23 | { 24 | return server_location_; 25 | } 26 | 27 | const flight::FlightDescriptor &descriptor() const 28 | { 29 | return descriptor_; 30 | } 31 | 32 | private: 33 | const string server_location_; 34 | const flight::FlightDescriptor descriptor_; 35 | }; 36 | 37 | } -------------------------------------------------------------------------------- /src/include/airport_macros.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "airport_flight_exception.hpp" 4 | 5 | #define AIRPORT_ARROW_ASSERT_OK_LOCATION(expr, location, message) \ 6 | for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); !_st.ok();) \ 7 | throw AirportFlightException(location, _st, ""); 8 | 9 | #define AIRPORT_ARROW_ASSERT_OK_LOCATION_DESCRIPTOR(expr, location, descriptor, message) \ 10 | for (::arrow::Status _st = ::arrow::internal::GenericToStatus((expr)); !_st.ok();) \ 11 | throw AirportFlightException(location, descriptor, _st, message); 12 | 13 | #define AIRPORT_ARROW_ASSERT_OK_CONTAINER(expr, container, message) \ 14 | AIRPORT_ARROW_ASSERT_OK_LOCATION_DESCRIPTOR(expr, (container)->server_location(), (container)->descriptor(), message) 15 | 16 | #define AIRPORT_ASSERT_OK_LOCATION_DESCRIPTOR(expr, location, descriptor, message) \ 17 | if (!expr) \ 18 | { \ 19 | throw AirportFlightException(location, descriptor, #expr, message); \ 20 | } 21 | 22 | #define AIRPORT_ASSERT_OK_CONTAINER(expr, container, message) \ 23 | AIRPORT_ASSERT_OK_LOCATION_DESCRIPTOR(expr, container->server_location(), container->descriptor(), message) 24 | 25 | #define AIRPORT_RETURN_IF_LOCATION_DESCRIPTOR(error_prefix, condition, status, location, descriptor, message, extra_message) \ 26 | do \ 27 | { \ 28 | if (ARROW_PREDICT_FALSE(condition)) \ 29 | { \ 30 | throw AirportFlightException(location, descriptor, status, message, {{"extra_details", string(extra_message)}}); \ 31 | } \ 32 | } while (0) 33 | 34 | #define AIRPORT_RETURN_IF_LOCATION(error_prefix, condition, status, location, message, extra_message) \ 35 | do \ 36 | { \ 37 | if (ARROW_PREDICT_FALSE(condition)) \ 38 | { \ 39 | throw AirportFlightException(location, status, message, {{"extra_details", string(extra_message)}}); \ 40 | } \ 41 | } while (0) 42 | 43 | #define AIRPORT_ASSIGN_OR_RAISE_IMPL_LOCATION_DESCRIPTOR(result_name, lhs, rexpr, location, descriptor, message) \ 44 | auto &&result_name = (rexpr); \ 45 | AIRPORT_RETURN_IF_LOCATION_DESCRIPTOR(error_prefix, !(result_name).ok(), (result_name).status(), location, descriptor, message, ARROW_STRINGIFY(rexpr)); \ 46 | lhs = std::move(result_name).ValueUnsafe(); 47 | 48 | #define AIRPORT_ASSIGN_OR_RAISE_IMPL_LOCATION(result_name, lhs, rexpr, location, message) \ 49 | auto &&result_name = (rexpr); \ 50 | AIRPORT_RETURN_IF_LOCATION(error_prefix, !(result_name).ok(), (result_name).status(), location, message, ARROW_STRINGIFY(rexpr)); \ 51 | lhs = std::move(result_name).ValueUnsafe(); 52 | 53 | #define AIRPORT_ASSIGN_OR_RAISE_LOCATION_DESCRIPTOR(lhs, rexpr, location, descriptor, message) \ 54 | AIRPORT_ASSIGN_OR_RAISE_IMPL_LOCATION_DESCRIPTOR(ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), lhs, rexpr, location, descriptor, message); 55 | 56 | #define AIRPORT_ASSIGN_OR_RAISE_CONTAINER(lhs, rexpr, container, message) \ 57 | AIRPORT_ASSIGN_OR_RAISE_IMPL_LOCATION_DESCRIPTOR(ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), lhs, rexpr, (container)->server_location(), (container)->descriptor(), message); 58 | 59 | #define AIRPORT_ASSIGN_OR_RAISE_LOCATION(lhs, rexpr, location, message) \ 60 | AIRPORT_ASSIGN_OR_RAISE_IMPL_LOCATION(ARROW_ASSIGN_OR_RAISE_NAME(_error_or_value, __COUNTER__), lhs, rexpr, location, message); 61 | 62 | #define AIRPORT_MSGPACK_UNPACK(destination_type, destination_name, source, location, message) \ 63 | destination_type destination_name; \ 64 | try \ 65 | { \ 66 | msgpack::object_handle oh = msgpack::unpack( \ 67 | reinterpret_cast(source.data()), \ 68 | source.size(), \ 69 | 0); \ 70 | msgpack::object obj = oh.get(); \ 71 | obj.convert(destination_name); \ 72 | } \ 73 | catch (const std::exception &e) \ 74 | { \ 75 | throw AirportFlightException(location, string(message) + " " + string(e.what())); \ 76 | } 77 | 78 | #define AIRPORT_MSGPACK_UNPACK_CONTAINER(destination_type, destination_name, source, container, message) \ 79 | AIRPORT_MSGPACK_UNPACK(destination_type, destination_name, source, container->server_location(), message); 80 | 81 | #define AIRPORT_MSGPACK_ACTION_SINGLE_PARAMETER(result_name, action_name, parameters) \ 82 | msgpack::sbuffer parameters_packed_buffer; \ 83 | msgpack::pack(parameters_packed_buffer, parameters); \ 84 | arrow::flight::Action result_name{ \ 85 | action_name, \ 86 | std::make_shared( \ 87 | reinterpret_cast(parameters_packed_buffer.data()), \ 88 | parameters_packed_buffer.size())}; 89 | -------------------------------------------------------------------------------- /src/include/airport_optimizer.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "duckdb/main/config.hpp" 3 | 4 | namespace duckdb 5 | { 6 | class AirportOptimizer 7 | { 8 | public: 9 | static void Optimize(OptimizerExtensionInput &input, unique_ptr &plan); 10 | }; 11 | 12 | } 13 | -------------------------------------------------------------------------------- /src/include/airport_request_headers.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include 4 | #include 5 | #include "arrow/flight/client.h" 6 | #include "airport_flight_stream.hpp" 7 | 8 | namespace duckdb 9 | { 10 | std::string airport_user_agent(); 11 | void airport_add_standard_headers(arrow::flight::FlightCallOptions &options, const std::string &server_location); 12 | void airport_add_authorization_header(arrow::flight::FlightCallOptions &options, const std::string &auth_token); 13 | void airport_add_flight_path_header(arrow::flight::FlightCallOptions &options, 14 | const arrow::flight::FlightDescriptor &descriptor); 15 | void airport_add_trace_id_header(arrow::flight::FlightCallOptions &options, 16 | const string &trace_id); 17 | 18 | void airport_add_normal_headers(arrow::flight::FlightCallOptions &options, 19 | const AirportTakeFlightParameters ¶ms, 20 | const std::string &trace_id, 21 | std::optional descriptor = std::nullopt); 22 | 23 | // Generate a random id that is used for request tracking. 24 | string airport_trace_id(); 25 | 26 | } -------------------------------------------------------------------------------- /src/include/airport_scalar_function.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb.hpp" 4 | #include "duckdb/common/arrow/schema_metadata.hpp" 5 | #include "duckdb/function/table/arrow.hpp" 6 | #include "duckdb/parser/parsed_data/create_table_info.hpp" 7 | #include "duckdb/parser/parser.hpp" 8 | 9 | #include "airport_request_headers.hpp" 10 | #include "airport_macros.hpp" 11 | #include "airport_secrets.hpp" 12 | #include "arrow/util/key_value_metadata.h" 13 | 14 | namespace duckdb 15 | { 16 | 17 | class AirportScalarFunctionInfo : public ScalarFunctionInfo, public AirportLocationDescriptor 18 | { 19 | private: 20 | const string function_name_; 21 | const std::shared_ptr output_schema_; 22 | const std::shared_ptr input_schema_; 23 | Catalog &catalog_; 24 | 25 | public: 26 | AirportScalarFunctionInfo( 27 | const string &name, 28 | const AirportLocationDescriptor &location, 29 | const std::shared_ptr &output_schema, 30 | const std::shared_ptr &input_schema, 31 | Catalog &catalog) 32 | : ScalarFunctionInfo(), 33 | AirportLocationDescriptor(location), 34 | function_name_(name), 35 | output_schema_(output_schema), 36 | input_schema_(input_schema), 37 | catalog_(catalog) 38 | { 39 | } 40 | 41 | ~AirportScalarFunctionInfo() override 42 | { 43 | } 44 | 45 | Catalog &catalog() const 46 | { 47 | return catalog_; 48 | } 49 | 50 | const string &function_name() const 51 | { 52 | return function_name_; 53 | } 54 | 55 | const bool input_schema_includes_any_types() const 56 | { 57 | for (int i = 0; i < input_schema_->num_fields(); ++i) 58 | { 59 | const auto &field = input_schema_->field(i); 60 | auto field_metadata = field->metadata(); 61 | 62 | if (field_metadata != nullptr && field_metadata->Contains("is_any_type")) 63 | { 64 | return true; 65 | } 66 | } 67 | return false; 68 | } 69 | 70 | const std::shared_ptr &output_schema() const 71 | { 72 | return output_schema_; 73 | } 74 | 75 | const std::shared_ptr &input_schema() const 76 | { 77 | return input_schema_; 78 | } 79 | }; 80 | 81 | void AirportScalarFunctionProcessChunk(DataChunk &args, ExpressionState &state, Vector &result); 82 | unique_ptr AirportScalarFunctionInitLocalState(ExpressionState &state, const BoundFunctionExpression &expr, FunctionData *bind_data); 83 | 84 | unique_ptr AirportScalarFunctionBind(ClientContext &context, ScalarFunction &bound_function, 85 | vector> &arguments); 86 | } -------------------------------------------------------------------------------- /src/include/airport_schema_utils.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "airport_extension.hpp" 3 | #include "duckdb.hpp" 4 | 5 | namespace duckdb 6 | { 7 | 8 | void AirportExamineSchema( 9 | ClientContext &context, 10 | const ArrowSchemaWrapper &schema_root, 11 | ArrowTableType *arrow_table, 12 | vector *return_types, 13 | vector *names, 14 | vector *duckdb_type_names, 15 | idx_t *rowid_column_index, 16 | bool skip_rowid_column); 17 | 18 | bool AirportFieldMetadataIsRowId(const char *metadata); 19 | } -------------------------------------------------------------------------------- /src/include/airport_secrets.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #include "duckdb/main/secret/secret_manager.hpp" 3 | 4 | namespace duckdb 5 | { 6 | unique_ptr AirportGetSecretByName(ClientContext &context, const string &secret_name); 7 | 8 | SecretMatch AirportGetSecretByPath(ClientContext &context, const string &path); 9 | 10 | string AirportAuthTokenForLocation(ClientContext &context, const string &server_location, const string &secret_name, const string &auth_token); 11 | 12 | } 13 | -------------------------------------------------------------------------------- /src/include/airport_take_flight.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "airport_extension.hpp" 4 | #include "duckdb.hpp" 5 | 6 | #include "airport_flight_stream.hpp" 7 | 8 | namespace duckdb 9 | { 10 | struct AirportArrowScanGlobalState : public GlobalTableFunctionState 11 | { 12 | // idx_t batch_index = 0; 13 | 14 | idx_t MaxThreads() const override 15 | { 16 | return endpoints_.size() == 0 ? 1 : endpoints_.size(); 17 | } 18 | 19 | bool CanRemoveFilterColumns() const 20 | { 21 | return !projection_ids_.empty(); 22 | } 23 | 24 | // If there is a list of endpoints this constructor it used. 25 | AirportArrowScanGlobalState(const vector &endpoints, 26 | const vector &projection_ids, 27 | const vector &scanned_types, 28 | const std::optional &input) 29 | : endpoints_(endpoints), 30 | projection_ids_(projection_ids), 31 | scanned_types_(scanned_types), 32 | init_input_(input) 33 | { 34 | if (init_input_) 35 | { 36 | if (init_input_->filters) 37 | { 38 | init_input_->filters = init_input_->filters->Copy(); 39 | } 40 | } 41 | } 42 | 43 | // There are cases where a list of endpoints isn't available, for example 44 | // the calls to DoExchange, so in that case don't set the endpoints. 45 | explicit AirportArrowScanGlobalState() 46 | { 47 | } 48 | 49 | const size_t total_endpoints() const 50 | { 51 | return endpoints_.size(); 52 | } 53 | 54 | const std::optional GetNextEndpoint() 55 | { 56 | size_t index = current_endpoint_.fetch_add(1, std::memory_order_relaxed); 57 | if (index < endpoints_.size()) 58 | { 59 | return endpoints_[index]; 60 | } 61 | return std::nullopt; 62 | } 63 | 64 | const vector &projection_ids() const 65 | { 66 | return projection_ids_; 67 | } 68 | 69 | const vector &scanned_types() const 70 | { 71 | return scanned_types_; 72 | } 73 | 74 | const std::optional &init_input() const 75 | { 76 | return init_input_; 77 | } 78 | 79 | private: 80 | vector endpoints_; 81 | std::atomic current_endpoint_ = 0; 82 | const vector projection_ids_; 83 | const vector scanned_types_; 84 | std::optional init_input_ = std::nullopt; 85 | }; 86 | 87 | shared_ptr AirportProduceArrowScan( 88 | const ArrowScanFunctionData &function, 89 | const vector &column_ids, 90 | const TableFilterSet *filters, 91 | atomic *progress, 92 | std::shared_ptr *last_app_metadata, 93 | const std::shared_ptr &schema, 94 | const AirportLocationDescriptor &location_descriptor, 95 | AirportArrowScanLocalState &local_state); 96 | 97 | void AirportTakeFlight(ClientContext &context, TableFunctionInput &data_p, DataChunk &output); 98 | 99 | unique_ptr AirportArrowScanInitGlobal(ClientContext &context, 100 | TableFunctionInitInput &input); 101 | 102 | unique_ptr AirportTakeFlightBindWithFlightDescriptor( 103 | const AirportTakeFlightParameters &take_flight_params, 104 | const arrow::flight::FlightDescriptor &descriptor, 105 | ClientContext &context, 106 | const TableFunctionBindInput &input, 107 | vector &return_types, 108 | vector &names, 109 | std::shared_ptr schema, 110 | const std::optional &table_function_parameters, 111 | const AirportTableEntry *table_entry); 112 | 113 | std::string AirportNameForField(const string &name, const idx_t col_idx); 114 | 115 | void AirportTakeFlightComplexFilterPushdown(ClientContext &context, LogicalGet &get, FunctionData *bind_data_p, 116 | vector> &filters); 117 | unique_ptr AirportTakeFlightCardinality(ClientContext &context, const FunctionData *data); 118 | unique_ptr AirportTakeFlightStatistics(ClientContext &context, const FunctionData *bind_data, column_t column_index); 119 | double AirportTakeFlightScanProgress(ClientContext &, const FunctionData *data, const GlobalTableFunctionState *global_state); 120 | 121 | unique_ptr AirportArrowScanInitLocal(ExecutionContext &context, 122 | TableFunctionInitInput &input, 123 | GlobalTableFunctionState *global_state_p); 124 | 125 | } 126 | -------------------------------------------------------------------------------- /src/include/storage/airport_catalog.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/catalog/catalog.hpp" 4 | #include "duckdb/function/table_function.hpp" 5 | #include "duckdb/common/enums/access_mode.hpp" 6 | #include "storage/airport_schema_set.hpp" 7 | 8 | namespace duckdb 9 | { 10 | class AirportSchemaEntry; 11 | 12 | struct AirportAttachParameters 13 | { 14 | AirportAttachParameters(const string &location, const string &auth_token, const string &secret_name, const string &criteria) 15 | : location_(location), auth_token_(auth_token), secret_name_(secret_name), criteria_(criteria) 16 | { 17 | } 18 | 19 | const string &location() const 20 | { 21 | return location_; 22 | } 23 | 24 | const string &auth_token() const 25 | { 26 | return auth_token_; 27 | } 28 | 29 | const string &secret_name() const 30 | { 31 | return secret_name_; 32 | } 33 | 34 | const string &criteria() const 35 | { 36 | return criteria_; 37 | } 38 | 39 | private: 40 | // The location of the flight server. 41 | string location_; 42 | // The authorization token to use. 43 | string auth_token_; 44 | // The name of the secret to use 45 | string secret_name_; 46 | // The criteria to pass to the flight server when listing flights. 47 | string criteria_; 48 | }; 49 | 50 | class AirportClearCacheFunction : public TableFunction 51 | { 52 | public: 53 | AirportClearCacheFunction(); 54 | 55 | static void ClearCacheOnSetting(ClientContext &context, SetScope scope, Value ¶meter); 56 | }; 57 | 58 | class AirportCatalog : public Catalog 59 | { 60 | public: 61 | explicit AirportCatalog(AttachedDatabase &db_p, const string &internal_name, AccessMode access_mode, 62 | AirportAttachParameters attach_params); 63 | ~AirportCatalog() override; 64 | 65 | public: 66 | void Initialize(bool load_builtin) override; 67 | string GetCatalogType() override 68 | { 69 | return "airport"; 70 | } 71 | 72 | bool SupportsTimeTravel() const override 73 | { 74 | return true; 75 | } 76 | 77 | optional_ptr CreateSchema(CatalogTransaction transaction, CreateSchemaInfo &info) override; 78 | 79 | void ScanSchemas(ClientContext &context, std::function callback) override; 80 | 81 | optional_ptr LookupSchema(CatalogTransaction transaction, 82 | const EntryLookupInfo &schema_lookup, 83 | OnEntryNotFound if_not_found) override; 84 | 85 | PhysicalOperator &PlanCreateTableAs(ClientContext &context, PhysicalPlanGenerator &planner, 86 | LogicalCreateTable &op, PhysicalOperator &plan) override; 87 | PhysicalOperator &PlanInsert(ClientContext &context, PhysicalPlanGenerator &planner, LogicalInsert &op, 88 | optional_ptr plan) override; 89 | PhysicalOperator &PlanDelete(ClientContext &context, PhysicalPlanGenerator &planner, LogicalDelete &op, 90 | PhysicalOperator &plan) override; 91 | 92 | PhysicalOperator &PlanUpdate(ClientContext &context, PhysicalPlanGenerator &planner, LogicalUpdate &op, 93 | PhysicalOperator &plan) override; 94 | 95 | unique_ptr BindCreateIndex(Binder &binder, CreateStatement &stmt, TableCatalogEntry &table, 96 | unique_ptr plan) override; 97 | 98 | DatabaseSize GetDatabaseSize(ClientContext &context) override; 99 | 100 | //! Whether or not this is an in-memory database 101 | bool InMemory() override; 102 | string GetDBPath() override; 103 | 104 | void ClearCache(); 105 | 106 | optional_idx GetCatalogVersion(ClientContext &context) override; 107 | 108 | std::optional GetTransactionIdentifier(); 109 | 110 | // Track what version of the catalog has been loaded. 111 | std::optional loaded_catalog_version; 112 | 113 | const string &internal_name() const 114 | { 115 | return internal_name_; 116 | } 117 | 118 | const std::shared_ptr &attach_parameters() const 119 | { 120 | return attach_parameters_; 121 | } 122 | 123 | const AccessMode &access_mode() const 124 | { 125 | return access_mode_; 126 | } 127 | 128 | private: 129 | void DropSchema(ClientContext &context, DropInfo &info) override; 130 | 131 | private: 132 | std::shared_ptr flight_client_; 133 | AccessMode access_mode_; 134 | std::shared_ptr attach_parameters_; 135 | string internal_name_; 136 | AirportSchemaSet schemas; 137 | string default_schema; 138 | }; 139 | } 140 | -------------------------------------------------------------------------------- /src/include/storage/airport_catalog_set.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/transaction/transaction.hpp" 4 | #include "duckdb/common/case_insensitive_map.hpp" 5 | #include "duckdb/common/mutex.hpp" 6 | 7 | namespace duckdb 8 | { 9 | struct DropInfo; 10 | class AirportSchemaEntry; 11 | class AirportTransaction; 12 | 13 | class AirportCatalogSet 14 | { 15 | public: 16 | AirportCatalogSet(Catalog &catalog); 17 | 18 | virtual optional_ptr GetEntry(ClientContext &context, const EntryLookupInfo &lookup_info); 19 | virtual void DropEntry(ClientContext &context, DropInfo &info); 20 | void Scan(ClientContext &context, const std::function &callback); 21 | virtual optional_ptr CreateEntry(unique_ptr entry); 22 | void ClearEntries(); 23 | 24 | void ReplaceEntry( 25 | const string &name, 26 | unique_ptr entry); 27 | 28 | protected: 29 | virtual void LoadEntries(ClientContext &context) = 0; 30 | 31 | void EraseEntryInternal(const string &name); 32 | 33 | protected: 34 | Catalog &catalog; 35 | mutex entry_lock; 36 | case_insensitive_map_t> entries; 37 | bool is_loaded; 38 | }; 39 | 40 | class AirportInSchemaSet : public AirportCatalogSet 41 | { 42 | public: 43 | explicit AirportInSchemaSet(AirportSchemaEntry &schema); 44 | 45 | optional_ptr CreateEntry(unique_ptr entry) override; 46 | 47 | protected: 48 | AirportSchemaEntry &schema; 49 | }; 50 | 51 | } // namespace duckdb 52 | -------------------------------------------------------------------------------- /src/include/storage/airport_catalog_set_base.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "storage/airport_catalog_set.hpp" 4 | #include "storage/airport_table_entry.hpp" 5 | 6 | struct AirportFunctionCatalogSchemaNameKey 7 | { 8 | std::string catalog_name; 9 | std::string schema_name; 10 | std::string name; 11 | 12 | // Define equality operator to compare two keys 13 | bool operator==(const AirportFunctionCatalogSchemaNameKey &other) const 14 | { 15 | return catalog_name == other.catalog_name && schema_name == other.schema_name && name == other.name; 16 | } 17 | }; 18 | 19 | namespace std 20 | { 21 | template <> 22 | struct hash 23 | { 24 | size_t operator()(const AirportFunctionCatalogSchemaNameKey &k) const 25 | { 26 | // Combine the hash of all 3 strings 27 | return hash()(k.catalog_name) ^ (hash()(k.schema_name) << 1) ^ (hash()(k.name) << 2); 28 | } 29 | }; 30 | } 31 | 32 | namespace duckdb 33 | { 34 | struct CreateTableInfo; 35 | class AirportResult; 36 | class AirportSchemaEntry; 37 | class AirportCurlPool; 38 | struct AirportTableInfo; 39 | 40 | class AirportCatalogSetBase : public AirportInSchemaSet 41 | { 42 | protected: 43 | AirportCurlPool &connection_pool_; 44 | string cache_directory_; 45 | 46 | public: 47 | explicit AirportCatalogSetBase(AirportCurlPool &connection_pool, AirportSchemaEntry &schema, const string &cache_directory) 48 | : AirportInSchemaSet(schema), connection_pool_(connection_pool), cache_directory_(cache_directory) 49 | { 50 | } 51 | }; 52 | 53 | } 54 | -------------------------------------------------------------------------------- /src/include/storage/airport_curl_pool.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | #ifndef CURLPOOL_H 3 | #define CURLPOOL_H 4 | 5 | #include 6 | #include 7 | #include 8 | #include 9 | 10 | namespace duckdb 11 | { 12 | class AirportCurlPool 13 | { 14 | public: 15 | // Constructor: Initializes a pool with the given number of CURL handles 16 | explicit AirportCurlPool(size_t size); 17 | 18 | // Destructor: Cleans up all CURL handles in the pool 19 | ~AirportCurlPool(); 20 | 21 | // Acquire a CURL handle from the pool 22 | CURL *acquire(); 23 | 24 | // Release a CURL handle back to the pool 25 | void release(CURL *handle); 26 | 27 | private: 28 | std::vector _pool; // Pool of CURL handles 29 | std::mutex _mutex; // Mutex for synchronizing access to the pool 30 | std::condition_variable _cv; // Condition variable for signaling availability of handles 31 | }; 32 | 33 | } 34 | 35 | #endif // CURLPOOL_H 36 | -------------------------------------------------------------------------------- /src/include/storage/airport_delete.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/execution/physical_operator.hpp" 4 | 5 | namespace duckdb 6 | { 7 | class AirportDelete : public PhysicalOperator 8 | { 9 | public: 10 | AirportDelete(LogicalOperator &op, TableCatalogEntry &table, idx_t rowid_index, bool return_chunk); 11 | 12 | //! The table to delete from 13 | TableCatalogEntry &table; 14 | idx_t rowid_index; 15 | bool return_chunk; 16 | 17 | public: 18 | // Source interface 19 | SourceResultType GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; 20 | 21 | bool IsSource() const override 22 | { 23 | return true; 24 | } 25 | 26 | public: 27 | // Sink interface 28 | unique_ptr GetGlobalSinkState(ClientContext &context) const override; 29 | unique_ptr GetLocalSinkState(ExecutionContext &context) const override; 30 | 31 | unique_ptr GetGlobalSourceState(ClientContext &context) const override; 32 | 33 | SinkResultType Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const override; 34 | SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 35 | OperatorSinkFinalizeInput &input) const override; 36 | 37 | bool IsSink() const override 38 | { 39 | return true; 40 | } 41 | 42 | bool ParallelSink() const override 43 | { 44 | return false; 45 | } 46 | 47 | string GetName() const override; 48 | InsertionOrderPreservingMap ParamsToString() const override; 49 | }; 50 | 51 | } // namespace duckdb 52 | -------------------------------------------------------------------------------- /src/include/storage/airport_delete_parameterized.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/execution/physical_operator.hpp" 4 | 5 | namespace duckdb 6 | { 7 | class AirportDeleteParameterized : public PhysicalOperator 8 | { 9 | public: 10 | AirportDeleteParameterized(LogicalOperator &op, TableCatalogEntry &table, PhysicalOperator &plan); 11 | 12 | //! The table to delete from 13 | TableCatalogEntry &table; 14 | string sql_filters; 15 | 16 | public: 17 | bool IsSource() const override 18 | { 19 | return true; 20 | } 21 | 22 | // Source interface 23 | SourceResultType GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; 24 | 25 | public: 26 | // Sink interface 27 | unique_ptr GetGlobalSinkState(ClientContext &context) const override; 28 | SinkResultType Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const override; 29 | SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 30 | OperatorSinkFinalizeInput &input) const override; 31 | 32 | bool IsSink() const override 33 | { 34 | return true; 35 | } 36 | 37 | bool ParallelSink() const override 38 | { 39 | return false; 40 | } 41 | 42 | string GetName() const override; 43 | InsertionOrderPreservingMap ParamsToString() const override; 44 | }; 45 | 46 | } // namespace duckdb 47 | -------------------------------------------------------------------------------- /src/include/storage/airport_exchange.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb.hpp" 4 | #include "duckdb/function/table/arrow.hpp" 5 | #include "airport_flight_stream.hpp" 6 | #include "airport_take_flight.hpp" 7 | #include "airport_table_entry.hpp" 8 | #include "airport_schema_utils.hpp" 9 | 10 | namespace duckdb 11 | { 12 | 13 | struct AirportExchangeTakeFlightBindData : public AirportTakeFlightBindData 14 | { 15 | 16 | public: 17 | explicit AirportExchangeTakeFlightBindData( 18 | stream_factory_produce_t scanner_producer_p, 19 | const string &trace_id, 20 | const int64_t estimated_records, 21 | const AirportTakeFlightParameters &take_flight_params_p, 22 | const std::optional &table_function_parameters_p, 23 | std::shared_ptr schema, 24 | const flight::FlightDescriptor &descriptor, AirportTableEntry *table_entry, 25 | shared_ptr dependency = nullptr) 26 | : AirportTakeFlightBindData( 27 | scanner_producer_p, 28 | trace_id, 29 | estimated_records, 30 | take_flight_params_p, 31 | table_function_parameters_p, 32 | schema, 33 | descriptor, 34 | table_entry, 35 | std::move(dependency)) 36 | { 37 | } 38 | 39 | mutable mutex lock; 40 | 41 | void examine_schema( 42 | ClientContext &context, 43 | bool skip_rowid_column) 44 | { 45 | AirportExamineSchema( 46 | context, 47 | schema_root, 48 | &arrow_table, 49 | &return_types_, 50 | &names_, 51 | nullptr, 52 | &rowid_column_index, 53 | skip_rowid_column); 54 | } 55 | 56 | const vector &names() const 57 | { 58 | return names_; 59 | } 60 | 61 | const vector &return_types() const 62 | { 63 | return return_types_; 64 | } 65 | 66 | private: 67 | // these are set in examine_schema(). 68 | vector names_; 69 | vector return_types_; 70 | }; 71 | 72 | // This is all of the state is needed to perform a ArrowScan on a resulting 73 | // DoExchange stream, this is useful for having RETURNING data work for 74 | // INSERT, DELETE or UPDATE. 75 | class AirportExchangeGlobalState 76 | { 77 | public: 78 | std::shared_ptr send_schema; 79 | 80 | std::unique_ptr scan_bind_data; 81 | std::unique_ptr reader; 82 | std::unique_ptr writer; 83 | 84 | duckdb::unique_ptr scan_table_function_input; 85 | 86 | duckdb::unique_ptr scan_global_state; 87 | duckdb::unique_ptr scan_local_state; 88 | 89 | vector send_types; 90 | vector send_names; 91 | }; 92 | 93 | void AirportExchangeGetGlobalSinkState(ClientContext &context, 94 | const TableCatalogEntry &table, 95 | const AirportTableEntry &airport_table, 96 | AirportExchangeGlobalState *global_state, 97 | const ArrowSchema &send_schema, 98 | const bool return_chunk, 99 | const string exchange_operation, 100 | const vector returning_column_names, 101 | const std::optional transaction_id); 102 | } -------------------------------------------------------------------------------- /src/include/storage/airport_insert.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/execution/physical_operator.hpp" 4 | #include "duckdb/common/index_vector.hpp" 5 | #include "duckdb/parser/statement/insert_statement.hpp" 6 | 7 | namespace duckdb 8 | { 9 | 10 | class AirportInsertGlobalState; 11 | class AirportInsertLocalState; 12 | class AirportInsert : public PhysicalOperator 13 | { 14 | public: 15 | //! INSERT INTO 16 | AirportInsert(LogicalOperator &op, 17 | TableCatalogEntry &table, 18 | physical_index_vector_t column_index_map, 19 | bool return_chunk, 20 | vector> bound_defaults, 21 | vector> bound_constraints); 22 | 23 | //! CREATE TABLE AS 24 | AirportInsert(LogicalOperator &op, SchemaCatalogEntry &schema, unique_ptr info, idx_t estimated_cardinality); 25 | 26 | //! The table to insert into 27 | optional_ptr insert_table; 28 | // optional_ptr table; 29 | 30 | //! The insert types 31 | vector insert_types; 32 | 33 | //! Table schema, in case of CREATE TABLE AS 34 | optional_ptr schema; 35 | //! Create table info, in case of CREATE TABLE AS 36 | unique_ptr info; 37 | //! column_index_map 38 | physical_index_vector_t column_index_map; 39 | 40 | bool return_chunk; 41 | 42 | //! The default expressions of the columns for which no value is provided 43 | vector> bound_defaults; 44 | //! The bound constraints for the table 45 | vector> bound_constraints; 46 | 47 | // For now always just throw errors. 48 | OnConflictAction action_type = OnConflictAction::THROW; 49 | 50 | public: 51 | // Source interface 52 | SourceResultType GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; 53 | 54 | bool IsSource() const override 55 | { 56 | return true; 57 | } 58 | 59 | protected: 60 | idx_t OnConflictHandling(TableCatalogEntry &table, 61 | ExecutionContext &context, 62 | AirportInsertGlobalState &gstate, 63 | AirportInsertLocalState &lstate, 64 | DataChunk &chunk) const; 65 | 66 | private: 67 | static void ResolveDefaults(const TableCatalogEntry &table, DataChunk &chunk, 68 | const physical_index_vector_t &column_index_map, 69 | ExpressionExecutor &default_executor, DataChunk &result); 70 | 71 | public: 72 | // Sink interface 73 | unique_ptr GetGlobalSinkState(ClientContext &context) const override; 74 | unique_ptr GetLocalSinkState(ExecutionContext &context) const override; 75 | 76 | unique_ptr GetGlobalSourceState(ClientContext &context) const override; 77 | 78 | SinkResultType Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const override; 79 | SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 80 | OperatorSinkFinalizeInput &input) const override; 81 | 82 | bool IsSink() const override 83 | { 84 | return true; 85 | } 86 | 87 | bool ParallelSink() const override 88 | { 89 | return false; 90 | } 91 | 92 | string GetName() const override; 93 | InsertionOrderPreservingMap ParamsToString() const override; 94 | }; 95 | 96 | } 97 | -------------------------------------------------------------------------------- /src/include/storage/airport_scalar_function_set.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "storage/airport_catalog_set.hpp" 4 | #include "storage/airport_table_entry.hpp" 5 | #include "storage/airport_catalog_set_base.hpp" 6 | 7 | namespace duckdb 8 | { 9 | class AirportScalarFunctionSet : public AirportCatalogSetBase 10 | { 11 | 12 | protected: 13 | void LoadEntries(ClientContext &context) override; 14 | 15 | public: 16 | explicit AirportScalarFunctionSet(AirportCurlPool &connection_pool, AirportSchemaEntry &schema, const string &cache_directory) : AirportCatalogSetBase(connection_pool, schema, cache_directory) 17 | { 18 | } 19 | ~AirportScalarFunctionSet() {} 20 | }; 21 | 22 | } -------------------------------------------------------------------------------- /src/include/storage/airport_schema_entry.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/catalog/catalog_entry/schema_catalog_entry.hpp" 4 | #include "storage/airport_table_set.hpp" 5 | #include "storage/airport_curl_pool.hpp" 6 | #include "storage/airport_scalar_function_set.hpp" 7 | #include "storage/airport_table_function_set.hpp" 8 | 9 | namespace duckdb 10 | { 11 | class AirportTransaction; 12 | 13 | class AirportSchemaEntry : public SchemaCatalogEntry 14 | { 15 | public: 16 | AirportSchemaEntry(Catalog &catalog, 17 | CreateSchemaInfo &info, 18 | AirportCurlPool &connection_pool, 19 | const string &cache_directory, 20 | const AirportAPISchema &schema_data); 21 | ~AirportSchemaEntry() override; 22 | 23 | public: 24 | optional_ptr CreateTable(CatalogTransaction transaction, BoundCreateTableInfo &info) override; 25 | optional_ptr CreateFunction(CatalogTransaction transaction, CreateFunctionInfo &info) override; 26 | optional_ptr CreateIndex(CatalogTransaction transaction, CreateIndexInfo &info, 27 | TableCatalogEntry &table) override; 28 | optional_ptr CreateView(CatalogTransaction transaction, CreateViewInfo &info) override; 29 | optional_ptr CreateSequence(CatalogTransaction transaction, CreateSequenceInfo &info) override; 30 | optional_ptr CreateTableFunction(CatalogTransaction transaction, 31 | CreateTableFunctionInfo &info) override; 32 | optional_ptr CreateCopyFunction(CatalogTransaction transaction, 33 | CreateCopyFunctionInfo &info) override; 34 | optional_ptr CreatePragmaFunction(CatalogTransaction transaction, 35 | CreatePragmaFunctionInfo &info) override; 36 | optional_ptr CreateCollation(CatalogTransaction transaction, CreateCollationInfo &info) override; 37 | optional_ptr CreateType(CatalogTransaction transaction, CreateTypeInfo &info) override; 38 | void Alter(CatalogTransaction transaction, AlterInfo &info) override; 39 | void Scan(ClientContext &context, CatalogType type, const std::function &callback) override; 40 | void Scan(CatalogType type, const std::function &callback) override; 41 | void DropEntry(ClientContext &context, DropInfo &info) override; 42 | 43 | optional_ptr LookupEntry(CatalogTransaction transaction, const EntryLookupInfo &lookup_info) override; 44 | 45 | const AirportSerializedContentsWithSHA256Hash &serialized_source() const 46 | { 47 | return schema_data_.source(); 48 | } 49 | 50 | private: 51 | AirportAPISchema schema_data_; 52 | 53 | AirportCatalogSet &GetCatalogSet(CatalogType type); 54 | AirportTableSet tables; 55 | AirportScalarFunctionSet scalar_functions; 56 | AirportTableFunctionSet table_functions; 57 | }; 58 | 59 | } // namespace duckdb 60 | -------------------------------------------------------------------------------- /src/include/storage/airport_schema_set.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "airport_catalog_set.hpp" 4 | #include "airport_curl_pool.hpp" 5 | #include "airport_catalog_api.hpp" 6 | 7 | namespace duckdb 8 | { 9 | struct CreateSchemaInfo; 10 | 11 | class AirportSchemaSet : public AirportCatalogSet 12 | { 13 | public: 14 | explicit AirportSchemaSet(Catalog &catalog); 15 | 16 | public: 17 | optional_ptr CreateSchema(ClientContext &context, CreateSchemaInfo &info); 18 | 19 | // Load the schemas of the entire set from a cached url if possible, useful for scans 20 | // when all schemas are requested. 21 | void LoadEntireSet(ClientContext &context); 22 | 23 | protected: 24 | void LoadEntries(ClientContext &context) override; 25 | 26 | private: 27 | AirportCurlPool connection_pool; 28 | 29 | unique_ptr collection; 30 | 31 | bool populated_entire_set = false; 32 | bool called_load_entries = false; 33 | mutex entry_lock; 34 | }; 35 | 36 | } // namespace duckdb 37 | -------------------------------------------------------------------------------- /src/include/storage/airport_table_entry.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "airport_catalog_api.hpp" 4 | #include "duckdb/catalog/catalog_entry/table_catalog_entry.hpp" 5 | #include "duckdb/parser/parsed_data/create_table_info.hpp" 6 | #include "storage/airport_table_set.hpp" 7 | 8 | namespace duckdb 9 | { 10 | class AirportTableSet; 11 | 12 | struct AirportTableInfo 13 | { 14 | AirportTableInfo() : create_info(make_uniq()) 15 | { 16 | } 17 | 18 | AirportTableInfo(const string &schema, const string &table) 19 | : create_info(make_uniq(string(), schema, table)) 20 | { 21 | } 22 | 23 | AirportTableInfo(const SchemaCatalogEntry &schema, const string &table) 24 | : create_info(make_uniq((SchemaCatalogEntry &)schema, table)) 25 | { 26 | } 27 | 28 | const string &GetTableName() const 29 | { 30 | return create_info->table; 31 | } 32 | 33 | const unique_ptr create_info; 34 | }; 35 | 36 | class AirportTableEntry : public TableCatalogEntry 37 | { 38 | public: 39 | AirportTableEntry(Catalog &catalog, SchemaCatalogEntry &schema, CreateTableInfo &info, const LogicalType &rowid_type); 40 | AirportTableEntry(Catalog &catalog, SchemaCatalogEntry &schema, AirportTableInfo &info, const LogicalType &rowid_type); 41 | 42 | unique_ptr table_data; 43 | 44 | virtual_column_map_t GetVirtualColumns() const override 45 | { 46 | virtual_column_map_t virtual_columns; 47 | if (rowid_type.id() != LogicalTypeId::SQLNULL) 48 | { 49 | virtual_columns.insert(make_pair(COLUMN_IDENTIFIER_ROW_ID, TableColumn("rowid", rowid_type))); 50 | } 51 | // virtual_columns.insert(make_pair(COLUMN_IDENTIFIER_EMPTY, TableColumn("", LogicalType::BOOLEAN))); 52 | 53 | return virtual_columns; 54 | } 55 | 56 | const LogicalType &GetRowIdType() const 57 | { 58 | return rowid_type; 59 | } 60 | 61 | public: 62 | unique_ptr GetStatistics(ClientContext &context, column_t column_id) override; 63 | 64 | TableFunction GetScanFunction(ClientContext &context, unique_ptr &bind_data) override; 65 | 66 | TableFunction GetScanFunction(ClientContext &context, unique_ptr &bind_data, const EntryLookupInfo &lookup) override; 67 | 68 | TableStorageInfo GetStorageInfo(ClientContext &context) override; 69 | 70 | unique_ptr AlterEntryDirect(ClientContext &context, AlterInfo &info); 71 | // void BindUpdateConstraints(Binder &binder, LogicalGet &get, LogicalProjection &proj, LogicalUpdate &update, 72 | // ClientContext &context) override; 73 | 74 | Catalog &GetCatalog() const 75 | { 76 | return catalog_; 77 | } 78 | 79 | private: 80 | //! A logical type for the rowid of this table. 81 | const LogicalType rowid_type; 82 | 83 | Catalog &catalog_; 84 | }; 85 | 86 | } // namespace duckdb 87 | -------------------------------------------------------------------------------- /src/include/storage/airport_table_function_set.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "storage/airport_catalog_set.hpp" 4 | #include "storage/airport_catalog_set_base.hpp" 5 | 6 | namespace duckdb 7 | { 8 | class AirportTableFunctionSet : public AirportCatalogSetBase 9 | { 10 | 11 | protected: 12 | void LoadEntries(ClientContext &context) override; 13 | 14 | public: 15 | explicit AirportTableFunctionSet(AirportCurlPool &connection_pool, AirportSchemaEntry &schema, const string &cache_directory) : AirportCatalogSetBase(connection_pool, schema, cache_directory) 16 | { 17 | } 18 | ~AirportTableFunctionSet() {} 19 | }; 20 | 21 | } // namespace duckdb 22 | -------------------------------------------------------------------------------- /src/include/storage/airport_table_set.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "storage/airport_catalog_set.hpp" 4 | #include "storage/airport_table_entry.hpp" 5 | 6 | #include "storage/airport_catalog_set_base.hpp" 7 | 8 | namespace duckdb 9 | { 10 | struct CreateTableInfo; 11 | class AirportResult; 12 | class AirportSchemaEntry; 13 | class AirportCurlPool; 14 | struct AirportTableInfo; 15 | 16 | class AirportTableSet : public AirportCatalogSetBase 17 | { 18 | public: 19 | explicit AirportTableSet(AirportCurlPool &connection_pool, AirportSchemaEntry &schema, const string &cache_directory) : AirportCatalogSetBase(connection_pool, schema, cache_directory) 20 | { 21 | } 22 | ~AirportTableSet() {} 23 | 24 | public: 25 | optional_ptr GetEntry(ClientContext &context, const EntryLookupInfo &lookup_info) override; 26 | 27 | optional_ptr CreateTable(ClientContext &context, BoundCreateTableInfo &info); 28 | 29 | static unique_ptr GetTableInfo(ClientContext &context, AirportSchemaEntry &schema, 30 | const string &table_name); 31 | optional_ptr RefreshTable(ClientContext &context, const string &table_name); 32 | 33 | void AlterTable(ClientContext &context, AlterTableInfo &info); 34 | 35 | protected: 36 | void LoadEntries(ClientContext &context) override; 37 | }; 38 | 39 | class AirportTableEntry; 40 | 41 | unique_ptr AirportCatalogEntryFromFlightInfo( 42 | std::unique_ptr flight_info, 43 | const std::string &server_location, 44 | SchemaCatalogEntry &schema_entry, 45 | Catalog &catalog, 46 | ClientContext &context); 47 | 48 | } // namespace duckdb 49 | -------------------------------------------------------------------------------- /src/include/storage/airport_transaction.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "airport_extension.hpp" 4 | #include "airport_catalog.hpp" 5 | #include "duckdb/transaction/transaction.hpp" 6 | 7 | namespace duckdb 8 | { 9 | class AirportCatalog; 10 | class AirportSchemaEntry; 11 | class AirportTableEntry; 12 | 13 | enum class AirportTransactionState 14 | { 15 | TRANSACTION_NOT_YET_STARTED, 16 | TRANSACTION_STARTED, 17 | TRANSACTION_FINISHED 18 | }; 19 | 20 | class AirportTransaction : public Transaction 21 | { 22 | public: 23 | AirportTransaction(AirportCatalog &airport_catalog, TransactionManager &manager, ClientContext &context); 24 | ~AirportTransaction() override; 25 | 26 | void Start(); 27 | void Commit(); 28 | void Rollback(); 29 | 30 | // UCConnection &GetConnection(); 31 | // unique_ptr Query(const string &query); 32 | static AirportTransaction &Get(ClientContext &context, Catalog &catalog); 33 | AccessMode GetAccessMode() const 34 | { 35 | return access_mode_; 36 | } 37 | 38 | // The identifier returned from the Arrow flight server. 39 | const std::optional &identifier() const 40 | { 41 | return identifier_; 42 | } 43 | 44 | vector> point_in_time_entries_; 45 | 46 | private: 47 | // The identifier returned from the Arrow flight server. 48 | std::optional identifier_; 49 | 50 | std::optional GetTransactionIdentifier(); 51 | 52 | AirportTransactionState transaction_state = AirportTransactionState::TRANSACTION_NOT_YET_STARTED; 53 | AccessMode access_mode_; 54 | 55 | // The name of the catalog where this transaction is running. 56 | std::string catalog_name; 57 | // Copied from the airport catalog, since it can't keep a reference. 58 | std::shared_ptr attach_parameters; 59 | }; 60 | 61 | } // namespace duckdb 62 | -------------------------------------------------------------------------------- /src/include/storage/airport_transaction_manager.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/transaction/transaction_manager.hpp" 4 | #include "storage/airport_catalog.hpp" 5 | #include "storage/airport_transaction.hpp" 6 | 7 | namespace duckdb 8 | { 9 | 10 | class AirportTransactionManager : public TransactionManager 11 | { 12 | public: 13 | AirportTransactionManager(AttachedDatabase &db_p, AirportCatalog &airport_catalog); 14 | 15 | Transaction &StartTransaction(ClientContext &context) override; 16 | ErrorData CommitTransaction(ClientContext &context, Transaction &transaction) override; 17 | void RollbackTransaction(Transaction &transaction) override; 18 | 19 | void Checkpoint(ClientContext &context, bool force = false) override; 20 | 21 | private: 22 | AirportCatalog &airport_catalog; 23 | mutex transaction_lock; 24 | reference_map_t> transactions; 25 | }; 26 | 27 | } // namespace duckdb 28 | -------------------------------------------------------------------------------- /src/include/storage/airport_update.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/execution/physical_operator.hpp" 4 | #include "duckdb/common/index_vector.hpp" 5 | 6 | namespace duckdb 7 | { 8 | 9 | class AirportUpdate : public PhysicalOperator 10 | { 11 | public: 12 | AirportUpdate( 13 | LogicalOperator &op, 14 | vector types, 15 | TableCatalogEntry &table, 16 | vector columns, vector> expressions, 17 | vector> bound_defaults, vector> bound_constraints, 18 | idx_t estimated_cardinality, bool return_chunk, 19 | bool update_is_del_and_insert); 20 | 21 | // LogicalOperator &op, TableCatalogEntry &table, vector columns, bool return_chunk); 22 | 23 | //! The table to delete from 24 | TableCatalogEntry &table; 25 | //! The set of columns to update 26 | vector columns; 27 | 28 | vector> expressions; 29 | vector> bound_defaults; 30 | vector> bound_constraints; 31 | bool update_is_del_and_insert; 32 | //! If the returning statement is present, return the whole chunk 33 | bool return_chunk; 34 | 35 | public: 36 | // Source interface 37 | SourceResultType GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; 38 | 39 | bool IsSource() const override 40 | { 41 | return true; 42 | } 43 | 44 | public: 45 | // Sink interface 46 | unique_ptr GetGlobalSinkState(ClientContext &context) const override; 47 | unique_ptr GetLocalSinkState(ExecutionContext &context) const override; 48 | 49 | unique_ptr GetGlobalSourceState(ClientContext &context) const override; 50 | 51 | SinkResultType Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const override; 52 | SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 53 | OperatorSinkFinalizeInput &input) const override; 54 | 55 | bool IsSink() const override 56 | { 57 | return true; 58 | } 59 | 60 | bool ParallelSink() const override 61 | { 62 | return false; 63 | } 64 | 65 | string GetName() const override; 66 | InsertionOrderPreservingMap ParamsToString() const override; 67 | 68 | private: 69 | vector send_types; 70 | // This is a list of all column names that will be send 71 | vector send_names; 72 | }; 73 | } 74 | -------------------------------------------------------------------------------- /src/include/storage/airport_update_parameterized.hpp: -------------------------------------------------------------------------------- 1 | #pragma once 2 | 3 | #include "duckdb/execution/physical_operator.hpp" 4 | 5 | namespace duckdb 6 | { 7 | class AirportUpdateParameterized : public PhysicalOperator 8 | { 9 | public: 10 | AirportUpdateParameterized(LogicalOperator &op, TableCatalogEntry &table, PhysicalOperator &plan); 11 | 12 | //! The table to delete from 13 | TableCatalogEntry &table; 14 | string sql_filters; 15 | 16 | public: 17 | bool IsSource() const override 18 | { 19 | return true; 20 | } 21 | 22 | // Source interface 23 | SourceResultType GetData(ExecutionContext &context, DataChunk &chunk, OperatorSourceInput &input) const override; 24 | 25 | public: 26 | // Sink interface 27 | unique_ptr GetGlobalSinkState(ClientContext &context) const override; 28 | SinkResultType Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const override; 29 | SinkFinalizeType Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 30 | OperatorSinkFinalizeInput &input) const override; 31 | 32 | bool IsSink() const override 33 | { 34 | return true; 35 | } 36 | 37 | bool ParallelSink() const override 38 | { 39 | return false; 40 | } 41 | 42 | string GetName() const override; 43 | InsertionOrderPreservingMap ParamsToString() const override; 44 | }; 45 | 46 | } // namespace duckdb 47 | -------------------------------------------------------------------------------- /src/storage/airport_catalog.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "storage/airport_catalog.hpp" 3 | #include "storage/airport_schema_entry.hpp" 4 | #include "storage/airport_transaction.hpp" 5 | #include "duckdb/storage/database_size.hpp" 6 | #include "duckdb/parser/parsed_data/drop_info.hpp" 7 | #include "duckdb/parser/parsed_data/create_schema_info.hpp" 8 | #include "duckdb/main/attached_database.hpp" 9 | #include "storage/airport_delete.hpp" 10 | #include "storage/airport_insert.hpp" 11 | #include "duckdb/planner/operator/logical_create_table.hpp" 12 | #include "duckdb/execution/physical_operator.hpp" 13 | #include 14 | #include 15 | #include 16 | #include "airport_macros.hpp" 17 | 18 | #include "airport_request_headers.hpp" 19 | 20 | namespace duckdb 21 | { 22 | 23 | AirportCatalog::AirportCatalog(AttachedDatabase &db_p, const string &internal_name, AccessMode access_mode, 24 | AirportAttachParameters attach_params) 25 | : Catalog(db_p), access_mode_(access_mode), attach_parameters_(std::make_shared(std::move(attach_params))), 26 | internal_name_(internal_name), schemas(*this) 27 | { 28 | flight_client_ = AirportAPI::FlightClientForLocation(this->attach_parameters_->location()); 29 | } 30 | 31 | AirportCatalog::~AirportCatalog() = default; 32 | 33 | void AirportCatalog::Initialize(bool load_builtin) 34 | { 35 | } 36 | 37 | struct AirportGetCatalogVersionParams 38 | { 39 | string catalog_name; 40 | MSGPACK_DEFINE_MAP(catalog_name); 41 | }; 42 | 43 | optional_idx AirportCatalog::GetCatalogVersion(ClientContext &context) 44 | { 45 | if (loaded_catalog_version.has_value() && loaded_catalog_version.value().is_fixed) 46 | { 47 | return loaded_catalog_version.value().catalog_version; 48 | } 49 | 50 | arrow::flight::FlightCallOptions call_options; 51 | airport_add_standard_headers(call_options, attach_parameters_->location()); 52 | airport_add_authorization_header(call_options, attach_parameters_->auth_token()); 53 | 54 | // Might want to cache this though if a server declares the server catalog will not change. 55 | 56 | auto &server_location = attach_parameters_->location(); 57 | 58 | AirportGetCatalogVersionParams params; 59 | params.catalog_name = internal_name_; 60 | 61 | AIRPORT_MSGPACK_ACTION_SINGLE_PARAMETER(action, "catalog_version", params); 62 | 63 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto action_results, 64 | flight_client_->DoAction(call_options, action), 65 | server_location, 66 | "calling catalog_version action"); 67 | 68 | // The only item returned is a serialized flight info. 69 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto serialized_catalog_version_buffer, 70 | action_results->Next(), 71 | server_location, 72 | "reading catalog_version action result"); 73 | 74 | AIRPORT_MSGPACK_UNPACK(AirportGetCatalogVersionResult, result, 75 | (*serialized_catalog_version_buffer->body), 76 | server_location, 77 | "File to parse msgpack encoded catalog_version response"); 78 | 79 | loaded_catalog_version = result; 80 | 81 | return result.catalog_version; 82 | } 83 | 84 | optional_ptr AirportCatalog::CreateSchema(CatalogTransaction transaction, CreateSchemaInfo &info) 85 | { 86 | if (info.on_conflict == OnCreateConflict::REPLACE_ON_CONFLICT) 87 | { 88 | DropInfo try_drop; 89 | try_drop.type = CatalogType::SCHEMA_ENTRY; 90 | try_drop.name = info.schema; 91 | try_drop.if_not_found = OnEntryNotFound::RETURN_NULL; 92 | try_drop.cascade = false; 93 | schemas.DropEntry(transaction.GetContext(), try_drop); 94 | } 95 | return schemas.CreateSchema(transaction.GetContext(), info); 96 | } 97 | 98 | void AirportCatalog::DropSchema(ClientContext &context, DropInfo &info) 99 | { 100 | return schemas.DropEntry(context, info); 101 | } 102 | 103 | void AirportCatalog::ScanSchemas(ClientContext &context, std::function callback) 104 | { 105 | // If there is a contents_url for all schemas make sure it is present and decompressed on the disk, so that the 106 | // schema loaders will grab it. 107 | 108 | schemas.LoadEntireSet(context); 109 | 110 | schemas.Scan(context, [&](CatalogEntry &schema) 111 | { callback(schema.Cast()); }); 112 | } 113 | 114 | optional_ptr AirportCatalog::LookupSchema(CatalogTransaction transaction, 115 | const EntryLookupInfo &schema_lookup, 116 | OnEntryNotFound if_not_found) 117 | { 118 | auto &schema_name = schema_lookup.GetEntryName(); 119 | if (schema_name == DEFAULT_SCHEMA) 120 | { 121 | if (if_not_found == OnEntryNotFound::RETURN_NULL) 122 | { 123 | // There really isn't a default way to handle this, so just return null. 124 | return nullptr; 125 | } 126 | throw CatalogException(schema_lookup.GetErrorContext(), "Schema with name \"%s\" not found", schema_name); 127 | } 128 | auto entry = schemas.GetEntry(transaction.GetContext(), schema_lookup); 129 | if (!entry && if_not_found != OnEntryNotFound::RETURN_NULL) 130 | { 131 | throw CatalogException(schema_lookup.GetErrorContext(), "Schema with name \"%s\" not found", schema_name); 132 | } 133 | return reinterpret_cast(entry.get()); 134 | } 135 | 136 | bool AirportCatalog::InMemory() 137 | { 138 | return false; 139 | } 140 | 141 | string AirportCatalog::GetDBPath() 142 | { 143 | return internal_name_; 144 | } 145 | 146 | DatabaseSize AirportCatalog::GetDatabaseSize(ClientContext &context) 147 | { 148 | DatabaseSize size; 149 | return size; 150 | } 151 | 152 | void AirportCatalog::ClearCache() 153 | { 154 | schemas.ClearEntries(); 155 | } 156 | 157 | PhysicalOperator &AirportCatalog::PlanCreateTableAs(ClientContext &context, 158 | PhysicalPlanGenerator &planner, 159 | LogicalCreateTable &op, 160 | PhysicalOperator &plan) 161 | { 162 | auto &insert = planner.Make(op, op.schema, std::move(op.info), false); 163 | insert.children.push_back(plan); 164 | return insert; 165 | } 166 | 167 | unique_ptr AirportCatalog::BindCreateIndex(Binder &binder, CreateStatement &stmt, TableCatalogEntry &table, 168 | unique_ptr plan) 169 | { 170 | throw NotImplementedException("AirportCatalog BindCreateIndex"); 171 | } 172 | 173 | } // namespace duckdb 174 | -------------------------------------------------------------------------------- /src/storage/airport_catalog_set.cpp: -------------------------------------------------------------------------------- 1 | #include "storage/airport_catalog_set.hpp" 2 | #include "storage/airport_transaction.hpp" 3 | #include "duckdb/parser/parsed_data/drop_info.hpp" 4 | #include "storage/airport_schema_entry.hpp" 5 | #include "airport_request_headers.hpp" 6 | #include "storage/airport_catalog.hpp" 7 | #include "airport_macros.hpp" 8 | #include 9 | #include 10 | 11 | namespace duckdb 12 | { 13 | 14 | AirportCatalogSet::AirportCatalogSet(Catalog &catalog) : catalog(catalog), is_loaded(false) 15 | { 16 | } 17 | 18 | optional_ptr AirportCatalogSet::GetEntry(ClientContext &context, const EntryLookupInfo &lookup_info) 19 | { 20 | lock_guard l(entry_lock); 21 | if (!is_loaded) 22 | { 23 | is_loaded = true; 24 | LoadEntries(context); 25 | } 26 | auto entry = entries.find(lookup_info.GetEntryName()); 27 | if (entry == entries.end()) 28 | { 29 | return nullptr; 30 | } 31 | return entry->second.get(); 32 | } 33 | 34 | struct DropItemActionParameters 35 | { 36 | // the type of the item to drop, table, schema. 37 | std::string type; 38 | 39 | std::string catalog_name; 40 | std::string schema_name; 41 | std::string name; 42 | 43 | bool ignore_not_found; 44 | 45 | MSGPACK_DEFINE_MAP(type, catalog_name, schema_name, name, ignore_not_found) 46 | }; 47 | 48 | void AirportCatalogSet::DropEntry(ClientContext &context, DropInfo &info) 49 | { 50 | if (!is_loaded) 51 | { 52 | is_loaded = true; 53 | LoadEntries(context); 54 | } 55 | 56 | auto &airport_catalog = catalog.Cast(); 57 | arrow::flight::FlightCallOptions call_options; 58 | 59 | airport_add_standard_headers(call_options, airport_catalog.attach_parameters()->location()); 60 | airport_add_authorization_header(call_options, airport_catalog.attach_parameters()->auth_token()); 61 | 62 | auto flight_client = AirportAPI::FlightClientForLocation(airport_catalog.attach_parameters()->location()); 63 | 64 | // Common parameters 65 | DropItemActionParameters params; 66 | params.catalog_name = airport_catalog.internal_name(); 67 | params.schema_name = info.schema; 68 | params.name = info.name; 69 | params.ignore_not_found = (info.if_not_found == OnEntryNotFound::RETURN_NULL) ? true : false; 70 | 71 | std::string action_type; 72 | 73 | switch (info.type) 74 | { 75 | case CatalogType::TABLE_ENTRY: 76 | params.type = "table"; 77 | action_type = "drop_table"; 78 | break; 79 | case CatalogType::SCHEMA_ENTRY: 80 | params.type = "schema"; 81 | action_type = "drop_schema"; 82 | call_options.headers.emplace_back("airport-action-name", action_type); 83 | break; 84 | default: 85 | throw NotImplementedException("AirportCatalogSet::DropEntry for type"); 86 | } 87 | 88 | AIRPORT_MSGPACK_ACTION_SINGLE_PARAMETER(action, action_type, params); 89 | 90 | auto &server_location = airport_catalog.attach_parameters()->location(); 91 | 92 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto action_results, 93 | flight_client->DoAction(call_options, action), 94 | server_location, 95 | "airport_create_schema"); 96 | 97 | AIRPORT_ARROW_ASSERT_OK_LOCATION(action_results->Drain(), server_location, ""); 98 | 99 | D_ASSERT(entries.find(info.name) != entries.end()); 100 | 101 | EraseEntryInternal(info.name); 102 | } 103 | 104 | void AirportCatalogSet::EraseEntryInternal(const string &name) 105 | { 106 | lock_guard l(entry_lock); 107 | entries.erase(name); 108 | } 109 | 110 | void AirportCatalogSet::Scan(ClientContext &context, const std::function &callback) 111 | { 112 | lock_guard l(entry_lock); 113 | if (!is_loaded) 114 | { 115 | is_loaded = true; 116 | LoadEntries(context); 117 | } 118 | for (auto &entry : entries) 119 | { 120 | callback(*entry.second); 121 | } 122 | } 123 | 124 | optional_ptr AirportCatalogSet::CreateEntry(unique_ptr entry) 125 | { 126 | auto result = entry.get(); 127 | if (result->name.empty()) 128 | { 129 | throw CatalogException("AirportCatalogSet::CreateEntry called with empty name"); 130 | } 131 | // printf("Creating catalog entry\n"); 132 | entries.insert(make_pair(result->name, std::move(entry))); 133 | return result; 134 | } 135 | 136 | void AirportCatalogSet::ClearEntries() 137 | { 138 | lock_guard l(entry_lock); 139 | entries.clear(); 140 | is_loaded = false; 141 | } 142 | 143 | void AirportCatalogSet::ReplaceEntry( 144 | const string &name, 145 | unique_ptr entry) 146 | { 147 | lock_guard l(entry_lock); 148 | auto it = entries.find(name); 149 | if (it == entries.end()) 150 | { 151 | throw CatalogException("AirportCatalogSet::ReplaceEntry called with non-existing entry"); 152 | } 153 | 154 | if (entry->name == name) 155 | { 156 | it->second = std::move(entry); 157 | } 158 | else 159 | { 160 | entries.erase(name); 161 | CreateEntry(std::move(entry)); 162 | } 163 | } 164 | 165 | AirportInSchemaSet::AirportInSchemaSet(AirportSchemaEntry &schema) : AirportCatalogSet(schema.ParentCatalog()), schema(schema) 166 | { 167 | } 168 | 169 | optional_ptr AirportInSchemaSet::CreateEntry(unique_ptr entry) 170 | { 171 | if (!entry->internal) 172 | { 173 | entry->internal = schema.internal; 174 | } 175 | return AirportCatalogSet::CreateEntry(std::move(entry)); 176 | } 177 | 178 | } // namespace duckdb 179 | -------------------------------------------------------------------------------- /src/storage/airport_clear_cache.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "duckdb/parser/parsed_data/create_table_function_info.hpp" 3 | #include "duckdb/main/database_manager.hpp" 4 | #include "duckdb/main/attached_database.hpp" 5 | #include "storage/airport_catalog.hpp" 6 | 7 | namespace duckdb 8 | { 9 | 10 | struct ClearCacheFunctionData : public TableFunctionData 11 | { 12 | bool finished = false; 13 | }; 14 | 15 | static unique_ptr ClearCacheBind(ClientContext &context, TableFunctionBindInput &input, 16 | vector &return_types, vector &names) 17 | { 18 | 19 | auto result = make_uniq(); 20 | return_types.push_back(LogicalType::BOOLEAN); 21 | names.emplace_back("Success"); 22 | return result; 23 | } 24 | 25 | static void ClearAirportCaches(ClientContext &context) 26 | { 27 | auto databases = DatabaseManager::Get(context).GetDatabases(context); 28 | for (auto &db_ref : databases) 29 | { 30 | auto &db = db_ref.get(); 31 | auto &catalog = db.GetCatalog(); 32 | if (catalog.GetCatalogType() != "airport") 33 | { 34 | continue; 35 | } 36 | catalog.Cast().ClearCache(); 37 | } 38 | } 39 | 40 | static void ClearCacheFunction(ClientContext &context, TableFunctionInput &data_p, DataChunk &output) 41 | { 42 | auto &data = data_p.bind_data->CastNoConst(); 43 | if (data.finished) 44 | { 45 | return; 46 | } 47 | ClearAirportCaches(context); 48 | data.finished = true; 49 | } 50 | 51 | void AirportClearCacheFunction::ClearCacheOnSetting(ClientContext &context, SetScope scope, Value ¶meter) 52 | { 53 | ClearAirportCaches(context); 54 | } 55 | 56 | AirportClearCacheFunction::AirportClearCacheFunction() : TableFunction("airport_clear_cache", {}, ClearCacheFunction, ClearCacheBind) 57 | { 58 | } 59 | } // namespace duckdb 60 | -------------------------------------------------------------------------------- /src/storage/airport_curl_pool.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | 3 | #include 4 | #include 5 | #include 6 | #include 7 | #include "storage/airport_curl_pool.hpp" 8 | 9 | namespace duckdb 10 | { 11 | 12 | AirportCurlPool::AirportCurlPool(size_t size) 13 | { 14 | for (size_t i = 0; i < size; ++i) 15 | { 16 | CURL *handle = curl_easy_init(); 17 | if (handle) 18 | { 19 | _pool.push_back(handle); 20 | } 21 | } 22 | } 23 | 24 | AirportCurlPool::~AirportCurlPool() 25 | { 26 | for (auto handle : _pool) 27 | { 28 | curl_easy_cleanup(handle); 29 | } 30 | } 31 | 32 | CURL *AirportCurlPool::acquire() 33 | { 34 | std::unique_lock lock(_mutex); 35 | _cv.wait(lock, [this]() 36 | { return !_pool.empty(); }); 37 | CURL *handle = _pool.back(); 38 | _pool.pop_back(); 39 | return handle; 40 | } 41 | 42 | void AirportCurlPool::release(CURL *handle) 43 | { 44 | std::unique_lock lock(_mutex); 45 | _pool.push_back(handle); 46 | _cv.notify_one(); 47 | } 48 | 49 | } -------------------------------------------------------------------------------- /src/storage/airport_delete_parameterized.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | #include "storage/airport_delete_parameterized.hpp" 3 | #include "storage/airport_table_entry.hpp" 4 | #include "duckdb/planner/operator/logical_delete.hpp" 5 | #include "storage/airport_catalog.hpp" 6 | #include "storage/airport_transaction.hpp" 7 | #include "duckdb/planner/expression/bound_reference_expression.hpp" 8 | #include "duckdb/common/arrow/arrow_appender.hpp" 9 | #include "duckdb/common/arrow/arrow_converter.hpp" 10 | #include "duckdb/common/types/uuid.hpp" 11 | #include "duckdb/function/table/arrow/arrow_duck_schema.hpp" 12 | #include "duckdb/function/table/arrow.hpp" 13 | #include "airport_macros.hpp" 14 | #include "airport_request_headers.hpp" 15 | #include "airport_flight_exception.hpp" 16 | #include "airport_secrets.hpp" 17 | 18 | #include "duckdb/common/arrow/schema_metadata.hpp" 19 | 20 | #include "airport_flight_stream.hpp" 21 | #include "airport_take_flight.hpp" 22 | #include "storage/airport_exchange.hpp" 23 | #include "duckdb/execution/operator/filter/physical_filter.hpp" 24 | #include "duckdb/execution/operator/scan/physical_table_scan.hpp" 25 | 26 | namespace duckdb 27 | { 28 | 29 | //===--------------------------------------------------------------------===// 30 | // Plan 31 | //===--------------------------------------------------------------------===// 32 | static string ExtractFilters(PhysicalOperator &child, const string &statement) 33 | { 34 | // FIXME - all of this is pretty gnarly, we should provide a hook earlier on 35 | // in the planning process to convert this into a SQL statement 36 | if (child.type == PhysicalOperatorType::FILTER) 37 | { 38 | auto &filter = child.Cast(); 39 | auto result = ExtractFilters(child.children[0], statement); 40 | auto filter_str = filter.expression->ToString(); 41 | if (result.empty()) 42 | { 43 | return filter_str; 44 | } 45 | else 46 | { 47 | return result + " AND " + filter_str; 48 | } 49 | } 50 | else if (child.type == PhysicalOperatorType::TABLE_SCAN) 51 | { 52 | const auto &table_scan = child.Cast(); 53 | if (!table_scan.table_filters) 54 | { 55 | return string(); 56 | } 57 | throw NotImplementedException("Pushed down table filters not supported currently"); 58 | } 59 | else 60 | { 61 | throw NotImplementedException("Unsupported operator type %s in %s statement - only simple expressions " 62 | "(e.g. %s " 63 | "FROM tbl WHERE x=y) are supported in AirportParameterizedDelete", 64 | PhysicalOperatorToString(child.type), statement, statement); 65 | } 66 | } 67 | 68 | AirportDeleteParameterized::AirportDeleteParameterized(LogicalOperator &op, TableCatalogEntry &table, PhysicalOperator &plan) 69 | : PhysicalOperator(PhysicalOperatorType::EXTENSION, op.types, 1), table(table), 70 | sql_filters(ExtractFilters(plan, "DELETE")) 71 | { 72 | printf("Got SQL filters: %s\n", sql_filters.c_str()); 73 | } 74 | 75 | //===--------------------------------------------------------------------===// 76 | // States 77 | //===--------------------------------------------------------------------===// 78 | class AirportDeleteParameterizedGlobalState : public GlobalSinkState 79 | { 80 | public: 81 | explicit AirportDeleteParameterizedGlobalState() : affected_rows(0) 82 | { 83 | } 84 | 85 | idx_t affected_rows; 86 | }; 87 | 88 | unique_ptr AirportDeleteParameterized::GetGlobalSinkState(ClientContext &context) const 89 | { 90 | printf("AirportDeleteParameterized::GetGlobalSinkState\n"); 91 | return make_uniq(); 92 | } 93 | 94 | //===--------------------------------------------------------------------===// 95 | // Sink 96 | //===--------------------------------------------------------------------===// 97 | SinkResultType AirportDeleteParameterized::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const 98 | { 99 | printf("AirportDeleteParameterized::Sink\n"); 100 | return SinkResultType::FINISHED; 101 | } 102 | 103 | //===--------------------------------------------------------------------===// 104 | // Finalize 105 | //===--------------------------------------------------------------------===// 106 | SinkFinalizeType AirportDeleteParameterized::Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 107 | OperatorSinkFinalizeInput &input) const 108 | { 109 | auto &gstate = input.global_state.Cast(); 110 | gstate.affected_rows = 0; 111 | // auto &transaction = MySQLTransaction::Get(context, table.catalog); 112 | // auto &connection = transaction.GetConnection(); 113 | // auto result = connection.Query(query); 114 | // gstate.affected_rows = result->AffectedRows(); 115 | printf("AirportDeleteParameterized::Finalize\n"); 116 | return SinkFinalizeType::READY; 117 | } 118 | 119 | //===--------------------------------------------------------------------===// 120 | // GetData 121 | //===--------------------------------------------------------------------===// 122 | SourceResultType AirportDeleteParameterized::GetData(ExecutionContext &context, DataChunk &chunk, 123 | OperatorSourceInput &input) const 124 | { 125 | auto &insert_gstate = sink_state->Cast(); 126 | chunk.SetCardinality(1); 127 | chunk.SetValue(0, 0, Value::BIGINT(insert_gstate.affected_rows)); 128 | printf("AirportDeleteParameterized::GetData\n"); 129 | return SourceResultType::FINISHED; 130 | } 131 | 132 | //===--------------------------------------------------------------------===// 133 | // Helpers 134 | //===--------------------------------------------------------------------===// 135 | string AirportDeleteParameterized::GetName() const 136 | { 137 | return "AIRPORT_DELETE_PARAMETERIZED"; 138 | } 139 | 140 | InsertionOrderPreservingMap AirportDeleteParameterized::ParamsToString() const 141 | { 142 | InsertionOrderPreservingMap result; 143 | result["Table Name"] = table.name; 144 | return result; 145 | } 146 | 147 | // string ConstructDeleteStatement(LogicalDelete &op, PhysicalOperator &child) 148 | // { 149 | // string result = "DELETE FROM "; 150 | // result += MySQLUtils::WriteIdentifier(op.table.schema.name); 151 | // result += "."; 152 | // result += MySQLUtils::WriteIdentifier(op.table.name); 153 | // auto filters = ExtractFilters(child, "DELETE"); 154 | // if (!filters.empty()) 155 | // { 156 | // result += " WHERE " + filters; 157 | // } 158 | // return result; 159 | // } 160 | 161 | // unique_ptr MySQLCatalog::PlanDelete(ClientContext &context, LogicalDelete &op, 162 | // unique_ptr plan) 163 | // { 164 | // if (op.return_chunk) 165 | // { 166 | // throw BinderException("RETURNING clause not yet supported for deletion of a MySQL table"); 167 | // } 168 | 169 | // auto result = make_uniq(op, "DELETE", op.table, ConstructDeleteStatement(op, *plan)); 170 | // result->children.push_back(std::move(plan)); 171 | // return std::move(result); 172 | // } 173 | 174 | // string ConstructUpdateStatement(LogicalUpdate &op, PhysicalOperator &child) 175 | // { 176 | // // FIXME - all of this is pretty gnarly, we should provide a hook earlier on 177 | // // in the planning process to convert this into a SQL statement 178 | // string result = "UPDATE"; 179 | // result += MySQLUtils::WriteIdentifier(op.table.schema.name); 180 | // result += "."; 181 | // result += MySQLUtils::WriteIdentifier(op.table.name); 182 | // result += " SET "; 183 | // if (child.type != PhysicalOperatorType::PROJECTION) 184 | // { 185 | // throw NotImplementedException("MySQL Update not supported - Expected the " 186 | // "child of an update to be a projection"); 187 | // } 188 | // auto &proj = child.Cast(); 189 | // for (idx_t c = 0; c < op.columns.size(); c++) 190 | // { 191 | // if (c > 0) 192 | // { 193 | // result += ", "; 194 | // } 195 | // auto &col = op.table.GetColumn(op.table.GetColumns().PhysicalToLogical(op.columns[c])); 196 | // result += MySQLUtils::WriteIdentifier(col.GetName()); 197 | // result += " = "; 198 | // if (op.expressions[c]->type == ExpressionType::VALUE_DEFAULT) 199 | // { 200 | // result += "DEFAULT"; 201 | // continue; 202 | // } 203 | // if (op.expressions[c]->type != ExpressionType::BOUND_REF) 204 | // { 205 | // throw NotImplementedException("MySQL Update not supported - Expected a bound reference expression"); 206 | // } 207 | // auto &ref = op.expressions[c]->Cast(); 208 | // result += proj.select_list[ref.index]->ToString(); 209 | // } 210 | // result += " "; 211 | // auto filters = ExtractFilters(*child.children[0], "UPDATE"); 212 | // if (!filters.empty()) 213 | // { 214 | // result += " WHERE " + filters; 215 | // } 216 | // return result; 217 | // } 218 | 219 | // unique_ptr MySQLCatalog::PlanUpdate(ClientContext &context, LogicalUpdate &op, 220 | // unique_ptr plan) 221 | // { 222 | // if (op.return_chunk) 223 | // { 224 | // throw BinderException("RETURNING clause not yet supported for updates of a MySQL table"); 225 | // } 226 | // auto result = make_uniq(op, "UPDATE", op.table, ConstructUpdateStatement(op, *plan)); 227 | // result->children.push_back(std::move(plan)); 228 | // return std::move(result); 229 | // } 230 | 231 | } // namespace duckdb 232 | -------------------------------------------------------------------------------- /src/storage/airport_exchange.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | #include "storage/airport_delete.hpp" 3 | #include "storage/airport_table_entry.hpp" 4 | #include "duckdb/planner/operator/logical_delete.hpp" 5 | #include "storage/airport_catalog.hpp" 6 | #include "storage/airport_transaction.hpp" 7 | #include "duckdb/planner/expression/bound_reference_expression.hpp" 8 | #include "duckdb/common/arrow/arrow_appender.hpp" 9 | #include "duckdb/common/arrow/arrow_converter.hpp" 10 | #include "duckdb/common/types/uuid.hpp" 11 | #include "duckdb/function/table/arrow/arrow_duck_schema.hpp" 12 | #include "duckdb/function/table/arrow.hpp" 13 | #include "airport_macros.hpp" 14 | #include "airport_request_headers.hpp" 15 | #include "airport_flight_exception.hpp" 16 | #include "airport_secrets.hpp" 17 | #include "airport_flight_stream.hpp" 18 | #include "airport_take_flight.hpp" 19 | #include "storage/airport_exchange.hpp" 20 | #include "airport_schema_utils.hpp" 21 | #include "duckdb/common/arrow/schema_metadata.hpp" 22 | 23 | #include 24 | 25 | namespace duckdb 26 | { 27 | 28 | static int findIndex(const std::vector &vec, const std::string &target) 29 | { 30 | auto it = std::find(vec.begin(), vec.end(), target); 31 | 32 | if (it == vec.end()) 33 | { 34 | throw std::runtime_error("String not found in vector"); 35 | } 36 | 37 | return std::distance(vec.begin(), it); 38 | } 39 | 40 | void AirportExchangeGetGlobalSinkState(ClientContext &context, 41 | const TableCatalogEntry &table, 42 | const AirportTableEntry &airport_table, 43 | AirportExchangeGlobalState *global_state, 44 | const ArrowSchema &send_schema, 45 | const bool return_chunk, 46 | const string exchange_operation, 47 | const vector destination_chunk_column_names, 48 | const std::optional transaction_id) 49 | { 50 | AIRPORT_ASSIGN_OR_RAISE_CONTAINER( 51 | global_state->send_schema, 52 | arrow::ImportSchema((ArrowSchema *)&send_schema), 53 | airport_table.table_data, 54 | ""); 55 | 56 | const auto &server_location = airport_table.table_data->server_location(); 57 | const auto &descriptor = airport_table.table_data->descriptor(); 58 | 59 | // global_state->flight_descriptor = descriptor; 60 | 61 | auto auth_token = AirportAuthTokenForLocation(context, server_location, "", ""); 62 | 63 | D_ASSERT(airport_table.table_data != nullptr); 64 | 65 | auto flight_client = AirportAPI::FlightClientForLocation(server_location); 66 | 67 | auto trace_uuid = airport_trace_id(); 68 | 69 | arrow::flight::FlightCallOptions call_options; 70 | airport_add_standard_headers(call_options, server_location); 71 | airport_add_authorization_header(call_options, auth_token); 72 | airport_add_trace_id_header(call_options, trace_uuid); 73 | 74 | // Indicate that we are doing a delete. 75 | call_options.headers.emplace_back("airport-operation", exchange_operation); 76 | 77 | if (transaction_id.has_value() && !transaction_id.value().empty()) 78 | { 79 | call_options.headers.emplace_back("airport-transaction-id", transaction_id.value()); 80 | } 81 | 82 | // Indicate if the caller is interested in data being returned. 83 | call_options.headers.emplace_back("return-chunks", return_chunk ? "1" : "0"); 84 | 85 | airport_add_flight_path_header(call_options, descriptor); 86 | 87 | AIRPORT_ASSIGN_OR_RAISE_CONTAINER( 88 | auto exchange_result, 89 | flight_client->DoExchange(call_options, descriptor), 90 | airport_table.table_data, ""); 91 | 92 | // Tell the server the schema that we will be using to write data. 93 | AIRPORT_ARROW_ASSERT_OK_CONTAINER( 94 | exchange_result.writer->Begin(global_state->send_schema), 95 | airport_table.table_data, 96 | "Begin schema"); 97 | 98 | D_ASSERT(exchange_result.reader != nullptr); 99 | D_ASSERT(exchange_result.writer != nullptr); 100 | // Now that there is a reader stream and a writer stream, we want to reuse the Arrow 101 | // scan code as much as possible, but the problem is it assumes its being called as 102 | // part of a table returning function, with the life cycle of bind, init global, init local 103 | // and scan. 104 | // 105 | // But we can simulate most of that here. 106 | 107 | AIRPORT_ASSIGN_OR_RAISE_CONTAINER(auto read_schema, 108 | exchange_result.reader->GetSchema(), 109 | airport_table.table_data, 110 | ""); 111 | 112 | auto scan_bind_data = make_uniq( 113 | (stream_factory_produce_t)&AirportCreateStream, 114 | trace_uuid, 115 | -1, 116 | AirportTakeFlightParameters(server_location, context), 117 | std::nullopt, 118 | read_schema, 119 | descriptor, 120 | nullptr); 121 | 122 | vector column_ids; 123 | 124 | if (return_chunk) 125 | { 126 | // printf("Schema of reader stream is:\n----------\n%s\n---------\n", read_schema->ToString().c_str()); 127 | 128 | vector reading_arrow_column_names; 129 | 130 | const auto column_count = (idx_t)scan_bind_data->schema_root.arrow_schema.n_children; 131 | 132 | reading_arrow_column_names.reserve(column_count); 133 | 134 | std::optional rowid_column_name = std::nullopt; 135 | 136 | for (idx_t col_idx = 0; 137 | col_idx < column_count; col_idx++) 138 | { 139 | const auto &schema_item = *scan_bind_data->schema_root.arrow_schema.children[col_idx]; 140 | if (!schema_item.release) 141 | { 142 | throw InvalidInputException("airport_exchange: released schema passed"); 143 | } 144 | 145 | // Check to see if the column is marked as the rowid column. 146 | if (AirportFieldMetadataIsRowId(schema_item.metadata)) 147 | { 148 | rowid_column_name = schema_item.name; 149 | } 150 | 151 | reading_arrow_column_names.push_back(AirportNameForField(schema_item.name, col_idx)); 152 | } 153 | 154 | column_ids.reserve(destination_chunk_column_names.size()); 155 | 156 | for (size_t output_index = 0; output_index < destination_chunk_column_names.size(); output_index++) 157 | { 158 | auto found_index = findIndex(reading_arrow_column_names, destination_chunk_column_names[output_index]); 159 | if (exchange_operation != "update") 160 | { 161 | column_ids.push_back(found_index); 162 | } 163 | else 164 | { 165 | // This is right for outputs, because it allowed the read chunk to happen. 166 | if (rowid_column_name.has_value() && 167 | destination_chunk_column_names[output_index] == rowid_column_name) 168 | { 169 | column_ids.push_back(COLUMN_IDENTIFIER_ROW_ID); 170 | } 171 | else 172 | { 173 | column_ids.push_back(output_index); 174 | } 175 | } 176 | // printf("Output data chunk column %s (type=%s) (%d) comes from arrow column index index %d\n", 177 | // destination_chunk_column_names[output_index].c_str(), 178 | // arrow_types[found_index].c_str(), 179 | // output_index, 180 | // found_index); 181 | } 182 | 183 | scan_bind_data->examine_schema(context, true); 184 | } 185 | 186 | // For each index in the arrow table, the column_ids is asked what 187 | // where to map that column, the row id can be expressed there. 188 | 189 | // There shouldn't be any projection ids. 190 | vector projection_ids; 191 | 192 | // Now to initialize the Arrow scan from the reader stream we need to do the steps 193 | // that the normal table returning function does. 194 | 195 | // bind 196 | // init global state 197 | // init local state 198 | // scan... 199 | 200 | // Init the global state. 201 | 202 | // Retain the global state. 203 | global_state->scan_global_state = make_uniq(); 204 | 205 | // Now simulate the init input. 206 | auto fake_init_input = TableFunctionInitInput( 207 | &scan_bind_data->Cast(), 208 | column_ids, 209 | projection_ids, 210 | nullptr); 211 | 212 | // Local init. 213 | 214 | D_ASSERT(exchange_result.reader != nullptr); 215 | auto current_chunk = make_uniq(); 216 | auto scan_local_state = make_uniq( 217 | std::move(current_chunk), 218 | context, 219 | std::move(exchange_result.reader), 220 | fake_init_input); 221 | scan_local_state->set_stream(AirportProduceArrowScan( 222 | scan_bind_data->CastNoConst(), 223 | column_ids, 224 | nullptr, 225 | nullptr, // No progress reporting. 226 | &scan_bind_data->last_app_metadata, 227 | scan_bind_data->schema(), 228 | *scan_bind_data, 229 | *scan_local_state)); 230 | 231 | scan_local_state->column_ids = fake_init_input.column_ids; 232 | scan_local_state->filters = fake_init_input.filters.get(); 233 | 234 | global_state->scan_local_state = std::move(scan_local_state); 235 | 236 | // Create a parameter is the commonly passed to the other functions. 237 | global_state->scan_bind_data = std::move(scan_bind_data); 238 | global_state->writer = std::move(exchange_result.writer); 239 | 240 | global_state->scan_table_function_input = make_uniq( 241 | global_state->scan_bind_data.get(), 242 | global_state->scan_local_state.get(), 243 | global_state->scan_global_state.get()); 244 | } 245 | } -------------------------------------------------------------------------------- /src/storage/airport_scalar_function_set.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | #include "duckdb/catalog/catalog_entry/scalar_function_catalog_entry.hpp" 3 | #include "duckdb/catalog/catalog_entry/table_function_catalog_entry.hpp" 4 | #include "duckdb/catalog/dependency_list.hpp" 5 | #include "duckdb/common/arrow/arrow_appender.hpp" 6 | #include "duckdb/common/arrow/arrow_converter.hpp" 7 | #include "duckdb/common/arrow/schema_metadata.hpp" 8 | #include "duckdb/common/types/uuid.hpp" 9 | #include "duckdb/function/table/arrow.hpp" 10 | #include "duckdb/parser/constraints/check_constraint.hpp" 11 | #include "duckdb/parser/constraints/list.hpp" 12 | #include "duckdb/parser/constraints/not_null_constraint.hpp" 13 | #include "duckdb/parser/constraints/unique_constraint.hpp" 14 | #include "duckdb/parser/expression/constant_expression.hpp" 15 | #include "duckdb/parser/parsed_data/create_function_info.hpp" 16 | #include "duckdb/parser/parsed_data/create_scalar_function_info.hpp" 17 | #include "duckdb/parser/parsed_data/create_table_function_info.hpp" 18 | #include "duckdb/parser/parsed_data/create_table_info.hpp" 19 | #include "duckdb/parser/parsed_data/create_table_info.hpp" 20 | #include "duckdb/parser/parser.hpp" 21 | #include "duckdb/planner/expression/bound_function_expression.hpp" 22 | #include "duckdb/planner/parsed_data/bound_create_table_info.hpp" 23 | #include 24 | #include 25 | #include 26 | #include 27 | #include 28 | #include "airport_flight_stream.hpp" 29 | #include "airport_request_headers.hpp" 30 | #include "airport_macros.hpp" 31 | #include "airport_macros.hpp" 32 | #include "airport_scalar_function.hpp" 33 | #include "airport_secrets.hpp" 34 | #include "airport_take_flight.hpp" 35 | #include "storage/airport_catalog_api.hpp" 36 | #include "storage/airport_catalog.hpp" 37 | #include "storage/airport_curl_pool.hpp" 38 | #include "storage/airport_exchange.hpp" 39 | #include "storage/airport_schema_entry.hpp" 40 | #include "storage/airport_table_set.hpp" 41 | #include "storage/airport_transaction.hpp" 42 | #include "airport_schema_utils.hpp" 43 | #include "storage/airport_alter_parameters.hpp" 44 | #include "storage/airport_scalar_function_set.hpp" 45 | 46 | #include "duckdb/main/extension_util.hpp" 47 | 48 | namespace duckdb 49 | { 50 | 51 | // Given an Arrow schema return a vector of the LogicalTypes for that schema. 52 | static vector AirportSchemaToLogicalTypes( 53 | ClientContext &context, 54 | std::shared_ptr schema, 55 | const string &server_location, 56 | const flight::FlightDescriptor &flight_descriptor) 57 | { 58 | ArrowSchemaWrapper schema_root; 59 | 60 | AIRPORT_ARROW_ASSERT_OK_LOCATION_DESCRIPTOR( 61 | ExportSchema(*schema, &schema_root.arrow_schema), 62 | server_location, 63 | flight_descriptor, 64 | "ExportSchema"); 65 | 66 | vector return_types; 67 | auto &config = DBConfig::GetConfig(context); 68 | 69 | const idx_t column_count = (idx_t)schema_root.arrow_schema.n_children; 70 | 71 | return_types.reserve(column_count); 72 | 73 | for (idx_t col_idx = 0; 74 | col_idx < column_count; col_idx++) 75 | { 76 | auto &schema_item = *schema_root.arrow_schema.children[col_idx]; 77 | if (!schema_item.release) 78 | { 79 | throw InvalidInputException("AirportSchemaToLogicalTypes: released schema passed"); 80 | } 81 | auto arrow_type = ArrowType::GetArrowLogicalType(config, schema_item); 82 | 83 | if (schema_item.dictionary) 84 | { 85 | auto dictionary_type = ArrowType::GetArrowLogicalType(config, *schema_item.dictionary); 86 | arrow_type->SetDictionary(std::move(dictionary_type)); 87 | } 88 | 89 | // Indicate that the field should select any type. 90 | bool is_any_type = false; 91 | if (schema_item.metadata != nullptr) 92 | { 93 | auto column_metadata = ArrowSchemaMetadata(schema_item.metadata); 94 | if (!column_metadata.GetOption("is_any_type").empty()) 95 | { 96 | is_any_type = true; 97 | } 98 | } 99 | 100 | if (is_any_type) 101 | { 102 | // This will be sorted out in the bind of the function. 103 | return_types.push_back(LogicalType::ANY); 104 | } 105 | else 106 | { 107 | return_types.emplace_back(arrow_type->GetDuckType()); 108 | } 109 | } 110 | return return_types; 111 | } 112 | 113 | void AirportScalarFunctionSet::LoadEntries(ClientContext &context) 114 | { 115 | // auto &transaction = AirportTransaction::Get(context, catalog); 116 | 117 | auto &airport_catalog = catalog.Cast(); 118 | 119 | // TODO: handle out-of-order columns using position property 120 | auto curl = connection_pool_.acquire(); 121 | auto contents = AirportAPI::GetSchemaItems( 122 | curl, 123 | catalog.GetDBPath(), 124 | schema.name, 125 | schema.serialized_source(), 126 | cache_directory_, 127 | airport_catalog.attach_parameters()); 128 | 129 | connection_pool_.release(curl); 130 | 131 | // printf("AirportScalarFunctionSet loading entries\n"); 132 | // printf("Total functions: %lu\n", tables_and_functions.second.size()); 133 | 134 | // There can be functions with the same name. 135 | std::unordered_map> functions_by_name; 136 | 137 | for (auto &function : contents->scalar_functions) 138 | { 139 | AirportFunctionCatalogSchemaNameKey function_key{function.catalog_name(), function.schema_name(), function.name()}; 140 | functions_by_name[function_key].emplace_back(function); 141 | } 142 | 143 | for (const auto &pair : functions_by_name) 144 | { 145 | ScalarFunctionSet flight_func_set(pair.first.name); 146 | 147 | // FIXME: need a way to specify the function stability. 148 | for (const auto &function : pair.second) 149 | { 150 | auto input_types = AirportSchemaToLogicalTypes(context, function.input_schema(), function.server_location(), function.descriptor()); 151 | 152 | auto output_types = AirportSchemaToLogicalTypes(context, function.schema(), function.server_location(), function.descriptor()); 153 | D_ASSERT(output_types.size() == 1); 154 | 155 | auto scalar_func = ScalarFunction(input_types, output_types[0], 156 | AirportScalarFunctionProcessChunk, 157 | AirportScalarFunctionBind, 158 | nullptr, 159 | nullptr, 160 | AirportScalarFunctionInitLocalState, 161 | LogicalTypeId::INVALID, 162 | duckdb::FunctionStability::VOLATILE, 163 | duckdb::FunctionNullHandling::DEFAULT_NULL_HANDLING, 164 | nullptr); 165 | scalar_func.function_info = make_uniq(function.name(), 166 | function, 167 | function.schema(), 168 | function.input_schema(), 169 | catalog); 170 | 171 | flight_func_set.AddFunction(scalar_func); 172 | } 173 | 174 | CreateScalarFunctionInfo info = CreateScalarFunctionInfo(flight_func_set); 175 | info.catalog = pair.first.catalog_name; 176 | info.schema = pair.first.schema_name; 177 | 178 | info.internal = true; 179 | 180 | auto function_entry = make_uniq_base( 181 | catalog, 182 | schema, 183 | info.Cast()); 184 | 185 | CreateEntry(std::move(function_entry)); 186 | } 187 | } 188 | 189 | } -------------------------------------------------------------------------------- /src/storage/airport_schema_entry.cpp: -------------------------------------------------------------------------------- 1 | #include "storage/airport_schema_entry.hpp" 2 | #include "storage/airport_table_entry.hpp" 3 | #include "storage/airport_transaction.hpp" 4 | #include "duckdb/parser/parsed_data/create_view_info.hpp" 5 | #include "duckdb/parser/parsed_data/create_index_info.hpp" 6 | #include "duckdb/planner/parsed_data/bound_create_table_info.hpp" 7 | #include "duckdb/parser/parsed_data/drop_info.hpp" 8 | #include "duckdb/parser/constraints/list.hpp" 9 | #include "duckdb/common/unordered_set.hpp" 10 | #include "duckdb/parser/parsed_data/alter_info.hpp" 11 | #include "duckdb/parser/parsed_data/alter_table_info.hpp" 12 | #include "duckdb/parser/parsed_expression_iterator.hpp" 13 | #include "storage/airport_curl_pool.hpp" 14 | #include "airport_request_headers.hpp" 15 | #include "storage/airport_catalog.hpp" 16 | #include "airport_macros.hpp" 17 | #include 18 | #include 19 | 20 | namespace duckdb 21 | { 22 | 23 | AirportSchemaEntry::AirportSchemaEntry(Catalog &catalog, 24 | CreateSchemaInfo &info, AirportCurlPool &connection_pool, const string &cache_directory, 25 | const AirportAPISchema &schema_data) 26 | : SchemaCatalogEntry(catalog, info), schema_data_(schema_data), tables(connection_pool, *this, cache_directory), scalar_functions(connection_pool, *this, cache_directory), table_functions(connection_pool, *this, cache_directory) 27 | { 28 | } 29 | 30 | AirportSchemaEntry::~AirportSchemaEntry() 31 | { 32 | } 33 | 34 | AirportTransaction &GetAirportTransaction(CatalogTransaction transaction) 35 | { 36 | if (!transaction.transaction) 37 | { 38 | throw InternalException("No transaction in GetAirportTransaction!?"); 39 | } 40 | return transaction.transaction->Cast(); 41 | } 42 | 43 | optional_ptr AirportSchemaEntry::CreateTable(CatalogTransaction transaction, BoundCreateTableInfo &info) 44 | { 45 | const auto &base_info = info.Base(); 46 | auto table_name = base_info.table; 47 | if (base_info.on_conflict == OnCreateConflict::REPLACE_ON_CONFLICT) 48 | { 49 | throw NotImplementedException("REPLACE ON CONFLICT in CreateTable"); 50 | } 51 | return tables.CreateTable(transaction.GetContext(), info); 52 | } 53 | 54 | optional_ptr AirportSchemaEntry::CreateFunction(CatalogTransaction transaction, CreateFunctionInfo &info) 55 | { 56 | throw BinderException("Airport databases do not support creating functions"); 57 | } 58 | 59 | void AirportUnqualifyColumnRef(ParsedExpression &expr) 60 | { 61 | if (expr.type == ExpressionType::COLUMN_REF) 62 | { 63 | auto &colref = expr.Cast(); 64 | auto name = std::move(colref.column_names.back()); 65 | colref.column_names = {std::move(name)}; 66 | return; 67 | } 68 | ParsedExpressionIterator::EnumerateChildren(expr, AirportUnqualifyColumnRef); 69 | } 70 | 71 | optional_ptr AirportSchemaEntry::CreateIndex(CatalogTransaction transaction, CreateIndexInfo &info, 72 | TableCatalogEntry &table) 73 | { 74 | throw NotImplementedException("CreateIndex"); 75 | } 76 | 77 | string GetAirportCreateView(const CreateViewInfo &info) 78 | { 79 | throw NotImplementedException("GetCreateView"); 80 | } 81 | 82 | optional_ptr AirportSchemaEntry::CreateView(CatalogTransaction transaction, CreateViewInfo &info) 83 | { 84 | if (info.sql.empty()) 85 | { 86 | throw BinderException("Cannot create view in Airport that originated from an " 87 | "empty SQL statement"); 88 | } 89 | if (info.on_conflict == OnCreateConflict::REPLACE_ON_CONFLICT || 90 | info.on_conflict == OnCreateConflict::IGNORE_ON_CONFLICT) 91 | { 92 | auto current_entry = GetEntry(transaction, CatalogType::VIEW_ENTRY, info.view_name); 93 | if (current_entry) 94 | { 95 | if (info.on_conflict == OnCreateConflict::IGNORE_ON_CONFLICT) 96 | { 97 | return current_entry; 98 | } 99 | throw NotImplementedException("REPLACE ON CONFLICT in CreateView"); 100 | } 101 | } 102 | // auto &airport_transaction = GetAirportTransaction(transaction); 103 | // uc_transaction.Query(GetAirportCreateView(info)); 104 | return tables.RefreshTable(transaction.GetContext(), info.view_name); 105 | } 106 | 107 | optional_ptr AirportSchemaEntry::CreateType(CatalogTransaction transaction, CreateTypeInfo &info) 108 | { 109 | throw BinderException("Airport databases do not support creating types"); 110 | } 111 | 112 | optional_ptr AirportSchemaEntry::CreateSequence(CatalogTransaction transaction, CreateSequenceInfo &info) 113 | { 114 | throw BinderException("Airport databases do not support creating sequences"); 115 | } 116 | 117 | optional_ptr AirportSchemaEntry::CreateTableFunction(CatalogTransaction transaction, 118 | CreateTableFunctionInfo &info) 119 | { 120 | throw BinderException("Airport databases do not support creating table functions"); 121 | } 122 | 123 | optional_ptr AirportSchemaEntry::CreateCopyFunction(CatalogTransaction transaction, 124 | CreateCopyFunctionInfo &info) 125 | { 126 | throw BinderException("Airport databases do not support creating copy functions"); 127 | } 128 | 129 | optional_ptr AirportSchemaEntry::CreatePragmaFunction(CatalogTransaction transaction, 130 | CreatePragmaFunctionInfo &info) 131 | { 132 | throw BinderException("Airport databases do not support creating pragma functions"); 133 | } 134 | 135 | optional_ptr AirportSchemaEntry::CreateCollation(CatalogTransaction transaction, CreateCollationInfo &info) 136 | { 137 | throw BinderException("Airport databases do not support creating collations"); 138 | } 139 | 140 | void AirportSchemaEntry::Alter(CatalogTransaction transaction, AlterInfo &info) 141 | { 142 | if (info.type != AlterType::ALTER_TABLE) 143 | { 144 | throw BinderException("Only altering tables is supported for now"); 145 | } 146 | auto &alter = info.Cast(); 147 | tables.AlterTable(transaction.GetContext(), alter); 148 | } 149 | 150 | bool CatalogTypeIsSupported(CatalogType type) 151 | { 152 | switch (type) 153 | { 154 | case CatalogType::SCALAR_FUNCTION_ENTRY: 155 | case CatalogType::TABLE_ENTRY: 156 | case CatalogType::TABLE_FUNCTION_ENTRY: 157 | return true; 158 | 159 | default: 160 | return false; 161 | } 162 | } 163 | 164 | void AirportSchemaEntry::Scan(ClientContext &context, CatalogType type, 165 | const std::function &callback) 166 | { 167 | if (!CatalogTypeIsSupported(type)) 168 | { 169 | return; 170 | } 171 | 172 | GetCatalogSet(type).Scan(context, callback); 173 | } 174 | void AirportSchemaEntry::Scan(CatalogType type, const std::function &callback) 175 | { 176 | throw NotImplementedException("Scan without context not supported"); 177 | } 178 | 179 | void AirportSchemaEntry::DropEntry(ClientContext &context, DropInfo &info) 180 | { 181 | switch (info.type) 182 | { 183 | case CatalogType::TABLE_ENTRY: 184 | { 185 | tables.DropEntry(context, info); 186 | break; 187 | } 188 | default: 189 | throw NotImplementedException("AirportSchemaEntry::DropEntry for type"); 190 | } 191 | } 192 | 193 | optional_ptr AirportSchemaEntry::LookupEntry(CatalogTransaction transaction, const EntryLookupInfo &lookup_info) 194 | { 195 | if (!CatalogTypeIsSupported(lookup_info.GetCatalogType())) 196 | { 197 | return nullptr; 198 | } 199 | return GetCatalogSet(lookup_info.GetCatalogType()).GetEntry(transaction.GetContext(), lookup_info); 200 | } 201 | 202 | AirportCatalogSet &AirportSchemaEntry::GetCatalogSet(CatalogType type) 203 | { 204 | switch (type) 205 | { 206 | case CatalogType::TABLE_ENTRY: 207 | case CatalogType::VIEW_ENTRY: 208 | return tables; 209 | case CatalogType::SCALAR_FUNCTION_ENTRY: 210 | return scalar_functions; 211 | case CatalogType::TABLE_FUNCTION_ENTRY: 212 | return table_functions; 213 | default: 214 | string error_message = "Airport: Type not supported for GetCatalogSet: " + CatalogTypeToString(type); 215 | throw NotImplementedException(error_message); 216 | } 217 | } 218 | 219 | } // namespace duckdb 220 | -------------------------------------------------------------------------------- /src/storage/airport_schema_set.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "storage/airport_schema_set.hpp" 3 | #include "storage/airport_catalog.hpp" 4 | #include "storage/airport_transaction.hpp" 5 | #include "storage/airport_catalog_api.hpp" 6 | #include "storage/airport_schema_entry.hpp" 7 | #include "duckdb/parser/parsed_data/create_schema_info.hpp" 8 | #include "duckdb/catalog/catalog.hpp" 9 | #include "duckdb/common/file_system.hpp" 10 | 11 | #include "airport_request_headers.hpp" 12 | #include "airport_macros.hpp" 13 | #include 14 | #include "msgpack.hpp" 15 | 16 | namespace duckdb 17 | { 18 | // Set the connection pool size. 19 | AirportSchemaSet::AirportSchemaSet(Catalog &catalog) : AirportCatalogSet(catalog), connection_pool(32) 20 | { 21 | } 22 | 23 | static bool IsInternalTable(const string &catalog, const string &schema) 24 | { 25 | if (schema == "information_schema") 26 | { 27 | return true; 28 | } 29 | return false; 30 | } 31 | 32 | static string DuckDBHomeDirectory(ClientContext &context) 33 | { 34 | auto &fs = FileSystem::GetFileSystem(context); 35 | 36 | string home_directory = fs.GetHomeDirectory(); 37 | // exception if the home directory does not exist, don't create whatever we think is home 38 | if (!fs.DirectoryExists(home_directory)) 39 | { 40 | throw IOException("Can't find the home directory at '%s'\nSpecify a home directory using the SET " 41 | "home_directory='/path/to/dir' option.", 42 | home_directory); 43 | } 44 | string cache_path = home_directory; 45 | cache_path = fs.JoinPath(cache_path, ".duckdb"); 46 | return cache_path; 47 | } 48 | 49 | void AirportSchemaSet::LoadEntireSet(ClientContext &context) 50 | { 51 | lock_guard l(entry_lock); 52 | 53 | if (called_load_entries == false) 54 | { 55 | // We haven't called load entries yet. 56 | LoadEntries(context); 57 | called_load_entries = true; 58 | } 59 | } 60 | 61 | void AirportSchemaSet::LoadEntries(ClientContext &context) 62 | { 63 | if (called_load_entries) 64 | { 65 | return; 66 | } 67 | // printf("Calling LoadEntries on AirportSchemaSet, catalog basically\n"); 68 | 69 | auto &airport_catalog = catalog.Cast(); 70 | string cache_path = DuckDBHomeDirectory(context); 71 | 72 | // catalog.GetName() is the catalog name even if its been aliased. 73 | // airport_catalog.internal_name() is the name of the database as passed ot attach. 74 | auto returned_collection = AirportAPI::GetSchemas(airport_catalog.internal_name(), airport_catalog.attach_parameters()); 75 | 76 | airport_catalog.loaded_catalog_version = returned_collection->version_info; 77 | 78 | collection = std::move(returned_collection); 79 | 80 | std::unordered_set seen_schema_names; 81 | 82 | // So the database can have all of its schemas sent at the top level. 83 | // 84 | // It can return a URL or an inline serialization of all saved schemas 85 | // 86 | // When the individual schemas are loaded they will be loaded through the 87 | // cached content that is already present on the disk, or if the schema 88 | // is serialized inline that will be used. 89 | // 90 | if (!populated_entire_set && 91 | !collection->source.sha256.empty() && 92 | (collection->source.serialized.has_value() || collection->source.url.has_value())) 93 | { 94 | auto cache_path = DuckDBHomeDirectory(context); 95 | 96 | // Populate the on-disk schema cache from the catalog while contents_url. 97 | auto curl = connection_pool.acquire(); 98 | AirportAPI::PopulateCatalogSchemaCacheFromURLorContent(curl, *collection, airport_catalog.internal_name(), cache_path); 99 | connection_pool.release(curl); 100 | } 101 | populated_entire_set = true; 102 | 103 | for (const auto &schema : collection->schemas) 104 | { 105 | CreateSchemaInfo info; 106 | 107 | if (schema.schema_name().empty()) 108 | { 109 | throw InvalidInputException("Airport: catalog '%s' contained a schema with an empty name", airport_catalog.internal_name()); 110 | } 111 | if (!(seen_schema_names.find(schema.schema_name()) == seen_schema_names.end())) 112 | { 113 | throw InvalidInputException("Airport: catalog '%s' contained two or more schemas named %s", airport_catalog.internal_name(), schema.schema_name().c_str()); 114 | } 115 | 116 | seen_schema_names.insert(schema.schema_name()); 117 | 118 | info.schema = schema.schema_name(); 119 | info.internal = IsInternalTable(schema.catalog_name(), schema.schema_name()); 120 | auto schema_entry = make_uniq(catalog, info, connection_pool, cache_path, schema); 121 | 122 | // Since these are DuckDB attributes, we need to copy them manually. 123 | schema_entry->comment = schema.comment(); 124 | for (auto &tag : schema.tags()) 125 | { 126 | schema_entry->tags[tag.first] = tag.second; 127 | } 128 | // printf("Creating schema %s\n", schema.schema_name.c_str()); 129 | CreateEntry(std::move(schema_entry)); 130 | } 131 | 132 | called_load_entries = true; 133 | } 134 | 135 | struct AirportCreateSchemaParameters 136 | { 137 | string catalog_name; 138 | string schema; 139 | 140 | std::optional comment; 141 | unordered_map tags; 142 | 143 | MSGPACK_DEFINE_MAP(catalog_name, schema, comment, tags) 144 | }; 145 | 146 | optional_ptr 147 | AirportSchemaSet::CreateSchema(ClientContext &context, CreateSchemaInfo &info) 148 | { 149 | auto &airport_catalog = catalog.Cast(); 150 | 151 | arrow::flight::FlightCallOptions call_options; 152 | 153 | airport_add_standard_headers(call_options, airport_catalog.attach_parameters()->location()); 154 | airport_add_authorization_header(call_options, airport_catalog.attach_parameters()->auth_token()); 155 | 156 | call_options.headers.emplace_back("airport-action-name", "create_schema"); 157 | 158 | auto flight_client = AirportAPI::FlightClientForLocation(airport_catalog.attach_parameters()->location()); 159 | 160 | AirportCreateSchemaParameters params; 161 | params.catalog_name = airport_catalog.internal_name(); 162 | params.schema = info.schema; 163 | if (!info.comment.IsNull()) 164 | { 165 | params.comment = info.comment.ToString(); 166 | } 167 | // for (auto &tag : info.tags()) 168 | // { 169 | // params.tags[tag.first] = tag.second; 170 | // } 171 | 172 | auto &server_location = airport_catalog.attach_parameters()->location(); 173 | 174 | AIRPORT_MSGPACK_ACTION_SINGLE_PARAMETER(action, "create_schema", params); 175 | 176 | std::unique_ptr action_results; 177 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(action_results, 178 | flight_client->DoAction(call_options, action), 179 | server_location, 180 | "airport_create_schema"); 181 | 182 | // We need to load the serialized scheam ifnormation from the server call. 183 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto msgpack_serialized_response, action_results->Next(), server_location, ""); 184 | 185 | if (msgpack_serialized_response == nullptr) 186 | { 187 | throw AirportFlightException(server_location, "Failed to obtain schema data from Arrow Flight create_schema RPC"); 188 | } 189 | 190 | const auto &body_buffer = msgpack_serialized_response.get()->body; 191 | AIRPORT_MSGPACK_UNPACK(AirportSerializedContentsWithSHA256Hash, contents, 192 | (*body_buffer), 193 | server_location, 194 | "File to parse msgpack encoded object from create_schema response"); 195 | 196 | AIRPORT_ARROW_ASSERT_OK_LOCATION(action_results->Drain(), server_location, ""); 197 | 198 | unordered_map empty; 199 | auto real_entry = AirportAPISchema( 200 | airport_catalog.internal_name(), 201 | info.schema, 202 | "", 203 | empty, 204 | contents); 205 | 206 | string cache_path = DuckDBHomeDirectory(context); 207 | 208 | auto schema_entry = make_uniq(catalog, info, connection_pool, cache_path, real_entry); 209 | 210 | return CreateEntry(std::move(schema_entry)); 211 | } 212 | 213 | } // namespace duckdb 214 | -------------------------------------------------------------------------------- /src/storage/airport_transaction.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "storage/airport_transaction.hpp" 3 | #include "storage/airport_catalog.hpp" 4 | #include "duckdb/parser/parsed_data/create_view_info.hpp" 5 | #include "duckdb/catalog/catalog_entry/index_catalog_entry.hpp" 6 | #include "duckdb/catalog/catalog_entry/view_catalog_entry.hpp" 7 | #include "airport_request_headers.hpp" 8 | #include "airport_macros.hpp" 9 | #include 10 | #include 11 | #include 12 | #include "airport_macros.hpp" 13 | 14 | namespace duckdb 15 | { 16 | 17 | AirportTransaction::AirportTransaction(AirportCatalog &airport_catalog, TransactionManager &manager, ClientContext &context) 18 | : Transaction(manager, context), access_mode_(airport_catalog.access_mode()), 19 | catalog_name(airport_catalog.internal_name()), 20 | attach_parameters(airport_catalog.attach_parameters()) 21 | { 22 | } 23 | 24 | AirportTransaction::~AirportTransaction() = default; 25 | 26 | struct GetTransactionIdentifierResult 27 | { 28 | std::optional identifier; 29 | MSGPACK_DEFINE_MAP(identifier) 30 | }; 31 | 32 | struct AirportCreateTransactionParameters 33 | { 34 | std::string catalog_name; 35 | MSGPACK_DEFINE_MAP(catalog_name) 36 | }; 37 | 38 | std::optional AirportTransaction::GetTransactionIdentifier() 39 | { 40 | auto &server_location = attach_parameters->location(); 41 | auto flight_client = AirportAPI::FlightClientForLocation(attach_parameters->location()); 42 | 43 | arrow::flight::FlightCallOptions call_options; 44 | airport_add_standard_headers(call_options, attach_parameters->location()); 45 | airport_add_authorization_header(call_options, attach_parameters->auth_token()); 46 | 47 | AirportCreateTransactionParameters params; 48 | params.catalog_name = catalog_name; 49 | AIRPORT_MSGPACK_ACTION_SINGLE_PARAMETER(action, "create_transaction", params); 50 | 51 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto action_results, 52 | flight_client->DoAction(call_options, action), 53 | server_location, 54 | "calling create_transaction action"); 55 | 56 | // The only item returned is a serialized flight info. 57 | AIRPORT_ASSIGN_OR_RAISE_LOCATION(auto result_buffer, 58 | action_results->Next(), 59 | server_location, 60 | "reading create_transaction action result"); 61 | 62 | AIRPORT_MSGPACK_UNPACK( 63 | GetTransactionIdentifierResult, result, 64 | (*(result_buffer->body)), 65 | server_location, 66 | "File to parse msgpack encoded create_transaction response"); 67 | 68 | AIRPORT_ARROW_ASSERT_OK_LOCATION(action_results->Drain(), server_location, ""); 69 | 70 | return result.identifier; 71 | } 72 | 73 | void AirportTransaction::Start() 74 | { 75 | transaction_state = AirportTransactionState::TRANSACTION_NOT_YET_STARTED; 76 | // Get an identifier from the server that will be passed as airport-transaction-id 77 | // in requests. 78 | identifier_ = GetTransactionIdentifier(); 79 | } 80 | void AirportTransaction::Commit() 81 | { 82 | if (transaction_state == AirportTransactionState::TRANSACTION_STARTED) 83 | { 84 | transaction_state = AirportTransactionState::TRANSACTION_FINISHED; 85 | } 86 | } 87 | void AirportTransaction::Rollback() 88 | { 89 | if (transaction_state == AirportTransactionState::TRANSACTION_STARTED) 90 | { 91 | transaction_state = AirportTransactionState::TRANSACTION_FINISHED; 92 | } 93 | } 94 | 95 | AirportTransaction &AirportTransaction::Get(ClientContext &context, Catalog &catalog) 96 | { 97 | return Transaction::Get(context, catalog).Cast(); 98 | } 99 | 100 | } // namespace duckdb 101 | -------------------------------------------------------------------------------- /src/storage/airport_transaction_manager.cpp: -------------------------------------------------------------------------------- 1 | #include "airport_extension.hpp" 2 | #include "storage/airport_transaction_manager.hpp" 3 | #include "duckdb/main/attached_database.hpp" 4 | 5 | namespace duckdb 6 | { 7 | 8 | AirportTransactionManager::AirportTransactionManager(AttachedDatabase &db_p, AirportCatalog &airport_catalog) 9 | : TransactionManager(db_p), airport_catalog(airport_catalog) 10 | { 11 | } 12 | 13 | Transaction &AirportTransactionManager::StartTransaction(ClientContext &context) 14 | { 15 | auto transaction = make_uniq(airport_catalog, *this, context); 16 | transaction->Start(); 17 | auto &result = *transaction; 18 | lock_guard l(transaction_lock); 19 | transactions[result] = std::move(transaction); 20 | return result; 21 | } 22 | 23 | ErrorData AirportTransactionManager::CommitTransaction(ClientContext &context, Transaction &transaction) 24 | { 25 | auto &airport_transaction = transaction.Cast(); 26 | airport_transaction.Commit(); 27 | lock_guard l(transaction_lock); 28 | transactions.erase(transaction); 29 | return ErrorData(); 30 | } 31 | 32 | void AirportTransactionManager::RollbackTransaction(Transaction &transaction) 33 | { 34 | auto &airport_transaction = transaction.Cast(); 35 | airport_transaction.Rollback(); 36 | lock_guard l(transaction_lock); 37 | transactions.erase(transaction); 38 | } 39 | 40 | void AirportTransactionManager::Checkpoint(ClientContext &context, bool force) 41 | { 42 | // auto &transaction = AirportTransaction::Get(context, db.GetCatalog()); 43 | // auto &db = transaction.GetConnection(); 44 | // db.Execute("CHECKPOINT"); 45 | } 46 | 47 | } // namespace duckdb 48 | -------------------------------------------------------------------------------- /src/storage/airport_update_parameterized.cpp: -------------------------------------------------------------------------------- 1 | #include "duckdb.hpp" 2 | #include "storage/airport_update_parameterized.hpp" 3 | #include "storage/airport_table_entry.hpp" 4 | #include "duckdb/planner/operator/logical_update.hpp" 5 | #include "storage/airport_catalog.hpp" 6 | #include "storage/airport_transaction.hpp" 7 | #include "duckdb/planner/expression/bound_reference_expression.hpp" 8 | #include "duckdb/common/arrow/arrow_appender.hpp" 9 | #include "duckdb/common/arrow/arrow_converter.hpp" 10 | #include "duckdb/common/types/uuid.hpp" 11 | #include "duckdb/function/table/arrow/arrow_duck_schema.hpp" 12 | #include "duckdb/function/table/arrow.hpp" 13 | #include "airport_macros.hpp" 14 | #include "airport_request_headers.hpp" 15 | #include "airport_flight_exception.hpp" 16 | #include "airport_secrets.hpp" 17 | 18 | #include "duckdb/common/arrow/schema_metadata.hpp" 19 | 20 | #include "airport_flight_stream.hpp" 21 | #include "airport_take_flight.hpp" 22 | #include "storage/airport_exchange.hpp" 23 | #include "duckdb/execution/operator/filter/physical_filter.hpp" 24 | #include "duckdb/execution/operator/scan/physical_table_scan.hpp" 25 | #include "duckdb/execution/operator/projection/physical_projection.hpp" 26 | 27 | namespace duckdb 28 | { 29 | 30 | static string ExtractFilters(PhysicalOperator &child, const string &statement) 31 | { 32 | // Check for FILTER operator type 33 | if (child.type == PhysicalOperatorType::FILTER) 34 | { 35 | auto &filter = child.Cast(); 36 | string result = ExtractFilters(child.children[0], statement); 37 | return result.empty() ? filter.expression->ToString() : result + " AND " + filter.expression->ToString(); 38 | } 39 | 40 | // Check for TABLE_SCAN operator type 41 | if (child.type == PhysicalOperatorType::TABLE_SCAN) 42 | { 43 | const auto &table_scan = child.Cast(); 44 | if (!table_scan.table_filters) 45 | { 46 | return {}; 47 | } 48 | throw NotImplementedException("Pushed down table filters not supported currently"); 49 | } 50 | 51 | if (child.type == PhysicalOperatorType::PROJECTION) 52 | { 53 | // auto &proj = child.Cast(); 54 | return ExtractFilters(child.children[0], statement); 55 | } 56 | 57 | // Handle unsupported operator types 58 | throw NotImplementedException( 59 | "Unsupported operator type %s in %s statement - only simple expressions " 60 | "(e.g. %s FROM tbl WHERE x=y) are supported in AirportParameterizedUpdate", 61 | PhysicalOperatorToString(child.type), statement, statement); 62 | } 63 | 64 | static std::pair ConstructUpdateStatement(LogicalUpdate &op, PhysicalOperator &child) 65 | { 66 | if (child.type != PhysicalOperatorType::PROJECTION) 67 | { 68 | throw NotImplementedException( 69 | "Airport Parameterized Update not supported - Expected the child of an update to be a projection"); 70 | } 71 | 72 | auto &proj = child.Cast(); 73 | string expressions; 74 | 75 | for (idx_t c = 0; c < op.columns.size(); ++c) 76 | { 77 | if (c > 0) 78 | { 79 | expressions += ", "; 80 | } 81 | 82 | const auto &col = op.table.GetColumn(op.table.GetColumns().PhysicalToLogical(op.columns[c])); 83 | expressions += col.GetName() + " = "; 84 | 85 | switch (op.expressions[c]->type) 86 | { 87 | case ExpressionType::VALUE_DEFAULT: 88 | expressions += "DEFAULT"; 89 | break; 90 | case ExpressionType::BOUND_REF: 91 | { 92 | const auto &ref = op.expressions[c]->Cast(); 93 | expressions += proj.select_list[ref.index]->ToString(); 94 | break; 95 | } 96 | default: 97 | throw NotImplementedException( 98 | "Airport Parameterized Update not supported - Expected a bound reference expression"); 99 | } 100 | } 101 | 102 | auto filters = ExtractFilters(child.children[0], "UPDATE"); 103 | return {expressions, filters}; 104 | } 105 | 106 | AirportUpdateParameterized::AirportUpdateParameterized(LogicalOperator &op, TableCatalogEntry &table, PhysicalOperator &plan) 107 | : PhysicalOperator(PhysicalOperatorType::EXTENSION, op.types, 1), table(table) 108 | { 109 | auto result = ConstructUpdateStatement(op.Cast(), plan); 110 | printf("Got expressions: %s\n", result.first.c_str()); 111 | printf("Got filters: %s\n", result.second.c_str()); 112 | } 113 | 114 | //===--------------------------------------------------------------------===// 115 | // States 116 | //===--------------------------------------------------------------------===// 117 | class AirportUpdateParameterizedGlobalState : public GlobalSinkState 118 | { 119 | public: 120 | explicit AirportUpdateParameterizedGlobalState() : affected_rows(0) 121 | { 122 | } 123 | 124 | idx_t affected_rows; 125 | }; 126 | 127 | unique_ptr AirportUpdateParameterized::GetGlobalSinkState(ClientContext &context) const 128 | { 129 | printf("AirportUpdateParameterized::GetGlobalSinkState\n"); 130 | return make_uniq(); 131 | } 132 | 133 | //===--------------------------------------------------------------------===// 134 | // Sink 135 | //===--------------------------------------------------------------------===// 136 | SinkResultType AirportUpdateParameterized::Sink(ExecutionContext &context, DataChunk &chunk, OperatorSinkInput &input) const 137 | { 138 | printf("AirportUpdateParameterized::Sink\n"); 139 | return SinkResultType::FINISHED; 140 | } 141 | 142 | //===--------------------------------------------------------------------===// 143 | // Finalize 144 | //===--------------------------------------------------------------------===// 145 | SinkFinalizeType AirportUpdateParameterized::Finalize(Pipeline &pipeline, Event &event, ClientContext &context, 146 | OperatorSinkFinalizeInput &input) const 147 | { 148 | auto &gstate = input.global_state.Cast(); 149 | gstate.affected_rows = 0; 150 | // auto &transaction = MySQLTransaction::Get(context, table.catalog); 151 | // auto &connection = transaction.GetConnection(); 152 | // auto result = connection.Query(query); 153 | // gstate.affected_rows = result->AffectedRows(); 154 | printf("AirportUpdateParameterized::Finalize\n"); 155 | return SinkFinalizeType::READY; 156 | } 157 | 158 | //===--------------------------------------------------------------------===// 159 | // GetData 160 | //===--------------------------------------------------------------------===// 161 | SourceResultType AirportUpdateParameterized::GetData(ExecutionContext &context, DataChunk &chunk, 162 | OperatorSourceInput &input) const 163 | { 164 | auto &insert_gstate = sink_state->Cast(); 165 | chunk.SetCardinality(1); 166 | chunk.SetValue(0, 0, Value::BIGINT(insert_gstate.affected_rows)); 167 | printf("AirportUpdateParameterized::GetData\n"); 168 | return SourceResultType::FINISHED; 169 | } 170 | 171 | //===--------------------------------------------------------------------===// 172 | // Helpers 173 | //===--------------------------------------------------------------------===// 174 | string AirportUpdateParameterized::GetName() const 175 | { 176 | return "AIRPORT_UPDATE_PARAMETERIZED"; 177 | } 178 | 179 | InsertionOrderPreservingMap AirportUpdateParameterized::ParamsToString() const 180 | { 181 | InsertionOrderPreservingMap result; 182 | result["Table Name"] = table.name; 183 | return result; 184 | } 185 | 186 | } // namespace duckdb 187 | -------------------------------------------------------------------------------- /test/README.md: -------------------------------------------------------------------------------- 1 | # Testing this extension 2 | This directory contains all the tests for this extension. The `sql` directory holds tests that are written as [SQLLogicTests](https://duckdb.org/dev/sqllogictest/intro.html). DuckDB aims to have most its tests in this format as SQL statements, so for the airport extension, this should probably be the goal too. 3 | 4 | The root makefile contains targets to build and run all of these tests. To run the SQLLogicTests: 5 | ```bash 6 | make test 7 | ``` 8 | or 9 | ```bash 10 | make test_debug 11 | ``` -------------------------------------------------------------------------------- /test/sql/airport.test: -------------------------------------------------------------------------------- 1 | # name: test/sql/airport.test 2 | # description: test airport extension 3 | # group: [airport] 4 | 5 | # Before we load the extension, this will fail 6 | statement error 7 | SELECT * from airport_flights('hello'); 8 | ---- 9 | Catalog Error: Table Function with name airport_flights does not exist! 10 | 11 | # Require statement will ensure this test is run with this extension loaded 12 | require airport 13 | -------------------------------------------------------------------------------- /vcpkg.json: -------------------------------------------------------------------------------- 1 | { 2 | "vcpkg-configuration": { 3 | "overlay-ports": [ 4 | "./extension-ci-tools/vcpkg_ports" 5 | ] 6 | }, 7 | "dependencies": [ 8 | { 9 | "name": "msgpack" 10 | }, 11 | { 12 | "name": "curl", 13 | "features": [ 14 | "http2", 15 | "openssl" 16 | ] 17 | }, 18 | { 19 | "name": "arrow", 20 | "default-features": false, 21 | "features": [ 22 | "filesystem", 23 | "flight" 24 | ] 25 | } 26 | ] 27 | } --------------------------------------------------------------------------------