sunank200 commented on code in PR #37307:
URL: https://github.com/apache/airflow/pull/37307#discussion_r1566585987
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,106 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
+ shards: int | None = None,
+ pods: int | None = None,
+ pod_type: str | None = "p1.x1",
+ metadata_config: dict | None = None,
+ source_collection: str | None = None,
+ environment: str | None = None,
+ ) -> PodSpec:
+ return PodSpec(
+ environment=environment or self.environment,
+ replicas=replicas,
+ shards=shards,
+ pods=pods,
+ pod_type=pod_type,
+ metadata_config=metadata_config,
+ source_collection=source_collection,
+ )
+
+ def get_serverless_spec_obj(self, cloud, region: str | None = None) ->
ServerlessSpec:
+ return ServerlessSpec(cloud=cloud, region=region or self.region)
+
def create_index(
+ self,
Review Comment:
What about other parameters which was there previously?
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,106 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
+ shards: int | None = None,
+ pods: int | None = None,
+ pod_type: str | None = "p1.x1",
+ metadata_config: dict | None = None,
+ source_collection: str | None = None,
+ environment: str | None = None,
+ ) -> PodSpec:
+ return PodSpec(
+ environment=environment or self.environment,
+ replicas=replicas,
+ shards=shards,
+ pods=pods,
+ pod_type=pod_type,
+ metadata_config=metadata_config,
+ source_collection=source_collection,
+ )
+
+ def get_serverless_spec_obj(self, cloud, region: str | None = None) ->
ServerlessSpec:
+ return ServerlessSpec(cloud=cloud, region=region or self.region)
+
def create_index(
+ self,
Review Comment:
What about other parameters which was there previously?
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -64,43 +69,73 @@ def get_ui_field_behaviour(cls) -> dict[str, Any]:
"""Return custom field behaviour."""
return {
"hidden_fields": ["port", "schema"],
- "relabeling": {"login": "Pinecone Environment", "password":
"Pinecone API key"},
+ "relabeling": {
+ "login": "Pinecone Environment",
+ "host": "Pinecone Host",
+ "password": "Pinecone API key",
+ },
}
- def __init__(self, conn_id: str = default_conn_name) -> None:
+ def __init__(
+ self, conn_id: str = default_conn_name, environment: str | None =
None, region: str | None = None
+ ) -> None:
self.conn_id = conn_id
- self.get_conn()
-
- def get_conn(self) -> None:
- pinecone_connection = self.get_connection(self.conn_id)
- api_key = pinecone_connection.password
- pinecone_environment = pinecone_connection.login
- pinecone_host = pinecone_connection.host
- extras = pinecone_connection.extra_dejson
+ self._environment = environment
+ self._region = region
+ self.conn = self.get_conn()
+
+ @property
+ def api_key(self) -> str:
+ key = self.conn.password
+ if not key:
+ raise LookupError("Pinecone API Key not found in connection")
+ return key
+
+ @cached_property
+ def environment(self) -> str:
+ if self._environment:
+ return self._environment
+ env = self.conn.login
+ if not env:
+ raise LookupError("Pinecone environment not found in connection")
+ return env
+
+ @cached_property
+ def region(self) -> str:
+ if self._region:
+ return self._region
+ region = self.conn.extra_dejson.get("region")
+ if not region:
+ raise LookupError("Pinecone region not found in connection")
+ return region
+
+ @cached_property
+ def pc(self) -> Pinecone:
+ """Pinecone object to interact with Pinecone."""
Review Comment:
Can this method be named better?
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,106 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
+ shards: int | None = None,
+ pods: int | None = None,
+ pod_type: str | None = "p1.x1",
+ metadata_config: dict | None = None,
+ source_collection: str | None = None,
+ environment: str | None = None,
+ ) -> PodSpec:
+ return PodSpec(
Review Comment:
doc-string missing
##########
airflow/providers/pinecone/hooks/pinecone.py:
##########
@@ -135,119 +170,106 @@ def upsert(
**kwargs,
)
- @staticmethod
+ def get_pod_spec_obj(
+ self,
+ replicas: int | None = None,
+ shards: int | None = None,
+ pods: int | None = None,
+ pod_type: str | None = "p1.x1",
+ metadata_config: dict | None = None,
+ source_collection: str | None = None,
+ environment: str | None = None,
+ ) -> PodSpec:
+ return PodSpec(
+ environment=environment or self.environment,
+ replicas=replicas,
+ shards=shards,
+ pods=pods,
+ pod_type=pod_type,
+ metadata_config=metadata_config,
+ source_collection=source_collection,
+ )
+
+ def get_serverless_spec_obj(self, cloud, region: str | None = None) ->
ServerlessSpec:
+ return ServerlessSpec(cloud=cloud, region=region or self.region)
+
def create_index(
+ self,
Review Comment:
What about other parameters which was there previously? Is it not required?
--
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
