SNOW-2241405: dbapi ingest as stringtype when came across unsupported type (#3751)

Author: sfc-gh-yuwang

CHANGELOG.md

@@ -20,6 +20,7 @@
 
 #### Improvements
 
+- Unsupported types in `DataFrameReader.dbapi`(PuPr) are ingested as `StringType` now.
 - Improved error message to list available columns when dataframe cannot resolve given column name.
 - Added a new option `cacheResult` to `DataFrameReader.xml` that allows users to cache the result of the XML reader to a temporary table after calling `xml`. It helps improve performance when subsequent operations are performed on the same DataFrame.
 

src/snowflake/snowpark/_internal/data_source/drivers/databricks_driver.py

@@ -18,6 +18,7 @@
     VariantType,
     TimestampType,
     TimestampTimeZone,
+    StringType,
 )
 
 if TYPE_CHECKING:
@@ -57,7 +58,10 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
         all_columns = []
         for column_name, column_type, _ in schema:
             column_type = convert_map_to_use.get(column_type, column_type)
-            data_type = type_string_to_type_object(column_type)
+            try:
+                data_type = type_string_to_type_object(column_type)
+            except ValueError:
+                data_type = StringType()
             if column_type.lower() == "timestamp":
                 # by default https://docs.databricks.com/aws/en/sql/language-manual/data-types/timestamp-type
                 data_type = TimestampType(TimestampTimeZone.LTZ)

src/snowflake/snowpark/_internal/data_source/drivers/oracledb_driver.py

@@ -67,7 +67,7 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             oracledb.DB_TYPE_XMLTYPE: StringType,
             oracledb.DB_TYPE_OBJECT: VariantType,
             oracledb.DB_TYPE_VECTOR: VectorType,
-            oracledb.DB_TYPE_CURSOR: None,  # NOT SUPPORTED
+            # oracledb.DB_TYPE_CURSOR: None,  # NOT SUPPORTED
         }
 
         fields = []
@@ -77,10 +77,7 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             precision = column.precision
             scale = column.scale
             null_ok = column.null_ok
-            snow_type = convert_map_to_use.get(type_code, None)
-            if snow_type is None:
-                # TODO: SNOW-1912068 support types that we don't have now
-                raise NotImplementedError(f"oracledb type not supported: {type_code}")
+            snow_type = convert_map_to_use.get(type_code, StringType)
             if type_code == oracledb.DB_TYPE_TIMESTAMP_TZ:
                 data_type = snow_type(TimestampTimeZone.TZ)
             elif type_code == oracledb.DB_TYPE_TIMESTAMP_LTZ:

src/snowflake/snowpark/_internal/data_source/drivers/psycopg2_driver.py

@@ -189,15 +189,10 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
             try:
                 type_code = Psycopg2TypeCode(type_code)
             except ValueError:
-                raise NotImplementedError(
-                    f"Postgres type not supported: {type_code} for column: {name}"
-                )
-            snow_type = BASE_POSTGRES_TYPE_TO_SNOW_TYPE.get(type_code)
-            if snow_type is None:
-                raise NotImplementedError(
-                    f"Postgres type not supported: {type_code} for column: {name}"
-                )
-            if Psycopg2TypeCode(type_code) == Psycopg2TypeCode.NUMERICOID:
+                # not supported type is now handled as string type in below code
+                type_code = None
+            snow_type = BASE_POSTGRES_TYPE_TO_SNOW_TYPE.get(type_code, StringType)
+            if type_code == Psycopg2TypeCode.NUMERICOID:
                 if not self.validate_numeric_precision_scale(precision, scale):
                     logger.debug(
                         f"Snowpark does not support column"

src/snowflake/snowpark/_internal/data_source/drivers/pymsql_driver.py

@@ -157,9 +157,7 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
                 scale,
                 null_ok,
             ) = col
-            snow_type = BASE_PYMYSQL_TYPE_TO_SNOW_TYPE.get(type_code, None)
-            if snow_type is None:
-                raise NotImplementedError(f"mysql type not supported: {type_code}")
+            snow_type = BASE_PYMYSQL_TYPE_TO_SNOW_TYPE.get(type_code, StringType)
             if type_code in (PymysqlTypeCode.DECIMAL, PymysqlTypeCode.NEWDECIMAL):
                 # we did -2 here because what driver returned is precision + 2, mysql store + 2 precision internally
                 precision -= 2

src/snowflake/snowpark/_internal/data_source/drivers/pyodbc_driver.py

@@ -59,9 +59,7 @@ def to_snow_type(self, schema: List[Any]) -> StructType:
                 scale,
                 null_ok,
             ) = column
-            snow_type = BASE_PYODBC_TYPE_TO_SNOW_TYPE.get(type_code, None)
-            if snow_type is None:
-                raise NotImplementedError(f"sql server type not supported: {type_code}")
+            snow_type = BASE_PYODBC_TYPE_TO_SNOW_TYPE.get(type_code, StringType)
             if type_code in (int, decimal.Decimal):
                 if not self.validate_numeric_precision_scale(precision, scale):
                     logger.debug(

tests/integ/datasource/test_databricks.py

@@ -250,3 +250,11 @@ def test_unit_udtf_ingestion():
             else:
                 # Keep other types as is
                 assert value == expected_row[index]
+
+
+def test_unsupported_type():
+
+    schema = DatabricksDriver(
+        create_databricks_connection, DBMS_TYPE.DATABRICKS_DB
+    ).to_snow_type([("test_col", "unsupported_type", True)])
+    assert schema == StructType([StructField("TEST_COL", StringType(), nullable=True)])

tests/integ/datasource/test_mysql.py

@@ -14,6 +14,7 @@
     PymysqlTypeCode,
 )
 from snowflake.snowpark._internal.data_source.utils import DBMS_TYPE
