Skip to content

Commit

Permalink
Fixed create commands and added for all other connection types, added…
Browse files Browse the repository at this point in the history 
… get, list, and delete commands as well
  • Loading branch information
@jamesd-db
jamesd-db committed Jun 6, 2023
1 parent fec7aab commit a179195
Show file tree
Hide file tree
Showing 2 changed files with 302 additions and 20 deletions.
315 changes: 299 additions & 16 deletions databricks_cli/unity_catalog/connection_cli.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,14 +28,14 @@
from databricks_cli.click_types import JsonClickType
from databricks_cli.configure.config import provide_api_client, profile_option, debug_option
from databricks_cli.unity_catalog.api import UnityCatalogApi
from databricks_cli.unity_catalog.utils import del_none, hide, json_file_help, json_string_help, \
from databricks_cli.unity_catalog.utils import hide, json_file_help, json_string_help, \
mc_pretty_format
from databricks_cli.utils import eat_exceptions, CONTEXT_SETTINGS, json_cli_base

# These two options are shared among create and updates, so they are very common
def create_update_common_options(f):
@click.option('--read-only/--no-read-only', is_flag=True, default=None,
help='Whether the location is read-only')
help='Whether the connection is read-only')
@click.option('--comment', default=None,
help='Free-form text description.')
@functools.wraps(f)
Expand Down Expand Up @@ -70,17 +70,6 @@ def wrapper(*args, **kwargs):
f(*args, **kwargs)
return wrapper

# Workaround to prompt for password if user does not specifiy inline JSON or JSON file
# See https://stackoverflow.com/questions/32656571/

