Skip to content

Commit

Permalink
Prepend user-supplied user-agent (#122)
Browse files Browse the repository at this point in the history
* Prepend user-supplied user-agent

* Skip cloud tests

* Fix lint
  • Loading branch information
phillipleblanc authored Dec 17, 2024
1 parent 8a27c04 commit 87f6323
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 18 deletions.
15 changes: 1 addition & 14 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -33,20 +33,7 @@ client = Client(
api_key='API_KEY',
flight_url="grpc+tls://flight.spiceai.io"
)
data = client.query('SELECT * FROM eth.recent_blocks LIMIT 10;', timeout=5*60)
pd = data.read_pandas()
```

**Firecache Query (Available if firecache is enabled)**

```python
from spicepy import Client

client = Client(
api_key='API_KEY',
flight_url="grpc+tls://flight.spiceai.io"
)
data = client.fire_query('SELECT * FROM eth.recent_blocks LIMIT 10;', timeout=5*60)
data = client.query('SELECT * FROM taxi_trips LIMIT 10;', timeout=5*60)
pd = data.read_pandas()
```

Expand Down
9 changes: 6 additions & 3 deletions spicepy/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -62,10 +62,13 @@ def read_cert(self, tls_root_cert):
class _SpiceFlight:
@staticmethod
def _user_agent(custom_user_agent=None):
# headers kwargs claim to support Tuple[str, str], but it's actually Tuple[bytes, bytes] :|
# headers kwargs claim to support Tuple[str, str], but it's actually Tuple[bytes, bytes]
# Open issue in Arrow: https://github.com/apache/arrow/issues/35288
user_agent = custom_user_agent or config.SPICE_USER_AGENT
return (str.encode("user-agent"), str.encode(user_agent))

# Prepend the custom user agent (if provided) to the default user agent
if custom_user_agent:
return (str.encode("user-agent"), str.encode(f"{custom_user_agent} {config.SPICE_USER_AGENT}"))
return (str.encode("user-agent"), str.encode(config.SPICE_USER_AGENT))

def __init__(self, grpc: str, api_key: str, tls_root_certs, user_agent=None):
self._flight_client = flight.connect(grpc, tls_root_certs=tls_root_certs)
Expand Down
4 changes: 3 additions & 1 deletion tests/test_main.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,9 @@

# Skip cloud tests if TEST_SPICE_CLOUD is not set to true
def skip_cloud():
skip = os.environ.get("TEST_SPICE_CLOUD") != "true"
# skip = os.environ.get("TEST_SPICE_CLOUD") != "true"
# Skipping all cloud tests for now
skip = True
return pytest.mark.skipif(skip, reason="Cloud tests disabled")


Expand Down

0 comments on commit 87f6323

Please sign in to comment.