Skip to content

Commit

Permalink
Encapsulate observations query in sql client. (#1486)
Browse files Browse the repository at this point in the history
  • Loading branch information
keyurva authored Jan 6, 2025
1 parent c8d8fd7 commit f59b77f
Show file tree
Hide file tree
Showing 11 changed files with 230 additions and 77 deletions.
4 changes: 2 additions & 2 deletions cmd/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -208,7 +208,7 @@ func main() {
if err != nil {
log.Fatalf("Cannot open sqlite database from: %s: %v", *sqlitePath, err)
}
sqlClient.DB = client.DB
sqlClient.UseConnections(client)
defer sqlClient.Close()
}

Expand All @@ -220,7 +220,7 @@ func main() {
if err != nil {
log.Fatalf("Cannot open cloud sql database from %s: %v", *cloudSQLInstance, err)
}
sqlClient.DB = client.DB
sqlClient.UseConnections(client)
defer sqlClient.Close()
}
}
Expand Down
7 changes: 6 additions & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ require (
cloud.google.com/go/pubsub v1.33.0
cloud.google.com/go/secretmanager v1.11.1
cloud.google.com/go/storage v1.33.0
github.com/go-sql-driver/mysql v1.7.1
github.com/go-sql-driver/mysql v1.8.1
github.com/go-test/deep v1.0.7
github.com/golang/geo v0.0.0-20230421003525-6adc56603217
github.com/google/go-cmp v0.5.9
Expand All @@ -27,6 +27,11 @@ require (
modernc.org/sqlite v1.31.1
)

require (
filippo.io/edwards25519 v1.1.0 // indirect
github.com/jmoiron/sqlx v1.4.0 // indirect
)

require (
cloud.google.com/go v0.110.8 // indirect
cloud.google.com/go/compute v1.23.0 // indirect
Expand Down
8 changes: 8 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@ cloud.google.com/go/spanner v1.47.0 h1:aqiMP8dhsEXgn9K5EZBWxPG7dxIiyM2VaikqeU4it
cloud.google.com/go/spanner v1.47.0/go.mod h1:IXsJwVW2j4UKs0eYDqodab6HgGuA1bViSqW4uH9lfUI=
cloud.google.com/go/storage v1.33.0 h1:PVrDOkIC8qQVa1P3SXGpQvfuJhN2LHOoyZvWs8D2X5M=
cloud.google.com/go/storage v1.33.0/go.mod h1:Hhh/dogNRGca7IWv1RC2YqEn0c0G77ctA/OxflYkiD8=
filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA=
filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4=
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
github.com/JohnCGriffin/overflow v0.0.0-20211019200055-46fa312c352c h1:RGWPOewvKIROun94nF7v2cua9qP+thov/7M50KEoeSU=
github.com/andybalholm/brotli v1.0.5 h1:8uQZIdzKmjc/iuPu7O2ioW48L81FgatrcpfFmiq/cCs=
Expand Down Expand Up @@ -62,6 +64,8 @@ github.com/envoyproxy/protoc-gen-validate v1.0.2 h1:QkIBuU5k+x7/QXPvPPnWXWlCdaBF
github.com/envoyproxy/protoc-gen-validate v1.0.2/go.mod h1:GpiZQP3dDbg4JouG/NNS7QWXpgx6x8QiMKdmN72jogE=
github.com/go-sql-driver/mysql v1.7.1 h1:lUIinVbN1DY0xBg0eMOzmmtGoHwWBbvnWubQUrtU8EI=
github.com/go-sql-driver/mysql v1.7.1/go.mod h1:OXbVy3sEdcQ2Doequ6Z5BW6fXNQTmx+9S1MCJN5yJMI=
github.com/go-sql-driver/mysql v1.8.1 h1:LedoTUt/eveggdHS9qUFC1EFSa8bU2+1pZjSRpvNJ1Y=
github.com/go-sql-driver/mysql v1.8.1/go.mod h1:wEBSXgmK//2ZFJyE+qWnIsVGmvmEKlqwuVSjsCm7DZg=
github.com/go-test/deep v1.0.7 h1:/VSMRlnY/JSyqxQUzQLKVMAskpY/NZKFA5j2P+0pP2M=
github.com/go-test/deep v1.0.7/go.mod h1:QV8Hv/iy04NyLBxAdO9njL0iVPN1S4d/A3NVv1V36o8=
github.com/goccy/go-json v0.10.2 h1:CrxCmQqYDkv1z7lO7Wbh2HN93uovUHgrECaO5ZrCXAU=
Expand Down Expand Up @@ -129,6 +133,8 @@ github.com/jackc/pgproto3/v2 v2.3.2 h1:7eY55bdBeCz1F2fTzSz69QC+pG46jYq9/jtSPiJ5n
github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk=
github.com/jackc/pgtype v1.14.0 h1:y+xUdabmyMkJLyApYuPj38mW+aAIqCe5uuBB51rH3Vw=
github.com/jackc/pgx/v4 v4.18.1 h1:YP7G1KABtKpB5IHrO9vYwSrCOhs7p3uqhvhhQBptya0=
github.com/jmoiron/sqlx v1.4.0 h1:1PLqN7S1UYp5t4SrVVnt4nUVNemrDAtxlulVe+Qgm3o=
github.com/jmoiron/sqlx v1.4.0/go.mod h1:ZrZ7UsYB/weZdl2Bxg6jCRO9c3YHl8r3ahlKmRT4JLY=
github.com/klauspost/asmfmt v1.3.2 h1:4Ri7ox3EwapiOjCki+hw14RyKk201CN4rzyCJRFLpK4=
github.com/klauspost/asmfmt v1.3.2/go.mod h1:AG8TuvYojzulgDAMCnYn50l/5QV3Bs/tp6j0HLHbNSE=
github.com/klauspost/compress v1.17.0 h1:Rnbp4K9EjcDuVuHtd0dgA4qNuv9yKDYKK1ulpJwgrqM=
Expand All @@ -141,8 +147,10 @@ github.com/kr/pretty v0.3.0 h1:WgNl7dwNpEZ6jJ9k1snq4pZsg7DOEN8hP9Xw0Tsjwk0=
github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
github.com/lib/pq v1.10.9/go.mod h1:AlVN5x4E4T544tWzH6hKfbfQvm3HdbOxrmggDNAPY9o=
github.com/mattn/go-isatty v0.0.20 h1:xfD0iDuEKnDkl03q4limB+vH+GxLEtL/jb4xVJSWWEY=
github.com/mattn/go-isatty v0.0.20/go.mod h1:W+V8PltTTMOvKvAeJH7IuucS94S2C6jfK/D7dTCTo3Y=
github.com/mattn/go-sqlite3 v1.14.22/go.mod h1:Uh1q+B4BYcTPb+yiD3kU8Ct7aC0hY9fxUwlHK0RXw+Y=
github.com/microsoft/go-mssqldb v1.6.0 h1:mM3gYdVwEPFrlg/Dvr2DNVEgYFG7L42l+dGc67NNNpc=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8/go.mod h1:mC1jAcsrzbxHt8iiaC+zU4b1ylILSosueou12R++wfY=
Expand Down
2 changes: 1 addition & 1 deletion internal/server/v2/observation/contained_in.go
Original file line number Diff line number Diff line change
Expand Up @@ -223,7 +223,7 @@ func FetchContainedIn(
}
directResp, err := sqlquery.GetObservations(
ctx,
store.SQLClient.DB,
&store.SQLClient,
sqlProvenances,
variables,
childPlaces,
Expand Down
2 changes: 1 addition & 1 deletion internal/server/v2/observation/direct.go
Original file line number Diff line number Diff line change
Expand Up @@ -89,7 +89,7 @@ func FetchDirect(
}
sqlObservation, err := sqlquery.GetObservations(
ctx,
store.SQLClient.DB,
&store.SQLClient,
sqlProvenances,
variables,
entities,
Expand Down
31 changes: 24 additions & 7 deletions internal/sqldb/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (

"cloud.google.com/go/cloudsqlconn"
"github.com/go-sql-driver/mysql"
"github.com/jmoiron/sqlx"

_ "modernc.org/sqlite" // import the sqlite driver
)
Expand All @@ -50,35 +51,51 @@ const (
type SQLClient struct {
// Direct access to the DB will be disabled eventually (by making it private).
// It's exposed right now so we can incrementally encapsulate all SQL functionality in the client before disabling it.
DB *sql.DB
DB *sql.DB
dbx *sqlx.DB
}

// UseConnections uses connections from the src client to this client.
// This method is to workaround the fact that we currently need to maintain the client by value in the store but connections by reference.
// This method should be removed once the store maintains the client by reference.
func (sc *SQLClient) UseConnections(src *SQLClient) {
sc.DB = src.DB
sc.dbx = src.dbx
}

// Close closes the underlying database connection
func (sc *SQLClient) Close() error {
if sc.DB != nil {
return sc.DB.Close()
if sc.dbx != nil {
return sc.dbx.Close()
}
return nil
}

func IsConnected(sqlClient *SQLClient) bool {
return sqlClient != nil && sqlClient.dbx != nil
}

func NewSQLiteClient(sqlitePath string) (*SQLClient, error) {
db, err := newSQLiteConnection(sqlitePath)
if err != nil {
return nil, err
}
return newSQLClient(db), nil
return newSQLClient(db, sqliteDriver), nil
}

func NewCloudSQLClient(instanceName string) (*SQLClient, error) {
db, err := newCloudSQLConnection(instanceName)
if err != nil {
return nil, err
}
return newSQLClient(db), nil
return newSQLClient(db, mysqlDriver), nil
}

func newSQLClient(db *sql.DB) *SQLClient {
return &SQLClient{DB: db}
func newSQLClient(db *sql.DB, driver string) *SQLClient {
return &SQLClient{
DB: db,
dbx: sqlx.NewDb(db, driver),
}
}

func newSQLiteConnection(dbPath string) (*sql.DB, error) {
Expand Down
30 changes: 30 additions & 0 deletions internal/sqldb/model.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
// Copyright 2024 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Model objects related to the SQL database.
package sqldb

// Observation struct to represent a row in the observations table
type Observation struct {
Entity string `db:"entity"`
Variable string `db:"variable"`
Date string `db:"date"`
Value float64 `db:"value"`
Provenance string `db:"provenance"`
Unit string `db:"unit"`
ScalingFactor string `db:"scaling_factor"`
MeasurementMethod string `db:"measurement_method"`
ObservationPeriod string `db:"observation_period"`
Properties string `db:"properties"`
}
96 changes: 96 additions & 0 deletions internal/sqldb/query.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,96 @@
// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Queries executed by the SQLClient.
package sqldb

import (
"context"

"github.com/jmoiron/sqlx"
)

const (
latestDate = "LATEST"
)

// GetObservations retrieves observations from SQL given a list of variables and entities and a date.
func (sc *SQLClient) GetObservations(ctx context.Context, variables []string, entities []string, date string) ([]*Observation, error) {
var observations []*Observation
if len(variables) == 0 || len(entities) == 0 {
return observations, nil
}

var stmt statement

switch {
case date != "" && date != latestDate:
stmt = statement{
query: statements.getObsByVariableEntityAndDate,
args: map[string]interface{}{
"variables": variables,
"entities": entities,
"date": date,
},
}
default:
stmt = statement{
query: statements.getObsByVariableAndEntity,
args: map[string]interface{}{
"variables": variables,
"entities": entities,
},
}
}

err := sc.queryAndCollect(
ctx,
stmt,
&observations,
)
if err != nil {
return nil, err
}

return observations, nil
}

func (sc *SQLClient) queryAndCollect(
ctx context.Context,
stmt statement,
dest interface{},
) error {
// Convert named query and maps of args to placeholder query and list of args.
query, args, err := sqlx.Named(stmt.query, stmt.args)
if err != nil {
return err
}

// Expand slice values.
query, args, err = sqlx.In(query, args...)
if err != nil {
return err
}

// Transform query to the driver's placeholder type.
query = sc.dbx.Rebind(query)

return sc.dbx.SelectContext(ctx, dest, query, args...)
}

// statement struct includes the sql query and named args used to execute a sql query.
type statement struct {
query string
args map[string]interface{}
}
Loading

0 comments on commit f59b77f

Please sign in to comment.