sessionFromContext extracts the session from gRPC metadata. The session token is expected in the "x-duckgres-session" header.
(ctx context.Context)
| 158 | // sessionFromContext extracts the session from gRPC metadata. |
| 159 | // The session token is expected in the "x-duckgres-session" header. |
| 160 | func (h *FlightSQLHandler) sessionFromContext(ctx context.Context) (*Session, error) { |
| 161 | md, ok := metadata.FromIncomingContext(ctx) |
| 162 | if !ok { |
| 163 | return nil, status.Error(codes.Unauthenticated, "missing metadata") |
| 164 | } |
| 165 | |
| 166 | tokens := md.Get("x-duckgres-session") |
| 167 | if len(tokens) == 0 { |
| 168 | return nil, status.Error(codes.Unauthenticated, "missing x-duckgres-session header") |
| 169 | } |
| 170 | |
| 171 | session, ok := h.pool.GetSession(tokens[0]) |
| 172 | if !ok { |
| 173 | return nil, status.Error(codes.Unauthenticated, "session not found") |
| 174 | } |
| 175 | if h.pool.sharedWarmMode { |
| 176 | epochs := md.Get("x-duckgres-owner-epoch") |
| 177 | if len(epochs) == 0 { |
| 178 | return nil, status.Error(codes.Unauthenticated, "missing x-duckgres-owner-epoch header") |
| 179 | } |
| 180 | ownerEpoch, err := strconv.ParseInt(epochs[0], 10, 64) |
| 181 | if err != nil { |
| 182 | return nil, status.Error(codes.InvalidArgument, "invalid x-duckgres-owner-epoch header") |
| 183 | } |
| 184 | workerIDs := md.Get("x-duckgres-worker-id") |
| 185 | if len(workerIDs) == 0 { |
| 186 | return nil, status.Error(codes.Unauthenticated, "missing x-duckgres-worker-id header") |
| 187 | } |
| 188 | workerID, err := strconv.Atoi(workerIDs[0]) |
| 189 | if err != nil { |
| 190 | return nil, status.Error(codes.InvalidArgument, "invalid x-duckgres-worker-id header") |
| 191 | } |
| 192 | cpInstanceIDs := md.Get("x-duckgres-cp-instance-id") |
| 193 | if len(cpInstanceIDs) == 0 { |
| 194 | return nil, status.Error(codes.Unauthenticated, "missing x-duckgres-cp-instance-id header") |
| 195 | } |
| 196 | if err := h.pool.validateControlMetadata(server.WorkerControlMetadata{ |
| 197 | WorkerID: workerID, |
| 198 | OwnerEpoch: ownerEpoch, |
| 199 | CPInstanceID: cpInstanceIDs[0], |
| 200 | }); err != nil { |
| 201 | return nil, status.Errorf(codes.FailedPrecondition, "stale worker owner: %v", err) |
| 202 | } |
| 203 | } |
| 204 | session.lastUsed.Store(time.Now().UnixNano()) |
| 205 | |
| 206 | return session, nil |
| 207 | } |
| 208 | |
| 209 | // Custom action handlers (called via customActionServer.DoAction) |
| 210 |