Skip to content

Commit

Permalink
Fixed the hardcoded default namespace value for GCP Data Fusion links. (
Browse files Browse the repository at this point in the history
#35379)

* update

* Add backward compatibility and cleanup

* fix formatting
  • Loading branch information
synsh committed Feb 3, 2024
1 parent 4e9aabb commit 11564a0
Show file tree
Hide file tree
Showing 2 changed files with 21 additions and 6 deletions.
17 changes: 12 additions & 5 deletions airflow/providers/google/cloud/links/datafusion.py
Expand Up @@ -15,7 +15,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
"""This module contains Google Compute Engine links."""
"""This module contains Google Data Fusion links."""
from __future__ import annotations

from typing import TYPE_CHECKING, ClassVar
Expand All @@ -30,8 +30,8 @@

BASE_LINK = "https://console.cloud.google.com/data-fusion"
DATAFUSION_INSTANCE_LINK = BASE_LINK + "/locations/{region}/instances/{instance_name}?project={project_id}"
DATAFUSION_PIPELINES_LINK = "{uri}/cdap/ns/default/pipelines"
DATAFUSION_PIPELINE_LINK = "{uri}/pipelines/ns/default/view/{pipeline_name}"
DATAFUSION_PIPELINES_LINK = "{uri}/cdap/ns/{namespace}/pipelines"
DATAFUSION_PIPELINE_LINK = "{uri}/pipelines/ns/{namespace}/view/{pipeline_name}"


class BaseGoogleLink(BaseOperatorLink):
Expand All @@ -52,10 +52,13 @@ def get_link(
ti_key: TaskInstanceKey,
) -> str:
conf = XCom.get_value(key=self.key, ti_key=ti_key)

if not conf:
return ""
if self.format_str.startswith("http"):
return self.format_str.format(**conf)

# Add a default value for the 'namespace' parameter for backward compatibility.
conf.setdefault("namespace", "default")

return self.format_str.format(**conf)


Expand Down Expand Up @@ -98,13 +101,15 @@ def persist(
task_instance: BaseOperator,
uri: str,
pipeline_name: str,
namespace: str,
):
task_instance.xcom_push(
context=context,
key=DataFusionPipelineLink.key,
value={
"uri": uri,
"pipeline_name": pipeline_name,
"namespace": namespace,
},
)

Expand All @@ -121,11 +126,13 @@ def persist(
context: Context,
task_instance: BaseOperator,
uri: str,
namespace: str,
):
task_instance.xcom_push(
context=context,
key=DataFusionPipelinesLink.key,
value={
"uri": uri,
"namespace": namespace,
},
)
10 changes: 9 additions & 1 deletion airflow/providers/google/cloud/operators/datafusion.py
Expand Up @@ -537,6 +537,7 @@ def execute(self, context: Context) -> None:
task_instance=self,
uri=instance["serviceEndpoint"],
pipeline_name=self.pipeline_name,
namespace=self.namespace,
)
self.log.info("Pipeline %s created", self.pipeline_name)

Expand Down Expand Up @@ -705,7 +706,12 @@ def execute(self, context: Context) -> dict:
)
self.log.info("Pipelines: %s", pipelines)

DataFusionPipelinesLink.persist(context=context, task_instance=self, uri=service_endpoint)
DataFusionPipelinesLink.persist(
context=context,
task_instance=self,
uri=service_endpoint,
namespace=self.namespace,
)
return pipelines


Expand Down Expand Up @@ -825,6 +831,7 @@ def execute(self, context: Context) -> str:
task_instance=self,
uri=instance["serviceEndpoint"],
pipeline_name=self.pipeline_name,
namespace=self.namespace,
)

if self.deferrable:
Expand Down Expand Up @@ -954,6 +961,7 @@ def execute(self, context: Context) -> None:
task_instance=self,
uri=instance["serviceEndpoint"],
pipeline_name=self.pipeline_name,
namespace=self.namespace,
)
hook.stop_pipeline(
pipeline_name=self.pipeline_name,
Expand Down

0 comments on commit 11564a0

Please sign in to comment.