apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/caching/stores/sqlite.py

@@ -0,0 +1,282 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pathlib
+import sqlite3
+import threading
+
+from hamilton.caching.cache_key import decode_key
+from hamilton.caching.stores.base import MetadataStore
+
+
+class SQLiteMetadataStore(MetadataStore):
+    def __init__(
+        self,
+        path: str,
+        connection_kwargs: dict | None = None,
+    ) -> None:
+        self._directory = pathlib.Path(path).resolve()
+        self._directory.mkdir(parents=True, exist_ok=True)
+        self._path = self._directory.joinpath("metadata_store").with_suffix(".db")
+        self.connection_kwargs: dict = connection_kwargs or {}
+
+        self._thread_local = threading.local()
+
+        # creating tables at `__init__` prevents other methods from encountering
+        # `sqlite3.OperationalError` because tables are missing.
+        self._create_tables_if_not_exists()
+
+    def __getstate__(self) -> dict:
+        """Serialized `__init__` arguments required to initialize the
+        MetadataStore in a new thread or process.
+        """
+        state = {}
+        # NOTE kwarg `path` is not equivalent to `self._path`
+        state["path"] = self._directory
+        state["connection_kwargs"] = self.connection_kwargs
+        return state
+
+    def _get_connection(self) -> sqlite3.Connection:
+        if not hasattr(self._thread_local, "connection"):
+            self._thread_local.connection = sqlite3.connect(
+                str(self._path), check_same_thread=False, **self.connection_kwargs
+            )
+        return self._thread_local.connection
+
+    def _close_connection(self) -> None:
+        if hasattr(self._thread_local, "connection"):
+            self._thread_local.connection.close()
+            del self._thread_local.connection
+
+    @property
+    def connection(self) -> sqlite3.Connection:
+        """Connection to the SQLite database."""
+        return self._get_connection()
+
+    def __del__(self):
+        """Close the SQLite connection when the object is deleted"""
+        self._close_connection()
+
+    def _create_tables_if_not_exists(self) -> None:
+        """Create the tables necessary for the cache:
+
+        run_ids: queue of run_ids, ordered by start time.
+        history: queue of executed node; allows to query "latest" execution of a node
+        cache_metadata: information to determine if a node needs to be computed or not
+
+        In the table ``cache_metadata``, the ``cache_key`` is unique whereas
+        ``history`` allows duplicate.
+        """
+        cur = self.connection.cursor()
+
+        cur.execute(
+            """\
+            CREATE TABLE IF NOT EXISTS run_ids (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                run_id TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+            )
+            """
+        )
+        cur.execute(
+            """\
+            CREATE TABLE IF NOT EXISTS history (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                cache_key TEXT,
+                run_id TEXT,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+
+                FOREIGN KEY (cache_key) REFERENCES cache_metadata(cache_key)
+            );
+            """
+        )
+        cur.execute(
+            """\
+            CREATE TABLE IF NOT EXISTS cache_metadata (
+                cache_key TEXT PRIMARY KEY,
+                data_version TEXT NOT NULL,
+                node_name TEXT NOT NULL,
+                code_version TEXT NOT NULL,
+                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+
+                FOREIGN KEY (cache_key) REFERENCES history(cache_key)
+            );
+            """
+        )
+        self.connection.commit()
+
+    def initialize(self, run_id) -> None:
+        """Call initialize when starting a run. This will create database tables
+        if necessary.
+        """
+        cur = self.connection.cursor()
+        cur.execute("INSERT INTO run_ids (run_id) VALUES (?)", (run_id,))
+        self.connection.commit()
+
+    def __len__(self) -> int:
+        """Number of entries in cache_metadata"""
+        cur = self.connection.cursor()
+        cur.execute("SELECT COUNT(*) FROM cache_metadata")
+        return cur.fetchone()[0]
+
+    def set(
+        self,
+        *,
+        cache_key: str,
+        data_version: str,
+        run_id: str,
+        node_name: str = None,
+        code_version: str = None,
+        **kwargs,
+    ) -> None:
+        cur = self.connection.cursor()
+
+        # if the caller of ``.set()`` directly provides the ``node_name`` and ``code_version``,
+        # we can skip the decoding step.
+        if (node_name is None) or (code_version is None):
+            try:
+                decoded_key = decode_key(cache_key)
+            except BaseException as e:
+                raise ValueError(
+                    f"Failed decoding the cache_key: {cache_key}.\n",
+                    "The `cache_key` must be created by `hamilton.caching.cache_key.create_cache_key()` ",
+                    "if `node_name` and `code_version` are not provided.",
+                ) from e
+
+            node_name = decoded_key["node_name"]
+            code_version = decoded_key["code_version"]
+
+        cur.execute("INSERT INTO history (cache_key, run_id) VALUES (?, ?)", (cache_key, run_id))
+        cur.execute(
+            """\
+            INSERT OR IGNORE INTO cache_metadata (
+                cache_key, node_name, code_version, data_version
+            ) VALUES (?, ?, ?, ?)
+            """,
+            (cache_key, node_name, code_version, data_version),
+        )
+
+        self.connection.commit()
+
+    def get(self, cache_key: str) -> str | None:
+        cur = self.connection.cursor()
+        cur.execute(
+            """\
+            SELECT data_version
+            FROM cache_metadata
+            WHERE cache_key = ?
+            """,
+            (cache_key,),
+        )
+        result = cur.fetchone()
+
+        if result is None:
+            data_version = None
+        else:
+            data_version = result[0]
+
+        return data_version
+
+    def delete(self, cache_key: str) -> None:
+        """Delete metadata associated with ``cache_key``."""
+        cur = self.connection.cursor()
+        cur.execute("DELETE FROM cache_metadata WHERE cache_key = ?", (cache_key,))
+        self.connection.commit()
+
+    def delete_all(self) -> None:
+        """Delete all existing tables from the database"""
+        cur = self.connection.cursor()
+
+        for table_name in ["run_ids", "history", "cache_metadata"]:
+            cur.execute(f"DROP TABLE IF EXISTS {table_name};")
+
+        self.connection.commit()
+
+    def exists(self, cache_key: str) -> bool:
+        """boolean check if a ``data_version`` is found for ``cache_key``
+        If True, ``.get()`` should successfully retrieve the ``data_version``.
+        """
+        cur = self.connection.cursor()
+        cur.execute("SELECT cache_key FROM cache_metadata WHERE cache_key = ?", (cache_key,))
+        result = cur.fetchone()
+
+        return result is not None
+
+    def get_run_ids(self) -> list[str]:
+        """Return a list of run ids, sorted from oldest to newest start time."""
+        cur = self.connection.cursor()
+        cur.execute("SELECT run_id FROM run_ids ORDER BY id")
+        result = cur.fetchall()
+
+        return [r[0] for r in result]
+
+    def _run_exists(self, run_id: str) -> bool:
+        """Returns True if a run was initialized with ``run_id``, even
+        if the run recorded no node executions.
+        """
+        cur = self.connection.cursor()
+        cur.execute(
+            """\
+            SELECT EXISTS(
+                SELECT 1
+                FROM run_ids
+                WHERE run_id = ?
+            )
+            """,
+            (run_id,),
+        )
+        result = cur.fetchone()
+        # SELECT EXISTS returns 1 for True, i.e., `run_id` is found
+        return result[0] == 1
+
+    def get_run(self, run_id: str) -> list[dict]:
+        """Return a list of node metadata associated with a run.
+
+        :param run_id: ID of the run to retrieve
+        :return: List of node metadata which includes ``cache_key``, ``data_version``,
+            ``node_name``, and ``code_version``. The list can be empty if a run was initialized
+            but no nodes were executed.
+
+        :raises IndexError: if the ``run_id`` is not found in metadata store.
+        """
+        cur = self.connection.cursor()
+        if self._run_exists(run_id) is False:
+            raise IndexError(f"`run_id` not found in table `run_ids`: {run_id}")
+
+        cur.execute(
+            """\
+            SELECT
+                cache_metadata.cache_key,
+                cache_metadata.data_version,
+                cache_metadata.node_name,
+                cache_metadata.code_version
+            FROM history
+            JOIN cache_metadata ON history.cache_key = cache_metadata.cache_key
+            WHERE history.run_id = ?
+            """,
+            (run_id,),
+        )
+        results = cur.fetchall()
+        return [
+            dict(
+                cache_key=cache_key,
+                data_version=data_version,
+                node_name=node_name,
+                code_version=code_version,
+            )
+            for cache_key, data_version, node_name, code_version in results
+        ]

