Add NetworkPolicy Recommendation on Snowflake Backend #137

dreamtalen · 2022-10-28T21:23:56Z

In this commit, we add the NetworkPolicy Recommendation Application on the Snowflake backend. NetworkPolicy Recommendation is implemented as Snowflake UDFs and could be running on Snowflake warehouses.

We add two commands in theia-sf CLI tool: create-udfs and policy-recommendation. create-udfs is used to upload and create UDFs on Snowflake database then user can call policy-recommendation command to run NetworkPolicy Recommendation application.

NetworkPolicy Recommendation UDFs are written in Python and stored under udf/ directory.

Signed-off-by: Yongming Ding dyongming@vmware.com

antoninbas

some initial high-level comments, let me know what you think:

kubernetes.zip should not be checked-in. Can we version it (i.e., add the K8s version to the file name) and upload it to antrea.io?
I am wondering if we really need the "create-udfs" command here. I think it will be easier if everything can be configured by the "onboard" command at first.
I think the user should not have to provide --udf-dir and I think that the CLI binary should also be completely self-contained (i.e., not assume that the Github repository has been cloned). As a result, I believe that we should embed the UDF code into the antrea-sf binary (using go:embed). At the end of "onboarding", the CLI tool will upload the code to the Snowflake stage and create the function object.

dreamtalen · 2022-10-28T22:29:36Z

some initial high-level comments, let me know what you think:

kubernetes.zip should not be checked-in. Can we version it (i.e., add the K8s version to the file name) and upload it to antrea.io?

I am wondering if we really need the "create-udfs" command here. I think it will be easier if everything can be configured by the "onboard" command at first.

I think the user should not have to provide --udf-dir and I think that the CLI binary should also be completely self-contained (i.e., not assume that the Github repository has been cloned). As a result, I believe that we should embed the UDF code into the antrea-sf binary (using go:embed). At the end of "onboarding", the CLI tool will upload the code to the Snowflake stage and create the function object.

Thanks Antonin.
1: sure.
2 & 3: sounds good to me, including creating UDFs in onboard could save a step for users. Just wondering if create-udfs command is still valuable, if users want to create or modify udfs by themselves? Or we probably tell users to run onboard again in that case.

antoninbas · 2022-10-28T22:36:11Z

Just wondering if create-udfs command is still valuable, if users want to create or modify udfs by themselves? Or we probably tell users to run onboard again in that case.

If they modify a UDF (or if we modify a UDF across releases), the idea is to run onboard again (the antrea-sf binary will have changed at that point).
Custom UDFs would be a pretty advanced use case, so I would expect users who want to create their own UDFs to be able to use the Snowflake API and do everything themselves. That being said, we can consider introducing the create-udfs command later on if it is appropriate. Re-building and running onboard again should also work for that case as you pointed out.
I want to keep things as simple as possible for the "standard" use case: for each new Theia release, users can download antrea-sf and run antrea-sf onboard again to upgrade to the latest stack and functionalities, with no breakage or data loss.

snowflake/README.md

snowflake/cmd/policyRecommendation.go

antoninbas · 2022-11-29T19:36:42Z

snowflake/cmd/policyRecommendation.go

+		// stackName, stateBackendURL, secretsProviderURL, region, workdir are not provided here
+		// because we only uses snowflake client in this command.
+		mgr := infra.NewManager(logger, "", "", "", "", warehouseName, "", verbose)
+		rows, err := mgr.RunUdf(ctx, query, databaseName)


I really think you should find a way to avoid using the infra manager here. Does RunUdf really need to be a method of the manager? You could use the Snowflake client directly.

Agree, I was thinking reuseing RunUdf function in the future for other udf commands. Moved it under infra/udfs package now.

antoninbas · 2022-11-29T19:40:02Z

snowflake/pkg/utils/utils.go

