Athena refactor

Author: Maxteabag

sqlit/db/schema.py

@@ -439,7 +439,8 @@ def _get_oracle_role_options() -> tuple[SelectOption, ...]:
 )
 
 
-def _get_supabase_region_options() -> tuple[SelectOption, ...]:
+def _get_aws_region_options() -> tuple[SelectOption, ...]:
+    """AWS regions shared by Supabase, Athena, and other AWS-based services."""
     regions = (
         "us-east-1",
         "us-east-2",
@@ -470,7 +471,7 @@ def _get_supabase_region_options() -> tuple[SelectOption, ...]:
             name="supabase_region",
             label="Region",
             field_type=FieldType.DROPDOWN,
-            options=_get_supabase_region_options(),
+            options=_get_aws_region_options(),
             required=True,
             default="us-east-1",
         ),
@@ -553,38 +554,38 @@ def _get_supabase_region_options() -> tuple[SelectOption, ...]:
 )
 
 
+def _get_athena_auth_options() -> tuple[SelectOption, ...]:
+    return (
+        SelectOption("profile", "AWS Profile"),
+        SelectOption("keys", "Access Keys"),
+    )
+
+
+def _athena_auth_is_profile(v: dict) -> bool:
+    return v.get("athena_auth_method", "profile") == "profile"
+
+
+def _athena_auth_is_keys(v: dict) -> bool:
+    return v.get("athena_auth_method") == "keys"
+
+
 ATHENA_SCHEMA = ConnectionSchema(
     db_type="athena",
     display_name="AWS Athena",
     fields=(
         SchemaField(
             name="athena_region_name",
             label="Region",
+            field_type=FieldType.DROPDOWN,
+            options=_get_aws_region_options(),
             required=True,
             default="us-east-1",
         ),
-        SchemaField(
-            name="athena_work_group",
-            label="WorkGroup",
-            required=True,
-            default="primary",
-            description="Athena WorkGroup",
-        ),
-        SchemaField(
-            name="athena_s3_staging_dir",
-            label="S3 Staging Dir",
-            placeholder="s3://your-bucket/path/",
-            required=True,
-            description="S3 location for query results",
-        ),
         SchemaField(
             name="athena_auth_method",
-            label="Auth Method",
-            field_type=FieldType.SELECT,
-            options=(
-                SelectOption("profile", "AWS Profile"),
-                SelectOption("keys", "Access Keys"),
-            ),
+            label="Authentication",
+            field_type=FieldType.DROPDOWN,
+            options=_get_athena_auth_options(),
             default="profile",
         ),
         SchemaField(
@@ -594,15 +595,15 @@ def _get_supabase_region_options() -> tuple[SelectOption, ...]:
             required=True,
             default="default",
             description="AWS CLI profile name",
-            visible_when=lambda v: v.get("athena_auth_method") == "profile",
+            visible_when=_athena_auth_is_profile,
         ),
         SchemaField(
             name="username",
             label="Access Key",
             placeholder="AWS Access Key ID",
             required=True,
             group="credentials",
-            visible_when=lambda v: v.get("athena_auth_method") == "keys",
+            visible_when=_athena_auth_is_keys,
         ),
         SchemaField(
             name="password",
@@ -611,10 +612,25 @@ def _get_supabase_region_options() -> tuple[SelectOption, ...]:
             placeholder="AWS Secret Access Key",
             required=True,
             group="credentials",
-            visible_when=lambda v: v.get("athena_auth_method") == "keys",
+            visible_when=_athena_auth_is_keys,
+        ),
+        SchemaField(
+            name="athena_work_group",
+            label="WorkGroup",
+            required=True,
+            default="primary",
+            description="Athena WorkGroup",
+        ),
+        SchemaField(
+            name="athena_s3_staging_dir",
+            label="S3 Staging Dir",
+            placeholder="s3://your-bucket/path/",
+            required=True,
+            description="S3 location for query results",
         ),
     ),
     supports_ssh=False,
+    has_advanced_auth=True,
 )