SNOW-1763960: clean up todos in code base (#2091)

Author: sfc-gh-aling

src/snowflake/connector/aio/__init__.py

@@ -12,3 +12,9 @@
     SnowflakeCursor,
     DictCursor,
 ]
+
+
+async def connect(**kwargs) -> SnowflakeConnection:
+    conn = SnowflakeConnection(**kwargs)
+    await conn.connect()
+    return conn

src/snowflake/connector/aio/_connection.py

@@ -10,7 +10,6 @@
 import os
 import pathlib
 import sys
-import traceback
 import uuid
 from contextlib import suppress
 from io import StringIO
@@ -33,7 +32,6 @@
 from ..connection import DEFAULT_CONFIGURATION as DEFAULT_CONFIGURATION_SYNC
 from ..connection import SnowflakeConnection as SnowflakeConnectionSync
 from ..connection import _get_private_bytes_from_file
-from ..connection_diagnostic import ConnectionDiagnostic
 from ..constants import (
     _CONNECTIVITY_ERR_MSG,
     ENV_VAR_PARTNER,
@@ -519,7 +517,7 @@ def _init_connection_parameters(
         elif is_kwargs_empty:
             # connection_name is None and kwargs was empty when called
             ret_kwargs = _get_default_connection_params()
-        self.__set_error_attributes()  # TODO: error attributes async?
+        # TODO: SNOW-1770153 on self.__set_error_attributes()
         return ret_kwargs
 
     async def _cancel_query(
@@ -873,42 +871,9 @@ async def connect(self, **kwargs) -> None:
             self.__config(**self._conn_parameters)
 
         if self.enable_connection_diag:
-            exceptions_dict = {}
-            # TODO: we can make ConnectionDiagnostic async, do we need?
-            connection_diag = ConnectionDiagnostic(
-                account=self.account,
-                host=self.host,
-                connection_diag_log_path=self.connection_diag_log_path,
-                connection_diag_allowlist_path=(
-                    self.connection_diag_allowlist_path
-                    if self.connection_diag_allowlist_path is not None
-                    else self.connection_diag_whitelist_path
-                ),
-                proxy_host=self.proxy_host,
-                proxy_port=self.proxy_port,
-                proxy_user=self.proxy_user,
-                proxy_password=self.proxy_password,
+            raise NotImplementedError(
+                "Connection diagnostic is not supported in asyncio"
             )
-            try:
-                connection_diag.run_test()
-                await self.__open_connection()
-                connection_diag.cursor = self.cursor()
-            except Exception:
-                exceptions_dict["connection_test"] = traceback.format_exc()
-                logger.warning(
-                    f"""Exception during connection test:\n{exceptions_dict["connection_test"]} """
-                )
-            try:
-                connection_diag.run_post_test()
-            except Exception:
-                exceptions_dict["post_test"] = traceback.format_exc()
-                logger.warning(
-                    f"""Exception during post connection test:\n{exceptions_dict["post_test"]} """
-                )
-            finally:
-                connection_diag.generate_report()
-                if exceptions_dict:
-                    raise Exception(str(exceptions_dict))
         else:
             await self.__open_connection()
         self._telemetry = TelemetryClient(self._rest)
@@ -924,7 +889,10 @@ def cursor(
                 None,
                 DatabaseError,
                 {
-                    "msg": "Connection is closed",
+                    "msg": "Connection is closed.\nPlease establish the connection first by "
+                    "explicitly calling `await SnowflakeConnection.connect()` or "
+                    "using an async context manager: `async with SnowflakeConnection() as conn`. "
+                    "\nEnsure the connection is open before attempting any operations.",
                     "errno": ER_CONNECTION_IS_CLOSED,
                     "sqlstate": SQLSTATE_CONNECTION_NOT_EXISTS,
                 },

src/snowflake/connector/aio/_cursor.py

@@ -828,7 +828,7 @@ async def execute_async(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
 
     @property
     def errorhandler(self):
-        # check SNOW-1763103
+        # TODO: SNOW-1763103 for async error handler
         raise NotImplementedError(
             "Async Snowflake Python Connector does not support errorhandler. "
             "Please open a feature request issue in github if your want this feature: "
@@ -837,7 +837,7 @@ def errorhandler(self):
 
     @errorhandler.setter
     def errorhandler(self, value):
-        # check SNOW-1763103
+        # TODO: SNOW-1763103 for async error handler
         raise NotImplementedError(
             "Async Snowflake Python Connector does not support errorhandler. "
             "Please open a feature request issue in github if your want this feature: "
@@ -987,7 +987,6 @@ async def fetch_pandas_all(self, **kwargs: Any) -> DataFrame:
             await self._prefetch_hook()
         if self._query_result_format != "arrow":
             raise NotSupportedError
-        # # TODO: async telemetry
         await self._log_telemetry_job_data(
             TelemetryField.PANDAS_FETCH_ALL, TelemetryData.TRUE
         )
@@ -1022,7 +1021,6 @@ async def get_result_batches(self) -> list[ResultBatch] | None:
         """
         if self._result_set is None:
             return None
-        # TODO: async telemetry SNOW-1572217
         await self._log_telemetry_job_data(
             TelemetryField.GET_PARTITIONS_USED, TelemetryData.TRUE
         )

src/snowflake/connector/aio/_file_transfer_agent.py

@@ -118,7 +118,7 @@ async def execute(self) -> None:
                 # multichunk threshold
                 m.multipart_threshold = self._multipart_threshold
 
-        # TODO: https://snowflakecomputing.atlassian.net/browse/SNOW-1625364
+        # TODO: SNOW-1625364 for renaming client_prefetch_threads in asyncio
         logger.debug(f"parallel=[{self._parallel}]")
         if self._raise_put_get_error and not self._file_metadata:
             Error.errorhandler_wrapper(
@@ -238,7 +238,7 @@ def postprocess_done_cb(
         task_of_files = []
         for file_client in files:
             try:
-                # TODO: https://snowflakecomputing.atlassian.net/browse/SNOW-1708819
+                # TODO: SNOW-1708819 for code refactoring
                 res = (
                     await file_client.prepare_upload()
                     if is_upload

src/snowflake/connector/aio/_network.py

@@ -467,12 +467,9 @@ async def _post_request(
         _include_retry_params: bool = False,
     ) -> dict[str, Any]:
         full_url = f"{self.server_url}{url}"
-        # TODO: sync feature parity, probe connection
-        # if self._connection._probe_connection:
-        #     from pprint import pprint
-        #
-        #     ret = probe_connection(full_url)
-        #     pprint(ret)
+        if self._connection._probe_connection:
+            # TODO: SNOW-1572318 for probe connection
+            raise NotImplementedError("probe_connection is not supported in asyncio")
 
         ret = await self.fetch(
             "post",
@@ -716,8 +713,6 @@ async def _request_exec(
             else:
                 input_data = data
 
-            # TODO: aiohttp auth parameter works differently than requests.session.request
-            #  we can check if there's other aiohttp built-in mechanism to update this
             if HEADER_AUTHORIZATION_KEY in headers:
                 del headers[HEADER_AUTHORIZATION_KEY]
             if token != NO_TOKEN:
@@ -745,7 +740,7 @@ async def _request_exec(
                     if is_raw_text:
                         ret = await raw_ret.text()
                     elif is_raw_binary:
-                        # check SNOW-1738595 for is_raw_binary support
+                        # TODO: SNOW-1738595 for is_raw_binary support
                         raise NotImplementedError(
                             "reading raw binary data is not supported in asyncio connector,"
                             " please open a feature request issue in"

src/snowflake/connector/aio/auth/_auth.py

@@ -69,7 +69,7 @@ async def authenticate(
         timeout: int | None = None,
     ) -> dict[str, str | int | bool]:
         if mfa_callback or password_callback:
-            # check SNOW-1707210 for mfa_callback and password_callback support
+            # TODO: SNOW-1707210 for mfa_callback and password_callback support
             raise NotImplementedError(
                 "mfa_callback or password_callback is not supported in asyncio connector, please open a feature"
                 " request issue in github: https://github.com/snowflakedb/snowflake-connector-python/issues/new/choose"

test/integ/aio/test_connection_async.py

@@ -1517,8 +1517,8 @@ async def test_mock_non_existing_server(conn_cnx, caplog):
         )
 
 
-@pytest.mark.skip(
-    "SNOW-1759084 await anext(self._generator, None) does not execute code after yield"
+@pytest.mark.xfail(
+    reason="TODO: SNOW-1759084 await anext(self._generator, None) does not execute code after yield"
 )
 async def test_disable_telemetry(conn_cnx, caplog):
     # default behavior, closing connection, it will send telemetry
@@ -1528,7 +1528,8 @@ async def test_disable_telemetry(conn_cnx, caplog):
                 await (await cur.execute("select 1")).fetchall()
             assert (
                 len(conn._telemetry._log_batch) == 3
-            )  # 3 events are import package, fetch first, fetch last
+            )  # 3 events are `import package`, `fetch first`, it's missing `fetch last` because of SNOW-1759084
+
     assert "POST /telemetry/send" in caplog.text
     caplog.clear()
 

test/integ/aio/test_cursor_async.py

@@ -984,14 +984,14 @@ async def test_real_decimal(conn, db_parameters):
             assert rec["RATIO"] == decimal.Decimal("23.4"), "the decimal value"
 
 
-@pytest.mark.skip("SNOW-1763103")
+@pytest.mark.skip("SNOW-1763103 error handler async")
 async def test_none_errorhandler(conn_testaccount):
     c = conn_testaccount.cursor()
     with pytest.raises(errors.ProgrammingError):
         c.errorhandler = None
 
 
-@pytest.mark.skip("SNOW-1763103")
+@pytest.mark.skip("SNOW-1763103 error handler async")
 async def test_nope_errorhandler(conn_testaccount):
     def user_errorhandler(connection, cursor, errorclass, errorvalue):
         pass
@@ -1432,9 +1432,6 @@ async def test__log_telemetry_job_data(conn_cnx, caplog):
     ) in caplog.record_tuples
 
 
-@pytest.mark.skip(
-    reason="SNOW-1759076 Async for support in Cursor.get_result_batches()"
-)
 @pytest.mark.parametrize(
     "result_format,expected_chunk_type",
     (
@@ -1446,7 +1443,7 @@ async def test_resultbatch(
     conn_cnx,
     result_format,
     expected_chunk_type,
-    capture_sf_telemetry,
+    capture_sf_telemetry_async,
 ):
     """This test checks the following things:
     1. After executing a query can we pickle the result batches
@@ -1461,13 +1458,13 @@ async def test_resultbatch(
             "python_connector_query_result_format": result_format,
         }
     ) as con:
-        with capture_sf_telemetry.patch_connection(con) as telemetry_data:
-            with con.cursor() as cur:
-                cur.execute(
+        async with capture_sf_telemetry_async.patch_connection(con) as telemetry_data:
+            async with con.cursor() as cur:
+                await cur.execute(
                     f"select seq4() from table(generator(rowcount => {rowcount}));"
                 )
                 assert cur._result_set.total_row_index() == rowcount
-                pre_pickle_partitions = cur.get_result_batches()
+                pre_pickle_partitions = await cur.get_result_batches()
                 assert len(pre_pickle_partitions) > 1
                 assert pre_pickle_partitions is not None
                 assert all(
@@ -1481,7 +1478,7 @@ async def test_resultbatch(
     post_pickle_partitions: list[ResultBatch] = pickle.loads(pickle_str)
     total_rows = 0
     # Make sure the batches can be iterated over individually
-    async for it in post_pickle_partitions:
+    for it in post_pickle_partitions:
         print(it)
 
     for i, partition in enumerate(post_pickle_partitions):
@@ -1492,15 +1489,17 @@ async def test_resultbatch(
         else:
             assert partition.compressed_size is not None
             assert partition.uncompressed_size is not None
-        for row in partition:
+        # TODO: SNOW-1759076 Async for support in Cursor.get_result_batches()
+        for row in await partition.create_iter():
             col1 = row[0]
             assert col1 == total_rows
             total_rows += 1
     assert total_rows == rowcount
     total_rows = 0
     # Make sure the batches can be iterated over again
     for partition in post_pickle_partitions:
-        for row in partition:
+        # TODO: SNOW-1759076 Async for support in Cursor.get_result_batches()
+        for row in await partition.create_iter():
             col1 = row[0]
             assert col1 == total_rows
             total_rows += 1

test/integ/aio/test_dbapi_async.py

@@ -115,6 +115,7 @@ async def test_exceptions():
     assert issubclass(errors.NotSupportedError, errors.Error)
 
 
+@pytest.mark.skip("SNOW-1770153 for error as attribute on connection")
 async def test_exceptions_as_connection_attributes(conn_cnx):
     async with conn_cnx() as con:
         try:

test/integ/aio/test_errors_async.py

@@ -14,6 +14,7 @@
 from snowflake.connector.telemetry import TelemetryField
 
 
+@pytest.mark.skip("SNOW-1770153 for error as attribute on connection")
 async def test_error_classes(conn_cnx):
     """Error classes in Connector module, object."""
     # class