relationalai 0.13.0.dev0__py3-none-any.whl → 0.13.1__py3-none-any.whl

relationalai/util/snowflake_handler.py

@@ -0,0 +1,88 @@
+import queue
+import threading
+import time
+
+from logging import Handler, LogRecord
+import traceback
+from uuid import UUID
+
+from relationalai.util.span_tracker import SpanTracker, record_is_span
+
+
+class SnowflakeHandler(Handler):
+    """
+    This is a logging handler that inserts spans directly into a Snowflake table.
+
+    It uses a queue and a worker thread to buffer spans and then insert them in batches.
+    """
+
+    def __init__(self, trace_id: UUID, snowflake_conn, sleep_interval_s=2):
+        super().__init__()
+        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=None, log_span_attrs_as_str=True)
+        Handler.__init__(self)
+        self.snowflake_conn = snowflake_conn
+        self.queue = queue.Queue()
+        self.is_shut_down = False
+        self.sleep_interval_s = sleep_interval_s
+        self.worker_thread = threading.Thread(target=self._consume_loop)
+        self.worker_thread.start()
+        print(f'snowflake logger started with trace id {trace_id}')
+
+    def emit(self, record: LogRecord):
+        if record_is_span(record):
+            self.queue.put(record.msg)
+
+    def _consume_loop(self):
+        while True:
+            while not self.queue.empty():
+                try:
+                    batch = self._get_batch()
+                    self._send_batch(batch)
+                except Exception as e:
+                    print('snowflake logger error:', e)
+                    traceback.print_exc()
+            if self.is_shut_down:
+                return
+            time.sleep(self.sleep_interval_s)
+
+    def _get_batch(self):
+        batch = []
+        try:
+            while True:
+                item = self.queue.get_nowait()
+                batch.append(item)
+        except queue.Empty:
+            pass
+        return batch
+
+    def _send_batch(self, batch):
+        if not batch:
+            return
+
+        inserts = []
+
+        for event in batch:
+            if event["event"] == "span_start":
+                self.span_tracker.handle_span_start(event)
+            elif event["event"] == "span_end":
+                span_end_result = self.span_tracker.handle_span_end(event)
+                if span_end_result:
+                    formatted_span_msg = span_end_result[0]
+                    if formatted_span_msg:
+                        inserts.append(formatted_span_msg)
+        # execute
+        if inserts:
+            self.snowflake_conn.cursor().executemany(
+                """
+                INSERT INTO spans_raw (id, parent_id, trace_id, type, key, start_ts, finish_ts, attrs)
+                VALUES (%(id)s, %(parent_id)s, %(trace_id)s, %(type)s, %(key)s, %(start_timestamp)s, %(end_timestamp)s, %(attrs)s)
+                """,
+                inserts,
+            )
+
+    
+    def shut_down(self):
+        self.is_shut_down = True
+        print('snowflake logger: waiting to shut down...')
+        self.worker_thread.join()
+        print('snowflake logger: shut down')

relationalai/util/span_format_test.py

@@ -0,0 +1,43 @@
+import logging
+from typing import Any, Dict, List
+
+from relationalai.debugging import Span
+
+class SpanCollectorHandler(logging.Handler):
+    def __init__(self):
+        super().__init__()
+        self.root = None
+        self.nodes = {}
+        self.events = []
+
+    def emit(self, record):
+        if isinstance(record.msg, dict):
+            if record.msg['event'] == 'span_start':
+                span = record.msg['span']
+                self.nodes[str(span.id)] = span
+            else:
+                self.events.append(record.msg)
+
+def assert_valid_span_structure(spans: Dict[str, 'Span'], events: List[Dict[str, Any]]):
+    """
+    Assert that the tree of spans & events has a form that will
+    work for the debuggers.
+    """
+
+    transaction_created_events = [
+        event for event in events if event['event'] == 'transaction_created'
+    ]
+    transaction_spans = [
+        span for span in spans.values() if span.type == 'transaction'
+    ]
+
+    assert len(transaction_created_events) > 0
+    assert len(transaction_spans) > 0
+
+    # only one transaction_created event per transaction span
+    # TODO: Should this be equal?
+    assert len(transaction_created_events) <= len(transaction_spans)
+
+    # transaction_created spans have a txn_id attribute
+    for event in transaction_created_events:
+        assert 'txn_id' in event

