squirrels 0.4.1__py3-none-any.whl → 0.5.0__py3-none-any.whl

squirrels/_initializer.py

@@ -1,124 +1,181 @@
 from typing import Optional
 from datetime import datetime
-import inquirer, os, shutil
+from pathlib import Path
+import inquirer, os, shutil, secrets
 
 from . import _constants as c, _utils as u
 
-base_proj_dir = u.Path(os.path.dirname(__file__), c.PACKAGE_DATA_FOLDER, c.BASE_PROJECT_FOLDER)
+base_proj_dir = Path(os.path.dirname(__file__), c.PACKAGE_DATA_FOLDER, c.BASE_PROJECT_FOLDER)
+
+TMP_FOLDER = "tmp"
 
 
 class Initializer:
-    def __init__(self, *, overwrite: bool = False):
-        self.overwrite = overwrite
+    def __init__(self, *, project_name: Optional[str] = None, use_curr_dir: bool = False):
+        self.project_name = project_name if not use_curr_dir else None
+        self.use_curr_dir = use_curr_dir
 
-    def _path_exists(self, filepath: u.Path) -> bool:
+    def _path_exists(self, filepath: Path) -> bool:
         return os.path.exists(filepath)
     
-    def _files_have_same_content(self, file1: u.Path, file2: u.Path) -> bool:
+    def _files_have_same_content(self, file1: Path, file2: Path) -> bool:
         with open(file1, 'rb') as f1, open(file2, 'rb') as f2:
             return f1.read() == f2.read()
     
-    def _add_timestamp_to_filename(self, path: u.Path) -> u.Path:
+    def _add_timestamp_to_filename(self, path: Path) -> Path:
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
         new_filename = f"{path.stem}_{timestamp}{path.suffix}"
         return path.with_name(new_filename)
     
-    def _copy_file(self, filepath: u.Path, *, src_folder: str = ""):
-        src_path = u.Path(base_proj_dir, src_folder, filepath)
+    def _copy_file(self, filepath: Path, *, src_folder: str = "", src_file: Path | None = None):
+        src_file = src_file if src_file is not None else filepath
+        src_path = Path(base_proj_dir, src_folder, src_file)
         
-        dest_dir = os.path.dirname(filepath)
+        filepath2 = Path(self.project_name, filepath) if self.project_name else filepath
+        dest_dir = os.path.dirname(filepath2)
         if dest_dir != "":
             os.makedirs(dest_dir, exist_ok=True)
         
         perform_copy = True
-        if self._path_exists(filepath):
-            old_filepath = filepath
-            if self._files_have_same_content(src_path, filepath):
+        if self._path_exists(filepath2):
+            old_filepath = filepath2
+            if self._files_have_same_content(src_path, filepath2):
                 perform_copy = False
                 extra_msg = "Skipping... file contents is same as source"
-            elif self.overwrite:
-                extra_msg = "Overwriting file..."
             else:
-                filepath = self._add_timestamp_to_filename(old_filepath)
-                extra_msg = f'Creating file as "{filepath}" instead...'
+                filepath2 = self._add_timestamp_to_filename(old_filepath)
+                extra_msg = f'Creating file as "{filepath2}" instead...'
             print(f'File "{old_filepath}" already exists.', extra_msg)
         else:
-            print(f'Creating file "{filepath}"...')
+            print(f'Creating file "{filepath2}"...')
         
         if perform_copy:
-            shutil.copy(src_path, filepath)
+            shutil.copy(src_path, filepath2)
+
+    def _copy_macros_file(self, filepath: str):
+        self._copy_file(Path(c.MACROS_FOLDER, filepath))
+
+    def _copy_models_file(self, filepath: str):
+        self._copy_file(Path(c.MODELS_FOLDER, filepath))
+
+    def _copy_build_file(self, filepath: str):
+        self._copy_file(Path(c.MODELS_FOLDER, c.BUILDS_FOLDER, filepath))
 
     def _copy_dbview_file(self, filepath: str):
-        self._copy_file(u.Path(c.MODELS_FOLDER, c.DBVIEWS_FOLDER, filepath))
+        self._copy_file(Path(c.MODELS_FOLDER, c.DBVIEWS_FOLDER, filepath))
 
     def _copy_federate_file(self, filepath: str):
-        self._copy_file(u.Path(c.MODELS_FOLDER, c.FEDERATES_FOLDER, filepath))
+        self._copy_file(Path(c.MODELS_FOLDER, c.FEDERATES_FOLDER, filepath))
 
     def _copy_database_file(self, filepath: str):
-        self._copy_file(u.Path(c.DATABASE_FOLDER, filepath))
+        self._copy_file(Path(c.DATABASE_FOLDER, filepath))
     
     def _copy_pyconfig_file(self, filepath: str):
-        self._copy_file(u.Path(c.PYCONFIGS_FOLDER, filepath))
+        self._copy_file(Path(c.PYCONFIGS_FOLDER, filepath))
     
     def _copy_seed_file(self, filepath: str):
-        self._copy_file(u.Path(c.SEEDS_FOLDER, filepath))
+        self._copy_file(Path(c.SEEDS_FOLDER, filepath))
     
     def _copy_dashboard_file(self, filepath: str):
-        self._copy_file(u.Path(c.DASHBOARDS_FOLDER, filepath))
-
-    def _create_manifest_file(self, has_connections: bool, has_parameters: bool, has_dashboards: bool):
-        TMP_FOLDER = "tmp"
+        self._copy_file(Path(c.DASHBOARDS_FOLDER, filepath))
 
+    def _create_manifest_file(self, has_connections: bool, has_parameters: bool):
         def get_content(file_name: Optional[str]) -> str:
             if file_name is None:
                 return ""
             
-            yaml_path = u.Path(base_proj_dir, file_name)
-            return u.read_file(yaml_path)
+            yaml_path = Path(base_proj_dir, file_name)
+            return yaml_path.read_text()
         
         file_name_dict = {
             "parameters": c.PARAMETERS_YML_FILE if has_parameters else None, 
             "connections": c.CONNECTIONS_YML_FILE if has_connections else None,
-            "dashboards": c.DASHBOARDS_YML_FILE if has_dashboards else None
         }
         substitutions = {key: get_content(val) for key, val in file_name_dict.items()}
         
         manifest_template = get_content(c.MANIFEST_JINJA_FILE)
         manifest_content = u.render_string(manifest_template, **substitutions)
-        output_path = u.Path(base_proj_dir, TMP_FOLDER, c.MANIFEST_FILE)
-        with open(u.Path(output_path), "w") as f:
-            f.write(manifest_content)
+        output_path = Path(base_proj_dir, TMP_FOLDER, c.MANIFEST_FILE)
+        output_path.write_text(manifest_content)
         
-        self._copy_file(u.Path(c.MANIFEST_FILE), src_folder=TMP_FOLDER)
+        self._copy_file(Path(c.MANIFEST_FILE), src_folder=TMP_FOLDER)
+    
+    def _copy_dotenv_files(self, admin_password: str | None = None):
+        substitutions = {
+            "random_secret_key": secrets.token_hex(32),
+            "random_admin_password": admin_password if admin_password else secrets.token_urlsafe(8),
+        }
+
+        dotenv_path = Path(base_proj_dir, c.DOTENV_FILE)
+        contents = u.render_string(dotenv_path.read_text(), **substitutions)
+
+        output_path = Path(base_proj_dir, TMP_FOLDER, c.DOTENV_FILE)
+        output_path.write_text(contents)
+
+        self._copy_file(Path(c.DOTENV_FILE), src_folder=TMP_FOLDER)
+        self._copy_file(Path(c.DOTENV_FILE + ".example"))
+
+    def _copy_gitignore_file(self):
+        self._copy_file(Path(c.GITIGNORE_FILE), src_file=Path("gitignore"))
 
     def init_project(self, args):
