import os
from flask import Flask, request, jsonify
from pyhive import hive
def get_hive_connection():
"""Get a connection to Hive."""
host = os.environ.get("HIVESERVER2_HOST")
port = int(os.environ.get("HIVESERVER2_PORT", 10000))
user = os.environ.get("HIVESERVER2_USER")
password = os.environ.get("HIVESERVER2_PASSWORD")
return hive.connect(host=host, port=port, username=user, password=password)
def create_app():
app = Flask(__name__)
@app.route("/query", methods=["POST"])
def query():
"""Execute a query on Hive."""
data = request.get_json()
query_str = data.get("query")
if not query_str:
return jsonify({"error": "Missing query"}), 400
try:
conn = get_hive_connection()
cursor = conn.cursor()
cursor.execute(query_str)
results = cursor.fetchall()
column_names = [desc[0] for desc in cursor.description]
conn.close()
return jsonify({"columns": column_names, "rows": results})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/list_tables", methods=["GET"])
def list_tables():
"""List all tables in the current database."""
try:
conn = get_hive_connection()
cursor = conn.cursor()
cursor.execute("SHOW TABLES")
tables = [row[0] for row in cursor.fetchall()]
conn.close()
return jsonify({"tables": tables})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/describe_table", methods=["POST"])
def describe_table():
"""Describe the schema of a specific table."""
data = request.get_json()
table_name = data.get("table")
if not table_name:
return jsonify({"error": "Missing table name"}), 400
try:
conn = get_hive_connection()
cursor = conn.cursor()
cursor.execute(f"DESCRIBE {table_name}")
schema = [{"name": row[0], "type": row[1]} for row in cursor.fetchall()]
conn.close()
return jsonify({"schema": schema})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/list_databases", methods=["GET"])
def list_databases():
"""List all available databases."""
try:
conn = get_hive_connection()
cursor = conn.cursor()
cursor.execute("SHOW DATABASES")
databases = [row[0] for row in cursor.fetchall()]
conn.close()
return jsonify({"databases": databases})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/use_database", methods=["POST"])
def use_database():
"""Switch to a different database."""
data = request.get_json()
database_name = data.get("database")
if not database_name:
return jsonify({"error": "Missing database name"}), 400
try:
conn = get_hive_connection()
cursor = conn.cursor()
cursor.execute(f"USE {database_name}")
conn.close()
return jsonify({"message": f"Switched to database: {database_name}"})
except Exception as e:
return jsonify({"error": str(e)}), 500
@app.route("/status", methods=["GET"])
def status():
"""Get the current connection status."""
try:
conn = get_hive_connection()
conn.close()
return jsonify({"status": "ok"})
except Exception as e:
return jsonify({"status": "error", "message": str(e)}), 500
return app
def main():
"""Start the MCP server."""
app = create_app()
app.run(host="0.0.0.0", port=5000)
if __name__ == "__main__":
main()
