feat(share): add sharing using magic-wormhole

openadapt/config.py

@@ -8,27 +8,69 @@
     ...
 
 """
-
 import multiprocessing
 import os
 import pathlib
 
-from dotenv import load_dotenv
+from dotenv import find_dotenv, load_dotenv
 from loguru import logger
 
+ROOT_DIRPATH = pathlib.Path(__file__).parent.parent.resolve()
+ZIPPED_RECORDING_FOLDER_PATH = ROOT_DIRPATH / "data" / "zipped"
+
+ENV_FILE_PATH = (ROOT_DIRPATH / ".env").resolve()
+logger.info(f"{ENV_FILE_PATH=}")
+dotenv_file = find_dotenv()
+load_dotenv(dotenv_file)
+
+
+def read_env_file(file_path):
+    """Read the contents of an environment file.
+
+    Args:
+        file_path (str): The path to the environment file.
+
+    Returns:
+        dict: A dictionary containing the environment variables and their values.
+    """
+    env_vars = {}
+    with open(file_path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if line and not line.startswith("#"):
+                key, value = line.split("=", 1)
+                env_vars[key] = value.strip('"')
+    return env_vars
+
+
+env_file_path = ".env"
+env_vars = read_env_file(env_file_path)
+
+DB_FNAME = env_vars.get("DB_FNAME")
+
+
+def set_db_url(db_fname):
+    """Set the database URL based on the given database file name.
+
+    Args:
+        db_fname (str): The database file name.
+    """
+    global DB_FNAME, DB_FPATH, DB_URL
+    DB_FNAME = db_fname
+    DB_FPATH = ROOT_DIRPATH / DB_FNAME
+    DB_URL = f"sqlite:///{DB_FPATH}"
+    logger.info(f"{DB_URL=}")
+
 
 _DEFAULTS = {
     "CACHE_DIR_PATH": ".cache",
     "CACHE_ENABLED": True,
     "CACHE_VERBOSITY": 0,
     "DB_ECHO": False,
-    "DB_FNAME": "openadapt.db",
+    "DB_FNAME": DB_FNAME,
     "OPENAI_API_KEY": "<set your api key in .env>",
-    # "OPENAI_MODEL_NAME": "gpt-4",
     "OPENAI_MODEL_NAME": "gpt-3.5-turbo",
-    # may incur significant performance penalty
     "RECORD_READ_ACTIVE_ELEMENT_STATE": False,
-    # TODO: remove?
     "REPLAY_STRIP_ELEMENT_STATE": True,
     # IGNORES WARNINGS (PICKLING, ETC.)
     # TODO: ignore warnings by default on GUI
@@ -87,29 +129,43 @@
 
 
 def getenv_fallback(var_name):
-    rval = os.getenv(var_name) or _DEFAULTS.get(var_name)
+    """Get the value of an environment variable with fallback to default.
+
+    Args:
+        var_name (str): The name of the environment variable.
+
+    Returns:
+        str: The value of the environment variable or the default value if not found.
+
+    Raises:
+        ValueError: If the environment variable is not defined.
+    """
+    if var_name == "DB_FNAME":
+        rval = _DEFAULTS.get(var_name)
+    else:
+        rval = os.getenv(var_name) or _DEFAULTS.get(var_name)
     if rval is None:
         raise ValueError(f"{var_name=} not defined")
     return rval
 
 
-load_dotenv()
-
 for key in _DEFAULTS:
     val = getenv_fallback(key)
     locals()[key] = val
 
+
 ROOT_DIRPATH = pathlib.Path(__file__).parent.parent.resolve()
 DB_FPATH = ROOT_DIRPATH / DB_FNAME
 DB_URL = f"sqlite:///{DB_FPATH}"
 DIRNAME_PERFORMANCE_PLOTS = "performance"
+DB_ECHO = False
+DT_FMT = "%Y-%m-%d_%H-%M-%S"
 
 if multiprocessing.current_process().name == "MainProcess":
     for key, val in locals().items():
         if not key.startswith("_") and key.isupper():
             logger.info(f"{key}={val}")
 
-
 def filter_log_messages(data):
     """
     This function filters log messages by ignoring any message that contains a specific string.
@@ -127,4 +183,4 @@ def filter_log_messages(data):
     messages_to_ignore = [
         "Cannot pickle Objective-C objects",
     ]
-    return not any(msg in data["message"] for msg in messages_to_ignore)
+    return not any(msg in data["message"] for msg in messages_to_ignore)

openadapt/crud.py

@@ -1,7 +1,14 @@
+import os
+import time
+import shutil
+
+from datetime import datetime
 from loguru import logger
+from sqlalchemy.orm import sessionmaker
 import sqlalchemy as sa
 
-from openadapt.db import Session
+from openadapt.db import Session, get_base, get_engine, engine
+from openadapt import config
 from openadapt.models import (
     ActionEvent,
     Screenshot,
@@ -19,13 +26,11 @@
 window_events = []
 performance_stats = []
 
+
 def _insert(event_data, table, buffer=None):
     """Insert using Core API for improved performance (no rows are returned)"""
 
-    db_obj = {
-        column.name: None
-        for column in table.__table__.columns
-    }
+    db_obj = {column.name: None for column in table.__table__.columns}
     for key in db_obj:
         if key in event_data:
             val = event_data[key]
@@ -74,6 +79,7 @@ def insert_window_event(recording_timestamp, event_timestamp, event_data):
     }
     _insert(event_data, WindowEvent, window_events)
 
