Move xm_local flags into xmanager/xm_flags.py.

PiperOrigin-RevId: 614834009
Change-Id: I4ffc491b1592a7414591d256610f28862870b520
GitOrigin-RevId: 9b9a11fb98baf0386481de87dd6d3675611d35e1

xmanager/xm_flags.py

@@ -14,3 +14,44 @@
 """XManager Flags."""
 
 from absl import flags
+
+# -------------------- xm_local --------------------
+
+DB_YAML_CONFIG_PATH = flags.DEFINE_string(
+    'xm_db_yaml_config_path',
+    None,
+    """
+    Path of YAML config file containing DB connection details.
+
+    A valid config file contains two main entries:
+      `sql_connector`: must be one of [`sqlite`, `generic`, `cloudsql`]
+
+      `sql_connection_settings`: contains details about the connection URL.
+        These match the interface of `SqlConnectionSettings` and their
+        combination must form a valid `sqlalchemy` connection URL. Possible
+        fields are:
+          - backend, e.g. 'mysql', 'postgresql'
+          - db_name
+          - driver, e.g. 'pymysql', 'pg8000'
+          - username
+          - password
+          - host (instance connection name when using CloudSql)
+          - port
+    """,
+)
+
+UPGRADE_DB = flags.DEFINE_boolean(
+    'xm_upgrade_db',
+    False,
+    """
+    Specifies if XManager should update the database to the latest version.
+    It's recommended to take a back-up of the database before updating, since
+    migrations can fail/have errors. This is especially true
+    for non-transactional DDLs, where partial migrations can occur on
+    failure, leaving the database in a not well-defined state.
+    """,
+)
+
+BAZEL_COMMAND = flags.DEFINE_string(
+    'xm_bazel_command', 'bazel', 'A command that runs Bazel.',
+)

xmanager/xm_local/packaging/bazel_tools.py

@@ -20,18 +20,14 @@
 import subprocess
 from typing import Dict, List, Optional, Sequence, Tuple
 
-from absl import flags
 from xmanager import xm
+from xmanager import xm_flags
 from xmanager.bazel import client
 from xmanager.bazel import file_utils
 
 from google.protobuf.internal.decoder import _DecodeVarint32
 from xmanager.generated import build_event_stream_pb2 as bes_pb2
 
-_BAZEL_COMMAND = flags.DEFINE_string(
-    'xm_bazel_command', 'bazel', 'A command that runs Bazel.'
-)
-
 
 def _get_important_outputs(
     events: Sequence[bes_pb2.BuildEvent], labels: Sequence[str]
@@ -103,7 +99,7 @@ def _root_absolute_path() -> str:
   return (
       os.getenv('BUILD_WORKSPACE_DIRECTORY')
       or subprocess.run(
-          [_BAZEL_COMMAND.value, 'info', 'workspace'],
+          [xm_flags.BAZEL_COMMAND.value, 'info', 'workspace'],
           check=True,
           stdout=subprocess.PIPE,
           stderr=subprocess.PIPE,
@@ -130,7 +126,7 @@ def _build_multiple_targets(
   with file_utils.TemporaryFilePath() as bep_path:
     subprocess.run(
         [
-            _BAZEL_COMMAND.value,
+            xm_flags.BAZEL_COMMAND.value,
             'build',
             f'--build_event_binary_file={bep_path}',
             # Forces a GC at the end of the build and publishes value to BEP.
@@ -204,7 +200,7 @@ def fetch_kinds(self, labels: Sequence[str]) -> List[str]:
     # https://docs.bazel.build/versions/main/query.html#output-label_kind.
     stdout = subprocess.run(
         [
-            _BAZEL_COMMAND.value,
+            xm_flags.BAZEL_COMMAND.value,
             'query',
             f"'{' union '.join(labels)}'",
             '--output',

xmanager/xm_local/storage/database.py

@@ -18,54 +18,19 @@
 import tempfile
 from typing import Any, Dict, List, Optional, Type, TypeVar
 
-from absl import flags
 import alembic
 from alembic.config import Config
 import attr
 import sqlalchemy
 from xmanager import xm
+from xmanager import xm_flags
 from xmanager.generated import data_pb2
 import yaml
 
 from google.protobuf import text_format
 
 from google.cloud.sql.connector import Connector, IPTypes
 
-_DB_YAML_CONFIG_PATH = flags.DEFINE_string(
-    'xm_db_yaml_config_path',
-    None,
-    """
-    Path of YAML config file containing DB connection details.
-
-    A valid config file contains two main entries:
-      `sql_connector`: must be one of [`sqlite`, `generic`, `cloudsql`]
-
-      `sql_connection_settings`: contains details about the connection URL.
-        These match the interface of `SqlConnectionSettings` and their
-        combination must form a valid `sqlalchemy` connection URL. Possible
-        fields are:
-          - backend, e.g. 'mysql', 'postgresql'
-          - db_name
-          - driver, e.g. 'pymysql', 'pg8000'
-          - username
-          - password
-          - host (instance connection name when using CloudSql)
-          - port
-    """,
-)
-
-_UPGRADE_DB = flags.DEFINE_boolean(
-    'xm_upgrade_db',
-    False,
-    """
-    Specifies if XManager should update the database to the latest version.
-    It's recommended to take a back-up of the database before updating, since
-    migrations can fail/have errors. This is especially true
-    for non-transactional DDLs, where partial migrations can occur on
-    failure, leaving the database in a not well-defined state.
-    """,
-)
-
 
 @attr.s(auto_attribs=True)
 class WorkUnitResult:
@@ -256,7 +221,7 @@ def maybe_migrate_database_version(self):
     need_to_update = (
         db_version != self.latest_version_available() and db_version
     ) or legacy_sqlite_db
-    if need_to_update and not _UPGRADE_DB.value:
+    if need_to_update and not xm_flags.UPGRADE_DB.value:
       raise RuntimeError(
           f'Database is not up to date: current={self.database_version()}, '
           f'latest={self.latest_version_available()}. Take a backup of the '
@@ -404,9 +369,9 @@ def _validate_db_config(config: Dict[str, Any]) -> None:
 @functools.lru_cache()
 def _db_config() -> Dict[str, Any]:
   """Parses and validates YAML DB config file to a dict."""
-  if _DB_YAML_CONFIG_PATH.value is not None:
+  if xm_flags.DB_YAML_CONFIG_PATH.value is not None:
     db_config_file = xm.utils.resolve_path_relative_to_launcher(
-        _DB_YAML_CONFIG_PATH.value
+        xm_flags.DB_YAML_CONFIG_PATH.value
     )
     with open(db_config_file, 'r') as f:
       config = yaml.safe_load(f)