bump version to 0.4.4 - Fix Redshift bugs and test failures

chuck_data/clients/redshift.py

@@ -251,7 +251,7 @@ def get_statement_result(self, statement_id: str) -> Dict:
     # Database/Schema/Table metadata methods (parallel to DatabricksAPIClient)
     #
 
-    def list_databases(self, database: Optional[str] = None) -> List[str]:
+    def list_databases(self, database: Optional[str] = None) -> Dict:
         """
         List Redshift databases.
 
@@ -261,7 +261,7 @@ def list_databases(self, database: Optional[str] = None) -> List[str]:
             database: Database name to connect to for listing (uses default if not specified)
 
         Returns:
-            List of database names
+            Dictionary containing databases list in format: {"databases": [{"name": "db1"}, ...]}
 
         Raises:
             ValueError: If an error occurs
@@ -279,7 +279,8 @@ def list_databases(self, database: Optional[str] = None) -> List[str]:
                 params["WorkgroupName"] = self.workgroup_name
 
             response = self.redshift_data.list_databases(**params)
-            return [db for db in response.get("Databases", [])]
+            # Return in same format as Databricks for consistency
+            return {"databases": [{"name": db} for db in response.get("Databases", [])]}
 
         except ClientError as e:
             logging.debug(f"Error listing databases: {e}")
@@ -288,7 +289,7 @@ def list_databases(self, database: Optional[str] = None) -> List[str]:
             logging.debug(f"Connection error: {e}")
             raise ConnectionError(f"Connection error occurred: {e}")
 
-    def list_schemas(self, database: Optional[str] = None) -> List[str]:
+    def list_schemas(self, database: Optional[str] = None) -> Dict:
         """
         List schemas in a database.
 
@@ -298,7 +299,7 @@ def list_schemas(self, database: Optional[str] = None) -> List[str]:
             database: Database name (uses default if not specified)
 
         Returns:
-            List of schema names
+            Dictionary containing schemas list in format: {"schemas": [{"name": "schema1"}, ...]}
 
         Raises:
             ValueError: If an error occurs
@@ -316,7 +317,10 @@ def list_schemas(self, database: Optional[str] = None) -> List[str]:
                 params["WorkgroupName"] = self.workgroup_name
 
             response = self.redshift_data.list_schemas(**params)
-            return [schema for schema in response.get("Schemas", [])]
+            # Return in same format as Databricks for consistency
+            return {
+                "schemas": [{"name": schema} for schema in response.get("Schemas", [])]
+            }
 
         except ClientError as e:
             logging.debug(f"Error listing schemas: {e}")

chuck_data/commands/auth.py

@@ -98,6 +98,7 @@ def handle_logout(
         needs_api_client=False,
         visible_to_user=True,
         visible_to_agent=False,
+        provider="databricks",
     ),
     CommandDefinition(
         name="logout",

chuck_data/commands/database_selection.py

@@ -84,7 +84,10 @@ def handle_command(client: Optional[RedshiftAPIClient], **kwargs) -> CommandResu
             f"Looking for database matching '{database}'", tool_output_callback
         )
 
-        databases = client.list_databases()
+        databases_result = client.list_databases()
+        database_dicts = databases_result.get("databases", [])
+        databases = [db.get("name") for db in database_dicts]
+
         if not databases:
             return CommandResult(
                 False, message="No databases found in Redshift cluster/workgroup."

chuck_data/commands/help.py

@@ -14,6 +14,8 @@
     get_user_commands,
     TUI_COMMAND_MAP,
 )
+from chuck_data.config import get_data_provider
+from chuck_data.ui.help_formatter import format_help_text
 from .base import CommandResult
 
 
