Add optional region arg to "unity-catalog metastores create" (#547)

* Added region field * Line too long * Update databricks_cli/unity_catalog/metastore_cli.py Co-authored-by: Pieter Noordhuis <[email protected]>
databricks · Aug 26, 2022 · a574280 · a574280
1 parent 8bf8310
commit a574280
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/databricks_cli/unity_catalog/api.py b/databricks_cli/unity_catalog/api.py
@@ -31,8 +31,8 @@ def __init__(self, api_client):
 
     # Metastore APIs
 
-    def create_metastore(self, name, storage_root):
-        return self.client.create_metastore(name, storage_root)
+    def create_metastore(self, name, storage_root, region):
+        return self.client.create_metastore(name, storage_root, region)
 
     def list_metastores(self):
         return self.client.list_metastores()

diff --git a/databricks_cli/unity_catalog/metastore_cli.py b/databricks_cli/unity_catalog/metastore_cli.py
@@ -39,15 +39,18 @@
 @click.option('--name', required=True, help='Name of the new metastore.')
 @click.option('--storage-root', required=True,
               help='Storage root URL for the new metastore.')
+@click.option('--region', required=False,
+              help='Region for the new metastore. ' +
+                   'Only workspaces located in the same region can be assigned to this metastore.')
 @debug_option
 @profile_option
 @eat_exceptions
 @provide_api_client
-def create_metastore_cli(api_client, name, storage_root):
+def create_metastore_cli(api_client, name, storage_root, region):
     """
     Create new metastore.
     """
-    metastore_json = UnityCatalogApi(api_client).create_metastore(name, storage_root)
+    metastore_json = UnityCatalogApi(api_client).create_metastore(name, storage_root, region)
     click.echo(mc_pretty_format(metastore_json))
 
 

diff --git a/databricks_cli/unity_catalog/uc_service.py b/databricks_cli/unity_catalog/uc_service.py
@@ -31,11 +31,13 @@ def __init__(self, client):
 
     # Metastore Operations
 
-    def create_metastore(self, name, storage_root, headers=None):
+    def create_metastore(self, name, storage_root, region, headers=None):
         _data = {
             'name': name,
             'storage_root': storage_root,
         }
+        if region is not None:
+            _data['region'] = region
         return self.client.perform_query('POST', '/unity-catalog/metastores', data=_data,
                                          headers=headers)