#
# 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 S3 Tables operators."""
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Any, Literal
from botocore.exceptions import ClientError
from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
from airflow.providers.amazon.aws.hooks.s3_tables import S3TablesHook
from airflow.providers.amazon.aws.operators.base_aws import AwsBaseOperator
from airflow.providers.amazon.aws.utils.mixins import aws_template_fields
from airflow.utils.helpers import prune_dict
if TYPE_CHECKING:
from airflow.sdk import Context
[docs]
class S3TablesCreateTableOperator(AwsBaseOperator[AwsBaseHook]):
"""
Create a new table in an Amazon S3 Tables namespace.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:S3TablesCreateTableOperator`
:param table_bucket_arn: The ARN of the table bucket to create the table in. (templated)
:param namespace: The namespace to associate with the table. (templated)
:param table_name: The name of the table. (templated)
:param format: The table format. (templated) Currently only ``ICEBERG`` is supported.
:param metadata: Optional Iceberg schema metadata. (templated)
Example: ``{"iceberg": {"schema": {"fields": [{"name": "id", "type": "int", "required": True}]}}}``
: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 verify: Whether or not to verify SSL certificates. See:
https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html
:param botocore_config: Configuration dictionary (key-values) for botocore client. See:
https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html
"""
[docs]
template_fields: Sequence[str] = aws_template_fields(
"table_bucket_arn", "namespace", "table_name", "format", "metadata"
)
[docs]
template_fields_renderers = {"metadata": "json"}
[docs]
aws_hook_class = AwsBaseHook
def __init__(
self,
*,
table_bucket_arn: str,
namespace: str,
table_name: str,
format: str = "ICEBERG",
metadata: dict[str, Any] | None = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
[docs]
self.table_bucket_arn = table_bucket_arn
[docs]
self.namespace = namespace
[docs]
self.table_name = table_name
@property
def _hook_parameters(self):
return {**super()._hook_parameters, "client_type": "s3tables"}
[docs]
def execute(self, context: Context) -> str:
self.log.info(
"Creating S3 table %s in namespace %s (bucket %s)",
self.table_name,
self.namespace,
self.table_bucket_arn,
)
kwargs: dict[str, Any] = {
"tableBucketARN": self.table_bucket_arn,
"namespace": self.namespace,
"name": self.table_name,
"format": self.format,
}
if self.metadata:
kwargs["metadata"] = self.metadata
response = self.hook.conn.create_table(**kwargs)
table_arn = response["tableARN"]
self.log.info("Created table: %s", table_arn)
return table_arn
[docs]
class S3TablesDeleteTableOperator(AwsBaseOperator[AwsBaseHook]):
"""
Delete a table from an Amazon S3 Tables namespace.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:S3TablesDeleteTableOperator`
:param table_bucket_arn: The ARN of the table bucket containing the table. (templated)
:param namespace: The namespace of the table. (templated)
:param table_name: The name of the table to delete. (templated)
:param version_token: Optional version token for optimistic concurrency. (templated)
"""
[docs]
template_fields: Sequence[str] = aws_template_fields(
"table_bucket_arn", "namespace", "table_name", "version_token"
)
[docs]
aws_hook_class = AwsBaseHook
def __init__(
self,
*,
table_bucket_arn: str,
namespace: str,
table_name: str,
version_token: str | None = None,
**kwargs,
) -> None:
super().__init__(**kwargs)
[docs]
self.table_bucket_arn = table_bucket_arn
[docs]
self.namespace = namespace
[docs]
self.table_name = table_name
[docs]
self.version_token = version_token
@property
def _hook_parameters(self):
return {**super()._hook_parameters, "client_type": "s3tables"}
[docs]
def execute(self, context: Context) -> None:
self.log.info(
"Deleting S3 table %s from namespace %s (bucket %s)",
self.table_name,
self.namespace,
self.table_bucket_arn,
)
kwargs: dict[str, Any] = prune_dict(
{
"tableBucketARN": self.table_bucket_arn,
"namespace": self.namespace,
"name": self.table_name,
"versionToken": self.version_token,
}
)
self.hook.conn.delete_table(**kwargs)
self.log.info("Deleted table %s", self.table_name)
[docs]
class S3TablesDeleteNamespaceOperator(AwsBaseOperator[S3TablesHook]):
"""
Delete a namespace from an Amazon S3 Tables table bucket.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:S3TablesDeleteNamespaceOperator`
:param table_bucket_arn: The ARN of the table bucket. (templated)
:param namespace: The namespace to delete. (templated)
"""
[docs]
template_fields: Sequence[str] = aws_template_fields("table_bucket_arn", "namespace")
[docs]
aws_hook_class = S3TablesHook
def __init__(
self,
*,
table_bucket_arn: str,
namespace: str,
**kwargs,
) -> None:
super().__init__(**kwargs)
[docs]
self.table_bucket_arn = table_bucket_arn
[docs]
self.namespace = namespace
[docs]
def execute(self, context: Context) -> None:
self.log.info("Deleting namespace %s from %s", self.namespace, self.table_bucket_arn)
self.hook.conn.delete_namespace(tableBucketARN=self.table_bucket_arn, namespace=self.namespace)
self.log.info("Deleted namespace %s", self.namespace)
[docs]
class S3TablesCreateTableBucketOperator(AwsBaseOperator[S3TablesHook]):
"""
Create an Amazon S3 Tables table bucket.
A table bucket is the top-level container for S3 Tables namespaces and tables.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:S3TablesCreateTableBucketOperator`
:param table_bucket_name: The name of the table bucket. (templated)
:param encryption_configuration: Optional encryption configuration dict with
``sseAlgorithm`` and optional ``kmsKeyArn``. (templated)
:param if_exists: Behavior when a table bucket with the same name already exists.
``"fail"`` raises an error, ``"skip"`` returns the existing bucket ARN.
"""
[docs]
template_fields: Sequence[str] = aws_template_fields("table_bucket_name")
[docs]
template_fields_renderers = {"encryption_configuration": "json"}
[docs]
aws_hook_class = S3TablesHook
def __init__(
self,
*,
table_bucket_name: str,
encryption_configuration: dict[str, str] | None = None,
tags: dict[str, str] | None = None,
if_exists: Literal["fail", "skip"] = "skip",
**kwargs,
) -> None:
super().__init__(**kwargs)
[docs]
self.table_bucket_name = table_bucket_name
[docs]
self.encryption_configuration = encryption_configuration
[docs]
self.if_exists = if_exists
[docs]
def execute(self, context: Context) -> str:
self.log.info("Creating S3 Tables table bucket %s", self.table_bucket_name)
kwargs: dict[str, Any] = prune_dict(
{
"name": self.table_bucket_name,
"encryptionConfiguration": self.encryption_configuration,
"tags": self.tags,
}
)
try:
response = self.hook.conn.create_table_bucket(**kwargs)
bucket_arn = response["arn"]
except ClientError as e:
if e.response["Error"]["Code"] == "ConflictException" and self.if_exists == "skip":
self.log.info("Table bucket %s already exists, skipping.", self.table_bucket_name)
bucket_arn = self.hook.get_table_bucket_arn_by_name(self.table_bucket_name)
if bucket_arn is None:
raise
else:
raise
self.log.info("Table bucket %s: %s", self.table_bucket_name, bucket_arn)
return bucket_arn
[docs]
class S3TablesDeleteTableBucketOperator(AwsBaseOperator[S3TablesHook]):
"""
Delete an Amazon S3 Tables table bucket.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:S3TablesDeleteTableBucketOperator`
:param table_bucket_arn: The ARN of the table bucket to delete. (templated)
"""
[docs]
template_fields: Sequence[str] = aws_template_fields("table_bucket_arn")
[docs]
aws_hook_class = S3TablesHook
def __init__(
self,
*,
table_bucket_arn: str,
**kwargs,
) -> None:
super().__init__(**kwargs)
[docs]
self.table_bucket_arn = table_bucket_arn
[docs]
def execute(self, context: Context) -> None:
self.log.info("Deleting S3 Tables table bucket %s", self.table_bucket_arn)
self.hook.conn.delete_table_bucket(tableBucketARN=self.table_bucket_arn)
self.log.info("Deleted table bucket %s", self.table_bucket_arn)
[docs]
class S3TablesCreateNamespaceOperator(AwsBaseOperator[S3TablesHook]):
"""
Create a namespace in an Amazon S3 Tables table bucket.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:S3TablesCreateNamespaceOperator`
:param table_bucket_arn: The ARN of the table bucket. (templated)
:param namespace: The namespace name to create. (templated)
:param if_exists: Behavior when the namespace already exists.
``"fail"`` raises an error, ``"skip"`` logs and returns.
"""
[docs]
template_fields: Sequence[str] = aws_template_fields("table_bucket_arn", "namespace")
[docs]
aws_hook_class = S3TablesHook
def __init__(
self,
*,
table_bucket_arn: str,
namespace: str,
if_exists: Literal["fail", "skip"] = "skip",
**kwargs,
) -> None:
super().__init__(**kwargs)
[docs]
self.table_bucket_arn = table_bucket_arn
[docs]
self.namespace = namespace
[docs]
self.if_exists = if_exists
[docs]
def execute(self, context: Context) -> str:
self.log.info("Creating namespace %s in %s", self.namespace, self.table_bucket_arn)
try:
self.hook.conn.create_namespace(tableBucketARN=self.table_bucket_arn, namespace=[self.namespace])
except ClientError as e:
if e.response["Error"]["Code"] == "ConflictException" and self.if_exists == "skip":
self.log.info("Namespace %s already exists, skipping.", self.namespace)
else:
raise
self.log.info("Namespace %s created.", self.namespace)
return self.namespace