sqlsaber 0.30.2__py3-none-any.whl → 0.32.0__py3-none-any.whl

sqlsaber/cli/auth.py

@@ -11,10 +11,12 @@ from sqlsaber.config.api_keys import APIKeyManager
 from sqlsaber.config.auth import AuthConfigManager, AuthMethod
 from sqlsaber.config.oauth_tokens import OAuthTokenManager
 from sqlsaber.theme.manager import create_console
+from sqlsaber.config.logging import get_logger
 
 # Global instances for CLI commands
 console = create_console()
 config_manager = AuthConfigManager()
+logger = get_logger(__name__)
 
 # Create the authentication management CLI app
 auth_app = cyclopts.App(
@@ -46,7 +48,9 @@ def setup():
         )
         return success, provider
 
-    success, _ = asyncio.run(run_setup())
+    logger.info("auth.setup.start")
+    success, provider = asyncio.run(run_setup())
+    logger.info("auth.setup.complete", success=bool(success), provider=str(provider))
 
     if not success:
         console.print("\n[warning]No authentication configured.[/warning]")
@@ -59,6 +63,7 @@ def setup():
 @auth_app.command
 def status():
     """Show current authentication configuration and provider key status."""
+    logger.info("auth.status.start")
     auth_method = config_manager.get_auth_method()
 
     console.print("\n[bold blue]Authentication Status[/bold blue]")
@@ -68,6 +73,7 @@ def status():
         console.print(
             "Run [primary]saber auth setup[/primary] to configure authentication."
         )
+        logger.info("auth.status.none_configured")
         return
 
     # Show configured method summary
@@ -93,6 +99,7 @@ def status():
             console.print(f"> {provider}: [green]configured[/green]")
         else:
             console.print(f"> {provider}: [warning]not configured[/warning]")
+    logger.info("auth.status.complete", method=str(auth_method))
 
 
 @auth_app.command
@@ -108,6 +115,7 @@ def reset():
 
     if provider is None:
         console.print("[warning]Reset cancelled.[/warning]")
+        logger.info("auth.reset.cancelled_no_provider")
         return
 
     api_key_manager = APIKeyManager()
@@ -125,6 +133,7 @@ def reset():
         console.print(
             f"[warning]No stored credentials found for {provider}. Nothing to reset.[/warning]"
         )
+        logger.info("auth.reset.nothing_to_reset", provider=provider)
         return
 
     # Build confirmation message
@@ -142,6 +151,7 @@ def reset():
 
     if not confirmed:
         console.print("Reset cancelled.")
+        logger.info("auth.reset.cancelled_confirm", provider=provider)
         return
 
     # Perform deletions
@@ -151,11 +161,13 @@ def reset():
         try:
             keyring.delete_password(service, provider)
             console.print(f"Removed {provider} API key from keyring", style="green")
+            logger.info("auth.reset.api_key_removed", provider=provider)
         except keyring.errors.PasswordDeleteError:
             # Already absent; treat as success
             pass
         except Exception as e:
             console.print(f"Warning: Could not remove API key: {e}", style="warning")
+            logger.warning("auth.reset.api_key_remove_failed", provider=provider, error=str(e))
 
     # Optionally clear global auth method if removing Anthropic OAuth configuration
     if provider == "anthropic" and oauth_present:
@@ -170,8 +182,10 @@ def reset():
                 config["auth_method"] = None
                 config_manager._save_config(config)
                 console.print("Global auth method unset.", style="green")
+                logger.info("auth.reset.global_method_unset")
 
     console.print("\n[success]✓ Reset complete.[/success]")
+    logger.info("auth.reset.complete", provider=provider)
     console.print(
         "Environment variables are not modified by this command.", style="dim"
     )

sqlsaber/cli/commands.py

@@ -1,6 +1,7 @@
 """CLI command definitions and handlers."""
 
 import asyncio
+import os
 import sys
 from typing import Annotated
 