relationalai/util/span_tracker.py

@@ -0,0 +1,207 @@
+import json
+from logging import LogRecord
+import os
+import subprocess
+from typing import List, Tuple, Union, Dict, Any
+from uuid import UUID, uuid4
+
+from relationalai.debugging import Span, encode_log_message, filter_span_attrs, logger
+from relationalai.util.constants import SPAN_ATTR_ALLOW_LIST, SPAN_TYPES_KEYS
+from relationalai.util.format import default_serialize
+from relationalai.clients.config import Config
+
+TRACE_ID = uuid4()
+
+class SpanTracker:
+    """
+    Abstract base class for tracking spans produced as log messages. 
+    Provides methods to handle the start and end of spans, format span messages, and manage span attributes.
+    """
+
+    def __init__(self, trace_id: UUID, span_types_to_skip=None, log_span_attrs_as_str=True):
+        self.trace_id = trace_id
+        self._set_span_types_to_skip(span_types_to_skip)
+        self._log_span_attrs_as_str = log_span_attrs_as_str
+        self.open_spans = {}
+        self.skipped_spans = set()
+
+    def handle_span_start(self, msg: dict) -> Union[None,Span]:
+        span: Span = msg["span"]
+        assert_key_exists_for_span_type(span)
+        if span.type in self.span_types_to_skip:
+            self.skipped_spans.add(str(span.id))
+            return None
+        self.open_spans[str(span.id)] = span
+        return span
+
+    def handle_span_end(self, msg: dict) -> Union[Tuple[dict, str, dict], None]:
+        if str(msg["id"]) in self.skipped_spans:
+            return None
+        span = self.open_spans.pop(msg["id"])
+        joined_key_path = get_span_path_as_str(span)
+        end_attrs = get_span_allowed_attributes_values(span, msg["end_attrs"])
+        combined_attributes = {
+            **span.attrs, # span_start attrs
+            **end_attrs, # span_end attrs
+            "key_path": joined_key_path,
+            "elapsed_s": (span.end_timestamp - span.start_timestamp).total_seconds(),
+        }
+        filtered_attrs = filter_span_attrs(combined_attributes)
+        formatted_span_msg = format_span_msg(self.trace_id, span, filtered_attrs, attrs_as_str=self._log_span_attrs_as_str)
+        return formatted_span_msg, joined_key_path, filtered_attrs
+
+    def _set_span_types_to_skip(self, span_types_to_skip):
+        if span_types_to_skip is None:
+            self.span_types_to_skip = set()
+        elif isinstance(span_types_to_skip, str):
+            self.span_types_to_skip = {span_types_to_skip}
+        elif isinstance(span_types_to_skip, (list, tuple)):
+            self.span_types_to_skip = set(span_types_to_skip)
+        elif isinstance(span_types_to_skip, set):
+            self.span_types_to_skip = span_types_to_skip
+        else:
+            raise ValueError(f"Invalid span_types_to_skip: {span_types_to_skip}, type: {type(span_types_to_skip)}")
+
+def format_span_json(trace_id: UUID, span: Span, attrs: dict, attrs_as_str:bool=True) -> dict:
+    span_key = get_span_value_from_key(span)
+    if attrs_as_str:
+        a = json.dumps(attrs, default=default_serialize)
+    else:
+        a = encode_log_message(attrs)
+
+    span_json = span.to_json()
+    extra_attrs = {
+        "trace_id": trace_id,
+        "key": str(span_key),
+        "attrs": a
+    }
+    formatted_span_msg = {**span_json, **extra_attrs}
+    return formatted_span_msg
+
+#region Span formatting and filtering functions
+def format_span_msg(trace_id: UUID, span: Span, attrs: dict, attrs_as_str:bool=True) -> dict:
+    span_key = get_span_value_from_key(span)
+    if attrs_as_str:
+        a = json.dumps(attrs, default=default_serialize)
+    else:
+        a = encode_log_message(attrs)
+
+    span_json = span.to_json()
+    extra_attrs = {
+        "trace_id": str(trace_id) if attrs_as_str else trace_id,
+        "key": str(span_key) if attrs_as_str else span_key,
+        "attrs": a
+    }
+    formatted_span_msg = {**span_json, **extra_attrs}
+    return formatted_span_msg
+
+def get_span_allowed_attributes_values(span: Span, attrs=None) -> dict:
+    # For an input span with attrs (or instead for an `attrs` dict), return a dict with only the allowed attributes for the span type
+    attrs = attrs or span.attrs
+    out = {}
+    allowed_attributes = SPAN_ATTR_ALLOW_LIST.get(span.type, set())
+    for attr in allowed_attributes:
+        if attr in attrs:
+            out[attr] = attrs[attr]
+    return out
+
+def assert_key_exists_for_span_type(span: Span):
+    if span.type in SPAN_TYPES_KEYS:
+        key = SPAN_TYPES_KEYS[span.type]
+        if key not in span.attrs:
+            raise ValueError(f"Attribute {key} is required for spans of type {span.type}, but wasn't provided. Provided attributes:\n{span.attrs}")
+
+def get_span_key_path(span: Span) -> str:
+    # Return the key path of the current span
+    value = get_span_value_from_key(span)
+    if value is None:
+        return span.type
+    return f"{span.type}({value})"
+
+def get_span_path_as_list(span: Span) -> List[Span]:
+    # Return the path from the root to the current span
+    path = []
+    cur = span
+    while cur is not None:
+        path.append(cur)
+        cur = cur.parent
+    path.reverse()
+    return path
+
+def get_span_path_as_str(span: Span) -> str:
+    # Return a string representation of the path from the root to the current span
+    # formerly format_key_path function
+    return '.'.join([get_span_key_path(span) for span in get_span_path_as_list(span)])
+
+def get_span_value_from_key(span: Span):
+    if span.type in SPAN_TYPES_KEYS:
+        key = SPAN_TYPES_KEYS[span.type]
+        if key in span.attrs:
+            return span.attrs[key]
+        raise ValueError(f"Key {key} not found in span of type {span.type}\n{span.attrs}")
+    return None
+
+def record_is_span(record: LogRecord) -> bool:
+    return (
+        isinstance(record.msg, dict) and
+        "event" in record.msg and
+        record.msg["event"] in ("span_start", "span_end")
+    )
+
+def _get_local_git_info() -> Dict[str, str]:
+    try:
+        branch = subprocess.check_output(['git', 'rev-parse', '--abbrev-ref', 'HEAD']).strip().decode('utf-8')
+        sha = subprocess.check_output(['git', 'rev-parse', 'HEAD']).strip().decode('utf-8')
+        repo = subprocess.check_output(['git', 'config', '--get', 'remote.origin.url']).strip().decode('utf-8').rsplit('/', maxsplit=1)[-1].replace('.git', '')
+        return {
+            'commit': sha,
+            'branch': branch,
+            'repo': repo,
+        }
+    except subprocess.CalledProcessError as e:
+        logger.warning(f"Could not get git info: {e}")
+        return {
+            'commit': 'local',
+            'branch': 'local',
+            'repo': 'local',
+        }
+
+def _get_git_info() -> Dict[str, str]:
+    if os.getenv("GITHUB_REF_NAME") is not None:
+        return {
+            'commit': os.getenv("GITHUB_SHA", ""),
+            'branch': os.getenv("GITHUB_REF_NAME", ""),
+            'repo': os.getenv("GITHUB_REPOSITORY", ""),
+        }
+    return _get_local_git_info()
+
+def _get_masked_config() -> Dict[str, Any]:
+    """Return a config with all sensitive keys (secrets and passwords) removed, such that it can be safely logged"""
+    config = Config()
+    masked_props = {
+        k: v for k, v in config.props.items()
+        if not isinstance(v, str)
+        or not any(prop_key_pattern in k for prop_key_pattern in ["secret", "password"])
+    }
+    return masked_props
+
+def _get_cloud_provider() -> Dict[str, str]:
+    provider = os.getenv("RAI_CLOUD_PROVIDER")
+    if not provider:
+        config = Config()
+        provider = config.get("platform")
+    assert provider, "Could not retrieve cloud provider from environment or config"
+    return {"platform": provider}
+
+def get_root_span_attrs(get_full_config: bool = False) -> Dict[str, Any]:
+    git_info = _get_git_info()
+    cloud_provider = _get_cloud_provider()
+    masked_config = _get_masked_config() if get_full_config else {}
+    attrs = {
+        **git_info,
+        **cloud_provider,
+        **masked_config
+    }
+    return attrs
+
+#endregion

