Access row elements with row.x rather than row[x], compatible wit…

…h sqlalchemy 2.0. Resolves #993

PiperOrigin-RevId: 586493566

requirements.txt

@@ -6,4 +6,4 @@ portpicker>=1.3.1
 grpcio>=1.35.0
 grpcio-tools>=1.35.0
 googleapis-common-protos>=1.56.4
-sqlalchemy>=1.4,<=1.4.20
+sqlalchemy>=1.4

vizier/__init__.py

@@ -23,4 +23,4 @@
 
 sys.path.append(PROTO_ROOT)
 
-__version__ = "0.1.12"
+__version__ = "0.1.13"

vizier/_src/service/sql_datastore.py

@@ -18,7 +18,7 @@
 
 import collections
 import threading
-from typing import Callable, DefaultDict, Iterable, List, Optional
+from typing import Callable, Iterable, List, Optional
 
 from absl import logging
 import sqlalchemy as sqla
@@ -124,7 +124,7 @@ def load_study(self, study_name: str) -> study_pb2.Study:
     row = result.fetchone()
     if not row:
       raise NotFoundError('Failed to find study name: %s' % study_name)
-    return study_pb2.Study.FromString(row['serialized_study'])
+    return study_pb2.Study.FromString(row.serialized_study)
 
   def update_study(self, study: study_pb2.Study) -> resources.StudyResource:
     study_resource = resources.StudyResource.from_name(study.name)
@@ -190,9 +190,7 @@ def list_studies(self, owner_name: str) -> List[study_pb2.Study]:
         raise NotFoundError('Owner name %s does not exist.' % owner_name)
       result = self._connection.execute(lq).fetchall()
 
-    return [
-        study_pb2.Study.FromString(row['serialized_study']) for row in result
-    ]
+    return [study_pb2.Study.FromString(row.serialized_study) for row in result]
 
   def create_trial(self, trial: study_pb2.Trial) -> resources.TrialResource:
     trial_resource = resources.TrialResource.from_name(trial.name)
@@ -223,7 +221,7 @@ def get_trial(self, trial_name: str) -> study_pb2.Trial:
     row = result.fetchone()
     if not row:
       raise NotFoundError('Failed to find trial name: %s' % trial_name)
-    return study_pb2.Trial.FromString(row['serialized_trial'])
+    return study_pb2.Trial.FromString(row.serialized_trial)
 
   def update_trial(self, trial: study_pb2.Trial) -> resources.TrialResource:
     trial_resource = resources.TrialResource.from_name(trial.name)
@@ -269,9 +267,7 @@ def list_trials(self, study_name: str) -> List[study_pb2.Trial]:
         raise NotFoundError('Study name %s does not exist.' % study_name)
       result = self._connection.execute(lq)
 
-    return [
-        study_pb2.Trial.FromString(row['serialized_trial']) for row in result
-    ]
+    return [study_pb2.Trial.FromString(row.serialized_trial) for row in result]
 
   def delete_trial(self, trial_name: str) -> None:
     # Exist query
@@ -347,7 +343,7 @@ def get_suggestion_operation(
     row = result.fetchone()
     if not row:
       raise NotFoundError('Failed to find suggest op name: %s' % operation_name)
-    return operations_pb2.Operation.FromString(row['serialized_op'])
+    return operations_pb2.Operation.FromString(row.serialized_op)
 
   def update_suggestion_operation(
       self, operation: operations_pb2.Operation
@@ -407,8 +403,7 @@ def list_suggestion_operations(
       result = self._connection.execute(q)
 
     all_ops = [
-        operations_pb2.Operation.FromString(row['serialized_op'])
-        for row in result
+        operations_pb2.Operation.FromString(row.serialized_op) for row in result
     ]
     if filter_fn is not None:
       output_list = []
@@ -495,9 +490,7 @@ def get_early_stopping_operation(
       raise NotFoundError(
           'Failed to find early stopping op name: %s' % operation_name
       )
-    return vizier_oss_pb2.EarlyStoppingOperation.FromString(
-        row['serialized_op']
-    )
+    return vizier_oss_pb2.EarlyStoppingOperation.FromString(row.serialized_op)
 
   def update_early_stopping_operation(
       self, operation: vizier_oss_pb2.EarlyStoppingOperation
@@ -552,7 +545,7 @@ def update_metadata(
       row = study_result.fetchone()
       if not row:
         raise NotFoundError('No such study:', s_resource.name)
-      original_study = study_pb2.Study.FromString(row['serialized_study'])
+      original_study = study_pb2.Study.FromString(row.serialized_study)
 
       # Store Study-related metadata into the database.
       vz.metadata_util.merge_study_metadata(
@@ -565,9 +558,7 @@ def update_metadata(
       self._connection.execute(usq)
 
       # Split the trial-related metadata by Trial.
-      split_metadata: DefaultDict[str, List[datastore.UnitMetadataUpdate]] = (
-          collections.defaultdict(list)
-      )
+      split_metadata = collections.defaultdict(list)
       for md in trial_metadata:
         split_metadata[md.trial_id].append(md)
 
@@ -583,7 +574,7 @@ def update_metadata(
         row = trial_result.fetchone()
         if not row:
           raise NotFoundError('No such trial:', trial_name)
-        original_trial = study_pb2.Trial.FromString(row['serialized_trial'])
+        original_trial = study_pb2.Trial.FromString(row.serialized_trial)
 
         # Update Trial.
         vz.metadata_util.merge_trial_metadata(original_trial, md_list)