squirrels 0.1.0__py3-none-any.whl → 0.6.0.post0__py3-none-any.whl

squirrels/_initializer.py

@@ -1,110 +1,319 @@
-import inquirer, os, shutil
-
-from squirrels import _constants as c, _utils
-from squirrels._version import major_version
-
-base_proj_dir = _utils.join_paths(os.path.dirname(__file__), 'package_data', 'base_project')
-dataset_dir = _utils.join_paths('datasets', 'sample_dataset')
-
-
-class Initializer:
-    def __init__(self, overwrite: bool):
-        self.overwrite = overwrite
-
-    def _path_exists(self, filepath: str) -> bool:
-        if not self.overwrite and os.path.exists(filepath):
-            print(f'File "{filepath}" already exists. Creation skipped.')
-            return True
-        return False
-    
-    def _copy_file(self, filepath: str):
-        if not self._path_exists(filepath):
-            dest_dir = os.path.dirname(filepath)
-            if dest_dir != '':
-                os.makedirs(dest_dir, exist_ok=True)
-            src_path = _utils.join_paths(base_proj_dir, filepath)
-            shutil.copy(src_path, filepath)
-
-    def _copy_dataset_file(self, filepath: str):
-        self._copy_file(_utils.join_paths(dataset_dir, filepath))
-
-    def _copy_database_file(self, filepath: str):
-        self._copy_file(_utils.join_paths('database', filepath))
-
-    def _create_requirements_txt(self):
-        filename = 'requirements.txt'
-        if not self._path_exists(filename):
-            next_major_version = int(major_version) + 1
-            content = f'squirrels<{next_major_version}'
-            with open(filename, 'w') as f:
-                f.write(content)
-
-    def init_project(self, args):
-        options = ['core', 'db_view', 'connections', 'context', 'selections_cfg', 'final_view', 'sample_db']
-        answers = { x: getattr(args, x) for x in options }
-        if not any(answers.values()):
-            core_questions = [
-                inquirer.Confirm('core', 
-                                message="Include all core project files?",
-                                default=True)
-            ]
-            answers = inquirer.prompt(core_questions)
-            
-            if answers.get('core', False):
-                conditional_questions = [
-                    inquirer.List('db_view', 
-                                  message="What's the file format for the database view?",
-                                  choices=c.FILE_TYPE_CHOICES),
-                ]
-                answers.update(inquirer.prompt(conditional_questions))
-
-            remaining_questions = [
-                inquirer.Confirm('connections',
-                                message=f"Do you want to add a '{c.CONNECTIONS_FILE}' file?" ,
-                                default=False),
-                inquirer.Confirm('context',
-                                message=f"Do you want to add a '{c.CONTEXT_FILE}' file?" ,
-                                default=False),
-                inquirer.Confirm('selections_cfg',
-                                message=f"Do you want to add a '{c.SELECTIONS_CFG_FILE}' file?" ,
-                                default=False),
-                inquirer.List('final_view', 
-                            message="What's the file format for the final view (if any)?",
-                            choices=['none'] + c.FILE_TYPE_CHOICES),
-                inquirer.List('sample_db', 
-                            message="What sample sqlite database do you wish to use (if any)?",
-                            choices=['none'] + c.DATABASE_CHOICES)
-            ]
-            answers.update(inquirer.prompt(remaining_questions))
-
-        if answers.get('core', False):
-            self._copy_file('.gitignore')
-            self._copy_file(c.MANIFEST_FILE)
-            self._create_requirements_txt()
-            self._copy_dataset_file(c.PARAMETERS_FILE)
-            if answers.get('db_view') == 'py':
-                self._copy_dataset_file(c.DATABASE_VIEW_PY_FILE)
-            else:
-                self._copy_dataset_file(c.DATABASE_VIEW_SQL_FILE)
-        
-        if answers.get('connections', False):
-            self._copy_file(c.CONNECTIONS_FILE)
-        
-        if answers.get('context', False):
-            self._copy_dataset_file(c.CONTEXT_FILE)
-        
-        if answers.get('selections_cfg', False):
-            self._copy_dataset_file(c.SELECTIONS_CFG_FILE)
-        
-        final_view_format = answers.get('final_view')
-        if final_view_format == 'py':
-            self._copy_dataset_file(c.FINAL_VIEW_PY_NAME)
-        elif final_view_format == 'sql':
-            self._copy_dataset_file(c.FINAL_VIEW_SQL_NAME)
-
-        sample_db = answers.get('sample_db')
-        if sample_db == 'sample_database':
-            self._copy_database_file('sample_database.db')
-        elif sample_db == 'seattle_weather':
-            self._copy_database_file('seattle_weather.db')
-    
+from typing import Optional
+from datetime import datetime
+from pathlib import Path
+import inquirer, os, shutil, secrets
+
+from . import _constants as c, _utils as u
+
+base_proj_dir = Path(os.path.dirname(__file__), c.PACKAGE_DATA_FOLDER, c.BASE_PROJECT_FOLDER)
+
+TMP_FOLDER = "tmp"
+
+
+class Initializer:
+    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: Path) -> bool:
+        return os.path.exists(filepath)
+    
+    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: 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: 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)
+        
+        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(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"
+            else:
+                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 "{filepath2}"...')
+        
+        if perform_copy:
+            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(Path(c.MODELS_FOLDER, c.DBVIEWS_FOLDER, filepath))
+
+    def _copy_federate_file(self, filepath: str):
+        self._copy_file(Path(c.MODELS_FOLDER, c.FEDERATES_FOLDER, filepath))
+
+    def _copy_resource_file(self, filepath: str):
+        self._copy_file(Path(c.RESOURCES_FOLDER, filepath))
+    
+    def _copy_pyconfig_file(self, filepath: str):
+        self._copy_file(Path(c.PYCONFIGS_FOLDER, filepath))
+    
+    def _copy_seed_file(self, filepath: str):
+        self._copy_file(Path(c.SEEDS_FOLDER, filepath))
+    
+    def _copy_dashboard_file(self, filepath: str):
+        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 = 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,
+        }
+        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 = Path(base_proj_dir, TMP_FOLDER, c.MANIFEST_FILE)
+        output_path.write_text(manifest_content)
+        
+        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 = ["connections", "parameters", "build", "federate", "dashboard", "admin_password"]
+        CONNECTIONS, PARAMETERS, BUILD, FEDERATE, DASHBOARD, ADMIN_PASSWORD = options
+
+        # Add project name prompt if not provided
+        if self.project_name is None and not self.use_curr_dir:
+            questions = [
+                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):
+            """
+            If key is in answers dict as None, using `.get` on a dictionary will return None even if a default is provided.
+
+            For instance, the following prints None.
+            >>> test_dict = {"key": None}
+            >>> print(test_dict.get("key", "default"))
+
+            This function will return the default value if the key is in the dict with value None.
+            """
+            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)
+
+        parameters_format = get_answer(PARAMETERS, c.PYTHON_FORMAT)
+        parameters_use_yaml = (parameters_format == c.YML_FORMAT)
+        parameters_use_py = (parameters_format == c.PYTHON_FORMAT)
+
+        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"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"
+        elif federate_format == c.PYTHON_FILE_TYPE:
+            federate_file = c.FEDERATE_FILE_STEM + ".py"
+        else:
+            raise NotImplementedError(f"Federate model format '{federate_format}' not supported")
+
+        dashboards_enabled = get_answer(DASHBOARD, False)
+
+        self._copy_dotenv_files(admin_password)
+        self._create_manifest_file(connections_use_yaml, parameters_use_yaml)
+        
+        self._copy_gitignore_file()
+        
+        if connections_use_py:
+            self._copy_pyconfig_file(c.CONNECTIONS_FILE)
+        elif connections_use_yaml:
+            pass # already included in squirrels.yml
+        else:
+            raise NotImplementedError(f"Format '{connections_format}' not supported for configuring database connections")
+        
+        if parameters_use_py:
+            self._copy_pyconfig_file(c.PARAMETERS_FILE)
+        elif parameters_use_yaml:
+            pass # already included in squirrels.yml
+        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.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")
+        
+        self._copy_resource_file(c.EXPENSES_DB)
+        self._copy_file(Path(c.RESOURCES_FOLDER, c.PUBLIC_FOLDER, ".gitkeep"))
+        
+        print(f"\nSuccessfully created new Squirrels project!\n")
+    
+    def get_file(self, args):
+        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)
+        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 == 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}")
+            
+            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_resource_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,121 @@
+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
+from ._request_context import get_request_id
+
+
+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)
+        
+        # Append request ID if available
+        request_id = get_request_id()
+        request_id_str = f" [req_id: {request_id}]" if request_id else ""
+        formatted = formatted.replace("{request_id}", request_id_str)
+        
+        # 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)
+        request_id = get_request_id()
+        info = {
+            "timestamp": self.formatTime(record),
+            "level": record.levelname,
+            "message": record.getMessage(),
+            "request_id": request_id,
+        }
+        output = {
+            "data": record.__dict__.get("data", {}),
+            "info": info
+        }
+        return json.dumps(output)
+
+
+def get_logger(
+    project_path: str, log_to_file: bool | str, log_level: str, log_format: str, log_file_size_mb: float, log_file_backup_count: int
+) -> u.Logger:
+    logger = u.Logger(name=uuid4().hex, level=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]{request_id} %(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:
+        if isinstance(log_to_file, str):
+            log_file_path = Path(project_path, log_to_file, c.LOGS_FILE)
+        else:
+            log_file_path = Path(project_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=int(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