def run_pre_experiment_hook(self):
"""
Run pre-experiment hook by initiating a database connection for state retrieval
:return: None
"""
from cockroachdb.modules.connection import (
database,
initialize_cockroach_database,
)
database.initialize(
initialize_cockroach_database(hosts=self.db_hosts,
port=self.db_port))
def run(transaction: BaseTransaction):
"""
Thunk to transaction
:param transaction: transaction object
:return: None
"""
from cockroachdb.modules.connection import (
database,
initialize_cockroach_database,
)
database.initialize(initialize_cockroach_database())
transaction.run()
def run_pre_transaction_hook(self):
"""
Initiate new cockroach database connection
:return: None
"""
from cockroachdb.modules.connection import (
database,
initialize_cockroach_database,
)
database.initialize(
initialize_cockroach_database(
host_index=self.client_number % len(self.db_hosts),
hosts=self.db_hosts,
port=self.db_port,
))
def run_init_hook(self):
"""
Run initialization hook before processing row data
:return: database reference
"""
from cockroachdb.modules.connection import (
database,
initialize_cockroach_database,
)
database.initialize(
initialize_cockroach_database(
hosts=self.db_hosts, port=self.db_port
)
)
return dict(database=database)
from cockroachdb.modules.connection import (
initialize_cockroach_database,
database,
)
from common.logging_tool import console
if __name__ == "__main__":
cockroach_db = initialize_cockroach_database()
database.initialize(cockroach_db)
with database.atomic():
result = cockroach_db.execute_sql("SHOW USERS;")
console.log("Connection is successfully established")
console.log(result.fetchall())