+
 def insert_perf_stat(recording_timestamp, event_type, start_time, end_time):
     """
     Insert event performance stat into db
@@ -87,19 +93,20 @@ def insert_perf_stat(recording_timestamp, event_type, start_time, end_time):
     }
     _insert(event_perf_stat, PerformanceStat, performance_stats)
 
+
 def get_perf_stats(recording_timestamp):
     """
     return performance stats for a given recording
     """
 
     return (
-        db
-        .query(PerformanceStat)
+        db.query(PerformanceStat)
         .filter(PerformanceStat.recording_timestamp == recording_timestamp)
         .order_by(PerformanceStat.start_time)
         .all()
     )
 
+
 def insert_recording(recording_data):
     db_obj = Recording(**recording_data)
     db.add(db_obj)
@@ -109,28 +116,107 @@ def insert_recording(recording_data):
 
 
 def get_latest_recording():
-    return (
-        db
-        .query(Recording)
-        .order_by(sa.desc(Recording.timestamp))
-        .limit(1)
-        .first()
-    )
+    return db.query(Recording).order_by(sa.desc(Recording.timestamp)).limit(1).first()
+
+
+def get_recording_by_id(recording_id):
+    return db.query(Recording).filter_by(id=recording_id).first()
+
+
+def export_sql(recording_id):
+    """Export the recording data as SQL statements.
+
+    Args:
+        recording_id (int): The ID of the recording.
+
+    Returns:
+        str: The SQL statements to insert the recording into the output file.
+    """
+    engine = sa.create_engine(config.DB_URL)
+    Session = sessionmaker(bind=engine)
+    session = Session()
+
+    recording = get_recording_by_id(recording_id)
+
+    if recording:
+        sql = f"INSERT INTO recording VALUES ({recording.id}, {recording.timestamp}, {recording.monitor_width}, {recording.monitor_height}, {recording.double_click_interval_seconds}, {recording.double_click_distance_pixels}, '{recording.platform}', '{recording.task_description}')"
+        logger.info(f"Recording with ID {recording_id} exported successfully.")
+    else:
+        logger.info(f"No recording found with ID {recording_id}.")
+
+    return sql
+
+
+def create_db(recording_id, sql):
+    """Create a new database and import the recording data.
 
+    Args:
+        recording_id (int): The ID of the recording.
+        sql (str): The SQL statements to import the recording.
+
+    Returns:
+        tuple: A tuple containing the timestamp and the file path of the new database.
+    """
+    db.close()
+    db_fname = f"recording_{recording_id}.db"
+
+    t = time.time()
+    shutil.copyfile(config.ENV_FILE_PATH, f"{config.ENV_FILE_PATH}-{t}")
+    config.set_db_url(db_fname)
+
+    with open(config.ENV_FILE_PATH, "r") as f:
+        lines = f.readlines()
+    lines[1] = f"DB_FNAME={db_fname}\n"
+    with open(config.ENV_FILE_PATH, "w") as f:
+        f.writelines(lines)
+
+    engine = sa.create_engine(config.DB_URL)
+    Session = sessionmaker(bind=engine)
+    session = Session()
+    os.system("alembic upgrade head")
+    db.engine = engine
+
+    with engine.begin() as connection:
+        connection.execute(sql)
+
+    db_file_path = config.DB_FPATH.resolve()
+
+    return t, db_file_path
+
+
+def restore_db(timestamp):
+    """Restore the database to a previous state.
+
+    Args:
+        timestamp (float): The timestamp associated with the backup file.
+    """
+    backup_file = f"{config.ENV_FILE_PATH}-{timestamp}"
+    shutil.copyfile(backup_file, config.ENV_FILE_PATH)
+    config.set_db_url("openadapt.db")
+    db.engine = get_engine()
+
+
+def export_recording(recording_id):
+    """Export a recording by creating a new database, importing the recording, and then restoring the previous state.
+
+    Args:
+        recording_id (int): The ID of the recording to export.
+
+    Returns:
+        str: The file path of the new database.
+    """
+    sql = export_sql(recording_id)
+    t, db_file_path = create_db(recording_id, sql)
+    restore_db(t)
+    return db_file_path
 
 def get_recording(timestamp):
-    return (
-        db
-        .query(Recording)
-        .filter(Recording.timestamp == timestamp)
-        .first()
-    )
+    return db.query(Recording).filter(Recording.timestamp == timestamp).first()
 
 
 def _get(table, recording_timestamp):
     return (
-        db
-        .query(table)
+        db.query(table)
         .filter(table.recording_timestamp == recording_timestamp)
         .order_by(table.timestamp)
         .all()
@@ -150,11 +236,11 @@ def get_screenshots(recording, precompute_diffs=False):
 
     # TODO: store diffs
     if precompute_diffs:
-        logger.info(f"precomputing diffs...")
+        logger.info("precomputing diffs...")
         [(screenshot.diff, screenshot.diff_mask) for screenshot in screenshots]
 
     return screenshots
 
 
 def get_window_events(recording):
-    return _get(WindowEvent, recording.timestamp)
+    return _get(WindowEvent, recording.timestamp)