def deactivate_prompts(ctx, _, value):
if not value:
for p in ctx.command.params:
if isinstance(p, click.Option) and p.prompt is not None:
p.prompt = None
return value


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create mysql connection with CLI flags.')
@common_create_args
Expand All @@ -90,6 +79,7 @@ def deactivate_prompts(ctx, _, value):
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_mysql_cli(api_client, name, host, port, user,
Expand All @@ -102,22 +92,246 @@ def create_mysql_cli(api_client, name, host, port, user,
data = {
'name': name,
'connection_type': 'MYSQL',
'options_kvpairs': {'host': host, 'port': port, 'user': user, 'password': password},
'options': {'host': host, 'port': port, 'user': user, 'password': password},
'read_only': read_only,
'comment': comment,
}
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create postgresql connection with CLI flags.')
@common_create_args
@create_update_common_options
@click.option(
"--password", prompt=True, hide_input=True,
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_postgresql_cli(api_client, name, host, port, user,
read_only, comment, password):
"""
Create new postgresql connection.
"""
if (name is None) or (host is None) or (port is None) or (user is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'POSTGRESQL',
'options': {'host': host, 'port': port, 'user': user, 'password': password},
'read_only': read_only,
'comment': comment,
}
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))

@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create snowflake connection with CLI flags.')
@common_create_args
@create_update_common_options
@click.option('--sfwarehouse', default=None,
help='Snowflake warehouse name of new connection')
@click.option(
"--password", prompt=True, hide_input=True,
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_snowflake_cli(api_client, name, host, port, user, sfwarehouse,
read_only, comment, password):
"""
Create new snowflake connection.
"""
if (name is None) or (host is None) or (port is None) or (user is None) or \
(sfwarehouse is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'SNOWFLAKE',
'options': {'host': host, 'port': port, 'user': user, 'password': password,
'sfWarehouse': sfwarehouse},
'read_only': read_only,
'comment': comment,
}
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create redshift connection with CLI flags.')
@common_create_args
@create_update_common_options
@click.option(
"--password", prompt=True, hide_input=True,
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_redshift_cli(api_client, name, host, port, user,
read_only, comment, password):
"""
Create new redshift connection.
"""
if (name is None) or (host is None) or (port is None) or (user is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'REDSHIFT',
'options': {'host': host, 'port': port, 'user': user, 'password': password},
'read_only': read_only,
'comment': comment,
}
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create mysql connection with a JSON input.')
short_help='Create sqldw connection with CLI flags.')
@common_create_args
@create_update_common_options
@click.option('--trustservercert', is_flag=True, default=None,
help='Trust the server provided certificate')
@click.option(
"--password", prompt=True, hide_input=True,
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_sqldw_cli(api_client, name, host, port, user, trustservercert,
read_only, comment, password):
"""
Create new sqldw connection.
"""
if (name is None) or (host is None) or (port is None) or (user is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'SQLDW',
'options': {'host': host, 'port': port, 'user': user, 'password': password},
'read_only': read_only,
'comment': comment,
}
if trustservercert != None:
data['options']['trustServerCertificate'] = trustservercert
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create sqlserver connection with CLI flags.')
@common_create_args
@create_update_common_options
@click.option('--trustservercert', is_flag=True, default=None,
help='Trust the server provided certificate')
@click.option(
"--password", prompt=True, hide_input=True,
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_sqlserver_cli(api_client, name, host, port, user, trustservercert,
read_only, comment, password):
"""
Create new sqlserver connection.
"""
if (name is None) or (host is None) or (port is None) or (user is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'SQLSERVER',
'options': {'host': host, 'port': port, 'user': user, 'password': password},
'read_only': read_only,
'comment': comment,
}
if trustservercert != None:
data['options']['trustServerCertificate'] = trustservercert
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create databricks connection with CLI flags.')
@create_update_common_options
@click.option('--name', default=None,
help='Name of new connection')
@click.option('--host', default=None,
help='Host of new connection')
@click.option('--httppath', default=None,
help='HTTP path of new connection')
@click.option(
"--token", prompt=True, hide_input=True,
confirmation_prompt=True
)
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_databricks_cli(api_client, name, host, httppath, token,
read_only, comment):
"""
Create new databricks connection.
"""
if (name is None) or (host is None) or (httppath is None) or (token is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'DATABRICKS',
'options': {'host': host, 'httpPath': httppath, 'personalAccessToken': token},
'read_only': read_only,
'comment': comment,
}
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create online catalog connection with CLI flags.')
@common_create_args
@create_update_common_options
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_online_catalog_cli(api_client, name, host, port, user,
read_only, comment):
"""
Create new online catalog connection.
"""
if (name is None) or (host is None) or (port is None) or (user is None):
raise ValueError('Must provide all required connection parameters')
data = {
'name': name,
'connection_type': 'ONLINE_CATALOG',
'options': {'host': host, 'port': port, 'user': user},
'read_only': read_only,
'comment': comment,
}
con_json = UnityCatalogApi(api_client).create_connection(data)
click.echo(mc_pretty_format(con_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Create connection with a JSON input.')
@json_options
@debug_option
@eat_exceptions
@profile_option
@provide_api_client
def create_json(api_client, json_file, json):
'''
Create new mysql connection with JSON.
Create new connection with an inline JSON or JSON file input.
'''
if (json is None) and (json_file is None):
raise ValueError('Must either provide inline JSON or JSON file.')
Expand All @@ -126,6 +340,53 @@ def create_json(api_client, json_file, json):
UnityCatalogApi(api_client).create_connection(json),
encode_utf8=True)



@click.command(context_settings=CONTEXT_SETTINGS,
short_help='List connections.')
@debug_option
@profile_option
@eat_exceptions
@provide_api_client
def list_connections_cli(api_client, ):
"""
List connections.
"""
locs_json = UnityCatalogApi(api_client).list_connections()
click.echo(mc_pretty_format(locs_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Get a connection.')
@click.option('--name', required=True,
help='Name of the connection to get.')
@debug_option
@profile_option
@eat_exceptions
@provide_api_client
def get_connection_cli(api_client, name):
"""
Get a connection.
"""
loc_json = UnityCatalogApi(api_client).get_connection(name)
click.echo(mc_pretty_format(loc_json))


@click.command(context_settings=CONTEXT_SETTINGS,
short_help='Delete a connection.')
@click.option('--name', required=True,
help='Name of the connection to delete.')
@debug_option
@profile_option
@eat_exceptions
@provide_api_client
def delete_connection_cli(api_client, name):
"""
Delete a connection.
"""
UnityCatalogApi(api_client).delete_connection(name)


@click.group()
def create_group(): # pragma: no cover
pass
Expand All @@ -137,10 +398,32 @@ def connections_group(): # pragma: no cover

def register_connection_commands(cmd_group):
# Register deprecated "verb-noun" commands for backward compatibility.
cmd_group.add_command(hide(create_json), name='create-connection-json')
cmd_group.add_command(hide(create_mysql_cli), name='create-mysql-connection')
cmd_group.add_command(hide(create_postgresql_cli), name='create-postgresql-connection')
cmd_group.add_command(hide(create_snowflake_cli), name='create-snowflake-connection')
cmd_group.add_command(hide(create_redshift_cli), name='create-redshift-connection')
cmd_group.add_command(hide(create_sqldw_cli), name='create-sqldw-connection')
cmd_group.add_command(hide(create_sqlserver_cli), name='create-sqlserver-connection')
cmd_group.add_command(hide(create_databricks_cli), name='create-databricks-connection')
cmd_group.add_command(hide(create_online_catalog_cli), name='create-online-catalog-connection')




# Register command group.
create_group.add_command(create_mysql_cli, name='mysql')
create_group.add_command(create_postgresql_cli, name='postgresql')
create_group.add_command(create_snowflake_cli, name='snowflake')
create_group.add_command(create_redshift_cli, name='redshift')
create_group.add_command(create_sqldw_cli, name='sqldw')
create_group.add_command(create_sqlserver_cli, name='sqlserver')
create_group.add_command(create_databricks_cli, name='databricks')
create_group.add_command(create_online_catalog_cli, name='online-catalog')

connections_group.add_command(create_group, name='create')
connections_group.add_command(create_json, name='create-json')
connections_group.add_command(list_connections_cli, name='list')
connections_group.add_command(get_connection_cli, name='get')
connections_group.add_command(delete_connection_cli, name='delete')
cmd_group.add_command(connections_group, name='connection')
Loading

0 comments on commit a179195

Please sign in to comment.