refactor duplicated code, more tests

Author: joroscoSF

src/datacustomcode/io/reader/query_api.py

@@ -22,51 +22,21 @@
     Union,
 )
 
-import pandas.api.types as pd_types
-from pyspark.sql.types import (
-    BooleanType,
-    DoubleType,
-    LongType,
-    StringType,
-    StructField,
-    StructType,
-    TimestampType,
-)
 from salesforcecdpconnector.connection import SalesforceCDPConnection
 
 from datacustomcode.credentials import AuthType, Credentials
 from datacustomcode.io.reader.base import BaseDataCloudReader
 from datacustomcode.io.reader.sf_cli import SFCLIDataCloudReader
+from datacustomcode.io.reader.utils import _pandas_to_spark_schema
 
 if TYPE_CHECKING:
-    import pandas
     from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession
-    from pyspark.sql.types import AtomicType
+    from pyspark.sql.types import AtomicType, StructType
 
 logger = logging.getLogger(__name__)
 
 
 SQL_QUERY_TEMPLATE: Final = "SELECT * FROM {} LIMIT {}"
-PANDAS_TYPE_MAPPING = {
-    "object": StringType(),
-    "int64": LongType(),
-    "float64": DoubleType(),
-    "bool": BooleanType(),
-}
-
-
-def _pandas_to_spark_schema(
-    pandas_df: pandas.DataFrame, nullable: bool = True
-) -> StructType:
-    fields = []
-    for column, dtype in pandas_df.dtypes.items():
-        spark_type: AtomicType
-        if pd_types.is_datetime64_any_dtype(dtype):
-            spark_type = TimestampType()
-        else:
-            spark_type = PANDAS_TYPE_MAPPING.get(str(dtype), StringType())
-        fields.append(StructField(column, spark_type, nullable))
-    return StructType(fields)
 
 
 def create_cdp_connection(

src/datacustomcode/io/reader/sf_cli.py

@@ -25,47 +25,18 @@
 )
 
 import pandas as pd
-import pandas.api.types as pd_types
-from pyspark.sql.types import (
-    BooleanType,
-    DoubleType,
-    LongType,
-    StringType,
-    StructField,
-    StructType,
-    TimestampType,
-)
 import requests
 
 from datacustomcode.io.reader.base import BaseDataCloudReader
+from datacustomcode.io.reader.utils import _pandas_to_spark_schema
 
 if TYPE_CHECKING:
     from pyspark.sql import DataFrame as PySparkDataFrame, SparkSession
-    from pyspark.sql.types import AtomicType
+    from pyspark.sql.types import AtomicType, StructType
 
 logger = logging.getLogger(__name__)
 
 API_VERSION: Final = "v66.0"
-PANDAS_TYPE_MAPPING = {
-    "object": StringType(),
-    "int64": LongType(),
-    "float64": DoubleType(),
-    "bool": BooleanType(),
-}
-
-
-def _pandas_to_spark_schema(
-    pandas_df: pd.DataFrame, nullable: bool = True
-) -> StructType:
-    fields = []
-    for column, dtype in pandas_df.dtypes.items():
-        spark_type: AtomicType
-        if pd_types.is_datetime64_any_dtype(dtype):
-            spark_type = TimestampType()
-        else:
-            spark_type = PANDAS_TYPE_MAPPING.get(str(dtype), StringType())
-        fields.append(StructField(column, spark_type, nullable))
-    return StructType(fields)
 
 
 class SFCLIDataCloudReader(BaseDataCloudReader):

src/datacustomcode/io/reader/utils.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2025, Salesforce, Inc.
+# SPDX-License-Identifier: Apache-2
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas.api.types as pd_types
+from pyspark.sql.types import (
+    BooleanType,
+    DoubleType,
+    LongType,
+    StringType,
+    StructField,
+    StructType,
+    TimestampType,
+)
+
+if TYPE_CHECKING:
+    import pandas
+    from pyspark.sql.types import AtomicType
+
+PANDAS_TYPE_MAPPING = {
+    "object": StringType(),
+    "int64": LongType(),
+    "float64": DoubleType(),
+    "bool": BooleanType(),
+}
+
+
+def _pandas_to_spark_schema(
+    pandas_df: pandas.DataFrame, nullable: bool = True
+) -> StructType:
+    fields = []
+    for column, dtype in pandas_df.dtypes.items():
+        spark_type: AtomicType
+        if pd_types.is_datetime64_any_dtype(dtype):
+            spark_type = TimestampType()
+        else:
+            spark_type = PANDAS_TYPE_MAPPING.get(str(dtype), StringType())
+        fields.append(StructField(column, spark_type, nullable))
+    return StructType(fields)

tests/io/reader/test_query_api.py

@@ -21,8 +21,8 @@
 from datacustomcode.io.reader.query_api import (
     SQL_QUERY_TEMPLATE,
     QueryAPIDataCloudReader,
-    _pandas_to_spark_schema,
 )
+from datacustomcode.io.reader.utils import _pandas_to_spark_schema
 
 
 class TestPandasToSparkSchema: