Skip to content

feat: Supports hive engine for connecting to hiveserver2/kyuubi/spark thrift server #13

Description

@wForget

I have tried to support connecting to Apache kyuubi by adding the pyhive dependency.

Patch:

Index: backend/src/analytics_agent/engines/sqlalchemy/engine.py
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/backend/src/analytics_agent/engines/sqlalchemy/engine.py b/backend/src/analytics_agent/engines/sqlalchemy/engine.py
--- a/backend/src/analytics_agent/engines/sqlalchemy/engine.py	(revision 0555729bd138158c8ed9c66b5746492330c7f8ba)
+++ b/backend/src/analytics_agent/engines/sqlalchemy/engine.py	(revision 6d15f3a9d7b475f2fcb15f7a55566547e0aa6e3b)
@@ -60,7 +60,8 @@
             from sqlalchemy import create_engine
 
             url = self._build_url()
-            self._engine = create_engine(url)
+            connect_args = self._cfg.get("connect_args", {})
+            self._engine = create_engine(url, connect_args=connect_args)
             logger.info("[SQLAlchemy] engine created for url=%s", repr(url))
         return self._engine
 
@@ -77,7 +78,7 @@
             return str(v)
         return v
 
-    def _run_query(self, sql: str, limit: int | None = None) -> dict:
+    def _run_query(self, sql: str, limit: int | None = None, sql_allow_limit: bool = False) -> dict:
         from analytics_agent.config import settings
 
         try:
@@ -86,7 +87,7 @@
             return {"error": str(e), "columns": [], "rows": [], "truncated": False}
 
         effective_sql = sql.strip().rstrip(";")
-        if limit is not None and "LIMIT" not in effective_sql.upper():
+        if sql_allow_limit and limit is not None and "LIMIT" not in effective_sql.upper():
             effective_sql = f"{effective_sql} LIMIT {limit}"
 
         try:
@@ -109,11 +110,10 @@
         engine = self
 
         @tool
-        def execute_sql(sql: str) -> str:
+        def execute_sql(sql: str, sql_allow_limit: bool = True) -> str:
             """Execute a SQL query against the connected database. Returns JSON with columns and rows."""
             from analytics_agent.config import settings
-
-            result = engine._run_query(sql, limit=settings.sql_row_limit)
+            result = engine._run_query(sql, limit=settings.sql_row_limit, sql_allow_limit=sql_allow_limit)
             return orjson.dumps(result).decode()
 
         @tool
Index: backend/src/analytics_agent/engines/factory.py
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/backend/src/analytics_agent/engines/factory.py b/backend/src/analytics_agent/engines/factory.py
--- a/backend/src/analytics_agent/engines/factory.py	(revision 0555729bd138158c8ed9c66b5746492330c7f8ba)
+++ b/backend/src/analytics_agent/engines/factory.py	(revision 6d15f3a9d7b475f2fcb15f7a55566547e0aa6e3b)
@@ -16,6 +16,7 @@
     return {
         "snowflake": SnowflakeQueryEngine,
         "mysql": SQLAlchemyQueryEngine,
+        "hive": SQLAlchemyQueryEngine,
         "sqlite": SQLAlchemyQueryEngine,
         "postgresql": SQLAlchemyQueryEngine,
         "sqlalchemy": SQLAlchemyQueryEngine,
Index: pyproject.toml
IDEA additional info:
Subsystem: com.intellij.openapi.diff.impl.patch.CharsetEP
<+>UTF-8
===================================================================
diff --git a/pyproject.toml b/pyproject.toml
--- a/pyproject.toml	(revision 0555729bd138158c8ed9c66b5746492330c7f8ba)
+++ b/pyproject.toml	(revision 6d15f3a9d7b475f2fcb15f7a55566547e0aa6e3b)
@@ -45,6 +45,10 @@
     # Query engines
     "snowflake-connector-python[pandas,secure-local-storage]>=3.10.0",
     "pymysql>=1.1.0",
+    "pyhive @ git+https://github.com/apache/kyuubi.git@master#subdirectory=python",
+    "kerberos>=1.3.0",
+    "thrift>=0.10.0",
+    "thrift_sasl @ git+https://github.com/cloudera/thrift_sasl.git",
     # Observability (OTEL) — no-op unless OTEL_EXPORTER_OTLP_ENDPOINT is set
     "opentelemetry-sdk>=1.25.0",
     "opentelemetry-exporter-otlp-proto-http>=1.25.0",
@@ -80,6 +84,9 @@
 [tool.hatch.build.targets.wheel]
 packages = ["backend/src/analytics_agent"]
 
+[tool.hatch.metadata]
+allow-direct-references = true
+
 [tool.uv]
 package = true

config.yml:

engines:
  - type: hive
    name: kyuubi
    connection:
     dialect: hive
     host: ${KYUUBI_HOST}
     port: ${KYUUBI_PORT}
     database: test_db
     connect_args:
       auth: KERBEROS
       kerberos_service_name: ${KYUUBI_KERBEROS_SERVICE}

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions