heurist-api 0.1.2__py3-none-any.whl

heurist/cli/__main__.py

@@ -0,0 +1,227 @@
+"""
+CLI commands for extracting, transforming, and loading remote Heurist data.
+"""
+
+import importlib.metadata
+
+import click
+from heurist import PACKAGE_NAME
+from heurist.api.credentials import CredentialHandler
+from heurist.api.exceptions import MissingParameterException
+from heurist.cli.load import load_command
+from heurist.cli.records import rty_command
+from heurist.cli.schema import schema_command
+from heurist.utils.constants import DEFAULT_RECORD_GROUPS
+from rich.console import Console
+
+# This name must match the package name ('name' kwarg) in the TOML file.
+__identifier__ = importlib.metadata.version(PACKAGE_NAME)
+
+
+# =========================== #
+#     Main cli group
+# =========================== #
+@click.group(help="Group CLI command for connecting to the Heurist DB")
+@click.version_option(__identifier__)
+@click.option(
+    "-d",
+    "--database",
+    type=click.STRING,
+    help="Name of the Heurist database",
+)
+@click.option(
+    "-l",
+    "--login",
+    type=click.STRING,
+    help="Login name for the database user",
+)
+@click.option(
+    "-p",
+    "--password",
+    type=click.STRING,
+    help="Password for the database user",
+)
+@click.option(
+    "--debugging",
+    required=False,
+    default=False,
+    is_flag=True,
+    help="Whether to run in debug mode, default false.",
+)
+@click.pass_context
+def cli(ctx, database, login, password, debugging):
+    ctx.ensure_object(dict)
+    ctx.obj["DEBUGGING"] = debugging
+    try:
+        ctx.obj["CREDENTIALS"] = CredentialHandler(
+            database_name=database,
+            login=login,
+            password=password,
+        )
+    except MissingParameterException:
+        c = Console()
+        c.print(
+            "Login informaiton is missing."
+            "Please provide your credentials when prompted."
+            "\nTo quit, press Ctrl+C then Enter."
+        )
+        _database = click.prompt("Heurist database name")
+        _login = click.prompt("Heurist user login")
+        _password = click.prompt("Heurist login password")
+        c.print("Retrying the connection...")
+        ctx.obj["CREDENTIALS"] = CredentialHandler(
+            database_name=_database,
+            login=_login,
+            password=_password,
+        )
+        c.print("Success!", style="green")
+
+
+# =========================== #
+#     'record' command
+# =========================== #
+@cli.command("record", help="Get a JSON export of a certain record type.")
+@click.option(
+    "-t",
+    "--record-type",
+    help="The ID fo the record type",
+    type=click.INT,
+    required=True,
+)
+@click.option(
+    "-o",
+    "--outfile",
+    help="JSON file path.",
+    type=click.Path(file_okay=True, writable=True),
+    required=False,
+)
+@click.pass_obj
+def records(ctx, record_type, outfile):
+    credentials = ctx["CREDENTIALS"]
+    rty_command(credentials, record_type, outfile)
+
+
+# =========================== #
+#     'schema' command
+# =========================== #
+@cli.command(
+    "schema",
+    help="Generate documentation about the database schema.",
+)
+@click.option(
+    "-t",
+    "--output-type",
+    required=True,
+    type=click.Choice(["csv", "json"], case_sensitive=False),
+    help="Data format in which the schema will be described. \
+    csv = 1 CSV file for each record type. json = 1 file that \
+    lists all records together",
+)
+@click.option(
+    "-r",
+    "--record-group",
+    required=False,
+    type=click.STRING,
+    multiple=True,
+    default=["My record types"],
+    show_default=True,
+    help="Group name of the record types to be described. \
+        Can be declared multiple times for multiple groups.",
+)
+@click.option(
+    "-o",
+    "--outdir",
+    required=False,
+    type=click.Path(file_okay=False, dir_okay=True),
+    help="Path to the directory in which the files will be written. \
+        Defaults to name of the database + '_schema'.",
+)
+@click.pass_obj
+def doc(ctx, record_group, outdir, output_type):
+    # Get context variables
+    credentials = ctx["CREDENTIALS"]
+    debugging = ctx["DEBUGGING"]
+
+    # Run the doc command
+    schema_command(
+        credentials=credentials,
+        record_group=record_group,
+        outdir=outdir,
+        output_type=output_type,
+        debugging=debugging,
+    )
+
+
+# =========================== #
+#     'download' command
+# =========================== #
+@cli.command(
+    "download",
+    help="Export data of records of 1 or more record group types.",
+)
+@click.option(
+    "-r",
+    "--record-group",
+    required=False,
+    type=click.STRING,
+    multiple=True,
+    default=DEFAULT_RECORD_GROUPS,
+    help="Record group of the entities whose data is exported. \
+        Default: 'My record types'.",
+)
+@click.option(
+    "-u",
+    "--user",
+    required=False,
+    type=click.INT,
+    multiple=True,
+    help="User or users who created the records to be exported. \
+        Default: all users' records.",
+)
+@click.option(
+    "-f",
+    "--filepath",
+    required=True,
+    type=click.Path(
+        file_okay=True,
+        dir_okay=False,
+    ),
+    help="Path to the DuckDB database file in which the data will be written.",
+)
+@click.option(
+    "-o",
+    "--outdir",
+    required=False,
+    type=click.Path(
+        file_okay=False,
+        dir_okay=True,
+    ),
+    help="Directory in which CSV files of the dumped tabular data \
+        will be written.",
+)
+@click.pass_obj
+def load(ctx, filepath, record_group, user, outdir):
+    # Get context variable
+    credentials = ctx["CREDENTIALS"]
+    testing = ctx["DEBUGGING"]
+
+    # Run the dump command
+    if not testing:
+        load_command(
+            credentials=credentials,
+            duckdb_database_connection_path=filepath,
+            record_group=record_group,
+            user=user,
+            outdir=outdir,
+        )
+    else:
+        print(
+            "\nCannot run 'dump' command in debugging mode.\
+            \nClient must connect to a remote Heurist database.\n"
+        )
+        print("Exiting.")
+        exit()
+
+
+if __name__ == "__main__":
+    cli()

heurist/cli/load.py

@@ -0,0 +1,55 @@
+"""
+CLI command for downloading all requested record types' data.
+"""
+
+from pathlib import Path
+
+import duckdb
+from heurist.api.connection import HeuristAPIConnection
+from heurist.api.credentials import CredentialHandler
+from heurist.utils.constants import DEFAULT_RECORD_GROUPS
+from heurist.workflows import extract_transform_load
+
+
+def load_command(
+    credentials: CredentialHandler,
+    duckdb_database_connection_path: Path | str,
+    record_group: tuple = DEFAULT_RECORD_GROUPS,
+    user: tuple = (),
+    outdir: Path | None = None,
+):
+    # Run the ETL process
+    if isinstance(duckdb_database_connection_path, Path):
+        duckdb_database_connection_path = str(duckdb_database_connection_path)
+    with (
+        duckdb.connect(duckdb_database_connection_path) as conn,
+        HeuristAPIConnection(
+            db=credentials.get_database(),
+            login=credentials.get_login(),
+            password=credentials.get_password(),
+        ) as client,
+    ):
+        extract_transform_load(
+            client=client,
+            duckdb_connection=conn,
+            record_group_names=record_group,
+            user=user,
+        )
+
+    # Show the results of the created DuckDB database
+    with duckdb.connect(duckdb_database_connection_path) as new_conn:
+        tables = new_conn.sql("show tables;")
+        print("\nCreated the following tables")
+        print(tables)
+
+        # If writing to CSV files, write only tables of record types
+        if outdir:
+            outdir = Path(outdir)
+            outdir.mkdir(exist_ok=True)
+            for tup in tables.fetchall():
+                table_name = tup[0]
+                # Skip the schema tables
+                if table_name in ["rtg", "rst", "rty", "dty", "trm"]:
+                    continue
+                fp = outdir.joinpath(f"{table_name}.csv")
+                new_conn.table(table_name).sort("H-ID").write_csv(str(fp))

heurist/cli/records.py

@@ -0,0 +1,49 @@
+"""
+CLI command for downloading one requested record type's data.
+"""
+
+import json
+from pathlib import Path
+
+from heurist.api.connection import HeuristAPIConnection
+from heurist.api.credentials import CredentialHandler
+from rich.progress import (
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+
+
+def rty_command(
+    credentials: CredentialHandler,
+    rty: int,
+    outfile: Path | str | None,
+):
+    with (
+        Progress(
+            TextColumn("{task.description}"),
+            SpinnerColumn(),
+            TimeElapsedColumn(),
+        ) as p,
+        HeuristAPIConnection(
+            db=credentials.get_database(),
+            login=credentials.get_login(),
+            password=credentials.get_password(),
+        ) as client,
+    ):
+        _ = p.add_task(f"Get Records of type {rty}", total=1)
+        records = client.get_records(rty)
+    # If no records of this type have been entered, stop.
+    if len(records) == 0:
+        print("No records of this type have been entered.\nExiting program...")
+        exit()
+
+    # Else, write the records to a JSON file.
+    if not outfile:
+        outfile = f"RTY_{rty}.json"
+    if not isinstance(outfile, Path):
+        outfile = Path(outfile)
+    print(f"Writing results to: {outfile}")
+    with open(outfile, "w") as f:
+        json.dump(records, f, indent=4)

heurist/cli/schema.py

@@ -0,0 +1,94 @@
+"""
+CLI command for downloading details about a Heurist database schema.
+"""
+
+from pathlib import Path
+
+from heurist.api.connection import HeuristAPIConnection
+from heurist.api.credentials import CredentialHandler
+from heurist.database import TransformedDatabase
+from heurist.schema import output_csv, output_json
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+
+
+def get_database_schema(
+    record_groups: list,
+    credentials: CredentialHandler,
+    debugging: bool,
+) -> TransformedDatabase:
+    # If testing, load the mock database XML schema
+    if debugging:
+        from mock_data import DB_STRUCTURE_XML
+
+        db = TransformedDatabase(
+            hml_xml=DB_STRUCTURE_XML, record_type_groups=record_groups
+        )
+    # If not testing, request the database XML schema from the server
+    else:
+        with (
+            Progress(
+                TextColumn("{task.description}"),
+                SpinnerColumn(),
+                TimeElapsedColumn(),
+            ) as p,
+            HeuristAPIConnection(
+                db=credentials.get_database(),
+                login=credentials.get_login(),
+                password=credentials.get_password(),
+            ) as client,
+        ):
+            _ = p.add_task("Downloading schemas")
+            xml = client.get_structure()
+            db = TransformedDatabase(
+                hml_xml=xml,
+                record_type_groups=record_groups,
+            )
+    return db
+
+
+def schema_command(
+    credentials: CredentialHandler,
+    record_group: list,
+    outdir: str,
+    output_type: str,
+    debugging: bool = False,
+):
+    # Set up the output directory
+    if not outdir:
+        outdir = f"{credentials.get_database()}_schema"
+    DIR = Path(outdir)
+    DIR.mkdir(exist_ok=True)
+
+    # Get the database schema
+    db = get_database_schema(
+        record_groups=record_group,
+        credentials=credentials,
+        debugging=debugging,
+    )
+
+    # Describe each targeted record type
+    record_type_ids = list(db.pydantic_models.keys())
+    with Progress(
+        TextColumn("{task.description}"), BarColumn(), MofNCompleteColumn()
+    ) as p:
+        descriptions = []
+        t = p.add_task("Describing record types", total=len(record_type_ids))
+        for id in record_type_ids:
+            rel = db.describe_record_schema(rty_ID=id)
+            descriptions.append(rel)
+            p.advance(t)
+
+    # Output the descriptions according to the desired data format
+    if output_type == "csv":
+        output_csv(dir=DIR, descriptions=descriptions)
+
+    elif output_type == "json":
+        outfile = DIR.joinpath("recordTypes.json")
+        output_json(descriptions=descriptions, fp=outfile)

heurist/database/__init__.py

@@ -0,0 +1,3 @@
+from heurist.database.database import TransformedDatabase
+
+TransformedDatabase

heurist/database/basedb.py

@@ -0,0 +1,125 @@
+import duckdb
+import polars as pl
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from heurist.models.structural.hml_structure import HMLStructure
+from heurist.sql import RECORD_TYPE_SCHEMA
+from pydantic_xml import BaseXmlModel
+
+
+class HeuristDatabase:
+    """Base class for loading the original Heurist database structure."""
+
+    BASE_TABLES = [
+        ("rtg", "RecTypeGroups"),
+        ("rst", "RecStructure"),
+        ("rty", "RecTypes"),
+        ("dty", "DetailTypes"),
+        ("trm", "Terms"),
+    ]
+
+    def __init__(
+        self,
+        hml_xml: bytes,
+        conn: DuckDBPyConnection | None = None,
+        db: str = ":memory:",
+    ) -> None:
+        """
+        Create a DuckDB database connection and populate the DuckDB database with the
+        5 base tables that comprise the Heurist database structure.
+
+        Args:
+            hml_xml (bytes): Heurist database structure exported in XML format.
+            conn (DuckDBPyConnection | None, optional): A DuckDB database connection. \
+                Defaults to None.
+            db (str, optional): Path to the DuckDB database. Defaults to ":memory:".
+        """
+
+        hml_xml = self.trim_xml_bytes(xml=hml_xml)
+        if not conn:
+            conn = duckdb.connect(db)
+        self.conn = conn
+        # Load the Heurist database structure XML into a nested Pydantic data model
+        self.hml = HMLStructure.from_xml(hml_xml)
+
+        # Create generic tables
+        for t in self.BASE_TABLES:
+            name = t[0]
+            pydantic_model = getattr(getattr(self.hml, t[1]), t[0])
+            self.create(name, pydantic_model)
+
+    @classmethod
+    def trim_xml_bytes(cls, xml: bytes) -> bytes:
+        """
+        Remove any extra whitespace from a potentially malformatted XML.
+
+        Args:
+            xml (bytes): Heurist database structure exported XML format.
+
+        Returns:
+            bytes: Validated Heurist database structure in XML format.
+        """
+
+        return xml.decode("utf-8").strip().encode("utf-8")
+
+    def delete_existing_table(self, table_name: str) -> None:
+        """
+        If the table already exists in the DuckDB database, drop it.
+
+        Args:
+            table_name (str): Name of the table to potentially drop.
+        """
+
+        if self.conn.sql(
+            f"""
+SELECT *
+FROM duckdb_tables()
+WHERE table_name like '{table_name}'
+"""
+        ):
+            self.conn.sql("DROP TABLE {}".format(table_name))
+
+    def create(self, name: str, model: BaseXmlModel) -> None:
+        """Create an empty table in the DuckDB database connection
+        based on a Pydantic model.
+
+        Examples:
+            >>> # Set up the database class and parse a table model.
+            >>> from mock_data import DB_STRUCTURE_XML
+            >>> db = HeuristDatabase(hml_xml=DB_STRUCTURE_XML)
+            >>> model = db.hml.RecTypeGroups.rtg
+            >>>
+            >>> # Create a table for the Record Type Group (rtg) table model.
+            >>> db.create(name="rtg", model=model)
+            >>> shape = db.conn.table("rtg").fetchall()
+            >>> # The Record Type Group (rtg) table should have 11 columns.
+            >>> len(shape)
+            11
+
+        Args:
+            model (BaseXmlModel): A Pydantic XML model.
+        """
+
+        self.delete_existing_table(name)
+
+        # Convert the model to a dataframe and register it for duckdb
+        df = pl.DataFrame(model)
+        assert df.shape[0] > 1
+
+        # Create table from model dataframe
+        sql = "CREATE TABLE {} AS FROM df".format(name)
+        self.conn.sql(sql)
+
+    def describe_record_schema(self, rty_ID: int) -> DuckDBPyRelation:
+        """Join the tables 'dty' (detail), 'rst' (record structure), 'rty' (record type)
+        to get all the relevant information for a specific record type, plus add the
+        label and description of the section / separator associated with each detail
+        (if any).
+
+        Args:
+            rty_ID (int): ID of the targeted record type.
+
+        Returns:
+            DuckDBPyRelation: A DuckDB Python relation that can be queried or converted.
+        """
+
+        return self.conn.from_query(query=RECORD_TYPE_SCHEMA, params=[rty_ID])

heurist/database/database.py

@@ -0,0 +1,96 @@
+import pandas as pd
+from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from heurist.database.basedb import HeuristDatabase
+from heurist.models.dynamic import HeuristRecord
+from heurist.sql import RECORD_BY_GROUP_TYPE, RECORD_TYPE_METADATA
+from heurist.validators.record_validator import RecordValidator
+
+
+class TransformedDatabase(HeuristDatabase):
+    """Class for building and populating SQL tables with data collected and \
+    transformed from remote Heurist DB.
+    """
+
+    def __init__(
+        self,
+        hml_xml: bytes,
+        conn: DuckDBPyConnection | None = None,
+        db: str | None = ":memory:",
+        record_type_groups: list[str] = ["My record types"],
+    ) -> None:
+        super().__init__(hml_xml, conn, db)
+
+        # Create an empty index of targeted record types' Pydantic models
+        self.pydantic_models = {}
+
+        # Joining together the Heurist database's structural tables, construct an SQL
+        # statement that selects the ID and name of the record types that belong to one
+        # of the targeted record type groups
+        condition = "\nWHERE rtg.rtg_Name like '{}'".format(record_type_groups[0])
+        if len(record_type_groups) > 1:
+            for rtg in record_type_groups[1:]:
+                condition += " OR rtg.rtg_Name like '{}'".format(rtg)
+        query = RECORD_BY_GROUP_TYPE + condition
+
+        # Iterate through each targeted record type's ID and name
+        for rty_ID, rty_Name in self.conn.sql(query).fetchall():
+            # Using the ID, select the metadata of a record type's data fields (details)
+            rel = self.conn.sql(query=RECORD_TYPE_METADATA, params=[rty_ID])
+            # Using this metadata, create a dynamic Pydantic model for the record type
+            data_field_metadata = rel.pl().to_dicts()
+            model = HeuristRecord(
+                rty_ID=rty_ID,
+                rty_Name=rty_Name,
+                detail_metadata=data_field_metadata,
+            )
+            # Add the dynamic Pydantic model to the index of models
+            self.pydantic_models.update({rty_ID: model})
+
+    def insert_records(
+        self, record_type_id: int, records: list[dict]
+    ) -> DuckDBPyRelation | None:
+        # From the index of Pydantic models, get this record type's
+        # dynamically-created Pydantic model.
+        dynamic_model = self.pydantic_models[record_type_id].model
+        table_name = self.pydantic_models[record_type_id].table_name
+
+        # Prepare a list in which to store dictionaries of the validated record data.
+        model_dict_sequence = []
+
+        # Using the dynamically-created Pyandtic model, validate the metadata of
+        # all the records of this type.
+        validator = RecordValidator(
+            pydantic_model=dynamic_model,
+            records=records,
+            rty_ID=record_type_id,
+        )
+        for model in validator:
+            # Dump the validated record's data model to a dictionary.
+            model_dict = model.model_dump(by_alias=True)
+            # Add the dictionary representation of the validated data to the sequence.
+            model_dict_sequence.append(model_dict)
+
+        # If no records of this type have been created yet, skip it.
+        if len(model_dict_sequence) == 0:
+            return
+
+        # Transform the sequence of dictionaries into a Pandas dataframe
+        try:
+            df = pd.DataFrame(model_dict_sequence)
+            assert df.shape[1] > 0
+        except Exception as e:
+            from pprint import pprint
+
+            pprint(model_dict_sequence)
+            print(df)
+            print(records)
+            print(table_name)
+            raise e
+
+        # Delete any existing table for this record type.
+        self.delete_existing_table(table_name=table_name)
+
+        # From the dataframe, build a new table for the record type.
+        sql = f"""CREATE TABLE {table_name} AS FROM df"""
+        self.conn.sql(sql)
+        return self.conn.table(table_name=table_name)

heurist/models/dynamic/__init__.py

@@ -0,0 +1,3 @@
+from heurist.models.dynamic.create_model import HeuristRecord
+
+HeuristRecord