+from snowflake.snowpark.types import StructType, StructField, StringType
 from tests.resources.test_data_source_dir.test_mysql_data import (
     mysql_real_data,
     MysqlType,
@@ -140,14 +141,6 @@ def test_dbapi_batch_fetch(
         assert df.order_by("ID").collect() == expected_result
 
 
-def test_type_conversion():
-    invalid_type = MysqlType("ID", "UNKNOWN", None, None, None, None, False)
-    with pytest.raises(NotImplementedError, match="mysql type not supported"):
-        PymysqlDriver(create_connection_mysql, DBMS_TYPE.MYSQL_DB).to_snow_type(
-            [invalid_type]
-        )
-
-
 def test_pymysql_driver_coverage(caplog):
     mysql_driver = PymysqlDriver(create_connection_mysql, DBMS_TYPE.MYSQL_DB)
     mysql_driver.to_snow_type(
@@ -296,3 +289,11 @@ def test_server_side_cursor():
     assert isinstance(cursor, pymysql.cursors.SSCursor)
     cursor.close()
     conn.close()
+
+
+def test_unsupported_type():
+
+    schema = PymysqlDriver(create_connection_mysql, DBMS_TYPE.MYSQL_DB).to_snow_type(
+        [("test_col", "unsupported_type", None, None, 0, 0, True)]
+    )
+    assert schema == StructType([StructField("TEST_COL", StringType(), nullable=True)])

tests/integ/datasource/test_oracledb.py

@@ -5,6 +5,7 @@
 import logging
 import math
 import sys
+from collections import namedtuple
 
 import pytest
 
@@ -18,6 +19,7 @@
 from snowflake.snowpark._internal.data_source.utils import (
     DBMS_TYPE,
 )
+from snowflake.snowpark.types import StructType, StructField, StringType
 from tests.parameters import ORACLEDB_CONNECTION_PARAMETERS
 from tests.resources.test_data_source_dir.test_data_source_data import (
     OracleDBType,
@@ -135,14 +137,6 @@ def test_dbapi_batch_fetch(
         assert df.order_by("ID").collect() == expected_result
 
 
-def test_type_conversion():
-    invalid_type = OracleDBType("ID", "UNKNOWN", None, None, False)
-    with pytest.raises(NotImplementedError, match="oracledb type not supported"):
-        OracledbDriver(create_connection_oracledb, DBMS_TYPE.ORACLE_DB).to_snow_type(
-            [invalid_type]
-        )
-
-
 def test_oracledb_driver_coverage(caplog):
     oracledb_driver = OracledbDriver(create_connection_oracledb, DBMS_TYPE.ORACLE_DB)
     conn = oracledb_driver.prepare_connection(oracledb_driver.create_connection(), 0)
@@ -242,3 +236,15 @@ def test_double_quoted_column_name_oracledb(session, custom_schema):
         )
     ]
     assert df.schema == oracledb_double_quoted_schema
+
+
+def test_unsupported_type():
+    invalid_type = OracleDBType("ID", "UNKNOWN", None, None, False)
+    MockDescription = namedtuple(
+        "mock_description", ["name", "type_code", "precision", "scale", "null_ok"]
+    )
+
+    schema = OracledbDriver(
+        create_connection_oracledb, DBMS_TYPE.ORACLE_DB
+    ).to_snow_type([MockDescription("test_col", invalid_type, 0, 0, True)])
+    assert schema == StructType([StructField("TEST_COL", StringType(), nullable=True)])

tests/integ/datasource/test_postgres.py

@@ -369,18 +369,26 @@ def test_unit_psycopg2_driver_to_snow_type_mapping():
     assert result.fields[3].datatype.scale == 0
 
     # Test unsupported type code
-    with pytest.raises(NotImplementedError, match="Postgres type not supported"):
-        nonexisting_type_code = -1
-        Psycopg2Driver(create_postgres_connection, DBMS_TYPE.POSTGRES_DB).to_snow_type(
-            [("UNSUPPORTED_COL", nonexisting_type_code, None, None, None, None, True)]
-        )
+    nonexisting_type_code = -1
+    schema = Psycopg2Driver(
+        create_postgres_connection, DBMS_TYPE.POSTGRES_DB
+    ).to_snow_type(
+        [("UNSUPPORTED_COL", nonexisting_type_code, None, None, None, None, True)]
+    )
+    assert schema == StructType(
+        [StructField("UNSUPPORTED_COL", StringType(), nullable=True)]
+    )
 
     # Test unsupported type code
-    with pytest.raises(NotImplementedError, match="Postgres type not supported"):
-        unimplemented_code = Psycopg2TypeCode.ACLITEMOID
-        Psycopg2Driver(create_postgres_connection, DBMS_TYPE.POSTGRES_DB).to_snow_type(
-            [("UNSUPPORTED_COL", unimplemented_code, None, None, None, None, True)]
-        )
+    unimplemented_code = Psycopg2TypeCode.ACLITEMOID
+    schema = Psycopg2Driver(
+        create_postgres_connection, DBMS_TYPE.POSTGRES_DB
+    ).to_snow_type(
+        [("UNSUPPORTED_COL", unimplemented_code, None, None, None, None, True)]
+    )
+    assert schema == StructType(
+        [StructField("UNSUPPORTED_COL", StringType(), nullable=True)]
+    )
 
 
 def test_unit_generate_select_query():