From f3660d3c0d2b535c7af736de074800fef44df9e8 Mon Sep 17 00:00:00 2001 From: Gagan Date: Sat, 2 Nov 2024 21:56:36 -0600 Subject: [PATCH 1/7] apache druid datasource --- ibis-server/app/model/__init__.py | 9 +++++++-- ibis-server/app/model/data_source.py | 12 ++++++++++++ ibis-server/pyproject.toml | 1 + 3 files changed, 20 insertions(+), 2 deletions(-) diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index 0bd588d2a..40daf7e66 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -45,7 +45,9 @@ class QueryPostgresDTO(QueryDTO): class QuerySnowflakeDTO(QueryDTO): connection_info: SnowflakeConnectionInfo = connection_info_field - + +class QueryDruidDTO(QueryDTO): + connection_info: DruidConnectionInfo = connection_info_field class QueryTrinoDTO(QueryDTO): connection_info: ConnectionUrl | TrinoConnectionInfo = connection_info_field @@ -75,6 +77,10 @@ class ClickHouseConnectionInfo(BaseModel): user: SecretStr password: SecretStr +class DruidConnectionInfo(BaseModel): + host: SecretStr = Field(examples=["localhost"]) + port: SecretStr = Field(examples=[8082]) + path: str class MSSqlConnectionInfo(BaseModel): host: SecretStr @@ -129,7 +135,6 @@ class TrinoConnectionInfo(BaseModel): user: SecretStr | None = None password: SecretStr | None = None - ConnectionInfo = ( BigQueryConnectionInfo | CannerConnectionInfo diff --git a/ibis-server/app/model/data_source.py b/ibis-server/app/model/data_source.py index 43a2fdb60..761dd3899 100644 --- a/ibis-server/app/model/data_source.py +++ b/ibis-server/app/model/data_source.py @@ -13,6 +13,7 @@ CannerConnectionInfo, ClickHouseConnectionInfo, ConnectionInfo, + DruidConnectionInfo, MSSqlConnectionInfo, MySqlConnectionInfo, PostgresConnectionInfo, @@ -20,6 +21,7 @@ QueryCannerDTO, QueryClickHouseDTO, QueryDTO, + QueryDruidDTO, QueryMSSqlDTO, QueryMySqlDTO, QueryPostgresDTO, @@ -39,6 +41,7 @@ class DataSource(StrEnum): postgres = auto() snowflake = auto() trino = auto() + druid = auto() def get_connection(self, info: ConnectionInfo) -> BaseBackend: try: @@ -62,6 +65,7 @@ class DataSourceExtension(Enum): postgres = QueryPostgresDTO snowflake = QuerySnowflakeDTO trino = QueryTrinoDTO + druid = QueryDruidDTO def __init__(self, dto: QueryDTO): self.dto = dto @@ -107,6 +111,14 @@ def get_clickhouse_connection(info: ClickHouseConnectionInfo) -> BaseBackend: user=info.user.get_secret_value(), password=info.password.get_secret_value(), ) + + @staticmethod + def get_druid_connection(info: DruidConnectionInfo) -> BaseBackend: + return ibis.druid.connect( + host=info.host.get_secret_value(), + port=int(info.port.get_secret_value()), + path=info.path, + ) @staticmethod def get_mssql_connection(info: MSSqlConnectionInfo) -> BaseBackend: diff --git a/ibis-server/pyproject.toml b/ibis-server/pyproject.toml index 2c85b0e58..55f64b50f 100644 --- a/ibis-server/pyproject.toml +++ b/ibis-server/pyproject.toml @@ -51,6 +51,7 @@ markers = [ "bigquery: mark a test as a bigquery test", "canner: mark a test as a canner test", "clickhouse: mark a test as a clickhouse test", + "druid: mark a test as a druid test", "mssql: mark a test as a mssql test", "mysql: mark a test as a mysql test", "postgres: mark a test as a postgres test", From 029f01d56d77bfb4ba1e742f07f08c20c7bea91b Mon Sep 17 00:00:00 2001 From: Gagan Date: Sun, 3 Nov 2024 20:49:48 -0700 Subject: [PATCH 2/7] test cases --- .../tests/routers/v2/connector/test_druid.py | 92 +++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 ibis-server/tests/routers/v2/connector/test_druid.py diff --git a/ibis-server/tests/routers/v2/connector/test_druid.py b/ibis-server/tests/routers/v2/connector/test_druid.py new file mode 100644 index 000000000..d2e749144 --- /dev/null +++ b/ibis-server/tests/routers/v2/connector/test_druid.py @@ -0,0 +1,92 @@ +import base64 + +import orjson +import pandas as pd +import pytest +import sqlalchemy +from sqlalchemy import text +from testcontainers.core.container import DockerContainer +from testcontainers.core.image import DockerImage + +from tests.confest import file_path + +pytestmark = pytest.mark.druid + +base_url = "/v2/connector/druid" + +manifest = { + "catalog": "my_catalog", + "schema": "my_schema", + "models": [ + { + "name": "Orders", + "refSql": "select * from public.orders", + "columns": [ + {"name": "orderkey", "expression": "o_orderkey", "type": "integer"}, + {"name": "custkey", "expression": "o_custkey", "type": "integer"}, + { + "name": "orderstatus", + "expression": "o_orderstatus", + "type": "varchar", + }, + { + "name": "totalprice", + "expression": "o_totalprice", + "type": "float", + }, + {"name": "orderdate", "expression": "o_orderdate", "type": "date"}, + { + "name": "order_cust_key", + "expression": "concat(o_orderkey, '_', o_custkey)", + "type": "varchar", + }, + { + "name": "timestamp", + "expression": "cast('2024-01-01T23:59:59' as timestamp)", + "type": "timestamp", + }, + { + "name": "timestamptz", + "expression": "cast('2024-01-01T23:59:59' as timestamp with time zone)", + "type": "timestamp", + }, + { + "name": "test_null_time", + "expression": "cast(NULL as timestamp)", + "type": "timestamp", + }, + { + "name": "bytea_column", + "expression": "cast('abc' as bytea)", + "type": "bytea", + }, + ], + "primaryKey": "orderkey", + }, + ], +} + + +@pytest.fixture +def manifest_str(): + return base64.b64encode(orjson.dumps(manifest)).decode("utf-8") + + +@pytest.fixture(scope="module") +def docker(request) -> DockerContainer: + with DockerImage(tag="apache/druid:latest") as image: + with DockerContainer(str(image)) as container: + druid = container.start() + engine = sqlalchemy.create_engine(druid.get_connection_url()) + pd.read_parquet(file_path("resource/tpch/data/orders.parquet")).to_sql( + "orders", engine, index=False + ) + with engine.begin() as conn: + conn.execute( + text("COMMENT ON TABLE orders IS 'This is a table comment'") + ) + conn.execute( + text("COMMENT ON COLUMN orders.o_comment IS 'This is a comment'") + ) + request.addfinalizer(druid.stop) + return druid From 8770101c4e650bdf277c9633a9bd8fd0c119e8e7 Mon Sep 17 00:00:00 2001 From: Gagan Date: Tue, 5 Nov 2024 01:23:42 -0700 Subject: [PATCH 3/7] test cases --- .../tests/resource/druid/docker-compose.yml | 134 +++++++++++++++ ibis-server/tests/resource/druid/environment | 53 ++++++ .../tests/routers/v2/connector/test_druid.py | 160 +++++++++--------- 3 files changed, 267 insertions(+), 80 deletions(-) create mode 100644 ibis-server/tests/resource/druid/docker-compose.yml create mode 100644 ibis-server/tests/resource/druid/environment diff --git a/ibis-server/tests/resource/druid/docker-compose.yml b/ibis-server/tests/resource/druid/docker-compose.yml new file mode 100644 index 000000000..0e18a600f --- /dev/null +++ b/ibis-server/tests/resource/druid/docker-compose.yml @@ -0,0 +1,134 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# +version: "2.2" + +volumes: + metadata_data: {} + middle_var: {} + historical_var: {} + broker_var: {} + coordinator_var: {} + router_var: {} + druid_shared: {} + + +services: + postgres: + container_name: postgres + image: postgres:latest + ports: + - "5432:5432" + volumes: + - metadata_data:/var/lib/postgresql/data + environment: + - POSTGRES_PASSWORD=FoolishPassword + - POSTGRES_USER=druid + - POSTGRES_DB=druid + + # Need 3.5 or later for container nodes + zookeeper: + container_name: zookeeper + image: zookeeper:3.5.10 + ports: + - "2181:2181" + environment: + - ZOO_MY_ID=1 + + coordinator: + image: apache/druid:31.0.0 + container_name: coordinator + volumes: + - druid_shared:/opt/shared + - coordinator_var:/opt/druid/var + depends_on: + - zookeeper + - postgres + ports: + - "8081:8081" + command: + - coordinator + env_file: + - environment + + broker: + image: apache/druid:31.0.0 + container_name: broker + volumes: + - broker_var:/opt/druid/var + depends_on: + - zookeeper + - postgres + - coordinator + ports: + - "8082:8082" + command: + - broker + env_file: + - environment + + historical: + image: apache/druid:31.0.0 + container_name: historical + volumes: + - druid_shared:/opt/shared + - historical_var:/opt/druid/var + depends_on: + - zookeeper + - postgres + - coordinator + ports: + - "8083:8083" + command: + - historical + env_file: + - environment + + middlemanager: + image: apache/druid:31.0.0 + container_name: middlemanager + volumes: + - druid_shared:/opt/shared + - middle_var:/opt/druid/var + depends_on: + - zookeeper + - postgres + - coordinator + ports: + - "8091:8091" + - "8100-8105:8100-8105" + command: + - middleManager + env_file: + - environment + + router: + image: apache/druid:31.0.0 + container_name: router + volumes: + - router_var:/opt/druid/var + depends_on: + - zookeeper + - postgres + - coordinator + ports: + - "8888:8888" + command: + - router + env_file: + - environment diff --git a/ibis-server/tests/resource/druid/environment b/ibis-server/tests/resource/druid/environment new file mode 100644 index 000000000..e223b5f4c --- /dev/null +++ b/ibis-server/tests/resource/druid/environment @@ -0,0 +1,53 @@ +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +# + +# Java tuning +#DRUID_XMX=1g +#DRUID_XMS=1g +#DRUID_MAXNEWSIZE=250m +#DRUID_NEWSIZE=250m +#DRUID_MAXDIRECTMEMORYSIZE=6172m +DRUID_SINGLE_NODE_CONF=micro-quickstart + +druid_emitter_logging_logLevel=debug + +druid_extensions_loadList=["druid-histogram", "druid-datasketches", "druid-lookups-cached-global", "postgresql-metadata-storage", "druid-multi-stage-query", "druid-parquet-extensions"] + +druid_zk_service_host=zookeeper + +druid_metadata_storage_host= +druid_metadata_storage_type=postgresql +druid_metadata_storage_connector_connectURI=jdbc:postgresql://postgres:5432/druid +druid_metadata_storage_connector_user=druid +druid_metadata_storage_connector_password=FoolishPassword + +druid_coordinator_balancer_strategy=cachingCost + +druid_indexer_runner_javaOptsArray=["-server", "-Xmx1g", "-Xms1g", "-XX:MaxDirectMemorySize=3g", "-Duser.timezone=UTC", "-Dfile.encoding=UTF-8", "-Djava.util.logging.manager=org.apache.logging.log4j.jul.LogManager"] +druid_indexer_fork_property_druid_processing_buffer_sizeBytes=256MiB + +druid_storage_type=local +druid_storage_storageDirectory=/opt/shared/segments +druid_indexer_logs_type=file +druid_indexer_logs_directory=/opt/shared/indexing-logs + +druid_processing_numThreads=2 +druid_processing_numMergeBuffers=2 + +DRUID_LOG4J= diff --git a/ibis-server/tests/routers/v2/connector/test_druid.py b/ibis-server/tests/routers/v2/connector/test_druid.py index d2e749144..3dc8ca1b0 100644 --- a/ibis-server/tests/routers/v2/connector/test_druid.py +++ b/ibis-server/tests/routers/v2/connector/test_druid.py @@ -1,92 +1,92 @@ -import base64 +import json +import time -import orjson -import pandas as pd import pytest -import sqlalchemy -from sqlalchemy import text -from testcontainers.core.container import DockerContainer -from testcontainers.core.image import DockerImage - -from tests.confest import file_path +import requests +from testcontainers.compose import DockerCompose pytestmark = pytest.mark.druid base_url = "/v2/connector/druid" -manifest = { - "catalog": "my_catalog", - "schema": "my_schema", - "models": [ - { - "name": "Orders", - "refSql": "select * from public.orders", - "columns": [ - {"name": "orderkey", "expression": "o_orderkey", "type": "integer"}, - {"name": "custkey", "expression": "o_custkey", "type": "integer"}, - { - "name": "orderstatus", - "expression": "o_orderstatus", - "type": "varchar", - }, - { - "name": "totalprice", - "expression": "o_totalprice", - "type": "float", - }, - {"name": "orderdate", "expression": "o_orderdate", "type": "date"}, - { - "name": "order_cust_key", - "expression": "concat(o_orderkey, '_', o_custkey)", - "type": "varchar", - }, - { - "name": "timestamp", - "expression": "cast('2024-01-01T23:59:59' as timestamp)", - "type": "timestamp", - }, - { - "name": "timestamptz", - "expression": "cast('2024-01-01T23:59:59' as timestamp with time zone)", - "type": "timestamp", - }, - { - "name": "test_null_time", - "expression": "cast(NULL as timestamp)", - "type": "timestamp", - }, - { - "name": "bytea_column", - "expression": "cast('abc' as bytea)", - "type": "bytea", - }, - ], - "primaryKey": "orderkey", - }, - ], -} +def wait_for_druid_service(url, timeout=300, interval=5): + """Wait for the Druid service to be ready. -@pytest.fixture -def manifest_str(): - return base64.b64encode(orjson.dumps(manifest)).decode("utf-8") + :param url: The URL to check. + :param timeout: The maximum time to wait (in seconds). + :param interval: The interval between checks (in seconds). + :return: True if the service is ready, False if the timeout is reached. + """ + start_time = time.time() + while time.time() - start_time < timeout: + try: + response = requests.get(url) + if response.status_code == 200: + return True + except requests.ConnectionError: + pass + time.sleep(interval) + return False @pytest.fixture(scope="module") -def docker(request) -> DockerContainer: - with DockerImage(tag="apache/druid:latest") as image: - with DockerContainer(str(image)) as container: - druid = container.start() - engine = sqlalchemy.create_engine(druid.get_connection_url()) - pd.read_parquet(file_path("resource/tpch/data/orders.parquet")).to_sql( - "orders", engine, index=False - ) - with engine.begin() as conn: - conn.execute( - text("COMMENT ON TABLE orders IS 'This is a table comment'") - ) - conn.execute( - text("COMMENT ON COLUMN orders.o_comment IS 'This is a comment'") - ) - request.addfinalizer(druid.stop) - return druid +def druid(request) -> DockerCompose: + with DockerCompose( + "tests/resource/druid", compose_file_name="docker-compose.yml", wait=True + ) as compose: + druid_url = "http://localhost:8081/status" + if not wait_for_druid_service(druid_url): + compose.stop() + raise Exception("Druid service did not become ready in time") + + yield compose + + +def test_create_datasource(druid: DockerCompose): + url = "http://localhost:8081/druid/indexer/v1/task" + payload = json.dumps( + { + "type": "index_parallel", + "spec": { + "dataSchema": { + "dataSource": "orders", + "timestampSpec": {"column": "timestamp_column", "format": "auto"}, + "dimensionsSpec": { + "dimensions": ["dimension1", "dimension2", "dimension3"] + }, + "metricsSpec": [ + {"type": "count", "name": "count"}, + { + "type": "doubleSum", + "name": "metric1", + "fieldName": "metric1", + }, + ], + "granularitySpec": { + "type": "uniform", + "segmentGranularity": "day", + "queryGranularity": "none", + }, + }, + "ioConfig": { + "type": "index_parallel", + "inputSource": { + "type": "local", + "baseDir": "tests/resource/tpch/data", + "filter": "orders.parquet", + }, + "inputFormat": {"type": "parquet"}, + }, + "tuningConfig": { + "type": "index_parallel", + "maxRowsPerSegment": 5000000, + "maxRowsInMemory": 25000, + }, + }, + } + ) + headers = {"Content-Type": "application/json"} + + response = requests.request("POST", url, headers=headers, data=payload) + assert response.status_code == 200 From 6e00d00faedb284def88170b242ac9f9b3cd4911 Mon Sep 17 00:00:00 2001 From: Gagan Date: Tue, 5 Nov 2024 01:28:00 -0700 Subject: [PATCH 4/7] test cases --- ibis-server/app/model/__init__.py | 7 ++++++- ibis-server/app/model/data_source.py | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index 40daf7e66..03ef9a4f3 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -45,10 +45,12 @@ class QueryPostgresDTO(QueryDTO): class QuerySnowflakeDTO(QueryDTO): connection_info: SnowflakeConnectionInfo = connection_info_field - + + class QueryDruidDTO(QueryDTO): connection_info: DruidConnectionInfo = connection_info_field + class QueryTrinoDTO(QueryDTO): connection_info: ConnectionUrl | TrinoConnectionInfo = connection_info_field @@ -77,11 +79,13 @@ class ClickHouseConnectionInfo(BaseModel): user: SecretStr password: SecretStr + class DruidConnectionInfo(BaseModel): host: SecretStr = Field(examples=["localhost"]) port: SecretStr = Field(examples=[8082]) path: str + class MSSqlConnectionInfo(BaseModel): host: SecretStr port: SecretStr @@ -135,6 +139,7 @@ class TrinoConnectionInfo(BaseModel): user: SecretStr | None = None password: SecretStr | None = None + ConnectionInfo = ( BigQueryConnectionInfo | CannerConnectionInfo diff --git a/ibis-server/app/model/data_source.py b/ibis-server/app/model/data_source.py index 761dd3899..28a834b04 100644 --- a/ibis-server/app/model/data_source.py +++ b/ibis-server/app/model/data_source.py @@ -20,8 +20,8 @@ QueryBigQueryDTO, QueryCannerDTO, QueryClickHouseDTO, - QueryDTO, QueryDruidDTO, + QueryDTO, QueryMSSqlDTO, QueryMySqlDTO, QueryPostgresDTO, @@ -111,7 +111,7 @@ def get_clickhouse_connection(info: ClickHouseConnectionInfo) -> BaseBackend: user=info.user.get_secret_value(), password=info.password.get_secret_value(), ) - + @staticmethod def get_druid_connection(info: DruidConnectionInfo) -> BaseBackend: return ibis.druid.connect( From 129aa1ef1bf80bbd1df80215cb748b4bf796055a Mon Sep 17 00:00:00 2001 From: Gagan Date: Tue, 5 Nov 2024 01:29:18 -0700 Subject: [PATCH 5/7] test cases --- .../tests/resource/druid/docker-compose.yml | 18 ------------------ ibis-server/tests/resource/druid/environment | 19 ------------------- 2 files changed, 37 deletions(-) diff --git a/ibis-server/tests/resource/druid/docker-compose.yml b/ibis-server/tests/resource/druid/docker-compose.yml index 0e18a600f..0cc0fad0f 100644 --- a/ibis-server/tests/resource/druid/docker-compose.yml +++ b/ibis-server/tests/resource/druid/docker-compose.yml @@ -1,21 +1,3 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# version: "2.2" volumes: diff --git a/ibis-server/tests/resource/druid/environment b/ibis-server/tests/resource/druid/environment index e223b5f4c..3e1a8cc4e 100644 --- a/ibis-server/tests/resource/druid/environment +++ b/ibis-server/tests/resource/druid/environment @@ -1,22 +1,3 @@ -# -# Licensed to the Apache Software Foundation (ASF) under one -# or more contributor license agreements. See the NOTICE file -# distributed with this work for additional information -# regarding copyright ownership. The ASF licenses this file -# to you under the Apache License, Version 2.0 (the -# "License"); you may not use this file except in compliance -# with the License. You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, -# software distributed under the License is distributed on an -# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -# KIND, either express or implied. See the License for the -# specific language governing permissions and limitations -# under the License. -# - # Java tuning #DRUID_XMX=1g #DRUID_XMS=1g From 44badb55a674e0420a410881ef0cbc68b48f9521 Mon Sep 17 00:00:00 2001 From: Gagan Date: Sat, 2 Nov 2024 21:56:36 -0600 Subject: [PATCH 6/7] apache druid datasource --- ibis-server/app/model/__init__.py | 9 +++++++-- ibis-server/app/model/data_source.py | 9 +++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index 03ef9a4f3..444ad2b0d 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -45,7 +45,9 @@ class QueryPostgresDTO(QueryDTO): class QuerySnowflakeDTO(QueryDTO): connection_info: SnowflakeConnectionInfo = connection_info_field - + +class QueryDruidDTO(QueryDTO): + connection_info: DruidConnectionInfo = connection_info_field class QueryDruidDTO(QueryDTO): connection_info: DruidConnectionInfo = connection_info_field @@ -79,6 +81,10 @@ class ClickHouseConnectionInfo(BaseModel): user: SecretStr password: SecretStr +class DruidConnectionInfo(BaseModel): + host: SecretStr = Field(examples=["localhost"]) + port: SecretStr = Field(examples=[8082]) + path: str class DruidConnectionInfo(BaseModel): host: SecretStr = Field(examples=["localhost"]) @@ -139,7 +145,6 @@ class TrinoConnectionInfo(BaseModel): user: SecretStr | None = None password: SecretStr | None = None - ConnectionInfo = ( BigQueryConnectionInfo | CannerConnectionInfo diff --git a/ibis-server/app/model/data_source.py b/ibis-server/app/model/data_source.py index 28a834b04..c671e4e65 100644 --- a/ibis-server/app/model/data_source.py +++ b/ibis-server/app/model/data_source.py @@ -22,6 +22,7 @@ QueryClickHouseDTO, QueryDruidDTO, QueryDTO, + QueryDruidDTO, QueryMSSqlDTO, QueryMySqlDTO, QueryPostgresDTO, @@ -111,6 +112,14 @@ def get_clickhouse_connection(info: ClickHouseConnectionInfo) -> BaseBackend: user=info.user.get_secret_value(), password=info.password.get_secret_value(), ) + + @staticmethod + def get_druid_connection(info: DruidConnectionInfo) -> BaseBackend: + return ibis.druid.connect( + host=info.host.get_secret_value(), + port=int(info.port.get_secret_value()), + path=info.path, + ) @staticmethod def get_druid_connection(info: DruidConnectionInfo) -> BaseBackend: From 4812ca30553e16a51316c1e727750c54845e7cc0 Mon Sep 17 00:00:00 2001 From: Gagan Date: Tue, 5 Nov 2024 01:28:00 -0700 Subject: [PATCH 7/7] test cases --- ibis-server/app/model/__init__.py | 9 ++------- ibis-server/app/model/data_source.py | 3 +-- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index 444ad2b0d..03ef9a4f3 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -45,9 +45,7 @@ class QueryPostgresDTO(QueryDTO): class QuerySnowflakeDTO(QueryDTO): connection_info: SnowflakeConnectionInfo = connection_info_field - -class QueryDruidDTO(QueryDTO): - connection_info: DruidConnectionInfo = connection_info_field + class QueryDruidDTO(QueryDTO): connection_info: DruidConnectionInfo = connection_info_field @@ -81,10 +79,6 @@ class ClickHouseConnectionInfo(BaseModel): user: SecretStr password: SecretStr -class DruidConnectionInfo(BaseModel): - host: SecretStr = Field(examples=["localhost"]) - port: SecretStr = Field(examples=[8082]) - path: str class DruidConnectionInfo(BaseModel): host: SecretStr = Field(examples=["localhost"]) @@ -145,6 +139,7 @@ class TrinoConnectionInfo(BaseModel): user: SecretStr | None = None password: SecretStr | None = None + ConnectionInfo = ( BigQueryConnectionInfo | CannerConnectionInfo diff --git a/ibis-server/app/model/data_source.py b/ibis-server/app/model/data_source.py index c671e4e65..5e6a7a34e 100644 --- a/ibis-server/app/model/data_source.py +++ b/ibis-server/app/model/data_source.py @@ -22,7 +22,6 @@ QueryClickHouseDTO, QueryDruidDTO, QueryDTO, - QueryDruidDTO, QueryMSSqlDTO, QueryMySqlDTO, QueryPostgresDTO, @@ -112,7 +111,7 @@ def get_clickhouse_connection(info: ClickHouseConnectionInfo) -> BaseBackend: user=info.user.get_secret_value(), password=info.password.get_secret_value(), ) - + @staticmethod def get_druid_connection(info: DruidConnectionInfo) -> BaseBackend: return ibis.druid.connect(