refactors telemetry and cloud to use new Config

src/ploomber/cli/cloud.py

@@ -12,16 +12,15 @@
 import warnings
 from datetime import datetime
 from json import JSONDecodeError
-from pathlib import Path
 import http.client as httplib
 import click
 from functools import wraps
 
 import humanize
 
+from ploomber.exceptions import BaseException
 from ploomber.telemetry import telemetry
-from ploomber.telemetry.telemetry import check_dir_exist, CONF_DIR, \
-    DEFAULT_USER_CONF, read_conf_file, update_conf_file, parse_dag
+from ploomber.telemetry.telemetry import parse_dag, UserSettings
 
 CLOUD_APP_URL = 'ggeheljnx2.execute-api.us-east-1.amazonaws.com'
 PIPELINES_RESOURCE = '/prod/pipelines'
@@ -33,11 +32,7 @@ def get_key():
     This gets the user cloud api key, returns None if doesn't exist.
     config.yaml is the default user conf file to fetch from.
     """
-    user_conf_path = Path(check_dir_exist(CONF_DIR), DEFAULT_USER_CONF)
-    conf = read_conf_file(user_conf_path)
-    key = conf.get('cloud_key', None)
-
-    return key
+    return UserSettings().cloud_key
 
 
 @telemetry.log_call('set-key')
@@ -52,13 +47,11 @@ def set_key(user_key):
 def _set_key(user_key):
     # Validate key
     if not user_key or len(user_key) != 22:
-        warnings.warn("The API key is malformed.\n"
-                      "Please validate your key or contact the admin\n")
-        return
+        raise BaseException("The API key is malformed.\n"
+                            "Please validate your key or contact the admin.")
 
-    user_key_dict = {'cloud_key': user_key}
-    user_conf_path = Path(check_dir_exist(CONF_DIR), DEFAULT_USER_CONF)
-    update_conf_file(user_conf_path, user_key_dict)
+    settings = UserSettings()
+    settings.cloud_key = user_key
     click.secho("Key was stored")
 
 

src/ploomber/telemetry/telemetry.py

@@ -31,20 +31,19 @@
 import datetime
 import http.client as httplib
 import json
-import warnings
 import os
 from pathlib import Path
 import sys
-import uuid
+from uuid import uuid4
 from functools import wraps
 import platform
 
 import click
 import posthog
-import yaml
 import distro
 
 from ploomber.telemetry import validate_inputs
+from ploomber.config import Config
 from ploomber import __version__
 
 TELEMETRY_VERSION = '0.3'
@@ -56,18 +55,15 @@
 PLOOMBER_HOME_DIR = os.getenv("PLOOMBER_HOME_DIR")
 
 
-class Config:
-    def read(self):
-        return read_conf_file(self.path())
-
-    def write(self, data, error=False):
-        write_conf_file(self.path(), data, error=error)
-
-
 class UserSettings(Config):
     """User-customizable settings
     """
-    def path(self):
+    version_check_enabled: bool = True
+    cloud_key: str = None
+    stats_enabled: bool = True
+
+    @classmethod
+    def path(cls):
         return Path(check_dir_exist(CONF_DIR), DEFAULT_USER_CONF)
 
 
@@ -76,28 +72,16 @@ class Internal(Config):
     Internal file to store settings (not intended to be modified by the
     user)
     """
-    def path(self):
-        return Path(check_dir_exist(CONF_DIR), DEFAULT_PLOOMBER_CONF)
-
-
-def read_conf_file(conf_path):
-    try:
-        with conf_path.open("r") as file:
-            conf = yaml.safe_load(file)
-            return conf
-    except Exception as e:
-        warnings.warn(f"Can't read config file {e}")
-        return {}
+    last_version_check: str = None
+    uid: str
+    first_time: bool = True
 
+    @classmethod
+    def path(cls):
+        return Path(check_dir_exist(CONF_DIR), DEFAULT_PLOOMBER_CONF)
 
-def write_conf_file(conf_path, to_write, error=None):
-    try:  # Create for future runs
-        with conf_path.open("w") as file:
-            yaml.dump(to_write, file)
-    except Exception as e:
-        warnings.warn(f"Can't write to config file: {e}")
-        if error:
-            return e
+    def uid_default(self):
+        return str(uuid4())
 
 
 def python_version():
@@ -119,8 +103,9 @@ def is_online():
         conn.close()
 
 
-# Will output if the code is within a container
 def is_docker():
+    """Will output if the code is within a container
+    """
     try:
         cgroup = Path('/proc/self/cgroup')
         docker_env = Path('/.dockerenv')
@@ -184,7 +169,8 @@ def is_conda():
 def get_base_prefix_compat():
     """
     This function will find the pip virtualenv with different python versions.
-    Get base/real prefix, or sys.prefix if there is none."""
+    Get base/real prefix, or sys.prefix if there is none.
+    """
     return getattr(sys, "base_prefix", None) or sys.prefix or getattr(
         sys, "real_prefix", None)
 
@@ -294,23 +280,7 @@ def check_dir_exist(input_location=None):
     return p
 
 
-def check_uid():
-    """
-    Checks if local user id exists as a uid file, creates if not.
-    """
-    internal = Internal()
-    conf = internal.read()  # file already exist due to version check
-    if 'uid' not in conf.keys():
-        uid = str(uuid.uuid4())
-        res = internal.write({"uid": uid}, error=True)
-        if res:
-            return f"NO_UID {res}"
-        else:
-            return uid
-    return conf.get('uid', "NO_UID")
-
-
-def check_stats_enabled():
+def check_telemetry_enabled():
     """
     Check if the user allows us to use telemetry. In order of precedence:
 
