ob-flight-extension 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,35 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *$py.class
4
+ *.so
5
+ dist/
6
+ build/
7
+ *.egg-info/
8
+ *.egg
9
+ .eggs/
10
+ .venv/
11
+ venv/
12
+ .env
13
+ *.db
14
+ *.duckdb
15
+ *.sqlite3
16
+ .mypy_cache/
17
+ .ruff_cache/
18
+ .pytest_cache/
19
+ htmlcov/
20
+ .coverage
21
+ .coverage.*
22
+ *.log
23
+ .DS_Store
24
+ .claude/
25
+ security/
26
+ infra/
27
+ site/
28
+ scripts/
29
+ .mcpregistry_github_token
30
+ .mcpregistry_registry_token
31
+ _*/
32
+ blog/
33
+ review*
34
+ design/
35
+ examples/test*
@@ -0,0 +1,285 @@
1
+ # ob-flight-extension — Arrow Flight SQL Server for orionbelt-api
2
+
3
+ ## Purpose
4
+
5
+ Adds Arrow Flight SQL protocol to the existing orionbelt-api FastAPI server.
6
+ Runs as a background thread in the same Python process — no extra container,
7
+ no HTTP hop between Flight and the OB compilation engine.
8
+
9
+ Target deployment: on-premise Docker, LAN access, enterprise data teams.
10
+ Client tools: DBeaver, Tableau (via Flight JDBC .jar), Power BI (via ODBC bridge).
11
+
12
+ ---
13
+
14
+ ## Integration into orionbelt-api
15
+
16
+ ### 1. Install as optional dependency
17
+
18
+ In orionbelt-api's pyproject.toml add:
19
+ ```toml
20
+ [project.optional-dependencies]
21
+ flight = ["ob-flight-extension>=0.1", "pyarrow>=16.0"]
22
+ ```
23
+
24
+ ### 2. Hook into lifespan
25
+
26
+ In orionbelt/api/app.py, modify create_app():
27
+
28
+ ```python
29
+ from contextlib import asynccontextmanager
30
+ import os
31
+
32
+ @asynccontextmanager
33
+ async def lifespan(app: FastAPI):
34
+ if os.getenv("FLIGHT_ENABLED", "false").lower() == "true":
35
+ from ob_flight.startup import start_flight_background
36
+ flight_thread = start_flight_background()
37
+ yield
38
+ if os.getenv("FLIGHT_ENABLED", "false").lower() == "true":
39
+ from ob_flight.startup import stop_flight_server
40
+ stop_flight_server()
41
+
42
+ app = FastAPI(lifespan=lifespan)
43
+ ```
44
+
45
+ ### 3. Docker port
46
+
47
+ Add port 8815 to docker-compose.yml (see ARCHITECTURE.md).
48
+
49
+ ---
50
+
51
+ ## Module Map
52
+
53
+ | File | Responsibility |
54
+ |---|---|
55
+ | startup.py | start_flight_background() — launches server in daemon thread |
56
+ | server.py | OBFlightServer(FlightServerBase) — pyarrow server subclass |
57
+ | handlers.py | GetFlightInfo, DoGet, GetSchema, ListFlights implementations |
58
+ | catalog.py | Model introspection → Flight catalog (tables = data objects) |
59
+ | auth.py | ServerAuthHandler — token or no-auth modes |
60
+ | converters.py | DB result rows → pyarrow RecordBatch |
61
+ | db_router.py | Routes to correct native connector based on model dialect |
62
+
63
+ ---
64
+
65
+ ## Key Implementation: GetFlightInfo
66
+
67
+ This is the entry point. DBeaver calls this first with the query text.
68
+
69
+ ```python
70
+ def get_flight_info(self, context, descriptor):
71
+ query_bytes = descriptor.command
72
+ query_str = query_bytes.decode("utf-8")
73
+
74
+ # Resolve model from descriptor path (= DBeaver "database" field)
75
+ model_id = self._resolve_model(descriptor.path)
76
+ model = self._model_cache[model_id]
77
+
78
+ # Compile if OBML, pass through if plain SQL
79
+ if is_obml(query_str):
80
+ obml = yaml.safe_load(query_str)
81
+ pipeline = CompilationPipeline()
82
+ result = pipeline.compile(obml, model, dialect=model.dialect)
83
+ sql = result.sql
84
+ else:
85
+ sql = query_str
86
+
87
+ # Store sql+model in a short-lived ticket (UUID key → in-memory dict)
88
+ ticket_id = str(uuid.uuid4())
89
+ self._pending[ticket_id] = (sql, model_id)
90
+ ticket = pa.flight.Ticket(ticket_id.encode())
91
+
92
+ # Return schema + ticket (schema requires a dry-run or model metadata)
93
+ schema = self._infer_schema(sql, model_id)
94
+ endpoint = pa.flight.FlightEndpoint(ticket, [])
95
+ return pa.flight.FlightInfo(schema, descriptor, [endpoint], -1, -1)
96
+ ```
97
+
98
+ ---
99
+
100
+ ## Key Implementation: DoGet
101
+
102
+ DBeaver calls this after GetFlightInfo to stream the actual data.
103
+
104
+ ```python
105
+ def do_get(self, context, ticket):
106
+ ticket_id = ticket.ticket.decode("utf-8")
107
+ sql, model_id = self._pending.pop(ticket_id)
108
+ model = self._model_cache[model_id]
109
+
110
+ # Execute on native connector
111
+ conn = db_router.connect(model)
112
+ cursor = conn.cursor()
113
+ cursor.execute(sql)
114
+
115
+ # Stream as Arrow record batches
116
+ schema = pa.Schema.from_pandas(...) # or from cursor.description
117
+ def record_batch_generator():
118
+ while True:
119
+ rows = cursor.fetchmany(1000)
120
+ if not rows:
121
+ break
122
+ yield pa.RecordBatch.from_pydict(
123
+ {col: [r[i] for r in rows]
124
+ for i, col in enumerate(col_names)},
125
+ schema=schema,
126
+ )
127
+ return pa.flight.RecordBatchStream(
128
+ pa.RecordBatchReader.from_batches(schema, record_batch_generator())
129
+ )
130
+ ```
131
+
132
+ ---
133
+
134
+ ## Key Implementation: ListFlights (DBeaver schema browser)
135
+
136
+ DBeaver calls ListFlights to populate the schema tree on the left panel.
137
+ Map OB model data objects → Flight "tables":
138
+
139
+ ```python
140
+ def list_flights(self, context, criteria):
141
+ for model_id, model in self._model_cache.items():
142
+ for obj_name, obj in model.data_objects.items():
143
+ descriptor = pa.flight.FlightDescriptor.for_path(
144
+ model_id, obj_name
145
+ )
146
+ schema = self._schema_for_object(obj)
147
+ yield pa.flight.FlightInfo(
148
+ schema, descriptor, [], -1, -1
149
+ )
150
+ ```
151
+
152
+ This makes DBeaver show the OB model's data objects as browsable tables,
153
+ with correct column types visible before any query is run.
154
+
155
+ ---
156
+
157
+ ## db_router.py — Vendor Routing
158
+
159
+ Routes execution to the correct native connector based on the model's
160
+ declared dialect or an explicit env override:
161
+
162
+ ```python
163
+ VENDOR_MAP = {
164
+ "snowflake": "ob_snowflake",
165
+ "postgres": "ob_postgres",
166
+ "clickhouse": "ob_clickhouse",
167
+ "dremio": "ob_dremio",
168
+ "databricks": "ob_databricks",
169
+ }
170
+
171
+ def connect(model: SemanticModel) -> Connection:
172
+ dialect = model.dialect or os.getenv("DB_VENDOR", "snowflake")
173
+ module = importlib.import_module(VENDOR_MAP[dialect])
174
+ return module.connect(**_creds_from_env(dialect))
175
+ ```
176
+
177
+ Credentials come from environment variables per vendor:
178
+ SNOWFLAKE_ACCOUNT, SNOWFLAKE_USER, SNOWFLAKE_PASSWORD, ...
179
+ POSTGRES_HOST, POSTGRES_PORT, POSTGRES_USER, ...
180
+ etc.
181
+
182
+ ---
183
+
184
+ ## Startup Thread
185
+
186
+ ```python
187
+ import threading
188
+ import pyarrow.flight as flight
189
+
190
+ _server: flight.FlightServerBase | None = None
191
+
192
+ def start_flight_background() -> threading.Thread:
193
+ global _server
194
+ port = int(os.getenv("FLIGHT_PORT", "8815"))
195
+ _server = OBFlightServer(location=f"grpc://0.0.0.0:{port}")
196
+ _server.init() # load models, init cache
197
+ thread = threading.Thread(
198
+ target=_server.serve,
199
+ name="ob-flight-server",
200
+ daemon=True, # dies with the main process
201
+ )
202
+ thread.start()
203
+ return thread
204
+
205
+ def stop_flight_server() -> None:
206
+ if _server:
207
+ _server.shutdown()
208
+ ```
209
+
210
+ daemon=True is critical — ensures the Flight server thread dies cleanly
211
+ when the FastAPI process exits or is killed by Docker.
212
+
213
+ ---
214
+
215
+ ## Environment Variables
216
+
217
+ | Variable | Default | Description |
218
+ |---|---|---|
219
+ | FLIGHT_ENABLED | false | Set to "true" to activate |
220
+ | FLIGHT_PORT | 8815 | gRPC listen port |
221
+ | FLIGHT_AUTH_MODE | none | "none" or "token" |
222
+ | FLIGHT_API_TOKEN | | Static token for auth mode "token" |
223
+ | FLIGHT_DEFAULT_MODEL | | Fallback model_id if path empty |
224
+ | FLIGHT_PRELOAD_MODELS | | Comma-sep .obml.yaml paths, loaded at startup |
225
+ | DB_VENDOR | snowflake | Default vendor for db_router |
226
+
227
+ ---
228
+
229
+ ## DBeaver Configuration (end-user guide)
230
+
231
+ 1. New Connection → Apache Arrow Flight SQL
232
+ 2. Host: <docker-host-ip> Port: 8815
233
+ 3. Database: <model_name_or_id> (e.g. "orionbelt_1")
234
+ 4. Authentication: leave empty (FLIGHT_AUTH_MODE=none on LAN)
235
+ 5. Test Connection → should show "Connected"
236
+ 6. In SQL editor, write either:
237
+ - Plain SQL: SELECT * FROM ORDERS LIMIT 10
238
+ - OBML YAML:
239
+ select:
240
+ dimensions: [Region]
241
+ measures: [Revenue]
242
+ 7. Schema browser shows OB model data objects as tables
243
+
244
+ ---
245
+
246
+ ## Tableau Configuration (end-user guide)
247
+
248
+ 1. Download flight-sql-jdbc-driver-*.jar from Apache Arrow releases
249
+ 2. Tableau → Connect → Other Databases (JDBC)
250
+ 3. URL: jdbc:arrow-flight-sql://<host>:8815?useEncryption=false
251
+ 4. Username/Password: leave empty or use token
252
+ 5. Custom SQL uses same OBML YAML or plain SQL
253
+
254
+ ---
255
+
256
+ ## Build Order
257
+
258
+ Session 1: converters.py + tests (rows → Arrow RecordBatch, no Flight needed)
259
+ Session 2: server.py + startup.py skeleton (pyarrow server, no handlers yet)
260
+ Session 3: handlers.py GetFlightInfo + DoGet (core query path)
261
+ Session 4: catalog.py + ListFlights (DBeaver schema browser)
262
+ Session 5: auth.py + db_router.py + integration test with DBeaver
263
+
264
+ ---
265
+
266
+ ## Dependencies
267
+
268
+ ```toml
269
+ [project.dependencies]
270
+ pyarrow = ">=16.0"
271
+ pyyaml = ">=6.0"
272
+
273
+ [project.optional-dependencies]
274
+ snowflake = ["snowflake-connector-python>=3.0"]
275
+ postgres = ["adbc-driver-postgresql>=1.0"]
276
+ clickhouse = ["clickhouse-connect>=0.7"]
277
+ dremio = ["pyarrow>=16.0"] # Dremio uses Arrow Flight natively
278
+ databricks = ["databricks-sql-connector>=3.0"]
279
+ all = [
280
+ "snowflake-connector-python>=3.0",
281
+ "adbc-driver-postgresql>=1.0",
282
+ "clickhouse-connect>=0.7",
283
+ "databricks-sql-connector>=3.0",
284
+ ]
285
+ ```
@@ -0,0 +1,70 @@
1
+ Business Source License 1.1
2
+
3
+ Parameters
4
+
5
+ Licensor: RALFORION d.o.o.
6
+ Licensed Work: OrionBelt Semantic Layer 1.0
7
+ The Licensed Work is (c) 2025 RALFORION d.o.o.
8
+ Additional Use Grant: You may use the Licensed Work for any purpose,
9
+ including production use, PROVIDED THAT you may
10
+ not embed OrionBelt in a commercial product or
11
+ provide it as part of a commercial service
12
+ without purchasing a commercial license from
13
+ the Licensor. Contact licensing@ralforion.com
14
+ for commercial licensing inquiries.
15
+ Change Date: 2030-03-16
16
+ Change License: Apache License, Version 2.0
17
+
18
+ -----------------------------------------------------------------------------
19
+
20
+ Business Source License 1.1
21
+
22
+ License text copyright (c) 2017 MariaDB Corporation Ab, All Rights Reserved.
23
+ "Business Source License" is a trademark of MariaDB Corporation Ab.
24
+
25
+ Terms
26
+
27
+ The Licensor hereby grants you the right to copy, modify, create derivative
28
+ works, redistribute, and make non-production use of the Licensed Work. The
29
+ Licensor may make an Additional Use Grant, above, permitting limited
30
+ production use.
31
+
32
+ Effective on the Change Date, or the fourth anniversary of the first publicly
33
+ available distribution of a specific version of the Licensed Work under this
34
+ License, whichever comes first, the Licensor hereby grants you rights under
35
+ the terms of the Change License, and the rights granted in the paragraph
36
+ above terminate.
37
+
38
+ If your use of the Licensed Work does not comply with the requirements
39
+ currently in effect as described in this License, you must purchase a
40
+ commercial license from the Licensor, its affiliated entities, or authorized
41
+ resellers, or you must refrain from using the Licensed Work.
42
+
43
+ All copies of the original and modified Licensed Work, and derivative works
44
+ of the Licensed Work, are subject to this License. This License applies
45
+ separately for each version of the Licensed Work and the Change Date may vary
46
+ for each version of the Licensed Work released by Licensor.
47
+
48
+ You must conspicuously display this License on each original or modified copy
49
+ of the Licensed Work. If you receive the Licensed Work in original or
50
+ modified form from a third party, the terms and conditions set forth in this
51
+ License apply to your use of that work.
52
+
53
+ Any use of the Licensed Work in violation of this License will automatically
54
+ terminate your rights under this License for the current and all other
55
+ versions of the Licensed Work.
56
+
57
+ This License does not grant you any right in any trademark or logo of
58
+ Licensor or its affiliates (provided that you may use a trademark or logo of
59
+ Licensor as expressly required by this License).
60
+
61
+ TO THE EXTENT PERMITTED BY APPLICABLE LAW, THE LICENSED WORK IS PROVIDED ON
62
+ AN "AS IS" BASIS. LICENSOR HEREBY DISCLAIMS ALL WARRANTIES AND CONDITIONS,
63
+ EXPRESS OR IMPLIED, INCLUDING (WITHOUT LIMITATION) WARRANTIES OF
64
+ MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, NON-INFRINGEMENT, AND
65
+ TITLE.
66
+
67
+ MariaDB hereby grants you permission to use this License's text to license
68
+ your works, and to refer to it using the trademark "Business Source License",
69
+ as long as you comply with the Covenants of Licensor of MariaDB's Business Source
70
+ License Trademark Policy available at https://mariadb.com/bsl11/.
@@ -0,0 +1,53 @@
1
+ Metadata-Version: 2.4
2
+ Name: ob-flight-extension
3
+ Version: 0.1.0
4
+ Summary: Arrow Flight SQL server extension for orionbelt-api (DBeaver, Tableau, Power BI)
5
+ Author-email: Ralf Becher <info@orionbelt.ai>
6
+ License: BSL-1.1
7
+ License-File: LICENSE
8
+ Keywords: arrow,dbeaver,flight-sql,orionbelt,semantic-layer,tableau
9
+ Requires-Python: >=3.12
10
+ Requires-Dist: ob-driver-core>=0.1
11
+ Requires-Dist: pyarrow>=16.0
12
+ Requires-Dist: pyyaml>=6.0
13
+ Provides-Extra: all
14
+ Requires-Dist: ob-clickhouse>=0.1; extra == 'all'
15
+ Requires-Dist: ob-databricks>=0.1; extra == 'all'
16
+ Requires-Dist: ob-dremio>=0.1; extra == 'all'
17
+ Requires-Dist: ob-duckdb>=0.1; extra == 'all'
18
+ Requires-Dist: ob-postgres>=0.1; extra == 'all'
19
+ Requires-Dist: ob-snowflake>=0.1; extra == 'all'
20
+ Provides-Extra: clickhouse
21
+ Requires-Dist: ob-clickhouse>=0.1; extra == 'clickhouse'
22
+ Provides-Extra: databricks
23
+ Requires-Dist: ob-databricks>=0.1; extra == 'databricks'
24
+ Provides-Extra: dev
25
+ Requires-Dist: mypy>=1.10; extra == 'dev'
26
+ Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
27
+ Requires-Dist: pytest>=8.0; extra == 'dev'
28
+ Requires-Dist: ruff>=0.4; extra == 'dev'
29
+ Provides-Extra: dremio
30
+ Requires-Dist: ob-dremio>=0.1; extra == 'dremio'
31
+ Provides-Extra: duckdb
32
+ Requires-Dist: ob-duckdb>=0.1; extra == 'duckdb'
33
+ Provides-Extra: postgres
34
+ Requires-Dist: ob-postgres>=0.1; extra == 'postgres'
35
+ Provides-Extra: snowflake
36
+ Requires-Dist: ob-snowflake>=0.1; extra == 'snowflake'
37
+ Description-Content-Type: text/markdown
38
+
39
+ # ob-flight-extension
40
+
41
+ Arrow Flight SQL server extension for orionbelt-api (DBeaver, Tableau, Power BI)
42
+
43
+ Part of the [OrionBelt Semantic Layer](https://github.com/ralfbecher/orionbelt-semantic-layer) project.
44
+
45
+ ## Installation
46
+
47
+ ```bash
48
+ pip install ob-flight-extension
49
+ ```
50
+
51
+ ## License
52
+
53
+ BSL-1.1 — see [LICENSE](LICENSE) for details.
@@ -0,0 +1,15 @@
1
+ # ob-flight-extension
2
+
3
+ Arrow Flight SQL server extension for orionbelt-api (DBeaver, Tableau, Power BI)
4
+
5
+ Part of the [OrionBelt Semantic Layer](https://github.com/ralfbecher/orionbelt-semantic-layer) project.
6
+
7
+ ## Installation
8
+
9
+ ```bash
10
+ pip install ob-flight-extension
11
+ ```
12
+
13
+ ## License
14
+
15
+ BSL-1.1 — see [LICENSE](LICENSE) for details.
@@ -0,0 +1,39 @@
1
+ # docker-compose.yml — add to your existing orionbelt-api service
2
+ # Only two lines change: expose port 8815 and add FLIGHT_ENABLED=true
3
+
4
+ services:
5
+ orionbelt:
6
+ image: orionbelt-api:latest
7
+ build: .
8
+ ports:
9
+ - "8000:8000" # REST API (existing)
10
+ - "8815:8815" # Arrow Flight SQL (NEW)
11
+ environment:
12
+ # --- existing env vars unchanged ---
13
+ LOG_LEVEL: INFO
14
+ API_SERVER_HOST: "0.0.0.0"
15
+ API_SERVER_PORT: "8000"
16
+ SESSION_TTL_SECONDS: "1800"
17
+
18
+ # --- Arrow Flight SQL (NEW) ---
19
+ FLIGHT_ENABLED: "true"
20
+ FLIGHT_PORT: "8815"
21
+ FLIGHT_AUTH_MODE: "none" # use "token" for production
22
+ # FLIGHT_API_TOKEN: "changeme" # required when AUTH_MODE=token
23
+ FLIGHT_DEFAULT_MODEL: "" # fallback if DBeaver sends no path
24
+ FLIGHT_PRELOAD_MODELS: "" # comma-sep paths, loaded at startup
25
+
26
+ # --- DB credentials (used by db_router.py) ---
27
+ DB_VENDOR: snowflake # default vendor
28
+ SNOWFLAKE_ACCOUNT: "xy12345.eu-west-1"
29
+ SNOWFLAKE_USER: "your_user"
30
+ SNOWFLAKE_PASSWORD: "your_password"
31
+ SNOWFLAKE_DATABASE: "MY_DB"
32
+ SNOWFLAKE_SCHEMA: "PUBLIC"
33
+ SNOWFLAKE_WAREHOUSE: "COMPUTE_WH"
34
+ # POSTGRES_HOST: localhost
35
+ # POSTGRES_PORT: 5432
36
+ # POSTGRES_USER: postgres
37
+ # POSTGRES_PASSWORD: secret
38
+ # POSTGRES_DBNAME: mydb
39
+ restart: unless-stopped
@@ -0,0 +1,58 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "ob-flight-extension"
7
+ version = "0.1.0"
8
+ description = "Arrow Flight SQL server extension for orionbelt-api (DBeaver, Tableau, Power BI)"
9
+ license = { text = "BSL-1.1" }
10
+ readme = "README.md"
11
+ authors = [{ name = "Ralf Becher", email = "info@orionbelt.ai" }]
12
+ keywords = ["arrow", "flight-sql", "dbeaver", "tableau", "orionbelt", "semantic-layer"]
13
+ requires-python = ">=3.12"
14
+ dependencies = [
15
+ "pyarrow>=16.0",
16
+ "pyyaml>=6.0",
17
+ "ob-driver-core>=0.1",
18
+ ]
19
+
20
+ [project.optional-dependencies]
21
+ duckdb = ["ob-duckdb>=0.1"]
22
+ snowflake = ["ob-snowflake>=0.1"]
23
+ postgres = ["ob-postgres>=0.1"]
24
+ clickhouse = ["ob-clickhouse>=0.1"]
25
+ dremio = ["ob-dremio>=0.1"]
26
+ databricks = ["ob-databricks>=0.1"]
27
+ all = [
28
+ "ob-duckdb>=0.1",
29
+ "ob-snowflake>=0.1",
30
+ "ob-postgres>=0.1",
31
+ "ob-clickhouse>=0.1",
32
+ "ob-dremio>=0.1",
33
+ "ob-databricks>=0.1",
34
+ ]
35
+ dev = [
36
+ "pytest>=8.0",
37
+ "pytest-asyncio>=0.23",
38
+ "mypy>=1.10",
39
+ "ruff>=0.4",
40
+ ]
41
+
42
+ [tool.hatch.build.targets.wheel]
43
+ packages = ["src/ob_flight"]
44
+
45
+ [tool.ruff]
46
+ line-length = 100
47
+ target-version = "py312"
48
+
49
+ [tool.mypy]
50
+ python_version = "3.12"
51
+ strict = true
52
+ files = ["src/ob_flight"]
53
+
54
+ [tool.pytest.ini_options]
55
+ asyncio_mode = "auto"
56
+ testpaths = ["tests"]
57
+ markers = ["integration: requires running orionbelt-api + DBeaver"]
58
+ addopts = "-m 'not integration'"
@@ -0,0 +1,36 @@
1
+ """ob-flight-extension — Arrow Flight SQL server for OrionBelt Semantic Layer.
2
+
3
+ Adds Arrow Flight SQL protocol to the orionbelt-api FastAPI server.
4
+ Runs as a daemon thread in the same process — no HTTP hop for compilation.
5
+
6
+ Usage: Set FLIGHT_ENABLED=true when starting orionbelt-api.
7
+
8
+ Note: server and startup are NOT imported at package level to avoid
9
+ pulling in pyarrow.flight (heavy gRPC init) when only db_router is needed.
10
+ Import them explicitly: ``from ob_flight.startup import start_flight_background``
11
+ """
12
+
13
+ from __future__ import annotations
14
+
15
+ __all__ = [
16
+ "OBFlightServer",
17
+ "start_flight_background",
18
+ "stop_flight_server",
19
+ ]
20
+
21
+
22
+ def __getattr__(name: str) -> object:
23
+ """Lazy imports for server and startup — avoids pyarrow.flight at import time."""
24
+ if name == "OBFlightServer":
25
+ from ob_flight.server import OBFlightServer
26
+
27
+ return OBFlightServer
28
+ if name == "start_flight_background":
29
+ from ob_flight.startup import start_flight_background
30
+
31
+ return start_flight_background
32
+ if name == "stop_flight_server":
33
+ from ob_flight.startup import stop_flight_server
34
+
35
+ return stop_flight_server
36
+ raise AttributeError(f"module 'ob_flight' has no attribute {name!r}")
@@ -0,0 +1,55 @@
1
+ """Authentication handlers for the Flight SQL server."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import os
6
+ from typing import Any
7
+
8
+ import pyarrow.flight as flight
9
+
10
+
11
+ class NoopAuthHandler(flight.ServerAuthHandler):
12
+ """No authentication — accept all connections."""
13
+
14
+ def authenticate(self, outgoing: Any, incoming: Any) -> None:
15
+ """Accept all connections without authentication."""
16
+
17
+ def is_valid(self, token: bytes) -> str:
18
+ """All tokens are valid — return empty peer identity."""
19
+ return ""
20
+
21
+
22
+ class TokenAuthHandler(flight.ServerAuthHandler):
23
+ """Simple static bearer token authentication."""
24
+
25
+ def __init__(self, token: str) -> None:
26
+ super().__init__()
27
+ self._token = token.encode("utf-8")
28
+
29
+ def authenticate(self, outgoing: Any, incoming: Any) -> None:
30
+ """Validate the token from the client."""
31
+ buf = incoming.read()
32
+ if buf != self._token:
33
+ raise flight.FlightUnauthenticatedError("Invalid token")
34
+ outgoing.write(self._token)
35
+
36
+ def is_valid(self, token: bytes) -> str:
37
+ """Check if the token matches."""
38
+ if token != self._token:
39
+ raise flight.FlightUnauthenticatedError("Invalid token")
40
+ return "authenticated"
41
+
42
+
43
+ def create_auth_handler() -> flight.ServerAuthHandler:
44
+ """Create an auth handler based on environment variables.
45
+
46
+ FLIGHT_AUTH_MODE=none (default) -> NoopAuthHandler
47
+ FLIGHT_AUTH_MODE=token -> TokenAuthHandler(FLIGHT_API_TOKEN)
48
+ """
49
+ mode = os.getenv("FLIGHT_AUTH_MODE", "none").lower()
50
+ if mode == "token":
51
+ token = os.getenv("FLIGHT_API_TOKEN", "")
52
+ if not token:
53
+ raise ValueError("FLIGHT_AUTH_MODE=token requires FLIGHT_API_TOKEN to be set")
54
+ return TokenAuthHandler(token)
55
+ return NoopAuthHandler()