Lee-W commented on code in PR #31925:
URL: https://github.com/apache/airflow/pull/31925#discussion_r1251815031
##########
airflow/providers/google/cloud/hooks/spanner.py:
##########
@@ -70,6 +85,34 @@ def _get_client(self, project_id: str) -> Client:
)
return self._client
+ def _get_conn_params(self) -> SpannerConnectionParams:
+ """Extract spanner database connection parameters."""
+ extras = self.get_connection(self.gcp_conn_id).extra_dejson
+ project_id = get_field(extras, "project_id") or self.project_id
+ instance_id = get_field(extras, "instance_id")
+ database_id = get_field(extras, "database_id")
+ return SpannerConnectionParams(project_id, instance_id, database_id)
+
+ def get_uri(self) -> str:
+ """Override DbApiHook get_uri method for get_sqlalchemy_engine()."""
+ project_id, instance_id, database_id = self._get_conn_params()
+ if not instance_id or not database_id:
+ raise AirflowException("The instance_id or database_id were not
specified")
+ return
f"spanner+spanner:///projects/{project_id}/instances/{instance_id}/databases/{database_id}"
+
+ def get_sqlalchemy_engine(self, engine_kwargs=None):
Review Comment:
nitpick: Should we add a return type annotation here?
##########
airflow/providers/google/cloud/hooks/spanner.py:
##########
@@ -70,6 +85,34 @@ def _get_client(self, project_id: str) -> Client:
)
return self._client
+ def _get_conn_params(self) -> SpannerConnectionParams:
+ """Extract spanner database connection parameters."""
+ extras = self.get_connection(self.gcp_conn_id).extra_dejson
+ project_id = get_field(extras, "project_id") or self.project_id
+ instance_id = get_field(extras, "instance_id")
+ database_id = get_field(extras, "database_id")
+ return SpannerConnectionParams(project_id, instance_id, database_id)
+
+ def get_uri(self) -> str:
+ """Override DbApiHook get_uri method for get_sqlalchemy_engine()."""
+ project_id, instance_id, database_id = self._get_conn_params()
+ if not instance_id or not database_id:
Review Comment:
nitpick: Should we use `it not all([instance_id, database_id])`?
--
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: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
