I have tried to support connecting to Apache kyuubi by adding the pyhive dependency.
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
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}
I have tried to support connecting to Apache kyuubi by adding the pyhive dependency.
Patch:
config.yml: