def Drop(self, are_you_sure_about_this_flag: bool = False):
"""Drop the database, irreverisbly destroying it.
Be careful with this! After calling this method an a Database instance, no
further operations can be made on it, and any Sessions should be discarded.
Args:
are_you_sure_about_this_flag: You should be sure.
Raises:
ValueError: In case you're not 100% sure.
"""
if not are_you_sure_about_this_flag:
raise ValueError("Let's take a minute to think things over")
if self.url.startswith("mysql://"):
engine = sql.create_engine("/".join(self.url.split("/")[:-1]))
database = self.url.split("/")[-1].split("?")[0]
logging.Log(logging.GetCallingModuleName(), 1, "database %s",
database)
engine.execute(f"DROP DATABASE IF EXISTS `{database}`")
elif self.url == "sqlite://":
# In-memory databases do not dropping.
pass
elif self.url.startswith("sqlite:///"):
path = pathlib.Path(self.url[len("sqlite:///"):])
assert path.is_file()
path.unlink()
else:
raise NotImplementedError(
f"Unsupported operation DROP for database: '{self.url}'", )
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 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."
"""
logging.Log(
logging.GetCallingModuleName(),
level,
_MaybeColorizeLog(
shell.ShellEscapeCodes.YELLOW
if level > 1 else shell.ShellEscapeCodes.CYAN,
msg,
*args,
),
**kwargs,
)
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 GetVerbosity() -> int:
"""Get the verbosity level.
This can be set per-module using --vmodule flag.
"""
return logging.GetModuleVerbosity(logging.GetCallingModuleName())
