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.dtypes import json_serialize_value
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=json_serialize_value,
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])

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

189def get_pipe_id(
190    self,
191    pipe: mrsm.Pipe,
192    debug: bool = False,
193    **kwargs: Any
194) -> Union[str, int, None]:
195    """
196    Return the `_id` for the pipe if it exists.
197
198    Parameters
199    ----------
200    pipe: mrsm.Pipe
201        The pipe whose `_id` to fetch.
202
203    Returns
204    -------
205    The `_id` for the pipe's document or `None`.
206    """
207    pipe_key = get_pipe_key(pipe)
208    try:
209        return int(self.get(pipe_key))
210    except Exception:
211        pass
212    return None

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

252def edit_pipe(
253    self,
254    pipe: mrsm.Pipe,
255    debug: bool = False,
256    **kwargs: Any
257) -> mrsm.SuccessTuple:
258    """
259    Edit the attributes of the pipe.
260
261    Parameters
262    ----------
263    pipe: mrsm.Pipe
264        The pipe whose in-memory parameters must be persisted.
265
266    Returns
267    -------
268    A `SuccessTuple` indicating success.
269    """
270    pipe_id = pipe.get_id(debug=debug)
271    if pipe_id is None:
272        return False, f"{pipe} is not registered."
273
274    parameters_key = get_pipe_parameters_key(pipe)
275    parameters_str = json.dumps(pipe.parameters, separators=(',', ':'))
276    self.set(parameters_key, parameters_str)
277    return True, "Success"

280def pipe_exists(
281    self,
282    pipe: mrsm.Pipe,
283    debug: bool = False,
284    **kwargs: Any
285) -> bool:
286    """
287    Check whether a pipe's target table exists.
288
289    Parameters
290    ----------
291    pipe: mrsm.Pipe
292        The pipe to check whether its table exists.
293
294    Returns
295    -------
296    A `bool` indicating the table exists.
297    """
298    table_name = self.quote_table(pipe.target)
299    return self.client.exists(table_name) != 0