def create_db_client(dsn: str) -> Iterator[Client]:
with edgedb.create_client(
dsn=dsn,
max_concurrency=1,
tls_ca_file=config.DATABASE_TLS_CA_FILE,
) as client:
yield client
def db_client():
"""Blocking connection to the database."""
with edgedb.create_client(
config.DATABASE_DSN,
max_concurrency=1,
tls_ca_file=config.DATABASE_TLS_CA_FILE,
) as client:
yield client
def execute_conflict(self, name='counter2', options=None):
con_args = self.get_connect_args().copy()
con_args.update(database=self.get_database_name())
client2 = edgedb.create_client(**con_args)
self.addCleanup(client2.close)
barrier = Barrier(2)
lock = threading.Lock()
iterations = 0
def transaction1(client):
for tx in client.transaction():
nonlocal iterations
iterations += 1
with tx:
# This magic query makes the test more reliable for some
# reason. I guess this is because starting a transaction
# in EdgeDB (and/or Postgres) is accomplished somewhat
# lazily, i.e. only start transaction on the first query
# rather than on the `START TRANSACTION`.
tx.query("SELECT 1")
# Start both transactions at the same initial data.
# One should succeed other should fail and retry.
# On next attempt, the latter should succeed
barrier.ready()
lock.acquire()
res = tx.query_single('''
SELECT (
INSERT test::Counter {
name := <str>$name,
value := 1,
} UNLESS CONFLICT ON .name
ELSE (
UPDATE test::Counter
SET { value := .value + 1 }
)
).value
''',
name=name)
lock.release()
return res
client = self.client
if options:
client = client.with_retry_options(options)
client2 = client2.with_retry_options(options)
with futures.ThreadPoolExecutor(2) as pool:
f1 = pool.submit(transaction1, client)
f2 = pool.submit(transaction1, client2)
results = {f1.result(), f2.result()}
self.assertEqual(results, {1, 2})
self.assertEqual(iterations, 3)
def test_client_suggested_concurrency(self):
conargs = self.get_connect_args().copy()
conargs["database"] = self.get_database_name()
conargs["timeout"] = 120
client = edgedb.create_client(**conargs)
self.assertEqual(client.max_concurrency, 1)
client.ensure_connected()
self.assertGreater(client.max_concurrency, 1)
client.close()
client = edgedb.create_client(**conargs, max_concurrency=5)
self.assertEqual(client.max_concurrency, 5)
client.ensure_connected()
self.assertEqual(client.max_concurrency, 5)
client.close()
def _setup_database(sync_queries: Queries) -> None:
client = edgedb.create_client()
sync_queries.seeds.create_test_data(client)
yield
sync_queries.seeds.drop_test_data(client)
client.close()
def sync_client() -> edgedb.Executor:
client = edgedb.create_client()
yield client
client.close()
def _start_cluster(*, cleanup_atexit=True):
global _default_cluster
if isinstance(_default_cluster, Exception):
# when starting a server fails
# don't retry starting one for every TestCase
# because repeating the failure can take a while
raise _default_cluster
if _default_cluster:
return _default_cluster
try:
tmpdir = tempfile.TemporaryDirectory()
status_file = os.path.join(tmpdir.name, 'server-status')
# if running on windows adjust the path for WSL
status_file_unix = _get_wsl_path(status_file)
env = os.environ.copy()
# Make sure the PYTHONPATH of _this_ process does
# not interfere with the server's.
env.pop('PYTHONPATH', None)
edgedb_server = env.get('EDGEDB_SERVER_BINARY', 'edgedb-server')
args = [
edgedb_server,
"--temp-dir",
"--testmode",
f"--emit-server-status={status_file_unix}",
"--port=auto",
"--auto-shutdown",
"--bootstrap-command=ALTER ROLE edgedb { SET password := '******' }",
]
help_args = [edgedb_server, "--help"]
if sys.platform == 'win32':
help_args = ['wsl', '-u', 'edgedb'] + help_args
if "--generate-self-signed-cert" in subprocess.run(
help_args,
universal_newlines=True,
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
encoding="utf-8",
env=env,
cwd=tmpdir.name,
).stdout:
args.append("--generate-self-signed-cert")
if sys.platform == 'win32':
args = ['wsl', '-u', 'edgedb'] + args
if env.get('EDGEDB_DEBUG_SERVER'):
server_stdout = None
else:
server_stdout = subprocess.DEVNULL
p = subprocess.Popen(
args,
env=env,
cwd=tmpdir.name,
stdout=server_stdout,
stderr=subprocess.STDOUT,
)
for _ in range(250):
try:
with open(status_file, 'rb') as f:
for line in f:
if line.startswith(b'READY='):
break
else:
raise RuntimeError('not ready')
break
except Exception:
time.sleep(1)
else:
raise RuntimeError('server status file not found')
data = json.loads(line.split(b'READY=')[1])
con_args = dict(host='localhost', port=data['port'])
if 'tls_cert_file' in data:
if sys.platform == 'win32':
con_args['tls_ca_file'] = os.path.join(tmpdir.name,
"edbtlscert.pem")
subprocess.check_call([
'wsl', '-u', 'edgedb', 'cp', data['tls_cert_file'],
_get_wsl_path(con_args['tls_ca_file'])
])
else:
con_args['tls_ca_file'] = data['tls_cert_file']
client = edgedb.create_client(password='******', **con_args)
client.ensure_connected()
_default_cluster = {
'proc': p,
'client': client,
'con_args': con_args,
}
if 'tls_cert_file' in data:
# Keep the temp dir which we also copied the cert from WSL
_default_cluster['_tmpdir'] = tmpdir
atexit.register(client.close)
except Exception as e:
_default_cluster = e
raise e
return _default_cluster