def connect(ctx):
try:
conn = edgedb.connect(
host=ctx.obj['host'],
port=ctx.obj['port'],
user=ctx.obj['user'],
database=ctx.obj['database'],
admin=ctx.obj['admin'],
password=ctx.obj['password'],
)
except edgedb.AuthenticationError:
if (ctx.obj['password'] is None
and ctx.obj['password_prompt'] is not None):
password = ctx.obj['password_prompt']
conn = edgedb.connect(
host=ctx.obj['host'],
port=ctx.obj['port'],
user=ctx.obj['user'],
database=ctx.obj['database'],
admin=ctx.obj['admin'],
password=password,
)
else:
raise
ctx.obj['conn'] = conn
ctx.call_on_close(lambda: conn.close())
def execute_script(conn_args, data):
password_prompt = conn_args.pop('password_prompt', None)
try:
con = edgedb.connect(**conn_args)
except edgedb.AuthenticationError:
if password_prompt:
password = password_prompt()
con = edgedb.connect(**{**conn_args, 'password': password})
else:
raise
try:
queries = lexutils.split_edgeql(data)[0]
ctx = context.ReplContext()
for query in queries:
try:
ret = con.fetchall(query)
except Exception as ex:
render.render_exception(
ctx,
ex,
query=query)
return 1
else:
render.render_binary(
ctx,
ret,
max_width=80)
finally:
con.close()
def ensure_connection(self):
try:
if self.connection is None:
self.connection = edgedb.connect(**self.conn_args)
elif self.connection.is_closed():
self.connection = edgedb.connect(**self.conn_args)
except edgedb.AuthenticationError:
if (self.conn_args['password'] is None
and self._password_prompt is not None
and not self._password_prompted):
try:
password = self._password_prompt()
self._password_prompted = True
self.connection = edgedb.connect(
**{**self.conn_args, 'password': password})
except Exception as e:
self.connection = None
reason = str(e)
else:
self.conn_args = self.conn_args.set('password', password)
except Exception as e:
self.connection = None
reason = str(e)
if self.connection is None:
dbname = self.conn_args.get("database")
if not dbname:
dbname = 'EdgeDB'
print(f'Could not establish connection to {dbname}: {reason}')
exit(1)
self.connection.add_log_listener(self.on_edgedb_log_message)
def close(self) -> typing.NoReturn:
edgedb.connect(
dsn=self.dsn,
host=self.host,
port=self.port,
admin=self.admin,
user=self.user,
password=self.password,
database=self.database,
timeout=self.timeout,
).close()
def drop_and_load_db(schema, dsn, database):
drop_all_connection()
with closing(edgedb.connect(dsn, database="edgedb")) as connection:
with suppress(InvalidReferenceError):
connection.execute(f"DROP DATABASE {database}")
print("Creating the database...")
connection.execute(f"CREATE DATABASE {database}")
with closing(edgedb.connect(dsn, database=database)) as connection:
with open(schema) as schema_f:
content = schema_f.read()
connection.execute(content)
connection.execute("POPULATE MIGRATION")
print("Committing the schema...")
connection.execute("COMMIT MIGRATION")
def setup_test_database():
from gen3.server import config
from gen3.server.app import app
print("Connecting to root database for test database setup...")
conn = edgedb.connect(
host=config.DB_HOST,
port=config.DB_PORT,
user=config.DB_USER,
password=str(config.DB_PASSWORD),
admin=config.DB_ADMIN,
database=config.DB_DATABASE_ROOT,
timeout=config.DB_CONNECT_TIMEOUT,
)
try:
try:
print("Attempting to create test database...")
conn.execute(f"CREATE DATABASE {config.DB_DATABASE}")
print("Test database created.")
except edgedb.InternalServerError:
print("Test database existed, reusing.")
print("Running migration...")
with TestClient(app) as client:
client.get("/migrate")
print("Migration done.")
yield
if not config.TEST_KEEP_DB:
print("Dropping test database...")
conn.execute(f"DROP DATABASE {config.DB_DATABASE}")
print("Test database dropped.")
finally:
conn.close()
def new_connection(
self,
*,
database: Optional[str] = None,
**extra_con_args: Any,
) -> edgedb.BlockingIOConnection:
try:
return edgedb.connect(
host=self.host,
port=self.port,
user=self.user,
database=self.database if database is None else database,
admin=self.admin,
password=self.password,
**extra_con_args)
except edgedb.AuthenticationError as ex:
if self.password is None and self.allow_password_request:
password = self.get_password()
if password is None:
raise click.ClickException(str(ex))
return self.new_connection(database=database, **extra_con_args)
else:
raise click.ClickException(str(ex))
except Exception as ex:
raise click.ClickException(str(ex))
def test_connect_sync_01(self):
orig_conn_args = self.get_connect_args()
conn_args = orig_conn_args.copy()
conn_args['port'] = self.port
with self.assertRaisesRegex(
edgedb.ClientConnectionError,
f'(?s).*Is the server running.*port {self.port}.*'):
conn_args['host'] = '127.0.0.1'
edgedb.connect(**conn_args)
with self.assertRaisesRegex(
edgedb.ClientConnectionError,
f'(?s).*Is the server running.*port {self.port}.*'):
conn_args['host'] = orig_conn_args['host']
edgedb.connect(**conn_args)
def gen_meta_grammars(names):
"""Generate keywords, builtins, operators, etc. which can be used
for EdgeQL and SDL grammars.
NAME - at the moment there's only one option 'edgeql'
"""
if names:
for name in names:
if name not in NAMES:
die(f'{name} is not a valid NAME')
if len(names) > 2:
die(f'too many NAMES')
con = None
try:
con = edgedb.connect(user=edgedb_defines.EDGEDB_SUPERUSER,
database=edgedb_defines.EDGEDB_SUPERUSER_DB)
main(names, con)
except Exception as ex:
die(str(ex))
finally:
if con is not None:
con.close()
def get_con() -> BlockingIOConnection:
con = connect(
host=settings.EDGEDB_HOST,
database=settings.EDGEDB_DB,
user=settings.EDGEDB_USER,
password=settings.EDGEDB_PASSWORD,
)
return con
def execute_conflict(self, name='counter2', options=None):
con_args = self.get_connect_args().copy()
con_args.update(database=self.get_database_name())
con2 = edgedb.connect(**con_args)
self.addCleanup(con2.close)
barrier = Barrier(2)
lock = threading.Lock()
iterations = 0
def transaction1(con):
for tx in con.retrying_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_one('''
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
con = self.con
if options:
con = con.with_retry_options(options)
con2 = con2.with_retry_options(options)
with futures.ThreadPoolExecutor(2) as pool:
f1 = pool.submit(transaction1, con)
f2 = pool.submit(transaction1, con2)
results = {f1.result(), f2.result()}
self.assertEqual(results, {1, 2})
self.assertEqual(iterations, 3)
def setUp(self):
super().setUp()
cls = type(self)
cls.async_con = cls.con
conargs = cls.get_connect_args().copy()
conargs.update(dict(database=cls.async_con.dbname))
cls.con = edgedb.connect(**conargs)
def connect_sync(
self,
connection: typing.Optional[EdgeDBConnection] = None,
) -> edgedb.BlockingIOConnection:
return edgedb.connect(
dsn=self.dsn,
host=self.host,
port=self.port,
admin=bool(self.admin),
user=self.user,
password=self.password,
database=self.database,
timeout=self.timeout,
)
def command_connect(self, args):
new_db = args.strip()
new_args = self.conn_args.set('database', new_db)
try:
new_connection = edgedb.connect(**new_args)
if self.context.introspect_types:
self.introspect_db(new_connection)
except Exception:
print(f'Could not establish connection to {new_db!r}', flush=True)
return
self.connection.close()
self.connection = new_connection
self.conn_args = new_args
def _test_connection(self, timeout=60):
while True:
started = time.monotonic()
try:
conn = edgedb.connect(host=str(self._runstate_dir),
port=self._effective_port,
admin=True,
database='edgedb',
user='******',
timeout=timeout)
except (OSError, socket.error, TimeoutError,
edgedb.ClientConnectionError):
timeout -= (time.monotonic() - started)
if timeout > 0.05:
time.sleep(0.05)
timeout -= 0.05
continue
raise ClusterError(f'could not connect to edgedb-server '
f'within {timeout} seconds')
else:
conn.close()
return
universal_newlines=True,
text=True,
)
for line in process.stdout:
prefix, _, raw_data = line.partition(":")
if prefix == "EDGEDB_SERVER_DATA":
break
else:
raise ValueError("No EDGEDB_SERVER_DATA on stdout")
server_data = json.loads(raw_data)
# TO-DO(medium): admin=True is deprecated
connection = edgedb.connect(
host=server_data["runstate_dir"],
port=server_data["port"],
user="******",
database="edgedb",
admin=True,
)
bootstrap_connection(connection)
config.database.options = {
"host": server_data["runstate_dir"],
"port": server_data["port"],
"user": TEST_DATABASE_USER,
"password": TEST_DATABASE_PASSWORD,
}
def update_db(change_db_schema):
assert config.database.database == TEST_DATABASE_NAME
def connect(self, **kwargs):
connect_args = self.get_connect_args().copy()
connect_args.update(kwargs)
return edgedb.connect(**connect_args)
def get_new_connection(*args, **kwargs):
_apply_defaults(kwargs)
connection = edgedb.connect(*args, **kwargs)
return closing(connection)
def connect(ctx):
return edgedb.connect()
import edgedb
from edgeql_queries import from_path
conn = edgedb.connect()
queries = from_path("./queries.edgeql", async_driver=False)
user = queries.select_person_by_ip(conn, user_ip="127.0.0.1")
if user is None:
print("oops, no user was found")
else:
print("user:", user.username)
import edgedb
import datetime
conn = edgedb.connect("edgedb://edgedb@localhost/ambv")
conn.execute("""
CREATE TYPE User {
CREATE REQUIRED PROPERTY name -> str;
CREATE PROPERTY date_of_birth -> cal::local_date;
}
""")
conn.query("""
INSERT User {
name := <str>$name,
date_of_birth := <cal::local_date>$dob
}
""",
name="Lukasz Langa",
dob=datetime.date(1985, 3, 7))
conn.query("SELECT User {name, date_of_birth);")
def connect(cls, loop, cluster, database=None):
conargs = cluster.get_connect_args().copy()
conargs.update(dict(
user='******', database=database, port=conargs['port'] + 1))
return loop.run_until_complete(edgedb.connect(**conargs))
import edgedb
conn = edgedb.connect("edgedb://edgedb@localhost/edgedb")
user = conn.query_single(
"""
SELECT Person {
username
}
FILTER .ip = <str>$user_ip
LIMIT 1
""",
user_ip="127.0.0.1",
)
if user is None:
print("oops, no user was found")
else:
print("user:", user.username)
def connect(dsn, database, *args, **kwargs):
return closing(edgedb.connect(dsn=dsn, database=database, *args, **kwargs))
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 = (
re.sub(r'^([A-Z]):', lambda m: f'/mnt/{m.group(1)}', status_file)
.replace("\\", '/')
.lower()
)
args = [
os.environ.get('EDGEDB_SERVER_BINARY', 'edgedb-server'),
"--temp-dir",
"--testmode",
f"--emit-server-status={status_file_unix}",
"--port=auto",
"--auto-shutdown",
"--bootstrap-command=ALTER ROLE edgedb { SET password := '******' }",
]
if sys.platform == 'win32':
args = ['wsl', '-u', 'edgedb'] + args
env = os.environ.copy()
# Make sure the PYTHONPATH of _this_ process does
# not interfere with the server's.
env.pop('PYTHONPATH', None)
p = subprocess.Popen(args, env=env, cwd=tmpdir.name)
for i 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 = edgedb.connect(
host='localhost', port=data['port'], password='******')
_default_cluster = {
'proc': p,
'con': con,
'con_args': {
'host': 'localhost',
'port': data['port'],
}
}
atexit.register(con.close)
except Exception as e:
_default_cluster = e
raise e
return _default_cluster
