"""Amazon Timestream Module.""" from __future__ import annotations import itertools import logging import time from datetime import datetime from typing import TYPE_CHECKING, Any, Dict, List, Literal, cast import boto3 from botocore.config import Config import awswrangler.pandas as pd from awswrangler import _data_types, _utils, exceptions, s3 from awswrangler._config import apply_configs from awswrangler._distributed import engine from awswrangler._executor import _BaseExecutor, _get_executor from awswrangler.distributed.ray import ray_get from awswrangler.typing import TimestreamBatchLoadReportS3Configuration if TYPE_CHECKING: from mypy_boto3_timestream_write.client import TimestreamWriteClient _BATCH_LOAD_FINAL_STATES: list[str] = ["SUCCEEDED", "FAILED", "PROGRESS_STOPPED", "PENDING_RESUME"] _BATCH_LOAD_WAIT_POLLING_DELAY: float = 2 # SECONDS _TIME_UNITS_MAPPING = { "SECONDS": (9, 0), "MILLISECONDS": (6, 3), "MICROSECONDS": (3, 6), "NANOSECONDS": (0, 9), } _logger: logging.Logger = logging.getLogger(__name__) _TimeUnitLiteral = Literal["MILLISECONDS", "SECONDS", "MICROSECONDS", "NANOSECONDS"] def _df2list(df: pd.DataFrame) -> list[list[Any]]: """Extract Parameters.""" parameters: list[list[Any]] = df.values.tolist() for i, row in enumerate(parameters): for j, value in enumerate(row): if pd.isna(value): parameters[i][j] = None elif hasattr(value, "to_pydatetime"): parameters[i][j] = value.to_pydatetime() return parameters def _check_time_unit(time_unit: _TimeUnitLiteral) -> str: time_unit = time_unit if time_unit else "MILLISECONDS" if time_unit not in _TIME_UNITS_MAPPING.keys(): raise exceptions.InvalidArgumentValue( f"Invalid time unit: {time_unit}. Must be one of {_TIME_UNITS_MAPPING.keys()}." ) return time_unit def _format_timestamp(timestamp: int | datetime, time_unit: _TimeUnitLiteral) -> str: if isinstance(timestamp, int): return str(round(timestamp / pow(10, _TIME_UNITS_MAPPING[time_unit][0]))) if isinstance(timestamp, datetime): return str(round(timestamp.timestamp() * pow(10, _TIME_UNITS_MAPPING[time_unit][1]))) raise exceptions.InvalidArgumentType("`time_col` must be of type timestamp.") def _format_measure( measure_name: str, measure_value: Any, measure_type: str, time_unit: _TimeUnitLiteral ) -> dict[str, str]: return { "Name": measure_name, "Value": _format_timestamp(measure_value, time_unit) if measure_type == "TIMESTAMP" else str(measure_value), "Type": measure_type, } def _sanitize_common_attributes( common_attributes: dict[str, Any] | None, version: int, time_unit: _TimeUnitLiteral, measure_name: str | None, ) -> dict[str, Any]: common_attributes = {} if not common_attributes else common_attributes # Values in common_attributes take precedence common_attributes.setdefault("Version", version) common_attributes.setdefault("TimeUnit", _check_time_unit(common_attributes.get("TimeUnit", time_unit))) if "Time" not in common_attributes and common_attributes["TimeUnit"] == "NANOSECONDS": raise exceptions.InvalidArgumentValue("Python datetime objects do not support nanoseconds precision.") if "MeasureValue" in common_attributes and "MeasureValueType" not in common_attributes: raise exceptions.InvalidArgumentCombination( "MeasureValueType must be supplied alongside MeasureValue in common_attributes." ) if measure_name: common_attributes.setdefault("MeasureName", measure_name) elif "MeasureName" not in common_attributes: raise exceptions.InvalidArgumentCombination( "MeasureName must be supplied with the `measure_name` argument or in common_attributes." ) return common_attributes @engine.dispatch_on_engine def _write_batch( timestream_client: "TimestreamWriteClient" | None, database: str, table: str, common_attributes: dict[str, Any], cols_names: list[str | None], measure_cols: list[str | None], measure_types: list[str], dimensions_cols: list[str | None], batch: list[Any], ) -> list[dict[str, str]]: client_timestream = timestream_client if timestream_client else _utils.client(service_name="timestream-write") records: list[dict[str, Any]] = [] scalar = bool(len(measure_cols) == 1 and "MeasureValues" not in common_attributes) time_loc = 0 measure_cols_loc = 1 if cols_names[0] else 0 dimensions_cols_loc = 1 if len(measure_cols) == 1 else 1 + len(measure_cols) if all(cols_names): # Time and Measures are supplied in the data frame dimensions_cols_loc = 1 + len(measure_cols) elif all(v is None for v in cols_names[:2]): # Time and Measures are supplied in common_attributes dimensions_cols_loc = 0 time_unit = common_attributes["TimeUnit"] for row in batch: record: dict[str, Any] = {} if "Time" not in common_attributes: record["Time"] = _format_timestamp(row[time_loc], time_unit) if scalar and "MeasureValue" not in common_attributes: measure_value = row[measure_cols_loc] if pd.isnull(measure_value): continue record["MeasureValue"] = str(measure_value) elif not scalar and "MeasureValues" not in common_attributes: record["MeasureValues"] = [ _format_measure(measure_name, measure_value, measure_value_type, time_unit) # type: ignore[arg-type] for measure_name, measure_value, measure_value_type in zip( measure_cols, row[measure_cols_loc:dimensions_cols_loc], measure_types ) if not pd.isnull(measure_value) ] if len(record["MeasureValues"]) == 0: continue if "MeasureValueType" not in common_attributes: record["MeasureValueType"] = measure_types[0] if scalar else "MULTI" # Dimensions can be specified in both common_attributes and the data frame dimensions = ( [ {"Name": name, "DimensionValueType": "VARCHAR", "Value": str(value)} for name, value in zip(dimensions_cols, row[dimensions_cols_loc:]) ] if all(dimensions_cols) else [] ) if dimensions: record["Dimensions"] = dimensions if record: records.append(record) try: if records: _utils.try_it( f=client_timestream.write_records, ex=( client_timestream.exceptions.ThrottlingException, client_timestream.exceptions.InternalServerException, ), max_num_tries=5, DatabaseName=database, TableName=table, CommonAttributes=common_attributes, Records=records, ) except client_timestream.exceptions.RejectedRecordsException as ex: return cast(List[Dict[str, str]], ex.response["RejectedRecords"]) # type: ignore[typeddict-item] return [] @engine.dispatch_on_engine def _write_df( df: pd.DataFrame, executor: _BaseExecutor, database: str, table: str, common_attributes: dict[str, Any], cols_names: list[str | None], measure_cols: list[str | None], measure_types: list[str], dimensions_cols: list[str | None], boto3_session: boto3.Session | None, ) -> list[dict[str, str]]: timestream_client = _utils.client( service_name="timestream-write", session=boto3_session, botocore_config=Config(read_timeout=20, max_pool_connections=5000, retries={"max_attempts": 10}), ) batches: list[list[Any]] = _utils.chunkify(lst=_df2list(df=df), max_length=100) _logger.debug("Writing %d batches of data", len(batches)) return executor.map( _write_batch, # type: ignore[arg-type] timestream_client, itertools.repeat(database), itertools.repeat(table), itertools.repeat(common_attributes), itertools.repeat(cols_names), itertools.repeat(measure_cols), itertools.repeat(measure_types), itertools.repeat(dimensions_cols), batches, ) @_utils.validate_distributed_kwargs( unsupported_kwargs=["boto3_session"], ) def write( df: pd.DataFrame, database: str, table: str, time_col: str | None = None, measure_col: str | list[str | None] | None = None, dimensions_cols: list[str | None] | None = None, version: int = 1, time_unit: _TimeUnitLiteral = "MILLISECONDS", use_threads: bool | int = True, measure_name: str | None = None, common_attributes: dict[str, Any] | None = None, boto3_session: boto3.Session | None = None, ) -> list[dict[str, str]]: """Store a Pandas DataFrame into an Amazon Timestream table. Note ---- In case `use_threads=True`, the number of threads from os.cpu_count() is used. If the Timestream service rejects a record(s), this function will not throw a Python exception. Instead it will return the rejection information. Note ---- If ``time_col`` column is supplied, it must be of type timestamp. ``time_unit`` is set to MILLISECONDS by default. NANOSECONDS is not supported as python datetime objects are limited to microseconds precision. Parameters ---------- df Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html database Amazon Timestream database name. table Amazon Timestream table name. time_col DataFrame column name to be used as time. MUST be a timestamp column. measure_col DataFrame column name(s) to be used as measure. dimensions_cols List of DataFrame column names to be used as dimensions. version Version number used for upserts. Documentation https://docs.aws.amazon.com/timestream/latest/developerguide/API_WriteRecords.html. time_unit Time unit for the time column. MILLISECONDS by default. use_threads True to enable concurrent writing, False to disable multiple threads. If enabled, os.cpu_count() is used as the number of threads. If integer is provided, specified number is used. measure_name Name that represents the data attribute of the time series. Overrides ``measure_col`` if specified. common_attributes Dictionary of attributes shared across all records in the request. Using common attributes can optimize the cost of writes by reducing the size of request payloads. Values in ``common_attributes`` take precedence over all other arguments and data frame values. Dimension attributes are merged with attributes in record objects. Example: ``{"Dimensions": [{"Name": "device_id", "Value": "12345"}], "MeasureValueType": "DOUBLE"}``. boto3_session The default boto3 session will be used if **boto3_session** is ``None``. Returns ------- Rejected records. Possible reasons for rejection are described here: https://docs.aws.amazon.com/timestream/latest/developerguide/API_RejectedRecord.html Examples -------- Store a Pandas DataFrame into a Amazon Timestream table. >>> import awswrangler as wr >>> import pandas as pd >>> df = pd.DataFrame( >>> { >>> "time": [datetime.now(), datetime.now(), datetime.now()], >>> "dim0": ["foo", "boo", "bar"], >>> "dim1": [1, 2, 3], >>> "measure": [1.0, 1.1, 1.2], >>> } >>> ) >>> rejected_records = wr.timestream.write( >>> df=df, >>> database="sampleDB", >>> table="sampleTable", >>> time_col="time", >>> measure_col="measure", >>> dimensions_cols=["dim0", "dim1"], >>> ) >>> assert len(rejected_records) == 0 Return value if some records are rejected. >>> [ >>> { >>> 'ExistingVersion': 2, >>> 'Reason': 'The record version 1 is lower than the existing version 2. A ' >>> 'higher version is required to update the measure value.', >>> 'RecordIndex': 0 >>> } >>> ] """ measure_cols = measure_col if isinstance(measure_col, list) else [measure_col] measure_types: list[str] = ( _data_types.timestream_type_from_pandas(df.loc[:, measure_cols]) if all(measure_cols) else [] ) dimensions_cols = dimensions_cols if dimensions_cols else [dimensions_cols] # type: ignore[list-item] cols_names: list[str | None] = [time_col] + measure_cols + dimensions_cols measure_name = measure_name if measure_name else measure_cols[0] common_attributes = _sanitize_common_attributes(common_attributes, version, time_unit, measure_name) _logger.debug( "Writing to Timestream table %s in database %s\ncommon_attributes: %s\n, cols_names: %s\n, measure_types: %s", table, database, common_attributes, cols_names, measure_types, ) # User can supply arguments in one of two ways: # 1. With the `common_attributes` dictionary which takes precedence # 2. With data frame columns # However, the data frame cannot be completely empty. # So if all values in `cols_names` are None, an exception is raised. if any(cols_names): dfs = _utils.split_pandas_frame( df.loc[:, [c for c in cols_names if c]], _utils.ensure_cpu_count(use_threads=use_threads) ) else: raise exceptions.InvalidArgumentCombination( "At least one of `time_col`, `measure_col` or `dimensions_cols` must be specified." ) _logger.debug("Writing %d dataframes to Timestream table", len(dfs)) executor: _BaseExecutor = _get_executor(use_threads=use_threads) errors = list( itertools.chain( *ray_get( [ _write_df( df=df, executor=executor, database=database, table=table, common_attributes=common_attributes, cols_names=cols_names, measure_cols=measure_cols, measure_types=measure_types, dimensions_cols=dimensions_cols, boto3_session=boto3_session, ) for df in dfs ] ) ) ) return list(itertools.chain(*ray_get(errors))) @apply_configs def wait_batch_load_task( task_id: str, timestream_batch_load_wait_polling_delay: float = _BATCH_LOAD_WAIT_POLLING_DELAY, boto3_session: boto3.Session | None = None, ) -> dict[str, Any]: """ Wait for the Timestream batch load task to complete. Parameters ---------- task_id The ID of the batch load task. timestream_batch_load_wait_polling_delay Time to wait between two polling attempts. boto3_session The default boto3 session will be used if **boto3_session** is ``None``. Returns ------- Dictionary with the describe_batch_load_task response. Examples -------- >>> import awswrangler as wr >>> res = wr.timestream.wait_batch_load_task(task_id='task-id') Raises ------ exceptions.TimestreamLoadError Error message raised by failed task. """ timestream_client = _utils.client(service_name="timestream-write", session=boto3_session) response = timestream_client.describe_batch_load_task(TaskId=task_id) status = response["BatchLoadTaskDescription"]["TaskStatus"] while status not in _BATCH_LOAD_FINAL_STATES: time.sleep(timestream_batch_load_wait_polling_delay) response = timestream_client.describe_batch_load_task(TaskId=task_id) status = response["BatchLoadTaskDescription"]["TaskStatus"] _logger.debug("Task status: %s", status) if status != "SUCCEEDED": _logger.debug("Task response: %s", response) raise exceptions.TimestreamLoadError(response.get("ErrorMessage")) return response # type: ignore[return-value] @apply_configs @_utils.validate_distributed_kwargs( unsupported_kwargs=["boto3_session", "s3_additional_kwargs"], ) def batch_load( df: pd.DataFrame, path: str, database: str, table: str, time_col: str, dimensions_cols: list[str], measure_cols: list[str], measure_name_col: str, report_s3_configuration: TimestreamBatchLoadReportS3Configuration, time_unit: _TimeUnitLiteral = "MILLISECONDS", record_version: int = 1, timestream_batch_load_wait_polling_delay: float = _BATCH_LOAD_WAIT_POLLING_DELAY, keep_files: bool = False, use_threads: bool | int = True, boto3_session: boto3.Session | None = None, s3_additional_kwargs: dict[str, str] | None = None, ) -> dict[str, Any]: """Batch load a Pandas DataFrame into a Amazon Timestream table. Note ---- The supplied column names (time, dimension, measure) MUST match those in the Timestream table. Note ---- Only ``MultiMeasureMappings`` is supported. See https://docs.aws.amazon.com/timestream/latest/developerguide/batch-load-data-model-mappings.html Parameters ---------- df Pandas DataFrame. path S3 prefix to write the data. database Amazon Timestream database name. table Amazon Timestream table name. time_col Column name with the time data. It must be a long data type that represents the time since the Unix epoch. dimensions_cols List of column names with the dimensions data. measure_cols List of column names with the measure data. measure_name_col Column name with the measure name. report_s3_configuration Dictionary of the configuration for the S3 bucket where the error report is stored. https://docs.aws.amazon.com/timestream/latest/developerguide/API_ReportS3Configuration.html Example: {"BucketName": 'error-report-bucket-name'} time_unit Time unit for the time column. MILLISECONDS by default. record_version Record version. timestream_batch_load_wait_polling_delay Time to wait between two polling attempts. keep_files Whether to keep the files after the operation. use_threads True to enable concurrent requests, False to disable multiple threads. boto3_session The default boto3 session will be used if **boto3_session** is ``None``. s3_additional_kwargs Forwarded to S3 botocore requests. Returns ------- A dictionary of the batch load task response. Examples -------- >>> import awswrangler as wr >>> response = wr.timestream.batch_load( >>> df=df, >>> path='s3://bucket/path/', >>> database='sample_db', >>> table='sample_table', >>> time_col='time', >>> dimensions_cols=['region', 'location'], >>> measure_cols=['memory_utilization', 'cpu_utilization'], >>> report_s3_configuration={'BucketName': 'error-report-bucket-name'}, >>> ) """ path = path if path.endswith("/") else f"{path}/" if s3.list_objects(path=path, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs): raise exceptions.InvalidArgument( f"The received S3 path ({path}) is not empty. " "Please, provide a different path or use wr.s3.delete_objects() to clean up the current one." ) columns = [time_col, *dimensions_cols, *measure_cols, measure_name_col] try: s3.to_csv( df=df.loc[:, columns], path=path, index=False, dataset=True, mode="append", use_threads=use_threads, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, ) measure_types: list[str] = _data_types.timestream_type_from_pandas(df.loc[:, measure_cols]) return batch_load_from_files( path=path, database=database, table=table, time_col=time_col, dimensions_cols=dimensions_cols, measure_cols=measure_cols, measure_types=measure_types, report_s3_configuration=report_s3_configuration, time_unit=time_unit, measure_name_col=measure_name_col, record_version=record_version, timestream_batch_load_wait_polling_delay=timestream_batch_load_wait_polling_delay, boto3_session=boto3_session, ) finally: if not keep_files: _logger.debug("Deleting objects in S3 path: %s", path) s3.delete_objects( path=path, use_threads=use_threads, boto3_session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, ) @apply_configs def batch_load_from_files( path: str, database: str, table: str, time_col: str, dimensions_cols: list[str], measure_cols: list[str], measure_types: list[str], measure_name_col: str, report_s3_configuration: TimestreamBatchLoadReportS3Configuration, time_unit: _TimeUnitLiteral = "MILLISECONDS", record_version: int = 1, data_source_csv_configuration: dict[str, str | bool] | None = None, timestream_batch_load_wait_polling_delay: float = _BATCH_LOAD_WAIT_POLLING_DELAY, boto3_session: boto3.Session | None = None, ) -> dict[str, Any]: """Batch load files from S3 into a Amazon Timestream table. Note ---- The supplied column names (time, dimension, measure) MUST match those in the Timestream table. Note ---- Only ``MultiMeasureMappings`` is supported. See https://docs.aws.amazon.com/timestream/latest/developerguide/batch-load-data-model-mappings.html Parameters ---------- path S3 prefix to write the data. database Amazon Timestream database name. table Amazon Timestream table name. time_col Column name with the time data. It must be a long data type that represents the time since the Unix epoch. dimensions_cols List of column names with the dimensions data. measure_cols List of column names with the measure data. measure_name_col Column name with the measure name. report_s3_configuration Dictionary of the configuration for the S3 bucket where the error report is stored. https://docs.aws.amazon.com/timestream/latest/developerguide/API_ReportS3Configuration.html Example: {"BucketName": 'error-report-bucket-name'} time_unit Time unit for the time column. MILLISECONDS by default. record_version Record version. data_source_csv_configuration Dictionary of the data source CSV configuration. https://docs.aws.amazon.com/timestream/latest/developerguide/API_CsvConfiguration.html timestream_batch_load_wait_polling_delay Time to wait between two polling attempts. boto3_session The default boto3 session will be used if **boto3_session** is ``None``. Returns ------- A dictionary of the batch load task response. Examples -------- >>> import awswrangler as wr >>> response = wr.timestream.batch_load_from_files( >>> path='s3://bucket/path/', >>> database='sample_db', >>> table='sample_table', >>> time_col='time', >>> dimensions_cols=['region', 'location'], >>> measure_cols=['memory_utilization', 'cpu_utilization'], >>> report_s3_configuration={'BucketName': 'error-report-bucket-name'}, >>> ) """ timestream_client = _utils.client(service_name="timestream-write", session=boto3_session) bucket, prefix = _utils.parse_path(path=path) kwargs: dict[str, Any] = { "TargetDatabaseName": database, "TargetTableName": table, "DataModelConfiguration": { "DataModel": { "TimeColumn": time_col, "TimeUnit": _check_time_unit(time_unit), "DimensionMappings": [{"SourceColumn": c} for c in dimensions_cols], "MeasureNameColumn": measure_name_col, "MultiMeasureMappings": { "MultiMeasureAttributeMappings": [ {"SourceColumn": c, "MeasureValueType": t} for c, t in zip(measure_cols, measure_types) ], }, } }, "DataSourceConfiguration": { "DataSourceS3Configuration": {"BucketName": bucket, "ObjectKeyPrefix": prefix}, "DataFormat": "CSV", "CsvConfiguration": data_source_csv_configuration if data_source_csv_configuration else {}, }, "ReportConfiguration": {"ReportS3Configuration": report_s3_configuration}, "RecordVersion": record_version, } task_id = timestream_client.create_batch_load_task(**kwargs)["TaskId"] return wait_batch_load_task( task_id=task_id, timestream_batch_load_wait_polling_delay=timestream_batch_load_wait_polling_delay, boto3_session=boto3_session, )