meerschaum.connectors.valkey

 86    @property
 87    def client(self):
 88        """
 89        Return the Valkey client.
 90        """
 91        if '_client' in self.__dict__:
 92            return self.__dict__['_client']
 93
 94        valkey = mrsm.attempt_import('valkey')
 95
 96        if 'uri' in self.__dict__:
 97            self._client = valkey.Valkey.from_url(self.__dict__.get('uri'))
 98            return self._client
 99
100        optional_kwargs = {
101            key: self.__dict__.get(key)
102            for key in self.OPTIONAL_ATTRIBUTES
103            if key in self.__dict__
104        }
105        connection_kwargs = {
106            'host': self.host,
107            **optional_kwargs
108        }
109
110        self._client = valkey.Valkey(**connection_kwargs)
111        return self._client

113    @property
114    def URI(self) -> str:
115        """
116        Return the connection URI for this connector.
117        """
118        import urllib.parse
119
120        if 'uri' in self.__dict__:
121            return self.__dict__.get('uri')
122
123        uri = "valkey://"
124        if 'username' in self.__dict__:
125            uri += urllib.parse.quote_plus(self.username) + ':'
126
127        if 'password' in self.__dict__:
128            uri += urllib.parse.quote_plus(self.password) + '@'
129
130        if 'host' in self.__dict__:
131            uri += self.host
132
133        if 'port' in self.__dict__:
134            uri += f':{self.port}'
135
136        if 'db' in self.__dict__:
137            uri += f"/{self.db}"
138
139        if 'socket_timeout' in self.__dict__:
140            uri += f"?timeout={self.socket_timeout}s"
141
142        return uri

144    def set(self, key: str, value: Any, **kwargs: Any) -> None:
145        """
146        Set the `key` to `value`.
147        """
148        return self.client.set(key, value, **kwargs)

150    def get(self, key: str) -> Union[str, None]:
151        """
152        Get the value for `key`.
153        """
154        val = self.client.get(key)
155        if val is None:
156            return None
157
158        return val.decode('utf-8')

160    def test_connection(self) -> bool:
161        """
162        Return whether a connection may be established.
163        """
164        return self.client.ping()

166    @classmethod
167    def quote_table(cls, table: str) -> str:
168        """
169        Return a quoted key.
170        """
171        return shlex.quote(table)

173    @classmethod
174    def get_counter_key(cls, table: str) -> str:
175        """
176        Return the counter key for a given table.
177        """
178        table_name = cls.quote_table(table)
179        return f"{table_name}:counter"

181    def push_df(
182        self,
183        df: 'pd.DataFrame',
184        table: str,
185        datetime_column: Optional[str] = None,
186        debug: bool = False,
187    ) -> int:
188        """
189        Append a pandas DataFrame to a table.
190
191        Parameters
192        ----------
193        df: pd.DataFrame
194            The pandas DataFrame to append to the table.
195
196        table: str
197            The "table" name (root key).
198
199        datetime_column: Optional[str], default None
200            If provided, use this key as the datetime index.
201
202        Returns
203        -------
204        The current index counter value (how many docs have been pushed).
205        """
206        from meerschaum.utils.dataframe import to_json
207        docs_str = to_json(df)
208        docs = json.loads(docs_str)
209        return self.push_docs(
210            docs,
211            table,
212            datetime_column=datetime_column,
213            debug=debug,
214        )

