From 27997b2f778bb6d562a6565656d7e04afdeac2e2 Mon Sep 17 00:00:00 2001
From: Milan Bariya <52292922+MilanBariya@users.noreply.github.com>
Date: Mon, 21 Nov 2022 17:05:23 +0530
Subject: [PATCH] Fix: Revert Athena tests and types change (#8909)

---
 ingestion/src/metadata/utils/source_connections.py   |  4 +---
 ingestion/tests/unit/test_source_connection.py       | 12 +++++-------
 .../connections/database/athenaConnection.json       |  2 +-
 3 files changed, 7 insertions(+), 11 deletions(-)

diff --git a/ingestion/src/metadata/utils/source_connections.py b/ingestion/src/metadata/utils/source_connections.py
index ac857c9bedf..bfc7f37511b 100644
--- a/ingestion/src/metadata/utils/source_connections.py
+++ b/ingestion/src/metadata/utils/source_connections.py
@@ -476,9 +476,7 @@ def _(connection: AthenaConnection):
         url += ":"
     url += f"@athena.{connection.awsConfig.awsRegion}.amazonaws.com:443"
 
-    staging_url = connection.s3StagingDir.scheme + "://" + str(connection.s3StagingDir)
-
-    url += f"?s3_staging_dir={quote_plus(staging_url)}"
+    url += f"?s3_staging_dir={quote_plus(connection.s3StagingDir)}"
     if connection.workgroup:
         url += f"&work_group={connection.workgroup}"
     if connection.awsConfig.awsSessionToken:
diff --git a/ingestion/tests/unit/test_source_connection.py b/ingestion/tests/unit/test_source_connection.py
index 8e6be412359..cb590ffd604 100644
--- a/ingestion/tests/unit/test_source_connection.py
+++ b/ingestion/tests/unit/test_source_connection.py
@@ -11,8 +11,6 @@
 
 from unittest import TestCase
 
-from pydantic import AnyUrl
-
 from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
     AthenaConnection,
     AthenaScheme,
@@ -768,20 +766,20 @@ class SouceConnectionTest(TestCase):
             awsAccessKeyId="key", awsRegion="us-east-2", awsSecretAccessKey="secret_key"
         )
 
-        expected_url = "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443?s3_staging_dir=s3%3A%2F%2Fs3athena-postgres&work_group=primary"
+        expected_url = "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443?s3_staging_dir=s3%3A%2F%2Fpostgres%2Finput%2F&work_group=primary"
         athena_conn_obj = AthenaConnection(
             awsConfig=awsCreds,
-            s3StagingDir=AnyUrl("s3athena-postgres", scheme="s3"),
+            s3StagingDir="s3://postgres/input/",
             workgroup="primary",
             scheme=AthenaScheme.awsathena_rest,
         )
         assert expected_url == get_connection_url(athena_conn_obj)
 
-        # connection arguments with db
-        expected_url = "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443?s3_staging_dir=s3%3A%2F%2Fs3athena-postgres&work_group=primary"
+        # connection arguments witho db
+        expected_url = "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443?s3_staging_dir=s3%3A%2F%2Fpostgres%2Fintput%2F&work_group=primary"
         athena_conn_obj = AthenaConnection(
             awsConfig=awsCreds,
-            s3StagingDir=AnyUrl("s3athena-postgres", scheme="s3"),
+            s3StagingDir="s3://postgres/intput/",
             workgroup="primary",
             scheme=AthenaScheme.awsathena_rest,
         )
diff --git a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/athenaConnection.json b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/athenaConnection.json
index 5f11d87d9b7..8051551f3d4 100644
--- a/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/athenaConnection.json
+++ b/openmetadata-spec/src/main/resources/json/schema/entity/services/connections/database/athenaConnection.json
@@ -38,7 +38,7 @@
     },
     "s3StagingDir": {
       "title": "S3 Staging Directory",
-      "description": "S3 Staging Directory.",
+      "description": "S3 Staging Directory. Example: s3://postgres/input/",
       "type": "string",
       "format": "uri"
     },