relationalai/util/spans_file_handler.py

@@ -0,0 +1,72 @@
+from copy import deepcopy
+from logging import FileHandler
+import os
+import sys
+from uuid import UUID
+from typing import Dict, Union
+
+from relationalai.debugging import JsonFormatter
+from relationalai.util import get_timestamp
+from relationalai.util.format import str_to_bool
+from relationalai.util.span_tracker import SpanTracker, record_is_span
+
+# Span types that are automatically generated by PyRel but are not relevant for performance analysis
+SKIP_SPAN_TYPES = ['check_status']
+
+
+class SpansFileHandler(FileHandler):
+    """
+    This logger writes spans to a local .jsonl file inside the `spans_output` folder as they finish.
+    The primary goal is to ingest these spans into Snowflake for analytics.
+    """
+
+    def __init__(self, trace_id: UUID, filename: Union[str,None] = None):
+        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=SKIP_SPAN_TYPES, log_span_attrs_as_str=False)
+        if filename is None:
+            filename = self._get_default_log_file_path()
+            if not os.path.exists(os.path.dirname(filename)):
+                os.makedirs(os.path.dirname(filename))
+        FileHandler.__init__(self, filename=filename)
+        self.setFormatter(JsonFormatter())
+
+    def emit(self, record):
+        if record_is_span(record):
+            msg: Dict = record.msg # type: ignore
+            if msg["event"] == "span_start":
+                self.span_tracker.handle_span_start(msg)
+            if msg["event"] == "span_end":
+                span_end_result = self.span_tracker.handle_span_end(msg)
+                if span_end_result:
+                    formatted_span_msg = span_end_result[0]
+                    if formatted_span_msg is None: # This will happen for span_types_to_skip
+                        return
+                    record_copy = deepcopy(record)  # Create a deep copy so as to not interfere with other handlers
+                    record_copy.msg = formatted_span_msg
+                    FileHandler.emit(self, record_copy)
+
+    def _get_default_log_file_path(self) -> str:
+        now_str = get_timestamp().strftime("%Y%m%d_%H%M%S")
+        trace_id_trunc = str(self.span_tracker.trace_id)[:8]
+        spans_subfolder = os.getenv("SPANS_SUBFOLDER", "")
+        return os.path.join(
+            os.getcwd(), "output_spans", spans_subfolder, f"{now_str}__{trace_id_trunc}.jsonl"
+        )
+
+
+def is_collect_only():
+    return "pytest" in sys.modules and "--collect-only" in sys.argv
+
+
+def maybe_add_spans_file_handler(logger, trace_id:UUID, if_not_set: bool = False):
+    """
+    Add a SpansFileHandler to the logger if the 'ENABLE_SPANS_FILE_HANDLER' environment variable is set to True,
+    or if the 'if_not_set' parameter is True.
+    """
+    if str_to_bool(os.getenv("ENABLE_SPANS_FILE_HANDLER")) or if_not_set:
+        if is_collect_only():
+            print("Running in collect-only mode, skipping SpansFileHandler")
+        elif any(isinstance(handler, SpansFileHandler) for handler in logger.handlers):
+            print("logger already has a SpansFileHandler instance, won't add another one")
+        else:
+            logger.debug("Adding SpansFileHandler to logger")
+            logger.addHandler(SpansFileHandler(trace_id))