216    def push_docs(
217        self,
218        docs: List[Dict[str, Any]],
219        table: str,
220        datetime_column: Optional[str] = None,
221        debug: bool = False,
222    ) -> int:
223        """
224        Append a list of documents to a table.
225
226        Parameters
227        ----------
228        docs: List[Dict[str, Any]]
229            The docs to be pushed.
230            All keys and values will be coerced into strings.
231
232        table: str
233            The "table" name (root key).
234
235        datetime_column: Optional[str], default None
236            If set, create a sorted set with this datetime column as the index.
237            Otherwise push the docs to a list.
238
239        Returns
240        -------
241        The current index counter value (how many docs have been pushed).
242        """
243        from meerschaum.utils.misc import json_serialize_datetime
244        table_name = self.quote_table(table)
245        datetime_column_key = self.get_datetime_column_key(table)
246        remote_datetime_column = self.get(datetime_column_key)
247        datetime_column = datetime_column or remote_datetime_column
248        dateutil_parser = mrsm.attempt_import('dateutil.parser')
249
250        old_len = (
251            self.client.zcard(table_name)
252            if datetime_column
253            else self.client.scard(table_name)
254        )
255        for doc in docs:
256            original_dt_val = (
257                doc[datetime_column]
258                if datetime_column and datetime_column in doc
259                else 0
260            )
261            dt_val = (
262                dateutil_parser.parse(str(original_dt_val))
263                if not isinstance(original_dt_val, int)
264                else int(original_dt_val)
265            ) if datetime_column else None
266            ts = (
267                int(dt_val.replace(tzinfo=timezone.utc).timestamp())
268                if isinstance(dt_val, datetime)
269                else int(dt_val)
270            ) if datetime_column else None
271            doc_str = json.dumps(
272                doc,
273                default=(lambda x: json_serialize_datetime(x) if hasattr(x, 'tzinfo') else str(x)),
274                separators=(',', ':'),
275                sort_keys=True,
276            )
277            if datetime_column:
278                self.client.zadd(table_name, {doc_str: ts})
279            else:
280                self.client.sadd(table_name, doc_str)
281
282        if datetime_column:
283            self.set(datetime_column_key, datetime_column)
284        new_len = (
285            self.client.zcard(table_name)
286            if datetime_column
287            else self.client.scard(table_name)
288        )
289
290        return new_len - old_len

308    def get_datetime_column_key(self, table: str) -> str:
309        """
310        Return the key to store the datetime index for `table`.
311        """
312        table_name = self.quote_table(table)
313        return f'{table_name}:datetime_column'

315    def read(
316        self,
317        table: str,
318        begin: Union[datetime, int, str, None] = None,
319        end: Union[datetime, int, str, None] = None,
320        params: Optional[Dict[str, Any]] = None,
321        datetime_column: Optional[str] = None,
322        select_columns: Optional[List[str]] = None,
323        omit_columns: Optional[List[str]] = None,
324        debug: bool = False
325    ) -> Union['pd.DataFrame', None]:
326        """
327        Query the table and return the result dataframe.
328
329        Parameters
330        ----------
331        table: str
332            The "table" name to be queried.
333
334        begin: Union[datetime, int, str, None], default None
335            If provided, only return rows greater than or equal to this datetime.
336
337        end: Union[datetime, int, str, None], default None
338            If provided, only return rows older than this datetime.
339
340        params: Optional[Dict[str, Any]]
341            Additional Meerschaum filter parameters.
342
343        datetime_column: Optional[str], default None
344            If provided, use this column for the datetime index.
345            Otherwise infer from the table metadata.
346
347        select_columns: Optional[List[str]], default None
348            If provided, only return these columns.
349
350        omit_columns: Optional[List[str]], default None
351            If provided, do not include these columns in the result.
352
353        Returns
354        -------
355        A Pandas DataFrame of the result, or `None`.
356        """
357        from meerschaum.utils.dataframe import parse_df_datetimes, query_df
358        docs = self.read_docs(
359            table,
360            begin=begin,
361            end=end,
362            debug=debug,
363        )
364        df = parse_df_datetimes(docs)
365        datetime_column_key = self.get_datetime_column_key(table)
366        datetime_column = datetime_column or self.get(datetime_column_key)
367
368        return query_df(
369            df,
370            begin=(begin if datetime_column is not None else None),
371            end=(end if datetime_column is not None else None),
372            params=params,
373            datetime_column=datetime_column,
374            select_columns=select_columns,
375            omit_columns=omit_columns,
376            inplace=True,
377            reset_index=True,
378            debug=debug,
379        )

