11import requests
22import json
3- from src .config import app_config
43
4+ from starlette .exceptions import HTTPException
5+ from fastmcp .server .dependencies import get_http_request
56
6- def __set_organzation_id ():
7- """
8- Set the organization ID for the current session.
9- """
10- if not app_config .is_organization_selected ():
11- select_organization ()
7+ from src .config .config import get_settings
128
139
1410def __query_graphql_organizations ():
@@ -18,7 +14,9 @@ def __query_graphql_organizations():
1814 Returns:
1915 List of organizations with their IDs and names
2016 """
21- graphql_endpoint = app_config .settings .singlestore_graphql_public_endpoint
17+ settings = get_settings ()
18+
19+ graphql_endpoint = settings .graphql_public_endpoint
2220
2321 # GraphQL query for organizations
2422 query = """
@@ -32,7 +30,7 @@ def __query_graphql_organizations():
3230
3331 # Headers with authentication
3432 headers = {
35- "Authorization" : f"Bearer { app_config . get_auth_token ()} " ,
33+ "Authorization" : f"Bearer { __get_access_token ()} " ,
3634 "Content-Type" : "application/json" ,
3735 }
3836
@@ -72,56 +70,12 @@ def __query_graphql_organizations():
7270 raise ValueError (f"Failed to query organizations: { str (e )} " )
7371
7472
75- def select_organization ():
76- """
77- Query available organizations and prompt the user to select one.
78-
79- This must be called after authentication and before making other API calls.
80- Sets the organization ID and name in the app_config.
81-
82- Returns:
83- Dictionary with the selected organization ID and name
84- """
85-
86- print ("select_org: " , app_config .organization_id )
87- # If organization is already selected, return it
88- if app_config .is_organization_selected ():
89- return {
90- "orgID" : app_config .organization_id ,
91- "name" : app_config .organization_name ,
92- }
93-
94- # Get available organizations
95- organizations = __query_graphql_organizations ()
96-
97- if not organizations :
98- raise ValueError ("No organizations found. Please check your account access." )
99-
100- # If only one organization is available, select it automatically
101- if len (organizations ) == 1 :
102- org = organizations [0 ]
103- app_config .set_organization (org ["orgID" ], org ["name" ])
104-
105- return {
106- "orgID" : app_config .organization_id ,
107- "name" : app_config .organization_name ,
108- }
109-
110- # Create a formatted list of organizations for the user to choose from
111- org_list = "\n " .join (
112- [
113- f"{ i + 1 } . { org ['name' ]} (ID: { org ['orgID' ]} )"
114- for i , org in enumerate (organizations )
115- ]
116- )
117-
118- # This will be handled by the LLM to ask the user which organization to use
119- raise ValueError (
120- f"Multiple organizations found. Please ask the user to select one:\n { org_list } "
121- )
122-
123-
124- def __build_request (type : str , endpoint : str , params : dict = None , data : dict = None ):
73+ def __build_request (
74+ type : str ,
75+ endpoint : str ,
76+ params : dict = None ,
77+ data : dict = None ,
78+ ):
12579 """
12680 Make an API request to the SingleStore Management API.
12781
@@ -136,18 +90,19 @@ def __build_request(type: str, endpoint: str, params: dict = None, data: dict =
13690 """
13791 # Ensure an organization is selected before making API requests
13892
139- __set_organzation_id ()
93+ # __set_organzation_id()
94+
95+ settings = get_settings ()
14096
14197 def build_request_endpoint (endpoint : str , params : dict = None ):
142- url = f"{ app_config . settings .singlestore_api_base_url } /v1/{ endpoint } "
98+ url = f"{ settings .s2_api_base_url } /v1/{ endpoint } "
14399
144100 # Add organization ID as a query parameter
145101 if params is None :
146102 params = {}
147103
148- print (app_config .organization_id )
149- if app_config .organization_id :
150- params ["organizationID" ] = app_config .organization_id
104+ if settings .is_remote :
105+ params ["organizationID" ] = settings .org_id
151106
152107 if params and type == "GET" : # Only add query params for GET requests
153108 url += "?"
@@ -158,10 +113,14 @@ def build_request_endpoint(endpoint: str, params: dict = None):
158113
159114 # Headers with authentication
160115 headers = {
161- "Authorization" : f"Bearer { app_config .get_auth_token ()} " ,
162116 "Content-Type" : "application/json" ,
163117 }
164118
119+ access_token = __get_access_token ()
120+
121+ if access_token is not None :
122+ headers ["Authorization" ] = f"Bearer { access_token } "
123+
165124 request_endpoint = build_request_endpoint (endpoint , params )
166125
167126 # Default empty JSON body for POST/PUT requests if none provided
@@ -245,30 +204,14 @@ def __get_workspace_endpoint(
245204 return workspace ["endpoint" ]
246205
247206
248- def __get_project_id ():
249- """
250- Get the organization ID (project ID) from the management API.
251-
252- Returns:
253- str: The organization ID
254- """
255- # Get current organization info to extract the project ID
256- org_info = __build_request ("GET" , "organizations/current" )
257- project_id = org_info .get ("orgID" )
258-
259- if not project_id :
260- raise ValueError ("Could not retrieve organization ID from the API" )
261-
262- return project_id
263-
264-
265- def __get_user_id ():
207+ def __get_user_id () -> str :
266208 """
267209 Get the current user's ID from the management API.
268210
269211 Returns:
270212 str: The user ID
271213 """
214+
272215 # Get all users in the organization
273216 users = __build_request ("GET" , "users" )
274217
@@ -281,3 +224,44 @@ def __get_user_id():
281224 return user_id
282225
283226 raise ValueError ("Could not retrieve user ID from the API" )
227+
228+
229+ def __get_org_id () -> str :
230+ """
231+ Get the organization ID from the management API.
232+
233+ Returns:
234+ str: The organization ID
235+ """
236+ settings = get_settings ()
237+
238+ if settings .is_remote :
239+ return settings .org_id
240+ else :
241+ organization = __build_request ("GET" , "organizations/current" )
242+ if "orgID" in organization :
243+ return organization ["orgID" ]
244+ else :
245+ raise ValueError ("Could not retrieve organization ID from the API" )
246+
247+
248+ def __get_access_token () -> str :
249+ """
250+ Get the access token for the current session.
251+
252+ Returns:
253+ str: The access token
254+ """
255+ settings = get_settings ()
256+
257+ access_token : str
258+ if settings .is_remote :
259+ request = get_http_request ()
260+ access_token = request .headers .get ("Authorization" , "" ).replace ("Bearer " , "" )
261+ else :
262+ access_token = settings .api_key
263+
264+ if not access_token :
265+ raise HTTPException (401 , "Unauthorized: No access token provided" )
266+
267+ return access_token
0 commit comments