relationalai 0.13.5__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

v0/relationalai/tools/cli_helpers.py

@@ -0,0 +1,390 @@
+#pyright: reportPrivateImportUsage=false
+from __future__ import annotations
+import io
+import os
+import re
+import sys
+import requests
+import rich
+import click
+import functools
+import pytz
+
+from v0.relationalai.util.constants import TOP_LEVEL_PROFILE_NAME
+from .. import Resources
+from rich.table import Table
+from typing import Callable, Dict, Any
+from ..clients import config
+from click.core import Context
+from rich.console import Console
+from rich import box as rich_box
+from collections import defaultdict
+from packaging.version import Version
+from ..clients.config import ConfigFile
+from datetime import datetime, timedelta
+from click.formatting import HelpFormatter
+from ..clients.client import ResourcesBase
+from v0.relationalai.tools.constants import GlobalProfile
+from v0.relationalai.tools.cli_controls import divider
+from v0.relationalai.util.format import humanized_bytes, humanized_duration
+
+#--------------------------------------------------
+# Helpers
+#--------------------------------------------------
+
+@functools.cache
+def get_config(profile:str|Dict[str, Any]|None = None):
+    return config.Config(profile or GlobalProfile.get())
+
+@functools.cache
+def get_resource_provider(platform:str|None=None, _cfg:config.Config|None = None) -> ResourcesBase:
+    cfg = _cfg or get_config()
+    provider = Resources(config=cfg)
+    return provider
+
+def unexpand_user_path(path):
+    """Inverse of os.path.expanduser"""
+    home_dir = os.path.expanduser('~')
+    if path.startswith(home_dir):
+        return '~' + path[len(home_dir):]
+    return path
+
+def account_from_url(account_or_url:str):
+    regex = r"https://app.snowflake.com/([^/]+)/([^/]+)/?.*"
+    match = re.match(regex, account_or_url)
+    if match:
+        org = match.group(1)
+        account = match.group(2)
+        return f"{org}-{account}"
+    elif "app.snowflake.com" in account_or_url or "https://" in account_or_url:
+        raise ValueError("URL not of the form https://app.snowflake.com/[org]/[account]")
+    else:
+        return account_or_url
+
+def supports_platform(*platform_names: str):
+    def decorator(cmd: click.Command):
+        setattr(cmd, "__available_platforms", platform_names)
+        cb = cmd.callback
+        assert cb
+        def checked(*args, **kwargs):
+            assert cmd.name
+            assert_command_available(cmd.name, command_available(cmd))
+            return cb(*args, **kwargs)
+
+        cmd.callback = checked
+        return cmd
+    return decorator
+
+def command_available(cmd: click.Command) -> bool:
+    available_platforms = getattr(cmd, "__available_platforms", ())
+    platform = get_config().get("platform", "")
+    return not available_platforms or not platform or platform in available_platforms
+
+def assert_command_available(name: str, available: bool, plural=False):
+    if not available:
+        platform = get_config().get("platform", "")
+        rich.print(f"[yellow]{name} {'are' if plural else 'is'} not available for {platform}")
+        divider()
+        sys.exit(1)
+
+def coming_soon():
+    rich.print("[yellow]This isn't quite ready yet, but it's coming soon!")
+    divider()
+    sys.exit(1)
+
+def issue_top_level_profile_warning():
+    rich.print("[yellow]Warning: Using a top-level profile is not recommended")
+    rich.print("[yellow]Consider naming the profile by adding \\[profile.<name>] to your raiconfig.toml file\n")
+    rich.print("[yellow]Example:")
+    rich.print("[yellow]\\[profile.default]")
+    rich.print("[yellow]platform = \"snowflake\"")
+    rich.print("[yellow]account = ...")
+    divider()
+
+def ensure_config(profile:str|None=None) -> config.Config:
+    cfg = get_config(profile)
+    if not cfg.file_path:
+        rich.print("[yellow bold]No configuration file found.")
+        rich.print("To create one, run: [green bold]rai init[/green bold]")
+        divider()
+        sys.exit(1)
+    return cfg
+
+def latest_version(package_name):
+    """Get the current version of a package on PyPI."""
+    url = f"https://pypi.org/pypi/{package_name}/json"
+    response = requests.get(url)
+    if response.status_code == 200:
+        data = response.json()
+        version = data['info']['version']
+        return version
+    else:
+        return None
+
+def is_latest_cli_version():
+    from .. import __version__
+    latest_ver_str = latest_version("relationalai")
+    latest_ver = Version(latest_ver_str) if latest_ver_str else Version("0.0.0")
+    version = Version(__version__)
+    return version >= latest_ver, version, latest_ver
+
+#--------------------------------------------------
+# Validation
+#--------------------------------------------------
+
+EMPTY_STRING_REGEX = re.compile(r'^\S+$')
+ENGINE_NAME_REGEX = re.compile(r'^[a-zA-Z]\w{2,}$')
+COMPUTE_POOL_REGEX = re.compile(r'^[a-zA-Z][a-zA-Z0-9_]*$')
+PASSCODE_REGEX = re.compile(r'^(\d{6})?$')
+ENGINE_NAME_ERROR = "Min 3 chars, start with letter, only letters, numbers, underscores allowed."
+UUID = re.compile('[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}')
+
+def validate_engine_name(name:str) -> tuple[bool, str|None]:
+    if not ENGINE_NAME_REGEX.match(name):
+        return False, ENGINE_NAME_ERROR
+    return True, None
+
+#--------------------------------------------------
+# Tables
+#--------------------------------------------------
+
+def get_color_by_state(state: str) -> str:
+    if state and isinstance(state, str):
+        state_lower = state.lower()
+        if state_lower in ("aborted", "quarantined"):
+            return "red"
+        elif state_lower == "completed":
+            return "white"
+        elif state_lower in ("running", "cancelling", "syncing", "pending", "processing"):
+            return "bold yellow"
+        elif state_lower == "suspended":
+            return "dim"
+        else:
+            return ""
+    return ""
+
+def format_value(value) -> str:
+    if value is None:
+        return "N/A"
+    elif isinstance(value, (int, float)):
+        return f"{value}"
+    elif isinstance(value, list):
+        return ", ".join(map(str, value))
+    elif isinstance(value, bool):
+        return f"{value}"
+    elif isinstance(value, datetime):
+        return value.strftime("%Y-%m-%d %H:%M:%S")
+    elif isinstance(value, timedelta):
+        return humanized_duration(int(value.total_seconds() * 1000))
+    else:
+        return str(value)
+
+def format_row(key: str, value) -> dict:
+    result = {}
+    result[key] = value
+    if "status" or "state" in key.lower():
+        result["style"] = get_color_by_state(value)
+    if key == "query_size" and isinstance(value, int):
+        result[key] = humanized_bytes(value)
+    else:
+        result[key] = format_value(value)
+    return result
+
+def show_dictionary_table(dict, format_fn:Callable|None=None):
+    table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
+    table.add_column("Field")
+    table.add_column("Value")
+    for key, value in dict.items():
+        if format_fn:
+            row = format_fn(key, value)
+            table.add_row(key, row.get(key), style=row.get("style"))
+        else:
+            table.add_row(key, value)
+    rich.print(table)
+
+#--------------------------------------------------
+# Custom help printer
+#--------------------------------------------------
+
+
+class RichGroup(click.Group):
+    def format_help(self, ctx: Context, formatter: HelpFormatter) -> None:
+        is_latest, current_ver, latest_ver = is_latest_cli_version()
+
+        global PROFILE
+        # @NOTE: I couldn't find a sane way to access the --profile option from here, so insane it is.
+        if "--profile" in sys.argv:
+            ix = sys.argv.index("--profile") + 1
+            if ix < len(sys.argv):
+                PROFILE = sys.argv[ix]
+
+        profile = get_config().profile
+        if profile == TOP_LEVEL_PROFILE_NAME:
+            profile = "[yellow bold]None[/yellow bold]" if not get_config().get("platform", "") else "[ROOT]"
+
+        sio = io.StringIO()
+        console = Console(file=sio, force_terminal=True)
+        divider(console)
+        console.print(f"[bold]Welcome to [green]RelationalAI[/bold] ({current_ver})!")
+        console.print()
+        if not is_latest:
+            console.print(f"[yellow]A new version of RelationalAI is available: {latest_ver}[/yellow] ")
+            console.print()
+        console.print("rai [magenta]\\[options][/magenta] [cyan]command[/cyan]")
+
+        console.print()
+        console.print(f"[magenta]--profile[/magenta][dim] - which config profile to use (current: [/dim][cyan]{profile}[/cyan][dim])")
+
+        unavailable_commands = []
+        groups = defaultdict(list)
+        for command in self.commands.keys():
+            if ":" in command:
+                group, _, _ = command.rpartition(":")
+                groups[group].append(command)
+            else:
+                groups[""].append(command)
+
+        console.print()
+        for command in groups[""]:
+            console.print(f"[cyan]{command}[/cyan][dim] - {self.commands[command].help}")
+
+        for group, commands in groups.items():
+            if not group:
+                continue
+
+            empty = True
+            for command in commands:
+                if command_available(self.commands[command]):
+                    if empty:
+                        empty = False
+                        console.print()
+
+                    console.print(f"[cyan]{command}[/cyan][dim] - {self.commands[command].help}")
+                else:
+                    unavailable_commands.append(command)
+
+        if unavailable_commands:
+            platform = get_config().get("platform", "")
+            console.print()
+            console.print(f"[yellow]Not available on {platform}[/yellow]")
+            console.print()
+            for command in unavailable_commands:
+                console.print(f"[dim yellow]{command} - {self.commands[command].help}")
+
+        divider(console)
+        formatter.write(sio.getvalue())
+        sio.close()
+
+def filter_profiles_by_platform(config:ConfigFile, platform:str):
+    filtered_config = {}
+    for profile, props in config.get_combined_profiles().items():
+        if profile == TOP_LEVEL_PROFILE_NAME:
+            continue
+        if props.get("platform") == platform or (
+            props.get("platform") is None
+            and platform == "azure"
+        ):
+            filtered_config[profile] = props
+    return filtered_config
+
+#--------------------------------------------------
+# Imports list
+#--------------------------------------------------
+
+def show_imports(imports, showId=False):
+    ensure_config()
+    if len(imports):
+        table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
+        cols = {"#": "index"}
+        if showId and len(imports) and "id" in imports[0]:
+            cols["ID"] = "id"
+        if len(imports) and "name" in imports[0]:
+            cols["Name"] = "name"
+        if len(imports) and "model" in imports[0]:
+            cols["Model"] = "model"
+        if len(imports) and "created" in imports[0]:
+            cols["Created"] = "created"
+        if len(imports) and "creator" in imports[0]:
+            cols["Creator"] = "creator"
+        if len(imports) and "batches" in imports[0]:
+            cols["Batches"] = "batches"
+        if len(imports) and "status" in imports[0]:
+            cols["Status"] = "status"
+        if len(imports) and "errors" in imports[0]:
+            cols["Errors"] = "errors"
+
+        for label in cols.keys():
+            table.add_column(label)
+
+        for index, imp in enumerate(imports):
+            imp["index"] = f"{index+1}"
+            style = get_color_by_state(imp["status"])
+            if imp["created"]:
+                imp["created"] = format_value(imp["created"])
+            table.add_row(*[imp.get(attr, None) for attr in cols.values()], style=style)
+        rich.print(table)
+    else:
+        rich.print("[yellow]No imports found")
+
+#--------------------------------------------------
+# Transactions
+#--------------------------------------------------
+
+def show_transactions(transactions, limit):
+    if len(transactions):
+        table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
+        table.add_column("#")
+        table.add_column("ID")
+        table.add_column("Schema")
+        table.add_column("Engine")
+        table.add_column("State")
+        table.add_column("Created")
+        table.add_column("Duration", justify="right")
+
+        added = 0
+        for i, txn in enumerate(transactions):
+            if added >= limit:
+                break
+
+            state = txn.get("state", "")
+            duration = txn.get("duration")
+            created_on = txn.get("created_on")
+
+            if isinstance(created_on, int):
+                created_on = datetime.fromtimestamp(created_on / 1000, tz=pytz.utc)
+            if duration is None:
+                duration = (datetime.now(created_on.tzinfo) - created_on).total_seconds() * 1000
+
+            table.add_row(
+                f"{i+1}",
+                txn.get("id"),
+                txn.get("database", ""),
+                txn.get("engine", ""),
+                state,
+                created_on.strftime("%Y-%m-%d %H:%M:%S"),
+                humanized_duration(int(duration)),
+                style=get_color_by_state(state)
+            )
+            added += 1
+        rich.print(table)
+    else:
+        rich.print("[yellow]No transactions found")
+
+def show_engines(engines):
+    if len(engines):
+        table = Table(show_header=True, border_style="dim", header_style="bold", box=rich_box.SIMPLE_HEAD)
+        table.add_column("#")
+        table.add_column("Name")
+        table.add_column("Size")
+        table.add_column("State")
+        for index, engine in enumerate(engines):
+            table.add_row(f"{index+1}", engine.get("name"), engine.get("size"), engine.get("state"))
+        rich.print(table)
+
+def exit_with_error(message:str):
+    rich.print(message, file=sys.stderr)
+    exit_with_divider(1)
+
+def exit_with_divider(exit_code:int=0):
+    divider()
+    sys.exit(exit_code)

