Repository: astronomer/astro-sdk

Issue:
Add merge function to databricks
Please describe the feature you'd like to see

I'd like to see a feature that allows users to run the aql.merge function to merge two Databricks Delta tables together using the aql library. This would enable us to more easily and efficiently merge data from multiple tables into a single table for further processing and analysis.

Describe the solution you'd like

A function or method that takes the names of the two Delta tables to be merged, along with a connection string for a Databricks cluster, and runs the aql.merge function on the tables using the aql library. The function should accept all of the same arguments as the traditional aql.merge function, except for an additional conn_id argument that specifies the connection to use.


Here's an example of what the function signature might look like:

```
target_table = Table(name="foo", conn_id="my_databricks_conn")
source_table = Table(name="bar", conn_id="my_databricks_conn")

aql.merge(
    target_table=target_table,
    source_table=source_table,
    target_conflict_columns=["id"],
    columns=["id", "name"],
    if_conflicts="update",
)
```

**Are there any alternatives to this feature?**
Another alternative would be to use a tool like Apache Spark to merge the data from the two tables into a third table. This would require users to set up and configure a Spark cluster in their databricks instance, whereas the proposed function would be more lightweight and easier to use.

**Additional context**
Add any other context about the feature request here.

**Acceptance Criteria**

- [ ] All checks and tests in the CI should pass
- [ ] Unit tests (90% code coverage or more, [once available](https://github.com/astronomer/astro-sdk/issues/191))
- [ ] Integration tests (if the feature relates to a new database or external service)
- [ ] Example DAG
- [ ] Docstrings in [reStructuredText](https://peps.python.org/pep-0287/) for each of  methods, classes, functions and module-level attributes (including Example DAG on how it should be used)
- [ ] Exception handling in case of errors
- [ ] Logging (are we exposing useful information to the user? e.g. source and destination)
- [ ] Improve the documentation (README, Sphinx, and any other relevant)
- [ ] How to use Guide for the feature ([example](https://airflow.apache.org/docs/apache-airflow-providers-postgres/stable/operators/postgres_operator_howto_guide.html))



Discussion:

Resolving patch:
diff --git a/python-sdk/src/astro/databricks/delta.py b/python-sdk/src/astro/databricks/delta.py
--- a/python-sdk/src/astro/databricks/delta.py
+++ b/python-sdk/src/astro/databricks/delta.py
@@ -85,16 +85,6 @@ def create_table_using_native_schema_autodetection(self, table: BaseTable, file:
         # TODO Do we need to implement this function? It seems like databricks will handle schemas for us
         raise NotImplementedError("Not implemented yet.")
 
-    def merge_table(
-        self,
-        source_table: BaseTable,
-        target_table: BaseTable,
-        source_to_target_columns_map: dict[str, str],
-        target_conflict_columns: list[str],
-        if_conflicts: MergeConflictStrategy = "exception",
-    ) -> None:
-        raise NotImplementedError("We do not yet support merge for databricks")
-
     def schema_exists(self, schema: str) -> bool:
         # Schemas do not need to be created for delta, so we can assume this is true
         return True
@@ -312,3 +302,98 @@ def load_pandas_dataframe_to_table(
             file = File(path=t + f"/dataframe_{uuid.uuid4()}.parquet", filetype=FileType.PARQUET)
             file.create_from_dataframe(source_dataframe)
             self.load_file_to_table(input_file=file, output_table=target_table, if_exists=if_exists)
+
+    @staticmethod
+    def get_merge_initialization_query(parameters: tuple) -> str:
+        """
+        Handles database-specific logic to handle constraints
+        for Delta. setting constraints is not required for merging tables in Delta so we only return "1"
+        """
+        return "SELECT 1"
+
+    def merge_table(
+        self,
+        source_table: BaseTable,
+        target_table: BaseTable,
+        source_to_target_columns_map: dict[str, str],
+        target_conflict_columns: list[str],
+        if_conflicts: MergeConflictStrategy = "exception",
+    ):
+        """
+        Merge the source table rows into a destination table.
+        The argument `if_conflicts` allows the user to define how to handle conflicts.
+
+        :param source_table: Contains the rows to be merged to the target_table
+        :param target_table: Contains the destination table in which the rows will be merged
+        :param source_to_target_columns_map: Dict of target_table columns names to source_table columns names
+        :param target_conflict_columns: List of cols where we expect to have a conflict while combining
+        :param if_conflicts: The strategy to be applied if there are conflicts.
+        """
+        statement, params = self._build_merge_sql(
+            source_table=source_table,
+            target_table=target_table,
+            source_to_target_columns_map=source_to_target_columns_map,
+            target_conflict_columns=target_conflict_columns,
+            if_conflicts=if_conflicts,
+        )
+        self.run_sql(sql=statement, parameters=params)
+
+    def _build_merge_sql(
+        self,
+        source_table: BaseTable,
+        target_table: BaseTable,
+        source_to_target_columns_map: dict[str, str],
+        target_conflict_columns: list[str],
+        if_conflicts: MergeConflictStrategy = "exception",
+    ):
+        """Build the SQL statement for Merge operation"""
+        # TODO: Simplify this function
+        source_table_name = source_table.name
+        target_table_name = target_table.name
+
+        source_cols = source_to_target_columns_map.keys()
+        target_cols = source_to_target_columns_map.values()
+
+        (
+            source_table_identifier,
+            source_table_param,
+        ) = self.get_sqlalchemy_template_table_identifier_and_parameter(source_table, "source_table")
+
+        (
+            target_table_identifier,
+            target_table_param,
+        ) = self.get_sqlalchemy_template_table_identifier_and_parameter(target_table, "target_table")
+        merge_target_dict = {
+            f"`target_table`.`{x}`": f"`target_table`.`{x}`" for i, x in enumerate(target_conflict_columns)
+        }
+        merge_source_dict = {
+            f"`source_table`.`{x}`": f"`source_table`.`{x}`" for i, x in enumerate(target_conflict_columns)
+        }
+        merge_clauses = " AND ".join(
+            f"{k}={v}" for k, v in zip(merge_target_dict.keys(), merge_source_dict.keys())
+        )
+        statement = (
+            f"merge into {target_table_identifier} as `target_table` "
+            f"using {source_table_identifier} as `source_table`"
+            f" on {merge_clauses}"
+        )
+
+        values_to_check = [target_table_name, source_table_name]
+        values_to_check.extend(source_cols)
+        values_to_check.extend(target_cols)
+        if if_conflicts == "update":
+            merge_statement = ",".join(
+                [f"target_table.{t} = source_table.{s}" for s, t in source_to_target_columns_map.items()]
+            )
+            statement += f" when matched then UPDATE SET {merge_statement}"
+        target_columns = ",".join(f"target_table.{t}" for t in target_cols)
+        append_columns = ",".join(f"source_table.{s}" for s in source_cols)
+        statement += f" when not matched then insert({target_columns}) values ({append_columns})"
+
+        params = {
+            **merge_target_dict,
+            **merge_source_dict,
+            "source_table": source_table_param,
+            "target_table": target_table_param,
+        }
+        return statement, params


Test patch:
diff --git a/python-sdk/tests/benchmark/dags/benchmark_gcs_to_big_query.py b/python-sdk/tests/benchmark/dags/benchmark_gcs_to_big_query.py
--- a/python-sdk/tests/benchmark/dags/benchmark_gcs_to_big_query.py
+++ b/python-sdk/tests/benchmark/dags/benchmark_gcs_to_big_query.py
@@ -96,7 +96,7 @@
 )
 
 delete_test_dataset = BigQueryDeleteDatasetOperator(
-    task_id="delete_airflow_test_dataset",
+    task_id="delete_airflow_dataset",
     dataset_id=DATASET_NAME,
     delete_contents=True,
     dag=dag,
diff --git a/python-sdk/tests/databricks_tests/test_databricks_api_util.py b/python-sdk/tests/databricks_tests/test_databricks_api_util.py
--- a/python-sdk/tests/databricks_tests/test_databricks_api_util.py
+++ b/python-sdk/tests/databricks_tests/test_databricks_api_util.py
@@ -27,7 +27,7 @@ def raise_not_exist_error(unneeded_arg):
 
 
 @mock.patch("databricks_cli.secrets.api.SecretApi.delete_scope")
-@mock.patch("databricks_cli.sdk.api_client.ApiClient")
+@mock.patch("databricks_cli.sdk.api_client.ApiClient", autospec=True)
 def test_delete_scope_http_error_arbitary(mock_api_client, mock_delete_secret):
     """Raise error when delete_scope throw error"""
     mock_delete_secret.raiseError.side_effect = raise_arb_http_error
@@ -37,7 +37,7 @@ def test_delete_scope_http_error_arbitary(mock_api_client, mock_delete_secret):
 
 
 @mock.patch("databricks_cli.secrets.api.SecretApi.delete_scope")
-@mock.patch("databricks_cli.sdk.api_client.ApiClient")
+@mock.patch("databricks_cli.sdk.api_client.ApiClient", autospec=True)
 def test_delete_scope_http_error_non_existent(mock_api_client, mock_delete_secret):
     """Do not raise an error when trying to delete non-existent scope and the error code is RESOURCE_DOES_NOT_EXIST"""
     mock_delete_secret.side_effect = raise_not_exist_error
diff --git a/python-sdk/tests_integration/databricks_tests/test_databricks_api_util.py b/python-sdk/tests_integration/databricks_tests/test_databricks_api_util.py
--- a/python-sdk/tests_integration/databricks_tests/test_databricks_api_util.py
+++ b/python-sdk/tests_integration/databricks_tests/test_databricks_api_util.py
@@ -1,24 +1,128 @@
 import pathlib
+from unittest import mock
+from unittest.mock import Mock, call
 
 import pytest
+from requests.exceptions import HTTPError
 
-from astro.constants import Database
-from astro.databricks.api_utils import delete_secret_scope
+from astro.databricks.api_utils import (
+    create_and_run_job,
+    create_secrets,
+    delete_secret_scope,
+    load_file_to_dbfs,
+)
 
 CWD = pathlib.Path(__file__).parent
 
 
-@pytest.mark.integration
-@pytest.mark.parametrize(
-    "database_table_fixture",
-    [
-        {
-            "database": Database.DELTA,
-        }
-    ],
-    indirect=True,
-    ids=["delta"],
-)
-def test_delete_scope_nonexistent(database_table_fixture):
-    db, _ = database_table_fixture
-    delete_secret_scope("non-existent-scope", api_client=db.api_client)
+def raise_arb_http_error(unneeded_arg):
+    arbitrary_http = HTTPError()
+    arbitrary_http.response = Mock(json=lambda: {"error_code": "FOOBAR"})
+    raise arbitrary_http
+
+
+def raise_not_exist_error(unneeded_arg):
+    non_existent_http = HTTPError()
+    non_existent_http.response = Mock(json=lambda: {"error_code": "RESOURCE_DOES_NOT_EXIST"})
+
+
+@mock.patch("databricks_cli.secrets.api.SecretApi.delete_scope")
+@mock.patch("databricks_cli.sdk.api_client.ApiClient")
+def test_delete_scope_http_error_arbitary(mock_api_client, mock_delete_secret):
+    """Raise error when delete_scope throw error"""
+    mock_delete_secret.raiseError.side_effect = raise_arb_http_error
+    mock_delete_secret.side_effect = raise_arb_http_error
+    with pytest.raises(HTTPError):
+        delete_secret_scope("non-existent-scope", api_client=mock_api_client)
+
+
+@mock.patch("databricks_cli.secrets.api.SecretApi.delete_scope")
+@mock.patch("databricks_cli.sdk.api_client.ApiClient")
+def test_delete_scope_http_error_non_existent(mock_api_client, mock_delete_secret):
+    """Do not raise an error when trying to delete non-existent scope and the error code is RESOURCE_DOES_NOT_EXIST"""
+    mock_delete_secret.side_effect = raise_not_exist_error
+    delete_secret_scope("non-existent-scope", api_client=mock_api_client)
+
+
+@mock.patch("databricks_cli.sdk.api_client.ApiClient", autospec=True)
+def test_create_and_run_job(mock_api_client):
+    create_and_run_job(
+        api_client=mock_api_client,
+        file_to_run="/foo/bar.py",
+        databricks_job_name="my-db-job",
+        existing_cluster_id="foobar",
+    )
+    mock_api_client.perform_query.return_value = {"job_id": 123}
+    calls = [
+        call.perform_query(
+            "POST",
+            "/jobs/create",
+            data={
+                "name": "my-db-job",
+                "spark_python_task": {"python_file": "/foo/bar.py"},
+                "existing_cluster_id": "foobar",
+            },
+            headers=None,
+            version=None,
+        ),
+    ]
+    mock_api_client.assert_has_calls(calls)
+
+
+@mock.patch("databricks_cli.sdk.api_client.ApiClient", autospec=True)
+def test_load_to_dbfs(mock_api_client):
+    load_file_to_dbfs(
+        local_file_path=CWD / "__init__.py", file_name="my_table.py", api_client=mock_api_client
+    )
+    calls = [
+        call.perform_query("POST", "/dbfs/mkdirs", data={"path": "dbfs:/mnt/pyscripts/"}, headers=None),
+        call.perform_query(
+            "POST",
+            "/dbfs/delete",
+            data={"path": "dbfs:/mnt/pyscripts/my_table.py", "recursive": False},
+            headers=None,
+        ),
+        call.perform_query(
+            "POST",
+            "/dbfs/put",
+            data={"path": "dbfs:/mnt/pyscripts/my_table.py", "overwrite": True},
+            headers={"Content-Type": None},
+            files={"file": mock.ANY},
+        ),
+    ]
+    mock_api_client.assert_has_calls(calls)
+
+
+@mock.patch("databricks_cli.sdk.api_client.ApiClient", autospec=True)
+def test_create_secrets(mock_api_client):
+    create_secrets(scope_name="my-scope", filesystem_secrets={"foo": "bar"}, api_client=mock_api_client)
+    calls = [
+        call.perform_query(
+            "POST",
+            "/secrets/scopes/create",
+            data={"scope": "my-scope"},
+            headers=None,
+        ),
+        call.perform_query(
+            "POST",
+            "/secrets/put",
+            data={"scope": "my-scope", "key": "astro_sdk_foo", "string_value": "bar"},
+            headers=None,
+        ),
+    ]
+    mock_api_client.assert_has_calls(calls)
+
+
+@mock.patch("databricks_cli.sdk.api_client.ApiClient", autospec=True)
+def test_delete_secret_scope(mock_api_client):
+    delete_secret_scope(scope_name="my-scope", api_client=mock_api_client)
+    mock_api_client.assert_has_calls(
+        [
+            call.perform_query(
+                "POST",
+                "/secrets/scopes/delete",
+                data={"scope": "my-scope"},
+                headers=None,
+            )
+        ]
+    )
diff --git a/python-sdk/tests_integration/sql/operators/test_merge.py b/python-sdk/tests_integration/sql/operators/test_merge.py
--- a/python-sdk/tests_integration/sql/operators/test_merge.py
+++ b/python-sdk/tests_integration/sql/operators/test_merge.py
@@ -113,9 +113,10 @@ def set_compare(l1, l2):
         {"database": Database.REDSHIFT},
         {"database": Database.POSTGRES},
         {"database": Database.SQLITE},
+        {"database": Database.DELTA},
     ],
     indirect=True,
-    ids=["snowflake", "bigquery", "redshift", "postgres", "sqlite"],
+    ids=["snowflake", "bigquery", "redshift", "postgres", "sqlite", "delta"],
 )
 @pytest.mark.parametrize(
     "multiple_tables_fixture",
@@ -150,9 +151,10 @@ def test_merge(database_table_fixture, multiple_tables_fixture, sample_dag, merg
         {"database": Database.SNOWFLAKE},
         {"database": Database.BIGQUERY},
         {"database": Database.REDSHIFT},
+        {"database": Database.DELTA},
     ],
     indirect=True,
-    ids=["snowflake", "bigquery", "redshift"],
+    ids=["snowflake", "bigquery", "redshift", "delta"],
 )
 @pytest.mark.parametrize(
     "multiple_tables_fixture",