potiuk commented on code in PR #74191:
URL: https://github.com/apache/airflow/pull/74191#discussion_r4182725398
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py:
##########
@@ -655,3 +663,49 @@ def get_openlineage_facets_on_complete(self, _):
job_facets={"sql":
SQLJobFacet(query=SQLParser.normalize_sql(self._sql))},
run_facets=run_facets,
)
+
+
+class DatabricksCopyIntoAssetOperator(DatabricksCopyIntoOperator):
+ """
+ Run ``COPY INTO`` and declare the target Unity Catalog table as an asset
outlet.
+
+ Accepts every :class:`DatabricksCopyIntoOperator` argument. ``table_name``
stays templated.
+ ``unity_table`` is static, so the asset is known when the Dag is parsed.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksCopyIntoAssetOperator`
+
+ :param unity_table: Static identity of the ``COPY INTO`` target table.
When ``outlets`` is
+ omitted, the outlets are ``[unity_table.to_asset()]``. When
``outlets`` is passed, including
+ ``outlets=[]``, it is kept as is, so callers that want extra outlets
include
+ ``unity_table.to_asset()`` among them. At execution, the rendered
``table_name`` is resolved
+ with ``catalog`` and ``schema`` and must equal ``unity_table``.
Otherwise ``ValueError`` is
+ raised and no SQL runs.
+ """
+
+ def __init__(self, *, unity_table: UnityTableIdentity, **kwargs) -> None:
+ if "outlets" not in kwargs:
+ kwargs["outlets"] = [unity_table.to_asset()]
+ super().__init__(**kwargs)
+ self.unity_table = unity_table
+
+ def execute(self, context: Context) -> Any:
+ expected = _CopyIntoTarget(
+ catalog=self.unity_table.catalog, schema=self.unity_table.schema,
table=self.unity_table.table
+ )
+ target = _resolve_copy_into_target(self.table_name,
catalog=self._catalog, schema=self._schema)
+ if target != expected:
Review Comment:
This compares catalog, schema and table names exactly, but Unity Catalog
names are case-insensitive and stored in lowercase, so `Main.Default.Users` and
`main.default.users` are the same table. Two things go wrong:
- A `table_name` that differs only in case from `unity_table`, e.g.
`MAIN.default.users` or a template that renders upper-case, raises `ValueError`
and the load never runs, although it targets the right table.
- A `UnityTableIdentity` with mixed-case names publishes a mixed-case asset
URI (`databricks://host/Main/...`). Other producers and consumers of the same
table will use the lowercase form, so Dags scheduled on it never trigger.
Could you lowercase `catalog`, `schema` and `table` in
`UnityTableIdentity.__post_init__`, and compare the resolved target
case-insensitively here?
##########
providers/databricks/src/airflow/providers/databricks/operators/databricks_sql.py:
##########
@@ -655,3 +663,49 @@ def get_openlineage_facets_on_complete(self, _):
job_facets={"sql":
SQLJobFacet(query=SQLParser.normalize_sql(self._sql))},
run_facets=run_facets,
)
+
+
+class DatabricksCopyIntoAssetOperator(DatabricksCopyIntoOperator):
+ """
+ Run ``COPY INTO`` and declare the target Unity Catalog table as an asset
outlet.
+
+ Accepts every :class:`DatabricksCopyIntoOperator` argument. ``table_name``
stays templated.
+ ``unity_table`` is static, so the asset is known when the Dag is parsed.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:DatabricksCopyIntoAssetOperator`
+
+ :param unity_table: Static identity of the ``COPY INTO`` target table.
When ``outlets`` is
+ omitted, the outlets are ``[unity_table.to_asset()]``. When
``outlets`` is passed, including
+ ``outlets=[]``, it is kept as is, so callers that want extra outlets
include
+ ``unity_table.to_asset()`` among them. At execution, the rendered
``table_name`` is resolved
+ with ``catalog`` and ``schema`` and must equal ``unity_table``.
Otherwise ``ValueError`` is
+ raised and no SQL runs.
+ """
+
+ def __init__(self, *, unity_table: UnityTableIdentity, **kwargs) -> None:
+ if "outlets" not in kwargs:
+ kwargs["outlets"] = [unity_table.to_asset()]
+ super().__init__(**kwargs)
+ self.unity_table = unity_table
+
+ def execute(self, context: Context) -> Any:
+ expected = _CopyIntoTarget(
+ catalog=self.unity_table.catalog, schema=self.unity_table.schema,
table=self.unity_table.table
+ )
+ target = _resolve_copy_into_target(self.table_name,
catalog=self._catalog, schema=self._schema)
+ if target != expected:
+ raise ValueError(
+ f"COPY INTO target {target._asdict()} resolved from
table_name={self.table_name!r}, "
+ f"catalog={self._catalog!r}, schema={self._schema!r} does not
match "
+ f"unity_table {expected._asdict()}."
+ )
+
+ hook = self._get_hook()
+ if hook.host != self.unity_table.host:
Review Comment:
Same issue for the host. `_parse_host` lowercases only URL-form hosts
(`urlsplit().hostname`) and leaves a bare hostname as written. So a connection
with `https://My-Workspace.cloud.databricks.com` becomes
`my-workspace.cloud.databricks.com`, while
`UnityTableIdentity(host="My-Workspace.cloud.databricks.com")` keeps its
capitals, and this check fails for the same workspace. The asset URI also
depends on how the user capitalised the host. Lowercasing the host in
`UnityTableIdentity.__post_init__` (and comparing case-insensitively here)
would fix both.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