@@ -26,10 +28,14 @@ def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandRe
         **kwargs: No parameters required
     """
     try:
-        from chuck_data.ui.help_formatter import format_help_text
 
-        user_commands = get_user_commands()
-        help_text = format_help_text(user_commands, TUI_COMMAND_MAP)
+        # Get current provider to filter commands appropriately
+        current_provider = get_data_provider()
+
+        user_commands = get_user_commands(provider=current_provider)
+        help_text = format_help_text(
+            user_commands, TUI_COMMAND_MAP, provider=current_provider
+        )
         return CommandResult(True, data={"help_text": help_text})
     except Exception as e:
         logging.error(f"Error generating help: {e}", exc_info=True)

chuck_data/commands/list_databases.py

@@ -36,7 +36,8 @@ def handle_command(client: Optional[RedshiftAPIClient], **kwargs: Any) -> Comman
 
     try:
         # List databases in Redshift
-        databases = client.list_databases()
+        databases_result = client.list_databases()
+        databases = databases_result.get("databases", [])
 
         if not databases:
             return CommandResult(
@@ -50,13 +51,8 @@ def handle_command(client: Optional[RedshiftAPIClient], **kwargs: Any) -> Comman
                 },
             )
 
-        # Format database information for display
-        formatted_databases = []
-        for db_name in databases:
-            formatted_database = {
-                "name": db_name,
-            }
-            formatted_databases.append(formatted_database)
+        # Format database information for display (databases already have "name" key)
+        formatted_databases = databases
 
         return CommandResult(
             True,

chuck_data/commands/list_redshift_schemas.py

@@ -49,7 +49,8 @@ def handle_command(client: Optional[RedshiftAPIClient], **kwargs: Any) -> Comman
 
     try:
         # List schemas in the database
-        schemas = client.list_schemas(database=database)
+        schemas_result = client.list_schemas(database=database)
+        schemas = schemas_result.get("schemas", [])
 
         if not schemas:
             return CommandResult(
@@ -64,11 +65,11 @@ def handle_command(client: Optional[RedshiftAPIClient], **kwargs: Any) -> Comman
                 },
             )
 
-        # Format schema information for display
+        # Format schema information for display (add database to each schema)
         formatted_schemas = []
-        for schema_name in schemas:
+        for schema_dict in schemas:
             formatted_schema = {
-                "name": schema_name,
+                "name": schema_dict.get("name"),
                 "database": database,
             }
             formatted_schemas.append(formatted_schema)

chuck_data/commands/list_tables.py

@@ -128,6 +128,37 @@ def handle_command(
                             }
                             logging.info(f"Table {table_name}: {column_count} columns")
 
+                    # If the SQL query succeeded but returned no rows, fall back to describe_table
+                    # This can happen if pg_table_def is not accessible or empty
+                    if not table_metadata:
+                        logging.warning(
+                            f"SQL query to pg_table_def returned 0 rows, falling back to describe_table for each table"
+                        )
+                        for table in result_tables:
+                            table_name = table.get("name")
+                            try:
+                                table_details = client.describe_table(
+                                    database=database,
+                                    schema=schema_name,
+                                    table=table_name,
+                                )
+                                columns = table_details.get("ColumnList", [])
+                                table_metadata[table_name] = {
+                                    "column_count": len(columns),
+                                    "row_count": "-",
+                                }
+                                logging.info(
+                                    f"Table {table_name}: {len(columns)} columns (from describe_table)"
+                                )
+                            except Exception as e2:
+                                logging.error(
+                                    f"Failed to fetch columns for table {table_name}: {str(e2)}"
+                                )
+                                table_metadata[table_name] = {
+                                    "column_count": 0,
+                                    "row_count": "-",
+                                }
+                    else:
                         # Note: Redshift doesn't provide table creation/modification timestamps
                         # in accessible system tables without special permissions (STL_DDLTEXT requires elevated access)
 
@@ -138,7 +169,7 @@ def handle_command(
                         )
                         for table_name in table_metadata.keys():
                             try:
-                                count_sql = f'SELECT COUNT(*) as row_count FROM "{database}"."{schema_name}"."{table_name}"'
+                                count_sql = f'SELECT COUNT(*) as row_count FROM "{schema_name}"."{table_name}"'
                                 count_result = client.execute_sql(
                                     count_sql, database=database
                                 )

chuck_data/commands/redshift_schema_selection.py

@@ -100,7 +100,10 @@ def handle_command(
             tool_output_callback,
         )
 
-        schemas = client.list_schemas(database=database)
+        schemas_result = client.list_schemas(database=database)
+        schema_dicts = schemas_result.get("schemas", [])
+        schemas = [s.get("name") for s in schema_dicts]
+
         if not schemas:
             return CommandResult(
                 False,

chuck_data/commands/redshift_status.py

@@ -59,7 +59,8 @@ def handle_command(client: Optional[RedshiftAPIClient], **kwargs) -> CommandResu
         if client:
             try:
                 # Test connection by attempting to list databases
-                databases = client.list_databases()
+                databases_result = client.list_databases()
+                databases = databases_result.get("databases", [])
                 data["connection_status"] = (
                     f"Connected (found {len(databases)} database(s))."
                 )

chuck_data/commands/wizard/steps.py

@@ -1224,7 +1224,8 @@ def handle_input(self, input_text: str, state: WizardState) -> StepResult:
                 )
                 serverless_config = {**client_config, "workgroup_name": identifier}
                 client = RedshiftAPIClient(**serverless_config)
-                databases = client.list_databases()
+                databases_result = client.list_databases()
+                databases = databases_result.get("databases", [])
                 logging.info(
                     f"Successfully connected to Redshift Serverless workgroup. Found {len(databases)} databases."
                 )
@@ -1242,7 +1243,8 @@ def handle_input(self, input_text: str, state: WizardState) -> StepResult:
                     )
                     cluster_config = {**client_config, "cluster_identifier": identifier}
                     client = RedshiftAPIClient(**cluster_config)
-                    databases = client.list_databases()
+                    databases_result = client.list_databases()
+                    databases = databases_result.get("databases", [])
                     logging.info(
                         f"Successfully connected to Redshift provisioned cluster. Found {len(databases)} databases."
                     )

chuck_data/commands/workspace_selection.py

@@ -72,4 +72,5 @@ def handle_command(client: Optional[DatabricksAPIClient], **kwargs) -> CommandRe
     tui_aliases=["/select-workspace"],
     visible_to_user=True,
     visible_to_agent=False,  # Agent doesn't need to select workspace
+    provider="databricks",
 )

chuck_data/config.py

@@ -202,15 +202,57 @@ def get_config(self) -> ChuckConfig:
         return self.load()
 
     def needs_setup(self) -> bool:
-        """Check if first-time setup is needed based on missing critical configuration"""
+        """Check if first-time setup is needed based on missing critical configuration.
+
+        This function is provider-aware and only checks for configs relevant to the
+        configured data provider. If no provider is set, it returns True (needs setup).
+        """
         config = self.load()
-        critical_configs = [
-            config.amperity_token,
-            config.databricks_token,
-            config.workspace_url,
-            config.active_model,
-        ]
-        return any(item is None or item == "" for item in critical_configs)
+
+        # FIRST: Check if logged in (amperity_token)
+        # This is required regardless of provider
+        if not config.amperity_token or config.amperity_token == "":
+            return True
+
+        # Always check for model
+        if not config.active_model:
+            return True
+
+        # If no data provider is set, we definitely need setup
+        provider = config.data_provider
+        if not provider:
+            return True
+
+        # Check provider-specific configs
+        if provider == "databricks":
+            # Databricks requires workspace_url and databricks_token
+            workspace_url = getattr(config, "workspace_url", None)
+            databricks_token = getattr(config, "databricks_token", None)
+
+            critical_configs = [workspace_url, databricks_token]
+            return any(item is None or item == "" for item in critical_configs)
+
+        elif provider == "aws_redshift":
+            # Redshift requires AWS configs and (optionally) amperity_token
+            # Check required AWS configs
+            aws_region = getattr(config, "aws_region", None)
+            if not aws_region or aws_region == "":
+                return True
+
+            # Either cluster_identifier or workgroup_name must be set
+            cluster_id = getattr(config, "redshift_cluster_identifier", None)
+            workgroup = getattr(config, "redshift_workgroup_name", None)
+
+            has_cluster_or_workgroup = (cluster_id and cluster_id != "") or (
+                workgroup and workgroup != ""
+            )
+            if not has_cluster_or_workgroup:
+                return True
+
+            return False
+
+        # Unknown provider - needs setup
+        return True
 
     def update(self, **kwargs) -> bool:
         """Update configuration values"""

chuck_data/llm/factory.py

@@ -74,10 +74,20 @@ def _get_provider_config(provider_name: str) -> dict:
             from chuck_data.config import get_config_manager
 
             chuck_config = get_config_manager().get_config()
+
+            # Get provider-specific config (if any)
             if hasattr(chuck_config, "llm_provider_config"):
                 provider_configs = chuck_config.llm_provider_config or {}
                 config = provider_configs.get(provider_name, {})
                 logger.debug(f"Loaded config for {provider_name}")
+
+            # Add active_model to config if set (overrides provider-specific model_id)
+            if hasattr(chuck_config, "active_model") and chuck_config.active_model:
+                config["model_id"] = chuck_config.active_model
+                logger.debug(
+                    f"Using active_model from config: {chuck_config.active_model}"
+                )
+
         except Exception as e:
             logger.debug(f"Could not load provider config: {e}")