Repository: astronomer/astro-sdk

Issue:
Issue due to removal of region from Snowflake URI
**Version**: `astro==0.4.0`

**Context**
The way we are building Snowflake URIs seems to be prone to errors.

Recently, a user had reported that Snowflake had deprecated the usage of the region in the host URI, and we made this change:
https://github.com/astro-projects/astro/commit/541038c81bef4ac5b9ed6ede87aeacf4b8d014a1#diff-4498bebc277444fc4d9662514a3a2e2eacdcdfd139e160abd6fdb16b0e43a854L31

Although the user was happy with the change, other users who were being able to successfully use Astro started experiencing the error:
```
snowflake.connector.errors.ForbiddenError: 250001 (08001): None: Failed to connect to DB. Verify the account name is correct: GP123.snowflakecomputing.com:443. 000403: 403: HTTP 403: Forbidden
```
The goal of this ticket is to resolve the issue so both users are able to use Astro, if possible removing any complexity within Astro.

**Additional details**

Based on the [Snowflake documentation](https://docs.snowflake.com/en/user-guide/admin-account-identifier.html#using-an-account-locator-as-an-identifier):
```
For example, if your account locator is xy12345:

    If the account is located in the AWS US West (Oregon) region, no additional segments are required and the URL would be xy12345.snowflakecomputing.com.

    If the account is located in the AWS US East (Ohio) region, additional segments are required and the URL would be xy12345.us-east-2.aws.snowflakecomputing.com.
```

One approach we can consider taking is to remove altogether the need for the `TempSnowflakeHook` class and trust the user to set the Snowflake host which works for them.


Discussion:

Resolving patch:
diff --git a/src/astro/sql/operators/agnostic_save_file.py b/src/astro/sql/operators/agnostic_save_file.py
--- a/src/astro/sql/operators/agnostic_save_file.py
+++ b/src/astro/sql/operators/agnostic_save_file.py
@@ -23,10 +23,11 @@
 from airflow.hooks.base import BaseHook
 from airflow.models import BaseOperator, DagRun, TaskInstance
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+from airflow.providers.postgres.hooks.postgres import PostgresHook
 from google.cloud.storage import Client
 from smart_open import open
 
-from astro.sql.operators.temp_hooks import TempPostgresHook, TempSnowflakeHook
+from astro.sql.operators.temp_hooks import TempSnowflakeHook
 from astro.sql.table import Table
 from astro.utils.cloud_storage_creds import gcs_client, s3fs_creds
 from astro.utils.schema_util import get_table_name
@@ -81,7 +82,7 @@ def execute(self, context):
 
         # Select database Hook based on `conn` type
         input_hook = {
-            "postgres": TempPostgresHook(
+            "postgres": PostgresHook(
                 postgres_conn_id=input_table.conn_id, schema=input_table.database
             ),
             "snowflake": TempSnowflakeHook(
diff --git a/src/astro/sql/operators/temp_hooks.py b/src/astro/sql/operators/temp_hooks.py
--- a/src/astro/sql/operators/temp_hooks.py
+++ b/src/astro/sql/operators/temp_hooks.py
@@ -13,13 +13,8 @@
 See the License for the specific language governing permissions and
 limitations under the License.
 """
-import warnings
-from typing import Dict, Optional, Sequence, Union
 from urllib.parse import quote_plus
 
-from airflow.hooks.dbapi import DbApiHook
-from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
-from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 
@@ -37,27 +32,3 @@ def get_uri(self) -> str:
             "?warehouse={warehouse}&role={role}&authenticator={authenticator}"
         )
         return uri.format(**conn_config)
-
-
-class TempPostgresHook(PostgresHook):
-    """
-    Temporary class to get around a bug in the snowflakehook when creating URIs
-    """
-
-    def get_uri(self) -> str:
-        """
-        Extract the URI from the connection.
-
-        :return: the extracted uri.
-        """
-        conn = self.get_connection(getattr(self, self.conn_name_attr))
-        login = ""
-        if conn.login:
-            login = f"{quote_plus(conn.login)}:{quote_plus(conn.password)}@"
-        host = conn.host
-        if conn.port is not None:
-            host += f":{conn.port}"
-        uri = f"postgresql://{login}{host}/"
-        if self.schema:
-            uri += self.schema
-        return uri
diff --git a/src/astro/utils/load_dataframe.py b/src/astro/utils/load_dataframe.py
--- a/src/astro/utils/load_dataframe.py
+++ b/src/astro/utils/load_dataframe.py
@@ -16,11 +16,12 @@
 from typing import Optional, Union
 
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+from airflow.providers.postgres.hooks.postgres import PostgresHook
 from pandas import DataFrame
 from pandas.io.sql import SQLDatabase
 from snowflake.connector.pandas_tools import write_pandas
 
-from astro.sql.operators.temp_hooks import TempPostgresHook, TempSnowflakeHook
+from astro.sql.operators.temp_hooks import TempSnowflakeHook
 from astro.utils.schema_util import set_schema_query
 
 
@@ -36,9 +37,9 @@ def move_dataframe_to_sql(
     chunksize=None,
 ):
     # Select database Hook based on `conn` type
-    hook: Union[TempPostgresHook, TempSnowflakeHook] = {  # type: ignore
-        "postgresql": TempPostgresHook(postgres_conn_id=conn_id, schema=database),
-        "postgres": TempPostgresHook(postgres_conn_id=conn_id, schema=database),
+    hook: Union[PostgresHook, TempSnowflakeHook] = {  # type: ignore
+        "postgresql": PostgresHook(postgres_conn_id=conn_id, schema=database),
+        "postgres": PostgresHook(postgres_conn_id=conn_id, schema=database),
         "snowflake": TempSnowflakeHook(
             snowflake_conn_id=conn_id,
             database=database,


Test patch:
diff --git a/tests/operators/test_agnostic_load_file.py b/tests/operators/test_agnostic_load_file.py
--- a/tests/operators/test_agnostic_load_file.py
+++ b/tests/operators/test_agnostic_load_file.py
@@ -36,6 +36,8 @@
 from airflow.models import DAG, DagRun
 from airflow.models import TaskInstance as TI
 from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+
+# from astro.sql.operators.temp_hooks import PostgresHook
 from airflow.providers.postgres.hooks.postgres import PostgresHook
 from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
 from airflow.utils import timezone
@@ -47,7 +49,6 @@
 
 # Import Operator
 from astro.sql.operators.agnostic_load_file import AgnosticLoadFile, load_file
-from astro.sql.operators.temp_hooks import TempPostgresHook
 from astro.sql.table import Table, TempTable
 from tests.operators import utils as test_utils
 
@@ -234,7 +235,7 @@ def test_unique_task_id_for_same_path(self):
     def test_aql_local_file_to_postgres_no_table_name(self):
         OUTPUT_TABLE_NAME = "expected_table_from_csv"
 
-        self.hook_target = TempPostgresHook(
+        self.hook_target = PostgresHook(
             postgres_conn_id="postgres_conn", schema="pagila"
         )
 
@@ -299,7 +300,7 @@ def test_aql_local_file_to_bigquery_no_table_name(self):
     def test_aql_overwrite_existing_table(self):
         OUTPUT_TABLE_NAME = "expected_table_from_csv"
 
-        self.hook_target = TempPostgresHook(
+        self.hook_target = PostgresHook(
             postgres_conn_id="postgres_conn", schema="pagila"
         )
 
@@ -341,7 +342,7 @@ def test_aql_overwrite_existing_table(self):
     def test_aql_s3_file_to_postgres(self):
         OUTPUT_TABLE_NAME = "expected_table_from_s3_csv"
 
-        self.hook_target = TempPostgresHook(
+        self.hook_target = PostgresHook(
             postgres_conn_id="postgres_conn", schema="pagila"
         )
 
@@ -373,7 +374,7 @@ def test_aql_s3_file_to_postgres(self):
     def test_aql_s3_file_to_postgres_no_table_name(self):
         OUTPUT_TABLE_NAME = "test_dag_load_file_homes_csv_2"
 
-        self.hook_target = TempPostgresHook(
+        self.hook_target = PostgresHook(
             postgres_conn_id="postgres_conn", schema="pagila"
         )
 
@@ -407,7 +408,7 @@ def test_aql_s3_file_to_postgres_no_table_name(self):
     def test_aql_s3_file_to_postgres_specify_schema(self):
         OUTPUT_TABLE_NAME = "expected_table_from_s3_csv"
 
-        self.hook_target = TempPostgresHook(
+        self.hook_target = PostgresHook(
             postgres_conn_id="postgres_conn", schema="pagila"
         )
 
diff --git a/tests/operators/test_postgres_append.py b/tests/operators/test_postgres_append.py
--- a/tests/operators/test_postgres_append.py
+++ b/tests/operators/test_postgres_append.py
@@ -31,7 +31,6 @@
 import unittest.mock
 
 import pandas as pd
-import pytest
 from airflow.executors.debug_executor import DebugExecutor
 from airflow.models import DAG, DagRun
 from airflow.models import TaskInstance as TI
@@ -44,7 +43,6 @@
 
 # Import Operator
 import astro.sql as aql
-from astro.sql.operators.temp_hooks import TempPostgresHook
 from astro.sql.table import Table
 
 # from tests.operators import utils as test_utils
@@ -185,7 +183,7 @@ def test_append(self):
 
     def test_append_all_fields(self):
 
-        hook = TempPostgresHook(postgres_conn_id="postgres_conn", schema="pagila")
+        hook = PostgresHook(postgres_conn_id="postgres_conn", schema="pagila")
 
         drop_table(table_name="test_main", postgres_conn=hook.get_conn())
         drop_table(table_name="test_append", postgres_conn=hook.get_conn())