eladkal commented on a change in pull request #22331:
URL: https://github.com/apache/airflow/pull/22331#discussion_r829198970
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -352,3 +355,44 @@ def get_repo_by_path(self, path: str) -> Optional[str]:
return str(result['object_id'])
return None
+
+ def import_notebook(self, notebook_name: str, raw_code: str, language:
str, overwrite: bool = True, format: str = 'SOURCE'):
+ """
+ Import a local notebook from Airflow into Databricks FS. Notebooks
saved to /Shared/airflow dbfs
+
+ Utility function to call the ``2.0/workspace/import`` endpoint.
+
+ :param notebook_name: String name of notebook on Databricks FS
+ :param raw_code: String of non-encoded code
+ :param language: Use one of the following strings 'SCALA', 'PYTHON',
'SQL', OR 'R'
+ :param overwrite: Boolean flag specifying whether to overwrite
existing object. It is true by default
+ :return: full dbfs notebook path
+ """
+ #enforce language options
+ language_options = ['SCALA', 'PYTHON', 'SQL', 'R']
+ if language.upper() not in language_options:
+ raise ValueError(f"results: language must be one of the following:
{str(language_options)}")
+
+ # enforce format options
+ format_options = ['SOURCE', 'HTML', 'JUPYTER', 'DBC']
+ if format.upper() not in format_options:
+ raise ValueError(f"results: format must be one of the following:
{str(format_options)}")
Review comment:
Same question/concern
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -352,3 +355,44 @@ def get_repo_by_path(self, path: str) -> Optional[str]:
return str(result['object_id'])
return None
+
+ def import_notebook(self, notebook_name: str, raw_code: str, language:
str, overwrite: bool = True, format: str = 'SOURCE'):
+ """
+ Import a local notebook from Airflow into Databricks FS. Notebooks
saved to /Shared/airflow dbfs
+
+ Utility function to call the ``2.0/workspace/import`` endpoint.
+
+ :param notebook_name: String name of notebook on Databricks FS
+ :param raw_code: String of non-encoded code
+ :param language: Use one of the following strings 'SCALA', 'PYTHON',
'SQL', OR 'R'
+ :param overwrite: Boolean flag specifying whether to overwrite
existing object. It is true by default
+ :return: full dbfs notebook path
+ """
+ #enforce language options
+ language_options = ['SCALA', 'PYTHON', 'SQL', 'R']
+ if language.upper() not in language_options:
+ raise ValueError(f"results: language must be one of the following:
{str(language_options)}")
Review comment:
What happens if you pass unsupported language to databricks? Will the
API tell you that the language is not supported?
I would prefer not to maintain our own list of languages because this means
that if in the future new languages will be supported users will be blocked
from using it due to enforcement from Airflow side.
##########
File path: airflow/providers/databricks/hooks/databricks.py
##########
@@ -352,3 +355,44 @@ def get_repo_by_path(self, path: str) -> Optional[str]:
return str(result['object_id'])
return None
+
+ def import_notebook(self, notebook_name: str, raw_code: str, language:
str, overwrite: bool = True, format: str = 'SOURCE'):
+ """
+ Import a local notebook from Airflow into Databricks FS. Notebooks
saved to /Shared/airflow dbfs
+
+ Utility function to call the ``2.0/workspace/import`` endpoint.
+
+ :param notebook_name: String name of notebook on Databricks FS
+ :param raw_code: String of non-encoded code
+ :param language: Use one of the following strings 'SCALA', 'PYTHON',
'SQL', OR 'R'
+ :param overwrite: Boolean flag specifying whether to overwrite
existing object. It is true by default
+ :return: full dbfs notebook path
+ """
+ #enforce language options
+ language_options = ['SCALA', 'PYTHON', 'SQL', 'R']
+ if language.upper() not in language_options:
+ raise ValueError(f"results: language must be one of the following:
{str(language_options)}")
+
+ # enforce format options
+ format_options = ['SOURCE', 'HTML', 'JUPYTER', 'DBC']
+ if format.upper() not in format_options:
+ raise ValueError(f"results: format must be one of the following:
{str(format_options)}")
+
+ # encode notebook
+ encodedBytes = base64.b64encode(raw_code.encode("utf-8"))
+ encodedStr = str(encodedBytes, "utf-8")
+
+ # create parent directory if not exists
+ self._do_api_call(WORKSPACE_MKDIR_ENDPOINT, {'path':
"/Shared/airflow"})
+
+ # upload notebook
+ json = {
+ 'path': f'/Shared/airflow/{notebook_name}',
+ 'content': encodedStr,
+ 'language': language,
+ 'overwrite': str(overwrite).lower(),
+ 'format': format
+ }
+ self._do_api_call(WORKSPACE_IMPORT_ENDPOINT, json)
+
+ return f'/Shared/airflow/{notebook_name}'
Review comment:
I don't know databricks but this return is strange to me.
The user provided `notebook_name` as parameter to this function. What is the
benefit in returning this string?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscr...@airflow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
