def DEFINE_database(name: str,
database_class,
default: Optional[str],
help: str,
must_exist: bool = False,
validator: Callable[[Any], bool] = None):
"""Registers a flag whose value is a sqlutil.Database class.
Unlike other DEFINE_* functions, the value produced by this flag is not an
instance of the value, but a lambda that will instantiate a database of the
requested type. This flag value must be called (with no arguments) in order to
instantiate a database.
Args:
name: The name of the flag.
database_class: The subclass of sqlutil.Database which is to be instantiated
when this value is called, using the URL declared in 'default'.
default: The default URL of the database. This is a required value.
help: The help string.
must_exist: If True, require that the database exists. Else, the database is
created if it does not exist.
"""
parser = flags_parsers.DatabaseParser(database_class,
must_exist=must_exist)
serializer = absl_flags.ArgumentSerializer()
absl_flags.DEFINE(parser,
name,
default,
help,
absl_flags.FLAGS,
serializer,
module_name=logging.GetCallingModuleName())
if validator:
RegisterFlagValidator(name, validator)
def GetOrAdd(session: sql.orm.session.Session,
model,
defaults: typing.Dict[str, object] = None,
**kwargs):
"""Instantiate a mapped database object.
If the object is not in the database,
add it. Note that no change is written to disk until commit() is called on the
session.
Args:
session: The database session.
model: The database table class.
defaults: Default values for mapped objects.
kwargs: The values for the table row.
Returns:
An instance of the model class, with the values specified.
"""
instance = session.query(model).filter_by(**kwargs).first()
if not instance:
params = {
k: v
for k, v in kwargs.items()
if not isinstance(v, sql.sql.expression.ClauseElement)
}
params.update(defaults or {})
instance = model(**params)
session.add(instance)
logging.Log(logging.GetCallingModuleName(), 5, 'New record: %s(%s)',
model.__name__, params)
return instance
def DEFINE_input_path(name: str,
default: Union[None, str, pathlib.Path],
help: str,
required: bool = False,
is_dir: bool = False,
validator: Callable[[pathlib.Path], bool] = None):
"""Registers a flag whose value is an input path.
An "input path" is a path to a file or directory that exists. The parsed value
is a pathlib.Path instance. Flag parsing will fail if the value of this flag
is not a path to an existing file or directory.
Args:
name: The name of the flag.
default: The default value for the flag. While None is a legal value, it
will fail during parsing - input paths are required flags.
help: The help string.
is_dir: If true, require the that the value be a directory. Else, require
that the value be a file. Parsing will fail if this is not the case.
"""
parser = flags_parsers.PathParser(must_exist=True, is_dir=is_dir)
serializer = absl_flags.ArgumentSerializer()
absl_flags.DEFINE(parser,
name,
default,
help,
absl_flags.FLAGS,
serializer,
module_name=logging.GetCallingModuleName())
if required:
absl_flags.mark_flag_as_required(name)
if validator:
RegisterFlagValidator(name, validator)
def Flush(self) -> None:
"""Commit all buffered mapped objects to database."""
failures = ResilientAddManyAndCommit(self._db, self._to_commit)
if len(failures):
logging.Log(logging.GetCallingModuleName(), 1,
'BufferedDatabaseWriter failed to commit %d objects',
len(failures))
self._to_commit = []
self._last_commit = time.time()
def Log(level: int, msg, *args, **kwargs):
"""Logs a message at the given level.
Per-module verbose level. The argument has to contain a comma-separated
list of <module name>=<log level>. <module name> is a glob pattern (e.g., "
"gfs* for all modules whose name starts with \"gfs\"), matched against the "
"filename base (that is, name ignoring .py). <log level> overrides any "
"value given by --v."
"""
calling_module = logging.GetCallingModuleName()
logging.Log(calling_module, level, msg, *args, **kwargs)
def DEFINE_list(name: str,
default: Optional[List[Any]],
help: str,
required: bool = False,
validator: Callable[[List[Any]], bool] = None):
"""Registers a flag whose value must be a list."""
absl_flags.DEFINE_list(name,
default,
help,
module_name=logging.GetCallingModuleName())
if required:
absl_flags.mark_flag_as_required(name)
if validator:
RegisterFlagValidator(name, validator)
def ResilientAddManyAndCommit(db: Database, mapped: typing.Iterable[Base]):
"""Attempt to commit all mapped objects and return those that fail.
This method creates a session and commits the given mapped objects.
In case of error, this method will recurse up to O(log(n)) times, committing
as many objects that can be as possible.
Args:
db: The database to add the objects to.
mapped: A sequence of objects to commit.
Returns:
Any items in `mapped` which could not be committed, if any. Relative order
of items is preserved.
"""
failures = []
if not mapped:
return failures
mapped = list(mapped)
try:
with db.Session(commit=True) as session:
session.add_all(mapped)
except sql.exc.SQLAlchemyError as e:
logging.Log(
logging.GetCallingModuleName(),
1,
'Caught error while committing %d mapped objects: %s',
len(mapped),
e,
)
# Divide and conquer. If we're committing only a single object, then a
# failure to commit it means that we can do nothing other than return it.
# Else, divide the mapped objects in half and attempt to commit as many of
# them as possible.
if len(mapped) == 1:
return mapped
else:
mid = int(len(mapped) / 2)
left = mapped[:mid]
right = mapped[mid:]
failures += ResilientAddManyAndCommit(db, left)
failures += ResilientAddManyAndCommit(db, right)
return failures
def DEFINE_float(name: str,
default: Optional[float],
help: str,
required: bool = False,
lower_bound: Optional[float] = None,
upper_bound: Optional[float] = None,
validator: Callable[[float], bool] = None):
"""Registers a flag whose value must be a float."""
absl_flags.DEFINE_float(name,
default,
help,
module_name=logging.GetCallingModuleName(),
lower_bound=lower_bound,
upper_bound=upper_bound)
if required:
absl_flags.mark_flag_as_required(name)
if validator:
RegisterFlagValidator(name, validator)
def DEFINE_output_path(name: str,
default: Union[None, str, pathlib.Path],
help: str,
required: bool = False,
is_dir: bool = False,
exist_ok: bool = True,
must_exist: bool = False,
validator: Callable[[pathlib.Path], bool] = None):
"""Registers a flag whose value is an output path.
An "output path" is a path to a file or directory that may or may not already
exist. The parsed value is a pathlib.Path instance. The idea is that this flag
can be used to specify paths to files or directories that will be created
during program execution. However, note that specifying an output path does
not guarantee that the file will be produced.
Args:
name: The name of the flag.
default: The default value for the flag. While None is a legal value, it
will fail during parsing - output paths are required flags.
help: The help string.
is_dir: If true, require the that the value be a directory. Else, require
that the value be a file. Parsing will fail if the path already exists and
is of the incorrect type.
exist_ok: If False, require that the path not exist, else parsing will fail.
must_exist: If True, require that the path exists, else parsing will fail.
"""
parser = flags_parsers.PathParser(must_exist=must_exist,
exist_ok=exist_ok,
is_dir=is_dir)
serializer = absl_flags.ArgumentSerializer()
absl_flags.DEFINE(parser,
name,
default,
help,
absl_flags.FLAGS,
serializer,
module_name=logging.GetCallingModuleName())
if required:
absl_flags.mark_flag_as_required(name)
if validator:
RegisterFlagValidator(name, validator)
def LogIf(level: int, condition, msg, *args, **kwargs):
if condition:
calling_module = logging.GetCallingModuleName()
logging.Log(calling_module, level, msg, *args, **kwargs)
def GetVerbosity() -> int:
"""Get the verbosity level.
This can be set per-module using --vmodule flag.
"""
return logging.GetVerbosity(logging.GetCallingModuleName())