hamilton/caching/stores/utils.py

@@ -0,0 +1,40 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import pathlib
+
+
+def get_directory_size(directory: str) -> float:
+    """Get the size of the content of a directory in bytes."""
+    total_size = 0
+    for p in pathlib.Path(directory).rglob("*"):
+        if p.is_file():
+            total_size += p.stat().st_size
+
+    return total_size
+
+
+def readable_bytes_size(n_bytes: float) -> str:
+    """Convert a number of bytes to a human-readable unit."""
+    labels = ["B", "KB", "MB", "GB", "TB"]
+    exponent = 0
+
+    while n_bytes > 1024.0:
+        n_bytes /= 1024.0
+        exponent += 1
+
+    return f"{n_bytes:.2f} {labels[exponent]}"

hamilton/cli/__init__.py

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

hamilton/cli/__main__.py

@@ -0,0 +1,328 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import dataclasses
+import json
+import logging
+import os
+import warnings
+from collections.abc import Callable
+from pathlib import Path
+from pprint import pprint
+from typing import Annotated, Any
+
+import typer
+
+# silence UserWarning: 'PYARROW_IGNORE_TIMEZONE'
+with warnings.catch_warnings():
+    warnings.filterwarnings("ignore", category=UserWarning)
+    from hamilton import driver
+
+from hamilton.cli import commands
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class Response:
+    command: str
+    success: bool
+    message: Any
+
+
+class CliState:
+    verbose: bool | None = None
+    json_out: bool | None = None
+    dr: driver.Driver | None = None
+    name: str | None = None
+
+
+cli = typer.Typer(rich_markup_mode="rich")
+state = CliState()
+
+MODULES_ANNOTATIONS = Annotated[
+    list[Path],
+    typer.Argument(
+        help="Paths to Hamilton modules",
+        exists=True,
+        dir_okay=False,
+        readable=True,
+        resolve_path=True,
+    ),
+]
+
+NAME_ANNOTATIONS = Annotated[
+    str | None,
+    typer.Option("--name", "-n", help="Name of the dataflow. Default: Derived from MODULES."),
+]
+
+CONTEXT_ANNOTATIONS = Annotated[
+    Path | None,
+    typer.Option(
+        "--context",
+        "-ctx",
+        help="Path to Driver context file [.json, .py]",
+        exists=True,
+        dir_okay=False,
+        readable=True,
+        resolve_path=True,
+    ),
+]
+
+VIZ_OUTPUT_ANNOTATIONS = Annotated[
+    Path,
+    typer.Option(
+        "--output",
+        "-o",
+        help="Output path of visualization. If path is a directory, use NAME for file name.",
+        dir_okay=True,
+        writable=True,
+        resolve_path=True,
+    ),
+]
+
+
+# TODO add `experiments` for `hamilton.plugins.h_experiments`
+# TODO add `dataflows` submenu to manage locally installed dataflows
+# TODO add `init` to load project template
+# callback() creates entrypoint for `hamilton` without command
+@cli.callback()
+def main(
+    ctx: typer.Context,
+    verbose: Annotated[
+        bool,
+        typer.Option(
+            help="Output all intermediary commands",
+            rich_help_panel="Output format",
+        ),
+    ] = False,
+    json_out: Annotated[
+        bool,
+        typer.Option(
+            help="Output JSON for programmatic use (e.g., CI)",
+            rich_help_panel="Output format",
+        ),
+    ] = False,
+):
+    """Hamilton CLI"""
+    state.verbose = verbose
+    state.json_out = json_out
+    logger.debug(f"verbose set to {verbose}")
+    logger.debug(f"json_out set to {json_out}")
+
+
+def _try_command(cmd: Callable, **cmd_kwargs) -> Any:
+    """Try a command and raise errors to Typer and exit CLI"""
+    cmd_name = cmd.__name__
+    try:
+        logger.debug(f"calling commands.{cmd_name}")
+        result = cmd(**cmd_kwargs)
+    except Exception as e:
+        response = Response(
+            command=cmd_name, success=False, message={"error": str(type(e)), "details": str(e)}
+        )
+        logger.error(dataclasses.asdict(response))
+        raise typer.Exit(code=1) from e
+
+    return result
+
+
+def _response_handler(ctx: typer.Context, response: Response) -> None:
+    """Handle how to display response"""
+    if (ctx.info_name == response.command) or state.verbose:
+        if state.json_out is True:
+            print(json.dumps(dataclasses.asdict(response)))
+        else:
+            pprint(response.message)
+
+
+@cli.command()
+def build(
+    ctx: typer.Context,
+    modules: MODULES_ANNOTATIONS,
+    name: NAME_ANNOTATIONS = None,
+    context_path: CONTEXT_ANNOTATIONS = None,
+):
+    """Build a single Driver with MODULES"""
+    state.dr = _try_command(cmd=commands.build, modules=modules, context_path=context_path)
+
+    if name:
+        state.name = name
+    else:
+        state.name = "_".join([str(Path(m).stem) for m in modules])[:40]
+
+    _response_handler(
+        ctx=ctx,
+        response=Response(
+            command="build",
+            success=True,
+            message={"modules": [p.stem for p in modules]},
+        ),
+    )
+
+
+@cli.command()
+def diff(
+    ctx: typer.Context,
+    modules: MODULES_ANNOTATIONS,
+    name: NAME_ANNOTATIONS = None,
+    context_path: CONTEXT_ANNOTATIONS = None,
+    output_file_path: VIZ_OUTPUT_ANNOTATIONS = Path("./"),
+    git_reference: Annotated[
+        str,
+        typer.Option(
+            help="[link=https://git-scm.com/book/en/v2/Git-Internals-Git-References]git reference[/link] to compare to"
+        ),
+    ] = "HEAD",
+    view: Annotated[
+        bool,
+        typer.Option(
+            "--view",
+            "-v",
+            help="Generate a dataflow diff visualization",
+        ),
+    ] = False,
+):
+    """Diff between the current MODULES and their specified GIT_REFERENCE"""
+    if state.dr is None:
+        ctx.invoke(version, ctx=ctx, modules=modules, name=name, context_path=context_path)
+
+    # default value isn't set to None to let Typer properly resolve the path
+    # then, we change the file name
+    if output_file_path.is_dir():
+        output_file_path.mkdir(parents=True, exist_ok=True)
+        output_file_path = output_file_path.joinpath(f"diff_{state.name}.png")
+
+    diff = _try_command(
+        cmd=commands.diff,
+        current_dr=state.dr,
+        modules=modules,
+        git_reference=git_reference,
+        view=view,
+        output_file_path=output_file_path,
+        context_path=context_path,
+    )
+    _response_handler(
+        ctx=ctx,
+        response=Response(
+            command="diff",
+            success=True,
+            message=diff,
+        ),
+    )
+
+
+@cli.command()
+def validate(
+    ctx: typer.Context,
+    modules: MODULES_ANNOTATIONS,
+    context_path: CONTEXT_ANNOTATIONS,
+    name: NAME_ANNOTATIONS = None,
+):
+    """Validate DATAFLOW execution for the given CONTEXT"""
+    if state.dr is None:
+        ctx.invoke(build, ctx=ctx, modules=modules, name=name, context_path=context_path)
+
+    validated_context = _try_command(commands.validate, dr=state.dr, context_path=context_path)
+    _response_handler(
+        ctx=ctx,
+        response=Response(
+            command="validate",
+            success=True,
+            message=validated_context,
+        ),
+    )
+
+
+@cli.command()
+def version(
+    ctx: typer.Context,
+    modules: MODULES_ANNOTATIONS,
+    name: NAME_ANNOTATIONS = None,
+    context_path: CONTEXT_ANNOTATIONS = None,
+):
+    """Version NODES and DATAFLOW from dataflow with MODULES"""
+    if state.dr is None:
+        ctx.invoke(build, ctx=ctx, modules=modules, name=name, context_path=context_path)
+
+    dataflow_version = _try_command(cmd=commands.version, dr=state.dr)
+    _response_handler(
+        ctx=ctx,
+        response=Response(
+            command="version",
+            success=True,
+            message=dataflow_version,
+        ),
+    )
+
+
+@cli.command()
+def view(
+    ctx: typer.Context,
+    modules: MODULES_ANNOTATIONS,
+    name: NAME_ANNOTATIONS = None,
+    context_path: CONTEXT_ANNOTATIONS = None,
+    output_file_path: VIZ_OUTPUT_ANNOTATIONS = Path("./"),
+):
+    """Build and visualize dataflow with MODULES"""
+    if state.dr is None:
+        ctx.invoke(build, ctx=ctx, modules=modules, name=name, context_path=context_path)
+
+    if output_file_path.is_dir():
+        output_file_path.mkdir(parents=True, exist_ok=True)
+        output_file_path = output_file_path.joinpath(f"dag_{state.name}.png")
+
+    _try_command(cmd=commands.view, dr=state.dr, output_file_path=output_file_path)
+    _response_handler(
+        ctx=ctx,
+        response=Response(command="view", success=True, message={"path": str(output_file_path)}),
+    )
+
+
+@cli.command()
+def ui(
+    ctx: typer.Context,
+    port: int = 8241,
+    base_dir: str = os.path.join(Path.home(), ".hamilton", "db"),
+    no_migration: bool = False,
+    no_open: bool = False,
+    settings_file: str = "mini",
+    config_file: str | None = None,
+):
+    """Runs the Hamilton UI on sqllite in port 8241"""
+    try:
+        from hamilton_ui import commands
+    except ImportError as e:
+        logger.error(
+            "hamilton[ui] not installed -- you have to install this to run the UI. "
+            'Run `pip install "sf-hamilton[ui]"` to install and get started with the UI!'
+        )
+        raise typer.Exit(code=1) from e
+
+    ctx.invoke(
+        commands.run,
+        port=port,
+        base_dir=base_dir,
+        no_migration=no_migration,
+        no_open=no_open,
+        settings_file=settings_file,
+        config_file=config_file,
+    )
+
+
+if __name__ == "__main__":
+    cli()