relationalai/util/tracing_handler.py

@@ -0,0 +1,34 @@
+from logging import Handler
+from typing import Dict
+from uuid import UUID
+
+from relationalai.util.span_tracker import SpanTracker, get_span_allowed_attributes_values, get_span_path_as_str, record_is_span
+
+
+class TracingHandler(Handler):
+    """
+    This logger prints spans to stdout as they start and finish.
+    """
+
+    def __init__(self, trace_id: UUID):
+        super().__init__()
+        self.span_tracker = SpanTracker(trace_id, span_types_to_skip=None)
+
+    def emit(self, record):
+        if record_is_span(record):
+            msg: Dict = record.msg # type: ignore
+            if msg["event"] == "span_start":
+                span_to_print = self.span_tracker.handle_span_start(msg)
+                if span_to_print is not None:
+                    attrs = get_span_allowed_attributes_values(span_to_print)
+                    print("start", get_span_path_as_str(span_to_print), attrs)
+            elif msg["event"] == "span_end":
+                span_end_result = self.span_tracker.handle_span_end(msg)
+                if span_end_result:
+                    joined_key_path, filtered_attrs = span_end_result[1], span_end_result[2]
+                    if joined_key_path is not None and filtered_attrs is not None:
+                        # key_path is already printed separately, and name is already in the key_path
+                        filtered_attrs = {k: v for k, v in filtered_attrs.items() if k not in ["name", "key_path"]}
+                        # Exclude class instances, which can have a large nested structure to print
+                        filtered_attrs = {k: v for k, v in filtered_attrs.items() if not hasattr(v, "__dict__")}
+                        print("  end", joined_key_path, filtered_attrs)

