fix tests

Author: m-abulazm

src/databricks/labs/lakebridge/reconcile/connectors/jdbc_reader.py

@@ -22,8 +22,7 @@ def _get_jdbc_reader(self, query, jdbc_url, driver, additional_options: dict | N
             .option("dbtable", f"({query}) tmp")
         )
         if isinstance(additional_options, dict):
-            for key, value in additional_options.items():
-                reader = reader.option(key, value)
+            reader = reader.options(**additional_options)
         return reader
 
     @staticmethod

src/databricks/labs/lakebridge/reconcile/connectors/oracle.py

@@ -64,9 +64,9 @@ def read_data(
         table_query = query.replace(":tbl", f"{schema}.{table}")
         try:
             if options is None:
-                return self.reader(table_query).options(**self._get_timestamp_options()).load()
+                return self.reader(table_query, self._get_timestamp_options()).load()
             reader_options = self._get_jdbc_reader_options(options) | self._get_timestamp_options()
-            df = self.reader(table_query).options(**reader_options).load()
+            df = self.reader(table_query, reader_options).load()
             logger.warning(f"Fetching data using query: \n`{table_query}`")
 
             # Convert all column names to lower case
@@ -107,12 +107,14 @@ def _get_timestamp_options() -> dict[str, str]:
             "HH24:MI:SS''');END;",
         }
 
-    def reader(self, query: str) -> DataFrameReader:
+    def reader(self, query: str, options: dict | None = None) -> DataFrameReader:
+        if options is None:
+            options = {}
         user = self._get_secret('user')
         password = self._get_secret('password')
         logger.debug(f"Using user: {user} to connect to Oracle")
         return self._get_jdbc_reader(
-            query, self.get_jdbc_url, OracleDataSource._DRIVER, {"user": user, "password": password}
+            query, self.get_jdbc_url, OracleDataSource._DRIVER, {**options, "user": user, "password": password}
         )
 
     def normalize_identifier(self, identifier: str) -> NormalizedIdentifier:

tests/unit/reconcile/connectors/test_oracle.py

@@ -76,9 +76,7 @@ def test_read_data_with_options():
     )
     spark.read.format().option().option.assert_called_with("driver", "oracle.jdbc.OracleDriver")
     spark.read.format().option().option().option.assert_called_with("dbtable", "(select 1 from data.employee) tmp")
-    spark.read.format().option().option().option().option.assert_called_with("user", "my_user")
-    spark.read.format().option().option().option().option().option.assert_called_with("password", "my_password")
-    jdbc_actual_args = spark.read.format().option().option().option().option().option().options.call_args.kwargs
+    jdbc_actual_args = spark.read.format().option().option().option().options.call_args.kwargs
     jdbc_expected_args = {
         "numPartitions": 50,
         "partitionColumn": "s_nationkey",
@@ -89,9 +87,11 @@ def test_read_data_with_options():
         "sessionInitStatement": r"BEGIN dbms_session.set_nls('nls_date_format', "
         r"'''YYYY-MM-DD''');dbms_session.set_nls('nls_timestamp_format', '''YYYY-MM-DD "
         r"HH24:MI:SS''');END;",
+        "user": "my_user",
+        "password": "my_password",
     }
     assert jdbc_actual_args == jdbc_expected_args
-    spark.read.format().option().option().option().option().option().options().load.assert_called_once()
+    spark.read.format().option().option().option().options().load.assert_called_once()
 
 
 def test_get_schema():
@@ -143,9 +143,7 @@ def test_read_data_exception_handling():
         filters=None,
     )
 
-    spark.read.format().option().option().option().option().option().options().load.side_effect = RuntimeError(
-        "Test Exception"
-    )
+    spark.read.format().option().option().option().options().load.side_effect = RuntimeError("Test Exception")
 
     # Call the read_data method with the Tables configuration and assert that a PySparkException is raised
     with pytest.raises(
@@ -160,7 +158,7 @@ def test_get_schema_exception_handling():
     engine, spark, ws, scope = initial_setup()
     ords = OracleDataSource(engine, spark, ws, scope)
 
-    spark.read.format().option().option().option().option().option().load.side_effect = RuntimeError("Test Exception")
+    spark.read.format().option().option().option().options().load.side_effect = RuntimeError("Test Exception")
 
     # Call the get_schema method with predefined table, schema, and catalog names and assert that a PySparkException
     # is raised

tests/unit/reconcile/connectors/test_sql_server.py

@@ -55,20 +55,7 @@ def test_get_jdbc_url_happy():
     url = data_source.get_jdbc_url
     # Assert that the URL is generated correctly
     assert url == (
-        """jdbc:sqlserver://my_host:777;databaseName=my_database;user=my_user;password=my_password;encrypt=true;trustServerCertificate=true;"""
-    )
-
-
-def test_get_jdbc_url_fail():
-    # initial setup
-    engine, spark, ws, scope = initial_setup()
-    ws.secrets.get_secret.side_effect = mock_secret
-    # create object for TSQLServerDataSource
-    data_source = TSQLServerDataSource(engine, spark, ws, scope)
-    url = data_source.get_jdbc_url
-    # Assert that the URL is generated correctly
-    assert url == (
-        """jdbc:sqlserver://my_host:777;databaseName=my_database;user=my_user;password=my_password;encrypt=true;trustServerCertificate=true;"""
+        """jdbc:sqlserver://my_host:777;databaseName=my_database;encrypt=true;trustServerCertificate=true;"""
     )
 
 
@@ -96,7 +83,7 @@ def test_read_data_with_options():
     spark.read.format.assert_called_with("jdbc")
     spark.read.format().option.assert_called_with(
         "url",
-        "jdbc:sqlserver://my_host:777;databaseName=my_database;user=my_user;password=my_password;encrypt=true;trustServerCertificate=true;",
+        "jdbc:sqlserver://my_host:777;databaseName=my_database;encrypt=true;trustServerCertificate=true;",
     )
     spark.read.format().option().option.assert_called_with("driver", "com.microsoft.sqlserver.jdbc.SQLServerDriver")
     spark.read.format().option().option().option.assert_called_with(
@@ -109,6 +96,8 @@ def test_read_data_with_options():
         "lowerBound": '0',
         "upperBound": "100",
         "fetchsize": 100,
+        "user": "my_user",
+        "password": "my_password",
     }
     assert actual_args == expected_args
     spark.read.format().option().option().option().options().load.assert_called_once()
@@ -166,7 +155,7 @@ def test_get_schema_exception_handling():
     engine, spark, ws, scope = initial_setup()
     data_source = TSQLServerDataSource(engine, spark, ws, scope)
 
-    spark.read.format().option().option().option().option().load.side_effect = RuntimeError("Test Exception")
+    spark.read.format().option().option().option().options().load.side_effect = RuntimeError("Test Exception")
 
     # Call the get_schema method with predefined table, schema, and catalog names and assert that a PySparkException
     # is raised