381    def read_docs(
382        self,
383        table: str,
384        begin: Union[datetime, int, str, None] = None,
385        end: Union[datetime, int, str, None] = None,
386        debug: bool = False,
387    ) -> List[Dict[str, str]]:
388        """
389        Return a list of previously pushed docs.
390
391        Parameters
392        ----------
393        table: str
394            The "table" name (root key) under which the docs were pushed.
395
396        begin: Union[datetime, int, str, None], default None
397            If provided and the table was created with a datetime index, only return documents
398            newer than this datetime.
399            If the table was not created with a datetime index and `begin` is an `int`,
400            return documents with a positional index greater than or equal to this value.
401
402        end: Union[datetime, int, str, None], default None
403            If provided and the table was created with a datetime index, only return documents
404            older than this datetime.
405            If the table was not created with a datetime index and `begin` is an `int`,
406            return documents with a positional index less than this value.
407
408        Returns
409        -------
410        A list of dictionaries, where all keys and values are strings.
411        """
412        from meerschaum.utils.dtypes import coerce_timezone
413        table_name = self.quote_table(table)
414        datetime_column_key = self.get_datetime_column_key(table)
415        datetime_column = self.get(datetime_column_key)
416
417        if debug:
418            dprint(f"Reading documents from '{table}' with {begin=}, {end=}")
419
420        if not datetime_column:
421            return [
422                json.loads(doc_bytes.decode('utf-8'))
423                for doc_bytes in self.client.smembers(table_name)
424            ]
425
426        dateutil_parser = mrsm.attempt_import('dateutil.parser')
427
428        if isinstance(begin, str):
429            begin = coerce_timezone(dateutil_parser.parse(begin))
430
431        if isinstance(end, str):
432            end = coerce_timezone(dateutil_parser.parse(end))
433
434        begin_ts = (
435            (
436                int(begin.replace(tzinfo=timezone.utc).timestamp())
437                if isinstance(begin, datetime)
438                else int(begin)
439            )
440            if begin is not None else '-inf'
441        )
442        end_ts = (
443            (
444                int(end.replace(tzinfo=timezone.utc).timestamp())
445                if isinstance(end, datetime)
446                else int(end)
447            )
448            if end is not None else '+inf'
449        )
450
451        if debug:
452            dprint(f"Reading documents with {begin_ts=}, {end_ts=}")
453
454        return [
455            json.loads(doc_bytes.decode('utf-8'))
456            for doc_bytes in self.client.zrangebyscore(
457                table_name,
458                begin_ts,
459                end_ts,
460                withscores=False,
461            )
462        ]

508    def drop_table(self, table: str, debug: bool = False) -> None:
509        """
510        Drop a "table" of documents.
511
512        Parameters
513        ----------
514        table: str
515            The "table" name (root key) to be deleted.
516        """
517        table_name = self.quote_table(table)
518        datetime_column_key = self.get_datetime_column_key(table)
519        self.client.delete(table_name)
520        self.client.delete(datetime_column_key)

522    @classmethod
523    def get_entity_key(cls, *keys: Any) -> str:
524        """
525        Return a joined key to set an entity.
526        """
527        if not keys:
528            raise ValueError("No keys to be joined.")
529
530        for key in keys:
531            if cls.KEY_SEPARATOR in str(key):
532                raise ValueError(f"Key cannot contain separator '{cls.KEY_SEPARATOR}'.")
533
534        return cls.KEY_SEPARATOR.join([str(key) for key in keys])