+)
+
+// Download a file from the given url to the current directory
+func DownloadFile(url string, filename string) error {


I think this function should use a Context object. Maybe do some unification with

theia/snowflake/pkg/infra/manager.go

Line 92 in 5a15bb8

func downloadAndUntar(ctx context.Context, logger logr.Logger, url string, dir string) error {

. You could improve this function and use it as part of the downloadAndUntar implementation?

Sure, merged it into downloadAndUntar function

antoninbas · 2022-11-29T19:43:51Z

snowflake/pkg/utils/timestamps/timestamps.go

+	"time"
+)
+
+func ParseTimestamp(t string, now time.Time, defaultT ...time.Time) (string, error) {


looking at usage of this function, I am not seeing defaultT being used anywhere. IMO, using a variadic function for this (passing an optional default value to the function) is also a bit strange. Something like this makes more sense IMO:

func ParseTimestamp(t string, now time.Time) (string, error) { return ParseTimestampWithDefault(t, now, now) } func ParseTimestampWithDefault(t string, now time.Time, defaultT time.Time) (string, error) { // ... }

Thanks for pointing out, I was copying this function from our gitlab repo. Remove the defaultT argument for now since we don't have any usecases.

antoninbas · 2022-11-30T23:51:17Z

snowflake/pkg/infra/udfs.go

@@ -0,0 +1,79 @@
+// Copyright 2022 Antrea Authors.


I am not sure this needs to be under pkg/infra at all. How about pkg/udfs?

Sounds good, moved it there.

antoninbas · 2022-11-30T23:51:40Z

snowflake/main.go

+
+// Embed the udfs directory here because go:embed doesn't support embeding in subpackages
+
+//go:embed udf/*


nit: maybe we could rename the directory to udfs/?

antoninbas · 2022-11-30T23:52:40Z

snowflake/main.go


 func main() {
+	infra.UdfFs = udfFs


I'd still avoid doing something like this. It can be confusing to change the value of a global variable which is in a different package. Maybe check out what I did for database migrations as a reference: https://github.com/antrea-io/theia/tree/main/snowflake/database

Thanks Antonin, I created udfs.go in a separate package like database migrations and also found some utils function like WriteMigrationsToDisk could be shared here, changed it to WriteEmbedDirToDisk and moved to utils package.

antoninbas · 2022-11-30T23:53:35Z

snowflake/pkg/utils/utils.go

+	"github.com/go-logr/logr"
+)
+
+func DownloadAndUntar(ctx context.Context, logger logr.Logger, url string, dir string, filename string, untar bool) error {


probably better to have 2 different functions: Download and DownloadAndUntar, and have the second one call the first one for its implementation.

Make senses, addressed.

antoninbas

some small comments, otherwise LGTM

antoninbas · 2022-12-02T02:39:23Z

snowflake/pkg/snowflake/snowflake.go

+	if !result {
+		_, err := c.db.ExecContext(multi_statement_context, query)
+		return nil, err
+	} else {
+		rows, err := c.db.QueryContext(multi_statement_context, query)
+		return rows, err
+	}


do you think we could have 2 different functions for the 2 different cases?

Sounds good to me, changed.

antoninbas · 2022-12-02T02:39:53Z

snowflake/pkg/snowflake/snowflake.go

+
+func (c *client) UseSchema(ctx context.Context, name string) error {
+	query := fmt.Sprintf("USE SCHEMA %s", name)
+	c.logger.Info("Snowflake query", "query", query)


the verbosity values for the logs don't seem consistent (see V(2) above)

Thanks, addressed.

antoninbas · 2022-12-02T02:43:09Z

snowflake/Makefile

@@ -5,6 +5,7 @@ all: bin

 .PHONY: bin
 bin:
+	make -C udfs/udfs/


IMO it doesn't make sense to have this in the bin target
You should define a new target and add it as a requirement for the all target so that it is built by default.

Right, thanks for pointing out.

antoninbas · 2022-12-02T02:44:14Z

snowflake/pkg/infra/manager.go

@@ -37,6 +36,8 @@ import (

 	"antrea.io/theia/snowflake/database"
 	sf "antrea.io/theia/snowflake/pkg/snowflake"
+	utils "antrea.io/theia/snowflake/pkg/utils"


nit: renaming the package doesn't seem necessary here

Thanks, addressed.

antoninbas · 2022-12-02T02:46:33Z

snowflake/pkg/utils/utils.go

+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package utils


nit: Quan pointed out in the past that utils was too generic for a package name. So I wonder if we could move these functions to pkg/utils/file/file.go (package can then be imported as fileutils).

Sounds good, changed.

antoninbas · 2022-12-02T02:48:13Z

snowflake/pkg/utils/utils.go

+	return nil
+}
+
+func WriteEmbedDirToDisk(ctx context.Context, logger logr.Logger, fsys fs.FS, embedPath string, dest string) error {


given that there is nothing specific to embed in this function implementation, I suggest you rename it to writeFSDirToDisk and rename embedPath to fsysPath.

Thanks, addressed.

In this commit, we add the NetworkPolicy Recommendation Application on the Snowflake backend. NetworkPolicy Recommendation is implemented as Snowflake UDFs and could be running on Snowflake warehouses. Users could start it using command policy-recommendation in theia-sf CLI tool. NetworkPolicy Recommendation UDFs are written in Python and stored under udf/ directory. Signed-off-by: Yongming Ding <dyongming@vmware.com>

Signed-off-by: Yongming Ding <dyongming@vmware.com>

antoninbas

when you are done with code changes, please make sure you run the whole thing end-to-end before we merge

snowflake/pkg/snowflake/snowflake.go

antoninbas · 2022-12-02T21:45:59Z

snowflake/pkg/utils/timestamps/timestamps.go

+	"time"
+)
+
+func ParseTimestamp(t string, now time.Time) (string, error) {


maybe we could add some table-driven unit tests for this function?

sounds good, added.

Signed-off-by: Yongming Ding <dyongming@vmware.com>

dreamtalen · 2022-12-06T00:37:39Z

when you are done with code changes, please make sure you run the whole thing end-to-end before we merge

Sure, I just tried e2e with Antrea Flow exporter & Aggregator locally.

antoninbas

LGTM

yanjunz97 · 2022-12-06T18:34:45Z

snowflake/cmd/policyRecommendation.go

+	policyRecommendationCmd.Flags().Int("isolationMethod", 1, `Network isolation preference. Currently we have 3 options:
+1: Recommending allow ANP/ACNP policies, with default deny rules only on Pods which have an allow rule applied
+2: Recommending allow ANP/ACNP policies, with default deny rules for whole cluster
+3: Recommending allow K8s NetworkPolicies only`)


Shall we use String instead of Int for isolationMethod to make it more readable as we do for policy-type in ClickHouse?

Make sense, addressed.

yanjunz97 · 2022-12-06T18:52:36Z

snowflake/cmd/policyRecommendation.go

+		fmt.Fprintf(&queryBuilder, `AND
+  flowEndSeconds >= '%s'
+`, endTime)


I might be wrong, but could you just check if this should be flowEndSeconds < '%s'?

Good catch! You're right

yanjunz97 · 2022-12-06T19:17:12Z

snowflake/pkg/infra/manager.go

@@ -488,3 +416,118 @@ func (m *Manager) Offboard(ctx context.Context) error {
 	_, err := m.run(ctx, true)
 	return err
 }
+
+func createUdfs(ctx context.Context, logger logr.Logger, databaseName string, warehouseName string, workdir string) error {
+	logger.Info("creating UDFs")


Just curious about whether we have a convention on whether to capital the first letter for logger.Info? I see both in code.

I prefer capitaling the first letter, let me fix that.

Signed-off-by: Yongming Ding <dyongming@vmware.com>

yanjunz97

LGTM

dreamtalen · 2022-12-06T21:19:47Z

/theia-test-e2e

dreamtalen added this to the Theia v0.4 release milestone Oct 28, 2022

dreamtalen force-pushed the sf-pr branch 4 times, most recently from fe283bf to de60b62 Compare October 28, 2022 22:05

antoninbas reviewed Oct 28, 2022

View reviewed changes

dreamtalen force-pushed the sf-pr branch 3 times, most recently from 8213ce4 to 8f45932 Compare November 14, 2022 22:34

dreamtalen requested review from heanlan and yanjunz97 November 17, 2022 19:01

antoninbas reviewed Nov 29, 2022

View reviewed changes

antoninbas reviewed Nov 30, 2022

View reviewed changes

dreamtalen force-pushed the sf-pr branch 2 times, most recently from 2b32bc9 to 10e73e5 Compare December 2, 2022 01:05

antoninbas reviewed Dec 2, 2022

View reviewed changes

Yongming Ding added 3 commits December 2, 2022 13:29

Address comments

ee68013

Signed-off-by: Yongming Ding <dyongming@vmware.com>

Address 2nd round comments

6d3889a

Signed-off-by: Yongming Ding <dyongming@vmware.com>

dreamtalen force-pushed the sf-pr branch from adc238c to 61ee524 Compare December 2, 2022 21:44

antoninbas reviewed Dec 2, 2022

View reviewed changes

Address 3rd round comments & add UTs

0756809

Signed-off-by: Yongming Ding <dyongming@vmware.com>

dreamtalen force-pushed the sf-pr branch from 61ee524 to 0756809 Compare December 6, 2022 00:35

antoninbas approved these changes Dec 6, 2022

View reviewed changes

yanjunz97 reviewed Dec 6, 2022

View reviewed changes

Address comments from yanjun

fe1ac86

Signed-off-by: Yongming Ding <dyongming@vmware.com>

yanjunz97 approved these changes Dec 6, 2022

View reviewed changes

dreamtalen merged commit af9d0fe into antrea-io:main Dec 6, 2022


		// Embed the udfs directory here because go:embed doesn't support embeding in subpackages

		//go:embed udf/*

Add NetworkPolicy Recommendation on Snowflake Backend #137

Add NetworkPolicy Recommendation on Snowflake Backend #137

Conversation

dreamtalen commented Oct 28, 2022

antoninbas left a comment

Choose a reason for hiding this comment

dreamtalen commented Oct 28, 2022

antoninbas commented Oct 28, 2022 • 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

antoninbas Nov 29, 2022 • 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

antoninbas left a 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

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

antoninbas left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dreamtalen commented Dec 6, 2022

antoninbas left a 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

Choose a reason for hiding this comment

yanjunz97 left a comment

Choose a reason for hiding this comment

dreamtalen commented Dec 6, 2022

antoninbas commented Oct 28, 2022 •

edited

Loading

antoninbas Nov 29, 2022 •

edited

Loading