Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion ibis-server/Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,8 @@ RUN apt-get -y install libpq-dev \
ENV VIRTUAL_ENV=/app/.venv \
PATH="/app/.venv/bin:$PATH" \
PYTHONPATH="/app" \
REMOTE_FUNCTION_LIST_PATH=/resources/function_list
REMOTE_FUNCTION_LIST_PATH=/resources/function_list \
REMOTE_WHITE_FUNCTION_LIST_PATH=/resources/white_function_list

# Set working directory before copying files
WORKDIR /app
Expand Down
22 changes: 22 additions & 0 deletions ibis-server/app/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,9 @@ def __init__(self):
load_dotenv(override=True)
self.wren_engine_endpoint = os.getenv("WREN_ENGINE_ENDPOINT")
self.remote_function_list_path = os.getenv("REMOTE_FUNCTION_LIST_PATH")
self.remote_white_function_list_path = os.getenv(
"REMOTE_WHITE_FUNCTION_LIST_PATH"
)
self.app_timeout_seconds = int(os.getenv("APP_TIMEOUT_SECONDS", "240"))
self.diagnose = False
self.init_logger()
Expand Down Expand Up @@ -66,9 +69,28 @@ def get_remote_function_list_path(self, data_source: str) -> str:
raise ValueError("Invalid data source path")
return path if os.path.isfile(path) else None

def get_remote_white_function_list_path(self, data_source: str) -> str:
if not self.remote_white_function_list_path:
return None
else:
base_path = os.path.normpath(self.remote_white_function_list_path)
path = os.path.normpath(os.path.join(base_path, f"{data_source}.csv"))
if not path.startswith(base_path):
raise ValueError("Invalid data source path")
return path if os.path.isfile(path) else None

def set_remote_function_list_path(self, path: str | None):
self.remote_function_list_path = path

def set_remote_white_function_list_path(self, path: str | None):
self.remote_white_function_list_path = path

def get_data_source_is_white_list(self, data_source: str) -> bool:
if not self.remote_white_function_list_path:
return False

return data_source in {"bigquery"}


config = Config()

Expand Down
17 changes: 16 additions & 1 deletion ibis-server/app/routers/v3/connector.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from typing import Annotated

import duckdb
from fastapi import APIRouter, Depends, Query, Response
from fastapi.responses import ORJSONResponse
from loguru import logger
Expand Down Expand Up @@ -372,8 +373,22 @@ def functions(
name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
):
file_path = get_config().get_remote_function_list_path(data_source)
white_function_list_path = get_config().get_remote_white_function_list_path(
data_source
)
is_white_list = get_config().get_data_source_is_white_list(data_source)
session_context = get_session_context(None, file_path)
func_list = [f.to_dict() for f in session_context.get_available_functions()]
if is_white_list:
func_list = (
duckdb.read_csv(
white_function_list_path,
header=True,
)
.to_df()
.to_dict("records")
)
else:
func_list = [f.to_dict() for f in session_context.get_available_functions()]
return ORJSONResponse(func_list)


Expand Down
Loading
Loading