From a19c5bceca11e2633ab99527bbda0f880d065a4c Mon Sep 17 00:00:00 2001
From: Ayush Shah <ayush@getcollate.io>
Date: Fri, 8 Oct 2021 00:55:03 +0530
Subject: [PATCH] Trino fix #705 (#706)

* Trino fix #705

* Update trino.json to point to default database
---
 ingestion/examples/workflows/trino.json          | 4 ++--
 ingestion/src/metadata/ingestion/source/trino.py | 9 +++++++--
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/ingestion/examples/workflows/trino.json b/ingestion/examples/workflows/trino.json
index 5749a98117f..989feb1c780 100644
--- a/ingestion/examples/workflows/trino.json
+++ b/ingestion/examples/workflows/trino.json
@@ -3,8 +3,8 @@
     "type": "trino",
     "config": {
       "service_name": "local_trino",
-      "host_port": "192.168.1.32:8080",
-      "database": "default"
+      "host_port": "localhost:8080",
+      "catalog": "system"
     }
   },
   "sink": {
diff --git a/ingestion/src/metadata/ingestion/source/trino.py b/ingestion/src/metadata/ingestion/source/trino.py
index 01d593bb36f..f1f43d4fde9 100644
--- a/ingestion/src/metadata/ingestion/source/trino.py
+++ b/ingestion/src/metadata/ingestion/source/trino.py
@@ -12,6 +12,7 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+from typing import Optional
 from urllib.parse import quote_plus
 
 from .sql_source import SQLSource, SQLConnectionConfig
@@ -22,6 +23,8 @@ class TrinoConfig(SQLConnectionConfig):
     host_port = "localhost:8080"
     scheme = "trino"
     service_type = "Trino"
+    catalog: str
+    schema_name: Optional[str]
 
     def get_connection_url(self):
         url = f"{self.scheme}://"
@@ -30,8 +33,10 @@ class TrinoConfig(SQLConnectionConfig):
             if self.password:
                 url += f":{quote_plus(self.password)}"
         url += f"{self.host_port}"
-        if self.database:
-            url += f"?schema={quote_plus(self.database)}"
+        if self.catalog:
+            url += f"/{quote_plus(self.catalog)}"
+            if self.schema_name:
+                url += f"/{quote_plus(self.schema_name)}"
         return url