Added Threshold Query Builder

pyproject.toml

@@ -132,7 +132,7 @@ branch = true
 parallel = true
 
 [tool.coverage.report]
-omit = ["src/databricks/labs/remorph/reconcile/*",
+omit = [
   "src/databricks/labs/remorph/coverage/*",
   "src/databricks/labs/remorph/helpers/execution_time.py",
   "__about__.py"]

src/databricks/labs/remorph/reconcile/connectors/data_source.py

@@ -1,3 +1,4 @@
+import re
 from abc import ABC, abstractmethod
 
 from databricks.sdk import WorkspaceClient  # pylint: disable-next=wrong-import-order
@@ -6,7 +7,6 @@
 from databricks.labs.remorph.reconcile.recon_config import (  # pylint: disable=ungrouped-imports
     JdbcReaderOptions,
     Schema,
-    Tables,
 )
 
 
@@ -20,11 +20,18 @@
         self.scope = scope
 
     @abstractmethod
-    def read_data(self, schema_name: str, catalog_name: str, query: str, table_conf: Tables) -> DataFrame:
+    def read_data(
+        self, catalog_name: str, schema_name: str, query: str, jdbc_reader_options: JdbcReaderOptions
+    ) -> DataFrame:
         return NotImplemented
 
     @abstractmethod
-    def get_schema(self, table_name: str, schema_name: str, catalog_name: str) -> list[Schema]:
+    def get_schema(
+        self,
+        catalog_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> list[Schema]:
         return NotImplemented
 
     def _get_jdbc_reader(self, query, jdbc_url, driver):
@@ -48,3 +55,15 @@
     def _get_secrets(self, key_name):
         key = self.source + '_' + key_name
         return self.ws.secrets.get_secret(self.scope, key)
+
+    @staticmethod
+    def _get_table_or_query(
+        catalog_name: str,
+        schema_name: str,
+        query: str,
+    ):
+        if re.search('select', query, re.IGNORECASE):
+            return query.format(catalog_name=catalog_name, schema_name=schema_name)
+        if catalog_name:
+            return catalog_name + "." + schema_name + "." + query
+        return schema_name + "." + query

src/databricks/labs/remorph/reconcile/connectors/databricks.py

@@ -1,15 +1,22 @@
 from pyspark.sql import DataFrame
 
 from databricks.labs.remorph.reconcile.connectors.data_source import DataSource
-from databricks.labs.remorph.reconcile.recon_config import Schema, Tables
+from databricks.labs.remorph.reconcile.recon_config import JdbcReaderOptions, Schema
 
 
 class DatabricksDataSource(DataSource):
-    def read_data(self, schema_name: str, catalog_name: str, query: str, table_conf: Tables) -> DataFrame:
+    def read_data(
+        self, catalog_name: str, schema_name: str, query: str, jdbc_reader_options: JdbcReaderOptions
+    ) -> DataFrame:
         # Implement Databricks-specific logic here
         return NotImplemented
 
-    def get_schema(self, table_name: str, schema_name: str, catalog_name: str) -> list[Schema]:
+    def get_schema(
+        self,
+        catalog_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> list[Schema]:
         # Implement Databricks-specific logic here
         return NotImplemented
 

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

@@ -3,7 +3,7 @@
 
 from databricks.labs.remorph.reconcile.connectors.data_source import DataSource
 from databricks.labs.remorph.reconcile.constants import SourceDriver
-from databricks.labs.remorph.reconcile.recon_config import Schema, Tables
+from databricks.labs.remorph.reconcile.recon_config import JdbcReaderOptions, Schema
 
 
 class OracleDataSource(DataSource):
@@ -16,16 +16,16 @@
             f":{self._get_secrets('port')}/{self._get_secrets('database')}"
         )
 
-    # TODO need to check schema_name,catalog_name is needed
-    def read_data(self, schema_name: str, catalog_name: str, query: str, table_conf: Tables) -> DataFrame:
+    def read_data(
+        self, catalog_name: str, schema_name: str, query: str, jdbc_reader_options: JdbcReaderOptions
+    ) -> DataFrame:
         try:
-            if table_conf.jdbc_reader_options is None:
-                return self.reader(query).options(**self._get_timestamp_options()).load()
+            table_query = self._get_table_or_query(catalog_name, schema_name, query)
+            if jdbc_reader_options is None:
+                return self.reader(table_query).options(**self._get_timestamp_options()).load()
             return (
-                self.reader(query)
-                .options(
-                    **self._get_jdbc_reader_options(table_conf.jdbc_reader_options) | self._get_timestamp_options()
-                )
+                self.reader(table_query)
+                .options(**self._get_jdbc_reader_options(jdbc_reader_options) | self._get_timestamp_options())
                 .load()
             )
         except PySparkException as e:
@@ -34,7 +34,12 @@
             )
             raise PySparkException(error_msg) from e
 
-    def get_schema(self, table_name: str, schema_name: str, catalog_name: str) -> list[Schema]:
+    def get_schema(
+        self,
+        catalog_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> list[Schema]:
         try:
             schema_query = self._get_schema_query(table_name, schema_name)
             schema_df = self.reader(schema_query).load()

src/databricks/labs/remorph/reconcile/connectors/snowflake.py

@@ -1,15 +1,22 @@
 from pyspark.sql import DataFrame
 
 from databricks.labs.remorph.reconcile.connectors.data_source import DataSource
-from databricks.labs.remorph.reconcile.recon_config import Schema, Tables
+from databricks.labs.remorph.reconcile.recon_config import JdbcReaderOptions, Schema
 
 
 class SnowflakeDataSource(DataSource):
-    def read_data(self, schema_name: str, catalog_name: str, query: str, table_conf: Tables) -> DataFrame:
+    def read_data(
+        self, catalog_name: str, schema_name: str, query: str, jdbc_reader_options: JdbcReaderOptions
+    ) -> DataFrame:
         # Implement Snowflake-specific logic here
         return NotImplemented
 
-    def get_schema(self, table_name: str, schema_name: str, catalog_name: str) -> list[Schema]:
+    def get_schema(
+        self,
+        catalog_name: str,
+        schema_name: str,
+        table_name: str,
+    ) -> list[Schema]:
         # Implement Snowflake-specific logic here
         return NotImplemented
 

src/databricks/labs/remorph/reconcile/execute.py

@@ -4,7 +4,7 @@
 from databricks.labs.blueprint.installation import Installation
 
 from databricks.labs.remorph.reconcile.connectors.data_source import DataSource
-from databricks.labs.remorph.reconcile.recon_config import TableRecon, Tables
+from databricks.labs.remorph.reconcile.recon_config import Table, TableRecon
 
 logger = logging.getLogger(__name__)
 
@@ -27,10 +27,10 @@ def __init__(self, source: DataSource, target: DataSource):
         self.source = source
         self.target = target
 
-    def compare_schemas(self, table_conf: Tables, schema_name: str, catalog_name: str) -> bool:
+    def compare_schemas(self, table_conf: Table, schema_name: str, catalog_name: str) -> bool:
         raise NotImplementedError
 
-    def compare_data(self, table_conf: Tables, schema_name: str, catalog_name: str) -> bool:
+    def compare_data(self, table_conf: Table, schema_name: str, catalog_name: str) -> bool:
         raise NotImplementedError