apache druid datasource

Canner · Nov 6, 2024 · eac8be8 · eac8be8
1 parent 1a84a6d
commit eac8be8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 2 deletions.
diff --git 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
@@ -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: