gcpdiag.queries.vertex

83class FeaturestoreStateEnum(enum.Enum):
84  """The possible states a Vertex AI featurestore can have.
85
86  https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.featurestores#state
87  """
88
89  STATE_UNSPECIFIED = 'STATE_UNSPECIFIED'
90  STABLE = 'STABLE'
91  UPDATING = 'UPDATING'
92
93  def __str__(self):
94    return str(self.value)

 97class Featurestore(models.Resource):
 98  """Represent a Vertex AI featurestore
 99
100  https://cloud.google.com/vertex-ai/docs/reference/rest/v1/projects.locations.featurestores#resource:-featurestore
101  """
102
103  _resource_data: dict
104
105  def __init__(self, project_id, resource_data):
106    super().__init__(project_id=project_id)
107    self._resource_data = resource_data
108
109  @property
110  def full_path(self) -> str:
111    """
112    The 'name' of the featurestore is already in the full path form
113    projects/{project}/locations/{location}/featurestores/{featurestore}.
114    """
115    return self._resource_data[NAME_KEY]
116
117  @property
118  def short_path(self) -> str:
119    path = self.full_path
120    path = re.sub(r'^projects/', '', path)
121    path = re.sub(r'/locations/', '/', path)
122    path = re.sub(r'/featurestores/', '/', path)
123    return path
124
125  @property
126  def name(self) -> str:
127    logging.info(self._resource_data)
128    return self._resource_data[NAME_KEY]
129
130  @property
131  def state(self) -> str:
132    logging.info(self._resource_data)
133    return self._resource_data[STATE_KEY]

136@caching.cached_api_call
137def get_featurestores(context: models.Context) -> Dict[str, Featurestore]:
138  featurestores: Dict[str, Featurestore] = {}
139  if not apis.is_enabled(context.project_id, 'aiplatform'):
140    return featurestores
141  for region in FEATURE_REGIONS[FEATURESTORES_KEY]:
142    featurestores_res: Dict[str, Featurestore] = {}
143    region_name = REGIONS[region]
144    logging.info(
145        'fetching list of Vertex AI featurestores in project %s for region %s',
146        context.project_id, region_name)
147    vertex_api = apis.get_api('aiplatform', 'v1', context.project_id,
148                              region_name)
149    query = vertex_api.projects().locations().featurestores().list(
150        parent=f'projects/{context.project_id}/locations/{region_name}')
151    try:
152      resp = query.execute(num_retries=config.API_RETRIES)
153      if FEATURESTORES_KEY not in resp:
154        continue
155      for resp_i in resp[FEATURESTORES_KEY]:
156        # verify that we have some minimal data that we expect
157        if NAME_KEY not in resp_i:
158          raise RuntimeError(
159              'missing featurestore name in projects.locations.featurestores.list response'
160          )
161        i = Featurestore(project_id=context.project_id, resource_data=resp_i)
162        featurestores_res[i.full_path] = i
163        if featurestores:
164          featurestores.update(featurestores_res)
165        else:
166          featurestores = featurestores_res
167    except googleapiclient.errors.HttpError as err:
168      raise utils.GcpApiError(err) from err
169  return featurestores