feat: use default session connection #87

GarrettWu · 2023-10-04T19:30:20Z

Thank you for opening a Pull Request! Before submitting your PR, there are a few things you can do to make sure it goes smoothly:

Make sure to open an issue as a bug/issue before writing your code! That way we can discuss the change, evaluate designs, and agree on the general idea
Ensure the tests and linter pass
Code coverage does not decrease (if any source code was changed)
Appropriate docs were updated (if necessary)

Fixes #<issue_number_goes_here> 🦕
http://b/303283207

bigframes/session.py

bigframes/remote_function.py

shobsi · 2023-10-10T21:44:14Z

bigframes/clients.py

@@ -162,3 +164,25 @@ def _get_service_account_if_connection_exists(
            pass

        return service_account
+
+
+def get_connection_name_full(


Would it be tidier to make it a class method in BqConnectionManager?

I think the function is only a helper that does string manipulation. Instead of the BqConnectionManager which contains the states of the clients. So separate them.

IMHO it would still make sense. A @classmethod would mean that it does not depend on the instance level state. But conceptually it would fit in nicely - BqConnectionManager sounds like the entity which is supposed to know the low level details of a connection and can provide helper functions about the same.

That could work. Will address in another PR.

shobsi · 2023-10-10T21:49:54Z

tests/system/small/test_remote_function.py

@@ -465,6 +465,36 @@ def square(x):
    assert_pandas_df_equal_ignore_ordering(bf_result, pd_result)


+@pytest.mark.flaky(retries=2, delay=120)
+def test_remote_function_default_connection(scalars_dfs, dataset_id):
+    @rf.remote_function([int], int, dataset=dataset_id)


direct version is deprecated, recommend testing session.remote_function or bigframes.pandas.remote_function interface.

python-bigquery-dataframes/bigframes/remote_function.py

Line 609 in 228aeba

This is an internal method. Please use :func:`bigframes.pandas.remote_function` instead.

Changed to bigframes.pandas.remote_function.

But in that case, if all the use cases are through session.remote_function or bigframes.pandas.remote_function, then it will be always the case that rf.remote_function's input session is not None. And we won't need to resolve to bpd.get_global_session() in rf.remote_function.

I'll leave it for you to resolve.

I think it's okay, the goal of this test is to assert that default connection takes effect irrespective of the session.

We can add a separate large test test_remote_function_direct_defaults_to_global_session specifically for that change.

tests/unit/test_clients.py

shobsi · 2023-10-11T07:08:31Z

bigframes/clients.py

@@ -162,3 +164,25 @@ def _get_service_account_if_connection_exists(
            pass

        return service_account
+
+
+def get_connection_name_full(


IMHO it would still make sense. A @classmethod would mean that it does not depend on the instance level state. But conceptually it would fit in nicely - BqConnectionManager sounds like the entity which is supposed to know the low level details of a connection and can provide helper functions about the same.

shobsi · 2023-10-11T07:12:44Z

bigframes/clients.py

+
+
+def get_connection_name_full(
+    connection_name: Optional[str], default_project: str, default_location: str


awkward to have optional argument followed by mandatory arguments, would be nice to have it in the end

Optional only means it can be value None. It is OK to put it at the front, as it is the most important param of the function.

Not really "optional", which means it has a default value. Then it has to be at the end.

shobsi · 2023-10-11T07:31:58Z

bigframes/ml/llm.py

@@ -48,7 +49,14 @@ def __init__(
        connection_name: Optional[str] = None,
    ):
        self.session = session or bpd.get_global_session()
-        self.connection_name = connection_name or self.session._bq_connection
+
+        connection_name = connection_name or self.session._bq_connection


I thought about this a bit, we are resolving between 3 potential sources of connection:

provided by the user to llm, i.e. connection_name arg

defined by the user in the session, i.e. self.session._bq_connection

the default defined in clients.py

How about we set the default in Session.__init__ like below?

self._bq_connection = context.bq_connection or "bigframes-default-connection"

We are assuming other session defaults there, such as self._location = "US".

Oh yea, I did this way at first place, but then I found remote_function doesn't always have a session object. So I had to put the default away from session. Then we added default session...

We can move it to session.

shobsi · 2023-10-11T07:34:17Z

bigframes/clients.py

@@ -162,3 +164,25 @@ def _get_service_account_if_connection_exists(
            pass

        return service_account
+
+
+def get_connection_name_full(


[nit] since it takes a connection and returns a connection after resolving the format, would prefer naming it resolve_full_connection_name

Will address in another PR.

feat: use default session connection

4d17229

GarrettWu requested review from shobsi and jiaxunwu October 4, 2023 19:30

GarrettWu requested review from a team as code owners October 4, 2023 19:30

product-auto-label bot added size: xs Pull request size is extra small. api: bigquery Issues related to the googleapis/python-bigquery-dataframes API. labels Oct 4, 2023

GarrettWu marked this pull request as draft October 4, 2023 19:35

jiaxunwu approved these changes Oct 4, 2023

View reviewed changes

shobsi reviewed Oct 4, 2023

View reviewed changes

bigframes/session.py Outdated Show resolved Hide resolved

shobsi reviewed Oct 4, 2023

View reviewed changes

bigframes/session.py Outdated Show resolved Hide resolved

update docs

b0e9215

product-auto-label bot added size: s Pull request size is small. and removed size: xs Pull request size is extra small. labels Oct 4, 2023

Merge remote-tracking branch 'github/main' into garrettwu-conn

dd05552

GarrettWu marked this pull request as ready for review October 9, 2023 04:38

GarrettWu and others added 3 commits October 9, 2023 10:14

Merge branch 'main' into garrettwu-conn

6739428

supplement partial connection_name

4ee895e

Merge remote-tracking branch 'github/main' into garrettwu-conn

dcf93fa

product-auto-label bot added size: m Pull request size is medium. and removed size: s Pull request size is small. labels Oct 10, 2023

GarrettWu and others added 3 commits October 10, 2023 00:43

fix test

c79c86e

fix test

62ac080

Merge branch 'main' into garrettwu-conn

cd6c86a

shobsi reviewed Oct 10, 2023

View reviewed changes

bigframes/remote_function.py Show resolved Hide resolved

shobsi reviewed Oct 10, 2023

View reviewed changes

bigframes/remote_function.py Outdated Show resolved Hide resolved

shobsi reviewed Oct 10, 2023

View reviewed changes

tests/unit/test_clients.py Outdated Show resolved Hide resolved

Merge remote-tracking branch 'github/main' into garrettwu-conn

71ec5b0

resolve comments

aa42ff7

product-auto-label bot added size: l Pull request size is large. and removed size: m Pull request size is medium. labels Oct 11, 2023

shobsi reviewed Oct 11, 2023

View reviewed changes

GarrettWu merged commit 4ae4ef9 into main Oct 11, 2023

GarrettWu deleted the garrettwu-conn branch October 11, 2023 16:49

release-please bot mentioned this pull request Oct 11, 2023

chore(main): release 0.7.0 #86

Merged

GarrettWu mentioned this pull request Oct 12, 2023

chore: address last comments of PR#87 #102

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: use default session connection #87

feat: use default session connection #87

GarrettWu commented Oct 4, 2023

shobsi Oct 10, 2023

GarrettWu Oct 11, 2023

shobsi Oct 11, 2023 •

edited

Loading

GarrettWu Oct 11, 2023

shobsi Oct 10, 2023

GarrettWu Oct 11, 2023

shobsi Oct 11, 2023

shobsi Oct 11, 2023 •

edited

Loading

shobsi Oct 11, 2023

GarrettWu Oct 11, 2023

shobsi Oct 11, 2023

GarrettWu Oct 11, 2023

shobsi Oct 11, 2023

GarrettWu Oct 11, 2023



		def get_connection_name_full(
		connection_name: Optional[str], default_project: str, default_location: str

feat: use default session connection #87

feat: use default session connection #87

Conversation

GarrettWu commented Oct 4, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shobsi Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shobsi Oct 11, 2023 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

shobsi Oct 11, 2023 •

edited

Loading

shobsi Oct 11, 2023 •

edited

Loading