@@ -321,24 +291,18 @@ def check_stats_enabled():
         return os.environ['PLOOMBER_STATS_ENABLED'].lower() == 'true'
 
     settings = UserSettings()
-    # Check if local config exists
-    config_path = settings.path()
-    if not config_path.exists():
-        settings.write({"stats_enabled": True})
-        return True
-    else:  # read and return config
-        conf = settings.read()
-        return conf.get('stats_enabled', True)
+    return settings.stats_enabled
 
 
 def check_first_time_usage():
     """
     The function checks for first time usage if the conf file exists and the
     uid file doesn't exist.
     """
-    config_path = UserSettings().path()
-    uid_conf = Internal().read()
-    return config_path.exists() and 'uid' not in uid_conf.keys()
+    internal = Internal()
+    first_time = internal.first_time
+    internal.first_time = False
+    return first_time
 
 
 def get_latest_version():
@@ -359,20 +323,14 @@ def get_latest_version():
         conn.close()
 
 
-def update_conf_file(conf_path, to_write, error=None):
-    conf = read_conf_file(conf_path)
-    new_conf = dict(conf, **to_write)
-    write_conf_file(conf_path, new_conf, error)
-
-
 def is_cloud_user():
     """
         The function checks if the cloud api key is set for the user.
         Checks if the cloud_key is set in the User conf file (config.yaml).
         returns True/False accordingly.
         """
-    conf = UserSettings().read()
-    return conf.get('cloud_key', False)
+    settings = UserSettings()
+    return settings.cloud_key
 
 
 def check_version():
@@ -382,31 +340,17 @@ def check_version():
     If it's not the latest, notifies the user and saves the metadata to conf
     Alerting every 2 days on stale versions
     """
-    # Read conf file
-    today = datetime.datetime.now()
-    conf = UserSettings().read()
-
-    internal = Internal()
-    version_path = internal.path()
-    # Update version conf if not there
-    if not version_path.exists():
-        version = {'last_version_check': today}
-    else:
-        version = internal.read()
-        if 'last_version_check' not in version.keys():
-            version['last_version_check'] = today
-
-    internal.write(version)
+    settings = UserSettings()
 
-    # Check if the flag was disabled
-    if conf and 'version_check_enabled' in conf.keys() \
-            and not conf['version_check_enabled']:
+    if not settings.version_check_enabled:
         return
 
+    internal = Internal()
+    now = datetime.datetime.now()
+
     # Check if we already notified in the last 2 days
-    last_message = version['last_version_check']
-    diff = (today - last_message).days
-    if diff < 2:
+    if internal.last_version_check and (now -
+                                        internal.last_version_check).days < 2:
         return
 
     # check latest version (this is an expensive call since it hits pypi.org)
@@ -424,8 +368,7 @@ def check_version():
         fg='yellow')
 
     # Update latest check date
-    version['last_version_check'] = today
-    internal.write(version)
+    internal.last_version_check = now
 
 
 def _get_telemetry_info():
@@ -435,19 +378,18 @@ def _get_telemetry_info():
     for first time installation.
     """
     # Check if telemetry is enabled, if not skip, else check for uid
-    telemetry_enabled = check_stats_enabled()
+    telemetry_enabled = check_telemetry_enabled()
 
     # Check latest version
     check_version()
 
     if telemetry_enabled:
-
         # Check first time install
         is_install = check_first_time_usage()
 
         # if not uid, create
-        uid = check_uid()
-        return telemetry_enabled, uid, is_install
+        internal = Internal()
+        return telemetry_enabled, internal.uid, is_install
     else:
         return False, '', False
 
@@ -469,11 +411,14 @@ def log_api(action, client_time=None, total_runtime=None, metadata=None):
     """
     metadata = metadata or {}
 
-    event_id = uuid.uuid4()
+    event_id = uuid4()
+
     if client_time is None:
         client_time = datetime.datetime.now()
 
     (telemetry_enabled, uid, is_install) = _get_telemetry_info()
+
+    # NOTE: this should not happen anymore
     if 'NO_UID' in uid:
         metadata['uid_issue'] = uid
         uid = None

tests/cli/test_cloud.py

@@ -91,7 +91,6 @@ def test_write_key_no_conf_file(tmp_directory, monkeypatch):
     with full_path.open("r") as file:
         conf = yaml.safe_load(file)
 
-    assert len(conf.keys()) == 1
     assert key_name in conf.keys()
     assert key_val in conf[key_name]
 
@@ -112,25 +111,19 @@ def test_overwrites_api_key(write_sample_conf):
     assert another_val in conf[key_name]
 
 
-def test_api_key_well_formatted(write_sample_conf):
-    with pytest.warns(Warning) as record:
-        cloud.set_key(None)
+@pytest.mark.parametrize('arg', [None, '12345'])
+def test_api_key_well_formatted(write_sample_conf, arg):
+    with pytest.raises(BaseException) as excinfo:
+        cloud.set_key(arg)
 
-    if not record:
-        pytest.fail("Expected a user warning!")
-
-    with pytest.warns(Warning) as record:
-        cloud.set_key("12345")
-
-    if not record:
-        pytest.fail("Expected a user warning!")
+    assert 'The API key is malformed' in str(excinfo.value)
 
 
 def test_get_api_key(write_sample_conf, capsys):
     key_val = "TEST_KEY12345678987654"
     runner = CliRunner()
     result = runner.invoke(set_key, args=[key_val], catch_exceptions=False)
-    assert 'Key was stored\n' == result.stdout
+    assert 'Key was stored\n' in result.stdout
 
     result = runner.invoke(get_key, catch_exceptions=False)
     assert key_val in result.stdout