You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "ellisms (via GitHub)" <gi...@apache.org> on 2026/03/26 23:46:15 UTC
[PR] Neptune analytics [airflow]
ellisms opened a new pull request, #64274:
URL: https://github.com/apache/airflow/pull/64274
<!-- SPDX-License-Identifier: Apache-2.0
https://www.apache.org/licenses/LICENSE-2.0 -->
<!--
Thank you for contributing!
Please provide above a brief description of the changes made in this pull request.
Write a good git commit message following this guide: http://chris.beams.io/posts/git-commit/
Please make sure that your code changes are covered with tests.
And in case of new features or big changes remember to adjust the documentation.
Feel free to ping (in general) for the review if you do not see reaction for a few days
(72 Hours is the minimum reaction time you can expect from volunteers) - we sometimes miss notifications.
In case of an existing issue, reference it using one of the following:
* closes: #ISSUE
* related: #ISSUE
-->
---
closes: #60975
Added Neptune Analytics operators to create and delete graphs, create and delete private endpoints, and create/cancel import tasks.
##### Was generative AI tooling used to co-author this PR?
<!--
If generative AI tooling has been used in the process of authoring this PR, please
change below checkbox to `[X]` followed by the name of the tool, uncomment the "Generated-by".
-->
- [ x ] Yes (please specify the tool below)
- Generated-by: [Kiro] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions). Kiro used to generate unit tests, not the operator code.
<!--
Generated-by: [Kiro] following [the guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#gen-ai-assisted-contributions)
-->
---
* Read the **[Pull Request Guidelines](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-guidelines)** for more information. Note: commit author/co-author name and email in commits become permanently public when merged.
* For fundamental code changes, an Airflow Improvement Proposal ([AIP](https://cwiki.apache.org/confluence/display/AIRFLOW/Airflow+Improvement+Proposals)) is needed.
* When adding dependency, check compliance with the [ASF 3rd Party License Policy](https://www.apache.org/legal/resolved.html#category-x).
* For significant user-facing changes create newsfragment: `{pr_number}.significant.rst`, in [airflow-core/newsfragments](https://github.com/apache/airflow/tree/main/airflow-core/newsfragments). You can add this file in a follow-up commit after the PR is created so you know the PR number.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3066501486
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Neptune graph % complete", self.graph_id)
+
Review Comment:
`execute_complete` logs with an invalid `%` format string ("Neptune graph % complete"), which will raise a formatting error at runtime. Update the log message to use a proper placeholder (e.g. `%s`) and also validate the deferred `event` (raise `AirflowException` on non-success) so deferrable failures don’t get treated as success.
```suggestion
if event is None:
raise AirflowException("No event received while waiting for Neptune graph creation to complete.")
if event.get("status") != "success":
raise AirflowException(
event.get(
"message",
f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
)
)
self.log.info("Neptune graph %s complete", self.graph_id)
```
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Neptune graph % complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise AirflowException(f"Private endpoint failed to create for graph {self.graph_identifier}")
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
Review Comment:
The deferrable triggers emit `event['status']` and error messages on failure, but `execute_complete` here ignores the event entirely and returns success unconditionally. Please validate the event (e.g. require `status == 'success'` and raise `AirflowException` otherwise) so deferrable endpoint creation failures correctly fail the task.
##########
providers/amazon/docs/operators/neptune_analytics.rst:
##########
@@ -0,0 +1,148 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+========================
+Amazon Neptune Analytics
+========================
+
+`Amazon Neptune Analytics <https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html>`__ is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+Prerequisite Tasks
+------------------
+
+.. include:: ../_partials/prerequisite_tasks.rst
+
+Generic Parameters
+------------------
+
+.. include:: ../_partials/generic_parameters.rst
+
+Operators
+---------
+
+.. _howto/operator:NeptuneCreateGraphOperator:
+
+Create a new Neptune Graph
+==========================
+
+To create a new Neptune Analytics Graph, you can use
+:class:`~airflow.providers.amazon.aws.operators.neptune.NeptuneCreateGraphOperator`.
+This operator can be run in deferrable mode by passing ``deferrable=True`` as a parameter. This requires
+the aiobotocore module to be installed.
+
Review Comment:
The `:class:` references point to `airflow.providers.amazon.aws.operators.neptune.*`, but the Neptune Analytics operators in this PR live in `airflow.providers.amazon.aws.operators.neptune_analytics`. Update these references throughout this doc so the rendered documentation links to the correct classes.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Neptune graph % complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise AirflowException(f"Private endpoint failed to create for graph {self.graph_identifier}")
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise AirflowException(f"Failed to delete private endpoint {endpoint_id}")
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and event.get("status") == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
Review Comment:
`execute_complete` doesn’t validate the deferred `event` status (and effectively becomes a no-op on failure). For deferrable runs, please check `event['status']` and raise on non-success so trigger failures don’t silently pass.
```suggestion
if event is None:
raise AirflowException("Private graph endpoint deletion did not return an event.")
status = event.get("status")
if status != "success":
error_message = event.get("message") or event.get("error") or str(event)
raise AirflowException(f"Failed to delete private endpoint: {error_message}")
vpc_endpoint_id = event.get("endpoint_id", "Unknown")
self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
```
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Neptune graph % complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise AirflowException(f"Private endpoint failed to create for graph {self.graph_identifier}")
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise AirflowException(f"Failed to delete private endpoint {endpoint_id}")
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and event.get("status") == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise AirflowException(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ if event:
+ graph_id = event.get("graph_id", "Unknown")
+
+ self.log.info("Neptune graph %s deleted", graph_id)
+
+
+class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph and imports data into it.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics,
+ you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphWithImportOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param source: The source from which to import data. Can be an S3 URI or Neptune database snapshot.
+ :param role_arn: The ARN of the IAM role that Neptune Analytics can assume to access the data source.
+ :param blank_node_handling: The method to handle blank nodes in the dataset. Options include 'convertToIri' or other handling strategies.
+ :param parquet_type: The type of Parquet files in the data source (if applicable).
+ :param format: The format of the data to be imported (e.g., 'csv', 'opencypher', 'ntriples', 'nquads', 'rdfxml', 'turtle').
+ :param min_provisioned_memory: The minimum provisioned memory for the graph in GBs.
+ :param max_provisioned_memory: The maximum provisioned memory for the graph in GBs.
+ :param fail_on_error: If True, the import will fail if any errors are encountered. If False, the import will continue despite errors.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled. (default: False)
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param import_options: Contains options for controlling the import process.
+ :param wait_for_completion: Whether to wait for the graph to be created and data imported. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 60)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be created and data imported.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "source", "role_arn", "kms_key"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (
+ NeptuneImportTaskLink(),
+ NeptuneGraphLink(),
+ )
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ source: str,
+ role_arn: str,
+ blank_node_handling: str | None = None,
+ parquet_type: str | None = None,
+ format: str | None = None,
+ min_provisioned_memory: int | None = None,
+ max_provisioned_memory: int | None = None,
+ fail_on_error: bool | None = None,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool | None = None,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ import_options: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.source = source
+ self.role_arn = role_arn
+ self.blank_node_handling = blank_node_handling
+ self.parquet_type = parquet_type
+ self.format = format
+ self.min_provisioned_memory = min_provisioned_memory
+ self.max_provisioned_memory = max_provisioned_memory
+ self.fail_on_error = fail_on_error
+ self.public_connectivity = public_connectivity
+ self.replica_count = replica_count
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.import_options = import_options
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s with import", self.graph_name)
+
+ # Build the import options
+ import_options = {
+ "neptune-analytics:blank-node-handling": self.blank_node_handling,
+ "neptune-analytics:parquet-type": self.parquet_type,
+ }
+
+ # Remove None values from import_options
+ import_options = {k: v for k, v in import_options.items() if v is not None}
+
+ # Merge with user-provided import_options
+ if self.import_options:
+ import_options.update(self.import_options)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "source": self.source,
+ "roleArn": self.role_arn,
+ **{
+ k: v
+ for k, v in {
+ "format": self.format,
+ "minProvisionedMemory": self.min_provisioned_memory,
+ "maxProvisionedMemory": self.max_provisioned_memory,
+ "failOnError": self.fail_on_error,
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ "importOptions": import_options if import_options else None,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph_using_import_task(**create_params)
+
+ self.log.info("Graph %s import task in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("graphId", None)
+ import_task_id = response.get("taskId")
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="defer_wait_for_task",
+ kwargs={"import_task_id": import_task_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ # Once the graph is available, wait for the task to complete
+
+ self.log.info("Waiting for import task %s", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def defer_wait_for_task(
+ self, import_task_id: str, context: Context, event: dict[str, Any] | None = None
+ ) -> None:
+ """Defers for import task completion."""
+ self.log.info("Deferring for import task %s completion", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
Review Comment:
`defer_wait_for_task` has the wrong parameter order for a deferral resume callback. Airflow calls the next method as `method(context, **next_kwargs)` (with `event` injected into `next_kwargs`), so having `import_task_id` as the first positional arg will cause `TypeError: multiple values for argument 'import_task_id'` at resume time. Change the signature to accept `context` first (and `event` as a kwarg), and keep `import_task_id` as a kwarg-only/optional parameter.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Neptune graph % complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise AirflowException(f"Private endpoint failed to create for graph {self.graph_identifier}")
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise AirflowException(f"Failed to delete private endpoint {endpoint_id}")
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and event.get("status") == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise AirflowException(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ if event:
+ graph_id = event.get("graph_id", "Unknown")
+
+ self.log.info("Neptune graph %s deleted", graph_id)
+
+
+class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph and imports data into it.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics,
+ you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphWithImportOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param source: The source from which to import data. Can be an S3 URI or Neptune database snapshot.
+ :param role_arn: The ARN of the IAM role that Neptune Analytics can assume to access the data source.
+ :param blank_node_handling: The method to handle blank nodes in the dataset. Options include 'convertToIri' or other handling strategies.
+ :param parquet_type: The type of Parquet files in the data source (if applicable).
+ :param format: The format of the data to be imported (e.g., 'csv', 'opencypher', 'ntriples', 'nquads', 'rdfxml', 'turtle').
+ :param min_provisioned_memory: The minimum provisioned memory for the graph in GBs.
+ :param max_provisioned_memory: The maximum provisioned memory for the graph in GBs.
+ :param fail_on_error: If True, the import will fail if any errors are encountered. If False, the import will continue despite errors.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled. (default: False)
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param import_options: Contains options for controlling the import process.
+ :param wait_for_completion: Whether to wait for the graph to be created and data imported. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 60)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be created and data imported.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "source", "role_arn", "kms_key"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (
+ NeptuneImportTaskLink(),
+ NeptuneGraphLink(),
+ )
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ source: str,
+ role_arn: str,
+ blank_node_handling: str | None = None,
+ parquet_type: str | None = None,
+ format: str | None = None,
+ min_provisioned_memory: int | None = None,
+ max_provisioned_memory: int | None = None,
+ fail_on_error: bool | None = None,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool | None = None,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ import_options: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.source = source
+ self.role_arn = role_arn
+ self.blank_node_handling = blank_node_handling
+ self.parquet_type = parquet_type
+ self.format = format
+ self.min_provisioned_memory = min_provisioned_memory
+ self.max_provisioned_memory = max_provisioned_memory
+ self.fail_on_error = fail_on_error
+ self.public_connectivity = public_connectivity
+ self.replica_count = replica_count
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.import_options = import_options
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s with import", self.graph_name)
+
+ # Build the import options
+ import_options = {
+ "neptune-analytics:blank-node-handling": self.blank_node_handling,
+ "neptune-analytics:parquet-type": self.parquet_type,
+ }
+
+ # Remove None values from import_options
+ import_options = {k: v for k, v in import_options.items() if v is not None}
+
+ # Merge with user-provided import_options
+ if self.import_options:
+ import_options.update(self.import_options)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "source": self.source,
+ "roleArn": self.role_arn,
+ **{
+ k: v
+ for k, v in {
+ "format": self.format,
+ "minProvisionedMemory": self.min_provisioned_memory,
+ "maxProvisionedMemory": self.max_provisioned_memory,
+ "failOnError": self.fail_on_error,
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ "importOptions": import_options if import_options else None,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph_using_import_task(**create_params)
+
+ self.log.info("Graph %s import task in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("graphId", None)
+ import_task_id = response.get("taskId")
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="defer_wait_for_task",
+ kwargs={"import_task_id": import_task_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ # Once the graph is available, wait for the task to complete
+
+ self.log.info("Waiting for import task %s", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def defer_wait_for_task(
+ self, import_task_id: str, context: Context, event: dict[str, Any] | None = None
+ ) -> None:
+ """Defers for import task completion."""
+ self.log.info("Deferring for import task %s completion", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Import complete for graph %s", self.graph_id)
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Starts a bulk data import task to load data into an empty Neptune graph.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneStartImportTaskOperator`
+
+ :param graph_identifier: Graph Id of target Neptune Graph
+ :param role_arn: IAM role ARN granting access to source data
+ :param source: URL identifying the source data location.
+ :param blank_node_handling: Method to handle blank nodes in dataset.
+ :param fail_on_error: If set to true, the task halts when an import error is encountered. If set to false, the task skips the data that caused the error and continues if possible.
+ :param format: Specifies the format of the Amazon S3 data to be imported.
+ :param import_options: Options on how to perform an import
+ :param parquet_type: Parquet type of import task
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "role_arn", "source", "import_options"
+ )
+ template_fields_renderers = {
+ "import_options": "json",
+ }
+ operator_extra_links = (NeptuneImportTaskLink(),)
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ role_arn: str,
+ source: str,
+ blank_node_handling: str | None = None,
+ fail_on_error: bool = True,
+ format: str | None = None,
+ import_options: dict | None = None,
+ parquet_type: str | None = "COLUMNAR",
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.role_arn = role_arn
+ self.source = source
+ self.blank_node_handling = blank_node_handling
+ self.fail_on_error = fail_on_error
+ self.format = format
+ self.import_options = import_options
+ self.parquet_type = parquet_type
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Starting data import to graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ "roleArn": self.role_arn,
+ "source": self.source,
+ **{
+ k: v
+ for k, v in {
+ "blankNodeHandling": self.blank_node_handling,
+ "failOnError": self.fail_on_error,
+ "format": self.format,
+ "importOptions": self.import_options,
+ "parquetType": self.parquet_type,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.start_import_task(**create_params)
+ import_task_id = response.get("taskId")
+ self.log.info("Import task %s started for graph %s", import_task_id, self.graph_identifier)
+
+ # Create the console link
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s completes", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to complete", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": import_task_id, "graph_id": self.graph_identifier}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ task_id = ""
+ if event:
+ task_id = event.get("import_task_id", "")
+
+ return {"graph_id": self.graph_identifier, "import_task_id": task_id}
+
+
+class NeptuneCancelImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Cancels an active Neptune Graph import task.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCancelImportTaskOperator`
+
+ :param import_task_id: Neptune Graph import task id to cancel.
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("import_task_id")
+
+ def __init__(
+ self,
+ import_task_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.import_task_id = import_task_id
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Cancelling import task %s", self.import_task_id)
+
+ response = self.hook.conn.cancel_import_task(taskIdentifier=self.import_task_id)
+
+ self.log.info("Import task %s status is %s", self.import_task_id, response.get("status", "Unknown"))
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s is cancelled", self.import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCancelledTrigger(
+ task_identifier=self.import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to be cancelled", self.import_task_id)
+ self.hook.get_waiter("import_task_cancelled").wait(
+ taskIdentifier=self.import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": self.import_task_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ task_id = ""
+ if event:
+ task_id = event.get("import_task_id", "")
+ self.log.info("Import task %s cancelled", task_id)
+
Review Comment:
`NeptuneCancelImportTaskOperator.execute_complete` looks for `event['import_task_id']`, but `NeptuneImportTaskCancelledTrigger` emits the id under the key `task_identifier` (its `return_key`). This means deferrable cancellations will always log/return an empty id. Align the trigger/operator to use the same key (prefer `import_task_id` for consistency) and validate `event['status']` to raise on errors.
```suggestion
if event is None:
raise AirflowException("No event received while waiting for Neptune import task cancellation.")
status = str(event.get("status", "")).lower()
if status not in {"success", "cancelled", "canceled"}:
message = event.get("message") or event.get("error") or f"Unexpected trigger status: {status!r}"
raise AirflowException(f"Error while waiting for Neptune import task cancellation: {message}")
task_id = event.get("import_task_id") or event.get("task_identifier", "")
self.log.info("Import task %s cancelled", task_id)
```
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
Review Comment:
The docstring says the operator returns a dictionary with “graph id and vpc id”, but `execute()` returns only `{'graph_id': ...}`. Please correct the return documentation (or return the additional fields if intended).
```suggestion
:return: dictionary with Neptune graph id
```
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,340 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+import time
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.common.compat.sdk import DAG, chain, task
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_neptune_analytics"
+
+ROLE_ARN_KEY = "ROLE_ARN"
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
+
+# Minimal OpenCypher CSV data for import testing.
+NODES_CSV = """~id,~label,name:String
+n1,Person,Alice
+n2,Person,Bob
+"""
+
+EDGES_CSV = """~id,~from,~to,~label
+e1,n1,n2,KNOWS
+"""
+
+NEPTUNE_ANALYTICS_TRUST_POLICY = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Principal": {"Service": "neptune-graph.amazonaws.com"},
+ "Action": "sts:AssumeRole",
+ }
+ ],
+ }
+)
+
+S3_READ_POLICY_DOCUMENT = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Action": ["s3:GetObject", "s3:ListBucket"],
+ "Resource": ["arn:aws:s3:::*", "arn:aws:s3:::*/*"],
+ }
+ ],
+ }
+)
+
+
+@task
+def create_neptune_import_role(role_name: str) -> str:
+ iam_client = boto3.client("iam")
+ iam_client.create_role(
+ RoleName=role_name,
+ AssumeRolePolicyDocument=NEPTUNE_ANALYTICS_TRUST_POLICY,
+ Description="Role for Neptune Analytics import system test",
+ )
+ iam_client.put_role_policy(
+ RoleName=role_name,
+ PolicyName="NeptuneAnalyticsS3Access",
+ PolicyDocument=S3_READ_POLICY_DOCUMENT,
+ )
+ role = iam_client.get_role(RoleName=role_name)
+ time.sleep(60) # Wait for IAM eventual consistency (role + inline policy propagation)
+ return role["Role"]["Arn"]
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_neptune_import_role(role_name: str) -> None:
+ iam_client = boto3.client("iam")
+ try:
+ iam_client.delete_role_policy(RoleName=role_name, PolicyName="NeptuneAnalyticsS3Access")
+ except iam_client.exceptions.NoSuchEntityException:
+ pass
+ try:
+ iam_client.delete_role(RoleName=role_name)
+ except iam_client.exceptions.NoSuchEntityException:
+ pass
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_graph_if_exists(graph_name: str) -> None:
+ """Safety net to clean up the graph in case a previous task failed."""
+ hook = NeptuneAnalyticsHook()
+ try:
+ # List graphs and find by name
+ paginator = hook.conn.get_paginator("list_graphs")
+ for page in paginator.paginate():
+ for graph in page.get("graphs", []):
+ if graph.get("name") == graph_name:
+ graph_id = graph["id"]
+ # Disable deletion protection if enabled
+ try:
+ hook.conn.update_graph(graphIdentifier=graph_id, deletionProtection=False)
+ except Exception:
+ pass
+ hook.conn.delete_graph(graphIdentifier=graph_id, skipSnapshot=True)
+ hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=graph_id,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ return
+ except Exception:
+ pass
+
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ env_id = test_context["ENV_ID"]
+ graph_name = f"{env_id}-graph"
+ import_graph_name = f"{env_id}-import-graph"
+ bucket_name = f"{env_id}-neptune-analytics"
+ import_role_name = f"{env_id}-neptune-import"
+ region = boto3.session.Session().region_name
+
+ # --- TEST SETUP ---
+
+ create_bucket = S3CreateBucketOperator(
+ task_id="create_bucket",
+ bucket_name=bucket_name,
+ )
+
+ upload_nodes = S3CreateObjectOperator(
+ task_id="upload_nodes",
+ s3_bucket=bucket_name,
+ s3_key="data/nodes.csv",
+ data=NODES_CSV,
+ replace=True,
+ )
+
+ upload_edges = S3CreateObjectOperator(
+ task_id="upload_edges",
+ s3_bucket=bucket_name,
+ s3_key="data/edges.csv",
+ data=EDGES_CSV,
+ replace=True,
+ )
+
+ create_role = create_neptune_import_role(import_role_name)
+
+ # --- TEST BODY ---
+
+ # [START howto_operator_neptune_analytics_create_graph]
+ create_graph = NeptuneCreateGraphOperator(
+ task_id="create_graph",
+ graph_name=graph_name,
+ vector_search_config={"dimension": 128},
+ provisioned_memory=32,
+ public_connectivity=True,
+ replica_count=0,
+ deletion_protection=False,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_create_graph]
+
+ # [START howto_operator_neptune_analytics_create_private_endpoint]
+ create_endpoint = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="create_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph)['graph_id']}}",
+ wait_for_completion=True,
+ )
+ # [END howto_operator_neptune_analytics_create_private_endpoint]
+
+ # [START howto_operator_neptune_analytics_delete_private_endpoint]
+ delete_endpoint = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="delete_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ vpc_id="{{ ti.xcom_pull(task_ids='create_endpoint')['vpc_id'] }}",
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_delete_private_endpoint]
+
+ # [START howto_operator_neptune_analytics_start_import_task]
+ start_import = NeptuneStartImportTaskOperator(
+ task_id="start_import",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ role_arn=create_role,
+ source=f"s3://{bucket_name}/data/",
+ format="CSV",
+ fail_on_error=True,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_start_import_task]
+
+ # [START howto_operator_neptune_analytics_cancel_import_task]
+ cancel_import = NeptuneCancelImportTaskOperator(
+ task_id="cancel_import",
+ import_task_id="{{ ti.xcom_pull(task_ids='start_import')['import_task_id']}}",
+ wait_for_completion=True,
+ aws_conn_id="aws_default",
+ )
+ # [END howto_operator_neptune_analytics_cancel_import_task]
+
+ # [START howto_operator_neptune_analytics_delete_graph]
+ delete_graph = NeptuneDeleteGraphOperator(
+ task_id="delete_graph",
+ graph_id="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ skip_snapshot=True,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_delete_graph]
+
+ # [START howto_operator_neptune_analytics_create_graph_with_import]
+ create_graph_with_import = NeptuneCreateGraphWithImportOperator(
+ task_id="create_graph_with_import",
+ graph_name=import_graph_name,
+ vector_search_config={"dimension": 128},
+ source=f"s3://{bucket_name}/data/",
+ role_arn=create_role,
+ format="CSV",
+ fail_on_error=True,
+ public_connectivity=True,
+ replica_count=0,
+ deletion_protection=False,
+ min_provisioned_memory=32,
+ max_provisioned_memory=32,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_create_graph_with_import]
+
+ # [START howto_operator_neptune_analytics_delete_import_graph]
+ delete_import_graph = NeptuneDeleteGraphOperator(
+ task_id="delete_import_graph",
+ graph_id="{{ ti.xcom_pull(task_ids='create_graph_with_import')['graph_id'] }}",
+ skip_snapshot=True,
+ wait_for_completion=True,
+ deferrable=False,
+ trigger_rule=TriggerRule.ALL_DONE,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_delete_import_graph]
+
+ # --- TEST TEARDOWN ---
+
+ delete_bucket = S3DeleteBucketOperator(
+ task_id="delete_bucket",
+ trigger_rule=TriggerRule.ALL_DONE,
+ bucket_name=bucket_name,
+ force_delete=True,
+ )
+
+ delete_role = delete_neptune_import_role(import_role_name)
+
+ cleanup_graph = delete_graph_if_exists(graph_name)
+ cleanup_import_graph = delete_graph_if_exists(import_graph_name)
+
+ chain(
+ # TEST SETUP
+ test_context,
+ create_bucket,
+ [upload_nodes, upload_edges],
+ create_role,
+ # TEST BODY: Create graph, import data, then delete
+ create_graph,
+ start_import,
+ delete_graph,
+ # TEST BODY: Create graph with import, then delete
+ create_graph_with_import,
+ delete_import_graph,
+ # TEST TEARDOWN
+ [cleanup_graph, cleanup_import_graph],
+ delete_bucket,
+ delete_role,
+ )
Review Comment:
`create_endpoint`, `delete_endpoint`, and `cancel_import` are defined but never wired into `chain()`. Unchained tasks will be scheduled independently (and `create_endpoint` depends on XComs from `create_graph`), leading to race/failures in the system test. Either add them to the dependency chain in the right order or remove them from the DAG if they’re only meant as docs snippets.
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,340 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import json
+import time
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.common.compat.sdk import DAG, chain, task
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_neptune_analytics"
+
+ROLE_ARN_KEY = "ROLE_ARN"
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
+
+# Minimal OpenCypher CSV data for import testing.
+NODES_CSV = """~id,~label,name:String
+n1,Person,Alice
+n2,Person,Bob
+"""
+
+EDGES_CSV = """~id,~from,~to,~label
+e1,n1,n2,KNOWS
+"""
+
+NEPTUNE_ANALYTICS_TRUST_POLICY = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Principal": {"Service": "neptune-graph.amazonaws.com"},
+ "Action": "sts:AssumeRole",
+ }
+ ],
+ }
+)
+
+S3_READ_POLICY_DOCUMENT = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Action": ["s3:GetObject", "s3:ListBucket"],
+ "Resource": ["arn:aws:s3:::*", "arn:aws:s3:::*/*"],
+ }
+ ],
+ }
+)
+
+
+@task
+def create_neptune_import_role(role_name: str) -> str:
+ iam_client = boto3.client("iam")
+ iam_client.create_role(
+ RoleName=role_name,
+ AssumeRolePolicyDocument=NEPTUNE_ANALYTICS_TRUST_POLICY,
+ Description="Role for Neptune Analytics import system test",
+ )
+ iam_client.put_role_policy(
+ RoleName=role_name,
+ PolicyName="NeptuneAnalyticsS3Access",
+ PolicyDocument=S3_READ_POLICY_DOCUMENT,
+ )
+ role = iam_client.get_role(RoleName=role_name)
+ time.sleep(60) # Wait for IAM eventual consistency (role + inline policy propagation)
+ return role["Role"]["Arn"]
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_neptune_import_role(role_name: str) -> None:
+ iam_client = boto3.client("iam")
+ try:
+ iam_client.delete_role_policy(RoleName=role_name, PolicyName="NeptuneAnalyticsS3Access")
+ except iam_client.exceptions.NoSuchEntityException:
+ pass
+ try:
+ iam_client.delete_role(RoleName=role_name)
+ except iam_client.exceptions.NoSuchEntityException:
+ pass
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_graph_if_exists(graph_name: str) -> None:
+ """Safety net to clean up the graph in case a previous task failed."""
+ hook = NeptuneAnalyticsHook()
+ try:
+ # List graphs and find by name
+ paginator = hook.conn.get_paginator("list_graphs")
+ for page in paginator.paginate():
+ for graph in page.get("graphs", []):
+ if graph.get("name") == graph_name:
+ graph_id = graph["id"]
+ # Disable deletion protection if enabled
+ try:
+ hook.conn.update_graph(graphIdentifier=graph_id, deletionProtection=False)
+ except Exception:
+ pass
+ hook.conn.delete_graph(graphIdentifier=graph_id, skipSnapshot=True)
+ hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=graph_id,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ return
+ except Exception:
+ pass
+
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ env_id = test_context["ENV_ID"]
+ graph_name = f"{env_id}-graph"
+ import_graph_name = f"{env_id}-import-graph"
+ bucket_name = f"{env_id}-neptune-analytics"
+ import_role_name = f"{env_id}-neptune-import"
+ region = boto3.session.Session().region_name
+
+ # --- TEST SETUP ---
+
+ create_bucket = S3CreateBucketOperator(
+ task_id="create_bucket",
+ bucket_name=bucket_name,
+ )
+
+ upload_nodes = S3CreateObjectOperator(
+ task_id="upload_nodes",
+ s3_bucket=bucket_name,
+ s3_key="data/nodes.csv",
+ data=NODES_CSV,
+ replace=True,
+ )
+
+ upload_edges = S3CreateObjectOperator(
+ task_id="upload_edges",
+ s3_bucket=bucket_name,
+ s3_key="data/edges.csv",
+ data=EDGES_CSV,
+ replace=True,
+ )
+
+ create_role = create_neptune_import_role(import_role_name)
+
+ # --- TEST BODY ---
+
+ # [START howto_operator_neptune_analytics_create_graph]
+ create_graph = NeptuneCreateGraphOperator(
+ task_id="create_graph",
+ graph_name=graph_name,
+ vector_search_config={"dimension": 128},
+ provisioned_memory=32,
+ public_connectivity=True,
+ replica_count=0,
+ deletion_protection=False,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_create_graph]
+
+ # [START howto_operator_neptune_analytics_create_private_endpoint]
+ create_endpoint = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="create_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph)['graph_id']}}",
Review Comment:
The Jinja template for `graph_identifier` is malformed (missing a closing quote/bracket) and will break DAG parsing/rendering. Fix the `ti.xcom_pull(...)` expression to be valid Jinja/Python indexing syntax.
```suggestion
graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
```
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/neptune_analytics.py:
##########
@@ -0,0 +1,282 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.triggers.base import AwsBaseWaiterTrigger
+
+if TYPE_CHECKING:
+ from airflow.providers.amazon.aws.hooks.base_aws import AwsGenericHook
+
+
+class NeptuneGraphAvailableTrigger(AwsBaseWaiterTrigger):
+ """
+ Triggers when a Neptune graph is available.
+
+ :param graph_id: Graph ID to poll.
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: AWS region name (example: us-east-1)
+ """
+
+ def __init__(
+ self,
+ *,
+ graph_id: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ **kwargs,
+ ) -> None:
+ super().__init__(
+ serialized_fields={"graph_id": graph_id},
+ waiter_name="graph_available",
+ waiter_args={"graphIdentifier": graph_id},
+ failure_message="Failed to create Neptune graph",
+ status_message="Status of Neptune graph is",
+ status_queries=["status"],
+ return_key="graph_id",
+ return_value=graph_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ **kwargs,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return NeptuneAnalyticsHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
+
+class NeptuneGraphPrivateEndpointAvailableTrigger(AwsBaseWaiterTrigger):
+ """
+ Triggers when a Neptune Graph private endpoint is available.
+
+ :param graph_id: Graph Id waiting for the endpoint
+ :param vpc_id: VPC id where endpoint is creating
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: AWS region name (example: us-east-1)
+ """
+
+ def __init__(
+ self,
+ *,
+ graph_id: str,
+ vpc_id: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ **kwargs,
+ ) -> None:
+ super().__init__(
+ serialized_fields={"graph_id": graph_id, "vpc_id": vpc_id},
+ waiter_name="private_graph_endpoint_available",
+ waiter_args={"graphIdentifier": graph_id, "vpcId": vpc_id},
+ failure_message="Failed to create Neptune graph endpoint",
+ status_message="Status of Neptune graph endpoint is",
+ status_queries=["status"],
+ return_key="graph_id",
+ return_value=graph_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ **kwargs,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return NeptuneAnalyticsHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
+
+class NeptuneGraphPrivateEndpointDeletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Triggers when a Neptune Graph private endpoint is deleted.
+
+ :param graph_id: Graph Id of the endpoint
+ :param vpc_id: VPC id where endpoint resides
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: AWS region name (example: us-east-1)
+ """
+
+ def __init__(
+ self,
+ *,
+ graph_id: str,
+ vpc_id: str,
+ endpoint_id: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ **kwargs,
+ ) -> None:
+ super().__init__(
+ serialized_fields={"graph_id": graph_id, "vpc_id": vpc_id, "endpoint_id": endpoint_id},
+ waiter_name="private_graph_endpoint_deleted",
+ waiter_args={"graphIdentifier": graph_id, "vpcId": vpc_id},
+ failure_message="Failed to delete Neptune graph endpoint",
+ status_message="Status of Neptune graph endpoint is",
+ status_queries=["status"],
+ return_key="endpoint_id",
+ return_value=endpoint_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ **kwargs,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return NeptuneAnalyticsHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
+
+class NeptuneGraphDeletedTrigger(AwsBaseWaiterTrigger):
+ """
+ Triggers when a Neptune Graph is deleted.
+
+ :param graph_id: Graph Id to be deleted
+ :param waiter_delay: The amount of time in seconds to wait between attempts.
+ :param waiter_max_attempts: The maximum number of attempts to be made.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ :param region_name: AWS region name (example: us-east-1)
+ """
+
+ def __init__(
+ self,
+ *,
+ graph_id: str,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ **kwargs,
+ ) -> None:
+ super().__init__(
+ serialized_fields={"graph_id": graph_id},
+ waiter_name="graph_deleted",
+ waiter_args={"graphIdentifier": graph_id},
+ failure_message="Failed to delete Neptune graph",
+ status_message="Status of Neptune graph is",
+ status_queries=["status"],
+ return_key="graph_id",
+ return_value=graph_id,
+ waiter_delay=waiter_delay,
+ waiter_max_attempts=waiter_max_attempts,
+ **kwargs,
+ )
+
+ def hook(self) -> AwsGenericHook:
+ return NeptuneAnalyticsHook(
+ aws_conn_id=self.aws_conn_id,
+ region_name=self.region_name,
+ verify=self.verify,
+ config=self.botocore_config,
+ )
+
Review Comment:
`NeptuneGraphDeletedTrigger` is used by `NeptuneDeleteGraphOperator` in deferrable mode, but there is no unit test covering its serialization/run behavior (unlike the other triggers). Add a trigger unit test similar to `TestNeptuneGraphAvailableTrigger` to prevent regressions.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
Review Comment:
This provider operator imports `AirflowException` from `airflow.exceptions`, while other Amazon provider operators consistently import it from `airflow.providers.common.compat.sdk` for Airflow 2/3 compatibility (e.g. `providers/amazon/src/airflow/providers/amazon/aws/operators/neptune.py:34`). Please switch to the compat import here as well.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4798126947
@ellisms — converting this to draft because it currently has merge conflicts with `main` that need resolving before a maintainer can review. Once you've rebased onto the latest `main` and pushed, please hit **Ready for review** again and it'll re-enter the queue. Thanks for the contribution!
See the [PR quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
<sub>This is an automated first-pass triage note drafted by an AI-assisted tool and may get things wrong; once you've addressed it, a real Apache Airflow maintainer takes the next look. ([why the first pass is automated](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md))</sub>
---
Drafted-by: Claude Code (Opus 4.8); reviewed by @potiuk before posting
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "shivaam (via GitHub)" <gi...@apache.org>.
shivaam commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3283419523
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None:
+ raise NeptuneGraphCreationFailedError(
+ "No event received while waiting for Neptune graph creation to complete."
+ )
+
+ if event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ event.get(
+ "message",
+ f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+ if event:
+ status = event["status"]
+ if status.lower() != "success":
+ raise NeptunePrivateEndpointCreationFailedError("Endpoint failed to create")
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ f"Failed to delete private endpoint {endpoint_id}"
+ )
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and (event.get("status") or "").lower() == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+ else:
+ raise NeptunePrivateEndpointDeletionFailedError("Endpoint failed to delete.")
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise NeptuneGraphDeletionFailedError(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ if event:
+ graph_id = event.get("graph_id", "Unknown")
+
+ self.log.info("Neptune graph %s deleted", graph_id)
+
+
+class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph and imports data into it.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics,
+ you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphWithImportOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param source: The source from which to import data. Can be an S3 URI or Neptune database snapshot.
+ :param role_arn: The ARN of the IAM role that Neptune Analytics can assume to access the data source.
+ :param blank_node_handling: The method to handle blank nodes in the dataset. Options include 'convertToIri' or other handling strategies.
+ :param parquet_type: The type of Parquet files in the data source (if applicable).
+ :param format: The format of the data to be imported (e.g., 'csv', 'opencypher', 'ntriples', 'nquads', 'rdfxml', 'turtle').
+ :param min_provisioned_memory: The minimum provisioned memory for the graph in GBs.
+ :param max_provisioned_memory: The maximum provisioned memory for the graph in GBs.
+ :param fail_on_error: If True, the import will fail if any errors are encountered. If False, the import will continue despite errors.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled. (default: False)
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param import_options: Contains options for controlling the import process.
+ :param wait_for_completion: Whether to wait for the graph to be created and data imported. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 60)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be created and data imported.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "source", "role_arn", "kms_key"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (
+ NeptuneImportTaskLink(),
+ NeptuneGraphLink(),
+ )
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ source: str,
+ role_arn: str,
+ blank_node_handling: str | None = None,
+ parquet_type: str | None = None,
+ format: str | None = None,
+ min_provisioned_memory: int | None = None,
+ max_provisioned_memory: int | None = None,
+ fail_on_error: bool | None = None,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool | None = None,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ import_options: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.source = source
+ self.role_arn = role_arn
+ self.blank_node_handling = blank_node_handling
+ self.parquet_type = parquet_type
+ self.format = format
+ self.min_provisioned_memory = min_provisioned_memory
+ self.max_provisioned_memory = max_provisioned_memory
+ self.fail_on_error = fail_on_error
+ self.public_connectivity = public_connectivity
+ self.replica_count = replica_count
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.import_options = import_options
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s with import", self.graph_name)
+
+ # Build the import options
+ import_options = {
+ "neptune-analytics:blank-node-handling": self.blank_node_handling,
+ "neptune-analytics:parquet-type": self.parquet_type,
+ }
+
+ # Remove None values from import_options
+ import_options = {k: v for k, v in import_options.items() if v is not None}
+
+ # Merge with user-provided import_options
+ if self.import_options:
+ import_options.update(self.import_options)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "source": self.source,
+ "roleArn": self.role_arn,
+ **{
+ k: v
+ for k, v in {
+ "format": self.format,
+ "minProvisionedMemory": self.min_provisioned_memory,
+ "maxProvisionedMemory": self.max_provisioned_memory,
+ "failOnError": self.fail_on_error,
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ "importOptions": import_options if import_options else None,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph_using_import_task(**create_params)
+
+ self.log.info("Graph %s import task in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("graphId", None)
+ import_task_id = response.get("taskId")
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="defer_wait_for_task",
+ kwargs={"import_task_id": import_task_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ # Once the graph is available, wait for the task to complete
+
+ self.log.info("Waiting for import task %s", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def defer_wait_for_task(
+ self, context: Context, event: dict[str, Any] | None = None, import_task_id: str | None = None
+ ) -> None:
+ """Defers for import task completion."""
+ if import_task_id:
+ self.log.info("Deferring for import task %s completion", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None or event.get("status") != "success":
+ message = (event or {}).get(
+ "message", f"Neptune graph {self.graph_id} import did not complete successfully"
+ )
+ raise NeptuneGraphCreationFailedError(message)
+ self.log.info("Import complete for graph %s", self.graph_id)
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Starts a bulk data import task to load data into an empty Neptune graph.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneStartImportTaskOperator`
+
+ :param graph_identifier: Graph Id of target Neptune Graph
+ :param role_arn: IAM role ARN granting access to source data
+ :param source: URL identifying the source data location.
+ :param blank_node_handling: Method to handle blank nodes in dataset.
+ :param fail_on_error: If set to true, the task halts when an import error is encountered. If set to false, the task skips the data that caused the error and continues if possible.
+ :param format: Specifies the format of the Amazon S3 data to be imported.
+ :param import_options: Options on how to perform an import
+ :param parquet_type: Parquet type of import task
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "role_arn", "source", "import_options"
+ )
+ template_fields_renderers = {
+ "import_options": "json",
+ }
+ operator_extra_links = (NeptuneImportTaskLink(),)
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ role_arn: str,
+ source: str,
+ blank_node_handling: str | None = None,
+ fail_on_error: bool = True,
+ format: str | None = None,
+ import_options: dict | None = None,
+ parquet_type: str | None = "COLUMNAR",
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.role_arn = role_arn
+ self.source = source
+ self.blank_node_handling = blank_node_handling
+ self.fail_on_error = fail_on_error
+ self.format = format
+ self.import_options = import_options
+ self.parquet_type = parquet_type
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Starting data import to graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ "roleArn": self.role_arn,
+ "source": self.source,
+ **{
+ k: v
+ for k, v in {
+ "blankNodeHandling": self.blank_node_handling,
+ "failOnError": self.fail_on_error,
+ "format": self.format,
+ "importOptions": self.import_options,
+ "parquetType": self.parquet_type,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.start_import_task(**create_params)
+ import_task_id = response.get("taskId")
+ self.log.info("Import task %s started for graph %s", import_task_id, self.graph_identifier)
+
+ # Create the console link
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s completes", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to complete", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": import_task_id, "graph_id": self.graph_identifier}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ task_id = ""
+ if event:
+ task_id = event.get("import_task_id", "")
Review Comment:
same issue. I think we need to check if the trigger event has returned success or error.\
NeptuneCreateGraphOperator.execute_complete (line 190) and NeptuneCancelImportTaskOperator.execute_complete (line 991) already do the check correctly
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3041296005
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
Review Comment:
I think this is a hallucination.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4201888474
@o-nikolas I think all issues are resolved.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3025348467
##########
providers/amazon/tests/unit/amazon/aws/triggers/test_neptune_analytics.py:
##########
@@ -0,0 +1,280 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+from botocore.exceptions import WaiterError
+
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.common.compat.sdk import AirflowException
+from airflow.triggers.base import TriggerEvent
+
+GRAPH_ID = "test-graph"
+VPC_ID = "test-vpc"
+ENDPOINT_ID = "test-endpoint"
+TASK_ID = "test-task-id"
+
+
+class TestNeptuneGraphAvailableTrigger:
+ def test_serialization(self):
+ """
+ Asserts that the TaskStateTrigger correctly serializes its arguments
Review Comment:
The test docstring mentions `TaskStateTrigger`, but this file is testing Neptune Analytics triggers. Please update the wording to reference `NeptuneGraphAvailableTrigger` (and similarly for other trigger tests if the wording was copy/pasted) to avoid confusion when reading test failures.
```suggestion
Asserts that the NeptuneGraphAvailableTrigger correctly serializes its arguments
```
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
Review Comment:
The `hook` fixture is defined but doesn't appear to be used anywhere in this test module. Please remove it to reduce noise, or refactor tests to use it (e.g., when you need a real hook under `moto`).
```suggestion
```
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
Review Comment:
Avoid importing `AirflowException` inside the test function (`test_create_endpoint_failed_status`). Keep imports at the top of the file or reuse the existing module-level import pattern for consistency and simpler linting.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
Review Comment:
Avoid importing `NeptuneGraphAvailableTrigger` inside `test_deferrable_defers_with_graph_available_trigger` (create-graph tests). Prefer module-level imports for consistency and easier maintenance.
##########
providers/amazon/tests/unit/amazon/aws/triggers/test_neptune_analytics.py:
##########
@@ -0,0 +1,280 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from unittest import mock
+from unittest.mock import AsyncMock
+
+import pytest
+from botocore.exceptions import WaiterError
+
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.common.compat.sdk import AirflowException
+from airflow.triggers.base import TriggerEvent
+
+GRAPH_ID = "test-graph"
+VPC_ID = "test-vpc"
+ENDPOINT_ID = "test-endpoint"
+TASK_ID = "test-task-id"
+
+
+class TestNeptuneGraphAvailableTrigger:
+ def test_serialization(self):
+ """
+ Asserts that the TaskStateTrigger correctly serializes its arguments
Review Comment:
The docstring in `TestNeptuneGraphAvailableTrigger.test_serialization` references `TaskStateTrigger`, but this test is for `NeptuneGraphAvailableTrigger`. Please update the docstring to match the actual trigger name to avoid confusion.
```suggestion
Asserts that the NeptuneGraphAvailableTrigger correctly serializes its arguments
```
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
Review Comment:
`from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink` is imported inside `test_operator_extra_links`. Please move this import to the module level (this test module already has a large import block), unless there is a specific circular-import reason.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=False,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is False
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=False,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_no_wait(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once()
+ mock_conn.get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_wait_for_completion(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+ mock_waiter = mock.MagicMock()
+ mock_conn.get_waiter.return_value = mock_waiter
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_conn.get_waiter.assert_called_once_with("graph_deleted")
+ mock_waiter.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_resource_not_found(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ # Simulate ResourceNotFoundException
+ error_response = {
+ "Error": {
+ "Code": "ResourceNotFoundException",
+ "Message": "Graph not found",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 404,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should not raise an exception, just log that graph not found
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_other_client_error(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ # Simulate other ClientError
+ error_response = {
+ "Error": {
+ "Code": "ValidationException",
+ "Message": "Invalid parameter",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 400,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should raise AirflowException for non-ResourceNotFoundException errors
+ with pytest.raises(AirflowException):
+ operator.execute(None)
+
+
+class TestNeptuneCreateGraphWithImportOperator:
+ IMPORT_TASK_ID = "import-task-12345"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ )
+
+ assert operator.graph_name == GRAPH_NAME
+ assert operator.vector_search_config == {"dimension": 128}
+ assert operator.source == SOURCE_S3_URI
+ assert operator.role_arn == ROLE_ARN
+ assert operator.blank_node_handling is None
+ assert operator.parquet_type is None
+ assert operator.format is None
+ assert operator.min_provisioned_memory is None
+ assert operator.max_provisioned_memory is None
+ assert operator.fail_on_error is None
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is None
+ assert operator.kms_key is None
+ assert operator.tags is None
+ assert operator.import_options is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.create_graph_using_import_task.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"dimension": 128},
+ source=SOURCE_S3_URI,
+ roleArn=ROLE_ARN,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_with_all_optional_params(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling="convertToIri",
+ parquet_type="COLUMNAR",
+ format="csv",
+ min_provisioned_memory=16,
+ max_provisioned_memory=32,
+ fail_on_error=True,
+ public_connectivity=True,
+ replica_count=2,
+ deletion_protection=True,
+ kms_key_id="test-kms-key",
+ tags={"env": "test"},
+ import_options={"custom-option": "value"},
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.blank_node_handling == "convertToIri"
+ assert operator.parquet_type == "COLUMNAR"
+ assert operator.format == "csv"
+ assert operator.min_provisioned_memory == 16
+ assert operator.max_provisioned_memory == 32
+ assert operator.fail_on_error is True
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 2
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-kms-key"
+ assert operator.tags == {"env": "test"}
+ assert operator.import_options == {"custom-option": "value"}
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ # Verify the call includes all parameters
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ assert call_args["graphName"] == GRAPH_NAME
+ assert call_args["vectorSearchConfiguration"] == {"dimension": 128}
+ assert call_args["source"] == SOURCE_S3_URI
+ assert call_args["roleArn"] == ROLE_ARN
+ assert call_args["format"] == "csv"
+ assert call_args["minProvisionedMemory"] == 16
+ assert call_args["maxProvisionedMemory"] == 32
+ assert call_args["failOnError"] is True
+ assert call_args["replicaCount"] == 2
+ assert call_args["publicConnectivity"] is True
+ assert call_args["deletionProtection"] is True
+ assert call_args["kmsKeyIdentifier"] == "test-kms-key"
+ assert call_args["tags"] == {"env": "test"}
+ # Check import options were merged
+ assert "neptune-analytics:blank-node-handling" in call_args["importOptions"]
+ assert call_args["importOptions"]["neptune-analytics:blank-node-handling"] == "convertToIri"
+ assert "neptune-analytics:parquet-type" in call_args["importOptions"]
+ assert call_args["importOptions"]["neptune-analytics:parquet-type"] == "COLUMNAR"
+ assert call_args["importOptions"]["custom-option"] == "value"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_import_options_handling(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling="convertToIri",
+ import_options={"another-option": "test"},
+ )
+
+ operator.execute(None)
+
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ # Verify import options were properly merged
+ assert call_args["importOptions"]["neptune-analytics:blank-node-handling"] == "convertToIri"
+ assert call_args["importOptions"]["another-option"] == "test"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_with_import_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ wait_for_completion=False,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert result == {"graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_with_import_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ # Should wait for both graph_available and import_task_successful
+ assert mock_hook_get_waiter.call_count == 2
+ mock_hook_get_waiter.assert_any_call("graph_available")
+ mock_hook_get_waiter.assert_any_call("import_task_successful")
+ assert result == {"graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_import_options_none_values_filtered(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ # Test that None values in blank_node_handling and parquet_type are filtered out
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling=None,
+ parquet_type=None,
+ )
+
+ operator.execute(None)
+
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ # importOptions should not be in call_args if all values are None
+ assert "importOptions" not in call_args
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_defer_wait_for_task(self, mock_conn):
+ """Test that defer_wait_for_task defers with the import task trigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneImportTaskCompleteTrigger
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.defer_wait_for_task(
+ import_task_id=self.IMPORT_TASK_ID,
+ context=None,
+ event={"status": "success"},
+ )
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneImportTaskCompleteTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn):
+ """Test that execute defers with graph_available trigger and passes import_task_id."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
Review Comment:
Avoid importing `NeptuneGraphAvailableTrigger` inside `test_deferrable_defers_with_graph_available_trigger`. Import at module scope to match other unit tests and reduce per-test boilerplate.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=False,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is False
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=False,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_no_wait(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once()
+ mock_conn.get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_wait_for_completion(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+ mock_waiter = mock.MagicMock()
+ mock_conn.get_waiter.return_value = mock_waiter
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_conn.get_waiter.assert_called_once_with("graph_deleted")
+ mock_waiter.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_resource_not_found(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ # Simulate ResourceNotFoundException
+ error_response = {
+ "Error": {
+ "Code": "ResourceNotFoundException",
+ "Message": "Graph not found",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 404,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should not raise an exception, just log that graph not found
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_other_client_error(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ from airflow.providers.common.compat.sdk import AirflowException
+
Review Comment:
There are additional function-local imports of `ClientError` and `AirflowException` in these tests. For consistency with other provider unit tests (and to keep linting/simple static analysis happy), please move these imports to the module level and reuse them across tests.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=False,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is False
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=False,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_no_wait(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once()
+ mock_conn.get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_wait_for_completion(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+ mock_waiter = mock.MagicMock()
+ mock_conn.get_waiter.return_value = mock_waiter
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_conn.get_waiter.assert_called_once_with("graph_deleted")
+ mock_waiter.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_resource_not_found(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ # Simulate ResourceNotFoundException
+ error_response = {
+ "Error": {
+ "Code": "ResourceNotFoundException",
+ "Message": "Graph not found",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 404,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should not raise an exception, just log that graph not found
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_other_client_error(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ from airflow.providers.common.compat.sdk import AirflowException
+
Review Comment:
Avoid importing `ClientError`/`AirflowException` inside the test (`test_delete_graph_other_client_error`). Prefer module-level imports like other operator unit tests to keep the module consistent and easier to maintain.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,961 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.exceptions import AirflowException
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id and vpc id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ self.log.info("Neptune graph % complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise AirflowException(f"Private endpoint failed to create for graph {self.graph_identifier}")
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise AirflowException(f"Failed to delete private endpoint {endpoint_id}")
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and event.get("status") == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise AirflowException(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ if event:
+ graph_id = event.get("graph_id", "Unknown")
+
+ self.log.info("Neptune graph % deleted", graph_id)
Review Comment:
Logging format string in `execute_complete` is broken: `"Neptune graph % deleted"` uses an invalid `%` placeholder and will raise a `TypeError` when this log line is executed. Use `%s` (or f-string) so the graph id is actually logged.
```suggestion
self.log.info("Neptune graph %s deleted", graph_id)
```
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=False,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is False
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=False,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_no_wait(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once()
+ mock_conn.get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_wait_for_completion(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+ mock_waiter = mock.MagicMock()
+ mock_conn.get_waiter.return_value = mock_waiter
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_conn.get_waiter.assert_called_once_with("graph_deleted")
+ mock_waiter.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_resource_not_found(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ # Simulate ResourceNotFoundException
+ error_response = {
+ "Error": {
+ "Code": "ResourceNotFoundException",
+ "Message": "Graph not found",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 404,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should not raise an exception, just log that graph not found
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_other_client_error(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ # Simulate other ClientError
+ error_response = {
+ "Error": {
+ "Code": "ValidationException",
+ "Message": "Invalid parameter",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 400,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should raise AirflowException for non-ResourceNotFoundException errors
+ with pytest.raises(AirflowException):
+ operator.execute(None)
+
+
+class TestNeptuneCreateGraphWithImportOperator:
+ IMPORT_TASK_ID = "import-task-12345"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ )
+
+ assert operator.graph_name == GRAPH_NAME
+ assert operator.vector_search_config == {"dimension": 128}
+ assert operator.source == SOURCE_S3_URI
+ assert operator.role_arn == ROLE_ARN
+ assert operator.blank_node_handling is None
+ assert operator.parquet_type is None
+ assert operator.format is None
+ assert operator.min_provisioned_memory is None
+ assert operator.max_provisioned_memory is None
+ assert operator.fail_on_error is None
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is None
+ assert operator.kms_key is None
+ assert operator.tags is None
+ assert operator.import_options is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.create_graph_using_import_task.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"dimension": 128},
+ source=SOURCE_S3_URI,
+ roleArn=ROLE_ARN,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_with_all_optional_params(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling="convertToIri",
+ parquet_type="COLUMNAR",
+ format="csv",
+ min_provisioned_memory=16,
+ max_provisioned_memory=32,
+ fail_on_error=True,
+ public_connectivity=True,
+ replica_count=2,
+ deletion_protection=True,
+ kms_key_id="test-kms-key",
+ tags={"env": "test"},
+ import_options={"custom-option": "value"},
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.blank_node_handling == "convertToIri"
+ assert operator.parquet_type == "COLUMNAR"
+ assert operator.format == "csv"
+ assert operator.min_provisioned_memory == 16
+ assert operator.max_provisioned_memory == 32
+ assert operator.fail_on_error is True
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 2
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-kms-key"
+ assert operator.tags == {"env": "test"}
+ assert operator.import_options == {"custom-option": "value"}
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ # Verify the call includes all parameters
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ assert call_args["graphName"] == GRAPH_NAME
+ assert call_args["vectorSearchConfiguration"] == {"dimension": 128}
+ assert call_args["source"] == SOURCE_S3_URI
+ assert call_args["roleArn"] == ROLE_ARN
+ assert call_args["format"] == "csv"
+ assert call_args["minProvisionedMemory"] == 16
+ assert call_args["maxProvisionedMemory"] == 32
+ assert call_args["failOnError"] is True
+ assert call_args["replicaCount"] == 2
+ assert call_args["publicConnectivity"] is True
+ assert call_args["deletionProtection"] is True
+ assert call_args["kmsKeyIdentifier"] == "test-kms-key"
+ assert call_args["tags"] == {"env": "test"}
+ # Check import options were merged
+ assert "neptune-analytics:blank-node-handling" in call_args["importOptions"]
+ assert call_args["importOptions"]["neptune-analytics:blank-node-handling"] == "convertToIri"
+ assert "neptune-analytics:parquet-type" in call_args["importOptions"]
+ assert call_args["importOptions"]["neptune-analytics:parquet-type"] == "COLUMNAR"
+ assert call_args["importOptions"]["custom-option"] == "value"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_import_options_handling(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling="convertToIri",
+ import_options={"another-option": "test"},
+ )
+
+ operator.execute(None)
+
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ # Verify import options were properly merged
+ assert call_args["importOptions"]["neptune-analytics:blank-node-handling"] == "convertToIri"
+ assert call_args["importOptions"]["another-option"] == "test"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_with_import_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ wait_for_completion=False,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert result == {"graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_with_import_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ # Should wait for both graph_available and import_task_successful
+ assert mock_hook_get_waiter.call_count == 2
+ mock_hook_get_waiter.assert_any_call("graph_available")
+ mock_hook_get_waiter.assert_any_call("import_task_successful")
+ assert result == {"graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_import_options_none_values_filtered(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ # Test that None values in blank_node_handling and parquet_type are filtered out
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling=None,
+ parquet_type=None,
+ )
+
+ operator.execute(None)
+
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ # importOptions should not be in call_args if all values are None
+ assert "importOptions" not in call_args
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_defer_wait_for_task(self, mock_conn):
+ """Test that defer_wait_for_task defers with the import task trigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneImportTaskCompleteTrigger
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.defer_wait_for_task(
+ import_task_id=self.IMPORT_TASK_ID,
+ context=None,
+ event={"status": "success"},
+ )
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneImportTaskCompleteTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn):
+ """Test that execute defers with graph_available trigger and passes import_task_id."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "defer_wait_for_task"
+ assert exc_info.value.kwargs == {"import_task_id": self.IMPORT_TASK_ID}
+
+
+TASK_ID = "import-task-id-12345"
+
+
+class TestNeptuneStartImportTaskOperator:
+ def test_template_fields(self):
+ fields = NeptuneStartImportTaskOperator.template_fields
+ assert "graph_identifier" in fields
+ assert "role_arn" in fields
+ assert "source" in fields
+ assert "import_options" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneStartImportTaskOperator.template_fields_renderers == {"import_options": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneImportTaskLink
+
+ assert len(NeptuneStartImportTaskOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneStartImportTaskOperator.operator_extra_links[0], NeptuneImportTaskLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.start_import_task.return_value = {
+ "taskId": TASK_ID,
+ "graphId": GRAPH_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneStartImportTaskOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ role_arn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.role_arn == ROLE_ARN
+ assert operator.source == SOURCE_S3_URI
+ assert operator.blank_node_handling is None
+ assert operator.fail_on_error is True
+ assert operator.format is None
+ assert operator.import_options is None
+ assert operator.parquet_type == "COLUMNAR"
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.start_import_task.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ roleArn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ failOnError=True,
+ parquetType="COLUMNAR",
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.start_import_task.return_value = {
+ "taskId": TASK_ID,
+ "graphId": GRAPH_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneStartImportTaskOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ role_arn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ blank_node_handling=None,
+ fail_on_error=False,
+ format="CSV",
+ import_options={"neptune.csv.allowEmptyStrings": True},
+ parquet_type=None,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.blank_node_handling is None
+ assert operator.fail_on_error is False
+ assert operator.format == "CSV"
+ assert operator.import_options == {"neptune.csv.allowEmptyStrings": True}
+ assert operator.parquet_type is None
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.start_import_task.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ roleArn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ failOnError=False,
+ format="CSV",
+ importOptions={"neptune.csv.allowEmptyStrings": True},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_start_import_no_wait(self, mock_get_waiter, mock_conn, mock_persist):
+ mock_conn.start_import_task.return_value = {
+ "taskId": TASK_ID,
+ "graphId": GRAPH_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneStartImportTaskOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ role_arn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ wait_for_completion=False,
+ )
+ result = operator.execute(None)
+
+ mock_get_waiter.assert_not_called()
+ assert result == {"import_task_id": TASK_ID, "graph_id": GRAPH_ID}
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_start_import_wait_for_completion(self, mock_get_waiter, mock_conn, mock_persist):
+ mock_conn.start_import_task.return_value = {
+ "taskId": TASK_ID,
+ "graphId": GRAPH_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneStartImportTaskOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ role_arn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_get_waiter.assert_called_once_with("import_task_successful")
+ assert result == {"import_task_id": TASK_ID, "graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneImportTaskLink.persist is called with the correct arguments."""
+ mock_conn.start_import_task.return_value = {
+ "taskId": TASK_ID,
+ "graphId": GRAPH_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneStartImportTaskOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ role_arn=ROLE_ARN,
+ source=SOURCE_S3_URI,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ import_task_id=TASK_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneImportTaskLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_import_task_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneImportTaskCompleteTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneImportTaskCompleteTrigger
+
Review Comment:
Avoid importing `NeptuneImportTaskCompleteTrigger` inside the test (`test_deferrable_defers_with_import_task_trigger`). Import it at module scope like the other trigger/operator symbols to keep imports consistent and easier to audit.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=False,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is False
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=False,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_no_wait(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once()
+ mock_conn.get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_wait_for_completion(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+ mock_waiter = mock.MagicMock()
+ mock_conn.get_waiter.return_value = mock_waiter
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_conn.get_waiter.assert_called_once_with("graph_deleted")
+ mock_waiter.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_resource_not_found(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ # Simulate ResourceNotFoundException
+ error_response = {
+ "Error": {
+ "Code": "ResourceNotFoundException",
+ "Message": "Graph not found",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 404,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should not raise an exception, just log that graph not found
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_other_client_error(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ # Simulate other ClientError
+ error_response = {
+ "Error": {
+ "Code": "ValidationException",
+ "Message": "Invalid parameter",
+ },
+ "ResponseMetadata": {
+ "HTTPStatusCode": 400,
+ },
+ }
+ mock_conn.delete_graph.side_effect = ClientError(error_response, "delete_graph")
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ # Should raise AirflowException for non-ResourceNotFoundException errors
+ with pytest.raises(AirflowException):
+ operator.execute(None)
+
+
+class TestNeptuneCreateGraphWithImportOperator:
+ IMPORT_TASK_ID = "import-task-12345"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ )
+
+ assert operator.graph_name == GRAPH_NAME
+ assert operator.vector_search_config == {"dimension": 128}
+ assert operator.source == SOURCE_S3_URI
+ assert operator.role_arn == ROLE_ARN
+ assert operator.blank_node_handling is None
+ assert operator.parquet_type is None
+ assert operator.format is None
+ assert operator.min_provisioned_memory is None
+ assert operator.max_provisioned_memory is None
+ assert operator.fail_on_error is None
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is None
+ assert operator.kms_key is None
+ assert operator.tags is None
+ assert operator.import_options is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.create_graph_using_import_task.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"dimension": 128},
+ source=SOURCE_S3_URI,
+ roleArn=ROLE_ARN,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_with_all_optional_params(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling="convertToIri",
+ parquet_type="COLUMNAR",
+ format="csv",
+ min_provisioned_memory=16,
+ max_provisioned_memory=32,
+ fail_on_error=True,
+ public_connectivity=True,
+ replica_count=2,
+ deletion_protection=True,
+ kms_key_id="test-kms-key",
+ tags={"env": "test"},
+ import_options={"custom-option": "value"},
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.blank_node_handling == "convertToIri"
+ assert operator.parquet_type == "COLUMNAR"
+ assert operator.format == "csv"
+ assert operator.min_provisioned_memory == 16
+ assert operator.max_provisioned_memory == 32
+ assert operator.fail_on_error is True
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 2
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-kms-key"
+ assert operator.tags == {"env": "test"}
+ assert operator.import_options == {"custom-option": "value"}
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ # Verify the call includes all parameters
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ assert call_args["graphName"] == GRAPH_NAME
+ assert call_args["vectorSearchConfiguration"] == {"dimension": 128}
+ assert call_args["source"] == SOURCE_S3_URI
+ assert call_args["roleArn"] == ROLE_ARN
+ assert call_args["format"] == "csv"
+ assert call_args["minProvisionedMemory"] == 16
+ assert call_args["maxProvisionedMemory"] == 32
+ assert call_args["failOnError"] is True
+ assert call_args["replicaCount"] == 2
+ assert call_args["publicConnectivity"] is True
+ assert call_args["deletionProtection"] is True
+ assert call_args["kmsKeyIdentifier"] == "test-kms-key"
+ assert call_args["tags"] == {"env": "test"}
+ # Check import options were merged
+ assert "neptune-analytics:blank-node-handling" in call_args["importOptions"]
+ assert call_args["importOptions"]["neptune-analytics:blank-node-handling"] == "convertToIri"
+ assert "neptune-analytics:parquet-type" in call_args["importOptions"]
+ assert call_args["importOptions"]["neptune-analytics:parquet-type"] == "COLUMNAR"
+ assert call_args["importOptions"]["custom-option"] == "value"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_import_options_handling(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling="convertToIri",
+ import_options={"another-option": "test"},
+ )
+
+ operator.execute(None)
+
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ # Verify import options were properly merged
+ assert call_args["importOptions"]["neptune-analytics:blank-node-handling"] == "convertToIri"
+ assert call_args["importOptions"]["another-option"] == "test"
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_with_import_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ wait_for_completion=False,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert result == {"graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_with_import_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ # Should wait for both graph_available and import_task_successful
+ assert mock_hook_get_waiter.call_count == 2
+ mock_hook_get_waiter.assert_any_call("graph_available")
+ mock_hook_get_waiter.assert_any_call("import_task_successful")
+ assert result == {"graph_id": GRAPH_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_import_options_none_values_filtered(self, mock_conn):
+ mock_conn.create_graph_using_import_task.return_value = {
+ "graphId": GRAPH_ID,
+ "taskId": self.IMPORT_TASK_ID,
+ "status": "IMPORTING",
+ }
+
+ # Test that None values in blank_node_handling and parquet_type are filtered out
+ operator = NeptuneCreateGraphWithImportOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"dimension": 128},
+ source=SOURCE_S3_URI,
+ role_arn=ROLE_ARN,
+ blank_node_handling=None,
+ parquet_type=None,
+ )
+
+ operator.execute(None)
+
+ call_args = mock_conn.create_graph_using_import_task.call_args[1]
+ # importOptions should not be in call_args if all values are None
+ assert "importOptions" not in call_args
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_defer_wait_for_task(self, mock_conn):
+ """Test that defer_wait_for_task defers with the import task trigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneImportTaskCompleteTrigger
+
Review Comment:
Avoid importing `NeptuneImportTaskCompleteTrigger` inside `test_defer_wait_for_task`. Prefer module-level imports like other tests in this repo to keep imports consistent and avoid hidden dependencies.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
Review Comment:
The `hook` fixture uses `moto.mock_aws()` and yields `NeptuneAnalyticsHook`, but it is never used in this test module. Please remove the fixture (and the `mock_aws`/`Generator` imports) or refactor tests to use it; otherwise it adds an unnecessary dependency and setup time.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
Review Comment:
Avoid importing `AirflowException` inside the test function (`test_delete_endpoint_failed_status`). Prefer module-level imports for consistency and to align with codebase style.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
Review Comment:
`mock_aws` (moto) is only used by the unused `hook` fixture. If you remove the fixture, please also drop this import to avoid an unnecessary dependency in unit tests.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
Review Comment:
`test_init_defaults` for `NeptuneDeleteGraphOperator` asserts `mock_conn.delete_graph.assert_called_once_with(...)`, but the operator implementation calls `self.hook.conn.delete_graph(...)`. In other tests in this module, `NeptuneAnalyticsHook.conn` is patched and assertions are made against the patched hook connection (not the boto client). Please align this test with the actual patch target (either patch `NeptuneAnalyticsHook.conn` and assert against that, or patch `NeptuneDeleteGraphOperator.hook/conn`), otherwise this test will fail.
##########
providers/amazon/tests/unit/amazon/aws/operators/test_neptune_analytics.py:
##########
@@ -0,0 +1,1314 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.common.compat.sdk import TaskDeferred
+
+GRAPH_NAME = "test_graph"
+GRAPH_ID = "test-graph-id"
+VPC_ID = "vpc-12345"
+SUBNET_IDS = ["subnet-1", "subnet-2"]
+SECURITY_GROUP_IDS = ["sg-1", "sg-2"]
+ENDPOINT_ID = "vpce-12345"
+SOURCE_S3_URI = "s3://my-bucket/my-data/"
+ROLE_ARN = "arn:aws:iam::123456789012:role/NeptuneImportRole"
+
+
+@pytest.fixture
+def hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneCreateGraphOperator:
+ def test_template_fields(self):
+ # Verify template_fields includes the expected fields
+ fields = NeptuneCreateGraphOperator.template_fields
+ assert "graph_name" in fields
+ assert "vector_search_config" in fields
+ assert "provisioned_memory" in fields
+
+ def test_template_fields_renderers(self):
+ assert NeptuneCreateGraphOperator.template_fields_renderers == {"vector_search_config": "json"}
+
+ def test_operator_extra_links(self):
+ from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink
+
+ assert len(NeptuneCreateGraphOperator.operator_extra_links) == 1
+ assert isinstance(NeptuneCreateGraphOperator.operator_extra_links[0], NeptuneGraphLink)
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ )
+
+ assert operator.public_connectivity is None
+ assert operator.replica_count is None
+ assert operator.deletion_protect is False
+ assert operator.kms_key is None
+ assert operator.tags is None
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ provisionedMemory=16,
+ deletionProtection=False,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ public_connectivity=True,
+ replica_count=3,
+ kms_key_id="test-key",
+ tags={"key1": "test"},
+ deletion_protection=True,
+ )
+
+ assert operator.public_connectivity is True
+ assert operator.replica_count == 3
+ assert operator.deletion_protect is True
+ assert operator.kms_key == "test-key"
+ assert operator.tags == {"key1": "test"}
+
+ operator.execute(None)
+
+ mock_conn.create_graph.assert_called_once_with(
+ graphName=GRAPH_NAME,
+ vectorSearchConfiguration={"test": 123},
+ replicaCount=3,
+ publicConnectivity=True,
+ provisionedMemory=16,
+ deletionProtection=True,
+ kmsKeyIdentifier="test-key",
+ tags={"key1": "test"},
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=False,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_graph_wait_for_completion(self, mock_hook_get_waiter, mock_conn, mock_persist):
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ provisioned_memory=16,
+ vector_search_config={"test": 123},
+ wait_for_completion=True,
+ )
+ resp = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("graph_available")
+ assert "graph_id" in resp
+ assert resp["graph_id"] == GRAPH_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_persist_called_with_correct_args(self, mock_conn):
+ """Test that NeptuneGraphLink.persist is called with the correct arguments."""
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ wait_for_completion=False,
+ )
+
+ mock_context = mock.MagicMock()
+ with mock.patch(
+ "airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist"
+ ) as mock_persist:
+ operator.execute(mock_context)
+
+ mock_persist.assert_called_once_with(
+ context=mock_context,
+ operator=operator,
+ region_name=mock.ANY,
+ aws_partition=mock.ANY,
+ graph_id=GRAPH_ID,
+ )
+
+ @mock.patch("airflow.providers.amazon.aws.operators.neptune_analytics.NeptuneGraphLink.persist")
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_deferrable_defers_with_graph_available_trigger(self, mock_conn, mock_persist):
+ """Test that deferrable mode defers with NeptuneGraphAvailableTrigger."""
+ from airflow.providers.amazon.aws.triggers.neptune_analytics import NeptuneGraphAvailableTrigger
+
+ mock_conn.create_graph.return_value = {"id": GRAPH_ID, "status": "CREATING"}
+
+ operator = NeptuneCreateGraphOperator(
+ task_id="test_task",
+ graph_name=GRAPH_NAME,
+ vector_search_config={"test": 123},
+ provisioned_memory=16,
+ deferrable=True,
+ )
+
+ with pytest.raises(TaskDeferred) as exc_info:
+ operator.execute(None)
+
+ trigger = exc_info.value.trigger
+ assert isinstance(trigger, NeptuneGraphAvailableTrigger)
+ assert exc_info.value.method_name == "execute_complete"
+
+
+class TestNeptuneCreatePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id is None
+ assert operator.subnet_ids is None
+ assert operator.vpc_security_group_ids is None
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ result = operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ )
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ assert result is not None
+ assert result["vpc_endpoint_id"] == ENDPOINT_ID
+ assert result["graph_id"] == GRAPH_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ subnet_ids=SUBNET_IDS,
+ vpc_security_group_ids=SECURITY_GROUP_IDS,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.subnet_ids == SUBNET_IDS
+ assert operator.vpc_security_group_ids == SECURITY_GROUP_IDS
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.create_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ subnetIds=SUBNET_IDS,
+ vpcSecurityGroupIds=SECURITY_GROUP_IDS,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_create_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ result = operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_available")
+ mock_hook_get_waiter.return_value.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_sets_vpc_id_from_response(self, mock_conn):
+ """When vpc_id is not provided, the operator should use the vpc_id from the API response."""
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "CREATING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ )
+
+ assert operator.vpc_id is None
+ result = operator.execute(None)
+
+ # vpc_id should be set from the create response
+ assert operator.vpc_id == VPC_ID
+ assert result["vpc_id"] == VPC_ID
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_create_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.create_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Private endpoint failed to create for graph {GRAPH_ID}"):
+ operator.execute(None)
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_execute_complete(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator.execute_complete(None, {"status": "success"})
+
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+ assert result == {"vpc_endpoint_id": ENDPOINT_ID, "graph_id": GRAPH_ID, "vpc_id": VPC_ID}
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_get_graph_endpoint_id(self, mock_conn):
+ mock_conn.get_private_graph_endpoint.return_value = {
+ "vpcEndpointId": ENDPOINT_ID,
+ }
+
+ operator = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="test_task", graph_identifier=GRAPH_ID, vpc_id=VPC_ID
+ )
+
+ result = operator._get_graph_endpoint_id()
+
+ assert result == ENDPOINT_ID
+ mock_conn.get_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+
+class TestNeptuneDeletePrivateGraphEndpointOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_identifier == GRAPH_ID
+ assert operator.vpc_id == VPC_ID
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_private_graph_endpoint.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ vpcId=VPC_ID,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_no_wait(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_delete_endpoint_wait_for_completion(self, mock_hook_get_waiter, mock_conn):
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "DELETING",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_hook_get_waiter.assert_called_once_with("private_graph_endpoint_deleted")
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_endpoint_failed_status(self, mock_conn):
+ from airflow.providers.common.compat.sdk import AirflowException
+
+ mock_conn.delete_private_graph_endpoint.return_value = {
+ "status": "FAILED",
+ "vpcEndpointId": ENDPOINT_ID,
+ "vpcId": VPC_ID,
+ }
+
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ with pytest.raises(AirflowException, match=f"Failed to delete private endpoint {ENDPOINT_ID}"):
+ operator.execute(None)
+
+ def test_execute_complete_success(self):
+ operator = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="test_task",
+ graph_identifier=GRAPH_ID,
+ vpc_id=VPC_ID,
+ )
+
+ event = {
+ "status": "success",
+ "endpoint_id": ENDPOINT_ID,
+ }
+
+ operator.execute_complete(None, event)
+
+ # Verify the method completes without error and logs the endpoint_id
+
+
+class TestNeptuneDeleteGraphOperator:
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_defaults(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is True
+ assert operator.wait_for_completion is True
+ assert operator.waiter_delay == 30
+ assert operator.waiter_max_attempts == 60
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=True,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_init_custom_args(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=False,
+ waiter_delay=60,
+ waiter_max_attempts=100,
+ )
+
+ assert operator.graph_id == GRAPH_ID
+ assert operator.skip_snapshot is False
+ assert operator.waiter_delay == 60
+ assert operator.waiter_max_attempts == 100
+
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ skipSnapshot=False,
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_no_wait(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=False,
+ )
+ operator.execute(None)
+
+ mock_conn.delete_graph.assert_called_once()
+ mock_conn.get_waiter.assert_not_called()
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_wait_for_completion(self, mock_conn):
+ mock_conn.delete_graph.return_value = {
+ "id": GRAPH_ID,
+ "name": GRAPH_NAME,
+ "status": "DELETING",
+ }
+ mock_waiter = mock.MagicMock()
+ mock_conn.get_waiter.return_value = mock_waiter
+
+ operator = NeptuneDeleteGraphOperator(
+ task_id="test_task",
+ graph_id=GRAPH_ID,
+ skip_snapshot=True,
+ wait_for_completion=True,
+ )
+ operator.execute(None)
+
+ mock_conn.get_waiter.assert_called_once_with("graph_deleted")
+ mock_waiter.wait.assert_called_once_with(
+ graphIdentifier=GRAPH_ID,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+
+ @mock.patch.object(NeptuneAnalyticsHook, "conn")
+ def test_delete_graph_resource_not_found(self, mock_conn):
+ from botocore.exceptions import ClientError
+
+ # Simulate ResourceNotFoundException
Review Comment:
Avoid importing `ClientError` inside the test (`test_delete_graph_resource_not_found`). Other unit test modules in this provider keep botocore exceptions imported at module level; moving this import up will keep style consistent.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3470655387
##########
providers/amazon/src/airflow/providers/amazon/aws/waiters/neptune_analytics.json:
##########
@@ -0,0 +1,38 @@
+{
+ "version": 2,
+ "waiters": {
+
+ "import_task_cancelled":{
+ "operation": "GetImportTask",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "SUCCEEDED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "CANCELLED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "ERROR_ENCOUNTERED",
+ "state": "error"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
Review Comment:
Okay, I don't fully grok that comment from @shivaam, but accepting FAILED as Successful just seemed odd to me. If you two have worked it out, you can resolve this thread.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3483457856
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
Review Comment:
good catch, fixed.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3462936046
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ # get the vpce id since it may not be returned immediately
+ endpoint_id = self.hook._get_graph_endpoint_id(graph_id=self.graph_identifier, vpc_id=self.vpc_id)
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ kwargs={"vpc_id": self.vpc_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptunePrivateEndpointCreationFailedError(
+ validated_event.get("message", "Endpoint failed to create")
+ )
+
+ graph_id = validated_event["value"]
+ vpc_id = validated_event["vpc_id"]
Review Comment:
Fixed. `value` was an error and `vpc_id` should be an argument, since its sent as a `kwarg` from the caller.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3283789903
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/neptune_analytics.py:
##########
@@ -0,0 +1,37 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class NeptuneAnalyticsHook(AwsBaseHook):
+ """
+ Interact with Amazon Neptune Analytics.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ def __init__(self, *args, **kwargs):
+ kwargs["client_type"] = "neptune-graph"
Review Comment:
Neptune analytics is how the service is identified in the AWS documentation. The boto3 service is neptune-graph, which causes confusion. I chose Neptune Analytics for all the operators since a customer is likely to use that name with searching.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "shivaam (via GitHub)" <gi...@apache.org>.
shivaam commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3283085275
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/neptune_analytics.py:
##########
@@ -0,0 +1,37 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+
+
+class NeptuneAnalyticsHook(AwsBaseHook):
+ """
+ Interact with Amazon Neptune Analytics.
+
+ Additional arguments (such as ``aws_conn_id``) may be specified and
+ are passed down to the underlying AwsBaseHook.
+
+ .. seealso::
+ - :class:`~airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook`
+ """
+
+ def __init__(self, *args, **kwargs):
+ kwargs["client_type"] = "neptune-graph"
Review Comment:
difference between neptune graph vs neptune analytics?
##########
airflow-core/src/airflow/migrations/versions/0113_3_3_0_add_retry_policy_fields_to_ti.py:
##########
@@ -30,7 +30,7 @@
is a metadata-only operation (no table rewrite).
Revision ID: b8f3e4a1d2c9
-Revises: 9fabad868fdb
+Revises: fde9ed84d07b
Create Date: 2026-04-16 12:00:00.000000
Review Comment:
Is this expected?
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,335 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.common.compat.sdk import DAG, chain, task
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_neptune_analytics"
+
+ROLE_ARN_KEY = "ROLE_ARN"
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
+
+# Minimal OpenCypher CSV data for import testing.
+NODES_CSV = """~id,~label,name:String
+n1,Person,Alice
+n2,Person,Bob
+"""
+
+EDGES_CSV = """~id,~from,~to,~label
+e1,n1,n2,KNOWS
+"""
+
+NEPTUNE_ANALYTICS_TRUST_POLICY = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Principal": {"Service": "neptune-graph.amazonaws.com"},
+ "Action": "sts:AssumeRole",
+ }
+ ],
+ }
+)
+
+S3_READ_POLICY_DOCUMENT = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Action": ["s3:GetObject", "s3:ListBucket"],
+ "Resource": ["arn:aws:s3:::*", "arn:aws:s3:::*/*"],
+ }
+ ],
+ }
+)
+
+
+@task
+def create_neptune_import_role(role_name: str) -> str:
+ iam_client = boto3.client("iam")
+ iam_client.create_role(
+ RoleName=role_name,
+ AssumeRolePolicyDocument=NEPTUNE_ANALYTICS_TRUST_POLICY,
+ Description="Role for Neptune Analytics import system test",
+ )
+ iam_client.put_role_policy(
+ RoleName=role_name,
+ PolicyName="NeptuneAnalyticsS3Access",
+ PolicyDocument=S3_READ_POLICY_DOCUMENT,
+ )
+ role = iam_client.get_role(RoleName=role_name)
+ time.sleep(60) # Wait for IAM eventual consistency (role + inline policy propagation)
+ return role["Role"]["Arn"]
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_neptune_import_role(role_name: str) -> None:
+ iam_client = boto3.client("iam")
+ with contextlib.suppress(iam_client.exceptions.NoSuchEntityException):
+ iam_client.delete_role_policy(RoleName=role_name, PolicyName="NeptuneAnalyticsS3Access")
Review Comment:
only deleting the policy not the role
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,335 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.common.compat.sdk import DAG, chain, task
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_neptune_analytics"
+
+ROLE_ARN_KEY = "ROLE_ARN"
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
+
+# Minimal OpenCypher CSV data for import testing.
+NODES_CSV = """~id,~label,name:String
+n1,Person,Alice
+n2,Person,Bob
+"""
+
+EDGES_CSV = """~id,~from,~to,~label
+e1,n1,n2,KNOWS
+"""
+
+NEPTUNE_ANALYTICS_TRUST_POLICY = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Principal": {"Service": "neptune-graph.amazonaws.com"},
+ "Action": "sts:AssumeRole",
+ }
+ ],
+ }
+)
+
+S3_READ_POLICY_DOCUMENT = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Action": ["s3:GetObject", "s3:ListBucket"],
+ "Resource": ["arn:aws:s3:::*", "arn:aws:s3:::*/*"],
+ }
+ ],
+ }
+)
+
+
+@task
+def create_neptune_import_role(role_name: str) -> str:
+ iam_client = boto3.client("iam")
+ iam_client.create_role(
+ RoleName=role_name,
+ AssumeRolePolicyDocument=NEPTUNE_ANALYTICS_TRUST_POLICY,
+ Description="Role for Neptune Analytics import system test",
+ )
+ iam_client.put_role_policy(
+ RoleName=role_name,
+ PolicyName="NeptuneAnalyticsS3Access",
+ PolicyDocument=S3_READ_POLICY_DOCUMENT,
+ )
+ role = iam_client.get_role(RoleName=role_name)
+ time.sleep(60) # Wait for IAM eventual consistency (role + inline policy propagation)
+ return role["Role"]["Arn"]
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_neptune_import_role(role_name: str) -> None:
+ iam_client = boto3.client("iam")
+ with contextlib.suppress(iam_client.exceptions.NoSuchEntityException):
+ iam_client.delete_role_policy(RoleName=role_name, PolicyName="NeptuneAnalyticsS3Access")
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_graph_if_exists(graph_name: str) -> None:
+ """Safety net to clean up the graph in case a previous task failed."""
+ hook = NeptuneAnalyticsHook()
+ with contextlib.suppress(Exception):
+ # List graphs and find by name
+ paginator = hook.conn.get_paginator("list_graphs")
+ for page in paginator.paginate():
+ for graph in page.get("graphs", []):
+ if graph.get("name") == graph_name:
+ graph_id = graph["id"]
+ # Disable deletion protection if enabled
+
+ hook.conn.update_graph(graphIdentifier=graph_id, deletionProtection=False)
+
+ hook.conn.delete_graph(graphIdentifier=graph_id, skipSnapshot=True)
+ hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=graph_id,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ return
+
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ env_id = test_context["ENV_ID"]
+ graph_name = f"{env_id}-graph"
+ import_graph_name = f"{env_id}-import-graph"
+ bucket_name = f"{env_id}-neptune-analytics"
+ import_role_name = f"{env_id}-neptune-import"
+ region = boto3.session.Session().region_name
+
+ # --- TEST SETUP ---
+
+ create_bucket = S3CreateBucketOperator(
+ task_id="create_bucket",
+ bucket_name=bucket_name,
+ )
+
+ upload_nodes = S3CreateObjectOperator(
+ task_id="upload_nodes",
+ s3_bucket=bucket_name,
+ s3_key="data/nodes.csv",
+ data=NODES_CSV,
+ replace=True,
+ )
+
+ upload_edges = S3CreateObjectOperator(
+ task_id="upload_edges",
+ s3_bucket=bucket_name,
+ s3_key="data/edges.csv",
+ data=EDGES_CSV,
+ replace=True,
+ )
+
+ create_role = create_neptune_import_role(import_role_name)
+
+ # --- TEST BODY ---
+
+ # [START howto_operator_neptune_analytics_create_graph]
+ create_graph = NeptuneCreateGraphOperator(
+ task_id="create_graph",
+ graph_name=graph_name,
+ vector_search_config={"dimension": 128},
+ provisioned_memory=32,
+ public_connectivity=True,
+ replica_count=0,
+ deletion_protection=False,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_create_graph]
+
+ # [START howto_operator_neptune_analytics_create_private_endpoint]
+ create_endpoint = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="create_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id']}}",
+ wait_for_completion=True,
+ )
+ # [END howto_operator_neptune_analytics_create_private_endpoint]
+
+ # [START howto_operator_neptune_analytics_delete_private_endpoint]
+ delete_endpoint = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="delete_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ vpc_id="{{ ti.xcom_pull(task_ids='create_endpoint')['vpc_id'] }}",
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_delete_private_endpoint]
+
+ # [START howto_operator_neptune_analytics_start_import_task]
+ start_import = NeptuneStartImportTaskOperator(
+ task_id="start_import",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ role_arn=create_role,
+ source=f"s3://{bucket_name}/data/",
+ format="CSV",
+ fail_on_error=True,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_start_import_task]
+
+ # [START howto_operator_neptune_analytics_cancel_import_task]
+ cancel_import = NeptuneCancelImportTaskOperator(
Review Comment:
wait_for_completion=True on start_import makes cancel_import impossible and by the time the cancel runs, the import is already SUCCEEDED and AWS returns ConflictException: Import task ... has completed and cannot be canceled.
Set this to wait_for_completion=False so cancel_import races a still-running task. Note the 2-row CSV completes in seconds even then, so consider a larger source (or accepting SUCCEEDED as a valid terminal state in the cancel waiter) to avoid an intermittent race on slow runners.
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,335 @@
+# Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Ran the system test against real AWS (us-west-2). Graph create, endpoint create, and import all worked. Two issues blocked the run from passing:
1. cancel_import always fails — start_import.wait_for_completion=True means the import is already SUCCEEDED when cancel runs (ConflictException: ... has completed and cannot be canceled).
2. Cleanup leaked a graph + VPC endpoint + IAM role: the delete_graph_if_exists safety net tries DeleteGraph without first deleting the attached private endpoint, and with contextlib.suppress(Exception): hides the resulting ConflictException. delete_neptune_import_role only deletes the inline policy, never the role itself. Had to clean up manually.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None:
+ raise NeptuneGraphCreationFailedError(
+ "No event received while waiting for Neptune graph creation to complete."
+ )
+
+ if event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ event.get(
+ "message",
+ f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
Review Comment:
self.defer() on line 298 raises TaskDeferred, so every line from there to the return is unreachable in deferrable mode — including this persist() call. Users running with deferrable=True won't see the VPC Endpoint link in the UI. Move the persist (and endpoint_url log line) above the defer block, using result["vpcEndpointId"] which is already available from the create_private_graph_endpoint response. NeptuneCreateGraphOperator already follows this pattern at line ~165.
##########
providers/amazon/tests/unit/amazon/aws/hooks/test_neptune_analytics.py:
##########
@@ -0,0 +1,46 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Generator
+from unittest import mock
+
+import pytest
+from moto import mock_aws
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+
+
+@pytest.fixture
+def neptune_hook() -> Generator[NeptuneAnalyticsHook, None, None]:
+ """Returns a NeptuneAnalyticsHook mocked with moto"""
+ with mock_aws():
+ yield NeptuneAnalyticsHook(aws_conn_id="aws_default")
+
+
+class TestNeptuneAnalyticsHook:
+ graph_id = "abc123"
+
+ def test_get_conn_returns_a_boto3_connection(self):
+ hook = NeptuneAnalyticsHook(aws_conn_id="aws_default")
+ assert hook.get_conn() is not None
+
+ @mock.patch.object(NeptuneAnalyticsHook, "get_waiter")
+ def test_wait_for_graph_availability(self, mock_get_waiter, neptune_hook: NeptuneAnalyticsHook):
+ waiter = mock_get_waiter("graph_available")
Review Comment:
you can potentially assert if the client name here. Not sure if this test adds any value.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None:
+ raise NeptuneGraphCreationFailedError(
+ "No event received while waiting for Neptune graph creation to complete."
+ )
+
+ if event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ event.get(
+ "message",
+ f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
Review Comment:
Other Amazon deferrable operators use validate_execute_complete_event(event) and pull identifiers from the event dict. I dont think we can use self.graph_id across the defer/resume boundary. Also, it is a good practice to keep all operators follow the same convention.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None:
+ raise NeptuneGraphCreationFailedError(
+ "No event received while waiting for Neptune graph creation to complete."
+ )
+
+ if event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ event.get(
+ "message",
+ f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+ if event:
+ status = event["status"]
+ if status.lower() != "success":
+ raise NeptunePrivateEndpointCreationFailedError("Endpoint failed to create")
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ f"Failed to delete private endpoint {endpoint_id}"
+ )
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and (event.get("status") or "").lower() == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+ else:
+ raise NeptunePrivateEndpointDeletionFailedError("Endpoint failed to delete.")
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise NeptuneGraphDeletionFailedError(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ if event:
+ graph_id = event.get("graph_id", "Unknown")
Review Comment:
The base trigger returns an event with status error.
event["status"] is never read. If the trigger emits {"status": "error", "graph_id": "g-foo", ...}, this method logs "Neptune graph g-foo deleted" and returns None. Task = SUCCESS.
##########
providers/amazon/src/airflow/providers/amazon/aws/links/ec2.py:
##########
@@ -44,3 +44,14 @@ class EC2InstanceDashboardLink(BaseAwsLink):
@staticmethod
def format_instance_id_filter(instance_ids: list[str]) -> str:
return ",:".join(instance_ids)
+
+
+class VpcEndpointLink(BaseAwsLink):
Review Comment:
Wondering if this should be in vpc.py
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,335 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.common.compat.sdk import DAG, chain, task
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_neptune_analytics"
+
+ROLE_ARN_KEY = "ROLE_ARN"
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
Review Comment:
doesnt look like we are using this as dag is creating it's own role
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
Review Comment:
Inconsistent names here. we can match the parm names
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None:
+ raise NeptuneGraphCreationFailedError(
+ "No event received while waiting for Neptune graph creation to complete."
+ )
+
+ if event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ event.get(
+ "message",
+ f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ endpoint_id = self._get_graph_endpoint_id()
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def _get_graph_endpoint_id(self):
+ """Return the vpc endpoint id for this graph."""
+ result = self.hook.conn.get_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+ return result.get("vpcEndpointId")
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ vpc_endpoint_id = self._get_graph_endpoint_id()
+ if event:
+ status = event["status"]
+ if status.lower() != "success":
+ raise NeptunePrivateEndpointCreationFailedError("Endpoint failed to create")
+
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ f"Failed to delete private endpoint {endpoint_id}"
+ )
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ vpc_endpoint_id = ""
+
+ if event and (event.get("status") or "").lower() == "success":
+ vpc_endpoint_id = event.get("endpoint_id", "Unknown")
+
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+ else:
+ raise NeptunePrivateEndpointDeletionFailedError("Endpoint failed to delete.")
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise NeptuneGraphDeletionFailedError(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ if event:
+ graph_id = event.get("graph_id", "Unknown")
+
+ self.log.info("Neptune graph %s deleted", graph_id)
+
+
+class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph and imports data into it.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics,
+ you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphWithImportOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param source: The source from which to import data. Can be an S3 URI or Neptune database snapshot.
+ :param role_arn: The ARN of the IAM role that Neptune Analytics can assume to access the data source.
+ :param blank_node_handling: The method to handle blank nodes in the dataset. Options include 'convertToIri' or other handling strategies.
+ :param parquet_type: The type of Parquet files in the data source (if applicable).
+ :param format: The format of the data to be imported (e.g., 'csv', 'opencypher', 'ntriples', 'nquads', 'rdfxml', 'turtle').
+ :param min_provisioned_memory: The minimum provisioned memory for the graph in GBs.
+ :param max_provisioned_memory: The maximum provisioned memory for the graph in GBs.
+ :param fail_on_error: If True, the import will fail if any errors are encountered. If False, the import will continue despite errors.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled. (default: False)
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param import_options: Contains options for controlling the import process.
+ :param wait_for_completion: Whether to wait for the graph to be created and data imported. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 60)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be created and data imported.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "source", "role_arn", "kms_key"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (
+ NeptuneImportTaskLink(),
+ NeptuneGraphLink(),
+ )
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ source: str,
+ role_arn: str,
+ blank_node_handling: str | None = None,
+ parquet_type: str | None = None,
+ format: str | None = None,
+ min_provisioned_memory: int | None = None,
+ max_provisioned_memory: int | None = None,
+ fail_on_error: bool | None = None,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool | None = None,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ import_options: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.source = source
+ self.role_arn = role_arn
+ self.blank_node_handling = blank_node_handling
+ self.parquet_type = parquet_type
+ self.format = format
+ self.min_provisioned_memory = min_provisioned_memory
+ self.max_provisioned_memory = max_provisioned_memory
+ self.fail_on_error = fail_on_error
+ self.public_connectivity = public_connectivity
+ self.replica_count = replica_count
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.import_options = import_options
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s with import", self.graph_name)
+
+ # Build the import options
+ import_options = {
+ "neptune-analytics:blank-node-handling": self.blank_node_handling,
+ "neptune-analytics:parquet-type": self.parquet_type,
+ }
+
+ # Remove None values from import_options
+ import_options = {k: v for k, v in import_options.items() if v is not None}
+
+ # Merge with user-provided import_options
+ if self.import_options:
+ import_options.update(self.import_options)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "source": self.source,
+ "roleArn": self.role_arn,
+ **{
+ k: v
+ for k, v in {
+ "format": self.format,
+ "minProvisionedMemory": self.min_provisioned_memory,
+ "maxProvisionedMemory": self.max_provisioned_memory,
+ "failOnError": self.fail_on_error,
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ "importOptions": import_options if import_options else None,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph_using_import_task(**create_params)
+
+ self.log.info("Graph %s import task in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("graphId", None)
+ import_task_id = response.get("taskId")
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="defer_wait_for_task",
+ kwargs={"import_task_id": import_task_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ # Once the graph is available, wait for the task to complete
+
+ self.log.info("Waiting for import task %s", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def defer_wait_for_task(
+ self, context: Context, event: dict[str, Any] | None = None, import_task_id: str | None = None
+ ) -> None:
+ """Defers for import task completion."""
+ if import_task_id:
+ self.log.info("Deferring for import task %s completion", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None or event.get("status") != "success":
+ message = (event or {}).get(
+ "message", f"Neptune graph {self.graph_id} import did not complete successfully"
+ )
+ raise NeptuneGraphCreationFailedError(message)
+ self.log.info("Import complete for graph %s", self.graph_id)
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Starts a bulk data import task to load data into an empty Neptune graph.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneStartImportTaskOperator`
+
+ :param graph_identifier: Graph Id of target Neptune Graph
+ :param role_arn: IAM role ARN granting access to source data
+ :param source: URL identifying the source data location.
+ :param blank_node_handling: Method to handle blank nodes in dataset.
+ :param fail_on_error: If set to true, the task halts when an import error is encountered. If set to false, the task skips the data that caused the error and continues if possible.
+ :param format: Specifies the format of the Amazon S3 data to be imported.
+ :param import_options: Options on how to perform an import
+ :param parquet_type: Parquet type of import task
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "role_arn", "source", "import_options"
+ )
+ template_fields_renderers = {
+ "import_options": "json",
+ }
+ operator_extra_links = (NeptuneImportTaskLink(),)
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ role_arn: str,
+ source: str,
+ blank_node_handling: str | None = None,
+ fail_on_error: bool = True,
+ format: str | None = None,
+ import_options: dict | None = None,
+ parquet_type: str | None = "COLUMNAR",
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.role_arn = role_arn
+ self.source = source
+ self.blank_node_handling = blank_node_handling
+ self.fail_on_error = fail_on_error
+ self.format = format
+ self.import_options = import_options
+ self.parquet_type = parquet_type
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Starting data import to graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ "roleArn": self.role_arn,
+ "source": self.source,
+ **{
+ k: v
+ for k, v in {
+ "blankNodeHandling": self.blank_node_handling,
+ "failOnError": self.fail_on_error,
+ "format": self.format,
+ "importOptions": self.import_options,
+ "parquetType": self.parquet_type,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.start_import_task(**create_params)
+ import_task_id = response.get("taskId")
+ self.log.info("Import task %s started for graph %s", import_task_id, self.graph_identifier)
+
+ # Create the console link
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s completes", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to complete", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": import_task_id, "graph_id": self.graph_identifier}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ task_id = ""
+ if event:
+ task_id = event.get("import_task_id", "")
Review Comment:
same issue. I think we need to check if the trigger event has returned success or error.\
NeptuneCreateGraphOperator.execute_complete (line 190) and NeptuneCancelImportTaskOperator.execute_complete (line 991) already do the check correctly — please apply the same pattern, ideally via
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3380280186
##########
providers/amazon/docs/operators/neptune_analytics.rst:
##########
@@ -0,0 +1,148 @@
+ .. Licensed to the Apache Software Foundation (ASF) under one
+ or more contributor license agreements. See the NOTICE file
+ distributed with this work for additional information
+ regarding copyright ownership. The ASF licenses this file
+ to you under the Apache License, Version 2.0 (the
+ "License"); you may not use this file except in compliance
+ with the License. You may obtain a copy of the License at
+
+ .. http://www.apache.org/licenses/LICENSE-2.0
+
+ .. Unless required by applicable law or agreed to in writing,
+ software distributed under the License is distributed on an
+ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ KIND, either express or implied. See the License for the
+ specific language governing permissions and limitations
+ under the License.
+
+========================
+Amazon Neptune Analytics
+========================
+
+`Amazon Neptune Analytics <https://docs.aws.amazon.com/neptune-analytics/latest/userguide/what-is-neptune-analytics.html>`__ is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+Prerequisite Tasks
+------------------
+
+.. include:: ../_partials/prerequisite_tasks.rst
+
+Generic Parameters
+------------------
+
+.. include:: ../_partials/generic_parameters.rst
+
+Operators
+---------
+
+.. _howto/operator:NeptuneCreateGraphOperator:
+
+Create a new Neptune Graph
+==========================
+
+To create a new Neptune Analytics Graph, you can use
+:class:`~airflow.providers.amazon.aws.operators.neptune.NeptuneCreateGraphOperator`.
+This operator can be run in deferrable mode by passing ``deferrable=True`` as a parameter. This requires
+the aiobotocore module to be installed.
+
Review Comment:
This is fixed.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3462943409
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ # get the vpce id since it may not be returned immediately
+ endpoint_id = self.hook._get_graph_endpoint_id(graph_id=self.graph_identifier, vpc_id=self.vpc_id)
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ kwargs={"vpc_id": self.vpc_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptunePrivateEndpointCreationFailedError(
+ validated_event.get("message", "Endpoint failed to create")
+ )
+
+ graph_id = validated_event["value"]
+ vpc_id = validated_event["vpc_id"]
+ vpc_endpoint_id = self.hook._get_graph_endpoint_id(graph_id=graph_id, vpc_id=vpc_id)
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": graph_id, "vpc_id": vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ f"Failed to delete private endpoint {endpoint_id}"
+ )
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ validated_event.get("message", "Endpoint failed to delete.")
+ )
+
+ vpc_endpoint_id = validated_event.get("endpoint_id", "Unknown")
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise NeptuneGraphDeletionFailedError(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ validated_event = validate_execute_complete_event(event)
+ graph_id = validated_event.get("graph_id")
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphDeletionFailedError(
+ validated_event.get("message", f"Neptune graph {graph_id} deletion failed")
+ )
+
+ self.log.info("Neptune graph %s deleted", validated_event.get("graph_id", graph_id))
+
+
+class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph and imports data into it.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics,
+ you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphWithImportOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param source: The source from which to import data. Can be an S3 URI or Neptune database snapshot.
+ :param role_arn: The ARN of the IAM role that Neptune Analytics can assume to access the data source.
+ :param blank_node_handling: The method to handle blank nodes in the dataset. Options include 'convertToIri' or other handling strategies.
+ :param parquet_type: The type of Parquet files in the data source (if applicable).
+ :param format: The format of the data to be imported (e.g., 'csv', 'opencypher', 'ntriples', 'nquads', 'rdfxml', 'turtle').
+ :param min_provisioned_memory: The minimum provisioned memory for the graph in GBs.
+ :param max_provisioned_memory: The maximum provisioned memory for the graph in GBs.
+ :param fail_on_error: If True, the import will fail if any errors are encountered. If False, the import will continue despite errors.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled. (default: False)
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param import_options: Contains options for controlling the import process.
+ :param wait_for_completion: Whether to wait for the graph to be created and data imported. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 60)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be created and data imported.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "source", "role_arn"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (
+ NeptuneImportTaskLink(),
+ NeptuneGraphLink(),
+ )
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ source: str,
+ role_arn: str,
+ blank_node_handling: str | None = None,
+ parquet_type: str | None = None,
+ format: str | None = None,
+ min_provisioned_memory: int | None = None,
+ max_provisioned_memory: int | None = None,
+ fail_on_error: bool | None = None,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool | None = None,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ import_options: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.source = source
+ self.role_arn = role_arn
+ self.blank_node_handling = blank_node_handling
+ self.parquet_type = parquet_type
+ self.format = format
+ self.min_provisioned_memory = min_provisioned_memory
+ self.max_provisioned_memory = max_provisioned_memory
+ self.fail_on_error = fail_on_error
+ self.public_connectivity = public_connectivity
+ self.replica_count = replica_count
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.import_options = import_options
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s with import", self.graph_name)
+
+ # Build the import options
+ import_options = {
+ "neptune-analytics:blank-node-handling": self.blank_node_handling,
+ "neptune-analytics:parquet-type": self.parquet_type,
+ }
+
+ # Remove None values from import_options
+ import_options = {k: v for k, v in import_options.items() if v is not None}
+
+ # Merge with user-provided import_options
+ if self.import_options:
+ import_options.update(self.import_options)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "source": self.source,
+ "roleArn": self.role_arn,
+ **{
+ k: v
+ for k, v in {
+ "format": self.format,
+ "minProvisionedMemory": self.min_provisioned_memory,
+ "maxProvisionedMemory": self.max_provisioned_memory,
+ "failOnError": self.fail_on_error,
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ "importOptions": import_options if import_options else None,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph_using_import_task(**create_params)
+
+ self.log.info("Graph %s import task in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("graphId", None)
+ import_task_id = response.get("taskId")
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="defer_wait_for_task",
+ kwargs={"import_task_id": import_task_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ # Once the graph is available, wait for the task to complete
+
+ self.log.info("Waiting for import task %s", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def defer_wait_for_task(
+ self, context: Context, event: dict[str, Any] | None = None, import_task_id: str | None = None
+ ) -> None:
+ """Defers for import task completion."""
+ validated_event = validate_execute_complete_event(event)
+ graph_id = validated_event.get("value")
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get("message", f"Neptune graph {graph_id} did not become available")
+ )
+
+ if import_task_id:
+ self.log.info("Deferring for import task %s completion", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ kwargs={"graph_id": graph_id},
+ )
+
+ def execute_complete(
+ self, context: Context, event: dict[str, Any] | None = None, graph_id: str | None = None
+ ) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message", f"Neptune graph {graph_id} import did not complete successfully"
+ )
+ )
+
+ self.log.info("Import complete for graph %s", graph_id)
+ return {"graph_id": graph_id}
+
+
+class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Starts a bulk data import task to load data into an empty Neptune graph.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneStartImportTaskOperator`
+
+ :param graph_identifier: Graph Id of target Neptune Graph
+ :param role_arn: IAM role ARN granting access to source data
+ :param source: URL identifying the source data location.
+ :param blank_node_handling: Method to handle blank nodes in dataset.
+ :param fail_on_error: If set to true, the task halts when an import error is encountered. If set to false, the task skips the data that caused the error and continues if possible.
+ :param format: Specifies the format of the Amazon S3 data to be imported.
+ :param import_options: Options on how to perform an import
+ :param parquet_type: Parquet type of import task
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "role_arn", "source", "import_options"
+ )
+ template_fields_renderers = {
+ "import_options": "json",
+ }
+ operator_extra_links = (NeptuneImportTaskLink(),)
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ role_arn: str,
+ source: str,
+ blank_node_handling: str | None = None,
+ fail_on_error: bool = True,
+ format: str | None = None,
+ import_options: dict | None = None,
+ parquet_type: str | None = "COLUMNAR",
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.role_arn = role_arn
+ self.source = source
+ self.blank_node_handling = blank_node_handling
+ self.fail_on_error = fail_on_error
+ self.format = format
+ self.import_options = import_options
+ self.parquet_type = parquet_type
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Starting data import to graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ "roleArn": self.role_arn,
+ "source": self.source,
+ **{
+ k: v
+ for k, v in {
+ "blankNodeHandling": self.blank_node_handling,
+ "failOnError": self.fail_on_error,
+ "format": self.format,
+ "importOptions": self.import_options,
+ "parquetType": self.parquet_type,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.start_import_task(**create_params)
+ import_task_id = response.get("taskId")
+ self.log.info("Import task %s started for graph %s", import_task_id, self.graph_identifier)
+
+ # Create the console link
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s completes", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to complete", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": import_task_id, "graph_id": self.graph_identifier}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneImportTaskFailedError(
+ validated_event.get("message", "Import task did not complete successfully")
+ )
+
+ task_id = validated_event.get("import_task_id", "")
+ self.log.info("Import task %s completed", task_id)
+ return {"graph_id": self.graph_identifier, "import_task_id": task_id}
+
+
+class NeptuneCancelImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Cancels an active Neptune Graph import task.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCancelImportTaskOperator`
+
+ :param import_task_id: Neptune Graph import task id to cancel.
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("import_task_id")
+
+ def __init__(
+ self,
+ import_task_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.import_task_id = import_task_id
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Cancelling import task %s", self.import_task_id)
+
+ response = self.hook.conn.cancel_import_task(taskIdentifier=self.import_task_id)
+
+ self.log.info("Import task %s status is %s", self.import_task_id, response.get("status", "Unknown"))
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s is cancelled", self.import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCancelledTrigger(
+ task_identifier=self.import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to be cancelled", self.import_task_id)
+ self.hook.get_waiter("import_task_cancelled").wait(
+ taskIdentifier=self.import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": self.import_task_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneImportTaskCancellationFailedError(
+ validated_event.get("message", "Error while waiting for Neptune import task cancellation")
+ )
+
+ task_id = validated_event.get("value", "")
Review Comment:
yep, fixed
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4836370609
Resolved conflicts, tested and resolved deferrable issues.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4663111066
Fixed the mypy error and resolved bot comments. Ready for re-review.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "eladkal (via GitHub)" <gi...@apache.org>.
eladkal commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3284627071
##########
airflow-core/src/airflow/migrations/versions/0113_3_3_0_add_retry_policy_fields_to_ti.py:
##########
@@ -30,7 +30,7 @@
is a metadata-only operation (no table rewrite).
Revision ID: b8f3e4a1d2c9
-Revises: 9fabad868fdb
+Revises: fde9ed84d07b
Create Date: 2026-04-16 12:00:00.000000
Review Comment:
Probably bad rebase. Just undo changes to this file
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3283847861
##########
providers/amazon/src/airflow/providers/amazon/aws/links/ec2.py:
##########
@@ -44,3 +44,14 @@ class EC2InstanceDashboardLink(BaseAwsLink):
@staticmethod
def format_instance_id_filter(instance_ids: list[str]) -> str:
return ",:".join(instance_ids)
+
+
+class VpcEndpointLink(BaseAwsLink):
Review Comment:
I wasn't sure about this. I opted for ec2.py because all the VPC-related actions fall under the ec2 client. I'd probably search under VPC before I searched under ec2 if I were looking for this. What's the preference?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4558070368
@ellisms Converting to **draft** — this PR doesn't yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
- :warning: **Unaddressed Copilot review** — 2 unresolved [Copilot review threads](https://github.com/apache/airflow/pull/64274/files) from 2026-04-02 and 2026-04-10 (55+ days old). Please address them or mark them resolved if no longer applicable.
- :x: **mypy (type checking)** — Failing: `MyPy providers checks`. See [docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
See the linked criteria for how to fix each item, then mark the PR "Ready for review". This is **not** a rejection — just an invitation to bring the PR up to standard. No rush.
---
_Note: This comment was drafted by an AI-assisted triage tool and may contain mistakes. Once you have addressed the points above, an Apache Airflow maintainer — a real person — will take the next look at your PR. We use this [two-stage triage process](https://github.com/apache/airflow/blob/main/contributing-docs/25_maintainer_pr_triage.md#why-the-first-pass-is-automated) so that our maintainers' limited time is spent where it matters most: the conversation with you._
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3462926650
##########
providers/amazon/src/airflow/providers/amazon/aws/waiters/neptune_analytics.json:
##########
@@ -0,0 +1,38 @@
+{
+ "version": 2,
+ "waiters": {
+
+ "import_task_cancelled":{
+ "operation": "GetImportTask",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "SUCCEEDED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "CANCELLED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "ERROR_ENCOUNTERED",
+ "state": "error"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
Review Comment:
Do you mean accepting FAILED as successful? If so, I made the change based on [this comment](https://github.com/apache/airflow/pull/64274#discussion_r3283359724). I'm on the fence about it and I could go either way. Would like other's input.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4165209646
Hey @ellisms, thanks for working on this one!
I see that you requested a review on this one but the build is still failing? Is there a reason you want an early review?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "potiuk (via GitHub)" <gi...@apache.org>.
potiuk commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4169930549
@ellisms This PR has been converted to **draft** because it does not yet meet our [Pull Request quality criteria](https://github.com/apache/airflow/blob/main/contributing-docs/05_pull_requests.rst#pull-request-quality-criteria).
**Issues found:**
- :x: **Pre-commit / static checks**: Failing: CI image checks / Static checks. Run `prek run --from-ref main` locally to find and fix issues. See [Pre-commit / static checks docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
- :x: **mypy (type checking)**: Failing: CI image checks / MyPy checks (mypy-providers). Run `prek --stage manual mypy-providers --all-files` locally to reproduce. You need `breeze ci-image build --python 3.10` for Docker-based mypy. See [mypy (type checking) docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst#mypy-checks).
- :x: **Other failing CI checks**: Failing: check-newsfragment-pr-number, Postgres tests: core / DB-core:Postgres:14:3.10:API...CLI, Postgres tests: core / DB-core:Postgres:14:3.10:Core...Serialization, MySQL tests: core / DB-core:MySQL:8.0:3.10:API...CLI, MySQL tests: core / DB-core:MySQL:8.0:3.10:Core...Serialization (+9 more). Run `prek run --from-ref main` locally to reproduce. See [static checks docs](https://github.com/apache/airflow/blob/main/contributing-docs/08_static_code_checks.rst).
> **Note:** Your branch is **123 commits behind `main`**. Some check failures may be caused by changes in the base branch rather than by your PR. Please rebase your branch and push again to get up-to-date CI results.
**What to do next:**
- The comment informs you what you need to do.
- Fix each issue, then mark the PR as "Ready for review" in the GitHub UI - but only after making sure that all the issues are fixed.
- There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates.
- Maintainers will then proceed with a normal review.
Converting a PR to draft is **not** a rejection — it is an invitation to bring the PR up to the project's standards so that maintainer review time is spent productively. There is no rush — take your time and work at your own pace. We appreciate your contribution and are happy to wait for updates. If you have questions, feel free to ask on the [Airflow Slack](https://s.apache.org/airflow-slack).
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3289985576
##########
providers/amazon/tests/system/amazon/aws/example_neptune_analytics.py:
##########
@@ -0,0 +1,335 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+import contextlib
+import json
+import time
+from datetime import datetime
+
+import boto3
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.operators.neptune_analytics import (
+ NeptuneCancelImportTaskOperator,
+ NeptuneCreateGraphOperator,
+ NeptuneCreateGraphWithImportOperator,
+ NeptuneCreatePrivateGraphEndpointOperator,
+ NeptuneDeleteGraphOperator,
+ NeptuneDeletePrivateGraphEndpointOperator,
+ NeptuneStartImportTaskOperator,
+)
+from airflow.providers.amazon.aws.operators.s3 import (
+ S3CreateBucketOperator,
+ S3CreateObjectOperator,
+ S3DeleteBucketOperator,
+)
+from airflow.providers.common.compat.sdk import DAG, chain, task
+
+try:
+ from airflow.sdk import TriggerRule
+except ImportError:
+ # Compatibility for Airflow < 3.1
+ from airflow.utils.trigger_rule import TriggerRule # type: ignore[no-redef,attr-defined]
+
+from system.amazon.aws.utils import SystemTestContextBuilder
+
+DAG_ID = "example_neptune_analytics"
+
+ROLE_ARN_KEY = "ROLE_ARN"
+
+sys_test_context_task = SystemTestContextBuilder().add_variable(ROLE_ARN_KEY).build()
+
+# Minimal OpenCypher CSV data for import testing.
+NODES_CSV = """~id,~label,name:String
+n1,Person,Alice
+n2,Person,Bob
+"""
+
+EDGES_CSV = """~id,~from,~to,~label
+e1,n1,n2,KNOWS
+"""
+
+NEPTUNE_ANALYTICS_TRUST_POLICY = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Principal": {"Service": "neptune-graph.amazonaws.com"},
+ "Action": "sts:AssumeRole",
+ }
+ ],
+ }
+)
+
+S3_READ_POLICY_DOCUMENT = json.dumps(
+ {
+ "Version": "2012-10-17",
+ "Statement": [
+ {
+ "Effect": "Allow",
+ "Action": ["s3:GetObject", "s3:ListBucket"],
+ "Resource": ["arn:aws:s3:::*", "arn:aws:s3:::*/*"],
+ }
+ ],
+ }
+)
+
+
+@task
+def create_neptune_import_role(role_name: str) -> str:
+ iam_client = boto3.client("iam")
+ iam_client.create_role(
+ RoleName=role_name,
+ AssumeRolePolicyDocument=NEPTUNE_ANALYTICS_TRUST_POLICY,
+ Description="Role for Neptune Analytics import system test",
+ )
+ iam_client.put_role_policy(
+ RoleName=role_name,
+ PolicyName="NeptuneAnalyticsS3Access",
+ PolicyDocument=S3_READ_POLICY_DOCUMENT,
+ )
+ role = iam_client.get_role(RoleName=role_name)
+ time.sleep(60) # Wait for IAM eventual consistency (role + inline policy propagation)
+ return role["Role"]["Arn"]
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_neptune_import_role(role_name: str) -> None:
+ iam_client = boto3.client("iam")
+ with contextlib.suppress(iam_client.exceptions.NoSuchEntityException):
+ iam_client.delete_role_policy(RoleName=role_name, PolicyName="NeptuneAnalyticsS3Access")
+
+
+@task(trigger_rule=TriggerRule.ALL_DONE)
+def delete_graph_if_exists(graph_name: str) -> None:
+ """Safety net to clean up the graph in case a previous task failed."""
+ hook = NeptuneAnalyticsHook()
+ with contextlib.suppress(Exception):
+ # List graphs and find by name
+ paginator = hook.conn.get_paginator("list_graphs")
+ for page in paginator.paginate():
+ for graph in page.get("graphs", []):
+ if graph.get("name") == graph_name:
+ graph_id = graph["id"]
+ # Disable deletion protection if enabled
+
+ hook.conn.update_graph(graphIdentifier=graph_id, deletionProtection=False)
+
+ hook.conn.delete_graph(graphIdentifier=graph_id, skipSnapshot=True)
+ hook.conn.get_waiter("graph_deleted").wait(
+ graphIdentifier=graph_id,
+ WaiterConfig={"Delay": 30, "MaxAttempts": 60},
+ )
+ return
+
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule="@once",
+ start_date=datetime(2021, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+
+ env_id = test_context["ENV_ID"]
+ graph_name = f"{env_id}-graph"
+ import_graph_name = f"{env_id}-import-graph"
+ bucket_name = f"{env_id}-neptune-analytics"
+ import_role_name = f"{env_id}-neptune-import"
+ region = boto3.session.Session().region_name
+
+ # --- TEST SETUP ---
+
+ create_bucket = S3CreateBucketOperator(
+ task_id="create_bucket",
+ bucket_name=bucket_name,
+ )
+
+ upload_nodes = S3CreateObjectOperator(
+ task_id="upload_nodes",
+ s3_bucket=bucket_name,
+ s3_key="data/nodes.csv",
+ data=NODES_CSV,
+ replace=True,
+ )
+
+ upload_edges = S3CreateObjectOperator(
+ task_id="upload_edges",
+ s3_bucket=bucket_name,
+ s3_key="data/edges.csv",
+ data=EDGES_CSV,
+ replace=True,
+ )
+
+ create_role = create_neptune_import_role(import_role_name)
+
+ # --- TEST BODY ---
+
+ # [START howto_operator_neptune_analytics_create_graph]
+ create_graph = NeptuneCreateGraphOperator(
+ task_id="create_graph",
+ graph_name=graph_name,
+ vector_search_config={"dimension": 128},
+ provisioned_memory=32,
+ public_connectivity=True,
+ replica_count=0,
+ deletion_protection=False,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_create_graph]
+
+ # [START howto_operator_neptune_analytics_create_private_endpoint]
+ create_endpoint = NeptuneCreatePrivateGraphEndpointOperator(
+ task_id="create_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id']}}",
+ wait_for_completion=True,
+ )
+ # [END howto_operator_neptune_analytics_create_private_endpoint]
+
+ # [START howto_operator_neptune_analytics_delete_private_endpoint]
+ delete_endpoint = NeptuneDeletePrivateGraphEndpointOperator(
+ task_id="delete_endpoint",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ vpc_id="{{ ti.xcom_pull(task_ids='create_endpoint')['vpc_id'] }}",
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_delete_private_endpoint]
+
+ # [START howto_operator_neptune_analytics_start_import_task]
+ start_import = NeptuneStartImportTaskOperator(
+ task_id="start_import",
+ graph_identifier="{{ ti.xcom_pull(task_ids='create_graph')['graph_id'] }}",
+ role_arn=create_role,
+ source=f"s3://{bucket_name}/data/",
+ format="CSV",
+ fail_on_error=True,
+ wait_for_completion=True,
+ deferrable=False,
+ waiter_delay=30,
+ waiter_max_attempts=60,
+ )
+ # [END howto_operator_neptune_analytics_start_import_task]
+
+ # [START howto_operator_neptune_analytics_cancel_import_task]
+ cancel_import = NeptuneCancelImportTaskOperator(
Review Comment:
I'll create a custom waiter that accepts SUCCEEDED as a valid terminal state.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3455806780
##########
providers/amazon/src/airflow/providers/amazon/aws/waiters/neptune_analytics.json:
##########
@@ -0,0 +1,38 @@
+{
+ "version": 2,
+ "waiters": {
+
+ "import_task_cancelled":{
+ "operation": "GetImportTask",
+ "delay": 30,
+ "maxAttempts": 60,
+ "acceptors": [
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "SUCCEEDED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "CANCELLED",
+ "state": "success"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
+ "expected": "ERROR_ENCOUNTERED",
+ "state": "error"
+ },
+ {
+ "matcher": "path",
+ "argument": "status",
Review Comment:
These don't seem to match?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
Review Comment:
This should come from the event? `self` won't contain this on resumption
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ # get the vpce id since it may not be returned immediately
+ endpoint_id = self.hook._get_graph_endpoint_id(graph_id=self.graph_identifier, vpc_id=self.vpc_id)
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
Review Comment:
You need to declare operator_extra_links for this to show up I believe, look for other examples in the codebase.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ # get the vpce id since it may not be returned immediately
+ endpoint_id = self.hook._get_graph_endpoint_id(graph_id=self.graph_identifier, vpc_id=self.vpc_id)
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ kwargs={"vpc_id": self.vpc_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptunePrivateEndpointCreationFailedError(
+ validated_event.get("message", "Endpoint failed to create")
+ )
+
+ graph_id = validated_event["value"]
+ vpc_id = validated_event["vpc_id"]
+ vpc_endpoint_id = self.hook._get_graph_endpoint_id(graph_id=graph_id, vpc_id=vpc_id)
+ return {"vpc_endpoint_id": vpc_endpoint_id, "graph_id": graph_id, "vpc_id": vpc_id}
+
+
+class NeptuneDeletePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeletePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC where endpoint resides
+ :param wait_for_completion: Whether to wait for the endpoint to be deleted. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_identifier", "vpc_id")
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> None:
+ self.log.info("Deleting private endpoint for graph %s", self.graph_identifier)
+
+ result = self.hook.conn.delete_private_graph_endpoint(
+ graphIdentifier=self.graph_identifier, vpcId=self.vpc_id
+ )
+
+ status = result.get("status")
+ endpoint_id = result.get("vpcEndpointId")
+
+ if status == "FAILED":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ f"Failed to delete private endpoint {endpoint_id}"
+ )
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint %s is deleted", endpoint_id)
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ endpoint_id=endpoint_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint %s is deleted", endpoint_id)
+ self.hook.get_waiter("private_graph_endpoint_deleted").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ self.log.info("Endpoint %s deleted", endpoint_id)
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> None:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptunePrivateEndpointDeletionFailedError(
+ validated_event.get("message", "Endpoint failed to delete.")
+ )
+
+ vpc_endpoint_id = validated_event.get("endpoint_id", "Unknown")
+ self.log.info("Endpoint id %s deleted", vpc_endpoint_id)
+
+
+class NeptuneDeleteGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Deletes an Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneDeleteGraphOperator`
+
+ :param graph_id: Name of Neptune graph to delete
+ :param skip_snapshot: Determines whether a final graph snapshot is created before the graph is deleted. If true is specified, no graph snapshot is created. If false is specified, a graph snapshot is created before the graph is deleted.
+ :param wait_for_completion: Whether to wait for the graph to delete. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be deleted.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("graph_id", "skip_snapshot")
+
+ def __init__(
+ self,
+ graph_id: str,
+ skip_snapshot: bool,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_id = graph_id
+ self.skip_snapshot = skip_snapshot
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context):
+ self.log.info("Deleting graph %s", self.graph_id)
+
+ try:
+ self.hook.conn.delete_graph(graphIdentifier=self.graph_id, skipSnapshot=self.skip_snapshot)
+ except ClientError as e:
+ # if not found, just exit because there is nothing to delete
+ if e.response["Error"]["Code"] == "ResourceNotFoundException":
+ self.log.info("Graph %s not found. Nothing to delete", self.graph_id)
+ return
+ raise NeptuneGraphDeletionFailedError(e.response["Error"])
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is deleted", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphDeletedTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+ if self.wait_for_completion:
+ self.log.info("Waiting to delete %s", self.graph_id)
+
+ self.hook.get_waiter("graph_deleted").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None):
+ validated_event = validate_execute_complete_event(event)
+ graph_id = validated_event.get("graph_id")
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphDeletionFailedError(
+ validated_event.get("message", f"Neptune graph {graph_id} deletion failed")
+ )
+
+ self.log.info("Neptune graph %s deleted", validated_event.get("graph_id", graph_id))
+
+
+class NeptuneCreateGraphWithImportOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph and imports data into it.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics,
+ you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphWithImportOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param source: The source from which to import data. Can be an S3 URI or Neptune database snapshot.
+ :param role_arn: The ARN of the IAM role that Neptune Analytics can assume to access the data source.
+ :param blank_node_handling: The method to handle blank nodes in the dataset. Options include 'convertToIri' or other handling strategies.
+ :param parquet_type: The type of Parquet files in the data source (if applicable).
+ :param format: The format of the data to be imported (e.g., 'csv', 'opencypher', 'ntriples', 'nquads', 'rdfxml', 'turtle').
+ :param min_provisioned_memory: The minimum provisioned memory for the graph in GBs.
+ :param max_provisioned_memory: The maximum provisioned memory for the graph in GBs.
+ :param fail_on_error: If True, the import will fail if any errors are encountered. If False, the import will continue despite errors.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled. (default: False)
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param import_options: Contains options for controlling the import process.
+ :param wait_for_completion: Whether to wait for the graph to be created and data imported. (default: True)
+ :param waiter_delay: Time in seconds to wait between status checks. (default: 30)
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion. (default: 60)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to be created and data imported.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "source", "role_arn"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (
+ NeptuneImportTaskLink(),
+ NeptuneGraphLink(),
+ )
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ source: str,
+ role_arn: str,
+ blank_node_handling: str | None = None,
+ parquet_type: str | None = None,
+ format: str | None = None,
+ min_provisioned_memory: int | None = None,
+ max_provisioned_memory: int | None = None,
+ fail_on_error: bool | None = None,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool | None = None,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ import_options: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.source = source
+ self.role_arn = role_arn
+ self.blank_node_handling = blank_node_handling
+ self.parquet_type = parquet_type
+ self.format = format
+ self.min_provisioned_memory = min_provisioned_memory
+ self.max_provisioned_memory = max_provisioned_memory
+ self.fail_on_error = fail_on_error
+ self.public_connectivity = public_connectivity
+ self.replica_count = replica_count
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.import_options = import_options
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s with import", self.graph_name)
+
+ # Build the import options
+ import_options = {
+ "neptune-analytics:blank-node-handling": self.blank_node_handling,
+ "neptune-analytics:parquet-type": self.parquet_type,
+ }
+
+ # Remove None values from import_options
+ import_options = {k: v for k, v in import_options.items() if v is not None}
+
+ # Merge with user-provided import_options
+ if self.import_options:
+ import_options.update(self.import_options)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "source": self.source,
+ "roleArn": self.role_arn,
+ **{
+ k: v
+ for k, v in {
+ "format": self.format,
+ "minProvisionedMemory": self.min_provisioned_memory,
+ "maxProvisionedMemory": self.max_provisioned_memory,
+ "failOnError": self.fail_on_error,
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ "importOptions": import_options if import_options else None,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph_using_import_task(**create_params)
+
+ self.log.info("Graph %s import task in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("graphId", None)
+ import_task_id = response.get("taskId")
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="defer_wait_for_task",
+ kwargs={"import_task_id": import_task_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+ # Once the graph is available, wait for the task to complete
+
+ self.log.info("Waiting for import task %s", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def defer_wait_for_task(
+ self, context: Context, event: dict[str, Any] | None = None, import_task_id: str | None = None
+ ) -> None:
+ """Defers for import task completion."""
+ validated_event = validate_execute_complete_event(event)
+ graph_id = validated_event.get("value")
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get("message", f"Neptune graph {graph_id} did not become available")
+ )
+
+ if import_task_id:
+ self.log.info("Deferring for import task %s completion", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ kwargs={"graph_id": graph_id},
+ )
+
+ def execute_complete(
+ self, context: Context, event: dict[str, Any] | None = None, graph_id: str | None = None
+ ) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message", f"Neptune graph {graph_id} import did not complete successfully"
+ )
+ )
+
+ self.log.info("Import complete for graph %s", graph_id)
+ return {"graph_id": graph_id}
+
+
+class NeptuneStartImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Starts a bulk data import task to load data into an empty Neptune graph.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneStartImportTaskOperator`
+
+ :param graph_identifier: Graph Id of target Neptune Graph
+ :param role_arn: IAM role ARN granting access to source data
+ :param source: URL identifying the source data location.
+ :param blank_node_handling: Method to handle blank nodes in dataset.
+ :param fail_on_error: If set to true, the task halts when an import error is encountered. If set to false, the task skips the data that caused the error and continues if possible.
+ :param format: Specifies the format of the Amazon S3 data to be imported.
+ :param import_options: Options on how to perform an import
+ :param parquet_type: Parquet type of import task
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "role_arn", "source", "import_options"
+ )
+ template_fields_renderers = {
+ "import_options": "json",
+ }
+ operator_extra_links = (NeptuneImportTaskLink(),)
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ role_arn: str,
+ source: str,
+ blank_node_handling: str | None = None,
+ fail_on_error: bool = True,
+ format: str | None = None,
+ import_options: dict | None = None,
+ parquet_type: str | None = "COLUMNAR",
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.role_arn = role_arn
+ self.source = source
+ self.blank_node_handling = blank_node_handling
+ self.fail_on_error = fail_on_error
+ self.format = format
+ self.import_options = import_options
+ self.parquet_type = parquet_type
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Starting data import to graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ "roleArn": self.role_arn,
+ "source": self.source,
+ **{
+ k: v
+ for k, v in {
+ "blankNodeHandling": self.blank_node_handling,
+ "failOnError": self.fail_on_error,
+ "format": self.format,
+ "importOptions": self.import_options,
+ "parquetType": self.parquet_type,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.start_import_task(**create_params)
+ import_task_id = response.get("taskId")
+ self.log.info("Import task %s started for graph %s", import_task_id, self.graph_identifier)
+
+ # Create the console link
+ import_task_url = NeptuneImportTaskLink.format_str.format(
+ import_task_id=import_task_id,
+ aws_domain=NeptuneImportTaskLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneImportTaskLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ import_task_id=import_task_id,
+ )
+ self.log.info("You can view this import task at : %s", import_task_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s completes", import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCompleteTrigger(
+ import_task_id=import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to complete", import_task_id)
+ self.hook.get_waiter("import_task_successful").wait(
+ taskIdentifier=import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": import_task_id, "graph_id": self.graph_identifier}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneImportTaskFailedError(
+ validated_event.get("message", "Import task did not complete successfully")
+ )
+
+ task_id = validated_event.get("import_task_id", "")
+ self.log.info("Import task %s completed", task_id)
+ return {"graph_id": self.graph_identifier, "import_task_id": task_id}
+
+
+class NeptuneCancelImportTaskOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Cancels an active Neptune Graph import task.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCancelImportTaskOperator`
+
+ :param import_task_id: Neptune Graph import task id to cancel.
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields("import_task_id")
+
+ def __init__(
+ self,
+ import_task_id: str,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.import_task_id = import_task_id
+ self.wait_for_completion = wait_for_completion
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+ self.deferrable = deferrable
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Cancelling import task %s", self.import_task_id)
+
+ response = self.hook.conn.cancel_import_task(taskIdentifier=self.import_task_id)
+
+ self.log.info("Import task %s status is %s", self.import_task_id, response.get("status", "Unknown"))
+
+ if self.deferrable:
+ self.log.info("Deferring until import task %s is cancelled", self.import_task_id)
+ self.defer(
+ trigger=NeptuneImportTaskCancelledTrigger(
+ task_identifier=self.import_task_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ aws_conn_id=self.aws_conn_id,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting for import task %s to be cancelled", self.import_task_id)
+ self.hook.get_waiter("import_task_cancelled").wait(
+ taskIdentifier=self.import_task_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"import_task_id": self.import_task_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneImportTaskCancellationFailedError(
+ validated_event.get("message", "Error while waiting for Neptune import task cancellation")
+ )
+
+ task_id = validated_event.get("value", "")
Review Comment:
Shouldn't the key be `import_task_id`?
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1025 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils import validate_execute_complete_event
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptuneImportTaskFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key_id,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ validated_event.get(
+ "message",
+ f"Neptune graph {validated_event.get('return_key')} creation did not complete successfully",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
+
+ return {"graph_id": self.graph_id}
+
+
+class NeptuneCreatePrivateGraphEndpointOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates a Neptune Graph private endpoint.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreatePrivateGraphEndpointOperator`
+
+ :param graph_identifier: Neptune Graph id
+ :param vpc_id: VPC to create endpoint in
+ :param subnet_ids: Subnets in which private graph endpoint ENIs are created
+ :param vpc_security_group_ids: Security groups to be attached to the private graph endpoint
+ :param wait_for_completion: Whether to wait for the endpoint to be available. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the endpoint to become available.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_identifier", "vpc_id", "subnet_ids", "vpc_security_group_ids"
+ )
+
+ def __init__(
+ self,
+ graph_identifier: str,
+ vpc_id: str | None = None,
+ subnet_ids: list[str] | None = None,
+ vpc_security_group_ids: list[str] | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_identifier = graph_identifier
+ self.vpc_id = vpc_id
+ self.subnet_ids = subnet_ids
+ self.vpc_security_group_ids = vpc_security_group_ids
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating private endpoint for graph %s", self.graph_identifier)
+
+ create_params = {
+ "graphIdentifier": self.graph_identifier,
+ **{
+ k: v
+ for k, v in {
+ "vpcId": self.vpc_id,
+ "subnetIds": self.subnet_ids,
+ "vpcSecurityGroupIds": self.vpc_security_group_ids,
+ }.items()
+ if v is not None
+ },
+ }
+
+ # create the endpoint
+
+ result = self.hook.conn.create_private_graph_endpoint(**create_params)
+ status = result.get("status", "Unknown")
+
+ self.log.info("Status of endpoint: %s", status)
+
+ if status in ["FAILED"]:
+ raise NeptunePrivateEndpointCreationFailedError(
+ f"Private endpoint failed to create for graph {self.graph_identifier}"
+ )
+
+ # if VPC not provided, use the one that is returned, which is the default VPC. Required for the waiter
+ self.vpc_id = result.get("vpcId", self.vpc_id)
+
+ # get the vpce id since it may not be returned immediately
+ endpoint_id = self.hook._get_graph_endpoint_id(graph_id=self.graph_identifier, vpc_id=self.vpc_id)
+
+ endpoint_url = VpcEndpointLink.format_str.format(
+ endpoint_id=endpoint_id,
+ aws_domain=VpcEndpointLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ VpcEndpointLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ endpoint_id=endpoint_id,
+ )
+ self.log.info("You can view this private endpoint at : %s", endpoint_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until endpoint is available")
+ self.defer(
+ trigger=NeptuneGraphPrivateEndpointAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_identifier,
+ vpc_id=self.vpc_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ kwargs={"vpc_id": self.vpc_id},
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until endpoint is available")
+ self.hook.get_waiter("private_graph_endpoint_available").wait(
+ graphIdentifier=self.graph_identifier,
+ vpcId=self.vpc_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"vpc_endpoint_id": endpoint_id, "graph_id": self.graph_identifier, "vpc_id": self.vpc_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ validated_event = validate_execute_complete_event(event)
+
+ if validated_event.get("status") != "success":
+ raise NeptunePrivateEndpointCreationFailedError(
+ validated_event.get("message", "Endpoint failed to create")
+ )
+
+ graph_id = validated_event["value"]
+ vpc_id = validated_event["vpc_id"]
Review Comment:
Where are these keys coming from? I don't see them returned by the NeptuneGraphPrivateEndpointAvailableTrigger?
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Neptune analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on PR #64274:
URL: https://github.com/apache/airflow/pull/64274#issuecomment-4172945177
Odd, all my prek checks passed locally and I rebased before pushing. Will check it out and fix.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3283989151
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/neptune_analytics.py:
##########
@@ -0,0 +1,1006 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING, Any
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.neptune_analytics import NeptuneAnalyticsHook
+from airflow.providers.amazon.aws.links.ec2 import VpcEndpointLink
+from airflow.providers.amazon.aws.links.neptune_analytics import NeptuneGraphLink, NeptuneImportTaskLink
+from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
+from airflow.providers.amazon.aws.triggers.neptune_analytics import (
+ NeptuneGraphAvailableTrigger,
+ NeptuneGraphDeletedTrigger,
+ NeptuneGraphPrivateEndpointAvailableTrigger,
+ NeptuneGraphPrivateEndpointDeletedTrigger,
+ NeptuneImportTaskCancelledTrigger,
+ NeptuneImportTaskCompleteTrigger,
+)
+from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
+from airflow.providers.common.compat.sdk import conf
+
+if TYPE_CHECKING:
+ from airflow.sdk import Context
+
+from airflow.providers.amazon.aws.exceptions import (
+ NeptuneGraphCreationFailedError,
+ NeptuneGraphDeletionFailedError,
+ NeptuneImportTaskCancellationFailedError,
+ NeptunePrivateEndpointCreationFailedError,
+ NeptunePrivateEndpointDeletionFailedError,
+)
+
+
+class NeptuneCreateGraphOperator(AwsBaseOperator[NeptuneAnalyticsHook]):
+ """
+ Creates an empty Amazon Neptune Graph database.
+
+ Neptune Analytics is a memory-optimized graph database engine for analytics. With Neptune Analytics, you can get insights and find trends by processing large amounts of graph data in seconds.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the guide:
+ :ref:`howto/operator:NeptuneCreateGraphOperator`
+
+ :param graph_name: Name of Neptune graph to create
+ :param vector_search_config: Specifies the number of dimensions for vector embeddings that will be loaded into the graph.
+ :param provisioned_memory: The provisioned memory-optimized Neptune Capacity Units (m-NCUs) to use for the graph.
+ :param public_connectivity: Specifies whether or not the graph can be reachable over the internet.
+ :param replica_count: The number of replicas in other AZs.
+ :param deletion_protection: Indicates whether or not to enable deletion protection on the graph.
+ The graph can't be deleted when deletion protection is enabled.
+ :param kms_key_id: Specifies a KMS key to use to encrypt data in the new graph.
+ :param tags: Specifies metadata tags to add to the graph.
+ :param wait_for_completion: Whether to wait for the graph to start. (default: True)
+ :param deferrable: If True, the operator will wait asynchronously for the graph to start.
+ This implies waiting for completion. This mode requires aiobotocore module to be installed.
+ (default: False)
+ :param waiter_delay: Time in seconds to wait between status checks.
+ :param waiter_max_attempts: Maximum number of attempts to check for job completion.
+ :param aws_conn_id: The Airflow connection used for AWS credentials.
+ If this is ``None`` or empty then the default boto3 behaviour is used. If
+ running Airflow in a distributed manner and aws_conn_id is None or
+ empty, then default boto3 configuration would be used (and must be
+ maintained on each worker node).
+ :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used.
+ :param botocore_config: Configuration dictionary (key-values) for botocore client. See:
+ https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
+ :return: dictionary with Neptune graph id
+ """
+
+ aws_hook_class = NeptuneAnalyticsHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "graph_name", "vector_search_config", "provisioned_memory"
+ )
+
+ template_fields_renderers = {
+ "vector_search_config": "json",
+ }
+
+ operator_extra_links = (NeptuneGraphLink(),)
+
+ def __init__(
+ self,
+ graph_name: str,
+ vector_search_config: dict,
+ provisioned_memory: int,
+ public_connectivity: bool | None = None,
+ replica_count: int | None = None,
+ deletion_protection: bool = False,
+ kms_key_id: str | None = None,
+ tags: dict | None = None,
+ wait_for_completion: bool = True,
+ waiter_delay: int = 30,
+ waiter_max_attempts: int = 60,
+ deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.graph_name = graph_name
+ self.vector_search_config = vector_search_config
+ self.replica_count = replica_count
+ self.provisioned_memory = provisioned_memory
+ self.public_connectivity = public_connectivity
+ self.deletion_protect = deletion_protection
+ self.kms_key = kms_key_id
+ self.tags = tags
+ self.wait_for_completion = wait_for_completion
+ self.deferrable = deferrable
+ self.waiter_delay = waiter_delay
+ self.waiter_max_attempts = waiter_max_attempts
+
+ def execute(self, context: Context) -> dict:
+ self.log.info("Creating graph %s", self.graph_name)
+
+ create_params = {
+ "graphName": self.graph_name,
+ "vectorSearchConfiguration": self.vector_search_config,
+ "provisionedMemory": self.provisioned_memory,
+ **{
+ k: v
+ for k, v in {
+ "replicaCount": self.replica_count,
+ "publicConnectivity": self.public_connectivity,
+ "deletionProtection": self.deletion_protect,
+ "kmsKeyIdentifier": self.kms_key,
+ "tags": self.tags,
+ }.items()
+ if v is not None
+ },
+ }
+
+ response = self.hook.conn.create_graph(**create_params)
+
+ self.log.info("Graph %s in status %s", self.graph_name, response.get("status", "Unknown"))
+ self.graph_id = response.get("id", None)
+
+ graph_url = NeptuneGraphLink.format_str.format(
+ graph_id=self.graph_id,
+ aws_domain=NeptuneGraphLink.get_aws_domain(self.hook.conn_partition),
+ region_name=self.hook.conn_region_name,
+ )
+
+ NeptuneGraphLink.persist(
+ context=context,
+ operator=self,
+ region_name=self.hook.conn_region_name,
+ aws_partition=self.hook.conn_partition,
+ graph_id=self.graph_id,
+ )
+ self.log.info("You can view this Neptune Graph at : %s", graph_url)
+
+ if self.deferrable:
+ self.log.info("Deferring until graph %s is available", self.graph_id)
+ self.defer(
+ trigger=NeptuneGraphAvailableTrigger(
+ aws_conn_id=self.aws_conn_id,
+ graph_id=self.graph_id,
+ waiter_delay=self.waiter_delay,
+ waiter_max_attempts=self.waiter_max_attempts,
+ ),
+ method_name="execute_complete",
+ )
+
+ if self.wait_for_completion:
+ self.log.info("Waiting until graph %s is available", self.graph_id)
+ self.hook.get_waiter("graph_available").wait(
+ graphIdentifier=self.graph_id,
+ WaiterConfig={"Delay": self.waiter_delay, "MaxAttempts": self.waiter_max_attempts},
+ )
+
+ return {"graph_id": self.graph_id}
+
+ def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]:
+ if event is None:
+ raise NeptuneGraphCreationFailedError(
+ "No event received while waiting for Neptune graph creation to complete."
+ )
+
+ if event.get("status") != "success":
+ raise NeptuneGraphCreationFailedError(
+ event.get(
+ "message",
+ f"Neptune graph {self.graph_id} creation did not complete successfully: {event}",
+ )
+ )
+
+ self.log.info("Neptune graph %s complete", self.graph_id)
Review Comment:
Ah, didn't realize there is a helper function. Will implement.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "ellisms (via GitHub)" <gi...@apache.org>.
ellisms commented on code in PR #64274:
URL: https://github.com/apache/airflow/pull/64274#discussion_r3283791944
##########
airflow-core/src/airflow/migrations/versions/0113_3_3_0_add_retry_policy_fields_to_ti.py:
##########
@@ -30,7 +30,7 @@
is a metadata-only operation (no table rewrite).
Revision ID: b8f3e4a1d2c9
-Revises: 9fabad868fdb
+Revises: fde9ed84d07b
Create Date: 2026-04-16 12:00:00.000000
Review Comment:
No, I don't know what this is.
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add Amazon Provider Neptune Analytics [airflow]
Posted by "o-nikolas (via GitHub)" <gi...@apache.org>.
o-nikolas merged PR #64274:
URL: https://github.com/apache/airflow/pull/64274
--
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-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org