-        options = ["core", "connections", "parameters", "dbview", "federate", "dashboard", "auth"]
-        _, CONNECTIONS, PARAMETERS, DBVIEW, FEDERATE, DASHBOARD, AUTH = options
+        options = ["connections", "parameters", "build", "federate", "dashboard", "admin_password"]
+        CONNECTIONS, PARAMETERS, BUILD, FEDERATE, DASHBOARD, ADMIN_PASSWORD = options
 
-        answers = { x: getattr(args, x) for x in options }
-        if not any(answers.values()):
+        # Add project name prompt if not provided
+        if self.project_name is None and not args.curr_dir:
             questions = [
-                inquirer.List(
-                    CONNECTIONS, message=f"How would you like to configure the database connections?", choices=c.CONF_FORMAT_CHOICES
-                ),
-                inquirer.List(
-                    PARAMETERS, message=f"How would you like to configure the parameters?", choices=c.CONF_FORMAT_CHOICES2
-                ),
-                inquirer.List(
-                    DBVIEW, message="What's the file format for the database view model?", choices=c.FILE_TYPE_CHOICES
-                ),
-                inquirer.List(
-                    FEDERATE, message="What's the file format for the federated model?", choices=c.FILE_TYPE_CHOICES
-                ),
-                inquirer.Confirm(
-                    DASHBOARD, message=f"Do you want to include a dashboard example?", default=False
-                ),
-                inquirer.Confirm(
-                    AUTH, message=f"Do you want to add the '{c.AUTH_FILE}' file to enable custom API authentication?", default=False
-                ),
+                inquirer.Text('project_name', message="What is your project folder name? (leave blank to create in current directory)")
             ]
             answers = inquirer.prompt(questions)
             assert isinstance(answers, dict)
