v0.21.1-fastapi

agent/tools/exesql.py

@@ -53,12 +53,13 @@ class ExeSQLParam(ToolParamBase):
         self.max_records = 1024
 
     def check(self):
-        self.check_valid_value(self.db_type, "Choose DB type", ['mysql', 'postgres', 'mariadb', 'mssql', 'IBM DB2'])
+        self.check_valid_value(self.db_type, "Choose DB type", ['mysql', 'postgres', 'mariadb', 'mssql', 'IBM DB2', 'trino'])
         self.check_empty(self.database, "Database name")
         self.check_empty(self.username, "database username")
         self.check_empty(self.host, "IP Address")
         self.check_positive_integer(self.port, "IP Port")
-        self.check_empty(self.password, "Database password")
+        if self.db_type != "trino":
+            self.check_empty(self.password, "Database password")
         self.check_positive_integer(self.max_records, "Maximum number of records")
         if self.database == "rag_flow":
             if self.host == "ragflow-mysql":
@@ -123,6 +124,45 @@ class ExeSQL(ToolBase, ABC):
                     r'PWD=' + self._param.password
             )
             db = pyodbc.connect(conn_str)
+        elif self._param.db_type == 'trino':
+            try:
+                import trino
+                from trino.auth import BasicAuthentication
+            except Exception:
+                raise Exception("Missing dependency 'trino'. Please install: pip install trino")
+
+            def _parse_catalog_schema(db: str):
+                if not db:
+                    return None, None
+                if "." in db:
+                    c, s = db.split(".", 1)
+                elif "/" in db:
+                    c, s = db.split("/", 1)
+                else:
+                    c, s = db, "default"
+                return c, s
+
+            catalog, schema = _parse_catalog_schema(self._param.database)
+            if not catalog:
+                raise Exception("For Trino, `database` must be 'catalog.schema' or at least 'catalog'.")
+
+            http_scheme = "https" if os.environ.get("TRINO_USE_TLS", "0") == "1" else "http"
+            auth = None
+            if http_scheme == "https" and self._param.password:
+                auth = BasicAuthentication(self._param.username, self._param.password)
+
+            try:
+                db = trino.dbapi.connect(
+                    host=self._param.host,
+                    port=int(self._param.port or 8080),
+                    user=self._param.username or "ragflow",
+                    catalog=catalog,
+                    schema=schema or "default",
+                    http_scheme=http_scheme,
+                    auth=auth
+                )
+            except Exception as e:
+                raise Exception("Database Connection Failed! \n" + str(e))
         elif self._param.db_type == 'IBM DB2':
             import ibm_db
             conn_str = (