relationalai-0.13.1.dist-info/METADATA

@@ -0,0 +1,74 @@
+Metadata-Version: 2.4
+Name: relationalai
+Version: 0.13.1
+Summary: RelationalAI Library and CLI
+Author-email: RelationalAI <support@relational.ai>
+License-File: LICENSE
+Requires-Python: >=3.10
+Requires-Dist: aiohttp
+Requires-Dist: click==8.2.1
+Requires-Dist: colorama
+Requires-Dist: cryptography
+Requires-Dist: gravis
+Requires-Dist: inquirerpy
+Requires-Dist: lqp==0.2.1
+Requires-Dist: nicegui==2.16.1
+Requires-Dist: numpy<2
+Requires-Dist: opentelemetry-api
+Requires-Dist: opentelemetry-exporter-otlp-proto-http
+Requires-Dist: opentelemetry-sdk
+Requires-Dist: packaging
+Requires-Dist: pandas
+Requires-Dist: pyarrow
+Requires-Dist: pyjwt
+Requires-Dist: requests
+Requires-Dist: rich
+Requires-Dist: snowflake-connector-python[secure-local-storage]
+Requires-Dist: snowflake-snowpark-python>=1.38.0
+Requires-Dist: toml
+Requires-Dist: tomlkit
+Requires-Dist: typing-extensions
+Requires-Dist: wcwidth
+Requires-Dist: websockets
+Provides-Extra: dev
+Requires-Dist: argcomplete; extra == 'dev'
+Requires-Dist: build; extra == 'dev'
+Requires-Dist: coverage; extra == 'dev'
+Requires-Dist: duckdb; extra == 'dev'
+Requires-Dist: faker; extra == 'dev'
+Requires-Dist: hupper; extra == 'dev'
+Requires-Dist: hypothesis; extra == 'dev'
+Requires-Dist: jupyter; extra == 'dev'
+Requires-Dist: logscope; extra == 'dev'
+Requires-Dist: pandas-stubs; extra == 'dev'
+Requires-Dist: pexpect; extra == 'dev'
+Requires-Dist: pip-tools; extra == 'dev'
+Requires-Dist: polars; extra == 'dev'
+Requires-Dist: prompt-toolkit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: pytest-cov; extra == 'dev'
+Requires-Dist: pytest-mock; extra == 'dev'
+Requires-Dist: pytest-snapshot; extra == 'dev'
+Requires-Dist: pytest-timeout; extra == 'dev'
+Requires-Dist: ruff; extra == 'dev'
+Requires-Dist: scipy; extra == 'dev'
+Requires-Dist: sqlite-utils; extra == 'dev'
+Requires-Dist: sympy; extra == 'dev'
+Requires-Dist: tablescope; extra == 'dev'
+Requires-Dist: watchdog; extra == 'dev'
+Provides-Extra: ea
+Requires-Dist: dataclasses-json; extra == 'ea'
+Requires-Dist: owlready2; extra == 'ea'
+Requires-Dist: parameterized; extra == 'ea'
+Requires-Dist: pysmt; extra == 'ea'
+Requires-Dist: simple-ddl-parser; extra == 'ea'
+Requires-Dist: xmltodict; extra == 'ea'
+Provides-Extra: legacy
+Requires-Dist: rai-sdk>=0.7.5; extra == 'legacy'
+Description-Content-Type: text/markdown
+
+# The RelationalAI Python Library
+
+The Python library for building and querying knowledge graphs with RelationalAI.
+
+For more information, visit our documentation at [https://relational.ai/docs](https://relational.ai/docs).