@@ -16,6 +17,7 @@ from sqlsaber.cli.threads import create_threads_app
 
 # Lazy imports - only import what's needed for CLI parsing
 from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.config.logging import get_logger, setup_logging
 from sqlsaber.theme.manager import create_console
 
 
@@ -42,6 +44,52 @@ app.command(create_threads_app(), name="threads")
 console = create_console()
 config_manager = DatabaseConfigManager()
 
+_MLFLOW_CONFIGURED = False
+
+
+def _maybe_configure_mlflow(log) -> bool:
+    """Enable mlflow autologging when environment variables are present."""
+    global _MLFLOW_CONFIGURED
+    if _MLFLOW_CONFIGURED:
+        return True
+
+    tracking_uri = os.getenv("MLFLOW_URI")
+    experiment = os.getenv("MLFLOW_EXP")
+    if not tracking_uri and not experiment:
+        return False
+
+    try:
+        import mlflow
+    except ModuleNotFoundError:
+        log.warning(
+            "mlflow.setup.skipped",
+            reason="mlflow package not installed",
+            uri=tracking_uri,
+            experiment=experiment,
+        )
+        return False
+
+    try:
+        mlflow.pydantic_ai.autolog()
+    except Exception:
+        log.warning("mlflow.autolog.failed", exc_info=True)
+    try:
+        if tracking_uri:
+            mlflow.set_tracking_uri(tracking_uri)
+        if experiment:
+            mlflow.set_experiment(experiment)
+    except Exception:
+        log.warning("mlflow.setup.failed", exc_info=True)
+        return False
+
+    _MLFLOW_CONFIGURED = True
+    log.info(
+        "mlflow.setup.enabled",
+        uri=tracking_uri,
+        experiment=experiment,
+    )
+    return True
+
 
 @app.meta.default
 def meta_handler(
@@ -111,6 +159,14 @@ def query(
     """
 
     async def run_session():
+        log = get_logger(__name__)
+        log.info(
+            "cli.session.start",
+            argv=sys.argv[1:],
+            database=database,
+            has_query=query_text is not None,
+            thinking=thinking,
+        )
         # Import heavy dependencies only when actually running a query
         # This is only done to speed up startup time
         from sqlsaber.agents import SQLSaberAgent
@@ -134,29 +190,39 @@ def query(
         # Check if onboarding is needed (only for interactive mode or when no database is configured)
         if needs_onboarding(database):
             # Run onboarding flow
+            log.debug("cli.onboarding.start")
             onboarding_success = await run_onboarding()
             if not onboarding_success:
                 # User cancelled or onboarding failed
                 raise CLIError(
                     "Setup incomplete. Please configure your database and try again."
                 )
+            log.info("cli.onboarding.complete", success=True)
 
         # Resolve database from CLI input
         try:
             resolved = resolve_database(database, config_manager)
             connection_string = resolved.connection_string
             db_name = resolved.name
+            log.info(
+                "db.resolve.success",
+                name=db_name,
+            )
         except DatabaseResolutionError as e:
+            log.error("db.resolve.error", error=str(e))
             raise CLIError(str(e))
 
         # Create database connection
         try:
             db_conn = DatabaseConnection(connection_string)
+            log.info("db.connection.created", db_type=type(db_conn).__name__)
         except Exception as e:
+            log.exception("db.connection.error", error=str(e))
             raise CLIError(f"Error creating database connection: {e}")
 
         # Create pydantic-ai agent instance with database name for memory context
         sqlsaber_agent = SQLSaberAgent(db_conn, db_name, thinking_enabled=thinking)
+        _maybe_configure_mlflow(log)
 
         try:
             if actual_query:
@@ -166,9 +232,11 @@ def query(
                 console.print(
                     f"[primary]Connected to:[/primary] {db_name} ({db_type})\n"
                 )
+                log.info("query.execute.start", db_name=db_name, db_type=db_type)
                 run = await streaming_handler.execute_streaming_query(
                     actual_query, sqlsaber_agent
                 )
+
                 # Persist non-interactive run as a thread snapshot so it can be resumed later
                 try:
                     if run is not None:
@@ -187,8 +255,10 @@ def query(
                         console.print(
                             f"[dim]You can continue this thread using:[/dim] saber threads resume {thread_id}"
                         )
+                        log.info("thread.save.success", thread_id=thread_id)
                 except Exception:
                     # best-effort persistence; don't fail the CLI on storage errors
+                    log.warning("thread.save.failed", exc_info=True)
                     pass
                 finally:
                     await threads.prune_threads()
@@ -200,16 +270,20 @@ def query(
         finally:
             # Clean up
             await db_conn.close()
+            log.info("db.connection.closed")
             console.print("\n[success]Goodbye![/success]")
 
     # Run the async function with proper error handling
     try:
         asyncio.run(run_session())
     except CLIError as e:
+        get_logger(__name__).error("cli.error", error=str(e))
         console.print(f"[error]Error:[/error] {e}")
         sys.exit(e.exit_code)
 
 
 def main():
     """Entry point for the CLI application."""
+    setup_logging()
+    get_logger(__name__).info("cli.start")
     app()

sqlsaber/cli/database.py

@@ -11,11 +11,13 @@ import questionary
 from rich.table import Table
 
 from sqlsaber.config.database import DatabaseConfig, DatabaseConfigManager
+from sqlsaber.config.logging import get_logger
 from sqlsaber.theme.manager import create_console
 
 # Global instances for CLI commands
 console = create_console()
 config_manager = DatabaseConfigManager()
+logger = get_logger(__name__)
 
 # Create the database management CLI app
 db_app = cyclopts.App(
@@ -78,6 +80,13 @@ def add(
     ] = True,
 ):
     """Add a new database connection."""
+    logger.info(
+        "db.add.start",
+        name=name,
+        type=type,
+        interactive=bool(interactive),
+        has_password=False,
+    )
 
     if interactive:
         # Interactive mode - prompt for all required fields
@@ -96,6 +105,7 @@ def add(
 
         if db_input is None:
             console.print("[warning]Operation cancelled[/warning]")
+            logger.info("db.add.cancelled")
             return
 
         # Extract values from db_input
@@ -116,6 +126,7 @@ def add(
                 console.print(
                     "[bold error]Error:[/bold error] Database file path is required for SQLite"
                 )
+                logger.error("db.add.missing_path", db_type="sqlite")
                 sys.exit(1)
             host = "localhost"
             port = 0
@@ -126,6 +137,7 @@ def add(
                 console.print(
                     "[bold error]Error:[/bold error] Database file path is required for DuckDB"
                 )
+                logger.error("db.add.missing_path", db_type="duckdb")
                 sys.exit(1)
             database = str(Path(database).expanduser().resolve())
             host = "localhost"
@@ -137,6 +149,7 @@ def add(
                 console.print(
                     "[bold error]Error:[/bold error] Host, database, and username are required"
                 )
+                logger.error("db.add.missing_fields")
                 sys.exit(1)
 
             if port is None:
@@ -173,12 +186,15 @@ def add(
         # Add the configuration
         config_manager.add_database(db_config, password if password else None)
         console.print(f"[green]Successfully added database connection '{name}'[/green]")
+        logger.info("db.add.success", name=name, type=type)
 
         # Set as default if it's the first one
         if len(config_manager.list_databases()) == 1:
             console.print(f"[blue]Set '{name}' as default database[/blue]")
+            logger.info("db.default.set", name=name)
 
     except Exception as e:
+        logger.exception("db.add.error", name=name, error=str(e))
         console.print(f"[bold error]Error adding database:[/bold error] {e}")
         sys.exit(1)
 
@@ -186,12 +202,14 @@ def add(
 @db_app.command
 def list():
     """List all configured database connections."""
+    logger.info("db.list.start")
     databases = config_manager.list_databases()
     default_name = config_manager.get_default_name()
 
     if not databases:
         console.print("[warning]No database connections configured[/warning]")
         console.print("Use 'sqlsaber db add <name>' to add a database connection")
+        logger.info("db.list.empty")
         return
 
     table = Table(title="Database Connections")
@@ -228,6 +246,7 @@ def list():
         )
 
     console.print(table)
+    logger.info("db.list.complete", count=len(databases))
 
 
 @db_app.command
@@ -237,10 +256,12 @@ def remove(
     ],
 ):
     """Remove a database connection."""
+    logger.info("db.remove.start", name=name)
     if not config_manager.get_database(name):
         console.print(
             f"[bold error]Error:[/bold error] Database connection '{name}' not found"
         )
+        logger.error("db.remove.not_found", name=name)
         sys.exit(1)
 
     if questionary.confirm(
@@ -250,13 +271,16 @@ def remove(
             console.print(
                 f"[green]Successfully removed database connection '{name}'[/green]"
             )
+            logger.info("db.remove.success", name=name)
         else:
             console.print(
                 f"[bold error]Error:[/bold error] Failed to remove database connection '{name}'"
             )
+            logger.error("db.remove.failed", name=name)
             sys.exit(1)
     else:
         console.print("Operation cancelled")
+        logger.info("db.remove.cancelled", name=name)
 
 
 @db_app.command
@@ -267,18 +291,22 @@ def set_default(
     ],
 ):
     """Set the default database connection."""
+    logger.info("db.default.start", name=name)
     if not config_manager.get_database(name):
         console.print(
             f"[bold error]Error:[/bold error] Database connection '{name}' not found"
         )
+        logger.error("db.default.not_found", name=name)
         sys.exit(1)
 
     if config_manager.set_default_database(name):
         console.print(f"[green]Successfully set '{name}' as default database[/green]")
+        logger.info("db.default.success", name=name)
     else:
         console.print(
             f"[bold error]Error:[/bold error] Failed to set '{name}' as default"
         )
+        logger.error("db.default.failed", name=name)
         sys.exit(1)
 
 
@@ -292,6 +320,7 @@ def test(
     ] = None,
 ):
     """Test a database connection."""
+    logger.info("db.test.start")
 
     async def test_connection():
         # Lazy import to keep CLI startup fast
@@ -303,6 +332,7 @@ def test(
                 console.print(
                     f"[bold error]Error:[/bold error] Database connection '{name}' not found"
                 )
+                logger.error("db.test.not_found", name=name)
                 sys.exit(1)
         else:
             db_config = config_manager.get_default_database()
@@ -313,6 +343,7 @@ def test(
                 console.print(
                     "Use 'sqlsaber db add <name>' to add a database connection"
                 )
+                logger.error("db.test.no_default")
                 sys.exit(1)
 
         console.print(f"[blue]Testing connection to '{db_config.name}'...[/blue]")
@@ -328,8 +359,16 @@ def test(
             console.print(
                 f"[green]✓ Connection to '{db_config.name}' successful[/green]"
             )
+            logger.info("db.test.success", name=db_config.name)
 
         except Exception as e:
+            logger.exception(
+                "db.test.failed",
+                name=(
+                    db_config.name if "db_config" in locals() and db_config else name
+                ),
+                error=str(e),
+            )
             console.print(f"[bold error]✗ Connection failed:[/bold error] {e}")
             sys.exit(1)
 

sqlsaber/cli/interactive.py

@@ -30,6 +30,7 @@ from sqlsaber.database import (
 from sqlsaber.database.schema import SchemaManager
 from sqlsaber.theme.manager import get_theme_manager
 from sqlsaber.threads import ThreadStorage
+from sqlsaber.config.logging import get_logger
 
 if TYPE_CHECKING:
     from sqlsaber.agents.pydantic_ai_agent import SQLSaberAgent
@@ -66,6 +67,7 @@ class InteractiveSession:
         self._threads = ThreadStorage()
         self._thread_id: str | None = initial_thread_id
         self.first_message = not self._thread_id
+        self.log = get_logger(__name__)
 
     def _history_path(self) -> Path:
         """Get the history file path, ensuring directory exists."""
@@ -240,6 +242,7 @@ class InteractiveSession:
 
     async def _execute_query_with_cancellation(self, user_query: str):
         """Execute a query with cancellation support."""
+        self.log.info("interactive.query.start", database=self.database_name)
         # Create cancellation token
         self.cancellation_token = asyncio.Event()
 
@@ -276,16 +279,18 @@ class InteractiveSession:
                             model_name=self.sqlsaber_agent.agent.model.model_name,
                         )
                         self.first_message = False
-                except Exception:
-                    pass
+                except Exception as e:
+                    self.log.warning("interactive.thread.save_failed", error=str(e))
                 finally:
                     await self._threads.prune_threads()
         finally:
             self.current_task = None
             self.cancellation_token = None
+            self.log.info("interactive.query.end")
 
     async def run(self):
         """Run the interactive session loop."""
+        self.log.info("interactive.start", database=self.database_name)
         self.show_welcome_message()
         await self.before_prompt_loop()
 
@@ -348,3 +353,4 @@ class InteractiveSession:
                 break
             except Exception as exc:
                 self.console.print(f"[error]Error:[/error] {exc}")
+                self.log.exception("interactive.error", error=str(exc))

sqlsaber/cli/memory.py

@@ -8,6 +8,7 @@ import questionary
 from rich.table import Table
 
 from sqlsaber.config.database import DatabaseConfigManager
+from sqlsaber.config.logging import get_logger
 from sqlsaber.memory.manager import MemoryManager
 from sqlsaber.theme.manager import create_console
 
@@ -15,6 +16,7 @@ from sqlsaber.theme.manager import create_console
 console = create_console()
 config_manager = DatabaseConfigManager()
 memory_manager = MemoryManager()
+logger = get_logger(__name__)
 
 # Create the memory management CLI app
 memory_app = cyclopts.App(
@@ -31,6 +33,7 @@ def _get_database_name(database: str | None = None) -> str:
             console.print(
                 f"[bold error]Error:[/bold error] Database connection '{database}' not found."
             )
+            logger.error("memory.db.not_found", database=database)
             sys.exit(1)
         return database
     else:
@@ -40,6 +43,7 @@ def _get_database_name(database: str | None = None) -> str:
                 "[bold error]Error:[/bold error] No database connections configured."
             )
             console.print("Use 'sqlsaber db add <name>' to add a database connection.")
+            logger.error("memory.db.none_configured")
             sys.exit(1)
         return db_config.name
 
@@ -57,14 +61,17 @@ def add(
 ):
     """Add a new memory for the specified database."""
     database_name = _get_database_name(database)
+    logger.info("memory.add.start", database=database_name)
 
     try:
         memory = memory_manager.add_memory(database_name, content)
         console.print(f"[green]✓ Memory added for database '{database_name}'[/green]")
         console.print(f"[dim]Memory ID:[/dim] {memory.id}")
         console.print(f"[dim]Content:[/dim] {memory.content}")
+        logger.info("memory.add.success", database=database_name, id=memory.id)
     except Exception as e:
         console.print(f"[bold error]Error adding memory:[/bold error] {e}")
+        logger.exception("memory.add.error", database=database_name, error=str(e))
         sys.exit(1)
 
 
@@ -80,6 +87,7 @@ def list(
 ):
     """List all memories for the specified database."""
     database_name = _get_database_name(database)
+    logger.info("memory.list.start", database=database_name)
 
     memories = memory_manager.get_memories(database_name)
 
@@ -88,6 +96,7 @@ def list(
             f"[warning]No memories found for database '{database_name}'[/warning]"
         )
         console.print("Use 'sqlsaber memory add \"<content>\"' to add memories")
+        logger.info("memory.list.empty", database=database_name)
         return
 
     table = Table(title=f"Memories for Database: {database_name}")
@@ -105,6 +114,7 @@ def list(
 
     console.print(table)
     console.print(f"\n[dim]Total memories: {len(memories)}[/dim]")
+    logger.info("memory.list.complete", database=database_name, count=len(memories))
 
 
 @memory_app.command
@@ -120,6 +130,7 @@ def show(
 ):
     """Show the full content of a specific memory."""
     database_name = _get_database_name(database)
+    logger.info("memory.show.start", database=database_name, id=memory_id)
 
     memory = memory_manager.get_memory_by_id(database_name, memory_id)
 
@@ -127,6 +138,7 @@ def show(
         console.print(
             f"[bold error]Error:[/bold error] Memory with ID '{memory_id}' not found for database '{database_name}'"
         )
+        logger.error("memory.show.not_found", database=database_name, id=memory_id)
         sys.exit(1)
 
     console.print(f"[bold]Memory ID:[/bold] {memory.id}")
@@ -149,6 +161,7 @@ def remove(
 ):
     """Remove a specific memory by ID."""
     database_name = _get_database_name(database)
+    logger.info("memory.remove.start", database=database_name, id=memory_id)
 
     # First check if memory exists
     memory = memory_manager.get_memory_by_id(database_name, memory_id)
@@ -156,6 +169,7 @@ def remove(
         console.print(
             f"[bold error]Error:[/bold error] Memory with ID '{memory_id}' not found for database '{database_name}'"
         )
+        logger.error("memory.remove.not_found", database=database_name, id=memory_id)
         sys.exit(1)
 
     # Show memory content before removal
@@ -166,10 +180,12 @@ def remove(
         console.print(
             f"[green]✓ Memory removed from database '{database_name}'[/green]"
         )
+        logger.info("memory.remove.success", database=database_name, id=memory_id)
     else:
         console.print(
             f"[bold error]Error:[/bold error] Failed to remove memory '{memory_id}'"
         )
+        logger.error("memory.remove.failed", database=database_name, id=memory_id)
         sys.exit(1)
 
 
@@ -192,6 +208,7 @@ def clear(
 ):
     """Clear all memories for the specified database."""
     database_name = _get_database_name(database)
+    logger.info("memory.clear.start", database=database_name, force=bool(force))
 
     # Count memories first
     memories_count = len(memory_manager.get_memories(database_name))
@@ -200,6 +217,7 @@ def clear(
         console.print(
             f"[warning]No memories to clear for database '{database_name}'[/warning]"
         )
+        logger.info("memory.clear.nothing", database=database_name)
         return
 
     if not force:
@@ -210,12 +228,14 @@ def clear(
 
         if not questionary.confirm("Are you sure you want to proceed?").ask():
             console.print("Operation cancelled")
+            logger.info("memory.clear.cancelled", database=database_name)
             return
 
     cleared_count = memory_manager.clear_memories(database_name)
     console.print(
         f"[green]✓ Cleared {cleared_count} memories for database '{database_name}'[/green]"
     )
+    logger.info("memory.clear.success", database=database_name, deleted=cleared_count)
 
 
 @memory_app.command
@@ -230,6 +250,7 @@ def summary(
 ):
     """Show memory summary for the specified database."""
     database_name = _get_database_name(database)
+    logger.info("memory.summary.start", database=database_name)
 
     summary = memory_manager.get_memories_summary(database_name)
 
@@ -240,6 +261,11 @@ def summary(
         console.print("\n[bold]Recent memories:[/bold]")
         for memory in summary["memories"][-5:]:  # Show last 5 memories
             console.print(f"[dim]{memory['timestamp']}[/dim] - {memory['content']}")
+    logger.info(
+        "memory.summary.complete",
+        database=database_name,
+        total=summary["total_memories"],
+    )
 
 
 def create_memory_app() -> cyclopts.App: