Operators (and Hooks) of the google-provider should allow custom api_endpoints #53596
Description
Description
Google's client library allows to set the api_endpoint. Unfortunately the Hooks and Operators in the google-provider do not forward this option. It would be great, if they would.
As a workaround we override the Hook class to set this option, and the Operator class to use our custom Hook, e.g.:
class CustomApiPubSubHook(PubSubHook):
def get_conn(self) -> PublisherClient:
if not self._client:
self._client = PublisherClient(
credentials=self.get_credentials(),
client_info=CLIENT_INFO,
# set api_endpoint:
client_options=ClientOptions(api_endpoint="pubsub-custom.googleapis.com"),
publisher_options=PublisherOptions(
enable_message_ordering=self.enable_message_ordering,
),
)
return self._client
class CustomApiPubSubPublishMessageOperator(PubSubPublishMessageOperator):
def execute(self, context: Context) -> None:
# use custom Hook
hook = CustomApiPubSubHook(
gcp_conn_id=self.gcp_conn_id,
impersonation_chain=self.impersonation_chain,
enable_message_ordering=self.enable_message_ordering,
)
self.log.info("Publishing to topic %s", self.topic)
hook.publish(project_id=self.project_id, topic=self.topic, messages=self.messages)
self.log.info("Published to topic %s", self.topic)This could be avoided, if the Hook and the Operator would accept an additional argument "api_endpoint" and the Operator would pass it to the Hook.
And, of course, this should be implemented in every Operator/Hook, not just Pub/Sub.
Use case/motivation
Due to our network setup we have to use a private service connect (PSC) to connect to the GCP. I.e. we have to use a different API endpoint.
Related issues
No response
Are you willing to submit a PR?
- Yes I am willing to submit a PR!
Code of Conduct
- I agree to follow this project's Code of Conduct