From 6d91964858684421a85405361791cc967a855e8e Mon Sep 17 00:00:00 2001 From: adishaa Date: Thu, 15 Jan 2026 10:25:13 -0800 Subject: [PATCH] Support additional search method --- resource_plan.csv | 9 ++- src/sagemaker_core/main/resources.py | 76 +++++++++++++++++++ src/sagemaker_core/main/shapes.py | 16 ++++ .../tools/additional_operations.json | 10 +++ src/sagemaker_core/tools/api_coverage.json | 2 +- 5 files changed, 110 insertions(+), 3 deletions(-) diff --git a/resource_plan.csv b/resource_plan.csv index 66884a61..7ba7fdd5 100644 --- a/resource_plan.csv +++ b/resource_plan.csv @@ -24,9 +24,11 @@ EndpointConfig,resource,"['create', 'get', 'get_all']","['delete', 'refresh']",[ Experiment,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update']",[],[],"['CreateExperiment', 'DeleteExperiment', 'DescribeExperiment', 'ListExperiments', 'UpdateExperiment']",[],[] FeatureGroup,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateFeatureGroup', 'DeleteFeatureGroup', 'DescribeFeatureGroup', 'ListFeatureGroups', 'UpdateFeatureGroup']","[{'name': 'FeatureGroupStatus', 'shape_name': 'FeatureGroupStatus'}]","['Creating', 'Created', 'CreateFailed', 'Deleting', 'DeleteFailed']" FeatureMetadata,resource,['get'],"['refresh', 'update']",[],[],"['DescribeFeatureMetadata', 'UpdateFeatureMetadata']",[],[] +FeatureStore,resource,[],[],[],[],[],[],[] FlowDefinition,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateFlowDefinition', 'DeleteFlowDefinition', 'DescribeFlowDefinition', 'ListFlowDefinitions']","[{'name': 'FlowDefinitionStatus', 'shape_name': 'FlowDefinitionStatus'}]","['Initializing', 'Active', 'Failed', 'Deleting']" Hub,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateHub', 'DeleteHub', 'DescribeHub', 'ListHubs', 'UpdateHub']","[{'name': 'HubStatus', 'shape_name': 'HubStatus'}]","['InService', 'Creating', 'Updating', 'Deleting', 'CreateFailed', 'UpdateFailed', 'DeleteFailed']" HubContent,resource,"['get', 'get_all', 'import']","['delete', 'refresh', 'update', 'wait_for_status']",[],[],"['DeleteHubContent', 'DescribeHubContent', 'ImportHubContent', 'ListHubContents', 'UpdateHubContent']","[{'name': 'SupportStatus', 'shape_name': 'HubContentSupportStatus'}]","['Supported', 'Deprecated', 'Restricted']" +HubContentPresignedUrls,resource,['create'],[],"['Hub', 'HubContent']",[],['CreateHubContentPresignedUrls'],[],[] HubContentReference,resource,['create'],"['delete', 'update']","['Hub', 'HubContent']",[],"['CreateHubContentReference', 'DeleteHubContentReference', 'UpdateHubContentReference']",[],[] HumanTaskUi,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateHumanTaskUi', 'DeleteHumanTaskUi', 'DescribeHumanTaskUi', 'ListHumanTaskUis']","[{'name': 'HumanTaskUiStatus', 'shape_name': 'HumanTaskUiStatus'}]","['Active', 'Deleting']" HyperParameterTuningJob,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'stop', 'wait', 'wait_for_delete']",[],[],"['CreateHyperParameterTuningJob', 'DeleteHyperParameterTuningJob', 'DescribeHyperParameterTuningJob', 'ListHyperParameterTuningJobs', 'StopHyperParameterTuningJob']","[{'name': 'HyperParameterTuningJobStatus', 'shape_name': 'HyperParameterTuningJobStatus'}]","['Completed', 'InProgress', 'Failed', 'Stopped', 'Stopping', 'Deleting', 'DeleteFailed']" @@ -37,6 +39,7 @@ InferenceExperiment,resource,"['create', 'get', 'get_all', 'start']","['delete', InferenceRecommendationsJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'wait', 'wait_for_delete']",[],[],"['CreateInferenceRecommendationsJob', 'DescribeInferenceRecommendationsJob', 'ListInferenceRecommendationsJobs', 'StopInferenceRecommendationsJob']","[{'name': 'Status', 'shape_name': 'RecommendationJobStatus'}]","['PENDING', 'IN_PROGRESS', 'COMPLETED', 'FAILED', 'STOPPING', 'STOPPED', 'DELETING', 'DELETED']" LabelingJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'wait']",[],[],"['CreateLabelingJob', 'DescribeLabelingJob', 'ListLabelingJobs', 'StopLabelingJob']","[{'name': 'LabelingJobStatus', 'shape_name': 'LabelingJobStatus'}]","['Initializing', 'InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped']" LineageGroup,resource,"['get', 'get_all']",['refresh'],[],[],"['DescribeLineageGroup', 'ListLineageGroups']",[],[] +MlflowApp,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateMlflowApp', 'DeleteMlflowApp', 'DescribeMlflowApp', 'ListMlflowApps', 'UpdateMlflowApp']","[{'name': 'Status', 'shape_name': 'MlflowAppStatus'}]","['Creating', 'Created', 'CreateFailed', 'Updating', 'Updated', 'UpdateFailed', 'Deleting', 'DeleteFailed', 'Deleted']" MlflowTrackingServer,resource,"['create', 'get', 'get_all', 'start']","['delete', 'refresh', 'stop', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateMlflowTrackingServer', 'DeleteMlflowTrackingServer', 'DescribeMlflowTrackingServer', 'ListMlflowTrackingServers', 'StartMlflowTrackingServer', 'StopMlflowTrackingServer', 'UpdateMlflowTrackingServer']","[{'name': 'TrackingServerStatus', 'shape_name': 'TrackingServerStatus'}]","['Creating', 'Created', 'CreateFailed', 'Updating', 'Updated', 'UpdateFailed', 'Deleting', 'DeleteFailed', 'Stopping', 'Stopped', 'StopFailed', 'Starting', 'Started', 'StartFailed', 'MaintenanceInProgress', 'MaintenanceComplete', 'MaintenanceFailed']" Model,resource,"['create', 'get', 'get_all']","['delete', 'refresh']",[],[],"['CreateModel', 'DeleteModel', 'DescribeModel', 'ListModels']",[],[] ModelBiasJobDefinition,resource,"['create', 'get', 'get_all']","['delete', 'refresh']",[],[],"['CreateModelBiasJobDefinition', 'DeleteModelBiasJobDefinition', 'DescribeModelBiasJobDefinition', 'ListModelBiasJobDefinitions']",[],[] @@ -57,17 +60,19 @@ PartnerAppPresignedUrl,resource,['create'],[],[],[],['CreatePartnerAppPresignedU Pipeline,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreatePipeline', 'DeletePipeline', 'DescribePipeline', 'ListPipelines', 'UpdatePipeline']","[{'name': 'PipelineStatus', 'shape_name': 'PipelineStatus'}]","['Active', 'Deleting']" PipelineExecution,resource,"['get', 'get_all', 'start']","['refresh', 'stop', 'update', 'wait_for_status']",[],[],"['DescribePipelineExecution', 'ListPipelineExecutions', 'StartPipelineExecution', 'StopPipelineExecution', 'UpdatePipelineExecution']","[{'name': 'PipelineExecutionStatus', 'shape_name': 'PipelineExecutionStatus'}]","['Executing', 'Stopping', 'Stopped', 'Failed', 'Succeeded']" PresignedDomainUrl,resource,['create'],[],"['Space', 'UserProfile']",[],['CreatePresignedDomainUrl'],[],[] +PresignedMlflowAppUrl,resource,['create'],[],[],[],['CreatePresignedMlflowAppUrl'],[],[] PresignedMlflowTrackingServerUrl,resource,['create'],[],[],[],['CreatePresignedMlflowTrackingServerUrl'],[],[] PresignedNotebookInstanceUrl,resource,['create'],[],['NotebookInstance'],[],['CreatePresignedNotebookInstanceUrl'],[],[] -ProcessingJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'wait']",[],[],"['CreateProcessingJob', 'DescribeProcessingJob', 'ListProcessingJobs', 'StopProcessingJob']","[{'name': 'ProcessingJobStatus', 'shape_name': 'ProcessingJobStatus'}]","['InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped']" +ProcessingJob,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'stop', 'wait']",[],[],"['CreateProcessingJob', 'DeleteProcessingJob', 'DescribeProcessingJob', 'ListProcessingJobs', 'StopProcessingJob']","[{'name': 'ProcessingJobStatus', 'shape_name': 'ProcessingJobStatus'}]","['InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped']" Project,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update', 'wait_for_status']",[],[],"['CreateProject', 'DeleteProject', 'DescribeProject', 'ListProjects', 'UpdateProject']","[{'name': 'ProjectStatus', 'shape_name': 'ProjectStatus'}]","['Pending', 'CreateInProgress', 'CreateCompleted', 'CreateFailed', 'DeleteInProgress', 'DeleteFailed', 'DeleteCompleted', 'UpdateInProgress', 'UpdateCompleted', 'UpdateFailed']" ResourceCatalog,resource,['get_all'],[],[],[],['ListResourceCatalogs'],[],[] SagemakerServicecatalogPortfolio,resource,[],[],[],[],[],[],[] +Session,resource,['start'],[],[],[],['StartSession'],[],[] Space,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update', 'wait_for_delete', 'wait_for_status']",[],[],"['CreateSpace', 'DeleteSpace', 'DescribeSpace', 'ListSpaces', 'UpdateSpace']","[{'name': 'Status', 'shape_name': 'SpaceStatus'}]","['Deleting', 'Failed', 'InService', 'Pending', 'Updating', 'Update_Failed', 'Delete_Failed']" StudioLifecycleConfig,resource,"['create', 'get', 'get_all']","['delete', 'refresh']",[],[],"['CreateStudioLifecycleConfig', 'DeleteStudioLifecycleConfig', 'DescribeStudioLifecycleConfig', 'ListStudioLifecycleConfigs']",[],[] SubscribedWorkteam,resource,"['get', 'get_all']",['refresh'],[],[],"['DescribeSubscribedWorkteam', 'ListSubscribedWorkteams']",[],[] Tag,resource,['get_all'],[],[],[],['ListTags'],[],[] -TrainingJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'update', 'wait']",[],[],"['CreateTrainingJob', 'DescribeTrainingJob', 'ListTrainingJobs', 'StopTrainingJob', 'UpdateTrainingJob']","[{'name': 'TrainingJobStatus', 'shape_name': 'TrainingJobStatus'}]","['InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped']" +TrainingJob,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'stop', 'update', 'wait', 'wait_for_delete']",[],[],"['CreateTrainingJob', 'DeleteTrainingJob', 'DescribeTrainingJob', 'ListTrainingJobs', 'StopTrainingJob', 'UpdateTrainingJob']","[{'name': 'TrainingJobStatus', 'shape_name': 'TrainingJobStatus'}]","['InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped', 'Deleting']" TrainingPlan,resource,"['create', 'get', 'get_all']","['refresh', 'wait_for_status']",[],[],"['CreateTrainingPlan', 'DescribeTrainingPlan', 'ListTrainingPlans']","[{'name': 'Status', 'shape_name': 'TrainingPlanStatus'}]","['Pending', 'Active', 'Scheduled', 'Expired', 'Failed']" TransformJob,resource,"['create', 'get', 'get_all']","['refresh', 'stop', 'wait']",['Model'],[],"['CreateTransformJob', 'DescribeTransformJob', 'ListTransformJobs', 'StopTransformJob']","[{'name': 'TransformJobStatus', 'shape_name': 'TransformJobStatus'}]","['InProgress', 'Completed', 'Failed', 'Stopping', 'Stopped']" Trial,resource,"['create', 'get', 'get_all']","['delete', 'refresh', 'update']",['Experiment'],[],"['CreateTrial', 'DeleteTrial', 'DescribeTrial', 'ListTrials', 'UpdateTrial']",[],[] diff --git a/src/sagemaker_core/main/resources.py b/src/sagemaker_core/main/resources.py index 84a64a45..49542eef 100644 --- a/src/sagemaker_core/main/resources.py +++ b/src/sagemaker_core/main/resources.py @@ -11326,6 +11326,82 @@ def update( return self +class FeatureStore(Base): + """ + Class representing resource FeatureStore + + """ + + @staticmethod + @Base.add_validate_call + def search( + resource: str, + search_expression: Optional[shapes.SearchExpression] = Unassigned(), + sort_by: Optional[str] = Unassigned(), + sort_order: Optional[str] = Unassigned(), + next_token: Optional[str] = Unassigned(), + max_results: Optional[int] = Unassigned(), + cross_account_filter_option: Optional[str] = Unassigned(), + visibility_conditions: Optional[List[shapes.VisibilityConditions]] = Unassigned(), + session: Optional[Session] = None, + region: Optional[str] = None, + ) -> Optional[shapes.SearchResponse]: + """ + Finds SageMaker resources that match a search query. + + Parameters: + resource: The name of the SageMaker resource to search for. + search_expression: A Boolean conditional statement. Resources must satisfy this condition to be included in search results. You must provide at least one subexpression, filter, or nested filter. The maximum number of recursive SubExpressions, NestedFilters, and Filters that can be included in a SearchExpression object is 50. + sort_by: The name of the resource property used to sort the SearchResults. The default is LastModifiedTime. + sort_order: How SearchResults are ordered. Valid values are Ascending or Descending. The default is Descending. + next_token: If more than MaxResults resources match the specified SearchExpression, the response includes a NextToken. The NextToken can be passed to the next SearchRequest to continue retrieving results. + max_results: The maximum number of results to return. + cross_account_filter_option: A cross account filter option. When the value is "CrossAccount" the search results will only include resources made discoverable to you from other accounts. When the value is "SameAccount" or null the search results will only include resources from your account. Default is null. For more information on searching for resources made discoverable to your account, see Search discoverable resources in the SageMaker Developer Guide. The maximum number of ResourceCatalogs viewable is 1000. + visibility_conditions: Limits the results of your search request to the resources that you can access. + session: Boto3 session. + region: Region name. + + Returns: + shapes.SearchResponse + + Raises: + botocore.exceptions.ClientError: This exception is raised for AWS service related errors. + The error message and error code can be parsed from the exception as follows: + ``` + try: + # AWS service call here + except botocore.exceptions.ClientError as e: + error_message = e.response['Error']['Message'] + error_code = e.response['Error']['Code'] + ``` + """ + + operation_input_args = { + "Resource": resource, + "SearchExpression": search_expression, + "SortBy": sort_by, + "SortOrder": sort_order, + "NextToken": next_token, + "MaxResults": max_results, + "CrossAccountFilterOption": cross_account_filter_option, + "VisibilityConditions": visibility_conditions, + } + # serialize the input request + operation_input_args = serialize(operation_input_args) + logger.debug(f"Serialized input request: {operation_input_args}") + + client = Base.get_sagemaker_client( + session=session, region_name=region, service_name="sagemaker" + ) + + logger.debug(f"Calling search API") + response = client.search(**operation_input_args) + logger.debug(f"Response: {response}") + + transformed_response = transform(response, "SearchResponse") + return shapes.SearchResponse(**transformed_response) + + class FlowDefinition(Base): """ Class representing resource FlowDefinition diff --git a/src/sagemaker_core/main/shapes.py b/src/sagemaker_core/main/shapes.py index be855cf4..5e0baf71 100644 --- a/src/sagemaker_core/main/shapes.py +++ b/src/sagemaker_core/main/shapes.py @@ -14233,6 +14233,22 @@ class TotalHits(Base): relation: Optional[str] = Unassigned() +class SearchResponse(Base): + """ + SearchResponse + + Attributes + ---------------------- + results: A list of SearchRecord objects. + next_token: If the result of the previous Search request was truncated, the response includes a NextToken. To retrieve the next set of results, use the token in the next request. + total_hits: The total number of matching results. + """ + + results: Optional[List[SearchRecord]] = Unassigned() + next_token: Optional[str] = Unassigned() + total_hits: Optional[TotalHits] = Unassigned() + + class TrainingPlanOffering(Base): """ TrainingPlanOffering diff --git a/src/sagemaker_core/tools/additional_operations.json b/src/sagemaker_core/tools/additional_operations.json index 85796128..39977c76 100644 --- a/src/sagemaker_core/tools/additional_operations.json +++ b/src/sagemaker_core/tools/additional_operations.json @@ -1,4 +1,14 @@ { + "FeatureStore": { + "Search": { + "operation_name": "Search", + "resource_name": "FeatureStore", + "method_name": "search", + "return_type": "SearchResponse", + "method_type": "static", + "service_name": "sagemaker" + } + }, "Cluster": { "DescribeClusterNode": { "operation_name": "DescribeClusterNode", diff --git a/src/sagemaker_core/tools/api_coverage.json b/src/sagemaker_core/tools/api_coverage.json index 570b6733..1f2a7749 100644 --- a/src/sagemaker_core/tools/api_coverage.json +++ b/src/sagemaker_core/tools/api_coverage.json @@ -1 +1 @@ -{"SupportedAPIs": 371, "UnsupportedAPIs": 17} \ No newline at end of file +{"SupportedAPIs": 372, "UnsupportedAPIs": 16} \ No newline at end of file