126def register_pipe(
127    self,
128    pipe: mrsm.Pipe,
129    debug: bool = False,
130    **kwargs: Any
131) -> SuccessTuple:
132    """
133    Insert the pipe's attributes into the internal `pipes` table.
134
135    Parameters
136    ----------
137    pipe: mrsm.Pipe
138        The pipe to be registered.
139
140    Returns
141    -------
142    A `SuccessTuple` of the result.
143    """
144    attributes = {
145        'connector_keys': str(pipe.connector_keys),
146        'metric_key': str(pipe.metric_key),
147        'location_key': str(pipe.location_key),
148    }
149    parameters_str = json.dumps(
150        pipe._attributes.get('parameters', {}),
151        separators=(',', ':'),
152    )
153
154    pipe_key = get_pipe_key(pipe)
155    parameters_key = get_pipe_parameters_key(pipe)
156
157    try:
158        existing_pipe_id = self.get(pipe_key)
159        if existing_pipe_id is not None:
160            return False, f"{pipe} is already registered."
161
162        pipe_id = self.client.incr(PIPES_COUNTER)
163        _ = self.push_docs(
164            [{'pipe_id': pipe_id, **attributes}],
165            PIPES_TABLE,
166            datetime_column='pipe_id',
167            debug=debug,
168        )
169        self.set(pipe_key, pipe_id)
170        self.set(parameters_key, parameters_str)
171
172    except Exception as e:
173        return False, f"Failed to register {pipe}:\n{e}"
174
175    return True, "Success"

178def get_pipe_id(
179    self,
180    pipe: mrsm.Pipe,
181    debug: bool = False,
182    **kwargs: Any
183) -> Union[str, int, None]:
184    """
185    Return the `_id` for the pipe if it exists.
186
187    Parameters
188    ----------
189    pipe: mrsm.Pipe
190        The pipe whose `_id` to fetch.
191
192    Returns
193    -------
194    The `_id` for the pipe's document or `None`.
195    """
196    pipe_key = get_pipe_key(pipe)
197    try:
198        return int(self.get(pipe_key))
199    except Exception:
200        pass
201    return None

204def get_pipe_attributes(
205    self,
206    pipe: mrsm.Pipe,
207    debug: bool = False,
208    **kwargs: Any
209) -> Dict[str, Any]:
210    """
211    Return the pipe's document from the internal `pipes` collection.
212
213    Parameters
214    ----------
215    pipe: mrsm.Pipe
216        The pipe whose attributes should be retrieved.
217
218    Returns
219    -------
220    The document that matches the keys of the pipe.
221    """
222    pipe_id = pipe.get_id(debug=debug)
223    if pipe_id is None:
224        return {}
225
226    parameters_key = get_pipe_parameters_key(pipe)
227    parameters_str = self.get(parameters_key)
228
229    parameters = json.loads(parameters_str) if parameters_str else {}
230
231    attributes = {
232        'connector_keys': pipe.connector_keys,
233        'metric_key': pipe.metric_key,
234        'location_key': pipe.location_key,
235        'parameters': parameters,
236        'pipe_id': pipe_id,
237    }
238    return attributes

241def edit_pipe(
242    self,
243    pipe: mrsm.Pipe,
244    debug: bool = False,
245    **kwargs: Any
246) -> mrsm.SuccessTuple:
247    """
248    Edit the attributes of the pipe.
249
250    Parameters
251    ----------
252    pipe: mrsm.Pipe
253        The pipe whose in-memory parameters must be persisted.
254
255    Returns
256    -------
257    A `SuccessTuple` indicating success.
258    """
259    pipe_id = pipe.get_id(debug=debug)
260    if pipe_id is None:
261        return False, f"{pipe} is not registered."
262
263    parameters_key = get_pipe_parameters_key(pipe)
264    parameters_str = json.dumps(pipe.parameters, separators=(',', ':'))
265    self.set(parameters_key, parameters_str)
266    return True, "Success"

269def pipe_exists(
270    self,
271    pipe: mrsm.Pipe,
272    debug: bool = False,
273    **kwargs: Any
274) -> bool:
275    """
276    Check whether a pipe's target table exists.
277
278    Parameters
279    ----------
280    pipe: mrsm.Pipe
281        The pipe to check whether its table exists.
282
283    Returns
284    -------
285    A `bool` indicating the table exists.
286    """
287    table_name = self.quote_table(pipe.target)
288    return self.client.exists(table_name) != 0