+            self.project_name = answers['project_name']
+
+        answers = { x: getattr(args, x) for x in options }
+        if answers.get(DASHBOARD) is not None:
+            answers[DASHBOARD] = (answers[DASHBOARD] == 'y') # convert 'y' or 'n' to boolean
+
+        if not args.use_defaults:
+            questions = []
+            if answers.get(CONNECTIONS) is None:
+                questions.append(
+                    inquirer.List(
+                        CONNECTIONS, message=f"How would you like to configure the database connections?", choices=c.CONF_FORMAT_CHOICES
+                    ),
+                )
+            if answers.get(PARAMETERS) is None:
+                questions.append(
+                    inquirer.List(
+                        PARAMETERS, message=f"How would you like to configure the parameters?", choices=c.CONF_FORMAT_CHOICES2
+                    ),
+                )
+            if answers.get(BUILD) is None:
+                questions.append(
+                    inquirer.List(
+                        BUILD, message="What's the file format for the build model?", choices=c.FILE_TYPE_CHOICES
+                    ),
+                )
+            if answers.get(FEDERATE) is None:
+                questions.append(
+                    inquirer.List(
+                        FEDERATE, message="What's the file format for the federated model?", choices=c.FILE_TYPE_CHOICES
+                    ),
+                )
+            if answers.get(DASHBOARD) is None:
+                questions.append(
+                    inquirer.Confirm(
+                        DASHBOARD, message=f"Do you want to include a dashboard example?", default=False
+                    ),
+                )
+            if answers.get(ADMIN_PASSWORD) is None:
+                questions.append(
+                    inquirer.Password(
+                        "admin_password", message="What's the admin password? (leave blank to generate a random one)"
+                    ),
+                )
+            more_answers = inquirer.prompt(questions)
+            assert isinstance(more_answers, dict)
+            answers.update(more_answers)
         
         def get_answer(key, default):
             """
@@ -133,6 +190,8 @@ class Initializer:
             answer = answers.get(key)
             return answer if answer is not None else default
         
+        admin_password = get_answer("admin_password", None)
+        
         connections_format = get_answer(CONNECTIONS, c.YML_FORMAT)
         connections_use_yaml = (connections_format == c.YML_FORMAT)
         connections_use_py = (connections_format == c.PYTHON_FORMAT)
@@ -141,14 +200,19 @@ class Initializer:
         parameters_use_yaml = (parameters_format == c.YML_FORMAT)
         parameters_use_py = (parameters_format == c.PYTHON_FORMAT)
 
-        dbview_format = get_answer(DBVIEW, c.SQL_FILE_TYPE)
-        if dbview_format == c.SQL_FILE_TYPE:
-            db_view_file = c.DBVIEW_FILE_STEM + ".sql"
-        elif dbview_format == c.PYTHON_FILE_TYPE:
-            db_view_file = c.DBVIEW_FILE_STEM + ".py"
+        build_config_file = c.BUILD_FILE_STEM + ".yml"
+        build_format = get_answer(BUILD, c.SQL_FILE_TYPE)
+        if build_format == c.SQL_FILE_TYPE:
+            build_file = c.BUILD_FILE_STEM + ".sql"
+        elif build_format == c.PYTHON_FILE_TYPE:
+            build_file = c.BUILD_FILE_STEM + ".py"
         else:
-            raise NotImplementedError(f"Dbview model format '{dbview_format}' not supported")
+            raise NotImplementedError(f"Build model format '{build_format}' not supported")
+
+        db_view_config_file = c.DBVIEW_FILE_STEM + ".yml"
+        db_view_file = c.DBVIEW_FILE_STEM + ".sql"
     
+        federate_config_file = c.FEDERATE_FILE_STEM + ".yml"
         federate_format = get_answer(FEDERATE, c.SQL_FILE_TYPE)
         if federate_format == c.SQL_FILE_TYPE:
             federate_file = c.FEDERATE_FILE_STEM + ".sql"
@@ -159,10 +223,10 @@ class Initializer:
 
         dashboards_enabled = get_answer(DASHBOARD, False)
 
-        self._create_manifest_file(connections_use_yaml, parameters_use_yaml, dashboards_enabled)
+        self._copy_dotenv_files(admin_password)
+        self._create_manifest_file(connections_use_yaml, parameters_use_yaml)
         
-        self._copy_file(u.Path(".gitignore"))
-        self._copy_file(u.Path(c.ENV_CONFIG_FILE))
+        self._copy_gitignore_file()
         
         if connections_use_py:
             self._copy_pyconfig_file(c.CONNECTIONS_FILE)
@@ -178,44 +242,77 @@ class Initializer:
         else:
             raise NotImplementedError(f"Format '{parameters_format}' not supported for configuring widget parameters")
         
+        self._copy_pyconfig_file(c.USER_FILE)
+
         self._copy_pyconfig_file(c.CONTEXT_FILE)
-        self._copy_seed_file(c.CATEGORY_SEED_FILE)
-        self._copy_seed_file(c.SUBCATEGORY_SEED_FILE)
+        self._copy_seed_file(c.SEED_CATEGORY_FILE_STEM + ".csv")
+        self._copy_seed_file(c.SEED_CATEGORY_FILE_STEM + ".yml")
+        self._copy_seed_file(c.SEED_SUBCATEGORY_FILE_STEM + ".csv")
+        self._copy_seed_file(c.SEED_SUBCATEGORY_FILE_STEM + ".yml")
 
+        self._copy_macros_file(c.MACROS_FILE)
+
+        self._copy_models_file(c.SOURCES_FILE)
+        self._copy_build_file(build_file)
+        self._copy_build_file(build_config_file)
         self._copy_dbview_file(db_view_file)
+        self._copy_dbview_file(db_view_config_file)
         self._copy_federate_file(federate_file)
+        self._copy_federate_file(federate_config_file)
         
         if dashboards_enabled:
             self._copy_dashboard_file(c.DASHBOARD_FILE_STEM + ".py")
+            self._copy_dashboard_file(c.DASHBOARD_FILE_STEM + ".yml")
         
-        if get_answer(AUTH, False):
-            self._copy_pyconfig_file(c.AUTH_FILE)
-
         self._copy_database_file(c.EXPENSES_DB)
         
-        print(f"\nSuccessfully created new Squirrels project in current directory!\n")
+        print(f"\nSuccessfully created new Squirrels project!\n")
     
     def get_file(self, args):
-        if args.file_name == c.ENV_CONFIG_FILE:
-            self._copy_file(u.Path(c.ENV_CONFIG_FILE))
-            print("PLEASE ENSURE THE FILE IS INCLUDED IN .gitignore")
+        if args.file_name == c.DOTENV_FILE:
+            self._copy_dotenv_files()
+            print(f"A random admin password was generated for your project. You can change it in the new {c.DOTENV_FILE} file.")
+            print()
+            print(f"IMPORTANT: Please ensure the {c.DOTENV_FILE} file is added to your {c.GITIGNORE_FILE} file.")
+            print(f"You may also run `sqrl get-file {c.GITIGNORE_FILE}` to add a sample {c.GITIGNORE_FILE} file to your project.")
+            print()
+        elif args.file_name == c.GITIGNORE_FILE:
+            self._copy_gitignore_file()
         elif args.file_name == c.MANIFEST_FILE:
-            self._create_manifest_file(not args.no_connections, args.parameters, args.dashboards)
-        elif args.file_name in (c.AUTH_FILE, c.CONNECTIONS_FILE, c.PARAMETERS_FILE, c.CONTEXT_FILE):
+            self._create_manifest_file(not args.no_connections, args.parameters)
+        elif args.file_name in (c.USER_FILE, c.CONNECTIONS_FILE, c.PARAMETERS_FILE, c.CONTEXT_FILE):
             self._copy_pyconfig_file(args.file_name)
-        elif args.file_name in (c.DBVIEW_FILE_STEM, c.FEDERATE_FILE_STEM):
-            if args.format == c.SQL_FILE_TYPE:
+        elif args.file_name == c.MACROS_FILE:
+            self._copy_macros_file(args.file_name)
+        elif args.file_name == c.SOURCES_FILE:
+            self._copy_models_file(args.file_name)
+        elif args.file_name in (c.BUILD_FILE_STEM, c.DBVIEW_FILE_STEM, c.FEDERATE_FILE_STEM):
+            if args.file_name == c.DBVIEW_FILE_STEM or args.format == c.SQL_FILE_TYPE:
                 extension = ".sql"
             elif args.format == c.PYTHON_FILE_TYPE:
                 extension = ".py"
             else:
                 raise NotImplementedError(f"Format '{args.format}' not supported for {args.file_name}")
-            copy_method = self._copy_dbview_file if args.file_name == c.DBVIEW_FILE_STEM else self._copy_federate_file
+            
+            if args.file_name == c.BUILD_FILE_STEM:
+                copy_method = self._copy_build_file
+            elif args.file_name == c.DBVIEW_FILE_STEM:
+                copy_method = self._copy_dbview_file
+            elif args.file_name == c.FEDERATE_FILE_STEM:
+                copy_method = self._copy_federate_file
+            else:
+                raise NotImplementedError(f"File '{args.file_name}' not supported")
+            
             copy_method(args.file_name + extension)
+            copy_method(args.file_name + ".yml")
         elif args.file_name == c.DASHBOARD_FILE_STEM:
             self._copy_dashboard_file(args.file_name + ".py")
+            self._copy_dashboard_file(args.file_name + ".yml")
         elif args.file_name in (c.EXPENSES_DB, c.WEATHER_DB):
             self._copy_database_file(args.file_name)
+        elif args.file_name in (c.SEED_CATEGORY_FILE_STEM, c.SEED_SUBCATEGORY_FILE_STEM):
+            self._copy_seed_file(args.file_name + ".csv")
+            self._copy_seed_file(args.file_name + ".yml")
         else:
             raise NotImplementedError(f"File '{args.file_name}' not supported")
         

squirrels/_logging.py

@@ -0,0 +1,115 @@
+from pathlib import Path
+from logging.handlers import RotatingFileHandler
+from uuid import uuid4
+import logging as l, json
+
+from . import _constants as c, _utils as u
+
+
+class _BaseFormatter(l.Formatter):
+    def _format_helper(self, level_for_print: str, record: l.LogRecord) -> str:
+        # Save original levelname
+        original_levelname = record.levelname
+
+        # Add padding to the levelname for printing
+        visible_length = len(record.levelname) + 1
+        padding_needed = max(1, 9 - visible_length)
+        padded_level = f"{level_for_print}:{' ' * padding_needed}"
+        record.levelname = padded_level
+        
+        # Format the message
+        formatted = super().format(record)
+        
+        # Restore original levelname
+        record.levelname = original_levelname
+        
+        return formatted
+
+
+class _ColoredFormatter(_BaseFormatter):
+    """Custom formatter that adds colors to log levels for terminal output"""
+    
+    # ANSI color codes
+    COLORS = {
+        'DEBUG': '\033[36m',      # Cyan
+        'INFO': '\033[32m',       # Green
+        'WARNING': '\033[33m',    # Yellow
+        'ERROR': '\033[31m',      # Red
+        'CRITICAL': '\033[35m',   # Magenta
+    }
+    RESET = '\033[0m'
+    BOLD = '\033[1m'
+    
+    def format(self, record: l.LogRecord) -> str:
+        # Add color to levelname with colon and padding
+        color = self.COLORS.get(record.levelname, '')
+        colored_level = f"{color}{record.levelname}{self.RESET}"
+        return self._format_helper(colored_level, record)
+
+
+class _PlainFormatter(_BaseFormatter):
+    """Custom formatter that adds colon to log levels for file output"""
+    
+    def format(self, record: l.LogRecord) -> str:
+        return self._format_helper(record.levelname, record)
+
+
+class _CustomJsonFormatter(l.Formatter):
+    def format(self, record: l.LogRecord) -> str:
+        super().format(record)
+        info = {
+            "timestamp": self.formatTime(record),
+            "project_id": record.name,
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "thread": record.thread,
+            "thread_name": record.threadName,
+            "process": record.process,
+            **record.__dict__.get("info", {})
+        }
+        output = {
+            "data": record.__dict__.get("data", {}),
+            "info": info
+        }
+        return json.dumps(output)
+
+
+def get_logger(
+    base_path: str, log_to_file: bool, log_level: str, log_format: str, log_file_size_mb: int, log_file_backup_count: int
+) -> u.Logger:
+    logger = u.Logger(name=uuid4().hex)
+    logger.setLevel(log_level.upper())
+
+    # Determine the formatter based on log_format
+    if log_format.lower() == "json":
+        stdout_formatter = _CustomJsonFormatter()
+        file_formatter = _CustomJsonFormatter()
+    elif log_format.lower() == "text":
+        # Use colored formatter for stdout, plain formatter with colon for file
+        format_string = "%(levelname)s [%(asctime)s] %(message)s"
+        stdout_formatter = _ColoredFormatter(format_string, datefmt="%Y-%m-%d %H:%M:%S")
+        file_formatter = _PlainFormatter(format_string, datefmt="%Y-%m-%d %H:%M:%S")
+    else:
+        raise ValueError("log_format must be either 'text' or 'json'")
+    
+    if log_to_file:
+        log_file_path = Path(base_path, c.LOGS_FOLDER, c.LOGS_FILE)
+        log_file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Rotating file handler
+        file_handler = RotatingFileHandler(
+            log_file_path,
+            maxBytes=log_file_size_mb * 1024 * 1024,
+            backupCount=log_file_backup_count
+        )
+        file_handler.setLevel(log_level.upper())
+        file_handler.setFormatter(file_formatter)
+        logger.addHandler(file_handler)
+    
+    else:
+        stdout_handler = l.StreamHandler()
+        stdout_handler.setLevel(log_level.upper())
+        stdout_handler.setFormatter(stdout_formatter)
+        logger.addHandler(stdout_handler)
+    
+    return logger