squirrels 0.5.0rc0__py3-none-any.whl → 0.5.1__py3-none-any.whl

squirrels/_initializer.py

@@ -1,35 +1,37 @@
 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 = 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, *, project_name: Optional[str] = None, overwrite: bool = False):
-        self.project_name = project_name
-        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)
         
-        filepath2 = u.Path(self.project_name, filepath) if self.project_name else 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)
@@ -40,8 +42,6 @@ class Initializer:
             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:
                 filepath2 = self._add_timestamp_to_filename(old_filepath)
                 extra_msg = f'Creating file as "{filepath2}" instead...'
@@ -53,38 +53,38 @@ class Initializer:
             shutil.copy(src_path, filepath2)
 
     def _copy_macros_file(self, filepath: str):
-        self._copy_file(u.Path(c.MACROS_FOLDER, filepath))
+        self._copy_file(Path(c.MACROS_FOLDER, filepath))
 
     def _copy_models_file(self, filepath: str):
-        self._copy_file(u.Path(c.MODELS_FOLDER, filepath))
+        self._copy_file(Path(c.MODELS_FOLDER, filepath))
 
     def _copy_build_file(self, filepath: str):
-        self._copy_file(u.Path(c.MODELS_FOLDER, c.BUILDS_FOLDER, filepath))
+        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))
+        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)
+            yaml_path = Path(base_proj_dir, file_name)
             return yaml_path.read_text()
         
         file_name_dict = {
@@ -95,10 +95,10 @@ class Initializer:
         
         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)
+        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 = {
@@ -106,52 +106,76 @@ class Initializer:
             "random_admin_password": admin_password if admin_password else secrets.token_urlsafe(8),
         }
 
-        dotenv_path = u.Path(base_proj_dir, c.DOTENV_FILE)
+        dotenv_path = Path(base_proj_dir, c.DOTENV_FILE)
         contents = u.render_string(dotenv_path.read_text(), **substitutions)
 
-        output_path = u.Path(base_proj_dir, TMP_FOLDER, c.DOTENV_FILE)
+        output_path = Path(base_proj_dir, TMP_FOLDER, c.DOTENV_FILE)
         output_path.write_text(contents)
 
-        self._copy_file(u.Path(c.DOTENV_FILE), src_folder=TMP_FOLDER)
-        self._copy_file(u.Path(c.DOTENV_FILE + ".example"))
+        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", "build", "federate", "dashboard"]
-        _, CONNECTIONS, PARAMETERS, BUILD, FEDERATE, DASHBOARD = options
+        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:
+        if self.project_name is None and not args.curr_dir:
             questions = [
-                inquirer.Text('project_name', message="What is your project name? (leave blank to create in current directory)")
+                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 not any(answers.values()):
-            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(
-                    BUILD, message="What's the file format for the build 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.Password(
-                    "admin_password", message="What's the admin password? (leave blank to generate a random one)"
-                ),
-            ]
-            answers = inquirer.prompt(questions)
-            assert isinstance(answers, dict)
+        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):
             """
@@ -177,7 +201,7 @@ class Initializer:
         parameters_use_py = (parameters_format == c.PYTHON_FORMAT)
 
         build_config_file = c.BUILD_FILE_STEM + ".yml"
-        build_format = get_answer(BUILD, c.PYTHON_FILE_TYPE)
+        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:
@@ -202,7 +226,7 @@ class Initializer:
         self._copy_dotenv_files(admin_password)
         self._create_manifest_file(connections_use_yaml, parameters_use_yaml)
         
-        self._copy_file(u.Path(c.GITIGNORE_FILE))
+        self._copy_gitignore_file()
         
         if connections_use_py:
             self._copy_pyconfig_file(c.CONNECTIONS_FILE)
@@ -242,7 +266,7 @@ class Initializer:
         
         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.DOTENV_FILE:
@@ -253,7 +277,7 @@ class Initializer:
             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_file(u.Path(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):

squirrels/_logging.py

@@ -0,0 +1,117 @@
+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.format(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(
+    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, 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:
+        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