v0/relationalai/tools/debugger.py

@@ -0,0 +1,183 @@
+import os
+import re
+from nicegui import ui
+import json
+
+from v0.relationalai.debugging import DEBUG_LOG_FILE
+
+
+last_mod_time = None
+current_json_objects = []
+active_ix = None
+active_item = None
+
+def set_item(ix):
+    global active_item, active_ix
+    if active_ix != ix:
+        active_ix = ix
+        active_item = current_json_objects[ix]
+        my_stuff.refresh()
+        details.refresh()
+
+def format_time(t):
+    if t > 1:
+        return f"{t:.1f}s"
+    elif t > 0.001:
+        return f"{t*1000:.1f}ms"
+    elif t > 0.0001:
+        return f"{t*1000:.2f}ms"
+    else:
+        return f"{t*1000000:.0f}us"
+
+def header(text):
+    return ui.label(text).style("font-size: 1.3em; font-weight: bold;")
+
+def replace_long_brace_contents(code_str):
+    def replacement(match):
+        string = match.group(0)
+        if len(string) > 300:
+            # Extract the first 50 and last 30 characters from the string within brackets
+            return '{' + string[1:51] + '...' + string[-31:-1] + '}'
+        else:
+            # If the string is not longer than 300 characters, return it unchanged
+            return string
+
+    # This regex matches sequences of characters wrapped in { }
+    brace_content_regex = r'\{[\s\S]*?\}'
+
+    # Use the sub method to replace the matched strings with the result of the replacement function
+    return re.sub(brace_content_regex, replacement, code_str)
+
+
+def code(c):
+    c = replace_long_brace_contents(c)
+    return ui.code(c).style("padding-right: 30px").classes("w-full")
+
+
+@ui.refreshable
+def details():
+    if active_item:
+        if active_item['event'] == "compilation":
+            header(f"{active_item['source']['file']}: {active_item['source']['line']}")
+            # with ui.row():
+            code(active_item["source"]["block"] or active_item["emitted"])
+            code(f"{active_item['emitted']}")
+            header("IR")
+            code(f"{active_item['task']}")
+            header("Rewritten")
+            code(f"{active_item['passes'][-1]['task']}")
+            header("Passes")
+            for p in active_item['passes']:
+                with ui.column().classes("w-full"):
+                    with ui.row():
+                        ui.label(p['name'])
+                        ui.label(f"({p['elapsed']*1000000:.0f} us)")
+                    code(p['task'])
+        elif active_item['event'] == "time":
+            ui.label(f"{active_item['type']} | {format_time(active_item['elapsed'])} | {active_item['results']['count']}")
+            vals = active_item['results']['values']
+            if len(vals):
+                keys = [k for k in vals[0].keys()]
+                columns = [{'name': k, 'label': k, 'field': k, "align": "left"} for k in keys]
+                ui.table(columns=columns, rows=vals)
+            if "code" in active_item:
+                header("Code")
+                code(active_item["code"])
+            # ui.label(f"{active_item['results']['values']}").style('white-space: pre;')
+
+@ui.refreshable
+def my_stuff():
+    total_time = 0
+    with ui.column():
+
+        for (ix, obj) in enumerate(current_json_objects):
+            me = None
+
+            if obj['event'] == "compilation":
+                me = code(obj["source"]["block"] or obj["emitted"])
+                total_time += obj["emit_time"]
+                for p in obj['passes']:
+                    total_time += p['elapsed']
+            elif obj['event'] == "time":
+                total_time += obj['elapsed']
+                with ui.row() as me:
+                    me.style("padding: 5px 10px; border: 1px solid #544; border-radius: 5px; background:#322;")
+                    ui.label(f"{obj['type']}")
+                    ui.label(format_time(obj['elapsed']))
+                    if obj.get('results'):
+                        ui.label(f"{obj['results']['count']}")
+            elif obj['event'] == "transaction_created":
+                with ui.row() as me:
+                    me.style("padding: 5px 10px; border: 1px solid #544; border-radius: 5px; background:#322;")
+                    ui.label("transaction")
+                    ui.label(obj["txn_id"])
+            elif obj.get('type'):
+                total_time += obj.get('elapsed', 0)
+                with ui.row() as me:
+                    me.style("padding: 5px 10px; border: 1px solid #544; border-radius: 5px; background:#322;")
+                    ui.label(f"{obj['type']}")
+                    ui.label(format_time(obj.get('elapsed', 0)))
+                    if obj.get('results'):
+                        ui.label(f"{obj['results']['count']}")
+            elif 'span' not in obj['event'] and obj['event'] != "profile_events":
+                with ui.row() as me:
+                    me.style("padding: 5px 10px; border: 1px solid #544; border-radius: 5px; background:#322;")
+                    ui.label(f"{obj['event']}")
+
+            if me:
+                me.classes("w-full")
+                me.on("click", lambda ix=ix: set_item(ix))
+                if ix != active_ix:
+                    me.style('opacity: 0.5')
+        with ui.row() as me:
+            me.style("padding: 5px 10px; border: 1px solid #445; border-radius: 5px; background:#223; opacity: 0.5;")
+            me.classes("w-full")
+            if total_time > 1:
+                ui.label(f"Total time: {total_time:.1f}s")
+            else:
+                ui.label(f"Total time: {total_time * 1000:.0f}ms")
+
+
+def poll():
+    global last_mod_time, active_item
+    global current_json_objects
+    # Check the last modification time of the file
+    try:
+        mod_time = os.path.getmtime(DEBUG_LOG_FILE)
+        if last_mod_time is None or mod_time > last_mod_time:
+            last_mod_time = mod_time
+            # File has changed, read and parse the new content
+            new_objects = []
+            with open(DEBUG_LOG_FILE, 'r') as file:
+                for line in file:
+                    try:
+                        # Parse each JSON object and add it to the list
+                        obj = json.loads(line)
+                        new_objects.append(obj)
+                    except json.JSONDecodeError as e:
+                        print(f"Error parsing JSON: {e}")
+            # Update the current JSON objects
+            current_json_objects = new_objects
+            if active_ix is not None and len(current_json_objects) > active_ix:
+                active_item = current_json_objects[active_ix]
+            # Refresh the UI with the new objects
+            my_stuff.refresh()
+            details.refresh()
+    except FileNotFoundError:
+        pass
+
+def main(host="0.0.0.0", port=8080):
+    ui.dark_mode().enable()
+    with ui.row():
+        with ui.column() as c:
+            c.style("cursor: pointer;")
+            my_stuff()
+        with ui.column() as c:
+            c.style("padding-left: 2em;")
+            details()
+
+    ui.timer(1, poll)
+    ui.run(reload=False, host=host, port=port)
+
+if __name__ == "__main__":
+    main()

