meerschaum 2.3.6__py3-none-any.whl → 2.4.0__py3-none-any.whl

meerschaum/connectors/valkey/_pipes.py

@@ -0,0 +1,844 @@
+#! /usr/bin/env python3
+# vim:fenc=utf-8
+
+"""
+Define pipes methods for `ValkeyConnector`.
+"""
+
+import json
+from datetime import datetime, timezone
+
+import meerschaum as mrsm
+from meerschaum.utils.typing import SuccessTuple, Any, Union, Optional, Dict, List, Tuple
+from meerschaum.utils.misc import json_serialize_datetime, string_to_dict
+from meerschaum.utils.warnings import warn
+from meerschaum.config.static import STATIC_CONFIG
+
+PIPES_TABLE: str = 'mrsm_pipes'
+PIPES_COUNTER: str = 'mrsm_pipes:counter'
+COLON: str = STATIC_CONFIG['valkey']['colon']
+
+
+def get_pipe_key(pipe: mrsm.Pipe) -> str:
+    """
+    Return the key to store a pipe's ID.
+    """
+    return f"mrsm_pipe:{pipe.connector_keys}:{pipe.metric_key}:{pipe.location_key}"
+
+
+def get_pipe_parameters_key(pipe: mrsm.Pipe) -> str:
+    """
+    Return the key to store a pipe's parameters.
+    """
+    return get_pipe_key(pipe) + ':parameters'
+
+
+def serialize_document(doc: Dict[str, Any]) -> str:
+    """
+    Return a serialized string for a document.
+
+    Parameters
+    ----------
+    doc: Dict[str, Any]
+        The document to be serialized.
+
+    Returns
+    -------
+    A serialized string for the document.
+    """
+    return json.dumps(
+        doc,
+        default=(lambda x: json_serialize_datetime(x) if hasattr(x, 'tzinfo') else str(x)),
+        separators=(',', ':'),
+        sort_keys=True,
+    )
+
+
+def get_document_key(
+    doc: Dict[str, Any],
+    indices: List[str],
+    table_name: Optional[str] = None,
+) -> str:
+    """
+    Return a serialized string for a document's indices only.
+
+    Parameters
+    ----------
+    doc: Dict[str, Any]
+        The document containing index values to be serialized.
+
+    indices: List[str]
+        The name of the indices to be serialized.
+
+    table_name: Optional[str], default None
+        If provided, prepend the table to the key.
+
+    Returns
+    -------
+    A serialized string of the document's indices.
+    """
+    from meerschaum.utils.dtypes import coerce_timezone
+    index_vals = {
+        key: (
+            str(val)
+            if not isinstance(val, datetime)
+            else str(int(coerce_timezone(val).replace(tzinfo=timezone.utc).timestamp()))
+        )
+        for key, val in doc.items()
+        if key in indices
+    } if indices else {}
+    indices_str = ((table_name + ':indices:') if table_name else '') + ','.join(
+        sorted(
+            [
+                f'{key}{COLON}{val}'
+                for key, val in index_vals.items()
+            ]
+        )
+    ) if indices else serialize_document(doc)
+    return indices_str
+
+
+def get_table_quoted_doc_key(
+    table_name: str,
+    doc: Dict[str, Any],
+    indices: List[str],
+    datetime_column: Optional[str] = None,
+) -> str:
+    """
+    Return the document string as stored in the underling set.
+    """
+    return json.dumps(
+        {
+            get_document_key(doc, indices, table_name): serialize_document(doc),
+            **(
+                {datetime_column: doc.get(datetime_column, 0)}
+                if datetime_column
+                else {}
+            )
+        },
+        sort_keys=True,
+        separators=(',', ':'),
+        default=(lambda x: json_serialize_datetime(x) if hasattr(x, 'tzinfo') else str(x)),
+    )
+
+
+def register_pipe(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> SuccessTuple:
+    """
+    Insert the pipe's attributes into the internal `pipes` table.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe to be registered.
+
+    Returns
+    -------
+    A `SuccessTuple` of the result.
+    """
+    attributes = {
+        'connector_keys': str(pipe.connector_keys),
+        'metric_key': str(pipe.metric_key),
+        'location_key': str(pipe.location_key),
+    }
+    parameters_str = json.dumps(
+        pipe._attributes.get('parameters', {}),
+        separators=(',', ':'),
+    )
+
+    pipe_key = get_pipe_key(pipe)
+    parameters_key = get_pipe_parameters_key(pipe)
+
+    try:
+        existing_pipe_id = self.get(pipe_key)
+        if existing_pipe_id is not None:
+            return False, f"{pipe} is already registered."
+
+        pipe_id = self.client.incr(PIPES_COUNTER)
+        _ = self.push_docs(
+            [{'pipe_id': pipe_id, **attributes}],
+            PIPES_TABLE,
+            datetime_column='pipe_id',
+            debug=debug,
+        )
+        self.set(pipe_key, pipe_id)
+        self.set(parameters_key, parameters_str)
+
+    except Exception as e:
+        return False, f"Failed to register {pipe}:\n{e}"
+
+    return True, "Success"
+
+
+def get_pipe_id(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> Union[str, int, None]:
+    """
+    Return the `_id` for the pipe if it exists.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe whose `_id` to fetch.
+
+    Returns
+    -------
+    The `_id` for the pipe's document or `None`.
+    """
+    pipe_key = get_pipe_key(pipe)
+    try:
+        return int(self.get(pipe_key))
+    except Exception:
+        pass
+    return None
+
+
+def get_pipe_attributes(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> Dict[str, Any]:
+    """
+    Return the pipe's document from the internal `pipes` collection.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe whose attributes should be retrieved.
+
+    Returns
+    -------
+    The document that matches the keys of the pipe.
+    """
+    pipe_id = pipe.get_id(debug=debug)
+    if pipe_id is None:
+        return {}
+
+    parameters_key = get_pipe_parameters_key(pipe)
+    parameters_str = self.get(parameters_key)
+
+    parameters = json.loads(parameters_str) if parameters_str else {}
+
+    attributes = {
+        'connector_keys': pipe.connector_keys,
+        'metric_key': pipe.metric_key,
+        'location_key': pipe.location_key,
+        'parameters': parameters,
+        'pipe_id': pipe_id,
+    }
+    return attributes
+
+
+def edit_pipe(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> mrsm.SuccessTuple:
+    """
+    Edit the attributes of the pipe.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe whose in-memory parameters must be persisted.
+
+    Returns
+    -------
+    A `SuccessTuple` indicating success.
+    """
+    pipe_id = pipe.get_id(debug=debug)
+    if pipe_id is None:
+        return False, f"{pipe} is not registered."
+
+    parameters_key = get_pipe_parameters_key(pipe)
+    parameters_str = json.dumps(pipe.parameters, separators=(',', ':'))
+    self.set(parameters_key, parameters_str)
+    return True, "Success"
+
+
+def pipe_exists(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> bool:
+    """
+    Check whether a pipe's target table exists.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe to check whether its table exists.
+
+    Returns
+    -------
+    A `bool` indicating the table exists.
+    """
+    table_name = self.quote_table(pipe.target)
+    return self.client.exists(table_name) != 0
+
+
+def drop_pipe(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> mrsm.SuccessTuple:
+    """
+    Drop a pipe's collection if it exists.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe to be dropped.
+
+    Returns
+    -------
+    A `SuccessTuple` indicating success.
+    """
+    try:
+        self.drop_table(pipe.target, debug=debug)
+    except Exception as e:
+        return False, f"Failed to drop {pipe}:\n{e}"
+
+    if 'valkey' not in pipe.parameters:
+        return True, "Success"
+
+    pipe.parameters['valkey']['dtypes'] = {}
+    if not pipe.temporary:
+        edit_success, edit_msg = pipe.edit(debug=debug)
+        if not edit_success:
+            return edit_success, edit_msg
+
+    return True, "Success"
+
+
+def delete_pipe(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> mrsm.SuccessTuple:
+    """
+    Delete a pipe's registration from the `pipes` collection.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe to be deleted.
+
+    Returns
+    -------
+    A `SuccessTuple` indicating success.
+    """
+    drop_success, drop_message = pipe.drop(debug=debug)
+    if not drop_success:
+        return drop_success, drop_message
+
+    pipe_id = self.get_pipe_id(pipe, debug=debug)
+    if pipe_id is None:
+        return False, f"{pipe} is not registered."
+
+    pipe_key = get_pipe_key(pipe)
+    parameters_key = get_pipe_parameters_key(pipe)
+    self.client.delete(pipe_key)
+    self.client.delete(parameters_key)
+    df = self.read(PIPES_TABLE, params={'pipe_id': pipe_id})
+    docs = df.to_dict(orient='records')
+    if docs:
+        doc = docs[0]
+        doc_str = json.dumps(
+            doc,
+            default=(lambda x: json_serialize_datetime(x) if hasattr(x, 'tzinfo') else str(x)),
+            separators=(',', ':'),
+            sort_keys=True,
+        )
+        self.client.zrem(PIPES_TABLE, doc_str)
+    return True, "Success"
+
+
+def get_pipe_data(
+    self,
+    pipe: mrsm.Pipe,
+    select_columns: Optional[List[str]] = None,
+    omit_columns: Optional[List[str]] = None,
+    begin: Union[datetime, int, None] = None,
+    end: Union[datetime, int, None] = None,
+    params: Optional[Dict[str, Any]] = None,
+    debug: bool = False,
+    **kwargs: Any
+) -> Union['pd.DataFrame', None]:
+    """
+    Query a pipe's target table and return the DataFrame.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe with the target table from which to read.
+
+    select_columns: Optional[List[str]], default None
+        If provided, only select these given columns.
+        Otherwise select all available columns (i.e. `SELECT *`).
+
+    omit_columns: Optional[List[str]], default None
+        If provided, remove these columns from the selection.
+
+    begin: Union[datetime, int, None], default None
+        The earliest `datetime` value to search from (inclusive).
+
+    end: Union[datetime, int, None], default None
+        The lastest `datetime` value to search from (exclusive).
+
+    params: Optional[Dict[str, str]], default None
+        Additional filters to apply to the query.
+
+    Returns
+    -------
+    The target table's data as a DataFrame.
+    """
+    if not pipe.exists(debug=debug):
+        return None
+
+    from meerschaum.utils.dataframe import query_df, parse_df_datetimes
+
+    valkey_dtypes = pipe.parameters.get('valkey', {}).get('dtypes', {})
+    dt_col = pipe.columns.get('datetime', None)
+    table_name = self.quote_table(pipe.target)
+    indices = [col for col in pipe.columns.values() if col]
+    ix_docs = [
+        string_to_dict(doc['ix'].replace(COLON, ':'))
+        for doc in self.read_docs(
+            pipe.target,
+            begin=begin,
+            end=end,
+            debug=debug,
+        )
+    ]
+    try:
+        docs_strings = [
+            self.get(get_document_key(
+                doc, indices, table_name
+            ))
+            for doc in ix_docs
+        ]
+    except Exception as e:
+        warn(f"Failed to fetch documents for {pipe}:\n{e}")
+        docs_strings = []
+
+    docs = [
+        json.loads(doc_str)
+        for doc_str in docs_strings
+        if doc_str
+    ]
+    ignore_dt_cols = [
+        col
+        for col, dtype in pipe.dtypes.items()
+        if 'datetime' not in str(dtype)
+    ]
+
+    df = parse_df_datetimes(
+        docs,
+        ignore_cols=ignore_dt_cols,
+        chunksize=kwargs.get('chunksize', None),
+        debug=debug,
+    )
+    for col, typ in valkey_dtypes.items():
+        try:
+            df[col] = df[col].astype(typ)
+        except Exception:
+            pass
+
+    df = pipe.enforce_dtypes(df, debug=debug)
+
+    if len(df) == 0:
+        return query_df(df, select_columns=select_columns, omit_columns=omit_columns)
+
+    return query_df(
+        df,
+        select_columns=select_columns,
+        omit_columns=omit_columns,
+        params=params,
+        begin=begin,
+        end=end,
+        datetime_column=dt_col,
+        inplace=True,
+        reset_index=True,
+    )
+
+
+def sync_pipe(
+    self,
+    pipe: mrsm.Pipe,
+    df: 'pd.DataFrame' = None,
+    check_existing: bool = True,
+    debug: bool = False,
+    **kwargs: Any
+) -> mrsm.SuccessTuple:
+    """
+    Upsert new documents into the pipe's collection.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe whose collection should receive the new documents.
+
+    df: Union['pd.DataFrame', Iterator['pd.DataFrame']], default None
+        The data to be synced.
+
+    check_existing: bool, default True
+        If `False`, do not check the documents against existing data and instead insert directly.
+
+    Returns
+    -------
+    A `SuccessTuple` indicating success.
+    """
+    dt_col = pipe.columns.get('datetime', None)
+    indices = [col for col in pipe.columns.values() if col]
+    table_name = self.quote_table(pipe.target)
+    is_dask = 'dask' in df.__module__
+    if is_dask:
+        df = df.compute()
+    upsert = pipe.parameters.get('upsert', False)
+
+    def _serialize_indices_docs(_docs):
+        return [
+            {
+                'ix': get_document_key(doc, indices),
+                **(
+                    {
+                        dt_col: doc.get(dt_col, 0)
+                    }
+                    if dt_col
+                    else {}
+                )
+            }
+            for doc in _docs
+        ]
+
+    valkey_dtypes = pipe.parameters.get('valkey', {}).get('dtypes', {})
+    new_dtypes = {
+        str(key): str(val)
+        for key, val in df.dtypes.items()
+        if str(key) not in valkey_dtypes
+    }
+    for col, typ in {c: v for c, v in valkey_dtypes.items()}.items():
+        if col in df.columns:
+            try:
+                df[col] = df[col].astype(typ)
+            except Exception:
+                valkey_dtypes[col] = 'string'
+                new_dtypes[col] = 'string'
+                df[col] = df[col].astype('string')
+
+    if new_dtypes:
+        valkey_dtypes.update(new_dtypes)
+        if 'valkey' not in pipe.parameters:
+            pipe.parameters['valkey'] = {}
+        pipe.parameters['valkey']['dtypes'] = valkey_dtypes
+        if not pipe.temporary:
+            edit_success, edit_msg = pipe.edit(debug=debug)
+            if not edit_success:
+                return edit_success, edit_msg
+
+    unseen_df, update_df, delta_df = (
+        pipe.filter_existing(df, include_unchanged_columns=True, debug=debug)
+        if check_existing and not upsert
+        else (df, None, df)
+    )
+    num_insert = len(unseen_df) if unseen_df is not None else 0
+    num_update = len(update_df) if update_df is not None else 0
+    msg = (
+        f"Inserted {num_insert}, updated {num_update} rows."
+        if not upsert
+        else f"Upserted {num_insert} rows."
+    )
+    if len(delta_df) == 0:
+        return True, msg
+
+    unseen_docs = unseen_df.to_dict(orient='records')
+    unseen_indices_docs = _serialize_indices_docs(unseen_docs)
+    unseen_ix_vals = {
+        get_document_key(doc, indices, table_name): serialize_document(doc)
+        for doc in unseen_docs
+    }
+    for key, val in unseen_ix_vals.items():
+        try:
+            self.set(key, val)
+        except Exception as e:
+            return False, f"Failed to set keys for {pipe}:\n{e}"
+
+    try:
+        self.push_docs(
+            unseen_indices_docs,
+            pipe.target,
+            datetime_column=dt_col,
+            debug=debug,
+        )
+    except Exception as e:
+        return False, f"Failed to push docs to '{pipe.target}':\n{e}"
+
+    update_docs = update_df.to_dict(orient='records') if update_df is not None else []
+    update_ix_docs = {
+        get_document_key(doc, indices, table_name): doc
+        for doc in update_docs
+    }
+    for key, doc in update_ix_docs.items():
+        try:
+            old_doc = json.loads(self.get(key))
+            old_doc.update(doc)
+            self.set(key, serialize_document(old_doc))
+        except Exception as e:
+            return False, f"Failed to set keys for {pipe}:\n{e}"
+
+    return True, msg
+
+
+def get_pipe_columns_types(
+    self,
+    pipe: mrsm.Pipe,
+    debug: bool = False,
+    **kwargs: Any
+) -> Dict[str, str]:
+    """
+    Return the data types for the columns in the target table for data type enforcement.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe whose target table contains columns and data types.
+
+    Returns
+    -------
+    A dictionary mapping columns to data types.
+    """
+    if not pipe.exists(debug=debug):
+        return {}
+
+    from meerschaum.utils.dtypes.sql import get_db_type_from_pd_type
+    return {
+        col: get_db_type_from_pd_type(typ)
+        for col, typ in pipe.parameters.get('valkey', {}).get('dtypes', {}).items()
+    }
+
+
+def clear_pipe(
+    self,
+    pipe: mrsm.Pipe,
+    begin: Union[datetime, int, None] = None,
+    end: Union[datetime, int, None] = None,
+    params: Optional[Dict[str, Any]] = None,
+    debug: bool = False,
+) -> mrsm.SuccessTuple:
+    """
+    Delete rows within `begin`, `end`, and `params`.
+
+    Parameters
+    ----------
+    pipe: mrsm.Pipe
+        The pipe whose rows to clear.
+
+    begin: Union[datetime, int, None], default None
+        If provided, remove rows >= `begin`.
+
+    end: Union[datetime, int, None], default None
+        If provided, remove rows < `end`.
+
+    params: Optional[Dict[str, Any]], default None
+        If provided, only remove rows which match the `params` filter.
+
+    Returns
+    -------
+    A `SuccessTuple` indicating success.
+    """
+    if begin is None and end is None and params is None:
+        return self.drop_pipe(pipe, debug=debug)
+
+    dt_col = pipe.columns.get('datetime', None)
+
+    existing_df = pipe.get_data(
+        begin=begin,
+        end=end,
+        params=params,
+        debug=debug,
+    )
+    if existing_df is None or len(existing_df) == 0:
+        return True, "Deleted 0 rows."
+
+    docs = existing_df.to_dict(orient='records')
+    table_name = self.quote_table(pipe.target)
+    indices = [col for col in pipe.columns.values() if col]
+    for doc in docs:
+        set_doc_key = get_document_key(doc, indices)
+        table_doc_key = get_document_key(doc, indices, table_name)
+        try:
+            if dt_col:
+                self.client.zrem(table_name, set_doc_key)
+            else:
+                self.client.srem(table_name, set_doc_key)
+            self.client.delete(table_doc_key)
+        except Exception as e:
+            return False, f"Failed to delete documents:\n{e}"
+    msg = (
+        f"Deleted {len(docs)} row"
+        + ('s' if len(docs) != 1 else '')
+        + '.'
+    )
+    return True, msg
+
+
+def get_sync_time(
+    self,
+    pipe: mrsm.Pipe,
+    newest: bool = True,
+    **kwargs: Any
+) -> Union[datetime, int, None]:
+    """
+    Return the newest (or oldest) timestamp in a pipe.
+    """
+    from meerschaum.utils.dtypes import are_dtypes_equal
+    dt_col = pipe.columns.get('datetime', None)
+    dt_typ = pipe.dtypes.get(dt_col, 'datetime64[ns]')
+    if not dt_col:
+        return None
+
+    dateutil_parser = mrsm.attempt_import('dateutil.parser')
+    table_name = self.quote_table(pipe.target)
+    try:
+        vals = (
+            self.client.zrevrange(table_name, 0, 0)
+            if newest
+            else self.client.zrange(table_name, 0, 0)
+        )
+        if not vals:
+            return None
+        val = vals[0]
+    except Exception:
+        return None
+
+    doc = json.loads(val)
+    dt_val = doc.get(dt_col, None)
+    if dt_val is None:
+        return None
+
+    try:
+        return (
+            int(dt_val)
+            if are_dtypes_equal(dt_typ, 'int')
+            else dateutil_parser.parse(str(dt_val)).replace(tzinfo=None)
+        )
+    except Exception as e:
+        warn(f"Failed to parse sync time for {pipe}:\n{e}")
+
+    return None
+
+
+def get_pipe_rowcount(
+    self,
+    pipe: mrsm.Pipe,
+    begin: Union[datetime, int, None] = None,
+    end: Union[datetime, int, None] = None,
+    params: Optional[Dict[str, Any]] = None,
+    debug: bool = False,
+    **kwargs: Any
+) -> Union[int, None]:
+    """
+    Return the number of documents in the pipe's set.
+    """
+    dt_col = pipe.columns.get('datetime', None)
+    table_name = self.quote_table(pipe.target)
+
+    if not pipe.exists():
+        return 0
+
+    try:
+        if begin is None and end is None and params is None:
+            return (
+                self.client.zcard(table_name)
+                if dt_col
+                else self.client.llen(table_name)
+            )
+    except Exception:
+        return None
+
+    df = pipe.get_data(begin=begin, end=end, params=params, debug=debug)
+    if df is None:
+        return 0
+
+    return len(df)
+
+
+def fetch_pipes_keys(
+    self,
+    connector_keys: Optional[List[str]] = None,
+    metric_keys: Optional[List[str]] = None,
+    location_keys: Optional[List[str]] = None,
+    tags: Optional[List[str]] = None,
+    params: Optional[Dict[str, Any]] = None,
+    debug: bool = False
+) -> Optional[List[Tuple[str, str, Optional[str]]]]:
+    """
+    Return the keys for the registered pipes.
+    """
+    from meerschaum.utils.dataframe import query_df
+    from meerschaum.utils.misc import separate_negation_values
+    try:
+        df = self.read(PIPES_TABLE, debug=debug)
+    except Exception:
+        return []
+
+    if df is None or len(df) == 0:
+        return []
+
+    query = {}
+    if connector_keys:
+        query['connector_keys'] = [str(k) for k in connector_keys]
+    if metric_keys:
+        query['metric_key'] = [str(k) for k in metric_keys]
+    if location_keys:
+        query['location_key'] = [str(k) for k in location_keys]
+    if params:
+        query.update(params)
+
+    df = query_df(df, query, inplace=True)
+
+    keys = [
+        (
+            doc['connector_keys'],
+            doc['metric_key'],
+            doc['location_key'],
+        )
+        for doc in df.to_dict(orient='records')
+    ]
+    if not tags:
+        return keys
+
+    tag_groups = [tag.split(',') for tag in tags]
+    in_ex_tag_groups = [separate_negation_values(tag_group) for tag_group in tag_groups]
+
+    filtered_keys = []
+    for ck, mk, lk in keys:
+        pipe = mrsm.Pipe(ck, mk, lk, instance=self)
+        pipe_tags = set(pipe.tags)
+        
+        include_pipe = True
+        for in_tags, ex_tags in in_ex_tag_groups:
+            all_in = all(tag in pipe_tags for tag in in_tags)
+            any_ex = any(tag in pipe_tags for tag in ex_tags)
+
+            if (not all_in) or any_ex:
+                include_pipe = False
+                continue
+
+        if include_pipe:
+            filtered_keys.append((ck, mk, lk))
+
+    return filtered_keys