Skip to content

Commit

Permalink
Wire up and revert some prints
Browse files Browse the repository at this point in the history
  • Loading branch information
odeke-em committed Jan 18, 2025
1 parent ac5d30e commit 7bbc4d3
Show file tree
Hide file tree
Showing 6 changed files with 23 additions and 33 deletions.
8 changes: 6 additions & 2 deletions google/cloud/spanner_v1/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -665,8 +665,12 @@ def inject_retry_header_control(api):

orig_getattribute = getattr(target, "__getattribute__")

def patched_getattribute(*args, **kwargs):
attr = orig_getattribute(*args, **kwargs)
def patched_getattribute(obj, key, *args, **kwargs):
if key.startswith("_"):
return orig_getattribute(obj, key, *args, **kwargs)

attr = orig_getattribute(obj, key, *args, **kwargs)
print("args", args, "attr.dir", dir(attr))

# 0. If we already patched it, we can return immediately.
if getattr(attr, "_patched", None) is not None:
Expand Down
29 changes: 10 additions & 19 deletions google/cloud/spanner_v1/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,6 @@
from google.cloud.spanner_v1._helpers import (
_metadata_with_prefix,
_metadata_with_leader_aware_routing,
AtomicCounter,
)
from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
from google.cloud.spanner_v1 import RequestOptions
Expand Down Expand Up @@ -228,20 +227,15 @@ def commit(
trace_attributes,
observability_options=observability_options,
):
attempt = AtomicCounter(0)
next_nth_request = database._next_nth_request

all_metadata = database.metadata_with_request_id(
next_nth_request,
attempt.increment(),
metadata,
)
method = functools.partial(
api.commit,
request=request,
metadata=all_metadata,
metadata=database.metadata_with_request_id(
database._next_nth_request,
1,
metadata,
),
)

deadline = time.time() + kwargs.get(
"timeout_secs", DEFAULT_RETRY_TIMEOUT_SECS
)
Expand Down Expand Up @@ -359,18 +353,15 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
trace_attributes,
observability_options=observability_options,
):
next_nth_request = database._next_nth_request
all_metadata = database.metadata_with_request_id(
next_nth_request,
0,
metadata,
)
method = functools.partial(
api.batch_write,
request=request,
metadata=all_metadata,
metadata=database.metadata_with_request_id(
database._next_nth_request,
1,
metadata,
),
)

response = _retry(
method,
allowed_exceptions={
Expand Down
1 change: 0 additions & 1 deletion google/cloud/spanner_v1/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -787,7 +787,6 @@ def execute_pdml():
query_options=query_options,
request_options=request_options,
)

method = functools.partial(
api.execute_streaming_sql,
metadata=self.metadata_with_request_id(
Expand Down
3 changes: 0 additions & 3 deletions google/cloud/spanner_v1/pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,10 @@
"""Pools managing shared Session objects."""

import datetime
import random
import queue
import time

from google.cloud.exceptions import NotFound
from google.api_core.exceptions import ServiceUnavailable
from google.cloud.spanner_v1 import BatchCreateSessionsRequest
from google.cloud.spanner_v1 import Session
from google.cloud.spanner_v1._helpers import (
Expand Down Expand Up @@ -253,7 +251,6 @@ def bind(self, database):
f"Creating {request.session_count} sessions",
span_event_attributes,
)

resp = api.batch_create_sessions(
request=request,
metadata=database.metadata_with_request_id(
Expand Down
4 changes: 1 addition & 3 deletions google/cloud/spanner_v1/services/spanner/transports/grpc.py
Original file line number Diff line number Diff line change
Expand Up @@ -316,9 +316,7 @@ def batch_create_sessions(
request_serializer=spanner.BatchCreateSessionsRequest.serialize,
response_deserializer=spanner.BatchCreateSessionsResponse.deserialize,
)
fn = self._stubs["batch_create_sessions"]
print("\033[32minvoking batch_create_sessionhex_id", hex(id(fn)), "\033[00m")
return fn
return self._stubs["batch_create_sessions"]

@property
def get_session(self) -> Callable[[spanner.GetSessionRequest], spanner.Session]:
Expand Down
11 changes: 6 additions & 5 deletions google/cloud/spanner_v1/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -205,16 +205,17 @@ def exists(self):
)
)

all_metadata = database.metadata_with_request_id(
database._next_nth_request, 1, metadata
)

observability_options = getattr(self._database, "observability_options", None)
with trace_call(
"CloudSpanner.GetSession", self, observability_options=observability_options
) as span:
try:
api.get_session(name=self.name, metadata=all_metadata)
api.get_session(
name=self.name,
metadata=database.metadata_with_request_id(
database._next_nth_request, 1, metadata
),
)
if span:
span.set_attribute("session_found", True)
except NotFound:
Expand Down

0 comments on commit 7bbc4d3

Please sign in to comment.