v0/relationalai/tools/debugger_client.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+import asyncio
+import logging
+import threading
+import websockets
+import queue
+from v0.relationalai import Resources, debugging
+from v0.relationalai.clients.config import Config
+
+from v0.relationalai.clients.profile_polling import ProfilePollerHandler
+
+class DummyRecord(logging.LogRecord):
+    def __init__(self, msg):
+        self.msg = msg
+
+#------------------------------------------------------------------------------
+# Logging Handler
+#------------------------------------------------------------------------------
+
+class WebSocketLoggingHandler(logging.Handler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        from v0.relationalai.analysis.mechanistic import Mechanism
+        self.Mechanism = Mechanism
+        self.queue: queue.Queue[str] = queue.Queue()
+
+    def start(self, url: str, config: Config):
+
+
+        # Catch any spans that are the programs parent that were added before the debug client was started.
+        open_spans = []
+        existing_parent = debugging.TIP_SPAN
+        while existing_parent:
+            open_spans.append(existing_parent)
+            existing_parent = existing_parent.parent
+        open_spans.reverse()
+
+        for span in open_spans:
+            self.queue.put_nowait(self.format_raw({"event": "span_start", "span": span}))
+
+        self.thread = threading.Thread(target=self.connect, args=(url,))
+        self.thread.start()
+
+    def emit(self, record):
+        d = record.msg
+        if (isinstance(d, dict) and
+            d["event"] == "span_start" and
+            "task" in d["span"].attrs and
+            "mech" not in d["span"].attrs):
+            d["span"].attrs["mech"] = self.Mechanism(d["span"].attrs["task"])
+        log_entry = self.format(record)
+        self.queue.put_nowait(log_entry)
+
+    def format_raw(self, msg):
+        return self.format(DummyRecord(msg))
+
+    def connect(self, url: str):
+        try:
+            try:
+                loop = asyncio.get_running_loop()
+                loop.run_until_complete(self.connect_async(url))
+            except RuntimeError:
+                asyncio.run(self.connect_async(url))
+        except (ConnectionRefusedError, websockets.WebSocketException, OSError, asyncio.TimeoutError):
+            pass
+
+    async def connect_async(self, url: str):
+        async with websockets.connect(url) as ws:
+            print(f"Connected to debugger at {url}...")
+            alive = True
+            while True:
+                if not threading.main_thread().is_alive():
+                    alive = False
+                    debugging.span_flush()
+
+                try:
+                    log_entry = self.queue.get(timeout=1)
+                    await ws.send(log_entry)
+                except queue.Empty:
+                    if not alive:
+                        break
+                except websockets.ConnectionClosedError:
+                    break
+
+#------------------------------------------------------------------------------
+# Enable debugging
+#------------------------------------------------------------------------------
+
+already_debugging = False
+def start_debugger_session(config: Config, host: str|None = None, port: int|None = None):
+    global already_debugging
+    if already_debugging:
+        return
+
+    if host is None:
+        host = "0.0.0.0"
+
+    if port is None:
+        port = 8080
+
+    ws_handler = WebSocketLoggingHandler()
+    ws_handler.setFormatter(debugging.JsonFormatter())
+    debugging.logger.addHandler(ws_handler)
+    ws_handler.start(f"ws://{host}:{port}/ws/program", config)
+    
+    resources = Resources(config=config)
+    debugging.logger.addHandler(ProfilePollerHandler(resources))
+
+    already_debugging = True